From fc542abb98159d77c2da864e611c21c47b5cdb94 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 07:14:55 +0000
Subject: [PATCH 01/51] [BDE] Add Block Diffusion Engine selection via config
 field

Introduce the Block Diffusion Engine (BDE), a DiffusionEngine subclass that
will own engine-level KV cache management for autoregressive / chunked
world-model diffusion (DreamZero and the AR-DiT family). This first change is
the engine-selection skeleton; no KV behavior yet.

- Add `engine_backend` field to OmniDiffusionConfig.
- DiffusionEngine.make_engine now dispatches via resolve_engine_class,
  mirroring DiffusionExecutor.get_class: accepts "default"/"bde", a
  DiffusionEngine subclass, or an import-path string. Both existing
  make_engine call sites pick this up unchanged.
- Add vllm_omni/bde with BDEEngine (behavior-preserving subclass for now;
  the KVCacheManager lifecycle is added in later phases).
- Route DreamZero (DreamZeroPipeline) to BDE by default at the config layer
  via default_engine_backend_for_model(), overridable by an explicit
  engine_backend. Other models are unaffected.
- Tests for the dispatcher and DreamZero defaulting (15 tests, all passing).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 README.md                                     |    1 -
 docs/models/supported_models.md               |    6 +-
 docs/serving/speech_api.md                    |    7 -
 docs/serving/videos_api.md                    |    7 +-
 .../examples/online_serving/image_to_video.md |   18 +-
 .../examples/online_serving/text_to_video.md  |   18 +-
 .../dreamzero/client_schedule.py              |  129 -
 .../dreamzero/export_prediction_video.py      |  136 +-
 .../online_serving/image_to_video/README.md   |   11 +-
 .../online_serving/text_to_video/README.md    |   11 +-
 recipes/NVIDIA/GR00T-N1.7.md                  |  101 -
 recipes/README.md                             |   11 +-
 recipes/mistralai/Voxtral-TTS.md              |   42 -
 requirements/common.txt                       |    2 -
 tests/bde/__init__.py                         |    1 +
 tests/bde/test_engine_routing.py              |   96 +
 tests/config/test_server_settings.py          |   84 -
 tests/diffusion/cache/test_cache_backends.py  |  115 +-
 tests/diffusion/cache/test_cache_dit.py       |   47 +-
 tests/diffusion/cache/test_stepcache.py       |  116 -
 tests/diffusion/models/gr00t/__init__.py      |    0
 tests/diffusion/models/gr00t/test_pipeline.py |  141 -
 .../models/ltx2/test_ltx2_3_pipeline.py       |   17 +-
 tests/diffusion/quantization/test_hsdp_fp8.py |  147 -
 tests/diffusion/test_diffusion_worker.py      |   25 +-
 tests/dreamzero/test_crossattn_cache.py       |  114 -
 tests/dreamzero/test_fused_qk_rms_norm.py     |  112 -
 tests/dreamzero/test_qkv_fusion.py            |  117 -
 tests/e2e/offline_inference/test_moss_tts.py  |  333 ++-
 .../test_moss_tts_realtime.py                 |  290 --
 .../test_dreamzero_expansion.py               |    2 +
 tests/e2e/online_serving/test_gr00t_openpi.py |  163 --
 .../openai_api/test_openpi_connection.py      |   41 -
 .../openai_api/test_storage_manager.py        |   94 -
 .../openai_api/test_video_server.py           |   51 +-
 tests/gr00t/__init__.py                       |    2 -
 tests/gr00t/openpi_client_helper.py           |  183 --
 .../models/moss_tts/test_per_request_codes.py |  176 --
 tests/test_omni_request.py                    |   89 -
 vllm_omni/bde/__init__.py                     |   12 +
 vllm_omni/bde/engine.py                       |   30 +
 vllm_omni/config/pipeline_registry.py         |    4 -
 vllm_omni/config/server_settings.py           |   58 -
 vllm_omni/deploy/Gr00tN1d7.yaml               |   41 -
 vllm_omni/deploy/dreamzero.yaml               |    7 +-
 vllm_omni/deploy/dreamzero_tp1_cfg2.yaml      |    7 +-
 vllm_omni/diffusion/cache/base.py             |    2 -
 .../diffusion/cache/cache_dit_backend.py      | 2401 ++++++++++++-----
 vllm_omni/diffusion/cache/selector.py         |   13 +-
 .../diffusion/cache/stepcache/__init__.py     |   66 -
 .../diffusion/cache/stepcache/backend.py      |  129 -
 vllm_omni/diffusion/cache/stepcache/config.py |   51 -
 vllm_omni/diffusion/cache/stepcache/state.py  |   51 -
 vllm_omni/diffusion/data.py                   |   44 +-
 vllm_omni/diffusion/diffusion_engine.py       |   50 +-
 vllm_omni/diffusion/distributed/hsdp.py       |    7 +-
 .../model_loader/diffusers_loader.py          |   27 +-
 .../models/bagel/bagel_transformer.py         |    9 -
 .../models/cosmos3/transformer_cosmos3.py     |   12 -
 .../diffusion/models/dreamid_omni/fusion.py   |    9 -
 .../models/dreamzero/causal_wan_model.py      |  127 +-
 .../models/dreamzero/pipeline_dreamzero.py    |  293 +-
 .../models/dreamzero/state_dreamzero.py       |  111 +-
 .../models/dreamzero/video_export_worker.py   |   25 +-
 .../dreamzero/wan_vae_feat_cache_patch.py     |   92 -
 .../ernie_image/ernie_image_transformer.py    |    8 -
 .../diffusion/models/flux/flux_transformer.py |    9 -
 .../models/flux2/flux2_transformer.py         |    9 -
 .../flux2_klein/flux2_klein_transformer.py    |    9 -
 .../models/glm_image/glm_image_transformer.py |   10 -
 vllm_omni/diffusion/models/gr00t/__init__.py  |    6 -
 .../models/gr00t/configs/__init__.py          |    8 -
 .../gr00t/configs/embodiment_configs.py       |   12 -
 .../models/gr00t/configs/gr00t_n1d7.py        |  109 -
 .../diffusion/models/gr00t/dataio/__init__.py |    2 -
 .../models/gr00t/dataio/embodiment_tags.py    |  191 --
 .../gr00t/dataio/state_action/__init__.py     |    2 -
 .../dataio/state_action/action_chunking.py    |  654 -----
 .../models/gr00t/dataio/state_action/pose.py  |  709 -----
 .../state_action/state_action_processor.py    |  643 -----
 .../diffusion/models/gr00t/dataio/types.py    |  112 -
 .../diffusion/models/gr00t/dataio/utils.py    |  187 --
 .../models/gr00t/modeling/__init__.py         |    7 -
 .../models/gr00t/modeling/gr00t_n1d7.py       |  491 ----
 .../models/gr00t/modeling/modules/__init__.py |    2 -
 .../models/gr00t/modeling/modules/dit.py      |  344 ---
 .../modules/embodiment_conditioned_mlp.py     |  178 --
 .../gr00t/modeling/processing_gr00t_n1d7.py   |  555 ----
 .../diffusion/models/gr00t/pipeline_gr00t.py  |  148 -
 vllm_omni/diffusion/models/gr00t/policy.py    |  434 ---
 .../models/helios/helios_transformer.py       |    9 -
 .../pipeline_hidream_i1_image.py              | 1059 ++++++++
 .../hunyuan_image3_transformer.py             |    9 -
 .../hunyuan_video_15_transformer.py           |    9 -
 .../longcat_image_transformer.py              |   10 -
 .../diffusion/models/ltx2/ltx2_transformer.py |   10 -
 .../diffusion/models/sd3/sd3_transformer.py   |    8 -
 .../sensenova_u1/pipeline_sensenova_u1.py     |    2 -
 .../sensenova_u1/sensenova_u1_transformer.py  |   10 -
 .../stable_audio/stable_audio_transformer.py  |    8 -
 vllm_omni/diffusion/quantization/__init__.py  |    2 -
 vllm_omni/diffusion/quantization/hsdp_fp8.py  |  112 -
 vllm_omni/diffusion/registry.py               |    5 -
 vllm_omni/diffusion/utils/flow_matching.py    |   28 -
 .../worker/diffusion_model_runner.py          |   24 +-
 .../diffusion/worker/diffusion_worker.py      |   15 +-
 vllm_omni/engine/async_omni_engine.py         |   10 +-
 vllm_omni/entrypoints/cli/serve.py            |    2 +-
 vllm_omni/entrypoints/openai/api_server.py    |   73 +-
 .../entrypoints/openai/serving_speech.py      |    1 +
 vllm_omni/entrypoints/openai/storage.py       |  170 +-
 .../entrypoints/openai/tts_adapters/base.py   |    3 +
 vllm_omni/entrypoints/openpi/connection.py    |   98 +-
 .../models/cosyvoice3/cosyvoice3.py           |   35 +-
 .../models/cosyvoice3/cosyvoice3_code2wav.py  |   73 +-
 .../models/cosyvoice3/pipeline.py             |    6 +-
 .../model_executor/models/cosyvoice3/utils.py |   22 +
 .../models/covo_audio/config_covo_audio.py    |    1 +
 .../models/covo_audio/token2wav.py            |    5 +-
 .../fish_speech/fish_speech_dac_decoder.py    |   18 +-
 .../models/fish_speech/fish_speech_fast_ar.py |   11 +-
 .../models/fish_speech/prompt_utils.py        |   10 +
 .../models/glm_tts/glm_tts_dit.py             |    4 +
 .../model_executor/models/gr00t/__init__.py   |    2 -
 .../model_executor/models/gr00t/pipeline.py   |   26 -
 .../configuration_higgs_audio_v2.py           |   12 +
 .../higgs_audio_v2/higgs_audio_v2_talker.py   |   35 +-
 .../models/higgs_audio_v3/ref_audio_cache.py  |  111 +
 .../models/ming_flash_omni/projectors.py      |   27 +
 .../models/ming_flash_omni/text_processing.py |   91 +
 .../models/ming_tts/config_ming_tts.py        |    2 +
 .../models/ming_tts/constants.py              |    1 +
 .../models/ming_tts/ming_tts.py               |    4 +
 .../models/ming_tts/ming_tts_llm.py           |   14 +-
 .../models/moss_tts/audio_tokenizer.py        |    2 +-
 .../moss_tts/modeling_moss_tts_talker.py      |  166 +-
 .../configuration_moss_tts_nano.py            |   60 +
 .../moss_tts_nano/modeling_moss_tts_nano.py   |    1 +
 .../models/omnivoice/omnivoice.py             |    3 +-
 .../models/omnivoice/omnivoice_generator.py   |   23 +-
 .../qwen3_tts/configuration_qwen3_tts.py      |   11 +-
 .../models/qwen3_tts/prompt_embeds_builder.py |    8 +-
 .../models/qwen3_tts/qwen3_tts_talker.py      |    7 +
 .../modeling_qwen3_tts_tokenizer_v2.py        |    5 +
 .../qwen3_tts/tokenizer_25hz/vq/core_vq.py    |    5 +
 .../models/step_audio2/step_audio2.py         |    8 +-
 .../step_audio2/step_audio2_constants.py      |   36 +
 .../models/step_audio2/step_audio2_thinker.py |   30 +-
 .../step_audio2/step_audio2_token2wav.py      |  121 +
 .../models/voxcpm2/minicpm4_paged.py          |   19 +-
 .../models/voxcpm2/voxcpm2_talker.py          |   43 +-
 .../models/voxtral_tts/voxtral_tts.py         |    3 +
 .../voxtral_tts_audio_generation.py           |   26 +-
 .../voxtral_tts_audio_tokenizer.py            |    4 +-
 vllm_omni/request.py                          |    7 +-
 .../transformers_utils/processors/ming.py     |    1 +
 vllm_omni/worker/gpu_model_runner.py          |   14 -
 157 files changed, 4319 insertions(+), 10044 deletions(-)
 delete mode 100644 examples/offline_inference/dreamzero/client_schedule.py
 delete mode 100644 recipes/NVIDIA/GR00T-N1.7.md
 create mode 100644 tests/bde/__init__.py
 create mode 100644 tests/bde/test_engine_routing.py
 delete mode 100644 tests/config/test_server_settings.py
 delete mode 100644 tests/diffusion/cache/test_stepcache.py
 delete mode 100644 tests/diffusion/models/gr00t/__init__.py
 delete mode 100644 tests/diffusion/models/gr00t/test_pipeline.py
 delete mode 100644 tests/diffusion/quantization/test_hsdp_fp8.py
 delete mode 100644 tests/dreamzero/test_crossattn_cache.py
 delete mode 100644 tests/dreamzero/test_fused_qk_rms_norm.py
 delete mode 100644 tests/dreamzero/test_qkv_fusion.py
 delete mode 100644 tests/e2e/offline_inference/test_moss_tts_realtime.py
 delete mode 100644 tests/e2e/online_serving/test_gr00t_openpi.py
 delete mode 100644 tests/entrypoints/openai_api/test_storage_manager.py
 delete mode 100644 tests/gr00t/__init__.py
 delete mode 100644 tests/gr00t/openpi_client_helper.py
 delete mode 100644 tests/model_executor/models/moss_tts/test_per_request_codes.py
 delete mode 100644 tests/test_omni_request.py
 create mode 100644 vllm_omni/bde/__init__.py
 create mode 100644 vllm_omni/bde/engine.py
 delete mode 100644 vllm_omni/config/server_settings.py
 delete mode 100644 vllm_omni/deploy/Gr00tN1d7.yaml
 delete mode 100644 vllm_omni/diffusion/cache/stepcache/__init__.py
 delete mode 100644 vllm_omni/diffusion/cache/stepcache/backend.py
 delete mode 100644 vllm_omni/diffusion/cache/stepcache/config.py
 delete mode 100644 vllm_omni/diffusion/cache/stepcache/state.py
 delete mode 100644 vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/__init__.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/configs/__init__.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/dataio/__init__.py
 delete mode 100755 vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/dataio/types.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/dataio/utils.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/modeling/__init__.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py
 delete mode 100644 vllm_omni/diffusion/models/gr00t/policy.py
 create mode 100644 vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py
 delete mode 100644 vllm_omni/diffusion/quantization/__init__.py
 delete mode 100644 vllm_omni/diffusion/quantization/hsdp_fp8.py
 delete mode 100644 vllm_omni/diffusion/utils/flow_matching.py
 delete mode 100644 vllm_omni/model_executor/models/gr00t/__init__.py
 delete mode 100644 vllm_omni/model_executor/models/gr00t/pipeline.py
 create mode 100644 vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py
 create mode 100644 vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py

diff --git a/README.md b/README.md
index e300b6b8d1e..cf1658b4574 100644
--- a/README.md
+++ b/README.md
@@ -68,7 +68,6 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
 - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
 - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
-- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
 
 ## Contributing
 
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index c224c1a6a2f..119f3bf683c 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -6,8 +6,6 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
 
 If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
 
-For deployment recipes (TTS and multimodal), see [recipes.vllm.ai](https://recipes.vllm.ai) and the in-repo [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes) directory.
-
 ## List of Supported Models
 
 <style>
@@ -24,7 +22,6 @@ th {
 | `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
 | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
-| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | |
 | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -66,14 +63,13 @@ th {
 | `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
 | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
 | `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
 | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
 | `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
 | `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
-| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
 | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
 | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md
index b70c282979f..935deba4f2d 100644
--- a/docs/serving/speech_api.md
+++ b/docs/serving/speech_api.md
@@ -9,13 +9,6 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
 
 See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano.
 
-!!! tip "Deployment recipes"
-    TTS deployment recipes are published at
-    [recipes.vllm.ai](https://recipes.vllm.ai) (e.g.
-    [Qwen3-TTS](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice),
-    [Higgs-Audio v3](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b)).
-    The in-repo runbooks live under [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes).
-
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
 ## Quick Start
diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md
index 95b5dbe7734..4d156bdf0f0 100644
--- a/docs/serving/videos_api.md
+++ b/docs/serving/videos_api.md
@@ -195,16 +195,13 @@ curl -X POST http://localhost:8091/v1/videos/sync \
 
 ## Storage
 
-Set `VLLM_OMNI_SERVER_STORAGE__PATH` to control where asynchronous video outputs are
+Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are
 stored:
 
 ```bash
-export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
 ```
 
-> `VLLM_OMNI_STORAGE_PATH` is deprecated and will be removed in a future release;
-> use `VLLM_OMNI_SERVER_STORAGE__PATH` instead.
-
 ## Model-Specific Examples
 
 For complete text-to-video, image-to-video, and model-specific video-to-video
diff --git a/docs/user_guide/examples/online_serving/image_to_video.md b/docs/user_guide/examples/online_serving/image_to_video.md
index 2dbc57633b8..781f0c2a5ed 100644
--- a/docs/user_guide/examples/online_serving/image_to_video.md
+++ b/docs/user_guide/examples/online_serving/image_to_video.md
@@ -84,26 +84,16 @@ curl -X POST http://localhost:8091/v1/videos/sync \
 Generated video files are stored on local disk by the async video API.
 Local file storage behavior can be controlled via the following environment variables:
 
-- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files
-  (default: `/tmp/storage`)
-- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete/open
-  operations (default: `4`)
-- `VLLM_OMNI_SERVER_STORAGE__FILE_TTL`: optional TTL for generated files in
-  seconds
-- `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL`: optional sweep interval in
-  seconds for enforcing file TTL; defaults to `300` when TTL is set
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
 
 Example:
 
 ```bash
-export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
-export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
-export VLLM_OMNI_SERVER_STORAGE__FILE_TTL=86400
-export VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL=300
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
 ```
 
-See also: [Configuration Options](../../../configuration/README.md)
-
 ## API Calls
 
 ### Method 1: Using curl
diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md
index d4d520d5e76..b918aac19d0 100644
--- a/docs/user_guide/examples/online_serving/text_to_video.md
+++ b/docs/user_guide/examples/online_serving/text_to_video.md
@@ -90,26 +90,16 @@ curl -X POST http://localhost:8091/v1/videos/sync \
 Generated video files are stored on local disk by the async video API.
 Local file storage behavior can be controlled via the following environment variables:
 
-- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files
-  (default: `/tmp/storage`)
-- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete/open
-  operations (default: `4`)
-- `VLLM_OMNI_SERVER_STORAGE__FILE_TTL`: optional TTL for generated files in
-  seconds
-- `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL`: optional sweep interval in
-  seconds for enforcing file TTL; defaults to `300` when TTL is set
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
 
 Example:
 
 ```bash
-export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
-export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
-export VLLM_OMNI_SERVER_STORAGE__FILE_TTL=86400
-export VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL=300
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
 ```
 
-See also: [Configuration Options](../../../configuration/README.md)
-
 ## API Calls
 
 ### Method 1: Using curl
diff --git a/examples/offline_inference/dreamzero/client_schedule.py b/examples/offline_inference/dreamzero/client_schedule.py
deleted file mode 100644
index e4dbad55a03..00000000000
--- a/examples/offline_inference/dreamzero/client_schedule.py
+++ /dev/null
@@ -1,129 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""DreamZero AR client frame schedule aligned with upstream ``test_client_AR.py``."""
-
-from __future__ import annotations
-
-import logging
-from typing import Any
-
-import numpy as np
-
-# Frame schedule constants (matching upstream debug_inference.py / test_client_AR.py).
-ACTION_HORIZON = 24
-DEFAULT_ACTION_DIM = 8
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-DEFAULT_NUM_AR_CHUNKS = 15
-
-CAMERA_FILES = {
-    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
-    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
-    "observation/wrist_image_left": "wrist_image_left.mp4",
-}
-
-
-def build_frame_schedule(
-    total_frames: int,
-    num_chunks: int,
-    *,
-    logger: logging.Logger | None = None,
-) -> list[list[int]]:
-    """Build the frame index schedule for multi-frame chunks.
-
-    Args:
-        total_frames: Number of frames available in each camera video.
-        num_chunks: Number of 4-frame chunks to schedule after the initial frame.
-
-    Returns:
-        A list of frame-index lists. Each inner list has four indices. The
-        returned list may be shorter than ``num_chunks`` when the videos run
-        out of frames (upstream stops early instead of erroring).
-    """
-    chunks: list[list[int]] = []
-    current_frame = 23
-    for _ in range(num_chunks):
-        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
-        if indices[-1] >= total_frames:
-            if logger is not None:
-                logger.info(
-                    "Frame %s >= %s, stopping at %s chunks",
-                    indices[-1],
-                    total_frames,
-                    len(chunks),
-                )
-            break
-        chunks.append(indices)
-        current_frame += ACTION_HORIZON
-    return chunks
-
-
-def make_obs_from_video(
-    camera_frames: dict[str, np.ndarray],
-    frame_indices: list[int],
-    *,
-    prompt: str,
-    session_id: str,
-) -> dict[str, Any]:
-    """Build an observation dict from real video frames.
-
-    For one frame each image key is ``(H, W, 3)``. For four frames each key is
-    ``(4, H, W, 3)``.
-    """
-    obs: dict[str, Any] = {}
-    for camera_key, all_frames in camera_frames.items():
-        selected = all_frames[frame_indices]
-        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
-
-    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
-    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
-    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-    obs["prompt"] = prompt
-    obs["session_id"] = session_id
-    return obs
-
-
-def build_ar_observations(
-    camera_frames: dict[str, np.ndarray],
-    *,
-    prompt: str,
-    session_id: str,
-    num_chunks: int = DEFAULT_NUM_AR_CHUNKS,
-    repeat_chunk_observations: bool = False,
-    logger: logging.Logger | None = None,
-) -> list[dict[str, Any]]:
-    """Build the AR observation sequence used by upstream ``test_client_AR.py``.
-
-    Step 0 sends frame ``[0]``. Each subsequent step sends one 4-frame chunk.
-
-    ``num_chunks`` counts only the 4-frame chunks after the initial frame, not
-    the total number of inferences.
-    """
-    if num_chunks < 0:
-        raise ValueError("num_chunks must be non-negative")
-
-    total_frames = min(frames.shape[0] for frames in camera_frames.values())
-    observations = [
-        make_obs_from_video(
-            camera_frames,
-            [0],
-            prompt=prompt,
-            session_id=session_id,
-        )
-    ]
-
-    chunk_schedule = build_frame_schedule(total_frames, num_chunks, logger=logger)
-    if repeat_chunk_observations and chunk_schedule and len(chunk_schedule) < num_chunks:
-        while len(chunk_schedule) < num_chunks:
-            chunk_schedule.append(chunk_schedule[-1])
-
-    for frame_indices in chunk_schedule:
-        observations.append(
-            make_obs_from_video(
-                camera_frames,
-                frame_indices,
-                prompt=prompt,
-                session_id=session_id,
-            )
-        )
-    return observations
diff --git a/examples/offline_inference/dreamzero/export_prediction_video.py b/examples/offline_inference/dreamzero/export_prediction_video.py
index 4c560418c2b..82f56ff26e5 100644
--- a/examples/offline_inference/dreamzero/export_prediction_video.py
+++ b/examples/offline_inference/dreamzero/export_prediction_video.py
@@ -11,11 +11,6 @@
 import cv2
 import numpy as np
 import torch
-from client_schedule import (
-    CAMERA_FILES,
-    DEFAULT_NUM_AR_CHUNKS,
-    build_ar_observations,
-)
 from PIL import Image
 
 from vllm_omni import Omni
@@ -25,6 +20,13 @@
 WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
 ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
 DEFAULT_SESSION_PREFIX = "dreamzero-export"
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+ACTION_HORIZON = 24
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
 
 
 def _parse_args() -> argparse.Namespace:
@@ -47,17 +49,6 @@ def _parse_args() -> argparse.Namespace:
         default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
     )
     parser.add_argument("--session-id", default=None)
-    parser.add_argument(
-        "--num-chunks",
-        type=int,
-        default=DEFAULT_NUM_AR_CHUNKS,
-        help="Number of 4-frame chunks to send after the initial frame (default: 15).",
-    )
-    parser.add_argument(
-        "--repeat-chunk-observations",
-        action="store_true",
-        help="Repeat the last valid chunk observation when bundled assets run out of frames.",
-    )
     parser.add_argument("--save-input-video", action="store_true")
     parser.add_argument("--save-gif", action="store_true")
     parser.add_argument("--save-actions", action="store_true")
@@ -93,25 +84,57 @@ def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
     return camera_frames
 
 
-def _build_observations(
-    video_dir: Path,
+def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
+    chunks: list[list[int]] = []
+    current_frame = 23
+    for _ in range(num_chunks):
+        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def _make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    *,
     prompt: str,
     session_id: str,
-    *,
-    num_chunks: int,
-    repeat_chunk_observations: bool,
-) -> tuple[dict[str, np.ndarray], list[dict]]:
+) -> dict:
+    obs: dict = {}
+    for camera_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]
+        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
     camera_frames = _load_camera_frames(video_dir)
-    observations = build_ar_observations(
-        camera_frames,
-        prompt=prompt,
-        session_id=session_id,
-        num_chunks=num_chunks,
-        repeat_chunk_observations=repeat_chunk_observations,
-    )
+    total_frames = min(frames.shape[0] for frames in camera_frames.values())
+    chunks = _build_frame_schedule(total_frames, 1)
+    observations = [
+        _make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
+    ]
+    if chunks:
+        observations.append(
+            _make_obs_from_video(
+                camera_frames,
+                chunks[0],
+                prompt=prompt,
+                session_id=session_id,
+            )
+        )
     if len(observations) < 2:
-        raise RuntimeError("Need at least two DreamZero observations to export a prediction video.")
-    return camera_frames, observations
+        raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
+    return camera_frames, observations[:2]
 
 
 def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
@@ -206,26 +229,6 @@ def _run_generation(
     return omni, outputs
 
 
-def _decode_accumulated_with_worker(omni: Omni, session_id: str) -> np.ndarray:
-    """Decode server-side accumulated latents when the pipeline supports it."""
-    stage_client = omni.engine.stage_clients[0]
-    engine = getattr(stage_client, "_engine", None)
-    if engine is None:
-        raise RuntimeError("DreamZero export requires inline diffusion stage access.")
-
-    decoded = engine.executor.collective_rpc(
-        "decode_accumulated_video_latents_to_uint8",
-        args=(session_id,),
-        unique_reply_rank=0,
-        exec_all_ranks=True,
-    )
-    if isinstance(decoded, torch.Tensor):
-        decoded = decoded.numpy()
-    if not isinstance(decoded, np.ndarray):
-        raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
-    return decoded
-
-
 def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
     stage_client = omni.engine.stage_clients[0]
     engine = getattr(stage_client, "_engine", None)
@@ -245,15 +248,6 @@ def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
     return decoded
 
 
-def _decode_video(omni: Omni, session_id: str, outputs: list[OmniRequestOutput]) -> np.ndarray:
-    """Decode full AR rollout by concatenating per-step latents from all outputs."""
-    latent_steps = [_extract_latents(output) for output in outputs]
-    latent_t = [step.shape[2] for step in latent_steps]
-    print(f"Decoding {len(latent_steps)} latent steps, latent_t={latent_t}, total_t={sum(latent_t)}")
-    full_latents = torch.cat(latent_steps, dim=2)
-    return _decode_with_worker(omni, full_latents)
-
-
 def main() -> None:
     args = _parse_args()
     session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
@@ -262,12 +256,6 @@ def main() -> None:
         video_dir=args.video_dir,
         prompt=args.prompt,
         session_id=session_id,
-        num_chunks=args.num_chunks,
-        repeat_chunk_observations=args.repeat_chunk_observations,
-    )
-    print(
-        f"Running {len(observations)} inferences "
-        f"(1 initial + {len(observations) - 1} chunks; requested {args.num_chunks} post-initial chunks)"
     )
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
@@ -285,19 +273,17 @@ def main() -> None:
             deploy_config_path=args.deploy_config,
             observations=observations,
         )
-        frames = _decode_video(omni, session_id, outputs)
+        latent_steps = [_extract_latents(output) for output in outputs]
+        full_latents = torch.cat(latent_steps, dim=2)
+        frames = _decode_with_worker(omni, full_latents)
     finally:
         if omni is not None:
             omni.close()
 
-    expected_frames = 8 * len(observations) + 1
     mp4_path = args.output_dir / f"{args.output_stem}.mp4"
 
     _write_mp4(mp4_path, frames, fps=args.fps)
-    print(
-        f"SAVED_MP4={mp4_path} frames={len(frames)} "
-        f"expected≈{expected_frames} duration≈{len(frames) / max(args.fps, 1):.1f}s"
-    )
+    print(f"SAVED_MP4={mp4_path}")
 
     if args.save_gif:
         gif_path = args.output_dir / f"{args.output_stem}.gif"
@@ -308,10 +294,8 @@ def main() -> None:
         npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
         np.savez(
             npz_path,
-            **{
-                f"step{index}": np.asarray(output.multimodal_output.get("actions"))
-                for index, output in enumerate(outputs)
-            },
+            step0=np.asarray(outputs[0].multimodal_output.get("actions")),
+            step1=np.asarray(outputs[1].multimodal_output.get("actions")),
         )
         print(f"SAVED_ACTIONS={npz_path}")
 
diff --git a/examples/online_serving/image_to_video/README.md b/examples/online_serving/image_to_video/README.md
index f18cea1c8ce..285eeb27983 100644
--- a/examples/online_serving/image_to_video/README.md
+++ b/examples/online_serving/image_to_video/README.md
@@ -120,17 +120,14 @@ Use `/v1/videos/sync` if you want to write the MP4 directly to a file. `POST /v1
 Generated video files are stored on local disk by the async video API.
 Local file storage behavior can be controlled via the following environment variables:
 
-- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files (default: `/tmp/storage`)
-- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
-
-`VLLM_OMNI_STORAGE_PATH` and `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` are deprecated and will be
-removed in a future release; use the names above instead.
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
 
 Example:
 
 ```bash
-export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
-export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
 ```
 
 ## API Calls
diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md
index 35b7f028490..c01e0602ff9 100644
--- a/examples/online_serving/text_to_video/README.md
+++ b/examples/online_serving/text_to_video/README.md
@@ -87,17 +87,14 @@ curl -X POST http://localhost:8091/v1/videos/sync \
 Generated video files are stored on local disk by the async video API.
 Local file storage behavior can be controlled via the following environment variables:
 
-- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files (default: `/tmp/storage`)
-- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
-
-`VLLM_OMNI_STORAGE_PATH` and `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` are deprecated and will be
-removed in a future release; use the names above instead.
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
 
 Example:
 
 ```bash
-export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
-export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
 ```
 
 ## API Calls
diff --git a/recipes/NVIDIA/GR00T-N1.7.md b/recipes/NVIDIA/GR00T-N1.7.md
deleted file mode 100644
index 91f676a524a..00000000000
--- a/recipes/NVIDIA/GR00T-N1.7.md
+++ /dev/null
@@ -1,101 +0,0 @@
-# GR00T-N1.7
-
-> NVIDIA Isaac GR00T-N1.7-3B robot VLA policy served over the OpenPI WebSocket protocol
-
-## Summary
-
-- Vendor: NVIDIA
-- Model: `nvidia/GR00T-N1.7-3B`
-- Task: Vision-Language-Action (VLA) inference for robot manipulation
-- Mode: Online serving via OpenPI WebSocket endpoint
-- Maintainer: timzsu
-
-## When to use this recipe
-
-Use this recipe when you need to serve GR00T-N1.7 as a real-time robot policy
-over the OpenPI WebSocket API. It configures the DROID embodiment
-(`OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT`) and exposes the standard DROID action
-keys (`eef_9d`, `gripper_position`, `joint_position`) with action horizon 40.
-
-## References
-
-- Upstream model: <https://huggingface.co/nvidia/GR00T-N1.7-3B>
-- Upstream codebase: <https://github.com/NVIDIA/Isaac-GR00T>
-- OpenPI client library: <https://github.com/Physical-Intelligence/openpi>
-- Pipeline: `vllm_omni.diffusion.models.gr00t.pipeline_gr00t.Gr00tN1d7Pipeline`
-- Deploy config: [`vllm_omni/deploy/Gr00tN1d7.yaml`](../../vllm_omni/deploy/Gr00tN1d7.yaml)
-- E2E test: [`tests/e2e/online_serving/test_gr00t_openpi.py`](../../tests/e2e/online_serving/test_gr00t_openpi.py)
-
-## Environment
-
-- OS: Linux
-- Python: 3.11+
-- Driver / runtime: NVIDIA CUDA
-- Hardware: 1 NVIDIA GPU. The upstream model card lists 16 GB+ VRAM for inference (e.g. RTX 4090, L40, H100); in practice this serving path uses ~6 GiB peak (bf16, TP=1, `max_num_seqs: 1`).
-- vLLM-Omni version or commit: use versions from your current checkout
-
-## Start server
-
-From repository root:
-
-```bash
-vllm serve nvidia/GR00T-N1.7-3B \
-  --omni \
-  --host 127.0.0.1 \
-  --port 8000 \
-  --served-model-name gr00t-n1d7 \
-  --stage-configs-path vllm_omni/deploy/Gr00tN1d7.yaml
-```
-
-Notes:
-
-- Only `max_num_seqs: 1` is supported (configured in the deploy YAML). The
-  policy is markovian (`state_history_length=1`, `reset()` returns `{}`) — the
-  reason for the cap is that `Gr00tPolicy` does its own per-sample
-  (un)batching inside `get_action` and is not integrated with vLLM's
-  continuous batching path.
-- This pipeline is a thin wrapper around the upstream HF policy
-  (`AutoModel.from_pretrained` + `enforce_eager`, `tensor_parallel_size=1`,
-  pipeline `load_weights` is a no-op). The standard diffusion accelerators
-  (SP, CFG, TeaCache, VAE tiling) do not transfer to a flow-matching action
-  policy at batch size 1, so a native-kernel port is intentionally out of
-  scope for this recipe — the value is the OpenPI serving integration, not
-  kernel-level acceleration.
-- The WebSocket endpoint is `ws://127.0.0.1:8000/v1/realtime/robot/openpi`.
-  The server handshake message (first frame after connect) is a msgpack-encoded
-  dict with `action_horizon`, `action_keys`, `embodiment_tag`, and
-  `needs_session_id`.
-
-## Verification
-
-```python
-from tests.gr00t.openpi_client_helper import run_policy_session, validate_session_result
-validate_session_result(run_policy_session(host="127.0.0.1", port=8000))
-```
-
-Or run the e2e test suite:
-
-```bash
-python -m pytest tests/e2e/online_serving/test_gr00t_openpi.py -v
-```
-
-The test sends a synthetic two-frame DROID observation and checks:
-
-- GR00T metadata contract: `image_resolution`, `action_horizon`, `action_keys`, `embodiment_tag`
-- Action shapes: `eef_9d (1,40,9)`, `gripper_position (1,40,1)`, `joint_position (1,40,7)`
-- All action values are finite float32
-- Reset response is `"reset successful"`
-
-## Notes
-
-- **Do not change the model-specific `policy_server_config` values.** `action_horizon`,
-  `action_keys`, and `supported_embodiments` are fixed by the GR00T-N1.7 checkpoint and
-  are validated against the loaded policy at startup (the server refuses to start on a
-  mismatch). Only `image_resolution` and `needs_session_id` are deployment knobs.
-- To switch embodiment, edit `embodiment_tag` under both `model_config` and
-  `policy_server_config` in `vllm_omni/deploy/Gr00tN1d7.yaml`. Supported values:
-  `OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT` (default), `XDOF`, `XDOF_SUBTASK`,
-  `REAL_G1`, `REAL_R1_PRO_SHARPA`, `LIBERO_PANDA`, `SIMPLER_ENV_GOOGLE`,
-  `SIMPLER_ENV_WIDOWX`.
-- GR00T weights are loaded directly by `Gr00tPolicy` via `AutoModel.from_pretrained`;
-  the pipeline's `load_weights` is intentionally a no-op.
diff --git a/recipes/README.md b/recipes/README.md
index e3ee2bc5639..1f5b5be5765 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -12,13 +12,6 @@ that is aligned with
 that external repository as a reference for structure rather than the place to
 add files for this repo. Use one Markdown file per model family by default.
 
-> **Note:** The vLLM-Omni TTS recipes (12 models across 8 providers) are also
-> published on the rendered recipes site at
-> [recipes.vllm.ai](https://recipes.vllm.ai) (see
-> [vllm-project/recipes#554](https://github.com/vllm-project/recipes/pull/554)).
-> Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`,
-> e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b).
-
 Example layout:
 
 ```text
@@ -37,12 +30,11 @@ recipes/
 | [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
 | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
 | [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
-| [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB |
 | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
-| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
+| [`inclusionAI/Ming-omni-tts-0.5B.md`](./inclusionAI/Ming-omni-tts-0.5B.md) | Offline + online dense Ming TTS/audio generation | 1x AMD `gfx942` (ROCm 7.2) |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
@@ -50,7 +42,6 @@ recipes/
 | [`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 |
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
 | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
-| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
diff --git a/recipes/mistralai/Voxtral-TTS.md b/recipes/mistralai/Voxtral-TTS.md
index aae48a6895a..595d5c29e5f 100644
--- a/recipes/mistralai/Voxtral-TTS.md
+++ b/recipes/mistralai/Voxtral-TTS.md
@@ -105,45 +105,3 @@ voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../exam
 ) or adding a `ref_audio` request parameter will fail
 because the public checkpoint does not provide the encoder weights needed to
 convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.`
-
-### 1x NVIDIA RTX A6000 48GB
-
-#### Environment
-
-- OS: Ubuntu 22.04
-- Python: 3.12
-- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
-- GPU: NVIDIA RTX A6000, 48 GB
-- vLLM version: 0.22.0
-- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
-- Required package: `pip install mistral-common`
-
-#### Command
-
-```bash
-vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
-```
-
-#### Verification
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "mistralai/Voxtral-4B-TTS-2603",
-    "input": "Artificial intelligence is transforming the way we interact with technology.",
-    "voice": "neutral_female",
-    "response_format": "wav"
-  }' -o voxtral_test.wav
-```
-
-Supported voice presets: `ar_male`, `casual_female`, `casual_male`, `cheerful_female`,
-`de_female`, `de_male`, `es_female`, `es_male`, `fr_female`, `fr_male`, `hi_female`,
-`hi_male`, `it_female`, `it_male`, `neutral_female`, `neutral_male`, `nl_female`,
-`nl_male`, `pt_female`, `pt_male`.
-
-#### Notes
-
-- Memory usage: Peak VRAM ~40.5 GiB during inference.
-- Generation time: ~8 seconds for ~200 words of English text.
-- Audio quality is high; mono 24 kHz output.
diff --git a/requirements/common.txt b/requirements/common.txt
index 6d9dd63ba40..469563d3ab0 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -28,5 +28,3 @@ aenum==3.1.16
 pyzmq>=25.0.0
 janus>=1.0.0
 msgpack>=1.0.0
-pydantic>=2.1.0
-pydantic-settings>=2.0.0
diff --git a/tests/bde/__init__.py b/tests/bde/__init__.py
new file mode 100644
index 00000000000..9881313609a
--- /dev/null
+++ b/tests/bde/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tests/bde/test_engine_routing.py b/tests/bde/test_engine_routing.py
new file mode 100644
index 00000000000..a53eb758c78
--- /dev/null
+++ b/tests/bde/test_engine_routing.py
@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Routing tests for BDE engine selection (Phase 1, PR-1).
+
+Covers ``DiffusionEngine.resolve_engine_class`` — the ``engine_backend``
+config-field dispatcher that routes a request to ``BDEEngine`` vs the base
+``DiffusionEngine``. Resolution is tested without constructing an engine
+(construction runs a dummy forward that needs a real model).
+"""
+
+from dataclasses import fields
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.bde.engine import BDEEngine
+from vllm_omni.diffusion.data import (
+    OmniDiffusionConfig,
+    default_engine_backend_for_model,
+)
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+
+def _cfg(backend):
+    return SimpleNamespace(engine_backend=backend)
+
+
+def test_default_resolves_to_diffusion_engine():
+    assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
+
+
+def test_missing_field_resolves_to_diffusion_engine():
+    # Backward compatibility: configs predating the engine_backend field.
+    assert DiffusionEngine.resolve_engine_class(SimpleNamespace()) is DiffusionEngine
+
+
+def test_none_resolves_to_diffusion_engine():
+    assert DiffusionEngine.resolve_engine_class(_cfg(None)) is DiffusionEngine
+
+
+def test_bde_key_resolves_to_bde_engine():
+    assert DiffusionEngine.resolve_engine_class(_cfg("bde")) is BDEEngine
+
+
+def test_subclass_type_is_returned():
+    assert DiffusionEngine.resolve_engine_class(_cfg(BDEEngine)) is BDEEngine
+
+
+def test_qualname_string_resolves():
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.bde.engine.BDEEngine"))
+    assert cls is BDEEngine
+
+
+def test_non_engine_type_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg(dict))
+
+
+def test_qualname_not_an_engine_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg("builtins.dict"))
+
+
+def test_bad_qualname_raises():
+    with pytest.raises(ValueError):
+        DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine"))
+
+
+def test_bde_engine_is_diffusion_engine_subclass():
+    assert issubclass(BDEEngine, DiffusionEngine)
+
+
+def test_config_engine_backend_field_defaults_to_default():
+    field = {f.name: f for f in fields(OmniDiffusionConfig)}["engine_backend"]
+    assert field.default == "default"
+
+
+# --- DreamZero defaults to the BDE engine -----------------------------------
+
+
+def test_dreamzero_pipeline_defaults_to_bde():
+    assert default_engine_backend_for_model("DreamZeroPipeline") == "bde"
+
+
+def test_unknown_model_defaults_to_none():
+    assert default_engine_backend_for_model("WanS2VPipeline") is None
+
+
+def test_none_model_defaults_to_none():
+    assert default_engine_backend_for_model(None) is None
+
+
+def test_dreamzero_default_routes_to_bde_engine():
+    # End-to-end of the routing: DreamZero's per-model default backend resolves
+    # to BDEEngine through the same dispatcher used by make_engine.
+    backend = default_engine_backend_for_model("DreamZeroPipeline")
+    assert DiffusionEngine.resolve_engine_class(_cfg(backend)) is BDEEngine
diff --git a/tests/config/test_server_settings.py b/tests/config/test_server_settings.py
deleted file mode 100644
index 53c15524ce6..00000000000
--- a/tests/config/test_server_settings.py
+++ /dev/null
@@ -1,84 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for ServerSettings env-var parameterization."""
-
-import pytest
-
-from vllm_omni.config.server_settings import ServerSettings
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-ENV_VARS = [
-    "VLLM_OMNI_SERVER_STORAGE__PATH",
-    "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY",
-    "VLLM_OMNI_SERVER_STORAGE__FILE_TTL",
-    "VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL",
-    "VLLM_OMNI_STORAGE_PATH",
-    "VLLM_OMNI_STORAGE_MAX_CONCURRENCY",
-]
-
-
-@pytest.fixture(autouse=True)
-def clean_env(monkeypatch):
-    for name in ENV_VARS:
-        monkeypatch.delenv(name, raising=False)
-
-
-def test_defaults_when_no_env_set():
-    storage = ServerSettings().storage
-
-    assert storage.path == "/tmp/storage"
-    assert storage.file_concurrency == 4
-    assert storage.file_ttl is None
-    assert storage.ttl_sweep_interval is None
-
-
-def test_nested_env_vars_populate_storage(monkeypatch):
-    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__PATH", "/var/storage")
-    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", "8")
-    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_TTL", "120")
-    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL", "30")
-
-    storage = ServerSettings().storage
-
-    assert storage.path == "/var/storage"
-    assert storage.file_concurrency == 8
-    assert storage.file_ttl == 120
-    assert storage.ttl_sweep_interval == 30
-
-
-def test_ttl_sweep_interval_defaults_when_ttl_set(monkeypatch):
-    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_TTL", "60")
-
-    storage = ServerSettings().storage
-
-    assert storage.file_ttl == 60
-    assert storage.ttl_sweep_interval == 300
-
-
-def test_deprecated_path_alias_populates_and_warns(monkeypatch):
-    monkeypatch.setenv("VLLM_OMNI_STORAGE_PATH", "/legacy/path")
-
-    with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_PATH"):
-        storage = ServerSettings().storage
-
-    assert storage.path == "/legacy/path"
-
-
-def test_deprecated_concurrency_alias_populates_and_warns(monkeypatch):
-    monkeypatch.setenv("VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "16")
-
-    with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_MAX_CONCURRENCY"):
-        storage = ServerSettings().storage
-
-    assert storage.file_concurrency == 16
-
-
-def test_new_env_var_wins_over_deprecated_alias(monkeypatch):
-    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__PATH", "/new/path")
-    monkeypatch.setenv("VLLM_OMNI_STORAGE_PATH", "/legacy/path")
-
-    with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_PATH"):
-        storage = ServerSettings().storage
-
-    assert storage.path == "/new/path"
diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py
index f5981b92748..021473ca5cc 100644
--- a/tests/diffusion/cache/test_cache_backends.py
+++ b/tests/diffusion/cache/test_cache_backends.py
@@ -13,12 +13,10 @@
 
 from unittest.mock import Mock, patch
 
-import cache_dit
 import pytest
-from cache_dit import ForwardPattern
 
 from vllm_omni.diffusion.cache.cache_dit_backend import (
-    CacheDiTAdapterConfig,
+    CUSTOM_DIT_ENABLERS,
     CacheDiTBackend,
 )
 from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
@@ -47,20 +45,13 @@ def test_init_with_config_object(self):
         assert backend.config.Fn_compute_blocks == 4
         assert backend.enabled is False
 
-    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-    def test_enable_single_transformer(self, mock_cache_dit, mock_block_adapter):
+    def test_enable_single_transformer(self, mock_cache_dit):
         """Test enabling cache-dit on single-transformer pipeline."""
         # Mock pipeline
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DiTPipeline"
         mock_transformer = Mock()
-        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
-            block_forward_patterns={
-                "layers": ForwardPattern.Pattern_0,
-            },
-        )
-
         mock_pipeline.transformer = mock_transformer
 
         # Mock cache_dit functions
@@ -75,20 +66,14 @@ def test_enable_single_transformer(self, mock_cache_dit, mock_block_adapter):
         assert backend._refresh_func is not None
         mock_cache_dit.enable_cache.assert_called_once()
 
-    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-    def test_refresh(self, mock_cache_dit, mock_block_adapter):
+    def test_refresh(self, mock_cache_dit):
         """Test refreshing cache context with SCM mask policy updates when num_inference_steps changes."""
         # Mock pipeline
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DiTPipeline"
         mock_transformer = Mock()
         mock_pipeline.transformer = mock_transformer
-        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
-            block_forward_patterns={
-                "layers": ForwardPattern.Pattern_0,
-            },
-        )
 
         # Mock cache_dit functions
         mock_cache_dit.enable_cache = Mock()
@@ -139,24 +124,23 @@ def test_refresh(self, mock_cache_dit, mock_block_adapter):
         assert "cache_config" in call_args[1]
         assert backend._last_num_inference_steps == 100
 
+    def test_hunyuan_custom_enabler_registered(self):
+        """Test HunyuanImage3 custom cache-dit enabler is registered."""
+        assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS
+
+    def test_dreamid_custom_enabler_registered(self):
+        """Test DreamIDOmni custom cache-dit enabler is registered."""
+        assert "DreamIDOmniPipeline" in CUSTOM_DIT_ENABLERS
+
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
-        """Test HunyuanImage3 uses pipeline.transformer for cache enable/refresh.
-
-        NOTE: HunyuanImage3 no longer has a custom enabler, so this tests against the generic path.
-        """
+        """Test HunyuanImage3 custom enabler uses pipeline.model for cache enable/refresh."""
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
         mock_pipeline.model = Mock()
         mock_pipeline.model.layers = Mock()
-        mock_pipeline.model._cache_dit_adapter_config = CacheDiTAdapterConfig(
-            block_forward_patterns={
-                "layers": ForwardPattern.Pattern_4,
-            },
-        )
-        # NOTE: pipe.transformer is pipe.model in HunyuanImage3Pipelines
-        mock_pipeline.transformer = mock_pipeline.model
+
         mock_cache_dit.enable_cache = Mock()
         mock_cache_dit.refresh_context = Mock()
 
@@ -168,9 +152,9 @@ def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mo
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.model
-        assert len(adapter_kwargs["blocks"]) == 1
-        assert adapter_kwargs["blocks"][0] == mock_pipeline.model.layers
-        assert adapter_kwargs["forward_pattern"][0] == ForwardPattern.Pattern_4
+        assert adapter_kwargs["blocks"] is mock_pipeline.model.layers
+        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_4
+        assert len(adapter_kwargs["params_modifiers"]) == 1
         mock_cache_dit.enable_cache.assert_called_once()
 
         backend.refresh(mock_pipeline, num_inference_steps=12)
@@ -179,23 +163,24 @@ def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mo
         assert call_args[0][0] is mock_pipeline.model
         assert call_args[1]["num_inference_steps"] == 12
 
+    def test_enable_hunyuan_pipeline_requires_model_layers(self):
+        """Test HunyuanImage3 enabler fails with a formatted pipeline class name."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
+        mock_pipeline.model = Mock(spec=[])
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+
+        with pytest.raises(ValueError, match="HunyuanImage3Pipeline"):
+            backend.enable(mock_pipeline)
+
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
-        """Test DreamID uses pipeline.transformer for cache enable/refresh.
-
-        NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
-        """
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
         mock_pipeline.transformer = Mock()
         mock_pipeline.transformer.fused_blocks = Mock()
-        mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
-            block_forward_patterns={
-                "fused_blocks": ForwardPattern.Pattern_0,
-            },
-            has_separate_cfg=True,
-        )
 
         mock_cache_dit.enable_cache = Mock()
         mock_cache_dit.refresh_context = Mock()
@@ -208,10 +193,10 @@ def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_bl
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.transformer
-        assert len(adapter_kwargs["blocks"]) == 1
-        assert adapter_kwargs["blocks"][0] == mock_pipeline.transformer.fused_blocks
-        assert adapter_kwargs["forward_pattern"][0] == ForwardPattern.Pattern_0
+        assert adapter_kwargs["blocks"] is mock_pipeline.transformer.fused_blocks
+        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_0
         assert adapter_kwargs["has_separate_cfg"] is True
+        assert len(adapter_kwargs["params_modifiers"]) == 1
         mock_cache_dit.enable_cache.assert_called_once()
 
         backend.refresh(mock_pipeline, num_inference_steps=12)
@@ -220,43 +205,15 @@ def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_bl
         assert call_args[0][0] is mock_pipeline.transformer
         assert call_args[1]["num_inference_steps"] == 12
 
-    @pytest.mark.parametrize("num_inference_steps", [1, 7])
-    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
-    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-    def test_refresh_scm_bypassed_for_unsupported_step_counts(
-        self, mock_cache_dit, mock_block_adapter, num_inference_steps
-    ):
-        """Ensure SCM is bypassed when num_inference_steps < 8 and not in (4, 6),
-        because cache_dit.steps_mask() raises for unsupported step count.
-        For these cases, we fall back to the non-SCM path to avoid crashing.
-        """
+    def test_enable_dreamid_pipeline_requires_fused_blocks(self):
         mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "DiTPipeline"
-        mock_transformer = Mock()
-        mock_pipeline.transformer = mock_transformer
-        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
-            block_forward_patterns={
-                "layers": ForwardPattern.Pattern_0,
-            },
-        )
-
-        mock_cache_dit.enable_cache = Mock()
-        mock_cache_dit.refresh_context = Mock()
-        mock_cache_dit.steps_mask = cache_dit.steps_mask
-
-        # Create a cache config with an scm policy & enable it
-        config = DiffusionCacheConfig(scm_steps_mask_policy="fast")
-        backend = CacheDiTBackend(config)
-        backend.enable(mock_pipeline)
+        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.transformer = Mock(spec=[])
 
-        backend.refresh(mock_pipeline, num_inference_steps=num_inference_steps)
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
 
-        mock_cache_dit.refresh_context.assert_called_once()
-        call_args = mock_cache_dit.refresh_context.call_args
-        # Ensure that we properly guard, i.e., cache config is filtered
-        assert call_args[0][0] == mock_transformer
-        assert call_args[1]["num_inference_steps"] == num_inference_steps
-        assert "cache_config" not in call_args[1]
+        with pytest.raises(ValueError, match="DreamIDOmniPipeline"):
+            backend.enable(mock_pipeline)
 
 
 class TestTeaCacheBackend:
diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py
index 0a05aeb761f..146084a508c 100644
--- a/tests/diffusion/cache/test_cache_dit.py
+++ b/tests/diffusion/cache/test_cache_dit.py
@@ -5,7 +5,6 @@
 Model specific tests for CacheDiT enablement.
 """
 
-import sys
 from unittest.mock import Mock, patch
 
 import pytest
@@ -13,49 +12,36 @@
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
 
 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
-from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
-from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
-from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
 from vllm_omni.platforms import current_omni_platform
 
-# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
-# the class even though the dependency may not be set up correctly; this is ok for these
-# tests because we just inspect it and never initialize the model.
-for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
-    sys.modules.setdefault(mod, Mock())
-# isort: split
-from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel  # noqa: E402
-
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-SEPARATE_CFG_TRANSFORMERS = [
-    DreamIdOmniModel,
-    LTX2VideoTransformer3DModel,
-    HeliosTransformer3DModel,
-    LongCatImageTransformer2DModel,
-    Cosmos3VFMTransformer,
+SEPARATE_CFG_ENABLERS = [
+    cd_backend.enable_cache_for_dreamid_omni,
+    cd_backend.enable_cache_for_ltx2,
+    cd_backend.enable_cache_for_helios,
+    cd_backend.enable_cache_for_wan22,
+    cd_backend.enable_cache_for_longcat_image,
+    cd_backend.enable_cache_for_cosmos3,
 ]
 
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
 
 
-@pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS)
-def test_cache_dit_configs_have_separate_cfg(transformer_model):
-    """Check models with separate CFG set has_separate_cfg=True in their configs."""
-    assert hasattr(transformer_model, "_cache_dit_adapter_config")
-    assert isinstance(transformer_model._cache_dit_adapter_config, CacheDiTAdapterConfig)
-    assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
-
-
+@pytest.mark.parametrize("enabler", SEPARATE_CFG_ENABLERS)
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
-    """Ensure that Wan22, which has a custom enabler, setts custom CFG correctly."""
+def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler):
+    """Ensure that custom enablers for models with separate CFG pass
+    the param through to cache_dit correctly.
+
+    Regression test for: https://github.com/vllm-project/vllm-omni/pull/2860
+    """
     mock_pipeline = Mock()
-    cd_backend.enable_cache_for_wan22(mock_pipeline, SAMPLE_CACHE_CONFIG)
+    enabler(mock_pipeline, SAMPLE_CACHE_CONFIG)
 
     mock_cache_dit.enable_cache.assert_called_once()
     adapter_kwargs = mock_block_adapter.call_args.kwargs
@@ -69,7 +55,6 @@ def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
     mock_pipeline = Mock()
     gen_layers = object()
     mock_pipeline.transformer.gen_layers = gen_layers
-    mock_pipeline.transformer._cache_dit_adapter_config = Cosmos3VFMTransformer._cache_dit_adapter_config
 
     cd_backend.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG)
 
diff --git a/tests/diffusion/cache/test_stepcache.py b/tests/diffusion/cache/test_stepcache.py
deleted file mode 100644
index 1fc3bbb5922..00000000000
--- a/tests/diffusion/cache/test_stepcache.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Unit tests for StepCacheBackend and StepCacheState."""
-
-from unittest.mock import Mock
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.cache.selector import get_cache_backend
-from vllm_omni.diffusion.cache.stepcache import (
-    STEP_CACHE_CONFIG_ATTR,
-    STEP_CACHE_STATE_ATTR,
-    StepCacheBackend,
-    StepCacheConfig,
-    StepCacheState,
-    is_stepcache_active,
-)
-from vllm_omni.diffusion.data import DiffusionCacheConfig
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_prev_predictions(sim: float, dim: int = 8) -> list[tuple[torch.Tensor]]:
-    v_prev = torch.randn(1, dim)
-    if sim >= 0.99:
-        v_last = v_prev.clone()
-    else:
-        v_last = torch.randn(1, dim)
-    return [(v_prev,), (v_last,)]
-
-
-class TestStepCacheConfig:
-    def test_from_diffusion_cache_config_respects_step_cache_dit_enabled(self):
-        enabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=True))
-        disabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=False))
-        assert enabled.enabled is True
-        assert disabled.enabled is False
-
-
-class TestStepCacheState:
-    def test_disabled_always_runs(self):
-        state = StepCacheState(StepCacheConfig(enabled=False))
-        assert state.should_run_step(_make_prev_predictions(0.99)) is True
-
-    def test_insufficient_history_always_runs(self):
-        state = StepCacheState(StepCacheConfig(min_history_steps=2))
-        assert state.should_run_step([(torch.randn(1, 4),)]) is True
-
-    def test_high_similarity_skips_with_countdown(self):
-        state = StepCacheState(
-            StepCacheConfig(
-                sim_thresholds=(0.95,),
-                skip_countdowns=(2,),
-            )
-        )
-        preds = _make_prev_predictions(0.99)
-        assert state.should_run_step(preds) is False
-        assert state.skip_countdown == 2
-        assert state.should_run_step(preds) is False
-        assert state.should_run_step(preds) is True
-
-    def test_low_similarity_runs(self):
-        state = StepCacheState(StepCacheConfig())
-        preds = _make_prev_predictions(0.5)
-        assert state.should_run_step(preds) is True
-
-    def test_trim_history(self):
-        state = StepCacheState(StepCacheConfig(max_history=2))
-        history: list[tuple[torch.Tensor]] = [(torch.randn(1, 2),) for _ in range(4)]
-        state.trim_history(history)
-        assert len(history) == 2
-
-
-class TestStepCacheBackend:
-    def test_enable_dreamzero_pipeline(self):
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "DreamZeroPipeline"
-
-        backend = StepCacheBackend(DiffusionCacheConfig())
-        backend.enable(mock_pipeline)
-
-        assert backend.enabled is True
-        cfg = getattr(mock_pipeline, STEP_CACHE_CONFIG_ATTR)
-        assert isinstance(cfg, StepCacheConfig)
-        assert is_stepcache_active(mock_pipeline)
-        state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR)
-        assert isinstance(state, StepCacheState)
-
-    def test_refresh_resets_state(self):
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "DreamZeroPipeline"
-        backend = StepCacheBackend(DiffusionCacheConfig())
-        backend.enable(mock_pipeline)
-        state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR)
-        state.skip_countdown = 3
-        backend.refresh(mock_pipeline, num_inference_steps=16, verbose=False)
-        assert state.skip_countdown == 0
-
-    def test_enable_unsupported_pipeline_raises(self):
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "QwenImagePipeline"
-        backend = StepCacheBackend(DiffusionCacheConfig())
-
-        with pytest.raises(ValueError, match="step_cache backend does not support"):
-            backend.enable(mock_pipeline)
-
-    def test_selector_returns_backend(self):
-        backend = get_cache_backend("step_cache", {"velocity_sim_thresholds": [0.9]})
-        assert isinstance(backend, StepCacheBackend)
-        assert backend.config.velocity_sim_thresholds == [0.9]
-
-    def test_legacy_backend_aliases(self):
-        assert isinstance(get_cache_backend("stepcache", {}), StepCacheBackend)
-        assert isinstance(get_cache_backend("step_cache_dit", {}), StepCacheBackend)
diff --git a/tests/diffusion/models/gr00t/__init__.py b/tests/diffusion/models/gr00t/__init__.py
deleted file mode 100644
index e69de29bb2d..00000000000
diff --git a/tests/diffusion/models/gr00t/test_pipeline.py b/tests/diffusion/models/gr00t/test_pipeline.py
deleted file mode 100644
index 72118c3f93d..00000000000
--- a/tests/diffusion/models/gr00t/test_pipeline.py
+++ /dev/null
@@ -1,141 +0,0 @@
-from types import SimpleNamespace
-
-import numpy as np
-import pytest
-
-from vllm_omni.diffusion.models.gr00t import pipeline_gr00t
-from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class FakeGr00tPolicy:
-    instances = []
-
-    def __init__(self, **kwargs):
-        self.kwargs = kwargs
-        self.reset_calls = 0
-        self.seen_obs = None
-        self.embodiment_tag = SimpleNamespace(value="fake_embodiment")
-        self.language_key = "annotation.language.language_instruction"
-        self.modality_configs = {
-            "action": SimpleNamespace(
-                delta_indices=[0, 1],
-                modality_keys=["arm", "gripper"],
-            )
-        }
-        self.processor = SimpleNamespace(
-            state_action_processor=SimpleNamespace(
-                norm_params={
-                    "fake_embodiment": {
-                        "action": {
-                            "arm": {"dim": np.array(2)},
-                            "gripper": {"dim": np.array(1)},
-                        }
-                    }
-                }
-            )
-        )
-        FakeGr00tPolicy.instances.append(self)
-
-    def get_action(self, obs):
-        self.seen_obs = obs
-        return {
-            "arm": np.array([[[1.0, 2.0]]], dtype=np.float64),
-            "gripper": [[[3.0]]],
-        }, {"latency_ms": 1.0}
-
-    def reset(self):
-        self.reset_calls += 1
-        return {"reset": True}
-
-
-@pytest.fixture(autouse=True)
-def fake_gr00t_policy(monkeypatch):
-    FakeGr00tPolicy.instances.clear()
-    monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", FakeGr00tPolicy)
-
-
-def _pipeline():
-    od_config = SimpleNamespace(
-        model="nvidia/GR00T-N1.7-3B",
-        model_config={
-            "embodiment_tag": "LIBERO_PANDA",
-            "strict": False,
-        },
-        custom_pipeline_args={},
-    )
-    return Gr00tN1d7Pipeline(od_config=od_config)
-
-
-def test_pipeline_initializes_local_policy():
-    pipeline = _pipeline()
-
-    policy = FakeGr00tPolicy.instances[0]
-    assert policy.kwargs["model_path"] == "nvidia/GR00T-N1.7-3B"
-    assert policy.kwargs["embodiment_tag"] == "LIBERO_PANDA"
-    assert policy.kwargs["strict"] is False
-    assert pipeline.weights_sources == ()
-    assert pipeline.load_weights(iter(())) == set()
-
-
-def test_forward_returns_dict_actions_in_output():
-    pipeline = _pipeline()
-    req = OmniDiffusionRequest(
-        prompts=["pick"],
-        request_id="req",
-        sampling_params=OmniDiffusionSamplingParams(
-            extra_args={
-                "robot_obs": {
-                    "images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)},
-                    "state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)},
-                    "prompt": "pick the cube",
-                    "session_id": "session-a",
-                },
-                "reset": True,
-            }
-        ),
-    )
-
-    output = pipeline.forward(req)
-
-    assert output.error is None
-    actions = output.output["actions"]
-    assert set(actions) == {"arm", "gripper"}
-    assert actions["arm"].dtype == np.float32
-    np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32))
-    policy = FakeGr00tPolicy.instances[0]
-    assert "video" in policy.seen_obs
-    assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
-    assert "images" not in policy.seen_obs
-    assert "prompt" not in policy.seen_obs
-    assert "session_id" not in policy.seen_obs
-    assert policy.reset_calls == 1
-
-
-def test_dummy_warmup_returns_shape_correct_zero_actions():
-    pipeline = _pipeline()
-    req = OmniDiffusionRequest(
-        prompts=["dummy run"],
-        request_id="dummy_req_id",
-        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-    )
-
-    output = pipeline.forward(req)
-
-    assert output.error is None
-    actions = output.output["actions"]
-    assert set(actions) == {"arm", "gripper"}
-    assert actions["arm"].shape == (1, 2, 2)
-    assert actions["gripper"].shape == (1, 2, 1)
-    assert not actions["arm"].any()
-    assert FakeGr00tPolicy.instances[0].seen_obs is None
-
-
-def test_reset_delegates_to_policy():
-    pipeline = _pipeline()
-
-    assert pipeline.reset() == {"reset": True}
-    assert FakeGr00tPolicy.instances[0].reset_calls == 1
diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
index 9f38f097acd..bf6a7563623 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
@@ -692,27 +692,16 @@ def test_post_process_funcs_registered(self):
             assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS"
             assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func"
 
-    def test_cache_dit_for_ltx2_does_not_have_custom_enablers_registered(self):
-        """Pipeline variants are *not* registered in CUSTOM_DIT_ENABLERS."""
+    def test_cache_dit_enablers_registered(self):
+        """Pipeline variants must be registered in CUSTOM_DIT_ENABLERS."""
         from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
 
-        # NOTE: We used to have custom enablers for this model, but refactored to handle
-        # it more generically. Now we only need to ensure it has git cache adapter config.
         expected = [
             "LTX23Pipeline",
             "LTX23ImageToVideoPipeline",
         ]
         for name in expected:
-            assert name not in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS"
-
-    def test_ltx2_transformer_has_dit_cache_config(self):
-        """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
-        from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
-        from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
-
-        adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
-        assert isinstance(adapter_config, CacheDiTAdapterConfig)
-        assert adapter_config.has_separate_cfg
+            assert name in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS"
 
 
 class TestVocoderSampleRateDetection:
diff --git a/tests/diffusion/quantization/test_hsdp_fp8.py b/tests/diffusion/quantization/test_hsdp_fp8.py
deleted file mode 100644
index 97279acf2f5..00000000000
--- a/tests/diffusion/quantization/test_hsdp_fp8.py
+++ /dev/null
@@ -1,147 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for HSDP/FSDP2 compatibility with online FP8 quantization."""
-
-import pytest
-import torch
-from torch import nn
-
-from vllm_omni.diffusion.quantization.hsdp_fp8 import (
-    _build_transposed_get_layer_params,
-    prepare_fp8_layers_for_fsdp,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-# --- shared test doubles ---
-
-
-class _ToyKernel:
-    """Minimal kernel that mirrors ``FP8ScaledMMLinearKernel``."""
-
-    def __init__(self):
-        self.layer_param_names = ("weight", "weight_scale", "input_scale", "input_scale_ub")
-
-    def _get_layer_params(self, layer):
-        w, w_s, x_s, x_s_ub = self.layer_param_names
-        return (
-            getattr(layer, w),
-            getattr(layer, w_s),
-            getattr(layer, x_s, None),
-            getattr(layer, x_s_ub, None),
-        )
-
-
-# --- helper to build a toy module simulating online-FP8 post-load state ---
-
-
-def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32):
-    """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
-    ``quant_method`` is an ``Fp8OnlineLinearMethod`` with a ``_ToyKernel``.
-    """
-    from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
-
-    # Create the quant-method instance without calling __init__ (avoids
-    # heavyweight upstream config dependencies).  isinstance() still works.
-    qm = object.__new__(Fp8OnlineLinearMethod)
-    qm.fp8_linear = _ToyKernel()
-
-    # randn does not support float8 on CPU; zeros does and the test only
-    # cares about shape/stride, not the actual weight values.
-    qweight = torch.zeros(out_features, in_features, dtype=torch.float8_e4m3fn)
-    weight = nn.Parameter(qweight.t())  # (in, out) non-contiguous column-major view
-
-    module = nn.Module()
-    module.quant_method = qm
-    module.weight = weight
-    module.weight_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
-    # Set input_scale / input_scale_ub to verify _get_layer_params
-    # passes them through unchanged.
-    module.input_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
-    module.input_scale_ub = nn.Parameter(torch.ones(1, dtype=torch.float32))
-
-    return module
-
-
-# --- tests ---
-
-
-def test_transposed_get_layer_params():
-    import types
-
-    kernel = _ToyKernel()
-
-    # original bound method + patch + re-bind, matching prepare_fp8_layers_for_fsdp
-    original_bound = kernel._get_layer_params
-    patched_func = _build_transposed_get_layer_params(original_bound)
-    kernel._get_layer_params = types.MethodType(patched_func, kernel)
-
-    module = _make_fp8_toy_module(16, 32)
-    # Reproduce the storage rewrite from prepare_fp8_layers_for_fsdp:
-    # weight was qweight.t()  →  .t() recovers qweight (row-major contiguous).
-    module.weight = nn.Parameter(module.weight.data.t(), requires_grad=False)
-    assert module.weight.is_contiguous(), "qweight should be row-major contiguous"
-
-    w, w_s, x_s, x_s_ub = kernel._get_layer_params(module)
-
-    # patched method transposes row-major → column-major for Cutlass
-    assert w.shape == (32, 16)
-    assert not w.is_contiguous()
-    # scales are passed through unchanged
-    assert w_s is module.weight_scale
-    assert x_s is module.input_scale
-    assert x_s_ub is module.input_scale_ub
-
-
-def test_rewrites_non_contiguous_weight():
-    module = _make_fp8_toy_module(16, 32)
-    assert not module.weight.is_contiguous()
-
-    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
-    assert n_rewritten_layers == 1
-    assert module.weight.is_contiguous()
-    assert module.weight.shape == (16, 32)
-
-
-def test_patched_get_layer_params_returns_transposed_view():
-    module = _make_fp8_toy_module(16, 32)
-    prepare_fp8_layers_for_fsdp(module)
-
-    kernel = module.quant_method.fp8_linear
-    w, w_s, x_s, x_s_ub = kernel._get_layer_params(module)
-
-    # Cutlass expects column-major B: (in, out), stride (1, in)
-    assert w.shape == (32, 16)
-    assert not w.is_contiguous()
-    assert w.stride() == (1, 32)
-
-    # scales pass through unchanged
-    assert w_s is module.weight_scale
-    assert x_s is module.input_scale
-    assert x_s_ub is module.input_scale_ub
-
-
-def test_contiguous_weights_skipped():
-    module = _make_fp8_toy_module(16, 32)
-    # make weight contiguous before calling
-    module.weight = nn.Parameter(module.weight.data.t().contiguous())
-    assert module.weight.is_contiguous()
-
-    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
-    assert n_rewritten_layers == 0
-
-
-def test_kernel_patched_only_once_per_instance():
-    module = _make_fp8_toy_module(16, 32)
-    kernel = module.quant_method.fp8_linear
-    original = kernel._get_layer_params
-
-    prepare_fp8_layers_for_fsdp(module)
-    first_patched = kernel._get_layer_params
-
-    # Second call with same kernel id must not stack patches
-    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
-    assert n_rewritten_layers == 0
-    assert kernel._get_layer_params is first_patched
-    assert kernel._get_layer_params is not original
diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py
index fb45cd41d0d..8ad6c273416 100644
--- a/tests/diffusion/test_diffusion_worker.py
+++ b/tests/diffusion/test_diffusion_worker.py
@@ -23,15 +23,6 @@
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
 
-def patch_cumem_allocator(mocker: MockerFixture):
-    mock_allocator_class = mocker.Mock()
-    mocker.patch(
-        "vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class",
-        return_value=mock_allocator_class,
-    )
-    return mock_allocator_class
-
-
 @pytest.fixture
 def mock_od_config(mocker: MockerFixture):
     """Create a mock OmniDiffusionConfig."""
@@ -127,7 +118,7 @@ def setup_allocator(self, mocker: MockerFixture):
         """
         Unified interception of Allocators, and provision of default security values.
         """
-        self.mock_allocator_class = patch_cumem_allocator(mocker)
+        self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
         self.mock_allocator = mocker.Mock()
         self.mock_allocator_class.get_instance.return_value = self.mock_allocator
         self.mock_allocator.get_current_usage.return_value = 4 * 1024**3
@@ -135,7 +126,7 @@ def setup_allocator(self, mocker: MockerFixture):
 
     def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
-        mock_allocator_class = patch_cumem_allocator(mocker)
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
         mock_platform.get_device_total_memory.return_value = 80 * 1024**3
@@ -167,7 +158,7 @@ def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
 
     def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 2 (offload all, save buffers)."""
-        mock_allocator_class = patch_cumem_allocator(mocker)
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3]
         mock_platform.get_device_total_memory.return_value = 80 * 1024**3
@@ -210,7 +201,7 @@ def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
 
     def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
         """Test that sleep validates memory was actually freed."""
-        mock_allocator_class = patch_cumem_allocator(mocker)
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_platform.get_free_memory.return_value = 10 * 1024**3
         mock_platform.get_device_total_memory.return_value = 80 * 1024**3
@@ -235,7 +226,7 @@ def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_wor
     def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep uses device-scoped fallback when NVML is unavailable."""
 
-        mock_allocator_class = patch_cumem_allocator(mocker)
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
         mock_get_process_memory.side_effect = [None, None]
@@ -261,7 +252,7 @@ class TestDiffusionWorkerWakeUp:
 
     def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up without saved buffers (level 1 sleep)."""
-        mock_allocator_class = patch_cumem_allocator(mocker)
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
 
         # Setup allocator mock
         mock_allocator = mocker.Mock()
@@ -281,7 +272,7 @@ def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):
 
     def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up with saved buffers (level 2 sleep)."""
-        mock_allocator_class = patch_cumem_allocator(mocker)
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
 
         # Setup allocator mock
         mock_allocator = mocker.Mock()
@@ -326,7 +317,7 @@ def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
 
     def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up only restores buffers that were saved."""
-        mock_allocator_class = patch_cumem_allocator(mocker)
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
 
         # Setup allocator mock
         mock_allocator = mocker.Mock()
diff --git a/tests/dreamzero/test_crossattn_cache.py b/tests/dreamzero/test_crossattn_cache.py
deleted file mode 100644
index 94361c090a0..00000000000
--- a/tests/dreamzero/test_crossattn_cache.py
+++ /dev/null
@@ -1,114 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Tests that DreamZero's cross-attention KV cache is actually init'd and reused.
-
-DreamZero's I2V cross-attention caches the text k/v and image k_img/v_img in a
-per-layer ``crossattn_cache`` dict: they are computed on the first forward of a
-session (``is_init`` False -> True) and reused on every later step, because the
-text/image context is session-invariant. Only the query (which depends on the
-per-step hidden state) is recomputed.
-
-These tests lock in that ``WanI2VCrossAttention.forward`` initialises the cache
-on the first call and reuses it (not recompute) on later calls, and recomputes
-every call when no cache is provided.
-
-CPU-only: ``WanI2VCrossAttention.__init__`` needs a tensor-parallel group (it
-builds ColumnParallel/RowParallel linears), so we bypass ``__init__`` and stub
-the projections/attention to count calls and exercise only the caching logic.
-"""
-
-from unittest.mock import MagicMock
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.models.dreamzero.causal_wan_model import WanI2VCrossAttention
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-_CLIP_LEN = 257  # image tokens the I2V cross-attn splits off the front of context
-
-
-def _make_cross_attn(n_heads: int = 2, head_dim: int = 4) -> WanI2VCrossAttention:
-    """A WanI2VCrossAttention whose submodules are call-counting stubs.
-
-    ``__init__`` is bypassed (it needs a TP group); we only wire what
-    ``forward`` touches. Each projection returns a [B, S, n*head_dim] tensor so
-    the subsequent ``unflatten(2, (n, head_dim))`` works, and is a MagicMock so
-    we can assert how many times it was invoked.
-    """
-    inner = n_heads * head_dim
-    module = WanI2VCrossAttention.__new__(WanI2VCrossAttention)
-    module.tp_num_heads = n_heads
-    module.head_dim = head_dim
-
-    def _proj() -> MagicMock:
-        return MagicMock(side_effect=lambda t: torch.zeros(t.shape[0], t.shape[1], inner))
-
-    module.q = _proj()
-    module.k = _proj()
-    module.v = _proj()
-    module.k_img = _proj()
-    module.v_img = _proj()
-    module.norm_q = lambda t: t
-    module.norm_k = lambda t: t
-    module.norm_k_img = lambda t: t
-    # attn returns the query unchanged ([B, S, n, d]); flatten(2) downstream is happy.
-    module.attn = lambda q, k, v: q
-    module.o = lambda t: t
-    return module
-
-
-def _fresh_cache() -> dict:
-    return {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None}
-
-
-def _inputs(inner: int = 8):
-    x = torch.randn(1, 3, inner)  # B=1, query_len=3
-    context = torch.randn(1, _CLIP_LEN + 5, inner)  # 257 image tokens + 5 text tokens
-    return x, context
-
-
-def test_cache_inits_on_first_call_and_reuses_on_second() -> None:
-    module = _make_cross_attn()
-    cache = _fresh_cache()
-    x, context = _inputs()
-
-    # First call: cache is initialised and every k/v projection runs once.
-    module.forward(x, context, crossattn_cache=cache)
-    assert cache["is_init"] is True
-    for key in ("k", "v", "k_img", "v_img"):
-        assert cache[key] is not None, f"{key} should be cached after first call"
-    assert module.k.call_count == 1
-    assert module.v.call_count == 1
-    assert module.k_img.call_count == 1
-    assert module.v_img.call_count == 1
-
-    cached = {key: cache[key] for key in ("k", "v", "k_img", "v_img")}
-
-    # Second call: k/v/k_img/v_img are read from cache, NOT recomputed.
-    module.forward(x, context, crossattn_cache=cache)
-    assert module.k.call_count == 1
-    assert module.v.call_count == 1
-    assert module.k_img.call_count == 1
-    assert module.v_img.call_count == 1
-    for key, tensor in cached.items():
-        assert cache[key] is tensor, f"{key} should be the same cached object"
-
-    # The query always depends on the per-step hidden state, so it is recomputed.
-    assert module.q.call_count == 2
-
-
-def test_no_cache_recomputes_every_call() -> None:
-    module = _make_cross_attn()
-    x, context = _inputs()
-
-    module.forward(x, context, crossattn_cache=None)
-    module.forward(x, context, crossattn_cache=None)
-
-    # Without a cache, the image/text k/v are recomputed on every call.
-    assert module.k_img.call_count == 2
-    assert module.v_img.call_count == 2
-    assert module.k.call_count == 2
-    assert module.v.call_count == 2
diff --git a/tests/dreamzero/test_fused_qk_rms_norm.py b/tests/dreamzero/test_fused_qk_rms_norm.py
deleted file mode 100644
index dbbb439400b..00000000000
--- a/tests/dreamzero/test_fused_qk_rms_norm.py
+++ /dev/null
@@ -1,112 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Equivalence tests for DreamZero's fused q/k qk-norm TP all-reduce.
-
-`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
-and issues a single tensor-parallel all-reduce instead of two. These tests lock
-in that it is numerically identical to applying `norm_q`/`norm_k` separately,
-and that it really issues exactly one (fused) all-reduce.
-
-CPU-only: the TP world size and the collective are mocked, so a single process
-simulates one rank (all-reduce is identity). This exercises the cat/split and
-count-scaling logic that a real multi-rank run depends on.
-"""
-
-from unittest.mock import patch
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.models.dreamzero.causal_wan_model import (
-    DistributedRMSNorm,
-    fused_qk_rms_norm,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-_MODULE = "vllm_omni.diffusion.models.dreamzero.causal_wan_model"
-
-
-def _make_norms(dim: int) -> tuple[DistributedRMSNorm, DistributedRMSNorm]:
-    """Two norms with distinct random weights so a swapped q/k slice is caught."""
-    torch.manual_seed(0)
-    norm_q = DistributedRMSNorm(dim)
-    norm_k = DistributedRMSNorm(dim)
-    norm_q.weight.data.normal_()
-    norm_k.weight.data.normal_()
-    return norm_q, norm_k
-
-
-def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
-    """Single-rank all-reduce: the reduction over one rank is the identity."""
-    return tensor
-
-
-def test_fused_matches_separate_tp1() -> None:
-    dim = 64
-    norm_q, norm_k = _make_norms(dim)
-    q = torch.randn(2, 5, dim)
-    k = torch.randn(2, 5, dim)
-
-    with patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1):
-        ref_q, ref_k = norm_q(q), norm_k(k)
-        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
-
-    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
-    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
-
-
-def test_fused_matches_separate_tp_gt1() -> None:
-    dim = 64
-    norm_q, norm_k = _make_norms(dim)
-    q = torch.randn(2, 5, dim)
-    k = torch.randn(2, 5, dim)
-
-    # tp_size=2 exercises the reduce branch and the `global_count *= tp_size`
-    # path in both the per-norm forward and the fused helper.
-    with (
-        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
-        patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce),
-    ):
-        ref_q, ref_k = norm_q(q), norm_k(k)
-        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
-
-    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
-    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
-
-
-def test_fused_issues_single_all_reduce_over_packed_pair() -> None:
-    dim = 64
-    norm_q, norm_k = _make_norms(dim)
-    q = torch.randn(2, 5, dim)
-    k = torch.randn(2, 5, dim)
-
-    with (
-        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
-        patch(
-            f"{_MODULE}.tensor_model_parallel_all_reduce",
-            side_effect=_identity_all_reduce,
-        ) as mock_ar,
-    ):
-        fused_qk_rms_norm(norm_q, norm_k, q, k)
-
-    # Exactly one collective, over a tensor whose last dim packs the q+k pair.
-    assert mock_ar.call_count == 1
-    packed = mock_ar.call_args[0][0]
-    assert packed.shape[-1] == 2
-
-
-def test_fused_identity_fallback_when_qk_norm_disabled() -> None:
-    # qk_norm=False makes the norms nn.Identity; the helper must fall back to
-    # plain passthrough without touching the collective.
-    norm_q, norm_k = torch.nn.Identity(), torch.nn.Identity()
-    q = torch.randn(2, 5, 64)
-    k = torch.randn(2, 5, 64)
-
-    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce") as mock_ar:
-        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
-
-    torch.testing.assert_close(fused_q, q, rtol=0, atol=0)
-    torch.testing.assert_close(fused_k, k, rtol=0, atol=0)
-    mock_ar.assert_not_called()
diff --git a/tests/dreamzero/test_qkv_fusion.py b/tests/dreamzero/test_qkv_fusion.py
deleted file mode 100644
index 16b78693f69..00000000000
--- a/tests/dreamzero/test_qkv_fusion.py
+++ /dev/null
@@ -1,117 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Equivalence test for DreamZero self-attn QKV fusion.
-
-CausalWanSelfAttention now projects q/k/v with a single QKVParallelLinear and
-splits the result, instead of three separate ColumnParallelLinear. The DreamZero
-checkpoint still ships separate q/k/v weights, so load_weights routes each into
-the packed `qkv` param with its shard id. This test proves the chain
-"load separate q/k/v via shard id -> single GEMM -> split" reproduces the three
-independent projections, given identical weights.
-
-CPU-only: a single-rank (world_size=1) tensor-parallel group is initialized and
-the GEMM dispatcher is forced to the default (CPU-compatible) path, following
-tests/diffusion/layers/test_adalayernorm.py.
-"""
-
-import os
-
-import pytest
-import torch
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.fixture(autouse=True)
-def _init_distributed():
-    """Minimal world_size=1 TP group required to build QKVParallelLinear."""
-    from vllm.distributed.parallel_state import (
-        cleanup_dist_env_and_memory,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ.setdefault("MASTER_ADDR", "localhost")
-    os.environ.setdefault("MASTER_PORT", "29507")
-    init_distributed_environment(
-        world_size=1,
-        rank=0,
-        local_rank=0,
-        distributed_init_method="env://",
-    )
-    initialize_model_parallel()
-    yield
-    cleanup_dist_env_and_memory()
-
-
-@pytest.fixture(autouse=True)
-def _force_default_gemm(monkeypatch):
-    """Force CPU-compatible GEMM dispatch for CPU test tensors."""
-    from vllm.model_executor.layers.utils import default_unquantized_gemm
-
-    monkeypatch.setattr(
-        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
-    )
-
-
-def test_fused_qkv_load_and_split_matches_separate_projections():
-    from vllm.model_executor.layers.linear import QKVParallelLinear
-
-    torch.manual_seed(0)
-    dim, head_dim, num_heads = 16, 4, 4  # inner_dim == dim, no GQA
-
-    qkv = QKVParallelLinear(
-        hidden_size=dim,
-        head_size=head_dim,
-        total_num_heads=num_heads,
-        bias=True,
-    )
-
-    # Distinct random weight + bias per projection, loaded the exact way
-    # DreamZero.load_weights does: param.weight_loader(param, tensor, shard_id).
-    refs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
-    for shard_id in ("q", "k", "v"):
-        w = torch.randn(dim, dim)
-        b = torch.randn(dim)
-        refs[shard_id] = (w, b)
-        qkv.weight.weight_loader(qkv.weight, w, shard_id)
-        qkv.bias.weight_loader(qkv.bias, b, shard_id)
-
-    x = torch.randn(2, 3, dim)
-    fused, _ = qkv(x)
-    q, k, v = fused.split([dim, dim, dim], dim=-1)
-
-    for got, shard_id in zip((q, k, v), ("q", "k", "v"), strict=True):
-        w, b = refs[shard_id]
-        ref = torch.nn.functional.linear(x, w, b)
-        torch.testing.assert_close(got, ref)
-
-
-def test_fused_qkv_shard_ids_are_not_transposed():
-    # Guard against q/k/v being routed to the wrong slice: give each projection a
-    # constant, distinct weight so a swapped shard id would surface immediately.
-    from vllm.model_executor.layers.linear import QKVParallelLinear
-
-    dim, head_dim, num_heads = 16, 4, 4
-    qkv = QKVParallelLinear(
-        hidden_size=dim,
-        head_size=head_dim,
-        total_num_heads=num_heads,
-        bias=True,
-    )
-
-    scale = {"q": 1.0, "k": 2.0, "v": 3.0}
-    for shard_id, s in scale.items():
-        qkv.weight.weight_loader(qkv.weight, torch.full((dim, dim), s), shard_id)
-        qkv.bias.weight_loader(qkv.bias, torch.zeros(dim), shard_id)
-
-    x = torch.ones(1, dim)
-    fused, _ = qkv(x)
-    q, k, v = fused.split([dim, dim, dim], dim=-1)
-
-    # F.linear(ones, full(s)) == s * dim per element.
-    torch.testing.assert_close(q, torch.full((1, dim), scale["q"] * dim))
-    torch.testing.assert_close(k, torch.full((1, dim), scale["k"] * dim))
-    torch.testing.assert_close(v, torch.full((1, dim), scale["v"] * dim))
diff --git a/tests/e2e/offline_inference/test_moss_tts.py b/tests/e2e/offline_inference/test_moss_tts.py
index 19a8ac6d6d7..505db1af718 100644
--- a/tests/e2e/offline_inference/test_moss_tts.py
+++ b/tests/e2e/offline_inference/test_moss_tts.py
@@ -1,240 +1,225 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B).
-
-Uses the standard omni_runner + pytestmark pattern (one module-scoped engine
-per file, skill invariant I4).  The realtime variant lives in the sibling file
-test_moss_tts_realtime.py.
-
-MOSS-VoiceGenerator synthesizes speech from a text *instruction* that describes
-the desired voice (e.g. "a warm female voice with an American accent").  It does
-NOT accept a reference audio clip.  The request format requires running
-AutoProcessor.from_pretrained once per call to encode the (text, instruction)
-pair into the (prompt_token_ids, codes.ref) grid that the MossTTSDelayModel
-talker expects — same as examples/offline_inference/text_to_speech/moss_tts/
-end2end.py:_build_unified_codes.
-
-No determinism test: VoiceGenerator produces variable-length output even with
-a fixed seed; waveform length reproducibility is not guaranteed.
+"""E2E offline inference tests for MOSS-TTS two-stage pipeline.
+
+Tests the MossTTSDelayModel path using MOSS-VoiceGenerator (1.7B) — the
+smallest MossTTSDelayModel variant — so the test runs on a single A10G or L4
+without requiring an 80 GB GPU.
+
+A second test class covers MOSS-TTS-Realtime (1.7B, MossTTSRealtime) to
+exercise the local-transformer generation path.
+
+Both classes use a module-scoped engine so the model is loaded once per test
+module execution.  Do NOT add a second engine fixture in this file — that
+triggers mid-module teardown races (see skill invariant I4).
 """
 
 from __future__ import annotations
 
-import gc
 import os
+import urllib.request
 
 import pytest
 import torch
 from vllm import SamplingParams
 
 from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni import Omni
 
 # ---------------------------------------------------------------------------
-# Constants
+# Shared constants
 # ---------------------------------------------------------------------------
 
-MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator"
-SAMPLE_RATE = 24_000
+# Selected by the nightly "TTS · Function Test" (-m "full_model and L4 and tts").
+# Without the tts + run-level marks these tests are collected but never run.
+# TODO: Fix this test
+# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
+
+SAMPLE_RATE = 24_000  # All MOSS-TTS full variants output 24 kHz
 
-# Stage 0 = AR talker; max_tokens=512 keeps tests fast (~20 s of audio).
-# Stage 1 = codec decoder; greedy, large context to hold all codec tokens.
-_STAGE0_PARAMS = SamplingParams(
+REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
+
+_DEFAULT_SAMPLING = SamplingParams(
     temperature=1.7,
     top_p=0.8,
     top_k=25,
-    max_tokens=512,
-    seed=42,
-    detokenize=False,
-)
-_STAGE1_PARAMS = SamplingParams(
-    temperature=0.0,
-    top_p=1.0,
-    top_k=-1,
-    max_tokens=65536,
+    max_tokens=512,  # short for tests (~20 s at 24 kHz, ~12.5 tokens/s)
     seed=42,
     detokenize=False,
 )
-_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS]
-
-# ---------------------------------------------------------------------------
-# Deploy config
-# ---------------------------------------------------------------------------
-
-
-def _get_test_config() -> str:
-    """Derive a CI-friendly config from moss_voice_generator.yaml.
-
-    Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom
-    for Stage 1 (0.12) on a shared L4/A10G.  max_num_seqs=1 keeps per-test
-    peak memory predictable.
-    """
-    return modify_stage_config(
-        get_deploy_config_path("moss_voice_generator.yaml"),
-        updates={
-            "stages": {
-                0: {
-                    "max_num_seqs": 1,
-                    "gpu_memory_utilization": 0.45,
-                },
-            },
-        },
-    )
-
-
-# ---------------------------------------------------------------------------
-# pytestmark — one engine for the whole module
-# ---------------------------------------------------------------------------
-
-pytestmark = [
-    pytest.mark.full_model,
-    pytest.mark.tts,
-    pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True),
-]
-
 
 # ---------------------------------------------------------------------------
-# Helpers
+# Session fixtures
 # ---------------------------------------------------------------------------
 
 
-def _build_request(text: str, instruction: str) -> dict:
-    """Build a VoiceGenerator request using AutoProcessor.
+@pytest.fixture(scope="session")
+def ref_audio_path(tmp_path_factory) -> str:
+    """Download the upstream reference clip once per session.
 
-    Calls proc.build_user_message + proc() to produce the unified-codes grid
-    (L, 1+n_vq) and splits it into prompt_token_ids + codes.ref, exactly as
-    end2end.py:_build_unified_codes does for the voice_generator variant.
-    The processor (including its CPU-resident audio_tokenizer) is freed before
-    returning so it does not compete with the running engine.
-
-    Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments where
-    the HF snapshot is not cached.
+    Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments.
     """
-    from transformers import AutoProcessor
-
+    cache_dir = tmp_path_factory.mktemp("moss_tts_ref")
+    target = cache_dir / "zh_1.wav"
     try:
-        proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True)
+        with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
+            target.write_bytes(resp.read())
     except Exception as exc:
-        msg = f"Cannot load AutoProcessor for {MODEL}: {exc}"
+        msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}"
         if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
             pytest.skip(msg)
         pytest.fail(msg)
+    if not target.exists() or target.stat().st_size == 0:
+        pytest.fail(f"Reference clip empty after download: {target}")
+    return str(target)
 
-    user_msg = proc.build_user_message(text=text, instruction=instruction)
-    batch = proc(conversations=[[user_msg]], mode="generation")
-    unified = batch["input_ids"][0]  # (L, 1+n_vq)
-    text_ids = unified[:, 0].tolist()
-    audio_codes = unified[:, 1:].contiguous().to(torch.int64)  # (L, n_vq)
-    del proc
-    gc.collect()
 
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _build_request(text: str, ref_audio_path: str, seed: int = 42) -> dict:
     return {
-        "prompt_token_ids": text_ids,
-        "additional_information": {"codes": {"ref": audio_codes}},
+        "prompt": "<|im_start|>assistant\n",
+        "additional_information": {
+            "task_type": ["voice_clone"],
+            "text": [text],
+            "mode": ["voice_clone"],
+            "prompt_audio_path": [ref_audio_path],
+            "seed": [seed],
+        },
     }
 
 
-def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]:
-    """Run one request and return (waveform_cpu, sample_rate).
-
-    Iterates all OmniRequestOutputs from generate() via the top-level
-    multimodal_output property (which aggregates from completion outputs).
-    AR-stage outputs have empty multimodal_output and are silently skipped.
-    With async_chunk=True and FINAL_ONLY, the codec stage consolidates its
-    chunks into a single tensor before yielding — but we handle list fallback
-    for robustness.
-    """
-    chunks: list[torch.Tensor] = []
-    sr_final = SAMPLE_RATE
-    for out in omni_runner.omni.generate(request, _SAMPLING):
-        mm = out.multimodal_output
-        if not mm:
-            continue
-        audio = mm.get("audio")
-        if audio is None:
-            audio = mm.get("model_outputs")
-        if audio is None:
-            continue
-        sr = mm.get("sr")
-        if sr is not None:
-            sr_final = int(sr.item() if hasattr(sr, "item") else sr)
-        if isinstance(audio, list):
-            audio = torch.cat(
-                [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
-                dim=0,
-            )
-        if isinstance(audio, torch.Tensor) and audio.numel() > 0:
-            chunks.append(audio.reshape(-1).cpu())
-    if not chunks:
-        raise AssertionError("No audio output received from generate()")
-    return torch.cat(chunks, dim=0), sr_final
+def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
+    """Run one request; return (waveform_cpu, sample_rate)."""
+    for stage_outputs in omni.generate(request, _DEFAULT_SAMPLING):
+        for req_output in stage_outputs.request_output:
+            mm = req_output.outputs[0].multimodal_output
+            assert mm is not None, "Expected non-None multimodal_output"
+            # Never use ``a or b`` on tensor values — a multi-element tensor
+            # raises on truthiness. Check ``is None`` explicitly and handle the
+            # pre-consolidation list form.
+            audio = mm.get("audio")
+            if audio is None:
+                audio = mm.get("model_outputs")
+            sr = mm.get("sr")
+            assert audio is not None, "Expected 'audio' or 'model_outputs' in multimodal_output"
+            if isinstance(audio, list):
+                audio = torch.cat(
+                    [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
+                    dim=0,
+                )
+            assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}"
+            return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
+    raise AssertionError("No stage outputs received")
 
 
 # ---------------------------------------------------------------------------
-# Tests
+# MossTTSDelayModel — MOSS-VoiceGenerator 1.7B
+# (tests the delay-pattern generation path on a small GPU)
 # ---------------------------------------------------------------------------
 
+_DELAY_MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator"
+
+
+@pytest.fixture(scope="module")
+def delay_engine():
+    return Omni(model=_DELAY_MODEL, stage_init_timeout=300)
 
-@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_moss_tts_english(omni_runner: OmniRunner) -> None:
-    """VoiceGenerator: English instruction produces non-empty 24 kHz audio."""
-    req = _build_request(
-        "Hello, this is a MOSS voice design test.",
-        "a warm female voice with an American accent",
-    )
-    audio, sr = _collect_audio(omni_runner, req)
+
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_delay_english(delay_engine, ref_audio_path):
+    """MossTTSDelayModel: English voice_clone produces non-empty 24 kHz audio."""
+    req = _build_request("Hello, this is a MOSS-TTS voice cloning test.", ref_audio_path)
+    audio, sr = _collect_audio(delay_engine, req)
 
     assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
     assert audio.numel() > 0, "Audio tensor is empty"
     assert not torch.all(audio == 0), "Audio is silence"
 
 
-@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_moss_tts_chinese(omni_runner: OmniRunner) -> None:
-    """VoiceGenerator: Chinese input produces non-empty audio."""
-    req = _build_request(
-        "你好，这是语音合成测试。",
-        "一个清晰温暖的女声",
-    )
-    audio, sr = _collect_audio(omni_runner, req)
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_delay_chinese(delay_engine, ref_audio_path):
+    """MossTTSDelayModel: Chinese input produces non-empty audio."""
+    req = _build_request("你好，这是语音合成测试。", ref_audio_path)
+    audio, sr = _collect_audio(delay_engine, req)
 
     assert sr == SAMPLE_RATE
     assert audio.numel() > 0
     assert not torch.all(audio == 0)
 
 
-@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_moss_tts_batch(omni_runner: OmniRunner) -> None:
-    """VoiceGenerator: batch of two requests each returns non-empty audio."""
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path):
+    """MossTTSDelayModel: same seed yields identical waveforms."""
+    req = _build_request("Reproducibility check.", ref_audio_path, seed=99)
+    audio1, _ = _collect_audio(delay_engine, req)
+    audio2, _ = _collect_audio(delay_engine, req)
+
+    assert audio1.shape == audio2.shape, "Shapes differ across identical seeds"
+    assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms differ across identical seeds"
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_delay_batch(delay_engine, ref_audio_path):
+    """MossTTSDelayModel: batch of two requests each returns non-empty audio."""
     requests = [
-        _build_request("First sentence.", "a warm female voice"),
-        _build_request("Second sentence.", "a young male voice"),
+        _build_request("First sentence.", ref_audio_path),
+        _build_request("Second sentence.", ref_audio_path),
     ]
     results: list[torch.Tensor] = []
-    for out in omni_runner.omni.generate(requests, _SAMPLING):
-        mm = out.multimodal_output
-        if not mm:
-            continue
-        audio = mm.get("audio")
-        if audio is None:
-            audio = mm.get("model_outputs")
-        if audio is None:
-            continue
-        if isinstance(audio, list):
-            audio = torch.cat(
-                [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
-                dim=0,
-            )
-        if isinstance(audio, torch.Tensor) and audio.numel() > 0:
+    for stage_outputs in delay_engine.generate(requests, [_DEFAULT_SAMPLING] * 2):
+        for req_output in stage_outputs.request_output:
+            mm = req_output.outputs[0].multimodal_output
+            assert mm is not None
+            audio = mm.get("audio")
+            if audio is None:
+                audio = mm.get("model_outputs")
+            assert audio is not None
+            if isinstance(audio, list):
+                audio = torch.cat(
+                    [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
+                    dim=0,
+                )
             results.append(audio.reshape(-1).cpu())
 
-    # With async_chunk streaming + FINAL_ONLY, each request produces one
-    # consolidated output. The >= 2 guard handles backward-compatibility if
-    # consolidation behavior changes.
-    assert len(results) >= 2, f"Expected at least 2 audio outputs, got {len(results)}"
+    assert len(results) == 2, f"Expected 2 outputs, got {len(results)}"
     for i, audio in enumerate(results):
-        assert audio.numel() > 0, f"Audio chunk[{i}] is empty"
+        assert audio.numel() > 0, f"Audio[{i}] is empty"
+
+
+# ---------------------------------------------------------------------------
+# MossTTSRealtime — MOSS-TTS-Realtime 1.7B
+# (tests the local-transformer generation path)
+# ---------------------------------------------------------------------------
+
+_REALTIME_MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
+
+
+@pytest.fixture(scope="module")
+def realtime_engine():
+    return Omni(model=_REALTIME_MODEL, stage_init_timeout=300)
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_realtime_english(realtime_engine, ref_audio_path):
+    """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
+    req = _build_request("This is a real-time TTS streaming test.", ref_audio_path)
+    audio, sr = _collect_audio(realtime_engine, req)
+
+    assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
+    assert audio.numel() > 0, "Audio tensor is empty"
+    assert not torch.all(audio == 0), "Audio is silence"
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_realtime_deterministic(realtime_engine, ref_audio_path):
+    """MossTTSRealtime: same seed yields identical waveforms."""
+    req = _build_request("Determinism check for realtime model.", ref_audio_path, seed=7)
+    audio1, _ = _collect_audio(realtime_engine, req)
+    audio2, _ = _collect_audio(realtime_engine, req)
+
+    assert audio1.shape == audio2.shape
+    assert torch.allclose(audio1, audio2, atol=1e-4)
diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py
deleted file mode 100644
index 540c8750edb..00000000000
--- a/tests/e2e/offline_inference/test_moss_tts_realtime.py
+++ /dev/null
@@ -1,290 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B).
-
-Uses the standard omni_runner + pytestmark pattern (one module-scoped engine
-per file, skill invariant I4).  The delay-model variant lives in the sibling
-file test_moss_tts.py.
-
-The request format mirrors end2end.py:_build_realtime_prompt exactly:
-  1. snapshot_download to locate the MossTTSRealtimeProcessor module.
-  2. MOSS-Audio-Tokenizer (separate HF repo) to encode the reference clip.
-  3. Build the (L, 17) int grid; split into prompt_token_ids + codes.ref +
-     ids.all (remaining text tokens fed one-per-step during decode).
-
-Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments.
-
-No determinism test: MossTTSRealtime uses a local depth transformer that is
-sensitive to async scheduling timing; waveform reproducibility is not
-guaranteed across back-to-back calls.
-"""
-
-from __future__ import annotations
-
-import gc
-import importlib.util
-import os
-import sys
-import urllib.request
-from pathlib import Path
-
-import numpy as np
-import pytest
-import soundfile as sf
-import torch
-from vllm import SamplingParams
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
-
-# ---------------------------------------------------------------------------
-# Constants
-# ---------------------------------------------------------------------------
-
-MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
-SAMPLE_RATE = 24_000
-
-REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
-
-_STAGE0_PARAMS = SamplingParams(
-    temperature=1.0,
-    top_p=0.95,
-    top_k=50,
-    max_tokens=512,
-    seed=42,
-    detokenize=False,
-)
-_STAGE1_PARAMS = SamplingParams(
-    temperature=0.0,
-    top_p=1.0,
-    top_k=-1,
-    max_tokens=65536,
-    seed=42,
-    detokenize=False,
-)
-_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS]
-
-# ---------------------------------------------------------------------------
-# Deploy config
-# ---------------------------------------------------------------------------
-
-
-def _get_test_config() -> str:
-    """Derive a CI-friendly config from moss_tts_realtime.yaml.
-
-    Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom
-    for Stage 1 (0.12) on a shared L4/A10G.  max_num_seqs=1 keeps per-test
-    peak memory predictable.
-    """
-    return modify_stage_config(
-        get_deploy_config_path("moss_tts_realtime.yaml"),
-        updates={
-            "stages": {
-                0: {
-                    "max_num_seqs": 1,
-                    "gpu_memory_utilization": 0.45,
-                },
-            },
-        },
-    )
-
-
-# ---------------------------------------------------------------------------
-# pytestmark — one engine for the whole module
-# ---------------------------------------------------------------------------
-
-pytestmark = [
-    pytest.mark.full_model,
-    pytest.mark.tts,
-    pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True),
-]
-
-# ---------------------------------------------------------------------------
-# Fixtures
-# ---------------------------------------------------------------------------
-
-
-@pytest.fixture(scope="module")
-def ref_audio_path(tmp_path_factory) -> str:
-    """Download the upstream reference clip once per test module.
-
-    Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments.
-    """
-    target = tmp_path_factory.mktemp("moss_tts_realtime_ref") / "zh_1.wav"
-    try:
-        with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
-            target.write_bytes(resp.read())
-    except Exception as exc:
-        msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}"
-        if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
-            pytest.skip(msg)
-        pytest.fail(msg)
-    if not target.exists() or target.stat().st_size == 0:
-        pytest.fail(f"Reference clip empty after download: {target}")
-    return str(target)
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
-    """Load and resample reference audio to (1, T) float32 at target_sr."""
-    data, sr = sf.read(path, always_2d=True)  # (T, C)
-    wav = torch.from_numpy(data.T.astype("float32"))  # (C, T)
-    if wav.shape[0] > 1:
-        wav = wav.mean(dim=0, keepdim=True)
-    if sr != target_sr:
-        import torchaudio
-
-        wav = torchaudio.functional.resample(wav, sr, target_sr)
-    return wav  # (1, T)
-
-
-def _build_request(ref_audio_path: str, text: str) -> dict:
-    """Build a MOSS-TTS-Realtime request.
-
-    Mirrors end2end.py:_build_realtime_prompt exactly:
-      1. Locate MossTTSRealtimeProcessor in the HF snapshot.
-      2. Encode ref audio via MOSS-Audio-Tokenizer (separate repo).
-      3. Build the (L, 17) int grid and return prompt_token_ids +
-         additional_information: {codes.ref, ids.all}.
-
-    Frees the codec model before returning to avoid competing with the
-    running vllm engine for GPU/CPU memory.
-    """
-    from huggingface_hub import snapshot_download
-    from transformers import AutoModel, AutoTokenizer
-
-    # Step 1: locate the realtime processor module in the snapshot.
-    try:
-        snap_dir = Path(snapshot_download(repo_id=MODEL))
-    except Exception as exc:
-        msg = f"Cannot locate snapshot for {MODEL}: {exc}"
-        if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
-            pytest.skip(msg)
-        pytest.fail(msg)
-
-    proc_module_path = snap_dir / "processing_mossttsrealtime.py"
-    if not proc_module_path.exists():
-        pytest.fail(f"Realtime processor module missing at {proc_module_path}")
-
-    mod_name = "_moss_tts_realtime_proc"
-    if mod_name not in sys.modules:
-        spec = importlib.util.spec_from_file_location(mod_name, proc_module_path)
-        mod = importlib.util.module_from_spec(spec)  # type: ignore[arg-type]
-        sys.modules[mod_name] = mod
-        spec.loader.exec_module(mod)  # type: ignore[union-attr]
-    else:
-        mod = sys.modules[mod_name]
-
-    tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
-    realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
-
-    # Step 2: encode the reference audio clip via MOSS-Audio-Tokenizer.
-    try:
-        codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).to("cpu").eval()
-    except Exception as exc:
-        msg = f"Cannot load MOSS-Audio-Tokenizer: {exc}"
-        if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
-            pytest.skip(msg)
-        pytest.fail(msg)
-
-    ref_wav = _load_ref_audio(ref_audio_path)
-    wav_2d = ref_wav if ref_wav.dim() == 2 else ref_wav.unsqueeze(0)
-    with torch.no_grad():
-        enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
-    codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
-    audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy()  # (T, 16)
-    del codec
-    gc.collect()
-
-    # Step 3: build the (L, 17) prefill grid.
-    system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens)
-
-    assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
-    assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
-    assistant_grid[:, 0] = assistant_header_ids
-
-    text_ids_only = tokenizer.encode(text, add_special_tokens=False)
-    if not text_ids_only:
-        pytest.fail(f"Empty text after tokenization: {text!r}")
-
-    PREFILL_MAX_TEXT = 12
-    cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
-    text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
-    text_grid[:, 0] = text_ids_only[:cur_len]
-    text_grid[-1, 1] = 1025  # audio_bos at the last prefilled text token
-
-    grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
-    prompt_token_ids = grid[:, 0].tolist()
-    audio_codes_tensor = torch.from_numpy(grid[:, 1:].astype(np.int64).copy())  # (L, 16)
-    remaining_text_ids = list(text_ids_only[cur_len:])
-
-    additional_info: dict = {"codes": {"ref": audio_codes_tensor}}
-    if remaining_text_ids:
-        additional_info["ids"] = {"all": remaining_text_ids}
-
-    return {
-        "prompt_token_ids": prompt_token_ids,
-        "additional_information": additional_info,
-    }
-
-
-def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]:
-    """Run one request and return (waveform_cpu, sample_rate)."""
-    chunks: list[torch.Tensor] = []
-    sr_final = SAMPLE_RATE
-    for out in omni_runner.omni.generate(request, _SAMPLING):
-        mm = out.multimodal_output
-        if not mm:
-            continue
-        audio = mm.get("audio")
-        if audio is None:
-            audio = mm.get("model_outputs")
-        if audio is None:
-            continue
-        sr = mm.get("sr")
-        if sr is not None:
-            sr_final = int(sr.item() if hasattr(sr, "item") else sr)
-        if isinstance(audio, list):
-            audio = torch.cat(
-                [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
-                dim=0,
-            )
-        if isinstance(audio, torch.Tensor) and audio.numel() > 0:
-            chunks.append(audio.reshape(-1).cpu())
-    if not chunks:
-        raise AssertionError("No audio output received from generate()")
-    return torch.cat(chunks, dim=0), sr_final
-
-
-# ---------------------------------------------------------------------------
-# Tests
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None:
-    """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
-    req = _build_request(ref_audio_path, "This is a real-time TTS streaming test.")
-    audio, sr = _collect_audio(omni_runner, req)
-
-    assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
-    assert audio.numel() > 0, "Audio tensor is empty"
-    assert not torch.all(audio == 0), "Audio is silence"
-
-
-@pytest.mark.advanced_model
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None:
-    """MossTTSRealtime: Chinese input produces non-empty audio."""
-    req = _build_request(ref_audio_path, "你好，这是语音合成测试。")
-    audio, sr = _collect_audio(omni_runner, req)
-
-    assert sr == SAMPLE_RATE
-    assert audio.numel() > 0
-    assert not torch.all(audio == 0)
diff --git a/tests/e2e/online_serving/test_dreamzero_expansion.py b/tests/e2e/online_serving/test_dreamzero_expansion.py
index 79429c1235a..bc321589132 100644
--- a/tests/e2e/online_serving/test_dreamzero_expansion.py
+++ b/tests/e2e/online_serving/test_dreamzero_expansion.py
@@ -28,6 +28,8 @@
         server_args=[
             "--deploy-config",
             "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml",
+            "--enforce-eager",
+            "--disable-log-stats",
         ],
         env_dict={
             "ATTENTION_BACKEND": "torch",
diff --git a/tests/e2e/online_serving/test_gr00t_openpi.py b/tests/e2e/online_serving/test_gr00t_openpi.py
deleted file mode 100644
index 40b8293d64f..00000000000
--- a/tests/e2e/online_serving/test_gr00t_openpi.py
+++ /dev/null
@@ -1,163 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end online serving test for GR00T N1.7 through the OpenPI robot endpoint."""
-
-import os
-
-import numpy as np
-import pytest
-
-from tests.gr00t import openpi_client_helper as openpi_client
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-MODEL = "nvidia/GR00T-N1.7-3B"
-
-pytest.importorskip("websockets")
-pytest.importorskip("openpi_client.msgpack_numpy")
-
-test_params = [
-    pytest.param(
-        OmniServerParams(
-            model=MODEL,
-            stage_config_path=get_deploy_config_path("Gr00tN1d7.yaml"),
-            server_args=["--disable-log-stats"],
-            env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1", "GR00T_NOISE_SEED": "42"},
-            init_timeout=1200,
-            stage_init_timeout=900,
-        ),
-        id="gr00t-n1d7-openpi",
-    )
-]
-
-# Reference values captured from the Isaac-GR00T ZMQ server (nvidia/GR00T-N1.7-3B,
-# OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT embodiment) using build_droid_observation()
-# inputs: zero images (256×256), identity eef_9d, zero gripper/joint, "pick up the object".
-# Outputs are bit-reproducible across resets and across runs (max_diff=0.0).
-_REF_EEF_9D = np.array(
-    [
-        [
-            0.014888550154864788,
-            -0.00039258378092199564,
-            -0.013574761338531971,
-            0.9999837875366211,
-            0.005678884219378233,
-            -0.00034708858584053814,
-            -0.005677700508385897,
-            0.9999783635139465,
-            0.0033210734836757183,
-        ],
-        [
-            0.020188070833683014,
-            -0.0003105341165792197,
-            -0.02232760563492775,
-            0.9997346997261047,
-            0.02301345206797123,
-            0.000983047066256404,
-            -0.02302766777575016,
-            0.9995648860931396,
-            0.018431292846798897,
-        ],
-        [
-            -0.007266733795404434,
-            -0.05537768080830574,
-            0.03667901083827019,
-            0.992686927318573,
-            0.1206541359424591,
-            0.003906540106981993,
-            -0.12024091929197311,
-            0.9853788614273071,
-            0.12070894986391068,
-        ],
-    ],
-    dtype=np.float32,
-)  # rows = step 0, step 4, step 39
-
-_REF_GRIPPER = np.array([[0.0], [0.0078125], [0.939453125]], dtype=np.float32)  # steps 0, 4, 39
-
-_REF_JOINT = np.array(
-    [
-        [
-            -0.0010484338272362947,
-            0.0014262489276006818,
-            -0.003565810853615403,
-            -3.846167237497866e-05,
-            -0.0002604846959002316,
-            0.008521700277924538,
-            -0.006872728932648897,
-        ],
-        [
-            -0.009435676969587803,
-            0.0021475711837410927,
-            -0.0031688229646533728,
-            -1.8328893929719925e-05,
-            0.0005945992306806147,
-            0.019159257411956787,
-            0.0009468861389905214,
-        ],
-        [
-            -0.02944088727235794,
-            -0.08419207483530045,
-            -0.0251418836414814,
-            0.00540524534881115,
-            0.04752273112535477,
-            -0.012884500436484814,
-            0.024298785254359245,
-        ],
-    ],
-    dtype=np.float32,
-)  # rows = step 0, step 4, step 39
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_gr00t_n1d7_openpi_online(omni_server) -> None:
-    result = openpi_client.run_policy_session(
-        host=omni_server.host,
-        port=omni_server.port,
-        session_id="gr00t-online-e2e",
-    )
-    openpi_client.validate_session_result(result)
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_gr00t_n1d7_openpi_precision(omni_server) -> None:
-    """Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs)."""
-    client = openpi_client.OpenPIWebsocketClient(host=omni_server.host, port=omni_server.port)
-    try:
-        obs = openpi_client.build_droid_observation(session_id="gr00t-precision-e2e")
-        client.reset({})
-        actions = client.infer(obs)
-    finally:
-        client.close()
-
-    steps = [0, 4, 39]
-    np.testing.assert_allclose(
-        actions["eef_9d"][0, steps, :],
-        _REF_EEF_9D,
-        atol=1e-2,
-        rtol=0.0,
-        err_msg="eef_9d action mismatch vs Isaac-GR00T reference",
-    )
-    np.testing.assert_allclose(
-        actions["gripper_position"][0, steps, :],
-        _REF_GRIPPER,
-        atol=1e-2,
-        rtol=0.0,
-        err_msg="gripper_position action mismatch vs Isaac-GR00T reference",
-    )
-    np.testing.assert_allclose(
-        actions["joint_position"][0, steps, :],
-        _REF_JOINT,
-        atol=1e-2,
-        rtol=0.0,
-        err_msg="joint_position action mismatch vs Isaac-GR00T reference",
-    )
diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py
index 398ae649415..83d17a28cd3 100644
--- a/tests/entrypoints/openai_api/test_openpi_connection.py
+++ b/tests/entrypoints/openai_api/test_openpi_connection.py
@@ -88,47 +88,6 @@ def test_unpack_accepts_msgpack_numpy_marker_dicts():
     np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32))
 
 
-def test_unpack_accepts_openpi_client_ndarray_markers():
-    image = np.arange(6, dtype=np.uint8).reshape(2, 3)
-    payload = {
-        "observation/exterior_image_0_left": {
-            b"__ndarray__": True,
-            b"dtype": image.dtype.str,
-            b"shape": image.shape,
-            b"data": image.tobytes(),
-        }
-    }
-
-    decoded = openpi_connection._unpack_numpy(payload)
-
-    np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image)
-    assert decoded["observation/exterior_image_0_left"].dtype == np.uint8
-
-
-def test_unpack_openpi_client_packed_observation():
-    openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy")
-    image = np.zeros((180, 320, 3), dtype=np.uint8)
-    obs = {
-        "observation/exterior_image_0_left": image,
-        "observation/joint_position": np.zeros(7, dtype=np.float32),
-    }
-
-    decoded = openpi_connection._unpack(openpi_msgpack.packb(obs))
-
-    np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image)
-    np.testing.assert_allclose(decoded["observation/joint_position"], obs["observation/joint_position"])
-
-
-def test_pack_uses_openpi_client_ndarray_markers():
-    openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy")
-    actions = np.asarray([[1.0, 2.0]], dtype=np.float32)
-
-    packed = openpi_connection._pack(actions)
-    decoded = openpi_msgpack.unpackb(packed)
-
-    np.testing.assert_allclose(decoded, actions)
-
-
 def test_unpack_leaves_user_dict_without_numpy_kind_marker_unchanged():
     payload = {
         "metadata": {
diff --git a/tests/entrypoints/openai_api/test_storage_manager.py b/tests/entrypoints/openai_api/test_storage_manager.py
deleted file mode 100644
index 30e50576672..00000000000
--- a/tests/entrypoints/openai_api/test_storage_manager.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for the OpenAI video storage managers.
-"""
-
-import asyncio
-import os
-import time
-
-import pytest
-
-from vllm_omni.entrypoints.openai import storage as storage_module
-from vllm_omni.entrypoints.openai.storage import FileStorageHandle, LocalStorageManager, LocalStorageTTLManager
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_local_storage_open_returns_file_handle_for_saved_key(tmp_path):
-    storage = LocalStorageManager(storage_path=str(tmp_path / "storage"))
-    save_context = asyncio.run(storage.save(b"video-bytes", "video-123"))
-
-    handle = asyncio.run(storage.open(save_context.key))
-
-    assert isinstance(handle, FileStorageHandle)
-    assert handle.path == os.path.join(storage.storage_path, "video-123")
-    with open(handle.path, "rb") as file:
-        assert file.read() == b"video-bytes"
-
-
-def test_local_storage_open_returns_none_for_missing_key(tmp_path):
-    storage = LocalStorageManager(storage_path=str(tmp_path / "storage"))
-
-    handle = asyncio.run(storage.open("missing-video"))
-
-    assert handle is None
-
-
-def test_local_storage_ttl_save_sets_expiration_metadata(tmp_path, monkeypatch):
-    storage = LocalStorageTTLManager(
-        storage_path=str(tmp_path / "storage"),
-        max_concurrency=1,
-        ttl_seconds=60,
-        sweep_interval_seconds=300,
-    )
-    monkeypatch.setattr(storage_module.time, "time", lambda: 1_700_000_000)
-
-    save_context = asyncio.run(storage.save(b"video-bytes", "video-ttl"))
-
-    assert save_context.key == "video-ttl"
-    assert save_context.created_at == 1_700_000_000
-    assert save_context.expires_at == 1_700_000_060
-
-
-def test_local_storage_ttl_sweeper_removes_expired_file(tmp_path):
-    storage = LocalStorageTTLManager(
-        storage_path=str(tmp_path / "storage"),
-        max_concurrency=1,
-        ttl_seconds=1,
-        sweep_interval_seconds=60,
-    )
-
-    async def setup_file() -> tuple[str, str]:
-        save_context = await storage.save(b"video-bytes", "video-expired")
-        file_path = storage.get_full_file_path(save_context.key)
-        return save_context.key, file_path
-
-    storage_key, file_path = asyncio.run(setup_file())
-    expired_mtime = time.time() - 10
-    os.utime(file_path, (expired_mtime, expired_mtime))
-    assert os.path.exists(file_path)
-
-    deleted = asyncio.run(storage._sweep_once(time.time() - 1))
-
-    assert deleted == 1
-    assert not os.path.exists(file_path)
-    assert asyncio.run(storage.open(storage_key)) is None
-
-
-def test_local_storage_ttl_sweeper_keeps_files_when_path_missing(tmp_path):
-    storage = LocalStorageTTLManager(
-        storage_path=str(tmp_path / "storage"),
-        max_concurrency=1,
-        ttl_seconds=1,
-        sweep_interval_seconds=60,
-    )
-
-    missing_path = storage.get_full_file_path("video-missing")
-    assert not os.path.exists(missing_path)
-
-    deleted = asyncio.run(storage._sweep_once(time.time() - 1))
-
-    assert deleted == 0
-    assert not os.path.exists(missing_path)
diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py
index dbf76639167..b54afd54a88 100644
--- a/tests/entrypoints/openai_api/test_video_server.py
+++ b/tests/entrypoints/openai_api/test_video_server.py
@@ -1020,31 +1020,6 @@ def test_invalid_lora_returns_400(test_client):
     assert "lora object" in failed["error"]["message"].lower()
 
 
-def test_failed_generation_awaits_storage_cleanup(test_client, isolated_video_backends, mocker: MockerFixture):
-    """Regression (merge seam): when async generation raises, the failure handler
-    must ``await _cleanup_video(video_id)`` (single-arg, async) and still record
-    FAILED. Upstream carried a sync ``_cleanup_video(video_id, output_path)`` whose
-    stale call in the generic handler sat outside the conflict markers; against the
-    PR's async storage manager that raised NameError before the FAILED update,
-    wedging the job in IN_PROGRESS and orphaning the artifact."""
-    _store, _tasks, storage = isolated_video_backends
-    delete_spy = mocker.spy(storage, "delete")
-    mocker.patch.object(
-        OmniOpenAIServingVideo,
-        "generate_video_bytes",
-        side_effect=RuntimeError("GPU exploded"),
-    )
-
-    response = test_client.post("/v1/videos", data={"prompt": "will fail"})
-    assert response.status_code == 200
-    video_id = response.json()["id"]
-
-    failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value)
-    assert failed["error"]["code"] == 500
-    assert "GPU exploded" in failed["error"]["message"]
-    delete_spy.assert_called_once_with(video_id)
-
-
 def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture):
     mocker.patch.object(
         OmniOpenAIServingVideo,
@@ -1203,7 +1178,7 @@ def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: Moc
     final = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
     file_name = final["file_name"]
     assert file_name is not None
-    file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, video_id)
+    file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, file_name)
     assert os.path.exists(file_path)
 
     delete_resp = test_client.delete(f"/v1/videos/{video_id}")
@@ -1214,30 +1189,6 @@ def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: Moc
     assert not os.path.exists(file_path)
 
 
-def test_download_completed_job_uses_storage_open_and_download_name(test_client, mocker: MockerFixture):
-    video_bytes = b"stored-video-data"
-    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
-        return_value=video_bytes,
-    )
-    create_resp = test_client.post("/v1/videos", data={"prompt": "Download this video"})
-    assert create_resp.status_code == 200
-    video_id = create_resp.json()["id"]
-
-    final = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
-    file_name = final["file_name"]
-    assert file_name == f"{video_id}.mp4"
-
-    storage_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, video_id)
-    assert os.path.exists(storage_path)
-
-    response = test_client.get(f"/v1/videos/{video_id}/content")
-    assert response.status_code == 200
-    assert response.content == video_bytes
-    assert response.headers["content-type"] == "video/mp4"
-    assert file_name in response.headers["content-disposition"]
-
-
 def test_delete_in_progress_job_cancels_task_and_removes_metadata(test_client):
     handler = BlockingVideoHandler()
     test_client.app.state.openai_serving_video = handler
diff --git a/tests/gr00t/__init__.py b/tests/gr00t/__init__.py
deleted file mode 100644
index 208f01a7cb5..00000000000
--- a/tests/gr00t/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/tests/gr00t/openpi_client_helper.py b/tests/gr00t/openpi_client_helper.py
deleted file mode 100644
index e5943ca3307..00000000000
--- a/tests/gr00t/openpi_client_helper.py
+++ /dev/null
@@ -1,183 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from dataclasses import dataclass
-from typing import Any
-
-import numpy as np
-
-try:
-    import websockets.sync.client as websockets_client
-except ImportError:  # pragma: no cover - optional e2e dependency
-    websockets_client = None
-
-try:
-    from openpi_client import msgpack_numpy
-except ImportError:  # pragma: no cover - optional e2e dependency
-    msgpack_numpy = None
-
-PING_INTERVAL_SECS = 300
-PING_TIMEOUT_SECS = 3600
-DEFAULT_HOST = "127.0.0.1"
-DEFAULT_PORT = 8000
-DEFAULT_PATH = "/v1/realtime/robot/openpi"
-DEFAULT_SESSION_ID = "gr00t-smoke"
-ACTION_KEYS = {"eef_9d", "gripper_position", "joint_position"}
-LANGUAGE_KEY = "annotation.language.language_instruction"
-
-
-def _identity_eef_9d_state() -> np.ndarray:
-    state = np.zeros((1, 1, 9), dtype=np.float32)
-    state[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
-    return state
-
-
-def require_dependencies() -> None:
-    missing = []
-    if websockets_client is None:
-        missing.append("websockets")
-    if msgpack_numpy is None:
-        missing.append("openpi-client")
-    if missing:
-        raise ModuleNotFoundError(f"GR00T OpenPI test dependencies are missing: {', '.join(missing)}")
-
-
-@dataclass(frozen=True)
-class Gr00tServerMetadata:
-    action_horizon: int
-    action_keys: set[str]
-    embodiment_tag: str
-    needs_session_id: bool
-
-    @classmethod
-    def from_dict(cls, payload: dict[str, Any]):
-        required_keys = ("action_horizon", "action_keys", "embodiment_tag", "needs_session_id")
-        missing_keys = [key for key in required_keys if key not in payload]
-        if missing_keys:
-            raise ValueError(f"Missing GR00T metadata keys: {missing_keys}")
-
-        return cls(
-            action_horizon=int(payload["action_horizon"]),
-            action_keys={str(key) for key in payload["action_keys"]},
-            embodiment_tag=str(payload["embodiment_tag"]),
-            needs_session_id=bool(payload["needs_session_id"]),
-        )
-
-
-class OpenPIWebsocketClient:
-    def __init__(
-        self,
-        *,
-        host: str = DEFAULT_HOST,
-        port: int = DEFAULT_PORT,
-        path: str = DEFAULT_PATH,
-    ) -> None:
-        require_dependencies()
-        self._uri = f"ws://{host}:{port}{path}"
-        self._packer = msgpack_numpy.Packer()
-        self._ws, self._server_metadata = self._connect()
-
-    def _connect(self):
-        conn = websockets_client.connect(
-            self._uri,
-            compression=None,
-            max_size=None,
-            ping_interval=PING_INTERVAL_SECS,
-            ping_timeout=PING_TIMEOUT_SECS,
-        )
-        metadata = msgpack_numpy.unpackb(conn.recv())
-        if not isinstance(metadata, dict):
-            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
-        return conn, metadata
-
-    def get_server_metadata(self) -> dict[str, Any]:
-        return dict(self._server_metadata)
-
-    def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]:
-        payload = dict(obs)
-        payload["endpoint"] = "infer"
-        self._ws.send(self._packer.pack(payload))
-        response = msgpack_numpy.unpackb(self._ws.recv())
-        if isinstance(response, dict) and response.get("type") == "error":
-            raise RuntimeError(f"Inference failed: {response['message']}")
-        if not isinstance(response, dict):
-            raise RuntimeError(f"Expected dict actions from GR00T OpenPI endpoint, got {type(response)!r}")
-        return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
-
-    def reset(self, reset_info: dict[str, Any] | None = None) -> str:
-        payload = dict(reset_info or {})
-        payload["endpoint"] = "reset"
-        self._ws.send(self._packer.pack(payload))
-        response = msgpack_numpy.unpackb(self._ws.recv())
-        if not isinstance(response, dict) or response.get("status") != "reset successful":
-            raise RuntimeError(f"Unexpected reset response: {response!r}")
-        return str(response["status"])
-
-    def close(self) -> None:
-        self._ws.close()
-
-
-def build_droid_observation(*, session_id: str = DEFAULT_SESSION_ID) -> dict[str, Any]:
-    return {
-        "session_id": session_id,
-        "video": {
-            "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
-            "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
-        },
-        "state": {
-            "eef_9d": _identity_eef_9d_state(),
-            "gripper_position": np.zeros((1, 1, 1), dtype=np.float32),
-            "joint_position": np.zeros((1, 1, 7), dtype=np.float32),
-        },
-        "language": {LANGUAGE_KEY: [["pick up the object"]]},
-    }
-
-
-def run_policy_session(
-    *,
-    host: str = DEFAULT_HOST,
-    port: int = DEFAULT_PORT,
-    path: str = DEFAULT_PATH,
-    session_id: str = DEFAULT_SESSION_ID,
-) -> dict[str, Any]:
-    client = OpenPIWebsocketClient(host=host, port=port, path=path)
-    try:
-        metadata = client.get_server_metadata()
-        actions = client.infer(build_droid_observation(session_id=session_id))
-        reset_status = client.reset({})
-        return {
-            "metadata": metadata,
-            "actions": actions,
-            "reset_status": reset_status,
-            "session_id": session_id,
-        }
-    finally:
-        client.close()
-
-
-def validate_session_result(result: dict[str, Any]) -> None:
-    metadata = Gr00tServerMetadata.from_dict(result["metadata"])
-    if metadata.embodiment_tag != "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT":
-        raise AssertionError(f"Unexpected embodiment_tag: {metadata.embodiment_tag}")
-    if not metadata.needs_session_id:
-        raise AssertionError("GR00T test expects needs_session_id metadata")
-    if metadata.action_keys != ACTION_KEYS:
-        raise AssertionError(f"Unexpected action keys: {metadata.action_keys}")
-    if result["reset_status"] != "reset successful":
-        raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}")
-
-    actions = result["actions"]
-    if set(actions) != ACTION_KEYS:
-        raise AssertionError(f"Unexpected action keys: {set(actions)}")
-    expected_shapes = {
-        "eef_9d": (1, metadata.action_horizon, 9),
-        "gripper_position": (1, metadata.action_horizon, 1),
-        "joint_position": (1, metadata.action_horizon, 7),
-    }
-    for key, expected_shape in expected_shapes.items():
-        if actions[key].shape != expected_shape:
-            raise AssertionError(f"Action {key} shape mismatch: expected {expected_shape}, got {actions[key].shape}")
-        if actions[key].dtype != np.float32:
-            raise AssertionError(f"Action {key} dtype mismatch: expected float32, got {actions[key].dtype}")
-        if not np.isfinite(actions[key]).all():
-            raise AssertionError(f"Action {key} contains non-finite values")
diff --git a/tests/model_executor/models/moss_tts/test_per_request_codes.py b/tests/model_executor/models/moss_tts/test_per_request_codes.py
deleted file mode 100644
index eabb4fdc02f..00000000000
--- a/tests/model_executor/models/moss_tts/test_per_request_codes.py
+++ /dev/null
@@ -1,176 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Regression tests for MOSS-TTS per-request code routing in the talker.
-
-Guards the cross-request audio-corruption class (RFC #4316 Issue 1, bug report
-#4355): under continuous batching every Stage-1 request must decode its *own*
-codes, not request 0's. The talker emits ``codes["audio"]`` as a batch-ordered
-**list** so the generic splitter (``to_payload_element``) routes
-``element[idx]`` to request ``idx``; the list must stay index-aligned with the
-batch (placeholder for skipped requests) or a shorter list silently falls back
-to entry 0 — the original corruption. The hidden-row mapping must also use the
-runner's real per-request token spans, not an equal split, or mixed
-prefill+decode steps sample each request's codes from the wrong rows.
-
-CPU-only and weight-free: ``make_omni_output`` is driven directly on a bare
-talker with ``_sample_audio_codes`` mocked to echo the hidden row it was given,
-so an assertion on the output value proves which row each request sampled from.
-"""
-
-from __future__ import annotations
-
-import functools
-
-import pytest
-
-torch = pytest.importorskip("torch")
-
-N_VQ = 4
-
-
-@functools.lru_cache(maxsize=1)
-def _delay_talker_cls():
-    """Defer talker import (pulls vLLM model_executor) until first use."""
-    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
-        MossTTSDelayTalkerForGeneration,
-    )
-
-    return MossTTSDelayTalkerForGeneration
-
-
-def _make_bare_delay_talker():
-    """A talker whose only behaviour is the routing/span logic under test.
-
-    ``_sample_audio_codes`` is replaced with a stub returning an (N_VQ,) tensor
-    filled with the value of the hidden row it received. With a ``hidden`` where
-    row r holds the scalar r, the emitted codes reveal which row each request
-    sampled.
-    """
-    cls = _delay_talker_cls()
-    talker = cls.__new__(cls)
-    talker._batch_state = None
-    talker.n_vq = N_VQ
-
-    def _stub_sample(last_h, state):
-        row_value = int(last_h.reshape(-1)[0].item())
-        return torch.full((N_VQ,), row_value, dtype=torch.long)
-
-    talker._sample_audio_codes = _stub_sample
-    return talker
-
-
-def _row_indexed_hidden(num_rows: int) -> torch.Tensor:
-    """(num_rows, 1) tensor where row r holds the scalar r."""
-    return torch.arange(num_rows, dtype=torch.float32).reshape(num_rows, 1)
-
-
-def _audio_payload(result):
-    return result.multimodal_outputs["codes"]["audio"]
-
-
-class TestPerRequestRouting:
-    def test_mixed_prefill_decode_routes_by_span(self):
-        """req 0 prefills 3 rows, req 1 decodes 1 row in the same step.
-
-        With real spans, req 0 samples its last row (2) and req 1 its last row
-        (3); the payload is a batch-aligned list, element[idx] per request.
-        """
-        talker = _make_bare_delay_talker()
-        hidden = _row_indexed_hidden(4)
-        info_dicts = [{}, {}]
-        spans = [(0, 3), (3, 4)]
-
-        result = talker.make_omni_output(
-            hidden,
-            runtime_additional_information=info_dicts,
-            request_token_spans=spans,
-        )
-        audio = _audio_payload(result)
-
-        assert isinstance(audio, list)
-        assert len(audio) == 2  # one entry per request, batch-aligned
-        assert audio[0].shape == (1, N_VQ)
-        assert torch.all(audio[0] == 2)  # req 0: last row of span [0,3)
-        assert torch.all(audio[1] == 3)  # req 1: last row of span [3,4)
-        # No cross-request leakage.
-        assert not torch.equal(audio[0], audio[1])
-
-    def test_no_request_samples_outside_its_span(self):
-        talker = _make_bare_delay_talker()
-        spans = [(0, 2), (2, 5), (5, 6)]
-        hidden = _row_indexed_hidden(6)
-        info_dicts = [{}, {}, {}]
-
-        result = talker.make_omni_output(
-            hidden,
-            runtime_additional_information=info_dicts,
-            request_token_spans=spans,
-        )
-        audio = _audio_payload(result)
-
-        assert len(audio) == 3
-        sampled = []
-        for i, (start, end) in enumerate(spans):
-            value = int(audio[i].reshape(-1)[0].item())
-            assert start <= value < end, f"req {i} sampled row {value} outside [{start},{end})"
-            sampled.append(value)
-        assert len(set(sampled)) == len(sampled)  # distinct rows, no sharing
-
-    def test_skipped_request_keeps_list_aligned(self):
-        """An empty-span request gets an empty placeholder at its own index, so
-        later requests are not shifted onto request 0's codes."""
-        talker = _make_bare_delay_talker()
-        # req 1 has an empty span (start == end) and must not displace req 2.
-        spans = [(0, 1), (1, 1), (1, 2)]
-        hidden = _row_indexed_hidden(2)
-        info_dicts = [{}, {}, {}]
-
-        result = talker.make_omni_output(
-            hidden,
-            runtime_additional_information=info_dicts,
-            request_token_spans=spans,
-        )
-        audio = _audio_payload(result)
-
-        assert len(audio) == 3
-        assert torch.all(audio[0] == 0)  # req 0: row 0
-        assert audio[1].numel() == 0  # req 1: empty placeholder, aligned
-        assert torch.all(audio[2] == 1)  # req 2: row 1, not shifted
-
-    def test_accumulates_into_per_request_history(self):
-        talker = _make_bare_delay_talker()
-        hidden = _row_indexed_hidden(2)
-        prior = torch.full((1, N_VQ), 7, dtype=torch.long)
-        info_dicts = [{"audio_codes": {"accumulated": prior}}, {}]
-        spans = [(0, 1), (1, 2)]
-
-        result = talker.make_omni_output(
-            hidden,
-            runtime_additional_information=info_dicts,
-            request_token_spans=spans,
-        )
-        audio = _audio_payload(result)
-
-        assert audio[0].shape == (2, N_VQ)
-        assert torch.all(audio[0][0] == 7)  # prior history kept
-        assert torch.all(audio[0][1] == 0)  # this step's row 0 appended
-        assert torch.all(audio[1] == 1)
-
-    def test_runs_without_spans_backward_compat(self):
-        """Missing request_token_spans must not crash (equal-split fallback)
-        and still yields a batch-aligned list."""
-        talker = _make_bare_delay_talker()
-        hidden = _row_indexed_hidden(2)
-        info_dicts = [{}, {}]
-
-        result = talker.make_omni_output(hidden, runtime_additional_information=info_dicts)
-        audio = _audio_payload(result)
-        assert isinstance(audio, list)
-        assert len(audio) == 2
-
-    def test_no_codes_returns_empty_payload(self):
-        """Zero hidden rows → no codes → empty multimodal payload (no list)."""
-        talker = _make_bare_delay_talker()
-        hidden = torch.zeros((0, 1), dtype=torch.float32)
-        result = talker.make_omni_output(hidden, runtime_additional_information=[{}], request_token_spans=[(0, 0)])
-        assert result.multimodal_outputs == {}
diff --git a/tests/test_omni_request.py b/tests/test_omni_request.py
deleted file mode 100644
index 09f893a0399..00000000000
--- a/tests/test_omni_request.py
+++ /dev/null
@@ -1,89 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Regression tests for OmniRequest substitutability with the base Request.
-
-vllm-omni rebinds ``vllm.v1.request.Request`` to ``OmniRequest`` at import
-time (see ``vllm_omni/patch.py``). vLLM core constructs ``Request`` positionally
-in some paths (notably ``vllm/v1/worker/gpu/warmup.py::warmup_kernels`` for V2
-model-runner architectures like Qwen3ForCausalLM). The omni-specific params
-must therefore be keyword-only so positional construction stays Liskov-
-substitutable with the base class — including base-style calls that pass
-``prompt_embeds`` (itself a positional-capable base param) positionally.
-"""
-
-import inspect
-
-import numpy as np
-import pytest
-import torch
-from vllm.sampling_params import SamplingParams
-from vllm.v1.request import Request
-
-from vllm_omni.engine import PromptEmbedsPayload
-from vllm_omni.request import OmniRequest
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_omni_params_are_keyword_only():
-    """The three omni params must be keyword-only after ``*args``.
-
-    Guards against re-introducing the import-time rebind bug where omni params
-    came first positionally and broke positional ``Request(...)`` construction.
-    """
-    params = inspect.signature(OmniRequest.__init__).parameters
-    for name in ("prompt_embeds", "external_req_id", "additional_information"):
-        assert params[name].kind is inspect.Parameter.KEYWORD_ONLY, name
-
-
-def test_positional_construction_matches_base_request():
-    """Reproduces the warmup_kernels call: Request(id, tokens, sp, pp)."""
-    req = OmniRequest("req-0", [1, 2, 3], SamplingParams(), None)
-
-    assert isinstance(req, Request)
-    assert req.request_id == "req-0"
-    assert req.prompt_token_ids == [1, 2, 3]
-    # Omni params default cleanly when constructed positionally.
-    assert req.external_req_id is None
-    assert req.additional_information is None
-    assert req.prompt_embeds_payload is None
-
-
-def test_positional_prompt_embeds_does_not_collide():
-    """Base-style positional call that includes ``prompt_embeds``.
-
-    ``prompt_embeds`` is a positional-capable base param (after
-    ``arrival_time``). A base-style call must not trip
-    ``got multiple values for argument 'prompt_embeds'`` — the subclass must
-    not also inject its keyword override when the value arrived positionally.
-    Positional order: request_id, prompt_token_ids, sampling_params,
-    pooling_params, client_index, arrival_time, prompt_embeds.
-    """
-    embeds = torch.arange(6, dtype=torch.float32).reshape(2, 3)
-    req = OmniRequest("req-pe", [1, 2], SamplingParams(), None, 0, None, embeds)
-
-    assert isinstance(req, Request)
-    assert torch.equal(req.prompt_embeds, embeds)
-    # Positional tensor is not a serialized payload.
-    assert req.prompt_embeds_payload is None
-    assert req.external_req_id is None
-
-
-def test_keyword_omni_params_round_trip():
-    """Keyword omni params are preserved; serialized embeds are decoded."""
-    arr = np.arange(6, dtype=np.float32).reshape(2, 3)
-    payload = PromptEmbedsPayload(data=arr.tobytes(), shape=[2, 3], dtype="float32")
-
-    req = OmniRequest(
-        request_id="req-1",
-        prompt_token_ids=[7, 8],
-        sampling_params=SamplingParams(),
-        pooling_params=None,
-        prompt_embeds=payload,
-        external_req_id="ext-1",
-    )
-
-    assert req.external_req_id == "ext-1"
-    # The serialized payload is retained and decoded into a tensor on the base.
-    assert req.prompt_embeds_payload is payload
-    assert torch.equal(req.prompt_embeds, torch.from_numpy(arr))
diff --git a/vllm_omni/bde/__init__.py b/vllm_omni/bde/__init__.py
new file mode 100644
index 00000000000..4dfcc40d2be
--- /dev/null
+++ b/vllm_omni/bde/__init__.py
@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Block Diffusion Engine (BDE).
+
+The AR-Diffusion engine: a ``DiffusionEngine`` subclass that adds engine-level KV
+cache management for autoregressive / chunked "world-model" diffusion models
+(DreamZero and the AR-DiT family). Selected via
+``OmniDiffusionConfig.engine_backend = "bde"``.
+"""
+
+from vllm_omni.bde.engine import BDEEngine
+
+__all__ = ["BDEEngine"]
diff --git a/vllm_omni/bde/engine.py b/vllm_omni/bde/engine.py
new file mode 100644
index 00000000000..69e0c1560c7
--- /dev/null
+++ b/vllm_omni/bde/engine.py
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The Block Diffusion Engine (BDE)."""
+
+from __future__ import annotations
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+logger = init_logger(__name__)
+
+
+class BDEEngine(DiffusionEngine):
+    """AR-Diffusion engine with engine-level KV cache management.
+
+    BDE serves autoregressive / chunked blockwise-causal diffusion models
+    (world models, AR-DiT) that materialize persistent attention KV. It reuses
+    vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` / ``BlockTables``)
+    as a library, driven from the engine rather than hand-rolled inside each
+    model. See ``BDE_doc/diffusion_kv_cache_management.md`` for the design and
+    ``BDE_doc/dreamzero_kv_phase1_plan.md`` for the rollout.
+
+    It is selected per model via ``OmniDiffusionConfig.engine_backend = "bde"``
+    (resolved by :meth:`DiffusionEngine.make_engine`), so models that do not opt
+    in keep using the base ``DiffusionEngine`` unchanged.
+
+    Phase 1 is a behavior-preserving subclass: it establishes the engine seam and
+    isolation boundary. The KV-cache lifecycle (pool construction, per-chunk
+    ``allocate_slots`` / ``free``) is added in later phases.
+    """
diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py
index 032bcbc1017..106b2bf12b8 100644
--- a/vllm_omni/config/pipeline_registry.py
+++ b/vllm_omni/config/pipeline_registry.py
@@ -77,10 +77,6 @@
         "vllm_omni.model_executor.models.glm_image.pipeline",
         "GLM_IMAGE_PIPELINE",
     ),
-    "Gr00tN1d7": (
-        "vllm_omni.model_executor.models.gr00t.pipeline",
-        "GR00T_N1D7_PIPELINE",
-    ),
     "hunyuan_image_3_moe": (
         "vllm_omni.model_executor.models.hunyuan_image3.pipeline",
         "HUNYUAN_IMAGE3_PIPELINE",
diff --git a/vllm_omni/config/server_settings.py b/vllm_omni/config/server_settings.py
deleted file mode 100644
index bb63712ba10..00000000000
--- a/vllm_omni/config/server_settings.py
+++ /dev/null
@@ -1,58 +0,0 @@
-import os
-import warnings
-from typing import Literal
-
-from pydantic import BaseModel, Field, field_validator, model_validator
-from pydantic_settings import BaseSettings, SettingsConfigDict
-
-
-def _resolve_deprecated_env(old_name: str, new_name: str, value):
-    old = os.environ.get(old_name)
-    if old is None:
-        return value
-    warnings.warn(
-        f"{old_name} is deprecated and will be removed in a future release; use {new_name} instead.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-    return value if new_name in os.environ else old
-
-
-class FileBackend(BaseModel):
-    type: Literal["file"] = "file"
-    path: str = Field(default="/tmp/storage", validate_default=True, description="Local path to store completed files.")
-    file_concurrency: int = Field(
-        default=4, validate_default=True, description="Maximum number of file operations permitted at a time"
-    )
-    file_ttl: int | None = Field(
-        default=None, description="Optional TTL (in seconds) configuration settings for locally stored files."
-    )
-    ttl_sweep_interval: int | None = Field(
-        default=None, description="Optional frequency (in seconds) to enforce file TTLs."
-    )
-
-    @field_validator("path", mode="before")
-    @classmethod
-    def _migrate_path(cls, value):
-        return _resolve_deprecated_env("VLLM_OMNI_STORAGE_PATH", "VLLM_OMNI_SERVER_STORAGE__PATH", value)
-
-    @field_validator("file_concurrency", mode="before")
-    @classmethod
-    def _migrate_file_concurrency(cls, value):
-        return _resolve_deprecated_env(
-            "VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", value
-        )
-
-    @model_validator(mode="after")
-    def set_default_ttl_sweep_interval(self) -> "FileBackend":
-        if self.file_ttl is not None and self.ttl_sweep_interval is None:
-            self.ttl_sweep_interval = 300
-        return self
-
-
-class ServerSettings(BaseSettings):
-    model_config = SettingsConfigDict(env_prefix="VLLM_OMNI_SERVER_", env_nested_delimiter="__")
-    storage: FileBackend = Field(default_factory=FileBackend)
-
-
-SERVER_SETTINGS_CONFIG = ServerSettings()
diff --git a/vllm_omni/deploy/Gr00tN1d7.yaml b/vllm_omni/deploy/Gr00tN1d7.yaml
deleted file mode 100644
index 0e90d8536b7..00000000000
--- a/vllm_omni/deploy/Gr00tN1d7.yaml
+++ /dev/null
@@ -1,41 +0,0 @@
-# GR00T N1.7 deploy: single diffusion stage.
-#
-# Topology is declared in vllm_omni/model_executor/models/gr00t/pipeline.py.
-# This default uses one GPU with TP=1 and CFG parallel disabled.
-
-pipeline: Gr00tN1d7
-# Required: async_chunk defaults to True, which the deploy validator rejects for a
-# single-stage pipeline (no next-stage processor). Must be false here.
-async_chunk: false
-
-stages:
-  - stage_id: 0
-    max_num_seqs: 1
-    enforce_eager: true
-    model_class_name: Gr00tN1d7Pipeline
-    model_config:
-      embodiment_tag: OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT
-      # OpenPI client handshake — model/embodiment-specific values for GR00T-N1.7.
-      # Sent verbatim to the client (the pipeline never reads these); they describe the
-      # action space / inputs the client must produce for this checkpoint + embodiment.
-      policy_server_config:
-        image_resolution: [180, 320]
-        needs_session_id: true
-        embodiment_tag: OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT
-        action_horizon: 40
-        action_keys:
-          - eef_9d
-          - gripper_position
-          - joint_position
-        supported_embodiments:
-          - OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT
-          - XDOF
-          - XDOF_SUBTASK
-          - REAL_G1
-          - REAL_R1_PRO_SHARPA
-          - REAL_R1_PRO_SHARPA_HUMAN
-          - REAL_R1_PRO_SHARPA_MAXINSIGHTS
-          - REAL_R1_PRO_SHARPA_MECKA
-          - LIBERO_PANDA
-          - SIMPLER_ENV_GOOGLE
-          - SIMPLER_ENV_WIDOWX
diff --git a/vllm_omni/deploy/dreamzero.yaml b/vllm_omni/deploy/dreamzero.yaml
index 079101ba0e0..e77afad8409 100644
--- a/vllm_omni/deploy/dreamzero.yaml
+++ b/vllm_omni/deploy/dreamzero.yaml
@@ -12,7 +12,7 @@ stages:
   - stage_id: 0
     devices: "0"
     max_num_seqs: 1
-    enforce_eager: false
+    enforce_eager: true
     model_class_name: DreamZeroPipeline
     model_config:
       default_robot_embodiment: roboarena
@@ -23,8 +23,3 @@ stages:
         needs_stereo_camera: false
         needs_session_id: true
         action_space: joint_position
-    cache_backend: step_cache
-    cache_config:
-      step_cache_dit_enabled: true
-      velocity_sim_thresholds: [0.95, 0.93]
-      velocity_skip_countdowns: [4, 2]
diff --git a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
index a453d43f3bf..ed9988c92a3 100644
--- a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
+++ b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
@@ -9,7 +9,7 @@ stages:
   - stage_id: 0
     devices: "0,1"
     max_num_seqs: 1
-    enforce_eager: false
+    enforce_eager: true
     model_class_name: DreamZeroPipeline
     parallel_config:
       tensor_parallel_size: 1
@@ -23,8 +23,3 @@ stages:
         needs_stereo_camera: false
         needs_session_id: true
         action_space: joint_position
-    cache_backend: step_cache
-    cache_config:
-      step_cache_dit_enabled: true
-      velocity_sim_thresholds: [0.95, 0.93]
-      velocity_skip_countdowns: [4, 2]
diff --git a/vllm_omni/diffusion/cache/base.py b/vllm_omni/diffusion/cache/base.py
index aafbf0ba7a3..d7807622497 100644
--- a/vllm_omni/diffusion/cache/base.py
+++ b/vllm_omni/diffusion/cache/base.py
@@ -13,8 +13,6 @@
    the cache-dit library. Inherits from CacheBackend. Used via cache_backend="cache_dit".
 2. TeaCacheBackend: Hook-based backend for TeaCache acceleration. Inherits from
    CacheBackend. Used via cache_backend="tea_cache".
-3. StepCacheBackend: Velocity cosine step-skipping for DreamZero. Inherits from
-   CacheBackend. Used via cache_backend="step_cache".
 
 All backends implement the same interface:
 - enable(pipeline): Enable cache on the pipeline
diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py
index 5add71daa61..ff1eea6aeaa 100644
--- a/vllm_omni/diffusion/cache/cache_dit_backend.py
+++ b/vllm_omni/diffusion/cache/cache_dit_backend.py
@@ -10,19 +10,11 @@
 import functools
 from collections.abc import Callable
 from contextlib import ExitStack
-from dataclasses import dataclass
-from typing import Any, Optional, TypeAlias
+from typing import Any, Optional
 
 import cache_dit
 import torch
-from cache_dit import (
-    BlockAdapter,
-    CalibratorConfig,
-    DBCacheConfig,
-    ForwardPattern,
-    ParamsModifier,
-    TaylorSeerCalibratorConfig,
-)
+from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifier, TaylorSeerCalibratorConfig
 from cache_dit.caching.block_adapters import FakeDiffusionPipeline
 from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
@@ -33,29 +25,10 @@
 
 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline
 
 logger = init_logger(__name__)
 
-RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None]
-
-
-@dataclass
-class CacheDiTAdapterConfig:
-    """Config for creating a Cache DiT's block adapter; to enable CacheDiT,
-    most models just need to define an instance of this class as a class
-    var in the DiT.
-    """
-
-    block_forward_patterns: dict[str, ForwardPattern]
-    has_separate_cfg: bool = False
-    cached_adapter_cls: type[CachedAdapter] | None = None
-    check_forward_pattern: bool = True
-
-
-# Registry of custom cache-dit enablers for specific models
-# Maps pipeline names to their cache-dit enablement functions
-CUSTOM_DIT_ENABLERS: dict[str, Callable] = {}
-
 
 # Small helper to centralize cache-dit summaries.
 def cache_summary(pipeline: Any, details: bool = True) -> None:
@@ -64,60 +37,14 @@ def cache_summary(pipeline: Any, details: bool = True) -> None:
         cache_dit.summary(pipeline.transformer_2, details=details)
 
 
-def default_get_pipeline_transformer(pipeline: Any) -> Any:
-    return pipeline.transformer
-
-
-def build_cache_context_refresh(
-    cache_config: DiffusionCacheConfig,
-    get_pipeline_transformer: Callable[[Any], Any] = default_get_pipeline_transformer,
-) -> RefreshCacheContextFunc:
-    """Build the cache context refresh func for a single Transformer."""
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
-
-        Args:
-            pipeline: The pipeline instance.
-            num_inference_steps: New number of inference steps.
-        """
-        transformer = get_pipeline_transformer(pipeline)
-
-        # Bypass SCM for step counts that don't support predefined masks (e.g., vLLM's 1-step dummy run)
-        scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6)
-
-        if cache_config.scm_steps_mask_policy is None or not scm_supported_steps:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-    return refresh_cache_context
-
-
-def _resolve_calibrator_config(cache_config: DiffusionCacheConfig) -> CalibratorConfig | None:
-    """Resolves the Calibrator subconfig For DiT cache."""
-    calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-    # In the future, more calibrators will likely be added to DiT Cache,
-    # e.g., focal; handle them generically here.
-    return calibrator
+# Registry of custom cache-dit enablers for specific models
+# Maps pipeline names to their cache-dit enablement functions
+# Models in this registry require custom handling (e.g., dual-transformer architectures)
+# Will be populated after function definitions
+CUSTOM_DIT_ENABLERS: dict[str, Callable] = {}
 
 
-def _build_db_cache_config(cache_config: DiffusionCacheConfig) -> DBCacheConfig:
+def _build_db_cache_config(cache_config: Any) -> DBCacheConfig:
     """Build DBCacheConfig with optional SCM (Step Computation Masking) support.
 
     Args:
@@ -126,8 +53,9 @@ def _build_db_cache_config(cache_config: DiffusionCacheConfig) -> DBCacheConfig:
     Returns:
         DBCacheConfig instance with SCM support if configured.
     """
+
     return DBCacheConfig(
-        # we will refresh the context when we get num_inference_steps in the first inference request
+        # we will refresh the context when gets num_inference_steps in the first inference request
         num_inference_steps=None,
         Fn_compute_blocks=cache_config.Fn_compute_blocks,
         Bn_compute_blocks=cache_config.Bn_compute_blocks,
@@ -140,100 +68,7 @@ def _build_db_cache_config(cache_config: DiffusionCacheConfig) -> DBCacheConfig:
     )
 
 
-def enable_cache_for_dit(
-    pipeline: Any,
-    cache_config: Any,
-    block_adapter: BlockAdapter | None = None,
-    adapter_cls: type[CachedAdapter] | None = None,
-) -> RefreshCacheContextFunc:
-    """Enable cache-dit for regular single-transformer DiT models.
-
-    Args:
-        pipeline: The diffusion pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-        block_adapter: Custom block adapters for specific model architectures.
-        adapter_cls: Custom cached adapter class for specific model architectures.
-
-    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
-    """
-    # Build DBCacheConfig with optional SCM support
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    # Build calibrator config if TaylorSeer is enabled
-    calibrator_config = _resolve_calibrator_config(cache_config)
-
-    logger.info(
-        f"Enabling cache-dit on transformer: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    # Enable cache-dit on the transformer
-    transformer = default_get_pipeline_transformer(pipeline)
-
-    # If we have a custom cached adapter subclass, call apply directly
-    if adapter_cls is not None:
-        adapter_cls.apply(
-            transformer if block_adapter is None else block_adapter,
-            cache_config=db_cache_config,
-            calibrator_config=calibrator_config,
-        )
-    else:
-        # Otherwise, call enable cache, which will call CachedAdapter.apply for us
-        cache_dit.enable_cache(
-            transformer if block_adapter is None else block_adapter,
-            cache_config=db_cache_config,
-            calibrator_config=calibrator_config,
-        )
-
-    return build_cache_context_refresh(cache_config)
-
-
-### Complex / custom enablers for DiT cache
-# NOTE (Alex): This case is rare; you should only really need to do this if you have a dual transformer
-# architecture, since it hasn't been handled generically yet, or if the model class has unique attributes
-# that it sets during Cache DiT enablement.
-#
-# For the vast majority of models, you should only need to add a _cache_dit_adapter_config attribute
-# to the Transformer class, which controls the forward pattern, whether or not we have separate CFG,
-# and so on.
-
-
-# from https://github.com/vipshop/cache-dit/pull/542
-def _split_wan22_inference_steps(pipeline, num_inference_steps: int) -> tuple[int, int]:
-    """Split inference steps into high-noise and low-noise steps for Wan2.2.
-
-    This is an internal helper function specific to Wan2.2's dual-transformer
-    architecture that uses boundary_ratio to determine the split point.
-
-    Args:
-        num_inference_steps: Total number of inference steps.
-
-    Returns:
-        A tuple of (num_high_noise_steps, num_low_noise_steps).
-    """
-    if pipeline.boundary_ratio is not None:
-        boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps
-    else:
-        boundary_timestep = None
-
-    # Set timesteps to calculate the split
-    device = next(pipeline.transformer.parameters()).device
-    pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
-
-    timesteps = pipeline.scheduler.timesteps
-    num_high_noise_steps = 0  # high-noise steps for transformer
-    for t in timesteps:
-        if boundary_timestep is None or t >= boundary_timestep:
-            num_high_noise_steps += 1
-    # low-noise steps for transformer_2
-    num_low_noise_steps = num_inference_steps - num_high_noise_steps
-    return num_high_noise_steps, num_low_noise_steps
-
-
-def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc:
+def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for Wan2.2 single or dual-transformer architecture.
 
     Wan2.2 can use single or dual transformers (transformer and transformer_2) that need
@@ -246,11 +81,10 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
     Returns:
         A refresh function that can be called to update cache context with new num_inference_steps.
     """
-    # Build DBCacheConfig with optional SCM support
-    db_cache_config = _build_db_cache_config(cache_config)
 
     if getattr(pipeline, "transformer_2", None) is None:
         logger.info("transformer_2 not found, enabling cache-dit for single transformer mode")
+        db_cache_config = _build_db_cache_config(cache_config)
         cache_dit.enable_cache(
             BlockAdapter(
                 transformer=pipeline.transformer,
@@ -263,7 +97,27 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
             ),
             cache_config=db_cache_config,
         )
-        return build_cache_context_refresh(cache_config)
+
+        def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+            """Refresh cache context for single transformer."""
+            if cache_config.scm_steps_mask_policy is None:
+                cache_dit.refresh_context(
+                    pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+                )
+            else:
+                cache_dit.refresh_context(
+                    pipeline.transformer,
+                    cache_config=DBCacheConfig().reset(
+                        num_inference_steps=num_inference_steps,
+                        steps_computation_mask=cache_dit.steps_mask(
+                            mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps
+                        ),
+                        steps_computation_policy=cache_config.scm_steps_policy,
+                    ),
+                    verbose=verbose,
+                )
+
+        return refresh_cache_context
 
     cache_dit.enable_cache(
         BlockAdapter(
@@ -296,11 +150,47 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheCont
             ],
             has_separate_cfg=True,
         ),
-        cache_config=db_cache_config,
+        cache_config=DBCacheConfig(
+            Fn_compute_blocks=cache_config.Fn_compute_blocks,
+            Bn_compute_blocks=cache_config.Bn_compute_blocks,
+            max_warmup_steps=cache_config.max_warmup_steps,
+            max_cached_steps=cache_config.max_cached_steps,
+            max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+            residual_diff_threshold=cache_config.residual_diff_threshold,
+            num_inference_steps=None,
+        ),
     )
 
-    refresh_trans_one = build_cache_context_refresh(cache_config)
-    refresh_trans_two = build_cache_context_refresh(cache_config, lambda pipeline: pipeline.transformer_2)
+    # from https://github.com/vipshop/cache-dit/pull/542
+    def _split_inference_steps(num_inference_steps: int) -> tuple[int, int]:
+        """Split inference steps into high-noise and low-noise steps for Wan2.2.
+
+        This is an internal helper function specific to Wan2.2's dual-transformer
+        architecture that uses boundary_ratio to determine the split point.
+
+        Args:
+            num_inference_steps: Total number of inference steps.
+
+        Returns:
+            A tuple of (num_high_noise_steps, num_low_noise_steps).
+        """
+        if pipeline.boundary_ratio is not None:
+            boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps
+        else:
+            boundary_timestep = None
+
+        # Set timesteps to calculate the split
+        device = next(pipeline.transformer.parameters()).device
+        pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        timesteps = pipeline.scheduler.timesteps
+        num_high_noise_steps = 0  # high-noise steps for transformer
+        for t in timesteps:
+            if boundary_timestep is None or t >= boundary_timestep:
+                num_high_noise_steps += 1
+        # low-noise steps for transformer_2
+        num_low_noise_steps = num_inference_steps - num_high_noise_steps
+        return num_high_noise_steps, num_low_noise_steps
 
     def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
         """Refresh cache context for both transformers with new num_inference_steps.
@@ -309,14 +199,50 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool
             pipeline: The Wan2.2 pipeline instance.
             num_inference_steps: New number of inference steps.
         """
-        num_high_noise_steps, num_low_noise_steps = _split_wan22_inference_steps(pipeline, num_inference_steps)
-        refresh_trans_one(pipeline, num_high_noise_steps, verbose)
-        refresh_trans_two(pipeline, num_low_noise_steps, verbose)
+
+        num_high_noise_steps, num_low_noise_steps = _split_inference_steps(num_inference_steps)
+        # Refresh context for high-noise transformer
+        if cache_config.scm_steps_mask_policy is None:
+            # cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_high_noise_steps, verbose=verbose)
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                num_inference_steps=num_high_noise_steps,
+                verbose=verbose,
+            )
+            cache_dit.refresh_context(
+                pipeline.transformer_2,
+                num_inference_steps=num_low_noise_steps,
+                verbose=verbose,
+            )
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_high_noise_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_high_noise_steps
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+            cache_dit.refresh_context(
+                pipeline.transformer_2,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_low_noise_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_low_noise_steps
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
 
     return refresh_cache_context
 
 
-def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc:
+def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for Wan2.2 S2V.
 
     S2V uses a single transformer, but unlike the other Wan2.2 variants its
@@ -390,550 +316,1635 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool
     return refresh_cache_context
 
 
-class BagelCachedContextManager(CachedContextManager):
-    """
-    Custom CachedContextManager for Bagel that safely handles NaiveCache objects
-    (mapped to encoder_hidden_states) by skipping tensor operations on them.
+def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for LongCatImage pipeline.
+
+    Args:
+        pipeline: The LongCatImage pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
     """
+    # Build DBCacheConfig for transformer
+    db_cache_config = _build_db_cache_config(cache_config)
 
-    @torch.compiler.disable
-    def apply_cache(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
-        prefix: str = "Bn",
-        encoder_prefix: str = "Bn_encoder",
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        # Allow Bn and Fn prefix to be used for residual cache.
-        if "Bn" in prefix:
-            hidden_states_prev = self.get_Bn_buffer(prefix)
-        else:
-            hidden_states_prev = self.get_Fn_buffer(prefix)
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-        assert hidden_states_prev is not None, f"{prefix}_buffer must be set before"
+    # Build ParamsModifier for transformer
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
 
-        if self.is_cache_residual():
-            hidden_states = hidden_states_prev + hidden_states
-        else:
-            # If cache is not residual, we use the hidden states directly
-            hidden_states = hidden_states_prev
+    logger.info(
+        f"Enabling cache-dit on LongCatImage transformer with BlockAdapter: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
 
-        hidden_states = hidden_states.contiguous()
+    # Enable cache-dit using BlockAdapter for transformer
+    cache_dit.enable_cache(
+        (
+            BlockAdapter(
+                transformer=pipeline.transformer,
+                blocks=[
+                    pipeline.transformer.transformer_blocks,
+                    pipeline.transformer.single_transformer_blocks,
+                ],
+                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+                params_modifiers=[modifier],
+                has_separate_cfg=True,
+            )
+        ),
+        cache_config=db_cache_config,
+    )
 
-        if encoder_hidden_states is not None:
-            if "Bn" in encoder_prefix:
-                encoder_hidden_states_prev = self.get_Bn_encoder_buffer(encoder_prefix)
-            else:
-                encoder_hidden_states_prev = self.get_Fn_encoder_buffer(encoder_prefix)
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
 
-            if encoder_hidden_states_prev is not None:
-                if self.is_encoder_cache_residual():
-                    # FIX: Check if encoder_hidden_states is a tensor before adding
-                    if isinstance(encoder_hidden_states, torch.Tensor) and isinstance(
-                        encoder_hidden_states_prev, torch.Tensor
-                    ):
-                        encoder_hidden_states = encoder_hidden_states_prev + encoder_hidden_states
-                else:
-                    # If encoder cache is not residual, we use the encoder hidden states directly
-                    encoder_hidden_states = encoder_hidden_states_prev
+        Args:
+            pipeline: The LongCatImage pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
 
-            # FIX: Check if encoder_hidden_states is a tensor before calling contiguous
-            if isinstance(encoder_hidden_states, torch.Tensor):
-                encoder_hidden_states = encoder_hidden_states.contiguous()
+    return refresh_cache_context
 
-        return hidden_states, encoder_hidden_states
 
+def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Flux.1-dev pipeline.
 
-class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2):
-    """
-    Custom CachedBlocks for Bagel that safely handles NaiveCache objects
-    by adding isinstance checks in call_Mn_blocks and compute_or_prune.
+    Args:
+        pipeline: The Flux pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+    Returns:
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
     """
+    # Build DBCacheConfig for transformer
+    db_cache_config = _build_db_cache_config(cache_config)
 
-    def call_Mn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
-        for block in self._Mn_blocks():
-            hidden_states = block(
-                hidden_states,
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
-            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
-
-        # compute hidden_states residual
-        hidden_states = hidden_states.contiguous()
-
-        hidden_states_residual = hidden_states - original_hidden_states
-
-        if (
-            encoder_hidden_states is not None
-            and original_encoder_hidden_states is not None
-            and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
-        ):
-            encoder_hidden_states = encoder_hidden_states.contiguous()
-            encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
-        else:
-            encoder_hidden_states_residual = None
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-        return (
-            hidden_states,
-            encoder_hidden_states,
-            hidden_states_residual,
-            encoder_hidden_states_residual,
-        )
+    # Build ParamsModifier for transformer
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
 
-    def compute_or_prune(
-        self,
-        block_id: int,  # Block index in the transformer blocks
-        # Below are the inputs to the block
-        block,  # The transformer block to be executed
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        # NOTE: Although Bagel likely won't use pruning, implementing safe version just in case.
-        # Copy-pasted from original but adding checks.
+    logger.info(
+        f"Enabling cache-dit on Flux transformer with BlockAdapter: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
 
-        original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
+    # Enable cache-dit using BlockAdapter for transformer
+    cache_dit.enable_cache(
+        (
+            BlockAdapter(
+                transformer=pipeline.transformer,
+                blocks=[
+                    pipeline.transformer.transformer_blocks,
+                    pipeline.transformer.single_transformer_blocks,
+                ],
+                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+                params_modifiers=[modifier],
+            )
+        ),
+        cache_config=db_cache_config,
+    )
 
-        can_use_prune = self._maybe_prune(
-            block_id,
-            hidden_states,
-            prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
-        )
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
 
-        torch._dynamo.graph_break()
-        if can_use_prune:
-            self.context_manager.add_pruned_step()
-            hidden_states, encoder_hidden_states = self.context_manager.apply_prune(
-                hidden_states,
-                encoder_hidden_states,
-                prefix=(
-                    f"{self.cache_prefix}_{block_id}_Bn_residual"
-                    if self.context_manager.is_cache_residual()
-                    else f"{self.cache_prefix}_Bn_hidden_states"
-                ),
-                encoder_prefix=(
-                    f"{self.cache_prefix}_{block_id}_Bn_encoder_residual"
-                    if self.context_manager.is_encoder_cache_residual()
-                    else f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states"
-                ),
-            )
-            torch._dynamo.graph_break()
+        Args:
+            pipeline: The Flux pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
         else:
-            # Normal steps: Compute the block and cache the residuals.
-            hidden_states = block(
-                hidden_states,
-                encoder_hidden_states,
-                *args,
-                **kwargs,
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
             )
-            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
-            if not self._skip_prune(block_id):
-                hidden_states = hidden_states.contiguous()
-                hidden_states_residual = hidden_states - original_hidden_states
 
-                if (
-                    encoder_hidden_states is not None
-                    and original_encoder_hidden_states is not None
-                    and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
-                ):
-                    encoder_hidden_states = encoder_hidden_states.contiguous()
-                    encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
-                else:
-                    encoder_hidden_states_residual = None
+    return refresh_cache_context
 
-                self.context_manager.set_Fn_buffer(
-                    original_hidden_states,
-                    prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
-                )
-                if self.context_manager.is_cache_residual():
-                    self.context_manager.set_Bn_buffer(
-                        hidden_states_residual,
-                        prefix=f"{self.cache_prefix}_{block_id}_Bn_residual",
-                    )
-                else:
-                    self.context_manager.set_Bn_buffer(
-                        hidden_states,
-                        prefix=f"{self.cache_prefix}_{block_id}_Bn_hidden_states",
-                    )
-                if encoder_hidden_states_residual is not None:
-                    if self.context_manager.is_encoder_cache_residual():
-                        self.context_manager.set_Bn_encoder_buffer(
-                            encoder_hidden_states_residual,
-                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_residual",
-                        )
-                    else:
-                        self.context_manager.set_Bn_encoder_buffer(
-                            encoder_hidden_states_residual,
-                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states",
-                        )
-            torch._dynamo.graph_break()
 
-        return hidden_states, encoder_hidden_states
+def enable_cache_for_flux2_klein(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for FLUX.2-klein-4B pipeline.
 
+    Args:
+        pipeline: The FLUX.2-klein-4B pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+    Returns:
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
+    """
+    # Build DBCacheConfig for transformer
+    db_cache_config = _build_db_cache_config(cache_config)
+    # The Fn_compute_blocks = 2 override is the most important decision here,
+    # and the rationale (quality degradation at Fn=1) only lives in flux2_klein.
+    db_cache_config.Fn_compute_blocks = 2
 
-class Wan22S2VCachedBlocks(CachedBlocks_Pattern_3_4_5):
-    """CacheDiT blocks wrapper that preserves S2V per-layer audio injection."""
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-    def _run_block(self, block_id: int, block: torch.nn.Module, hidden_states: torch.Tensor, *args, **kwargs):
-        hidden_states = block(hidden_states, *args, **kwargs)
-        hidden_states, new_encoder_hidden_states = self._process_block_outputs(hidden_states)
-        original_after_transformer_block = getattr(
-            self.transformer,
-            "_cache_dit_original_after_transformer_block",
-            getattr(self.transformer, "after_transformer_block", None),
-        )
-        if original_after_transformer_block is not None:
-            hidden_states = original_after_transformer_block(block_id, hidden_states)
-        return hidden_states, new_encoder_hidden_states
+    # Build ParamsModifier for transformer
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
 
-    def call_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        new_encoder_hidden_states = None
-        for block_id, block in enumerate(self.transformer_blocks):
-            hidden_states, new_encoder_hidden_states = self._run_block(
-                block_id,
-                block,
-                hidden_states,
-                *args,
-                **kwargs,
-            )
-        return hidden_states, new_encoder_hidden_states
+    logger.info(
+        f"Enabling cache-dit on Flux2-Klein transformer with BlockAdapter: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
 
-    def call_Fn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        new_encoder_hidden_states = None
-        for block_id, block in enumerate(self._Fn_blocks()):
-            hidden_states, new_encoder_hidden_states = self._run_block(
-                block_id,
-                block,
-                hidden_states,
-                *args,
-                **kwargs,
+    # Enable cache-dit using BlockAdapter for transformer
+    cache_dit.enable_cache(
+        (
+            BlockAdapter(
+                transformer=pipeline.transformer,
+                blocks=[
+                    pipeline.transformer.transformer_blocks,
+                    pipeline.transformer.single_transformer_blocks,
+                ],
+                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2],
+                params_modifiers=[modifier],
             )
-        return hidden_states, new_encoder_hidden_states
+        ),
+        cache_config=db_cache_config,
+    )
 
-    def call_Mn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        original_hidden_states = hidden_states
-        new_encoder_hidden_states = None
-        start_idx = self.context_manager.Fn_compute_blocks()
-        for block_id, block in enumerate(self._Mn_blocks(), start=start_idx):
-            hidden_states, new_encoder_hidden_states = self._run_block(
-                block_id,
-                block,
-                hidden_states,
-                *args,
-                **kwargs,
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
+
+        Args:
+            pipeline: The FLUX.2-klein-4B pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    Fn_compute_blocks=db_cache_config.Fn_compute_blocks,
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
             )
 
-        hidden_states = hidden_states.contiguous()
-        hidden_states_residual = hidden_states - original_hidden_states.to(hidden_states.device)
+    return refresh_cache_context
 
-        return (
-            hidden_states,
-            new_encoder_hidden_states,
-            hidden_states_residual,
-        )
 
-    def call_Bn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
+def enable_cache_for_stable_audio_open(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Stable Audio Open pipeline.
+
+    Args:
+        pipeline: The StableAudioPipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    # StableAudio is officially registered in CacheDiT as Pattern_3:
+    # https://github.com/vipshop/cache-dit/blob/69e82bd1/src/cache_dit/caching/block_adapters/__init__.py#L562
+    #
+    # Pattern_3 is required because StableAudioDiT uses cross-attention
+    # with static encoder_hidden_states that do not change inside the
+    # transformer block loop.
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=pipeline.transformer.transformer_blocks,
+            forward_pattern=ForwardPattern.Pattern_3,
+            params_modifiers=[
+                ParamsModifier(
+                    cache_config=db_cache_config,
+                    calibrator_config=calibrator_config,
+                )
+            ],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
+
+        Args:
+            pipeline: The StableAudioPipeline instance.
+            num_inference_steps: New number of inference steps.
+            verbose: Whether to log refresh operations.
+        """
+        # Bypass SCM for step counts that don't support predefined masks (e.g., vLLM's 1-step dummy run)
+        scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6)
+
+        if cache_config.scm_steps_mask_policy is None or not scm_supported_steps:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            updated_scm_config = DBCacheConfig().reset(
+                num_inference_steps=num_inference_steps,
+                steps_computation_mask=cache_dit.steps_mask(
+                    mask_policy=cache_config.scm_steps_mask_policy,
+                    total_steps=num_inference_steps,
+                ),
+                steps_computation_policy=cache_config.scm_steps_policy,
+            )
+
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=updated_scm_config,
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
+def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for StableDiffusion3Pipeline.
+
+    Args:
+        pipeline: The StableDiffusion3 pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+    """
+    # Build DBCacheConfig for transformer
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    # Build ParamsModifier for transformer
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on StableDiffusion3 transformer with BlockAdapter: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    # Enable cache-dit using BlockAdapter for transformer
+    cache_dit.enable_cache(
+        (
+            BlockAdapter(
+                transformer=pipeline.transformer,
+                blocks=pipeline.transformer.transformer_blocks,
+                forward_pattern=ForwardPattern.Pattern_1,
+                params_modifiers=[modifier],
+            )
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
+
+        Args:
+            pipeline: The LongCatImage pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
+def enable_cache_for_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for LTX2 pipelines (audio-video transformer blocks)."""
+    transformer = get_transformer_from_pipeline(pipeline)
+
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    blocks = transformer.transformer_blocks
+
+    logger.info(
+        f"Enabling cache-dit on LTX2 transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=blocks,
+            # LTX2 blocks return (hidden_states, audio_hidden_states)
+            forward_pattern=ForwardPattern.Pattern_0,
+            # Treat audio_hidden_states as encoder_hidden_states in Pattern_0
+            check_forward_pattern=False,
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = get_transformer_from_pipeline(pipeline)
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
+def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for regular single-transformer DiT models.
+
+    Args:
+        pipeline: The diffusion pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    # Build DBCacheConfig with optional SCM support
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    # Build calibrator config if TaylorSeer is enabled
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    logger.info(
+        f"Enabling cache-dit on transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    # Enable cache-dit on the transformer
+    transformer = get_transformer_from_pipeline(pipeline)
+    cache_dit.enable_cache(
+        transformer,
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
+
+        Args:
+            pipeline: The diffusion pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        transformer = get_transformer_from_pipeline(pipeline)
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
+def enable_cache_for_hunyuan_image3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for HunyuanImage3 pipeline.
+
+    HunyuanImage3 stores its main transformer stack at ``pipeline.model`` with
+    decoder blocks in ``pipeline.model.layers``.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    transformer = getattr(pipeline, "model", None)
+    if transformer is None or not hasattr(transformer, "layers"):
+        raise ValueError(
+            f"HunyuanImage3 cache-dit enabler expects pipeline.model.layers, got pipeline={pipeline.__class__.__name__}"
+        )
+
+    logger.info(
+        f"Enabling cache-dit on HunyuanImage3 model: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+    # HunyuanImage3 decoder layers are single-stream `hidden_states` blocks
+    # (`HunyuanImage3DecoderLayer.forward(hidden_states, ...)`) while caller
+    # still reads outputs as tuple (`layer_outputs[0]`). Pattern_4 matches this:
+    # hidden-only input with hidden-first tuple output.
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_4,
+            params_modifiers=[modifier],
+            check_forward_pattern=False,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = getattr(pipeline, "model", None)
+        if transformer is None:
+            raise ValueError("HunyuanImage3 cache-dit refresh expects pipeline.model to exist.")
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
+class BagelCachedContextManager(CachedContextManager):
+    """
+    Custom CachedContextManager for Bagel that safely handles NaiveCache objects
+    (mapped to encoder_hidden_states) by skipping tensor operations on them.
+    """
+
+    @torch.compiler.disable
+    def apply_cache(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        prefix: str = "Bn",
+        encoder_prefix: str = "Bn_encoder",
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # Allow Bn and Fn prefix to be used for residual cache.
+        if "Bn" in prefix:
+            hidden_states_prev = self.get_Bn_buffer(prefix)
+        else:
+            hidden_states_prev = self.get_Fn_buffer(prefix)
+
+        assert hidden_states_prev is not None, f"{prefix}_buffer must be set before"
+
+        if self.is_cache_residual():
+            hidden_states = hidden_states_prev + hidden_states
+        else:
+            # If cache is not residual, we use the hidden states directly
+            hidden_states = hidden_states_prev
+
+        hidden_states = hidden_states.contiguous()
+
+        if encoder_hidden_states is not None:
+            if "Bn" in encoder_prefix:
+                encoder_hidden_states_prev = self.get_Bn_encoder_buffer(encoder_prefix)
+            else:
+                encoder_hidden_states_prev = self.get_Fn_encoder_buffer(encoder_prefix)
+
+            if encoder_hidden_states_prev is not None:
+                if self.is_encoder_cache_residual():
+                    # FIX: Check if encoder_hidden_states is a tensor before adding
+                    if isinstance(encoder_hidden_states, torch.Tensor) and isinstance(
+                        encoder_hidden_states_prev, torch.Tensor
+                    ):
+                        encoder_hidden_states = encoder_hidden_states_prev + encoder_hidden_states
+                else:
+                    # If encoder cache is not residual, we use the encoder hidden states directly
+                    encoder_hidden_states = encoder_hidden_states_prev
+
+            # FIX: Check if encoder_hidden_states is a tensor before calling contiguous
+            if isinstance(encoder_hidden_states, torch.Tensor):
+                encoder_hidden_states = encoder_hidden_states.contiguous()
+
+        return hidden_states, encoder_hidden_states
+
+
+class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2):
+    """
+    Custom CachedBlocks for Bagel that safely handles NaiveCache objects
+    by adding isinstance checks in call_Mn_blocks and compute_or_prune.
+    """
+
+    def call_Mn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
+        for block in self._Mn_blocks():
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
+
+        # compute hidden_states residual
+        hidden_states = hidden_states.contiguous()
+
+        hidden_states_residual = hidden_states - original_hidden_states
+
+        if (
+            encoder_hidden_states is not None
+            and original_encoder_hidden_states is not None
+            and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
+        ):
+            encoder_hidden_states = encoder_hidden_states.contiguous()
+            encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
+        else:
+            encoder_hidden_states_residual = None
+
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            hidden_states_residual,
+            encoder_hidden_states_residual,
+        )
+
+    def compute_or_prune(
+        self,
+        block_id: int,  # Block index in the transformer blocks
+        # Below are the inputs to the block
+        block,  # The transformer block to be executed
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        # NOTE: Although Bagel likely won't use pruning, implementing safe version just in case.
+        # Copy-pasted from original but adding checks.
+
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
+
+        can_use_prune = self._maybe_prune(
+            block_id,
+            hidden_states,
+            prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
+        )
+
+        torch._dynamo.graph_break()
+        if can_use_prune:
+            self.context_manager.add_pruned_step()
+            hidden_states, encoder_hidden_states = self.context_manager.apply_prune(
+                hidden_states,
+                encoder_hidden_states,
+                prefix=(
+                    f"{self.cache_prefix}_{block_id}_Bn_residual"
+                    if self.context_manager.is_cache_residual()
+                    else f"{self.cache_prefix}_Bn_hidden_states"
+                ),
+                encoder_prefix=(
+                    f"{self.cache_prefix}_{block_id}_Bn_encoder_residual"
+                    if self.context_manager.is_encoder_cache_residual()
+                    else f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states"
+                ),
+            )
+            torch._dynamo.graph_break()
+        else:
+            # Normal steps: Compute the block and cache the residuals.
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
+            if not self._skip_prune(block_id):
+                hidden_states = hidden_states.contiguous()
+                hidden_states_residual = hidden_states - original_hidden_states
+
+                if (
+                    encoder_hidden_states is not None
+                    and original_encoder_hidden_states is not None
+                    and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
+                ):
+                    encoder_hidden_states = encoder_hidden_states.contiguous()
+                    encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
+                else:
+                    encoder_hidden_states_residual = None
+
+                self.context_manager.set_Fn_buffer(
+                    original_hidden_states,
+                    prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
+                )
+                if self.context_manager.is_cache_residual():
+                    self.context_manager.set_Bn_buffer(
+                        hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_{block_id}_Bn_residual",
+                    )
+                else:
+                    self.context_manager.set_Bn_buffer(
+                        hidden_states,
+                        prefix=f"{self.cache_prefix}_{block_id}_Bn_hidden_states",
+                    )
+                if encoder_hidden_states_residual is not None:
+                    if self.context_manager.is_encoder_cache_residual():
+                        self.context_manager.set_Bn_encoder_buffer(
+                            encoder_hidden_states_residual,
+                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_residual",
+                        )
+                    else:
+                        self.context_manager.set_Bn_encoder_buffer(
+                            encoder_hidden_states_residual,
+                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states",
+                        )
+            torch._dynamo.graph_break()
+
+        return hidden_states, encoder_hidden_states
+
+
+class Wan22S2VCachedBlocks(CachedBlocks_Pattern_3_4_5):
+    """CacheDiT blocks wrapper that preserves S2V per-layer audio injection."""
+
+    def _run_block(self, block_id: int, block: torch.nn.Module, hidden_states: torch.Tensor, *args, **kwargs):
+        hidden_states = block(hidden_states, *args, **kwargs)
+        hidden_states, new_encoder_hidden_states = self._process_block_outputs(hidden_states)
+        original_after_transformer_block = getattr(
+            self.transformer,
+            "_cache_dit_original_after_transformer_block",
+            getattr(self.transformer, "after_transformer_block", None),
+        )
+        if original_after_transformer_block is not None:
+            hidden_states = original_after_transformer_block(block_id, hidden_states)
+        return hidden_states, new_encoder_hidden_states
+
+    def call_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        for block_id, block in enumerate(self.transformer_blocks):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+        return hidden_states, new_encoder_hidden_states
+
+    def call_Fn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        for block_id, block in enumerate(self._Fn_blocks()):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+        return hidden_states, new_encoder_hidden_states
+
+    def call_Mn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        new_encoder_hidden_states = None
+        start_idx = self.context_manager.Fn_compute_blocks()
+        for block_id, block in enumerate(self._Mn_blocks(), start=start_idx):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+
+        hidden_states = hidden_states.contiguous()
+        hidden_states_residual = hidden_states - original_hidden_states.to(hidden_states.device)
+
+        return (
+            hidden_states,
+            new_encoder_hidden_states,
+            hidden_states_residual,
+        )
+
+    def call_Bn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
     ):
         new_encoder_hidden_states = None
         if self.context_manager.Bn_compute_blocks() == 0:
             return hidden_states, new_encoder_hidden_states
 
-        start_idx = len(self.transformer_blocks) - self.context_manager.Bn_compute_blocks()
-        for block_id, block in enumerate(self._Bn_blocks(), start=start_idx):
-            hidden_states, new_encoder_hidden_states = self._run_block(
-                block_id,
-                block,
+        start_idx = len(self.transformer_blocks) - self.context_manager.Bn_compute_blocks()
+        for block_id, block in enumerate(self._Bn_blocks(), start=start_idx):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+
+        return hidden_states, new_encoder_hidden_states
+
+
+class Wan22S2VCachedAdapter(CachedAdapter):
+    """CacheDiT adapter that uses Wan22S2VCachedBlocks for S2V audio injection.
+
+    Only overrides collect_unified_blocks to use Wan22S2VCachedBlocks (which
+    calls after_transformer_block per-layer internally). The base class
+    mock_transformer handles the forward wrapping — after_transformer_block is
+    permanently replaced with a no-op in enable_cache_for_wan22_s2v() to prevent
+    double injection.
+    """
+
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        BlockAdapter.assert_normalized(block_adapter)
+
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        Wan22S2VCachedBlocks(
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(unified_blocks_bind_context)
+
+        return total_cached_blocks
+
+
+class BagelCachedAdapter(CachedAdapter):
+    """
+    Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks.
+    """
+
+    @classmethod
+    def create_context(
+        cls,
+        block_adapter: BlockAdapter,
+        **context_kwargs,
+    ) -> tuple[list[str], list[dict[str, Any]]]:
+        # Override to use BagelCachedContextManager
+
+        BlockAdapter.assert_normalized(block_adapter)
+
+        if BlockAdapter.is_cached(block_adapter.pipe):
+            return block_adapter.pipe
+
+        # Check context_kwargs
+        context_kwargs = cls.check_context_kwargs(block_adapter, **context_kwargs)
+
+        # Each Pipeline should have it's own context manager instance.
+        cache_config: BasicCacheConfig = context_kwargs.get("cache_config", None)
+        assert cache_config is not None, "cache_config can not be None."
+
+        # Apply cache on pipeline: wrap cache context
+        pipe_cls_name = block_adapter.pipe.__class__.__name__
+
+        # USE CUSTOM CONTEXT MANAGER
+        context_manager = BagelCachedContextManager(
+            name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
+            persistent_context=isinstance(block_adapter.pipe, FakeDiffusionPipeline),
+        )
+
+        flatten_contexts, contexts_kwargs = cls.modify_context_params(block_adapter, **context_kwargs)
+
+        block_adapter.pipe._context_manager = context_manager  # instance level
+
+        if not context_manager.persistent_context:
+            original_call = block_adapter.pipe.__class__.__call__
+
+            @functools.wraps(original_call)
+            def new_call(self, *args, **kwargs):
+                with ExitStack() as stack:
+                    # cache context will be reset for each pipe inference
+                    for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
+                        stack.enter_context(
+                            context_manager.enter_context(
+                                context_manager.reset_context(
+                                    context_name,
+                                    **context_kwargs,
+                                ),
+                            )
+                        )
+                    outputs = original_call(self, *args, **kwargs)
+                    cls.apply_stats_hooks(block_adapter)
+                    return outputs
+
+            block_adapter.pipe.__class__.__call__ = new_call
+            block_adapter.pipe.__class__._original_call = original_call
+
+        else:
+            # Init persistent cache context for transformer
+            for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
+                context_manager.reset_context(
+                    context_name,
+                    **context_kwargs,
+                )
+
+        block_adapter.pipe.__class__._is_cached = True
+
+        cls.apply_params_hooks(block_adapter, contexts_kwargs)
+
+        return flatten_contexts, contexts_kwargs
+
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        # Override to use BagelCachedBlocks
+
+        BlockAdapter.assert_normalized(block_adapter)
+
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+        # Skipping isinstance check for ContextManager._supported_managers to avoid import issues
+
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+
+                # Directly instantiate BagelCachedBlocks
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        BagelCachedBlocks(
+                            # 0. Transformer blocks configuration
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            # 1. Cache/Prune context configuration
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(unified_blocks_bind_context)
+
+        return total_cached_blocks
+
+
+def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Bagel model (via OmniDiffusion pipeline).
+
+    Args:
+        pipeline: The OmniDiffusion pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    # Build DBCacheConfig
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    # Build calibrator config if TaylorSeer is enabled
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    # Access the transformer: BagelPipeline -> Qwen2MoTForCausalLM -> Qwen2MoTModel
+    # BagelPipeline has self.language_model which is Qwen2MoTForCausalLM
+    # Qwen2MoTForCausalLM has self.model which is Qwen2MoTModel
+    transformer = pipeline.language_model.model
+
+    logger.info(
+        f"Enabling cache-dit on Bagel transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    # Enable cache-dit on the transformer
+    # Pattern_0 corresponds to (hidden_states, encoder_hidden_states) input, output
+    # Custom adapter for Bagel to handle NaiveCache correctly
+    # from vllm_omni.diffusion.cache.bagel_cache_adapter import BagelCachedAdapter # No longer needed
+    BagelCachedAdapter.apply(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_0,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = pipeline.language_model.model
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
+class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5):
+    """
+    Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden
+    states during denoising.
+    """
+
+    @classmethod
+    def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool:
+        if kwargs.get("cache_dit_skip", False):
+            return False
+
+        # Prefix/text forwards either omit image_gen_indicators or update the
+        # DynamicCache. Denoising forwards are gen-only and use update_cache=False.
+        if kwargs.get("update_cache", True):
+            return False
+
+        exist_gen = kwargs.get("exist_gen")
+        exist_und = kwargs.get("exist_und")
+        if exist_gen is None or exist_und is None:
+            image_gen_indicators = kwargs.get("image_gen_indicators")
+            if image_gen_indicators is None:
+                return False
+            exist_gen = image_gen_indicators.any().item()
+            exist_und = (~image_gen_indicators).any().item()
+
+        return exist_gen and not exist_und
+
+    @staticmethod
+    def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+        kwargs = dict(kwargs)
+        kwargs.pop("cache_dit_skip", None)
+        return kwargs
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
+        block_kwargs = self._strip_cache_only_kwargs(kwargs)
+        if not self._is_denoising_call(kwargs):
+            hidden_states, new_encoder_hidden_states = self.call_blocks(
                 hidden_states,
                 *args,
-                **kwargs,
+                **block_kwargs,
+            )
+            return self._process_forward_outputs(hidden_states, new_encoder_hidden_states)
+
+        return super().forward(hidden_states, *args, **block_kwargs)
+
+
+class SensenovaCachedAdapter(CachedAdapter):
+    """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks."""
+
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        BlockAdapter.assert_normalized(block_adapter)
+
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        SensenovaCachedBlocks(
+                            # 0. Transformer blocks configuration
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            # 1. Cache/Prune context configuration
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(unified_blocks_bind_context)
+
+        return total_cached_blocks
+
+
+def enable_cache_for_sensenova_u1(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for SenseNova-U1 model.
+
+    Args:
+        pipeline: The SenseNova-U1 pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    transformer = getattr(getattr(pipeline, "language_model", None), "model", None)
+    if transformer is None or not hasattr(transformer, "layers"):
+        raise ValueError("SenseNovaU1Pipeline cache-dit expects pipeline.language_model.model.layers.")
+
+    # Build DBCacheConfig
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on SenseNova-U1 decoder layers: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    SensenovaCachedAdapter.apply(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_3,
+            params_modifiers=[modifier],
+            check_forward_pattern=True,
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = pipeline.language_model.model
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
             )
 
-        return hidden_states, new_encoder_hidden_states
+    return refresh_cache_context
 
 
-class Wan22S2VCachedAdapter(CachedAdapter):
-    """CacheDiT adapter that uses Wan22S2VCachedBlocks for S2V audio injection.
+def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Flux.2-dev pipeline.
 
-    Only overrides collect_unified_blocks to use Wan22S2VCachedBlocks (which
-    calls after_transformer_block per-layer internally). The base class
-    mock_transformer handles the forward wrapping — after_transformer_block is
-    permanently replaced with a no-op in enable_cache_for_wan22_s2v() to prevent
-    double injection.
+    Args:
+        pipeline: The Flux2 pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+    Returns:
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
     """
+    # Build DBCacheConfig for transformer
+    db_cache_config = _build_db_cache_config(cache_config)
 
-    @classmethod
-    def collect_unified_blocks(
-        cls,
-        block_adapter: BlockAdapter,
-        contexts_kwargs: list[dict],
-    ) -> list[dict[str, torch.nn.ModuleList]]:
-        BlockAdapter.assert_normalized(block_adapter)
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
-        assert hasattr(block_adapter.pipe, "_context_manager")
+    # Build ParamsModifier for transformer
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
 
-        for i in range(len(block_adapter.transformer)):
-            unified_blocks_bind_context = {}
-            for j in range(len(block_adapter.blocks[i])):
-                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
-                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
-                    [
-                        Wan22S2VCachedBlocks(
-                            block_adapter.blocks[i][j],
-                            transformer=block_adapter.transformer[i],
-                            forward_pattern=block_adapter.forward_pattern[i][j],
-                            check_forward_pattern=block_adapter.check_forward_pattern,
-                            check_num_outputs=block_adapter.check_num_outputs,
-                            cache_prefix=block_adapter.blocks_name[i][j],
-                            cache_context=block_adapter.unique_blocks_name[i][j],
-                            context_manager=block_adapter.pipe._context_manager,
-                            cache_type=cache_config.cache_type,
-                        )
-                    ]
-                )
+    logger.info(
+        f"Enabling cache-dit on Flux transformer with BlockAdapter: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
 
-            total_cached_blocks.append(unified_blocks_bind_context)
+    # Enable cache-dit using BlockAdapter for transformer
+    cache_dit.enable_cache(
+        (
+            BlockAdapter(
+                transformer=pipeline.transformer,
+                blocks=[
+                    pipeline.transformer.transformer_blocks,
+                    pipeline.transformer.single_transformer_blocks,
+                ],
+                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2],
+                params_modifiers=[modifier],
+            )
+        ),
+        cache_config=db_cache_config,
+    )
 
-        return total_cached_blocks
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
+
+        Args:
+            pipeline: The Flux2 pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
 
+    return refresh_cache_context
 
-class BagelCachedAdapter(CachedAdapter):
-    """
-    Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks.
+
+def enable_cache_for_glm_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for GlmImage pipeline.
+
+    Args:
+        pipeline: The GlmImage pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+    Returns:
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
     """
+    # Build DBCacheConfig for transformer
+    db_cache_config = _build_db_cache_config(cache_config)
 
-    @classmethod
-    def create_context(
-        cls,
-        block_adapter: BlockAdapter,
-        **context_kwargs,
-    ) -> tuple[list[str], list[dict[str, Any]]]:
-        # Override to use BagelCachedContextManager
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-        BlockAdapter.assert_normalized(block_adapter)
+    # Build ParamsModifier for transformer
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
 
-        if BlockAdapter.is_cached(block_adapter.pipe):
-            return block_adapter.pipe
+    logger.info(
+        f"Enabling cache-dit on GlmImage transformer with BlockAdapter: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
 
-        # Check context_kwargs
-        context_kwargs = cls.check_context_kwargs(block_adapter, **context_kwargs)
+    # Enable cache-dit using BlockAdapter for transformer
+    # Note: We don't use patch_functor here because it's designed for diffusers' GlmImage,
+    # and our vllm-omni implementation has a different forward signature.
+    # We use ForwardPattern.Pattern_0 because our block returns (hidden_states, encoder_hidden_states)
+    cache_dit.enable_cache(
+        (
+            BlockAdapter(
+                transformer=pipeline.transformer,
+                blocks=pipeline.transformer.transformer_blocks,
+                forward_pattern=ForwardPattern.Pattern_0,
+                params_modifiers=[modifier],
+                patch_functor=None,
+                has_separate_cfg=True,
+            )
+        ),
+        cache_config=db_cache_config,
+    )
 
-        # Each Pipeline should have it's own context manager instance.
-        cache_config: BasicCacheConfig = context_kwargs.get("cache_config", None)
-        assert cache_config is not None, "cache_config can not be None."
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
 
-        # Apply cache on pipeline: wrap cache context
-        pipe_cls_name = block_adapter.pipe.__class__.__name__
+        Args:
+            pipeline: The GlmImage pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
 
-        # USE CUSTOM CONTEXT MANAGER
-        context_manager = BagelCachedContextManager(
-            name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
-            persistent_context=isinstance(block_adapter.pipe, FakeDiffusionPipeline),
-        )
+    return refresh_cache_context
 
-        flatten_contexts, contexts_kwargs = cls.modify_context_params(block_adapter, **context_kwargs)
 
-        block_adapter.pipe._context_manager = context_manager  # instance level
+def enable_cache_for_dreamid_omni(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for DreamID-Omni fused pipeline.
 
-        if not context_manager.persistent_context:
-            original_call = block_adapter.pipe.__class__.__call__
+    Args:
+        pipeline: The DreamIDOmni pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
 
-            @functools.wraps(original_call)
-            def new_call(self, *args, **kwargs):
-                with ExitStack() as stack:
-                    # cache context will be reset for each pipe inference
-                    for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
-                        stack.enter_context(
-                            context_manager.enter_context(
-                                context_manager.reset_context(
-                                    context_name,
-                                    **context_kwargs,
-                                ),
-                            )
-                        )
-                    outputs = original_call(self, *args, **kwargs)
-                    cls.apply_stats_hooks(block_adapter)
-                    return outputs
+    Returns:
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
+    """
+    transformer = getattr(pipeline, "transformer", None)
+    fused_blocks = getattr(transformer, "fused_blocks", None)
+    if transformer is None or fused_blocks is None:
+        raise ValueError("DreamIDOmniPipeline cache-dit expects pipeline.transformer.fused_blocks.")
 
-            block_adapter.pipe.__class__.__call__ = new_call
-            block_adapter.pipe.__class__._original_call = original_call
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on DreamID-Omni fused blocks: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=fused_blocks,
+            forward_pattern=ForwardPattern.Pattern_0,
+            params_modifiers=[modifier],
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+    )
 
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = pipeline.transformer
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
         else:
-            # Init persistent cache context for transformer
-            for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
-                context_manager.reset_context(
-                    context_name,
-                    **context_kwargs,
-                )
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
 
-        block_adapter.pipe.__class__._is_cached = True
+    return refresh_cache_context
 
-        cls.apply_params_hooks(block_adapter, contexts_kwargs)
 
-        return flatten_contexts, contexts_kwargs
+def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for ERNIE-Image pipeline.
 
-    @classmethod
-    def collect_unified_blocks(
-        cls,
-        block_adapter: BlockAdapter,
-        contexts_kwargs: list[dict],
-    ) -> list[dict[str, torch.nn.ModuleList]]:
-        # Override to use BagelCachedBlocks
+    ERNIE-Image blocks have signature:
+        forward(x, rotary_pos_emb, temb, attention_mask) -> x
 
-        BlockAdapter.assert_normalized(block_adapter)
+    Where x is hidden_states (concatenated image + text tokens).
+    This matches Pattern_3 which expects:
+        - Input: hidden_states only
+        - Output: hidden_states only
+
+    Args:
+        pipeline: The ERNIE-Image pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
-        assert hasattr(block_adapter.pipe, "_context_manager")
-        # Skipping isinstance check for ContextManager._supported_managers to avoid import issues
+    logger.info(
+        f"Enabling cache-dit on ERNIE-Image transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
 
-        for i in range(len(block_adapter.transformer)):
-            unified_blocks_bind_context = {}
-            for j in range(len(block_adapter.blocks[i])):
-                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=pipeline.transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_3,
+            check_forward_pattern=False,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
 
-                # Directly instantiate BagelCachedBlocks
-                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
-                    [
-                        BagelCachedBlocks(
-                            # 0. Transformer blocks configuration
-                            block_adapter.blocks[i][j],
-                            transformer=block_adapter.transformer[i],
-                            forward_pattern=block_adapter.forward_pattern[i][j],
-                            check_forward_pattern=block_adapter.check_forward_pattern,
-                            check_num_outputs=block_adapter.check_num_outputs,
-                            # 1. Cache/Prune context configuration
-                            cache_prefix=block_adapter.blocks_name[i][j],
-                            cache_context=block_adapter.unique_blocks_name[i][j],
-                            context_manager=block_adapter.pipe._context_manager,
-                            cache_type=cache_config.cache_type,
-                        )
-                    ]
-                )
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
 
-            total_cached_blocks.append(unified_blocks_bind_context)
+    return refresh_cache_context
 
-        return total_cached_blocks
 
+def enable_cache_for_helios(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Helios pipeline.
 
-class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5):
-    """
-    Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden
-    states during denoising.
-    """
+    Helios extends Wan2.2 with multi-term memory patches and guidance cross-attention.
+    Its transformer blocks have the same single-output signature as Wan2.2 blocks
+    (returns hidden_states only), so we use ForwardPattern.Pattern_2.
 
-    @classmethod
-    def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool:
-        if kwargs.get("cache_dit_skip", False):
-            return False
+    Args:
+        pipeline: The HeliosPipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
 
-        # Prefix/text forwards either omit image_gen_indicators or update the
-        # DynamicCache. Denoising forwards are gen-only and use update_cache=False.
-        if kwargs.get("update_cache", True):
-            return False
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
 
-        exist_gen = kwargs.get("exist_gen")
-        exist_und = kwargs.get("exist_und")
-        if exist_gen is None or exist_und is None:
-            image_gen_indicators = kwargs.get("image_gen_indicators")
-            if image_gen_indicators is None:
-                return False
-            exist_gen = image_gen_indicators.any().item()
-            exist_und = (~image_gen_indicators).any().item()
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-        return exist_gen and not exist_und
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[pipeline.transformer.blocks],
+            forward_pattern=[ForwardPattern.Pattern_2],
+            params_modifiers=[
+                ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
+            ],
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
 
-    @staticmethod
-    def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
-        kwargs = dict(kwargs)
-        kwargs.pop("cache_dit_skip", None)
-        return kwargs
+    logger.info(
+        f"Enabling cache-dit on Helios transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
 
-    def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
-        block_kwargs = self._strip_cache_only_kwargs(kwargs)
-        if not self._is_denoising_call(kwargs):
-            hidden_states, new_encoder_hidden_states = self.call_blocks(
-                hidden_states,
-                *args,
-                **block_kwargs,
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
             )
-            return self._process_forward_outputs(hidden_states, new_encoder_hidden_states)
 
-        return super().forward(hidden_states, *args, **block_kwargs)
+    return refresh_cache_context
 
 
-class SensenovaCachedAdapter(CachedAdapter):
-    """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks."""
+def enable_cache_for_hunyuan_video_15(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for HunyuanVideo 1.5 pipeline.
 
-    @classmethod
-    def collect_unified_blocks(
-        cls,
-        block_adapter: BlockAdapter,
-        contexts_kwargs: list[dict],
-    ) -> list[dict[str, torch.nn.ModuleList]]:
-        BlockAdapter.assert_normalized(block_adapter)
+    HunyuanVideo 1.5 uses a single transformer with has_separate_cfg=True
+    (separate conditional/unconditional forward passes). The _sp_plan scatter
+    is applied at the transformer input boundary (empty-string key), so CacheDiT
+    sees already-sharded hidden_states throughout transformer_blocks.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
 
-        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
-        assert hasattr(block_adapter.pipe, "_context_manager")
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-        for i in range(len(block_adapter.transformer)):
-            unified_blocks_bind_context = {}
-            for j in range(len(block_adapter.blocks[i])):
-                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
-                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
-                    [
-                        SensenovaCachedBlocks(
-                            # 0. Transformer blocks configuration
-                            block_adapter.blocks[i][j],
-                            transformer=block_adapter.transformer[i],
-                            forward_pattern=block_adapter.forward_pattern[i][j],
-                            check_forward_pattern=block_adapter.check_forward_pattern,
-                            check_num_outputs=block_adapter.check_num_outputs,
-                            # 1. Cache/Prune context configuration
-                            cache_prefix=block_adapter.blocks_name[i][j],
-                            cache_context=block_adapter.unique_blocks_name[i][j],
-                            context_manager=block_adapter.pipe._context_manager,
-                            cache_type=cache_config.cache_type,
-                        )
-                    ]
-                )
+    logger.info(
+        f"Enabling cache-dit on HunyuanVideo15 transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
 
-            total_cached_blocks.append(unified_blocks_bind_context)
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=pipeline.transformer.transformer_blocks,
+            forward_pattern=ForwardPattern.Pattern_0,
+            check_forward_pattern=True,
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
 
-        return total_cached_blocks
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
 
 
-def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc:
+def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for Cosmos3.
 
     Cosmos3 has a dual-pathway architecture (UND + GEN) but only the GEN
@@ -948,6 +1959,41 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheCo
     Returns:
         A refresh function that can be called to update cache context with new num_inference_steps.
     """
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    logger.info(
+        f"Enabling cache-dit on Cosmos3 gen_layers: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[pipeline.transformer.gen_layers],
+            # Cosmos3 GEN blocks return only hidden_states.  Per-layer UND K/V
+            # conditioning uses the transformer's cache-dit fallback path.
+            forward_pattern=[ForwardPattern.Pattern_3],
+            params_modifiers=[
+                ParamsModifier(
+                    cache_config=db_cache_config,
+                    calibrator_config=calibrator_config,
+                ),
+            ],
+            check_forward_pattern=False,
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
     # The T2I denoising loop skips the unconditional forward outside the
     # guidance interval as a speed optimization. cache-dit distinguishes the
     # conditional vs unconditional passes purely by transformer-forward parity
@@ -955,8 +2001,25 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheCo
     # step accounting. Still do both cond/uncond CFG steps when cache-dit is active.
     # CFG is instead neutralized via scale=1.0 outside the interval.
     pipeline._cache_dit_requires_paired_cfg = True
-    block_adapter = CacheDiTBackend.maybe_build_block_adapter(pipeline)
-    return enable_cache_for_dit(pipeline, cache_config, block_adapter)
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
 
 
 # Register custom cache-dit enablers after function definitions
@@ -964,8 +2027,29 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheCo
     {
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
-        "Wan22TI2VPipeline": enable_cache_for_wan22,
         "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
+        "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3,
+        "FluxPipeline": enable_cache_for_flux,
+        "Flux2KleinPipeline": enable_cache_for_flux2_klein,
+        "LongCatImagePipeline": enable_cache_for_longcat_image,
+        "LongCatImageEditPipeline": enable_cache_for_longcat_image,
+        "StableAudioPipeline": enable_cache_for_stable_audio_open,
+        "StableDiffusion3Pipeline": enable_cache_for_sd3,
+        "LTX2Pipeline": enable_cache_for_ltx2,
+        "LTX2ImageToVideoPipeline": enable_cache_for_ltx2,
+        "LTX2TwoStagesPipeline": enable_cache_for_ltx2,
+        "LTX2ImageToVideoTwoStagesPipeline": enable_cache_for_ltx2,
+        "LTX23Pipeline": enable_cache_for_ltx2,
+        "LTX23ImageToVideoPipeline": enable_cache_for_ltx2,
+        "BagelPipeline": enable_cache_for_bagel,
+        "SenseNovaU1Pipeline": enable_cache_for_sensenova_u1,
+        "GlmImagePipeline": enable_cache_for_glm_image,
+        "Flux2Pipeline": enable_cache_for_flux2,
+        "DreamIDOmniPipeline": enable_cache_for_dreamid_omni,
+        "ErnieImagePipeline": enable_cache_for_ernie_image,
+        "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15,
+        "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15,
+        "HeliosPipeline": enable_cache_for_helios,
         "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
     }
 )
@@ -1008,43 +2092,6 @@ def __init__(self, cache_config: Any = None):
         self._refresh_func: Callable[[Any, int, bool], None] | None = None
         self._last_num_inference_steps: int | None = None
 
-    @staticmethod
-    def maybe_build_block_adapter(pipeline) -> BlockAdapter | None:
-        """If a module defines `_cache_dit_adapter_config`, build the corresponding
-        block adapter.
-        """
-        transformer = default_get_pipeline_transformer(pipeline)
-
-        adapter_cfg: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None)
-        if adapter_cfg is None:
-            return None
-
-        block_attrs, forward_pattern = zip(*(adapter_cfg.block_forward_patterns).items())
-        missing_attrs = [block_attr for block_attr in block_attrs if not hasattr(transformer, block_attr)]
-
-        if missing_attrs:
-            logger.warning("Missing Cache DiT block attributes: %s", missing_attrs)
-
-        block_adapter = BlockAdapter(
-            transformer=transformer,
-            blocks=[getattr(transformer, block_attr) for block_attr in block_attrs],
-            forward_pattern=list(forward_pattern),
-            has_separate_cfg=adapter_cfg.has_separate_cfg,
-            check_forward_pattern=adapter_cfg.check_forward_pattern,
-        )
-        return block_adapter
-
-    @staticmethod
-    def maybe_get_cached_adapter_cls(pipeline) -> type[CachedAdapter] | None:
-        """If a module has a custom cached adapter type registered, e.g., SenseNova, retrieve it
-        from the transformer's CacheDiTAdapterConfig."""
-        transformer = default_get_pipeline_transformer(pipeline)
-
-        adapter_cfg: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None)
-        if adapter_cfg is None:
-            return None
-        return adapter_cfg.cached_adapter_cls
-
     def enable(self, pipeline: Any) -> None:
         """Enable cache-dit on the pipeline if configured.
 
@@ -1063,12 +2110,8 @@ def enable(self, pipeline: Any) -> None:
             logger.info(f"Using custom cache-dit enabler for model: {pipeline_name}")
             self._refresh_func = CUSTOM_DIT_ENABLERS[pipeline_name](pipeline, self.config)
         else:
-            # Common case; either the model doesn't explicitly support dit cache yet,
-            # Or it defines its _cache_dit_adapter_config, which describes how we should
-            # create its block adapter.
-            block_adapter = self.maybe_build_block_adapter(pipeline)
-            adapter_cls = self.maybe_get_cached_adapter_cls(pipeline)
-            self._refresh_func = enable_cache_for_dit(pipeline, self.config, block_adapter, adapter_cls)
+            # For regular single-transformer models
+            self._refresh_func = enable_cache_for_dit(pipeline, self.config)
 
         self.enabled = True
         logger.info(f"Cache-dit enabled successfully on {pipeline_name}")
diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py
index 3b002abe0bf..e6c2e70e318 100644
--- a/vllm_omni/diffusion/cache/selector.py
+++ b/vllm_omni/diffusion/cache/selector.py
@@ -3,7 +3,6 @@
 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
 from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
-from vllm_omni.diffusion.cache.stepcache import StepCacheBackend
 from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
@@ -15,15 +14,14 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
     - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface
     - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface
     - mag_cache: Uses MagCacheBackend with enable()/refresh() interface
-    - step_cache: Uses StepCacheBackend for DreamZero velocity step skip
 
     Args:
-        cache_backend: Cache backend name ("cache_dit", "tea_cache",
-            "mag_cache", "step_cache", or None).
+        cache_backend: Cache backend name ("cache_dit", "tea_cache", "mag_cache", or None).
         cache_config: Cache configuration (dict or DiffusionCacheConfig instance).
 
     Returns:
-        Cache backend instance if cache_backend is set, None otherwise.
+        Cache backend instance (CacheDiTBackend, TeaCacheBackend, or MagCacheBackend)
+        if cache_backend is set, None otherwise.
 
     Raises:
         ValueError: If cache_backend is unsupported.
@@ -40,10 +38,7 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
         return TeaCacheBackend(cache_config)
     elif cache_backend == "mag_cache":
         return MagCacheBackend(cache_config)
-    elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"):
-        return StepCacheBackend(cache_config)
     else:
         raise ValueError(
-            f"Unsupported cache backend: {cache_backend}. "
-            "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'"
+            f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache', 'mag_cache'"
         )
diff --git a/vllm_omni/diffusion/cache/stepcache/__init__.py b/vllm_omni/diffusion/cache/stepcache/__init__.py
deleted file mode 100644
index a7e1852b0d4..00000000000
--- a/vllm_omni/diffusion/cache/stepcache/__init__.py
+++ /dev/null
@@ -1,66 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-StepCache: velocity-based step skipping for DreamZero-style DiT denoising.
-
-Unlike block-level Cache-DiT, this backend skips entire DiT forwards during the
-pipeline denoise loop when successive velocity predictions are highly aligned.
-
-Usage:
-    from vllm_omni import Omni
-
-    omni = Omni(
-        model="...",
-        cache_backend="step_cache",
-    )
-
-    # Alternative: environment variable
-    # export DIFFUSION_CACHE_BACKEND=step_cache
-"""
-
-from vllm_omni.diffusion.cache.stepcache.backend import (
-    CUSTOM_STEP_CACHE_DIT_ENABLERS,
-    CUSTOM_STEPCACHE_ENABLERS,
-    STEP_CACHE_CONFIG_ATTR,
-    STEP_CACHE_DIT_CONFIG_ATTR,
-    STEP_CACHE_DIT_STATE_ATTR,
-    STEP_CACHE_STATE_ATTR,
-    StepCacheBackend,
-    StepCacheDiTBackend,
-    StepCacheDiTConfig,
-    StepCacheDiTState,
-    enable_dreamzero_step_cache_dit,
-    enable_dreamzero_stepcache,
-    get_step_cache_dit_config,
-    get_step_cache_dit_state,
-    get_stepcache_config,
-    get_stepcache_state,
-    is_step_cache_dit_active,
-    is_stepcache_active,
-)
-from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
-from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
-
-__all__ = [
-    "CUSTOM_STEP_CACHE_DIT_ENABLERS",
-    "CUSTOM_STEPCACHE_ENABLERS",
-    "STEP_CACHE_CONFIG_ATTR",
-    "STEP_CACHE_DIT_CONFIG_ATTR",
-    "STEP_CACHE_DIT_STATE_ATTR",
-    "STEP_CACHE_STATE_ATTR",
-    "StepCacheBackend",
-    "StepCacheConfig",
-    "StepCacheDiTBackend",
-    "StepCacheDiTConfig",
-    "StepCacheDiTState",
-    "StepCacheState",
-    "enable_dreamzero_step_cache_dit",
-    "enable_dreamzero_stepcache",
-    "get_step_cache_dit_config",
-    "get_step_cache_dit_state",
-    "get_stepcache_config",
-    "get_stepcache_state",
-    "is_step_cache_dit_active",
-    "is_stepcache_active",
-]
diff --git a/vllm_omni/diffusion/cache/stepcache/backend.py b/vllm_omni/diffusion/cache/stepcache/backend.py
deleted file mode 100644
index 517d2369a3c..00000000000
--- a/vllm_omni/diffusion/cache/stepcache/backend.py
+++ /dev/null
@@ -1,129 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-StepCache backend implementation.
-
-This module provides the stepcache backend that implements the CacheBackend
-interface for DreamZero-style velocity-based step skipping in the denoise loop.
-"""
-
-from __future__ import annotations
-
-from typing import Any
-
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.cache.base import CacheBackend
-from vllm_omni.diffusion.cache.stepcache.config import (
-    STEP_CACHE_CONFIG_ATTR,
-    STEP_CACHE_STATE_ATTR,
-    StepCacheConfig,
-)
-from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
-from vllm_omni.diffusion.data import DiffusionCacheConfig
-
-logger = init_logger(__name__)
-
-# Backward-compatible aliases for older imports and pipeline attribute names.
-STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR
-STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR
-StepCacheDiTConfig = StepCacheConfig
-StepCacheDiTState = StepCacheState
-
-
-def is_stepcache_active(pipeline: Any) -> bool:
-    """Return True when stepcache is enabled on *pipeline*."""
-    cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None)
-    return isinstance(cfg, StepCacheConfig) and cfg.enabled
-
-
-def get_stepcache_config(pipeline: Any) -> StepCacheConfig | None:
-    cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None)
-    return cfg if isinstance(cfg, StepCacheConfig) else None
-
-
-def get_stepcache_state(pipeline: Any) -> StepCacheState | None:
-    state = getattr(pipeline, STEP_CACHE_STATE_ATTR, None)
-    return state if isinstance(state, StepCacheState) else None
-
-
-def is_step_cache_dit_active(pipeline: Any) -> bool:
-    return is_stepcache_active(pipeline)
-
-
-def get_step_cache_dit_config(pipeline: Any) -> StepCacheConfig | None:
-    return get_stepcache_config(pipeline)
-
-
-def get_step_cache_dit_state(pipeline: Any) -> StepCacheState | None:
-    return get_stepcache_state(pipeline)
-
-
-def enable_dreamzero_stepcache(pipeline: Any, config: DiffusionCacheConfig) -> None:
-    """Enable stepcache for DreamZeroPipeline."""
-    cache_config = StepCacheConfig.from_diffusion_cache_config(config)
-    setattr(pipeline, STEP_CACHE_CONFIG_ATTR, cache_config)
-    setattr(pipeline, STEP_CACHE_STATE_ATTR, StepCacheState(cache_config))
-
-    logger.info(
-        "stepcache enabled on DreamZeroPipeline (enabled=%s, thresholds=%s, countdowns=%s)",
-        cache_config.enabled,
-        cache_config.sim_thresholds,
-        cache_config.skip_countdowns,
-    )
-
-
-enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache
-
-CUSTOM_STEPCACHE_ENABLERS = {
-    "DreamZeroPipeline": enable_dreamzero_stepcache,
-}
-
-CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS
-
-
-class StepCacheBackend(CacheBackend):
-    """
-    Velocity cosine step-skipping cache backend for DreamZero.
-
-    Attaches :class:`StepCacheConfig` and :class:`StepCacheState` to supported
-    pipelines. The denoise loop calls :meth:`StepCacheState.should_run_step` to
-    decide whether to run ``predict_noise``.
-
-    Example:
-        >>> from vllm_omni.diffusion.data import DiffusionCacheConfig
-        >>> backend = StepCacheBackend(DiffusionCacheConfig())
-        >>> backend.enable(pipeline)
-        >>> backend.refresh(pipeline, num_inference_steps=16)
-    """
-
-    def enable(self, pipeline: Any) -> None:
-        """Enable stepcache on a supported pipeline."""
-        pipeline_type = pipeline.__class__.__name__
-
-        if pipeline_type in CUSTOM_STEPCACHE_ENABLERS:
-            logger.info("Using custom stepcache enabler for model: %s", pipeline_type)
-            CUSTOM_STEPCACHE_ENABLERS[pipeline_type](pipeline, self.config)
-        else:
-            raise ValueError(
-                f"step_cache backend does not support {pipeline_type}. Supported: {sorted(CUSTOM_STEPCACHE_ENABLERS)}"
-            )
-
-        self.enabled = True
-
-    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh stepcache state for a new generation."""
-        state = get_stepcache_state(pipeline)
-        if state is not None:
-            state.reset()
-            if verbose:
-                logger.debug(
-                    "stepcache state refreshed (num_inference_steps=%d)",
-                    num_inference_steps,
-                )
-        elif verbose and is_stepcache_active(pipeline):
-            logger.warning("stepcache config active but state not found on pipeline")
-
-
-StepCacheDiTBackend = StepCacheBackend
diff --git a/vllm_omni/diffusion/cache/stepcache/config.py b/vllm_omni/diffusion/cache/stepcache/config.py
deleted file mode 100644
index 00a7984ae7d..00000000000
--- a/vllm_omni/diffusion/cache/stepcache/config.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Configuration for step-level DiT velocity caching."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-from vllm_omni.diffusion.data import DiffusionCacheConfig
-
-# Pipeline attributes set by StepCacheBackend.enable().
-STEP_CACHE_CONFIG_ATTR = "_stepcache_config"
-STEP_CACHE_STATE_ATTR = "_stepcache_state"
-
-
-@dataclass(frozen=True)
-class StepCacheConfig:
-    """Runtime config for velocity-based step skipping.
-
-    Skips entire DiT forwards when successive video velocity predictions are
-    highly aligned (cosine similarity above configured thresholds).
-
-    Reference: DreamZero paper DiT Caching / dreamzero.git ``should_run_model``.
-    """
-
-    enabled: bool = True
-    min_history_steps: int = 2
-    max_history: int = 2
-    sim_thresholds: tuple[float, ...] = (0.95, 0.93)
-    skip_countdowns: tuple[int, ...] = (4, 2)
-
-    @classmethod
-    def from_diffusion_cache_config(cls, config: DiffusionCacheConfig) -> StepCacheConfig:
-        enabled = config.step_cache_dit_enabled
-
-        thresholds = tuple(config.velocity_sim_thresholds)
-        countdowns = tuple(config.velocity_skip_countdowns)
-        if len(thresholds) != len(countdowns):
-            raise ValueError(
-                "velocity_sim_thresholds and velocity_skip_countdowns must have the same length; "
-                f"got {len(thresholds)} and {len(countdowns)}."
-            )
-
-        return cls(
-            enabled=enabled,
-            min_history_steps=config.step_cache_dit_min_history,
-            max_history=config.step_cache_dit_max_history,
-            sim_thresholds=thresholds,
-            skip_countdowns=countdowns,
-        )
diff --git a/vllm_omni/diffusion/cache/stepcache/state.py b/vllm_omni/diffusion/cache/stepcache/state.py
deleted file mode 100644
index db786d4f210..00000000000
--- a/vllm_omni/diffusion/cache/stepcache/state.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Mutable state for step-level DiT velocity caching."""
-
-from __future__ import annotations
-
-import torch
-
-from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
-
-
-class StepCacheState:
-    """Per-generation mutable state (skip countdown)."""
-
-    def __init__(self, config: StepCacheConfig) -> None:
-        self.config = config
-        self.skip_countdown = 0
-
-    def reset(self) -> None:
-        self.skip_countdown = 0
-
-    def should_run_step(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> bool:
-        """Return True when the DiT forward should execute at this scheduler step."""
-        if not self.config.enabled:
-            return True
-
-        if len(prev_predictions) < self.config.min_history_steps:
-            return True
-
-        if self.skip_countdown > 1:
-            self.skip_countdown -= 1
-            return False
-        if self.skip_countdown == 1:
-            self.skip_countdown = 0
-            return True
-
-        v_last = prev_predictions[-1][0].flatten(1).float()
-        v_prev = prev_predictions[-2][0].flatten(1).float()
-        sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean()
-
-        for threshold, countdown in zip(self.config.sim_thresholds, self.config.skip_countdowns):
-            if sim > threshold:
-                self.skip_countdown = countdown
-                return False
-
-        return True
-
-    def trim_history(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> None:
-        while len(prev_predictions) > self.config.max_history:
-            prev_predictions.pop(0)
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 6a25c6dab02..3c1323bbc42 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -302,8 +302,6 @@ class DiffusionCacheConfig:
                     scm_steps_mask_policy, scm_steps_policy
         - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio,
                     mag_ratios, mag_calibrate
-        - step_cache: step_cache_dit_enabled, velocity_sim_thresholds,
-                          velocity_skip_countdowns, step_cache_dit_min_history
 
     Example:
         >>> # From dict (user-facing API) - partial config uses defaults for missing keys
@@ -377,13 +375,6 @@ class DiffusionCacheConfig:
     # "repeat" refreshes every force_refresh_step_hint steps.
     force_refresh_step_policy: str = "once"
 
-    # step_cache parameters [step_cache only] — DreamZero velocity schedule
-    step_cache_dit_enabled: bool = True
-    velocity_sim_thresholds: list[float] = field(default_factory=lambda: [0.95, 0.93])
-    velocity_skip_countdowns: list[int] = field(default_factory=lambda: [4, 2])
-    step_cache_dit_min_history: int = 2
-    step_cache_dit_max_history: int = 2
-
     # Additional parameters that may be passed but not explicitly defined
     _extra_params: dict[str, Any] = field(default_factory=dict, repr=False)
 
@@ -493,6 +484,23 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
     return None
 
 
+# Model classes that default to the BDE (AR-diffusion) engine.
+_BDE_ENGINE_MODEL_CLASSES = frozenset({"DreamZeroPipeline"})
+
+
+def default_engine_backend_for_model(model_class_name: str | None) -> str | None:
+    """Engine backend a model opts into by default.
+
+    Returns ``"bde"`` for AR-diffusion world models (DreamZero today) so they run
+    on the Block Diffusion Engine with engine-level KV cache management; returns
+    ``None`` to keep the base ``DiffusionEngine``. Extend
+    ``_BDE_ENGINE_MODEL_CLASSES`` as more AR-DiT models migrate onto BDE.
+    """
+    if model_class_name in _BDE_ENGINE_MODEL_CLASSES:
+        return "bde"
+    return None
+
+
 @dataclass
 class OmniDiffusionConfig:
     # Model and path configuration (for convenience)
@@ -538,6 +546,14 @@ class OmniDiffusionConfig:
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None
 
+    # Engine backend selection. Resolved by ``DiffusionEngine.make_engine``
+    # (mirrors ``distributed_executor_backend`` / ``DiffusionExecutor.get_class``):
+    #   "default" -> DiffusionEngine
+    #   "bde"     -> BDEEngine (AR-diffusion engine with engine-level KV cache
+    #                management); also accepts a DiffusionEngine subclass or an
+    #                import path string.
+    engine_backend: str = "default"
+
     # HuggingFace specific parameters
     trust_remote_code: bool = False
     revision: str | None = None
@@ -1053,10 +1069,6 @@ def enrich_config(self) -> None:
                         self.model_class_name = "WanS2VPipeline"
                     self.tf_model_config = TransformerConfig()
                     self.update_multimodal_support()
-                elif model_type == "Gr00tN1d7" or "Gr00tN1d7" in architectures:
-                    self.model_class_name = "Gr00tN1d7Pipeline"
-                    self.set_tf_model_config(TransformerConfig())
-                    self.update_multimodal_support()
                 elif model_type == "vla":
                     from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 
@@ -1064,6 +1076,12 @@ def enrich_config(self) -> None:
                         self.model_class_name = "DreamZeroPipeline"
                         self.set_tf_model_config(TransformerConfig())
                         self.update_multimodal_support()
+                        # DreamZero is an AR-diffusion world model: route it to the
+                        # BDE engine (engine-level KV cache management) unless the
+                        # caller explicitly selected an engine backend.
+                        bde_default = default_engine_backend_for_model(self.model_class_name)
+                        if bde_default is not None and self.engine_backend == "default":
+                            self.engine_backend = bde_default
                     else:
                         raise
                 elif architectures and len(architectures) == 1:
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index a7869b14022..97a5b21069a 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -17,6 +17,7 @@
 import PIL.Image
 import torch
 from vllm.logger import init_logger
+from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import (
@@ -411,20 +412,63 @@ def _handle_finished_requests(
             out = self._finalize_finished_request(rid, _output, missing_result_error)
             self._complete_future(fut, out)
 
+    @staticmethod
+    def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
+        """Resolve the engine class selected by ``config.engine_backend``.
+
+        Mirrors ``DiffusionExecutor.get_class``: accepts the string keys
+        ``"default"`` / ``"bde"``, a ``DiffusionEngine`` subclass, or an import
+        path string. Kept separate from :meth:`make_engine` so the selection is
+        testable without constructing an engine (which runs a dummy forward).
+
+        Args:
+            config: The configuration for the diffusion engine.
+
+        Returns:
+            The ``DiffusionEngine`` (sub)class to instantiate.
+        """
+        backend = getattr(config, "engine_backend", "default") or "default"
+
+        if isinstance(backend, type):
+            if not issubclass(backend, DiffusionEngine):
+                raise TypeError(f"engine_backend must be a DiffusionEngine subclass. Got {backend}.")
+            return backend
+        if backend == "default":
+            return DiffusionEngine
+        if backend == "bde":
+            # Lazy import to avoid a circular dependency (bde.engine imports this
+            # module), mirroring how DiffusionExecutor.get_class imports backends.
+            from vllm_omni.bde.engine import BDEEngine
+
+            return BDEEngine
+        if isinstance(backend, str):
+            try:
+                engine_class = resolve_obj_by_qualname(backend)
+            except (ImportError, ValueError) as e:
+                raise ValueError(
+                    f"Failed to load engine_backend '{backend}'. Ensure it is a valid "
+                    f"python path. Error: {e}"
+                ) from e
+            if not issubclass(engine_class, DiffusionEngine):
+                raise TypeError(f"engine_backend must resolve to a DiffusionEngine subclass. Got {engine_class}.")
+            return engine_class
+        raise ValueError(f"Unknown engine_backend: {backend!r}")
+
     @staticmethod
     def make_engine(
         config: OmniDiffusionConfig,
         scheduler: SchedulerInterface | None = None,
     ) -> DiffusionEngine:
-        """Factory method to create a DiffusionEngine instance.
+        """Factory method to create the engine selected by ``config.engine_backend``.
 
         Args:
             config: The configuration for the diffusion engine.
 
         Returns:
-            An instance of DiffusionEngine.
+            An instance of the resolved ``DiffusionEngine`` (sub)class.
         """
-        return DiffusionEngine(config, scheduler=scheduler)
+        engine_class = DiffusionEngine.resolve_engine_class(config)
+        return engine_class(config, scheduler=scheduler)
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         with self._cv:
diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py
index 3a056d6d45b..42675c7c095 100644
--- a/vllm_omni/diffusion/distributed/hsdp.py
+++ b/vllm_omni/diffusion/distributed/hsdp.py
@@ -140,13 +140,8 @@ def apply_hsdp_to_model(
         fs_rank,
     )
 
-    # When the model contains FP8 parameters (online quantization), let FSDP
-    # keep the original storage dtype on all-gather instead of casting to
-    # hsdp_config.param_dtype (typically bfloat16). FP8 GEMM kernels expect
-    # FP8 inputs; an implicit FP8 -> bf16 cast would silently break them.
-    has_fp8_params = any(p.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for p in model.parameters())
     mp_policy = MixedPrecisionPolicy(
-        param_dtype=None if has_fp8_params else hsdp_config.param_dtype,
+        param_dtype=hsdp_config.param_dtype,
         reduce_dtype=hsdp_config.reduce_dtype,
         output_dtype=hsdp_config.output_dtype,
         cast_forward_inputs=False,
diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py
index 1aa40aae5cf..b2339843ed8 100644
--- a/vllm_omni/diffusion/model_loader/diffusers_loader.py
+++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py
@@ -602,17 +602,10 @@ def _init_from_load_format(
             model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
             with set_current_diffusion_config(self.od_config):
                 model = model_cls(od_config=self.od_config)
-            # HSDP normally defers GPU placement to apply_hsdp_to_model to keep peak
-            # load-time memory on CPU. Online quantization (e.g. fp8) runs CUDA-only
-            # kernels inside load_weights via the layerwise loader, so when a quant
-            # config is set we initialize on the accelerator like the non-HSDP path;
-            # apply_hsdp_to_model shards GPU-resident params equally well.
-            hsdp_defer_to_cpu = is_hsdp and self.quant_config is None
-            if not hsdp_defer_to_cpu and target_device.type != "cpu":
+            if not is_hsdp and target_device.type != "cpu":
                 model.to(target_device)
         else:
-            hsdp_defer_to_cpu = is_hsdp and self.quant_config is None
-            device_ctx = contextlib.nullcontext() if hsdp_defer_to_cpu else target_device
+            device_ctx = target_device if not is_hsdp else contextlib.nullcontext()
             with device_ctx:
                 if load_format == "default":
                     model = initialize_model(self.od_config)
@@ -653,22 +646,6 @@ def _load_model_with_hsdp(
         model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
         self.load_weights(model)
 
-        # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights
-        # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel
-        # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params.
-        # Rewrite affected layers in-place to row-major contiguous storage and
-        # shift the .t() to GEMM-call time. Layers using other quant methods or
-        # already-contiguous weights are left untouched.
-        if self.quant_config is not None:
-            from vllm_omni.diffusion.quantization.hsdp_fp8 import (
-                prepare_fp8_layers_for_fsdp,
-            )
-
-            for _attr in ("transformer", "transformer_2"):
-                _trans = getattr(model, _attr, None)
-                if _trans is not None:
-                    prepare_fp8_layers_for_fsdp(_trans)
-
         # Collect all transformers to shard (some models have transformer_2 for MoE)
         transformers_to_shard = []
         transformer = getattr(model, "transformer", None)
diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py
index e6802596762..bc3eef49d68 100644
--- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py
+++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py
@@ -13,7 +13,6 @@
 import numpy as np
 import torch
 import torch.distributed as dist
-from cache_dit import ForwardPattern
 from torch import nn
 from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
 from transformers.models.qwen2.modeling_qwen2 import (
@@ -37,7 +36,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata as DiffusionAttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention
-from vllm_omni.diffusion.cache.cache_dit_backend import BagelCachedAdapter, CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
@@ -774,13 +772,6 @@ def forward(
 
 
 class Qwen2MoTModel(Qwen2PreTrainedModel):
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "layers": ForwardPattern.Pattern_0,
-        },
-        cached_adapter_cls=BagelCachedAdapter,
-    )
-
     _layerwise_offload_blocks_attrs = ["layers"]
 
     @staticmethod
diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
index 4543d692270..59086684628 100644
--- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
@@ -18,7 +18,6 @@
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
-from cache_dit import ForwardPattern
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
@@ -31,7 +30,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
@@ -966,16 +964,6 @@ class Cosmos3VFMTransformer(nn.Module):
     ``Attention`` layer (with Ulysses all-to-all) or plain SDPA accordingly.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        # Cosmos3 GEN blocks return only hidden_states.  Per-layer UND K/V
-        # conditioning uses the transformer's cache-dit fallback path.
-        block_forward_patterns={
-            "gen_layers": ForwardPattern.Pattern_3,
-        },
-        has_separate_cfg=True,
-        check_forward_pattern=False,
-    )
-
     _repeated_blocks = ["Cosmos3GenDecoderLayer"]
 
     _layerwise_offload_blocks_attrs = ["gen_layers"]
diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py
index acf24dd2f64..42393542712 100644
--- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py
+++ b/vllm_omni/diffusion/models/dreamid_omni/fusion.py
@@ -3,13 +3,11 @@
 
 import torch
 import torch.nn as nn
-from cache_dit import ForwardPattern
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -253,13 +251,6 @@ def forward(
 
 
 class FusionModel(nn.Module):
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "fused_blocks": ForwardPattern.Pattern_0,
-        },
-        has_separate_cfg=True,
-    )
-
     _layerwise_offload_blocks_attrs = ["fused_blocks"]
 
     packed_modules_mapping = {
diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
index ad3175e1ed5..e187054e2f6 100644
--- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
+++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
@@ -21,12 +21,11 @@
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
-    tensor_model_parallel_all_reduce,
+    get_tp_group,
 )
 from vllm.model_executor.layers.conv import Conv3dLayer
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
-    QKVParallelLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.utils import set_weight_attrs
@@ -175,82 +174,22 @@ def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> Non
             raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.")
         param.data.copy_(shard)
 
-    def _local_sum_sq(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
-        """Per-rank float32 activation, local sum-of-squares, and local width."""
-        x_float = x.float()
-        local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
-        return x_float, local_sum_sq, x.shape[-1]
-
-    def _scale(
-        self,
-        x_float: torch.Tensor,
-        global_sum_sq: torch.Tensor,
-        global_count: int,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        """Apply the (already reduced) RMS and the per-rank weight shard."""
-        mean_sq = global_sum_sq / global_count
-        return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
-
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         tp_size = get_tensor_model_parallel_world_size()
-        x_float, local_sum_sq, local_count = self._local_sum_sq(x)
+        x_float = x.float()
+        local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
+        local_count = x.shape[-1]
 
         if tp_size > 1:
-            # Use vLLM's collective (custom all-reduce / symmetric-mem fast path
-            # for small tensors) instead of raw torch.distributed.all_reduce, and
-            # take the return value so the custom-AR path (which may return a new
-            # buffer) is handled correctly. No .clone() needed.
-            global_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
+            global_sum_sq = local_sum_sq.clone()
+            torch.distributed.all_reduce(global_sum_sq, group=get_tp_group().device_group)
             global_count = local_count * tp_size
         else:
             global_sum_sq = local_sum_sq
             global_count = local_count
 
-        return self._scale(x_float, global_sum_sq, global_count, x)
-
-
-def fused_qk_rms_norm(
-    norm_q: nn.Module,
-    norm_k: nn.Module,
-    q: torch.Tensor,
-    k: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Apply q/k :class:`DistributedRMSNorm` with a SINGLE fused TP all-reduce.
-
-    WHY: self-attention norms q and k every step; on TP each norm issues its own
-    tiny all-reduce of the per-token sum-of-squares. For DreamZero's decode-like
-    per-step forward these latency-bound micro-collectives dominate, so we pack
-    both sum-of-squares into one tensor and reduce once (2 collectives → 1).
-
-    NUMERICALLY IDENTICAL to ``norm_q(q), norm_k(k)``: all-reduce is elementwise,
-    so packing along the last dim reduces each slice independently with the same
-    fp32 accumulation. Requires q and k to share the same shape (true for
-    self-attention — both come from the same hidden states). Falls back to
-    independent application when either norm is not a DistributedRMSNorm
-    (e.g. nn.Identity when qk_norm=False).
-    """
-    if not (isinstance(norm_q, DistributedRMSNorm) and isinstance(norm_k, DistributedRMSNorm)):
-        return norm_q(q), norm_k(k)
-
-    # The fused path reduces one packed sum-of-squares and reuses q's token width
-    # as the RMS count for BOTH q and k, so q and k must share the same shape.
-    # (Self-attention guarantees this: q and k are projected from the same x.)
-    assert q.shape == k.shape, "fused_qk_rms_norm requires q and k to have the same shape."
-
-    tp_size = get_tensor_model_parallel_world_size()
-    q_float, q_sum_sq, count = norm_q._local_sum_sq(q)
-    k_float, k_sum_sq, _ = norm_k._local_sum_sq(k)
-
-    if tp_size > 1:
-        packed = torch.cat([q_sum_sq, k_sum_sq], dim=-1)
-        packed = tensor_model_parallel_all_reduce(packed)
-        q_sum_sq, k_sum_sq = packed[..., 0:1], packed[..., 1:2]
-        count = count * tp_size
-
-    q_out = norm_q._scale(q_float, q_sum_sq, count, q)
-    k_out = norm_k._scale(k_float, k_sum_sq, count, k)
-    return q_out, k_out
+        mean_sq = global_sum_sq / global_count
+        return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
 
 
 # ── Projections ─────────────────────────────────────────────────────
@@ -398,35 +337,22 @@ def forward(
         context = context[:, 257:]
         n, d = self.tp_num_heads, self.head_dim
         q = self.norm_q(self.q(x)).unflatten(2, (n, d))
-        # context (text) and context_img (clip features) are constant within a
-        # session, so k/v and k_img/v_img are cached on first call and reused.
-        # context_img == img_emb(clip feature) is session-invariant (state.clip_feas
-        # is set once at current_start_frame==0 and only cleared on reset, which also
-        # rebuilds crossattn_cache). Caching k_img/v_img removes a per-step
-        # norm_k_img all-reduce plus the k_img/v_img projection GEMMs in steady state.
-        # Only q depends on the per-step x and is always recomputed.
         if crossattn_cache is not None:
             if not crossattn_cache["is_init"]:
                 crossattn_cache["is_init"] = True
                 k = self.norm_k(self.k(context)).unflatten(2, (n, d))
                 v = self.v(context).unflatten(2, (n, d))
-                k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
-                v_img = self.v_img(context_img).unflatten(2, (n, d))
                 crossattn_cache["k"] = k
                 crossattn_cache["v"] = v
-                crossattn_cache["k_img"] = k_img
-                crossattn_cache["v_img"] = v_img
             else:
                 k = crossattn_cache["k"]
                 v = crossattn_cache["v"]
-                k_img = crossattn_cache["k_img"]
-                v_img = crossattn_cache["v_img"]
         else:
             k = self.norm_k(self.k(context)).unflatten(2, (n, d))
             v = self.v(context).unflatten(2, (n, d))
-            k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
-            v_img = self.v_img(context_img).unflatten(2, (n, d))
         x = self.attn(q, k, v)
+        k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
+        v_img = self.v_img(context_img).unflatten(2, (n, d))
         img_x = self.attn(q, k_img, v_img)
         x = x.flatten(2)
         img_x = img_x.flatten(2)
@@ -477,21 +403,9 @@ def __init__(
         self.num_action_per_block = num_action_per_block
         self.num_state_per_block = num_state_per_block
         self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen
-        # Fused QKV projection: q/k/v all come from x, so a single column-parallel
-        # GEMM replaces three (fewer launches / better GEMM util in the decode-like
-        # per-step forward). No GQA here: total_num_kv_heads defaults to num_heads.
-        self.qkv = QKVParallelLinear(
-            hidden_size=dim,
-            head_size=self.head_dim,
-            total_num_heads=num_heads,
-            bias=True,
-        )
-        # The forward splits qkv into three equal q/k/v shards, which is only
-        # correct without GQA. total_num_kv_heads defaults to total_num_heads
-        # here; fail loud if a future checkpoint introduces GQA so the equal
-        # split does not silently misalign k/v.
-        if self.qkv.total_num_kv_heads != self.qkv.total_num_heads:
-            raise ValueError("Self-attn QKV fusion requires no GQA (total_num_kv_heads == total_num_heads).")
+        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
         self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
         self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
@@ -517,15 +431,9 @@ def forward(
         """Inference-only forward (KV cache path)."""
         n, d = self.tp_num_heads, self.head_dim
 
-        # Single fused QKV GEMM, then split into the per-rank q/k/v shards.
-        qkv, _ = self.qkv(x)
-        qk_size = self.tp_num_heads * self.head_dim
-        q, k, v = qkv.split([qk_size, qk_size, qk_size], dim=-1)
-        # Fuse q/k qk-norm into a single TP all-reduce (both come from x, same shape).
-        q, k = fused_qk_rms_norm(self.norm_q, self.norm_k, q, k)
-        q = q.unflatten(2, (n, d))
-        k = k.unflatten(2, (n, d))
-        v = v.unflatten(2, (n, d))
+        q = self.norm_q(self.q(x)).unflatten(2, (n, d))
+        k = self.norm_k(self.k(x)).unflatten(2, (n, d))
+        v = self.v(x).unflatten(2, (n, d))
 
         updated_kv_cache: torch.Tensor | None = None
 
@@ -926,7 +834,6 @@ def _forward_blocks(
         timestep_action: torch.Tensor | None,
         state: torch.Tensor | None,
         kv_cache: list[torch.Tensor],
-        crossattn_cache: list[dict] | None,
         current_start_frame: int,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
         x = x.flatten(start_dim=2).transpose(1, 2)
@@ -977,7 +884,6 @@ def _forward_blocks(
                 context=context,
                 action_register_length=action_register_length,
                 kv_cache=kv_cache[block_index] if kv_cache else None,
-                crossattn_cache=crossattn_cache[block_index] if crossattn_cache else None,
                 current_start_frame=current_start_frame,
             )
             updated_kv_caches.append(updated_kv_cache)
@@ -1035,7 +941,6 @@ def _forward_inference(
             timestep_action=timestep_action,
             state=state,
             kv_cache=kv_cache,
-            crossattn_cache=crossattn_cache,
             current_start_frame=current_start_frame,
         )
 
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index bbff1237364..30ea2a4cebd 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -23,10 +23,6 @@
 from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
-from vllm_omni.diffusion.cache.stepcache import (
-    get_stepcache_state,
-    is_stepcache_active,
-)
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
@@ -59,6 +55,10 @@
 MAX_DREAMZERO_SESSIONS = 64
 
 
+# ---------------------------------------------------------------------------
+# ---------------------------------------------------------------------------
+
+
 class VideoActionScheduler:
     """Wraps video + action schedulers into single .step() interface."""
 
@@ -205,9 +205,6 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         self._max_session_states = MAX_DREAMZERO_SESSIONS
         self.state = self._get_or_create_state("default")
 
-        # DiT step cache is configured by StepCacheBackend
-        # (cache_backend="step_cache") via pipeline._stepcache_config.
-
         # Keep runtime inference settings separate from the training-time config.
         self.num_inference_steps: int = model_config.get(
             "num_inference_steps",
@@ -309,27 +306,6 @@ def _load_repo_json(model_path: str, relative_path: str, local_files_only: bool)
 
     def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
         """Call CausalWanModel, return (video_pred, action_pred)."""
-        video_pred, action_pred = self._predict_noise_eager(kwargs)
-
-        if is_stepcache_active(self):
-            video_pred = video_pred.clone()
-            if action_pred is not None:
-                action_pred = action_pred.clone()
-
-        if action_pred is None:
-            batch_size = kwargs["hidden_states"].shape[0]
-            action_pred = torch.empty(
-                batch_size,
-                0,
-                self.transformer.action_dim,
-                device=video_pred.device,
-                dtype=video_pred.dtype,
-            )
-        return (video_pred, action_pred)
-
-    def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Eager DiT forward; also handles KV-cache write-back on prefill."""
-        self._cudagraph_mark_step_begin()
         video_pred, action_pred, updated_kv_caches = self.transformer(
             x=kwargs["hidden_states"],
             timestep=kwargs["timestep_video"],
@@ -351,118 +327,19 @@ def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor
             for i, kv in enumerate(updated_kv_caches):
                 state.update_kv_cache(i, kv, is_negative=is_neg)
 
-        return video_pred, action_pred
-
-    # -----------------------------------------------------------------------
-    # torch.compile setup (paper D.2 extended: encoders + VAE + DiT blocks)
-    # -----------------------------------------------------------------------
-
-    def setup_compile(self) -> None:
-        """Compile DreamZero encoders, VAE, and per-block DiT for inference.
-
-        Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False``
-        on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and
-        compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop).
-        DiT blocks use per-block ``fullgraph=True``.
-        """
-        if not torch.cuda.is_available():
-            logger.info("DreamZero setup_compile skipped: CUDA not available.")
-            return
-
-        from vllm_omni.diffusion.models.dreamzero.wan_vae_feat_cache_patch import (
-            apply_wan_vae_feat_cache_tensor_patch,
-        )
-
-        apply_wan_vae_feat_cache_tensor_patch()
-
-        compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
-        # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
-        dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
-
-        logger.info(
-            "DreamZero: torch.compile text/image/VAE encode + per-block DiT (encoders reduce-overhead, DiT default)."
-        )
-
-        try:
-            self.text_encoder.forward = torch.compile(self.text_encoder.forward, **compile_ro)
-        except Exception as exc:
-            logger.warning("DreamZero: text_encoder compile failed (%s); skipping.", exc)
-
-        try:
-            self.image_encoder.model.visual.forward = torch.compile(
-                self.image_encoder.model.visual.forward,
-                **compile_ro,
-            )
-        except Exception as exc:
-            logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
-
-        try:
-            self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
-        except Exception as exc:
-            logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
-
-        compiled_blocks = 0
-        for block in self.transformer.blocks:
-            try:
-                block.forward = torch.compile(block.forward, **dit_compile)
-                compiled_blocks += 1
-            except Exception as exc:
-                logger.warning(
-                    "DreamZero: transformer block %d compile failed (%s); leaving remaining eager.",
-                    compiled_blocks,
-                    exc,
-                )
-                break
-        if compiled_blocks:
-            logger.info(
-                "DreamZero: compiled %d/%d transformer blocks.",
-                compiled_blocks,
-                len(self.transformer.blocks),
-            )
-
-        self.warmup_compile()
-
-    def warmup_compile(self) -> None:
-        """Warm up compiled text/image/VAE paths before timed inference."""
-        if not torch.cuda.is_available():
-            return
-
-        device = next(self.text_encoder.parameters()).device
-        with torch.inference_mode():
-            try:
-                text_tokens = torch.zeros(1, 16, dtype=torch.long, device=device)
-                attention_mask = torch.ones_like(text_tokens)
-                self._encode_text(text_tokens, attention_mask)
-            except Exception as exc:
-                logger.warning("DreamZero compile warmup (text_encoder) skipped: %s", exc)
-
-            try:
-                image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
-                self._encode_image(image, self.num_frames, 180, 320)
-            except Exception as exc:
-                logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
-
-            try:
-                latent_h, latent_w = 180 // 8, 320 // 8
-                dummy_latent = torch.zeros(
-                    1,
-                    16,
-                    self.num_frame_per_block,
-                    latent_h * 2,
-                    latent_w * 2,
-                    dtype=self.vae.dtype,
-                    device=device,
-                )
-                mean = self.vae_latents_mean.to(device=device, dtype=self.vae.dtype)
-                inv_std = self.vae_latents_inv_std.to(device=device, dtype=self.vae.dtype)
-                dummy_latent_denorm = dummy_latent / inv_std + mean
-                self._cudagraph_mark_step_begin()
-                self.vae.decode(dummy_latent_denorm, return_dict=False)
-            except Exception as exc:
-                logger.warning("DreamZero compile warmup (vae decode) skipped: %s", exc)
-
-        torch.accelerator.synchronize(device)
-        logger.info("DreamZero compile warmup finished (text / image / vae decode).")
+        video_pred = video_pred.clone()
+        if action_pred is not None:
+            action_pred = action_pred.clone()
+        else:
+            batch_size = kwargs["hidden_states"].shape[0]
+            action_pred = torch.empty(
+                batch_size,
+                0,
+                self.transformer.action_dim,
+                device=video_pred.device,
+                dtype=video_pred.dtype,
+            )  # CFG-parallel-safe dummy action pred
+        return (video_pred, action_pred)
 
     def combine_cfg_noise(
         self,
@@ -511,20 +388,12 @@ def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor:
             videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4)
         return videos.to(dtype=torch.bfloat16)
 
-    @staticmethod
-    def _cudagraph_mark_step_begin() -> None:
-        try:
-            torch.compiler.cudagraph_mark_step_begin()
-        except Exception:
-            pass
-
     # -----------------------------------------------------------------------
     # Text encoding
     # -----------------------------------------------------------------------
 
     def _encode_text(self, text_tokens: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
         """Encode text prompt via UMT5."""
-        self._cudagraph_mark_step_begin()
         seq_lens = attention_mask.gt(0).sum(dim=1).long()
         prompt_emb = self.text_encoder(
             text_tokens,
@@ -553,7 +422,6 @@ def _encode_image(
         batch_size = image.shape[0]
 
         with torch.amp.autocast(dtype=torch.bfloat16, device_type=device.type):
-            self._cudagraph_mark_step_begin()
             clip_context = self.image_encoder.encode_image(image)
 
             msk = torch.ones(batch_size, num_frames, height // 8, width // 8, device=device)
@@ -601,23 +469,8 @@ def decode_video_latents(self, video_latents: torch.Tensor) -> torch.Tensor:
         inv_std = self.vae_latents_inv_std.to(device=vae_device, dtype=vae_dtype)
         latents = latents / inv_std + mean
         with torch.no_grad():
-            self._cudagraph_mark_step_begin()
             return self.vae.decode(latents, return_dict=False)[0]
 
-    def decode_accumulated_video_latents(self, session_id: str | None = None) -> torch.Tensor:
-        """Decode all AR-chunk latents accumulated for ``session_id``."""
-        state = self._get_or_create_state(session_id)
-        latents = state.get_concatenated_video_latents()
-        if latents is None:
-            session_key = str(session_id or "default")
-            raise RuntimeError(f"No accumulated video latents for session {session_key!r}.")
-        return self.decode_video_latents(latents)
-
-    def clear_accumulated_video_latents(self, session_id: str | None = None) -> None:
-        """Clear accumulated video latents for ``session_id`` without resetting KV state."""
-        state = self._get_or_create_state(session_id)
-        state.clear_video_latents()
-
     # -----------------------------------------------------------------------
     # KV cache prefill
     # -----------------------------------------------------------------------
@@ -782,12 +635,6 @@ def diffuse(
 
         noisy_input = video_latents
         noisy_input_action = action_latents
-        # ---- step_cache (StepCacheBackend / dreamzero.git) ----
-        _cached_flow_pred: torch.Tensor | None = None
-        _cached_flow_pred_action: torch.Tensor | None = None
-        _prev_predictions: list[tuple[torch.Tensor]] = []
-        _step_cache = get_stepcache_state(self) if is_stepcache_active(self) else None
-
         for index in range(len(timesteps_video)):
             video_timestep = timesteps_video[index]
             action_timestep = timesteps_action[index]
@@ -816,58 +663,46 @@ def diffuse(
             else:
                 y = state.ys[:, :, -self.num_frame_per_block :]
 
-            run_dit = _step_cache is None or _step_cache.should_run_step(_prev_predictions)
-            if run_dit:
-                positive_kwargs = dict(
+            positive_kwargs = dict(
+                hidden_states=noisy_input.transpose(1, 2),
+                timestep_video=timestep,
+                encoder_hidden_states=prompt_embeds,
+                kv_cache=state.get_kv_caches(False),
+                crossattn_cache=state.get_crossattn_caches(False),
+                y=y,
+                clip_feature=state.clip_feas,
+                action=noisy_input_action,
+                timestep_action=timestep_action,
+                is_negative=False,
+                **common_kwargs,
+            )
+
+            # Negative (uncond) kwargs
+            if do_true_cfg and negative_prompt_embeds is not None:
+                negative_kwargs = dict(
                     hidden_states=noisy_input.transpose(1, 2),
                     timestep_video=timestep,
-                    encoder_hidden_states=prompt_embeds,
-                    kv_cache=state.get_kv_caches(False),
-                    crossattn_cache=state.get_crossattn_caches(False),
+                    encoder_hidden_states=negative_prompt_embeds,
+                    kv_cache=state.get_kv_caches(True),
+                    crossattn_cache=state.get_crossattn_caches(True),
                     y=y,
                     clip_feature=state.clip_feas,
                     action=noisy_input_action,
                     timestep_action=timestep_action,
-                    is_negative=False,
+                    is_negative=True,
                     **common_kwargs,
                 )
-
-                if do_true_cfg and negative_prompt_embeds is not None:
-                    negative_kwargs = dict(
-                        hidden_states=noisy_input.transpose(1, 2),
-                        timestep_video=timestep,
-                        encoder_hidden_states=negative_prompt_embeds,
-                        kv_cache=state.get_kv_caches(True),
-                        crossattn_cache=state.get_crossattn_caches(True),
-                        y=y,
-                        clip_feature=state.clip_feas,
-                        action=noisy_input_action,
-                        timestep_action=timestep_action,
-                        is_negative=True,
-                        **common_kwargs,
-                    )
-                else:
-                    negative_kwargs = None
-
-                noise_pred = self.predict_noise_maybe_with_cfg(
-                    positive_kwargs=positive_kwargs,
-                    negative_kwargs=negative_kwargs,
-                    do_true_cfg=do_true_cfg,
-                    true_cfg_scale=self.cfg_scale,
-                    cfg_normalize=False,
-                )
-                flow_pred, flow_pred_action = noise_pred
-                _cached_flow_pred = flow_pred
-                _cached_flow_pred_action = flow_pred_action
-
-                _prev_predictions.append((flow_pred,))
-                if _step_cache is not None:
-                    _step_cache.trim_history(_prev_predictions)
             else:
-                # Reuse previous prediction (DiT step-skipping cache).
-                assert _cached_flow_pred is not None, "DiT cache: no cached prediction available for step skip."
-                flow_pred = _cached_flow_pred
-                flow_pred_action = _cached_flow_pred_action
+                negative_kwargs = None
+
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=self.cfg_scale,
+                cfg_normalize=False,
+            )
+            flow_pred, flow_pred_action = noise_pred
 
             latents = (noisy_input, noisy_input_action)
             t = (video_timestep, action_timestep)
@@ -976,12 +811,9 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
         # on the next inference request after websocket reset/session switch.
         if extra_args.get("reset", False):
             state.reset()
-        else:
-            reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
-            if reset_reason == "session":
-                state.reset()
-            elif reset_reason == "inference":
-                state.reset_inference_state()
+        # Auto-reset based on model state (before accumulation)
+        if state.should_reset(text_tokens, 0, self.transformer.local_attn_size):
+            state.reset()
         state.language = text_tokens
 
         # Frame accumulation: stitched single frame → multi-frame video
@@ -1127,8 +959,6 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             video_out = torch.cat([image, video_out], dim=1)
         state.current_start_frame += self.num_frame_per_block
 
-        state.append_video_latents(video_out)
-
         # q99 denorm: [-1,1] → real values
         action_out = self._denormalize_action(action_out.float(), embodiment_name)
 
@@ -1274,27 +1104,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     .replace("img_emb.proj.3.", "img_emb.fc2.")
                     .replace("img_emb.proj.4.", "img_emb.norm2.")
                 )
-
-                # Self-attn q/k/v are fused into a single QKVParallelLinear; route
-                # each separate checkpoint weight/bias to the packed `qkv` param
-                # with its shard id. cross_attn keeps separate q/k/v (q from x,
-                # k/v from context — not fusible), so it is left untouched here.
-                qkv_shard_id: str | None = None
-                for shard_id in ("q", "k", "v"):
-                    needle = f".self_attn.{shard_id}."
-                    if needle in new_name:
-                        new_name = new_name.replace(needle, ".self_attn.qkv.")
-                        qkv_shard_id = shard_id
-                        break
-
                 if new_name in params:
                     param = params[new_name]
-                    if qkv_shard_id is not None:
-                        # QKVParallelLinear.weight_loader needs the shard id.
-                        param.weight_loader(param, tensor, qkv_shard_id)
-                    else:
-                        weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                        weight_loader(param, tensor)
+                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                    weight_loader(param, tensor)
                     loaded.add(new_name)
                 elif new_name in buffers:
                     buffers[new_name].data.copy_(tensor)
diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
index 5815d93d4b1..0972d1b5a2c 100644
--- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
@@ -23,8 +23,7 @@ class DreamZeroState:
     Lifecycle:
         - Created once in DreamZeroPipeline.__init__()
         - Mutated every forward() call (frame append, KV cache grow)
-        - reset() on new session / language change
-        - reset_inference_state() on local_attn_size exceeded (KV only)
+        - reset() on new session / language change / local_attn_size exceeded
     """
 
     def __init__(self) -> None:
@@ -66,106 +65,50 @@ def accumulate_frames(self, stitched: np.ndarray) -> np.ndarray:
         self.call_count += 1
         return np.stack(frames, axis=0)  # (T, H, W, C)
 
-    def append_video_latents(self, video_out: torch.Tensor) -> None:
-        """Append one AR chunk of normalized video latents for later decode.
-
-        Args:
-            video_out: ``(B, T, C, H, W)`` tensor returned by the denoise loop
-                (before the ``transpose(1, 2)`` stored in ``DiffusionOutput``).
-        """
-        if video_out.dim() != 5:
-            raise ValueError(f"Expected 5D video_out, got shape {tuple(video_out.shape)}")
-        # Upstream ``torch.cat(self.video_across_time, dim=2)`` uses (B, C, T, H, W).
-        chunk = video_out.transpose(1, 2).detach().cpu()
-        self.video_latents_across_time.append(chunk)
-        logger.info(
-            "append_video_latents: chunk_shape=%s total_chunks=%d total_latent_t=%d",
-            tuple(chunk.shape),
-            len(self.video_latents_across_time),
-            int(sum(c.shape[2] for c in self.video_latents_across_time)),
-        )
-
-    def get_concatenated_video_latents(self) -> torch.Tensor | None:
-        """Return all accumulated chunks concatenated along the time dimension."""
-        if not self.video_latents_across_time:
-            return None
-        if len(self.video_latents_across_time) == 1:
-            return self.video_latents_across_time[0]
-        return torch.cat(self.video_latents_across_time, dim=2)
-
-    def clear_video_latents(self) -> None:
-        """Drop accumulated video latents without resetting KV/frame state."""
-        self.video_latents_across_time = []
-
     # ------------------------------------------------------------------
     # Reset / should_reset
     # ------------------------------------------------------------------
 
-    def reset(self, *, clear_video_latents: bool = True) -> None:
-        """Clear session state.
-
-        Args:
-            clear_video_latents: When ``False``, keep ``video_latents_across_time``
-                so offline/online export can decode the full AR rollout.
-        """
-        saved_video_latents = [] if clear_video_latents else list(self.video_latents_across_time)
-        self.stitched_buffer = deque(maxlen=FRAMES_PER_CHUNK)
-        self.call_count = 0
-        # Normalized VAE latents per AR chunk, each (B, C, T, H, W) on CPU.
-        # Concatenated along T before decode (matches upstream video_across_time).
-        self.video_latents_across_time = saved_video_latents
+    def reset(self) -> None:
+        """Clear all state."""
+        self.stitched_buffer: deque[np.ndarray] = deque(maxlen=FRAMES_PER_CHUNK)
+        self.call_count: int = 0
 
         # KV cache once robot-policy diffusion supports that integration.
-        self.kv_cache = None
-        self.kv_cache_neg = None
-        self.crossattn_cache = None
-        self.crossattn_cache_neg = None
-        self.current_start_frame = 0
-
-        self.clip_feas = None
-        self.ys = None
-        if clear_video_latents:
-            self.language = None
-
-    def reset_inference_state(self) -> None:
-        """Reset KV/frame state after local attention rolls without dropping video latents."""
-        self.reset(clear_video_latents=False)
-
-    def reset_reason(
-        self,
-        text_tokens: torch.Tensor | None,
-        num_video_frames: int,
-        local_attn_size: int,
-    ) -> str | None:
-        """Return why state should reset before the next forward(), if any."""
+        self.kv_cache: list[torch.Tensor] | None = None
+        self.kv_cache_neg: list[torch.Tensor] | None = None
+        self.crossattn_cache: list[dict[str, bool | torch.Tensor | None]] | None = None
+        self.crossattn_cache_neg: list[dict[str, bool | torch.Tensor | None]] | None = None
+        self.current_start_frame: int = 0
+
+        self.clip_feas: torch.Tensor | None = None
+        self.ys: torch.Tensor | None = None
+        self.language: torch.Tensor | None = None
+
+    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
+        """Determine if state should be reset before this forward()."""
         if self.language is None:
             logger.info("language is None, resetting")
-            return "session"
+            return True
 
         if text_tokens is not None and not torch.equal(self.language, text_tokens):
             logger.info("language changed, resetting")
-            return "session"
+            return True
 
         # NOTE: after accumulate_frames, num_video_frames is the accumulated T
         # (1 for first call, 4 for subsequent). Only reset on true single-frame
         # which happens when the stitched_buffer was cleared externally.
         if num_video_frames == 1 and self.call_count > 1:
             logger.info("single frame input after first call, resetting")
-            return "session"
+            return True
 
         if local_attn_size != -1 and self.current_start_frame >= local_attn_size:
             logger.info(
-                "current_start_frame %d >= local_attn_size %d, resetting inference state",
-                self.current_start_frame,
-                local_attn_size,
+                "current_start_frame %d >= local_attn_size %d, resetting", self.current_start_frame, local_attn_size
             )
-            return "inference"
-
-        return None
+            return True
 
-    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
-        """Determine if state should be reset before this forward()."""
-        return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None
+        return False
 
     # ------------------------------------------------------------------
     # KV cache management
@@ -188,12 +131,8 @@ def create_kv_caches(
             torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
         ]
 
-        self.crossattn_cache = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-        self.crossattn_cache_neg = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
+        self.crossattn_cache = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
+        self.crossattn_cache_neg = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
 
     def update_kv_cache(
         self,
diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
index d039e53c816..20c8cd8cdc7 100644
--- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
+++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
@@ -9,30 +9,13 @@
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
-    @staticmethod
-    def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
-        decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
-        decoded = decoded.clamp(-1, 1) * 0.5 + 0.5
-        return (decoded * 255.0).round().to(torch.uint8).cpu()
-
     def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor:
         if self.model_runner is None or self.model_runner.pipeline is None:
             raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
 
         with torch.inference_mode():
             decoded = self.model_runner.pipeline.decode_video_latents(video_latents)
-            return self._latents_to_uint8_frames(decoded)
-
-    def decode_accumulated_video_latents_to_uint8(self, session_id: str = "default") -> torch.Tensor:
-        """Decode all AR-chunk latents accumulated on the server for one session."""
-        if self.model_runner is None or self.model_runner.pipeline is None:
-            raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
-
-        with torch.inference_mode():
-            decoded = self.model_runner.pipeline.decode_accumulated_video_latents(session_id)
-            return self._latents_to_uint8_frames(decoded)
-
-    def clear_accumulated_video_latents(self, session_id: str = "default") -> None:
-        if self.model_runner is None or self.model_runner.pipeline is None:
-            raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
-        self.model_runner.pipeline.clear_accumulated_video_latents(session_id)
+            decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
+            decoded = decoded.clamp(-1, 1) * 0.5 + 0.5
+            decoded = (decoded * 255.0).round().to(torch.uint8).cpu()
+        return decoded
diff --git a/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py
deleted file mode 100644
index 57d9b866c56..00000000000
--- a/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Make diffusers Wan VAE feat_cache compile-friendly.
-
-diffusers ``WanResample`` stores the string sentinel ``"Rep"`` in feat_cache slots.
-``tensor == "Rep"`` returns a non-Tensor and breaks ``torch.compile(fullgraph=True)``.
-We replace it with a scalar int8 tensor marker (value 0).
-"""
-
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-_PATCHED = False
-
-
-def _is_feat_rep(entry: Any) -> bool:
-    if not isinstance(entry, torch.Tensor):
-        return False
-    return entry.numel() == 1 and entry.dtype == torch.int8
-
-
-def _make_feat_rep(x: torch.Tensor) -> torch.Tensor:
-    return torch.zeros(1, device=x.device, dtype=torch.int8)
-
-
-def _wan_resample_forward_tensorized(self, x, feat_cache=None, feat_idx=[0]):
-    """``WanResample.forward`` with tensor feat_cache sentinel (no ``"Rep"`` string)."""
-    from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T
-
-    b, c, t, h, w = x.size()
-    if self.mode == "upsample3d":
-        if feat_cache is not None:
-            idx = feat_idx[0]
-            if feat_cache[idx] is None:
-                feat_cache[idx] = _make_feat_rep(x)
-                feat_idx[0] += 1
-            else:
-                cache_x = x[:, :, -CACHE_T:, :, :].clone()
-                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and not _is_feat_rep(feat_cache[idx]):
-                    cache_x = torch.cat(
-                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
-                        dim=2,
-                    )
-                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and _is_feat_rep(feat_cache[idx]):
-                    cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
-                if _is_feat_rep(feat_cache[idx]):
-                    x = self.time_conv(x)
-                else:
-                    x = self.time_conv(x, feat_cache[idx])
-                feat_cache[idx] = cache_x
-                feat_idx[0] += 1
-
-                x = x.reshape(b, 2, c, t, h, w)
-                x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
-                x = x.reshape(b, c, t * 2, h, w)
-    t = x.shape[2]
-    x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
-    x = self.resample(x)
-    x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
-
-    if self.mode == "downsample3d":
-        if feat_cache is not None:
-            idx = feat_idx[0]
-            if feat_cache[idx] is None:
-                feat_cache[idx] = x.clone()
-                feat_idx[0] += 1
-            else:
-                cache_x = x[:, :, -1:, :, :].clone()
-                x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
-                feat_cache[idx] = cache_x
-                feat_idx[0] += 1
-    return x
-
-
-def apply_wan_vae_feat_cache_tensor_patch() -> None:
-    """Monkeypatch diffusers WanResample once per process."""
-    global _PATCHED
-    if _PATCHED:
-        return
-
-    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample
-
-    WanResample.forward = _wan_resample_forward_tensorized
-    _PATCHED = True
-    logger.info("DreamZero: applied Wan VAE feat_cache tensor sentinel patch.")
diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py
index 8398cb7aba2..aa9b0d8b79a 100644
--- a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py
+++ b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py
@@ -8,7 +8,6 @@
 
 import torch
 import torch.nn as nn
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -22,7 +21,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -461,12 +459,6 @@ def forward(self, x: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
 
 
 class ErnieImageTransformer2DModel(nn.Module):
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "layers": ForwardPattern.Pattern_3,
-        },
-        check_forward_pattern=False,
-    )
     _repeated_blocks = ["ErnieImageSharedAdaLNBlock"]
     _layerwise_offload_blocks_attrs = ["layers"]
 
diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py
index 9d89368bec0..88d8126c8ad 100644
--- a/vllm_omni/diffusion/models/flux/flux_transformer.py
+++ b/vllm_omni/diffusion/models/flux/flux_transformer.py
@@ -6,7 +6,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     CombinedTimestepGuidanceTextProjEmbeddings,
     CombinedTimestepTextProjEmbeddings,
@@ -26,7 +25,6 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -526,13 +524,6 @@ class FluxTransformer2DModel(nn.Module):
             The dimensions to use for the rotary positional embeddings.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_1,
-            "single_transformer_blocks": ForwardPattern.Pattern_1,
-        }
-    )
-
     # the small and frequently-repeated block(s) of a model
     # -- typically a transformer layer
     # used for torch compile optimizations
diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py
index 5a65db257da..0c6896c131d 100644
--- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py
+++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py
@@ -5,7 +5,6 @@
 from typing import TYPE_CHECKING, Any
 
 import torch
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     TimestepEmbedding,
     Timesteps,
@@ -26,7 +25,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -757,13 +755,6 @@ class Flux2Transformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_1,
-            "single_transformer_blocks": ForwardPattern.Pattern_2,
-        }
-    )
-
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
     _sp_plan = {
         "": {
diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
index cf2c9f78a19..a708a8ece6a 100644
--- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
+++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
@@ -21,7 +21,6 @@
 
 import torch
 import torch.nn as nn
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     TimestepEmbedding,
     Timesteps,
@@ -41,7 +40,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -770,13 +768,6 @@ class Flux2Transformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_1,
-            "single_transformer_blocks": ForwardPattern.Pattern_2,
-        }
-    )
-
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
index c933d4e7f31..56a7066e390 100644
--- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
+++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
@@ -9,7 +9,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from cache_dit import ForwardPattern
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.transformers.transformer_glm_image import GlmImageCombinedTimestepSizeEmbeddings
 from vllm.logger import init_logger
@@ -21,8 +20,6 @@
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
-
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
@@ -938,13 +935,6 @@ class GlmImageTransformer2DModel(CachedTransformer):
             `od_config.tf_model_config`.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_0,
-        },
-        has_separate_cfg=True,
-    )
-
     _repeated_blocks = ["GlmImageTransformerBlock"]
     # SP plan using GlmImagePrepare module for sharding hidden_states and RoPE together.
     # Similar to Qwen-Image's ImageRopePrepare, this creates a module boundary where
diff --git a/vllm_omni/diffusion/models/gr00t/__init__.py b/vllm_omni/diffusion/models/gr00t/__init__.py
deleted file mode 100644
index a10202a40bc..00000000000
--- a/vllm_omni/diffusion/models/gr00t/__init__.py
+++ /dev/null
@@ -1,6 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
-
-__all__ = ["Gr00tN1d7Pipeline"]
diff --git a/vllm_omni/diffusion/models/gr00t/configs/__init__.py b/vllm_omni/diffusion/models/gr00t/configs/__init__.py
deleted file mode 100644
index 9648b859751..00000000000
--- a/vllm_omni/diffusion/models/gr00t/configs/__init__.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-MODEL_CONFIG_TYPES: dict[str, type] = {}
-
-
-def register_model_config(shortname: str, configtype: type) -> None:
-    MODEL_CONFIG_TYPES[shortname] = configtype
diff --git a/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py b/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py
deleted file mode 100644
index 7d11a651cf6..00000000000
--- a/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py
+++ /dev/null
@@ -1,12 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from vllm_omni.diffusion.models.gr00t.dataio.types import (
-    ActionConfig,
-    ActionFormat,
-    ActionRepresentation,
-    ActionType,
-    ModalityConfig,
-)
-
-__all__ = ["ActionConfig", "ActionFormat", "ActionRepresentation", "ActionType", "ModalityConfig"]
diff --git a/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py
deleted file mode 100644
index d7d4527fb93..00000000000
--- a/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py
+++ /dev/null
@@ -1,109 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from transformers import PretrainedConfig
-
-from . import register_model_config
-
-
-def _default_diffusion_model_cfg() -> dict:
-    return {
-        "positional_embeddings": None,
-        "num_layers": 16,
-        "num_attention_heads": 32,
-        "attention_head_dim": 48,
-        "norm_type": "ada_norm",
-        "output_dim": 1024,
-        "interleave_self_attention": True,
-    }
-
-
-class Gr00tN1d7Config(PretrainedConfig):
-    model_type = "Gr00tN1d7"
-
-    def __init__(
-        self,
-        model_dtype: str = "bfloat16",
-        model_name: str = "nvidia/Cosmos-Reason2-2B",
-        backbone_model_type: str = "qwen",
-        model_revision: str | None = None,
-        backbone_embedding_dim: int = 2048,
-        select_layer: int = 12,
-        reproject_vision: bool = False,
-        use_flash_attention: bool = True,
-        load_bf16: bool = False,
-        image_crop_size: tuple | None = (230, 230),
-        image_target_size: tuple | None = (256, 256),
-        shortest_image_edge: int | None = None,
-        crop_fraction: float | None = None,
-        random_rotation_angle: int | None = None,
-        color_jitter_params: dict | None = None,
-        formalize_language: bool = True,
-        apply_sincos_state_encoding: bool = False,
-        use_percentiles: bool = True,
-        use_relative_action: bool = False,
-        max_state_dim: int = 132,
-        max_action_dim: int = 132,
-        action_horizon: int = 40,
-        hidden_size: int = 1024,
-        input_embedding_dim: int = 1536,
-        state_history_length: int = 1,
-        add_pos_embed: bool = True,
-        attn_dropout: float = 0.2,
-        use_vlln: bool = True,
-        max_seq_len: int = 1024,
-        use_alternate_vl_dit: bool = True,
-        attend_text_every_n_blocks: int = 2,
-        diffusion_model_cfg: dict | None = None,
-        vl_self_attention_cfg: dict | None = None,
-        num_inference_timesteps: int = 4,
-        num_timestep_buckets: int = 1000,
-        exclude_state: bool = False,
-        use_mean_std: bool = False,
-        max_num_embodiments: int = 32,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.model_dtype = model_dtype
-        self.model_name = model_name
-        self.backbone_model_type = backbone_model_type
-        self.model_revision = model_revision
-        self.backbone_embedding_dim = backbone_embedding_dim
-        self.select_layer = select_layer
-        self.reproject_vision = reproject_vision
-        self.use_flash_attention = use_flash_attention
-        self.load_bf16 = load_bf16
-        self.image_crop_size = image_crop_size
-        self.image_target_size = image_target_size
-        self.shortest_image_edge = shortest_image_edge
-        self.crop_fraction = crop_fraction
-        self.random_rotation_angle = random_rotation_angle
-        self.color_jitter_params = color_jitter_params
-        self.formalize_language = formalize_language
-        self.apply_sincos_state_encoding = apply_sincos_state_encoding
-        self.use_percentiles = use_percentiles
-        self.use_relative_action = use_relative_action
-        self.max_state_dim = max_state_dim
-        self.max_action_dim = max_action_dim
-        self.action_horizon = action_horizon
-        self.hidden_size = hidden_size
-        self.input_embedding_dim = input_embedding_dim
-        self.state_history_length = state_history_length
-        self.add_pos_embed = add_pos_embed
-        self.attn_dropout = attn_dropout
-        self.use_vlln = use_vlln
-        self.max_seq_len = max_seq_len
-        self.use_alternate_vl_dit = use_alternate_vl_dit
-        self.attend_text_every_n_blocks = attend_text_every_n_blocks
-        self.diffusion_model_cfg = (
-            diffusion_model_cfg if diffusion_model_cfg is not None else _default_diffusion_model_cfg()
-        )
-        self.vl_self_attention_cfg = vl_self_attention_cfg
-        self.num_inference_timesteps = num_inference_timesteps
-        self.num_timestep_buckets = num_timestep_buckets
-        self.exclude_state = exclude_state
-        self.use_mean_std = use_mean_std
-        self.max_num_embodiments = max_num_embodiments
-
-
-register_model_config("Gr00tN1d7", Gr00tN1d7Config)
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/__init__.py b/vllm_omni/diffusion/models/gr00t/dataio/__init__.py
deleted file mode 100644
index 208f01a7cb5..00000000000
--- a/vllm_omni/diffusion/models/gr00t/dataio/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py b/vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py
deleted file mode 100755
index 3ed0c9b0360..00000000000
--- a/vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py
+++ /dev/null
@@ -1,191 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from enum import Enum
-
-"""
-Embodiment tags are used to identify the robot embodiment in the data.
-
-Naming convention:
-<dataset>_<robot_name>
-
-If using multiple datasets, e.g. sim GR1 and real GR1, we can drop the dataset name and use only the robot name.
-"""
-
-
-class EmbodimentTag(Enum):
-    """Embodiment tags supported by the GR00T N1.7 checkpoint.
-
-    Pretrain tags (baked into the base model nvidia/GR00T-N1.7-3B, inference-ready):
-    - OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT -> "oxe_droid_relative_eef_relative_joint"
-    - XDOF                                  -> "xdof_relative_eef_relative_joint"
-    - XDOF_SUBTASK                          -> "xdof_relative_eef_relative_joint_subtask"
-    - REAL_G1                               -> "real_g1_relative_eef_relative_joints"
-    - REAL_R1_PRO_SHARPA                    -> "real_r1_pro_sharpa_relative_eef"
-    - REAL_R1_PRO_SHARPA_HUMAN              -> "real_r1_pro_sharpa_relative_eef_human"
-    - REAL_R1_PRO_SHARPA_MAXINSIGHTS        -> "real_r1_pro_sharpa_relative_eef_maxinsights"
-    - REAL_R1_PRO_SHARPA_MECKA              -> "real_r1_pro_sharpa_relative_eef_mecka"
-
-    Pre-registered posttrain tags (require finetuned checkpoint):
-    - UNITREE_G1           -> "unitree_g1_full_body_with_waist_height_nav_cmd"
-    - UNITREE_G1_SONIC     -> "unitree_g1_sonic"
-    - SIMPLER_ENV_GOOGLE   -> "simpler_env_google"
-    - SIMPLER_ENV_WIDOWX   -> "simpler_env_widowx"
-    - LIBERO_PANDA         -> "libero_sim"
-
-    Finetuning tag (for custom robots):
-    - NEW_EMBODIMENT       -> "new_embodiment"
-
-    Use ``EmbodimentTag.resolve(s)`` to look up a tag by name or value,
-    case-insensitively.
-    """
-
-    OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT = "oxe_droid_relative_eef_relative_joint"
-    """
-    The Open-X-Embodiment DROID robot with relative EEF and relative joint position actions.
-    """
-
-    XDOF = "xdof_relative_eef_relative_joint"
-    """
-    The generic X-DOF robot with relative EEF and relative joint position actions.
-    """
-
-    XDOF_SUBTASK = "xdof_relative_eef_relative_joint_subtask"
-    """
-    The generic X-DOF robot (subtask variant).
-    """
-
-    REAL_G1 = "real_g1_relative_eef_relative_joints"
-    """
-    Real-world Unitree G1 with relative EEF and relative joint actions.
-    """
-
-    REAL_R1_PRO_SHARPA = "real_r1_pro_sharpa_relative_eef"
-    """
-    Real-world R1 Pro Sharpa with relative EEF actions.
-    """
-
-    REAL_R1_PRO_SHARPA_HUMAN = "real_r1_pro_sharpa_relative_eef_human"
-    """
-    Real-world R1 Pro Sharpa with relative EEF actions (human teleop data).
-    """
-
-    REAL_R1_PRO_SHARPA_MAXINSIGHTS = "real_r1_pro_sharpa_relative_eef_maxinsights"
-    """
-    Real-world R1 Pro Sharpa with relative EEF actions (MaxInsights data, single-cam).
-    """
-
-    REAL_R1_PRO_SHARPA_MECKA = "real_r1_pro_sharpa_relative_eef_mecka"
-    """
-    Real-world R1 Pro Sharpa with relative EEF actions (Mecka data, single-cam).
-    """
-
-    UNITREE_G1 = "unitree_g1_full_body_with_waist_height_nav_cmd"
-    """
-    The Unitree G1 robot (sim, full-body with waist height and nav commands).
-    """
-
-    UNITREE_G1_SONIC = "unitree_g1_sonic"
-    """
-    The Unitree G1 robot with SONIC whole-body controller. VLA action space is SONIC latents.
-    """
-
-    SIMPLER_ENV_GOOGLE = "simpler_env_google"
-    """
-    The SimplerEnv Google robot.
-    """
-
-    SIMPLER_ENV_WIDOWX = "simpler_env_widowx"
-    """
-    The SimplerEnv WidowX robot.
-    """
-
-    LIBERO_PANDA = "libero_sim"
-    """
-    The LIBERO Panda robot (used for LIBERO-Goal, LIBERO-Object, LIBERO-Spatial, LIBERO-10).
-    """
-
-    # New embodiment during post-training
-    NEW_EMBODIMENT = "new_embodiment"
-    """
-    Any new embodiment.
-    """
-
-    @classmethod
-    def resolve(cls, tag: "str | EmbodimentTag") -> "EmbodimentTag":
-        """Resolve a string to an EmbodimentTag, case-insensitively.
-
-        Matches by enum **name** first (e.g. ``"xdof"`` -> ``XDOF``), then by
-        enum **value** (e.g. ``"xdof_relative_eef_relative_joint"`` -> ``XDOF``).
-
-        Raises:
-            ValueError: If *tag* does not match any known embodiment.
-        """
-        if isinstance(tag, cls):
-            return tag
-        key = tag.strip()
-        key_lower = key.lower()
-        # Match by enum name (case-insensitive)
-        for member in cls:
-            if member.name.lower() == key_lower:
-                return member
-        # Match by enum value (case-insensitive)
-        for member in cls:
-            if member.value.lower() == key_lower:
-                return member
-
-        def _fmt(tags):
-            return "\n".join(f"    {m.name:40s} -> {m.value}" for m in tags)
-
-        msg = (
-            f"Unknown embodiment tag: {tag!r}\n\n"
-            f"  Base model tags (work with nvidia/GR00T-N1.7-3B):\n"
-            f"{_fmt(PRETRAIN_TAGS)}\n\n"
-            f"  Posttrain tags (require a finetuned checkpoint):\n"
-            f"{_fmt(POSTTRAIN_TAGS)}\n\n"
-            f"  Finetuning-only tags (for custom robots):\n"
-            f"{_fmt(FINETUNE_ONLY_TAGS)}"
-        )
-        raise ValueError(msg)
-
-    @classmethod
-    def reverse_lookup(cls, value: str) -> "str":
-        """Map a tag value string back to its enum name, or return the value as-is."""
-        for member in cls:
-            if member.value == value:
-                return member.name
-        return value
-
-
-# Module-level tag category sets (cannot be Enum class attributes).
-PRETRAIN_TAGS: frozenset[EmbodimentTag] = frozenset(
-    {
-        EmbodimentTag.OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT,
-        EmbodimentTag.XDOF,
-        EmbodimentTag.XDOF_SUBTASK,
-        EmbodimentTag.REAL_G1,
-        EmbodimentTag.REAL_R1_PRO_SHARPA,
-        EmbodimentTag.REAL_R1_PRO_SHARPA_HUMAN,
-        EmbodimentTag.REAL_R1_PRO_SHARPA_MAXINSIGHTS,
-        EmbodimentTag.REAL_R1_PRO_SHARPA_MECKA,
-    }
-)
-"""Tags baked into the base model (nvidia/GR00T-N1.7-3B) — usable without finetuning."""
-
-POSTTRAIN_TAGS: frozenset[EmbodimentTag] = frozenset(
-    {
-        EmbodimentTag.UNITREE_G1,
-        EmbodimentTag.UNITREE_G1_SONIC,
-        EmbodimentTag.SIMPLER_ENV_GOOGLE,
-        EmbodimentTag.SIMPLER_ENV_WIDOWX,
-        EmbodimentTag.LIBERO_PANDA,
-    }
-)
-"""Tags that require a finetuned checkpoint."""
-
-FINETUNE_ONLY_TAGS: frozenset[EmbodimentTag] = frozenset(
-    {
-        EmbodimentTag.NEW_EMBODIMENT,
-    }
-)
-"""Tags for custom robots (finetuning only, not in any shipped checkpoint)."""
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py
deleted file mode 100644
index 208f01a7cb5..00000000000
--- a/vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py
deleted file mode 100644
index 4128d6cae6c..00000000000
--- a/vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py
+++ /dev/null
@@ -1,654 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from collections.abc import Sequence
-from typing import Generic, TypeVar
-
-import numpy as np
-from numpy.typing import NDArray
-from scipy import interpolate
-from scipy.spatial.transform import Rotation, Slerp
-
-from vllm_omni.diffusion.models.gr00t.dataio.state_action.pose import EndEffectorPose, JointPose, Pose
-from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat
-
-PoseType = TypeVar("PoseType", bound=Pose)
-
-
-class ActionChunk(Generic[PoseType]):
-    """
-    Abstract base class for robot action chunking.
-
-    This class provides common functionality for different action chunking types
-    including relative and delta action chunking computation with optional reference frames,
-    interpolation, and format conversion.
-    """
-
-    def __init__(
-        self,
-        poses: Sequence[PoseType],
-        times: Sequence[float] | NDArray[np.float64] | None = None,
-    ):
-        """
-        Initialize action chunking from a list of poses.
-
-        Args:
-            poses: Sequence of Pose objects
-            times: Optional sequence of timestamps for each pose. If None, assumes
-                   uniform spacing starting from 0 with step 1.0
-
-        Raises:
-            ValueError: If action chunking is empty or times length doesn't match poses
-        """
-        if not poses:
-            raise ValueError("ActionChunk must contain at least one pose")
-
-        self._poses: list[PoseType] = list(poses)
-
-        # Set up times
-        if times is None:
-            self._times = np.arange(len(poses), dtype=np.float64)
-        else:
-            if len(times) != len(poses):
-                raise ValueError("Number of times must match number of poses")
-            self._times = np.array(times, dtype=np.float64)
-
-    @property
-    def poses(self) -> list[PoseType]:
-        """Get the list of poses"""
-        return self._poses.copy()
-
-    @property
-    def times(self) -> NDArray[np.float64]:
-        """Get the timestamps"""
-        return self._times.copy()
-
-    @property
-    def num_poses(self) -> int:
-        """Get the number of poses in the action chunking"""
-        return len(self._poses)
-
-    def relative_chunking(self, reference_frame: PoseType | None = None) -> "ActionChunk[PoseType]":
-        """
-        Compute the relative action chunking with respect to a reference frame.
-
-        If reference_frame is None, uses the first pose in the action chunking as reference.
-        All poses are transformed to be relative to the reference frame.
-
-        Args:
-            reference_frame: Optional reference pose. If None, uses first pose.
-
-        Returns:
-            A new ActionChunk of the same type where all poses are relative to the reference frame.
-        """
-        if not self._poses:
-            return self.__class__([], times=[])
-
-        # Use the first pose as the reference if one is not provided.
-        ref_pose = reference_frame if reference_frame is not None else self._poses[0]
-
-        # Use the polymorphic subtraction defined in the Pose subclasses.
-        # The subtraction returns the same type as the operands
-        relative_poses: list[PoseType] = [pose - ref_pose for pose in self._poses]  # type: ignore[misc]
-
-        # Return a new instance of the same action chunking class
-        # (e.g., JointActionChunk or EndEffectorActionChunk)
-        return self.__class__(relative_poses, times=self.times)
-
-    def delta_chunking(self, reference_frame: PoseType | None = None) -> "ActionChunk[PoseType]":
-        """
-        Compute the delta action chunking where each pose represents the relative
-        transformation from the previous frame.
-
-        If reference_frame is provided, it is treated as the first frame, and the
-        first delta will be from reference_frame to the first pose in the action chunking.
-        Otherwise, the first pose in the delta action chunking will be the identity/zero transformation.
-
-        Args:
-            reference_frame: Optional reference pose to use as the first frame.
-
-        Returns:
-            A new ActionChunk of the same type where each pose is relative to the previous pose.
-        """
-        if not self._poses:
-            return self.__class__([], times=[])
-
-        delta_poses: list[PoseType] = []
-
-        # Determine the initial reference for the very first pose.
-        # If a reference_frame is given, the first delta is pose[0] - reference_frame.
-        # If not, the first delta is pose[0] - pose[0], resulting in an identity/zero pose.
-        prev_pose = reference_frame if reference_frame is not None else self._poses[0]
-
-        for current_pose in self._poses:
-            delta: PoseType = current_pose - prev_pose  # type: ignore[assignment]
-            delta_poses.append(delta)
-            prev_pose = current_pose  # Update the reference for the next step
-
-        return self.__class__(delta_poses, times=self.times.tolist())
-
-    def to_absolute_chunking(self, reference_frame: PoseType) -> "ActionChunk[PoseType]":
-        """
-        Convert a relative action chunking to an absolute action chunking by applying
-        the relative poses on top of a reference frame.
-
-        This is the inverse operation of relative_chunking(). Each relative pose
-        is composed with the reference frame to produce absolute poses.
-
-        Args:
-            reference_frame: The reference pose to apply the relative action chunking on top of.
-
-        Returns:
-            A new ActionChunk of the same type with absolute poses.
-
-        Raises:
-            NotImplementedError: Must be implemented by subclasses
-        """
-        raise NotImplementedError("Subclasses must implement to_absolute_chunking")
-
-    def interpolate(
-        self,
-        num_points: int | None = None,
-        times: NDArray[np.float64] | None = None,
-    ) -> "ActionChunk":
-        """
-        Interpolate the action chunking to generate intermediate poses.
-        Must be implemented by subclasses.
-
-        Args:
-            num_points: Number of evenly-spaced points to generate
-            times: Specific timestamps at which to interpolate
-
-        Returns:
-            A new ActionChunk with interpolated poses
-        """
-        raise NotImplementedError("Subclasses must implement interpolate")
-
-    def to(self, action_format: ActionFormat) -> NDArray[np.float64]:
-        """
-        Convert action chunking to the specified action format.
-        Must be implemented by subclasses.
-
-        Args:
-            action_format: The desired output format
-
-        Returns:
-            Array in the requested format
-
-        Raises:
-            NotImplementedError: If not implemented by subclass
-        """
-        raise NotImplementedError("Subclasses must implement to method")
-
-    def __len__(self) -> int:
-        """Return the number of poses in the action chunking"""
-        return len(self._poses)
-
-    def __getitem__(self, index: int) -> PoseType:
-        """Get a pose by index"""
-        return self._poses[index]
-
-    def __repr__(self) -> str:
-        return (
-            f"{self.__class__.__name__}(num_poses={len(self._poses)}, "
-            f"time_range=[{self._times[0]:.2f}, {self._times[-1]:.2f}])"
-        )
-
-
-class JointActionChunk(ActionChunk[JointPose]):
-    """
-    Represents action chunking in joint space as a sequence of joint configurations.
-
-    Examples:
-        # Create a joint action chunking
-        joint_poses = [
-            JointPose([0.0, 0.0, 0.0]),
-            JointPose([0.5, 0.5, 0.5]),
-            JointPose([1.0, 1.0, 1.0]),
-        ]
-        action_chunking = JointActionChunk(joint_poses)
-
-        # Get relative trajectory (all poses relative to first pose)
-        relative_traj = action_chunking.relative_chunking()
-
-        # Get relative trajectory with custom reference
-        reference = JointPose([0.1, 0.1, 0.1])
-        relative_traj = action_chunking.relative_chunking(reference_frame=reference)
-
-        # Get delta trajectory (incremental changes)
-        delta_traj = action_chunking.delta_chunking()
-
-        # Convert relative trajectory back to absolute
-        reference = JointPose([0.1, 0.1, 0.1])
-        absolute_traj = relative_traj.to_absolute_chunking(reference_frame=reference)
-
-        # Interpolate trajectory
-        interpolated = action_chunking.interpolate(num_points=10)
-
-        # Convert to desired format
-        from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat
-        array_data = action_chunking.to(ActionFormat.DEFAULT)  # Returns joint array
-    """
-
-    def __init__(
-        self,
-        poses: Sequence[JointPose],
-        times: Sequence[float] | NDArray[np.float64] | None = None,
-    ):
-        """
-        Initialize a joint trajectory from a list of joint poses.
-
-        Args:
-            poses: Sequence of JointPose objects
-            times: Optional sequence of timestamps for each pose
-
-        Raises:
-            TypeError: If poses are not all JointPose objects
-        """
-        # Validate all poses are JointPose
-        if not all(isinstance(p, JointPose) for p in poses):
-            raise TypeError("All poses must be JointPose objects for JointActionChunk")
-
-        super().__init__(poses, times)
-
-    def interpolate(
-        self,
-        num_points: int | None = None,
-        times: NDArray[np.float64] | None = None,
-    ) -> "JointActionChunk":
-        """
-        Interpolate the joint action chunking to generate intermediate configurations.
-
-        Uses linear interpolation for joint values.
-
-        Args:
-            num_points: Number of evenly-spaced points to generate (including endpoints).
-                       Only used if times is None.
-            times: Specific timestamps at which to interpolate. If provided,
-                  num_points is ignored.
-
-        Returns:
-            A new JointActionChunk with interpolated poses
-
-        Raises:
-            ValueError: If neither num_points nor times is provided, or if
-                       interpolation times are outside the trajectory range
-        """
-        if num_points is None and times is None:
-            raise ValueError("Must provide either num_points or times")
-
-        if len(self._poses) < 2:
-            raise ValueError("Need at least 2 poses for interpolation")
-
-        # Prepare data: extract joint values
-        timestamps = self._times.copy()
-        joint_values = np.array([pose.joints for pose in self._poses])  # (N, num_joints)
-
-        # Find and remove non-monotonic timestamps
-        drop_indices = [idx for idx in range(1, len(timestamps)) if timestamps[idx] <= timestamps[idx - 1]]
-
-        if drop_indices:
-            for idx in drop_indices:
-                print(
-                    f"Dropping timestamp pair - Previous: {timestamps[idx - 1]}, "
-                    f"Current: {timestamps[idx]} at index {idx}"
-                )
-            timestamps = np.delete(timestamps, drop_indices)
-            joint_values = np.delete(joint_values, drop_indices, axis=0)
-
-        # Check if we still have enough poses after cleanup
-        if len(timestamps) < 2:
-            raise ValueError("Need at least 2 poses with monotonic timestamps for interpolation")
-
-        # Create interpolator
-        joint_interp = interpolate.interp1d(timestamps, joint_values, kind="linear", axis=0)
-
-        # Generate interpolation times if not provided
-        if times is None:
-            assert num_points is not None  # Type narrowing for type checker
-            interp_times = np.linspace(timestamps[0], timestamps[-1], num_points)
-        else:
-            interp_times = np.array(times, dtype=np.float64)
-
-        # Check that interpolation times are within bounds
-        if np.any(interp_times < timestamps[0]) or np.any(interp_times > timestamps[-1]):
-            raise ValueError(f"Interpolation times must be within [{timestamps[0]}, {timestamps[-1]}]")
-
-        # Interpolate joint values
-        interp_joint_values = joint_interp(interp_times)
-
-        # Create interpolated poses
-        joint_names = self._poses[0].joint_names
-        interpolated_poses = [
-            JointPose(joints=interp_joint_values[i], joint_names=joint_names) for i in range(len(interp_times))
-        ]
-
-        return JointActionChunk(interpolated_poses, times=interp_times)
-
-    def to_array(self) -> NDArray[np.float64]:
-        """
-        Convert trajectory to array of joint values.
-
-        Returns:
-            Array with shape (N, num_joints) where N is the number of poses
-        """
-        return np.array([pose.joints for pose in self._poses])
-
-    def to_absolute_chunking(self, reference_frame: JointPose) -> "JointActionChunk":
-        """
-        Convert a relative joint action chunking to an absolute action chunking by adding
-        the relative joint positions to the reference frame.
-
-        This is the inverse operation of relative_chunking(). Each relative joint
-        configuration is added to the reference frame to produce absolute joint positions.
-
-        Args:
-            reference_frame: The reference joint pose to apply the relative trajectory on top of.
-
-        Returns:
-            A new JointActionChunk with absolute joint positions.
-
-        Raises:
-            ValueError: If joint dimensions don't match
-        """
-        if not self._poses:
-            return JointActionChunk([], times=[])
-
-        if len(self._poses[0].joints) != len(reference_frame.joints):
-            raise ValueError(
-                f"Cannot apply relative trajectory: "
-                f"joint dimensions don't match ({len(self._poses[0].joints)} vs {len(reference_frame.joints)})"
-            )
-
-        # Add each relative pose to the reference frame
-        absolute_poses: list[JointPose] = []
-        for relative_pose in self._poses:
-            absolute_joints = reference_frame.joints + relative_pose.joints
-            absolute_pose = JointPose(joints=absolute_joints, joint_names=reference_frame.joint_names)
-            absolute_poses.append(absolute_pose)
-
-        return JointActionChunk(absolute_poses, times=self.times)
-
-    def to(self, action_format: ActionFormat) -> NDArray[np.float64]:
-        """
-        Convert trajectory to the desired format.
-
-        Args:
-            action_format: The desired output format
-
-        Returns:
-            Array in the requested format
-
-        Raises:
-            ValueError: If the action format is not supported for joint trajectories
-        """
-        if action_format == ActionFormat.DEFAULT:
-            return self.to_array()
-        else:
-            raise ValueError(
-                f"ActionFormat {action_format} is not supported for JointActionChunk. "
-                f"Only {ActionFormat.DEFAULT} is supported."
-            )
-
-
-class EndEffectorActionChunk(ActionChunk[EndEffectorPose]):
-    """
-    Represents action chunking in Cartesian space as a sequence of end-effector poses.
-
-    Examples:
-        # Create an end-effector action chunking
-        ee_poses = [
-            EndEffectorPose(translation=[0, 0, 0], rotation=[1, 0, 0, 0],
-                          rotation_type="quat", rotation_order="wxyz"),
-            EndEffectorPose(translation=[1, 0, 0], rotation=[0.707, 0, 0, 0.707],
-                          rotation_type="quat", rotation_order="wxyz"),
-            EndEffectorPose(translation=[2, 0, 0], rotation=[0, 0, 0, 1],
-                          rotation_type="quat", rotation_order="wxyz"),
-        ]
-        action_chunking = EndEffectorActionChunk(ee_poses)
-
-        # Get relative trajectory (all poses relative to first pose)
-        relative_traj = action_chunking.relative_chunking()
-
-        # Get relative trajectory with custom reference frame
-        reference = EndEffectorPose(translation=[0.5, 0, 0], rotation=[1,0,0,0],
-                                   rotation_type="quat", rotation_order="wxyz")
-        relative_traj = action_chunking.relative_chunking(reference_frame=reference)
-
-        # Get delta trajectory
-        delta_traj = action_chunking.delta_chunking()
-
-        # Convert relative trajectory back to absolute
-        reference = EndEffectorPose(translation=[0.5, 0, 0], rotation=[1,0,0,0],
-                                   rotation_type="quat", rotation_order="wxyz")
-        absolute_traj = relative_traj.to_absolute_chunking(reference_frame=reference)
-
-        # Interpolate trajectory
-        interpolated = action_chunking.interpolate(num_points=10)
-
-        # Convert to desired format
-        from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat
-        homo_matrices = action_chunking.to(ActionFormat.DEFAULT)      # (N, 4, 4) homogeneous matrices
-        xyz_rot6d = action_chunking.to(ActionFormat.XYZ_ROT6D)        # (N, 9) xyz + rot6d
-        xyz_rotvec = action_chunking.to(ActionFormat.XYZ_ROTVEC)      # (N, 6) xyz + rotvec
-    """
-
-    def __init__(
-        self,
-        poses: Sequence[EndEffectorPose],
-        times: Sequence[float] | NDArray[np.float64] | None = None,
-    ):
-        """
-        Initialize an end-effector trajectory from a list of end-effector poses.
-
-        Args:
-            poses: Sequence of EndEffectorPose objects
-            times: Optional sequence of timestamps for each pose
-
-        Raises:
-            TypeError: If poses are not all EndEffectorPose objects
-        """
-        # Validate all poses are EndEffectorPose
-        if not all(isinstance(p, EndEffectorPose) for p in poses):
-            raise TypeError("All poses must be EndEffectorPose objects for EndEffectorActionChunk")
-
-        super().__init__(poses, times)
-
-    @classmethod
-    def from_array(cls, data: np.ndarray, action_format: ActionFormat) -> "EndEffectorActionChunk":
-        """
-        Create an EndEffectorActionChunk from a 2-D array using the specified action format.
-
-        This is the inverse of ``.to(action_format)``.
-
-        Args:
-            data: Array of shape (N, D) where D depends on the action_format.
-            action_format: The format that describes the layout of each row.
-
-        Returns:
-            EndEffectorActionChunk with N poses.
-        """
-        poses = [EndEffectorPose.from_action_format(row, action_format) for row in data]
-        return cls(poses)
-
-    def interpolate(
-        self,
-        num_points: int | None = None,
-        times: NDArray[np.float64] | None = None,
-    ) -> "EndEffectorActionChunk":
-        """
-        Interpolate the action chunking to generate intermediate poses.
-
-        Uses linear interpolation for translation and SLERP (Spherical Linear
-        Interpolation) for rotation.
-
-        Args:
-            num_points: Number of evenly-spaced points to generate (including endpoints).
-                       Only used if times is None.
-            times: Specific timestamps at which to interpolate. If provided,
-                  num_points is ignored.
-
-        Returns:
-            A new EndEffectorActionChunk with interpolated poses
-
-        Raises:
-            ValueError: If neither num_points nor times is provided, or if
-                       interpolation times are outside the trajectory range
-        """
-        if num_points is None and times is None:
-            raise ValueError("Must provide either num_points or times")
-
-        if len(self._poses) < 2:
-            raise ValueError("Need at least 2 poses for interpolation")
-
-        # Prepare data: extract positions and rotations
-        timestamps = self._times.copy()
-        homogeneous_matrices = np.array([pose.homogeneous for pose in self._poses])
-        positions = homogeneous_matrices[:, :3, 3]
-        rotations = Rotation.from_matrix(homogeneous_matrices[:, :3, :3])
-
-        # Find indices where timestamps are not monotonically increasing
-        drop_indices = [idx for idx in range(1, len(timestamps)) if timestamps[idx] <= timestamps[idx - 1]]
-
-        # Remove the problematic timestamps and corresponding data
-        if drop_indices:
-            for idx in drop_indices:
-                print(
-                    f"Dropping timestamp pair - Previous: {timestamps[idx - 1]}, "
-                    f"Current: {timestamps[idx]} at index {idx}"
-                )
-            timestamps = np.delete(timestamps, drop_indices)
-            positions = np.delete(positions, drop_indices, axis=0)
-            rotations = Rotation.from_matrix(np.delete(homogeneous_matrices[:, :3, :3], drop_indices, axis=0))
-
-        # Check if we still have enough poses after cleanup
-        if len(timestamps) < 2:
-            raise ValueError("Need at least 2 poses with monotonic timestamps for interpolation")
-
-        # Create interpolators
-        pos_interp = interpolate.interp1d(timestamps, positions, kind="linear", axis=0)
-        rot_interp = Slerp(timestamps, rotations)
-
-        # Generate interpolation times if not provided
-        if times is None:
-            assert num_points is not None  # Type narrowing for type checker
-            interp_times = np.linspace(timestamps[0], timestamps[-1], num_points)
-        else:
-            interp_times = np.array(times, dtype=np.float64)
-
-        # Check that interpolation times are within bounds
-        if np.any(interp_times < timestamps[0]) or np.any(interp_times > timestamps[-1]):
-            raise ValueError(f"Interpolation times must be within [{timestamps[0]}, {timestamps[-1]}]")
-
-        # Interpolate positions and rotations
-        interp_positions = pos_interp(interp_times)
-        interp_rotations = rot_interp(interp_times)
-
-        # Create interpolated poses
-        interpolated_poses = []
-        for i in range(len(interp_times)):
-            pose = EndEffectorPose(
-                translation=interp_positions[i],
-                rotation=interp_rotations[i].as_matrix(),
-                rotation_type="matrix",
-            )
-            interpolated_poses.append(pose)
-
-        return EndEffectorActionChunk(interpolated_poses, times=interp_times)
-
-    def to_homogeneous_matrices(self) -> NDArray[np.float64]:
-        """
-        Convert trajectory to array of homogeneous transformation matrices.
-
-        Returns:
-            Array of homogeneous matrices with shape (N, 4, 4) where N is the number of poses
-        """
-        return np.array([pose.homogeneous for pose in self._poses])
-
-    def to_translation_rot6d(self) -> NDArray[np.float64]:
-        """
-        Convert trajectory to array of translations and 6D rotations.
-
-        Returns:
-            Array with shape (N, 9) - 3 for xyz + 6 for rot6d
-        """
-        translations = np.array([pose.translation for pose in self._poses])  # (N, 3)
-        rotations = np.array([pose.rot6d for pose in self._poses])  # (N, 6)
-
-        # Concatenate translation and rotation
-        xyz_rot6d = np.concatenate([translations, rotations], axis=1)  # (N, 9)
-
-        return xyz_rot6d
-
-    def to_translation_rotvec(self) -> NDArray[np.float64]:
-        """
-        Convert trajectory to array of translations and rotation vectors.
-
-        Returns:
-            Array with shape (N, 6) - 3 for xyz + 3 for rotvec
-        """
-        translations = np.array([pose.translation for pose in self._poses])  # (N, 3)
-        rotations = np.array([pose.rotvec for pose in self._poses])  # (N, 3)
-
-        # Concatenate translation and rotation
-        xyz_rotvec = np.concatenate([translations, rotations], axis=1)  # (N, 6)
-
-        return xyz_rotvec
-
-    def to_absolute_chunking(self, reference_frame: EndEffectorPose) -> "EndEffectorActionChunk":
-        """
-        Convert a relative end-effector action chunking to an absolute action chunking by
-        composing each relative transformation with the reference frame.
-
-        This is the inverse operation of relative_chunking(). Each relative pose
-        represents a transformation that is applied on top of the reference frame
-        to produce absolute poses.
-
-        Args:
-            reference_frame: The reference end-effector pose to apply the relative trajectory on top of.
-
-        Returns:
-            A new EndEffectorActionChunk with absolute poses.
-        """
-        if not self._poses:
-            return EndEffectorActionChunk([], times=[])
-
-        # Get reference frame as homogeneous matrix
-        T_ref = reference_frame.homogeneous
-
-        # Compose each relative transformation with the reference frame
-        absolute_poses: list[EndEffectorPose] = []
-        for relative_pose in self._poses:
-            # Get relative transformation as homogeneous matrix
-            T_relative = relative_pose.homogeneous
-
-            # Compose transformations: T_absolute = T_ref @ T_relative
-            T_absolute = T_ref @ T_relative
-
-            # Create absolute pose from composed transformation
-            absolute_pose = EndEffectorPose(homogeneous=T_absolute)
-            absolute_poses.append(absolute_pose)
-
-        return EndEffectorActionChunk(absolute_poses, times=self.times)
-
-    def to(self, action_format: ActionFormat) -> NDArray[np.float64]:
-        """
-        Convert trajectory to the desired format.
-
-        Args:
-            action_format: The desired output format
-
-        Returns:
-            Array in the requested format
-
-        Raises:
-            ValueError: If the action format is not supported
-        """
-        if action_format == ActionFormat.DEFAULT:
-            return self.to_homogeneous_matrices()
-        elif action_format == ActionFormat.XYZ_ROT6D:
-            return self.to_translation_rot6d()
-        elif action_format == ActionFormat.XYZ_ROTVEC:
-            return self.to_translation_rotvec()
-        else:
-            raise ValueError(f"Unsupported action format: {action_format}")
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py
deleted file mode 100644
index 0e0849f17b4..00000000000
--- a/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py
+++ /dev/null
@@ -1,709 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from enum import Enum
-from typing import TypeVar
-
-import numpy as np
-from numpy.typing import NDArray
-from scipy.spatial.transform import Rotation
-
-from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat
-
-# TypeVar for self-type preservation in Pose operations
-PoseT = TypeVar("PoseT", bound="Pose")
-
-
-def invert_transformation(transform: NDArray[np.float64]) -> NDArray[np.float64]:
-    """
-    Invert a homogeneous transformation matrix.
-
-    Args:
-        transform: A 4x4 homogeneous transformation matrix
-
-    Returns:
-        The inverse of the transformation matrix (4x4)
-    """
-    R = transform[:3, :3]  # Extract the rotation matrix
-    t = transform[:3, 3]  # Extract the translation vector
-
-    # Inverse of the rotation matrix is its transpose (since it's orthogonal)
-    R_inv = R.T
-
-    # Inverse of the translation is -R_inv * t
-    t_inv = -R_inv @ t
-
-    # Construct the inverse transformation matrix
-    T_inv = np.eye(4)
-    T_inv[:3, :3] = R_inv
-    T_inv[:3, 3] = t_inv
-
-    return T_inv
-
-
-def relative_transformation(
-    base_transform: NDArray[np.float64], target_transform: NDArray[np.float64]
-) -> NDArray[np.float64]:
-    """
-    Compute the relative transformation between two poses.
-
-    Args:
-        base_transform: Initial 4x4 homogeneous transformation matrix
-        target_transform: Current 4x4 homogeneous transformation matrix
-
-    Returns:
-        The relative transformation matrix (4x4) from base_transform to target_transform
-    """
-    # Relative transformation is base_transform^{-1} * target_transform
-    T_relative = invert_transformation(base_transform) @ target_transform
-    return T_relative
-
-
-class RotationType(Enum):
-    """Supported rotation representation types"""
-
-    QUAT = "quat"
-    EULER = "euler"
-    ROTVEC = "rotvec"
-    MATRIX = "matrix"
-    ROT6D = "rot6d"
-
-
-class EulerOrder(Enum):
-    """Common Euler angle conventions"""
-
-    XYZ = "xyz"
-    ZYX = "zyx"
-    XZY = "xzy"
-    YXZ = "yxz"
-    YZX = "yzx"
-    ZXY = "zxy"
-
-
-class QuatOrder(Enum):
-    """Quaternion ordering conventions"""
-
-    WXYZ = "wxyz"  # scalar-first (w, x, y, z)
-    XYZW = "xyzw"  # scalar-last (x, y, z, w)
-
-
-class Pose:
-    """
-    Abstract base class for robot poses.
-
-    This class provides common functionality for different pose representations
-    including relative pose computation via the subtraction operator.
-    """
-
-    pose_type: str
-
-    def __sub__(self: PoseT, other: PoseT) -> PoseT:
-        """
-        Compute relative transformation between two poses.
-
-        For EndEffectorPose: Computes the relative transformation from other to self.
-        Result represents the transformation needed to go from other's frame to self's frame.
-
-        For JointPose: Computes the joint-space difference (self - other).
-
-        Args:
-            other: The reference pose to compute relative transformation from
-
-        Returns:
-            Relative pose (same type as self)
-
-        Raises:
-            TypeError: If poses are not of the same type
-
-        Examples:
-            # End-effector poses
-            pose1 = EndEffectorPose(translation=[1, 0, 0], rotation=[1,0,0,0],
-                                   rotation_type="quat", rotation_order="wxyz")
-            pose2 = EndEffectorPose(translation=[2, 0, 0], rotation=[1,0,0,0],
-                                   rotation_type="quat", rotation_order="wxyz")
-            relative = pose2 - pose1  # Transformation from pose1 to pose2
-
-            # Joint poses
-            joint1 = JointPose([0.0, 0.5, 1.0])
-            joint2 = JointPose([0.1, 0.6, 1.2])
-            joint_diff = joint2 - joint1  # Joint differences: [0.1, 0.1, 0.2]
-        """
-        if type(self) is not type(other):
-            raise TypeError(
-                f"Cannot compute relative transformation between different pose types: "
-                f"{type(self).__name__} and {type(other).__name__}"
-            )
-
-        return self._compute_relative(other)
-
-    def _compute_relative(self: PoseT, other: PoseT) -> PoseT:
-        """
-        Internal method to compute relative transformation.
-        Must be implemented by subclasses.
-
-        Args:
-            other: The reference pose
-
-        Returns:
-            Relative pose
-        """
-        raise NotImplementedError("Subclasses must implement _compute_relative")
-
-    def copy(self: PoseT) -> PoseT:
-        """
-        Create a deep copy of this pose.
-        Must be implemented by subclasses.
-
-        Returns:
-            New Pose instance with copied data
-        """
-        raise NotImplementedError("Subclasses must implement copy")
-
-
-class JointPose(Pose):
-    """
-    Represents a robot configuration in joint space.
-
-    This class stores joint angles/positions for a robot manipulator.
-    Unlike end-effector poses, joint poses represent the configuration
-    of all joints in the kinematic chain.
-
-    Examples:
-        # Create a 6-DOF joint configuration
-        joint_pose = JointPose(
-            joints=[0.0, -np.pi/4, np.pi/2, 0.0, np.pi/4, 0.0],
-            joint_names=["shoulder_pan", "shoulder_lift", "elbow",
-                        "wrist_1", "wrist_2", "wrist_3"]
-        )
-
-        # Create with default joint names
-        joint_pose = JointPose(joints=[0.0, 0.5, 1.0])
-
-        # Get as dictionary
-        joint_dict = joint_pose.to_dict()  # {"joint_0": 0.0, ...}
-
-        # Access individual joints
-        first_joint = joint_pose.joints[0]
-        num_joints = joint_pose.num_joints
-
-        # Compute relative joint displacement
-        joint1 = JointPose([0.0, 0.5, 1.0])
-        joint2 = JointPose([0.1, 0.6, 1.2])
-        relative = joint2 - joint1  # [0.1, 0.1, 0.2]
-    """
-
-    pose_type = "joint"
-
-    def __init__(
-        self,
-        joints: list | np.ndarray,
-        joint_names: list | None = None,
-    ):
-        """
-        Initialize a joint pose.
-
-        Args:
-            joints: Joint angles/positions as array-like of shape (n,)
-            joint_names: Optional list of names for each joint. If None,
-                        defaults to ["joint_0", "joint_1", ...]
-        """
-        super().__init__()
-        self.joints = np.array(joints, dtype=np.float64)
-
-        # Set defaults and validate joint_names
-        if joint_names is None:
-            self.joint_names = [f"joint_{i}" for i in range(len(self.joints))]
-        else:
-            if len(joint_names) != len(self.joints):
-                raise ValueError(
-                    f"Number of joint names ({len(joint_names)}) must match number of joints ({len(self.joints)})"
-                )
-            self.joint_names = joint_names
-
-    @property
-    def num_joints(self) -> int:
-        """
-        Get the number of joints.
-
-        Returns:
-            Number of joints in the configuration
-        """
-        return len(self.joints)
-
-    def to_dict(self) -> dict:
-        """
-        Convert joint configuration to dictionary.
-
-        Returns:
-            Dictionary mapping joint names to joint values
-        """
-        return dict(zip(self.joint_names, self.joints))
-
-    def _compute_relative(self, other):  # type: ignore[override]
-        """
-        Compute relative joint displacement.
-
-        Args:
-            other: Reference joint pose
-
-        Returns:
-            JointPose representing the joint-space difference (self - other)
-
-        Raises:
-            ValueError: If joint dimensions don't match
-        """
-        if len(self.joints) != len(other.joints):
-            raise ValueError(
-                f"Cannot compute relative joint pose: "
-                f"joint dimensions don't match ({len(self.joints)} vs {len(other.joints)})"
-            )
-
-        relative_joints = self.joints - other.joints
-        return JointPose(joints=relative_joints, joint_names=self.joint_names)
-
-    def copy(self) -> "JointPose":
-        """
-        Create a deep copy of this joint pose.
-
-        Returns:
-            New JointPose instance with copied data
-        """
-        return JointPose(
-            joints=self.joints.copy(),
-            joint_names=self.joint_names.copy(),
-        )
-
-    def __repr__(self) -> str:
-        if len(self.joints) <= 6:
-            joints_str = np.array2string(self.joints, precision=4, suppress_small=True)
-        else:
-            joints_str = f"[{self.joints[0]:.4f}, ..., {self.joints[-1]:.4f}] ({len(self.joints)} joints)"
-
-        return f"JointPose(joints={joints_str})"
-
-    def __eq__(self, other) -> bool:
-        if not isinstance(other, JointPose):
-            return False
-        return np.allclose(self.joints, other.joints) and self.joint_names == other.joint_names
-
-    def __getitem__(self, index) -> float | NDArray[np.float64]:
-        """Allow indexing: joint_pose[0] returns first joint value"""
-        return self.joints[index]
-
-    def __len__(self) -> int:
-        """Allow len(): len(joint_pose) returns number of joints"""
-        return len(self.joints)
-
-
-class EndEffectorPose(Pose):
-    """
-    Represents a single end-effector pose with translation and rotation components.
-
-    This class handles Cartesian space representations of robot end-effector poses,
-    supporting multiple rotation representations (quaternions, Euler angles, rotation
-    vectors, rotation matrices, etc.).
-
-    Examples:
-        # Create with quaternion (wxyz order)
-        pose = EndEffectorPose(
-            translation=[1.0, 2.0, 3.0],
-            rotation=[1.0, 0.0, 0.0, 0.0],
-            rotation_type="quat",
-            rotation_order="wxyz"
-        )
-
-        # Create with Euler angles (degrees by default)
-        pose = EndEffectorPose(
-            translation=[1, 2, 3],
-            rotation=[0, 0, 90],
-            rotation_type="euler",
-            rotation_order="xyz"
-        )
-
-        # Create with Euler angles in radians
-        pose = EndEffectorPose(
-            translation=[1, 2, 3],
-            rotation=[0, 0, np.pi/2],
-            rotation_type="euler",
-            rotation_order="xyz",
-            degrees=False
-        )
-
-        # Create from homogeneous matrix
-        H = np.eye(4)
-        H[:3, 3] = [1, 2, 3]
-        pose = EndEffectorPose(homogeneous=H)
-
-        # Convert between representations
-        quat_wxyz = pose.to_rotation("quat", "wxyz")
-        euler_zyx = pose.to_rotation("euler", "zyx")
-        rot6d = pose.to_rotation("rot6d")
-
-        # Compute relative transformation
-        pose1 = EndEffectorPose(translation=[1, 0, 0], rotation=[1,0,0,0],
-                               rotation_type="quat", rotation_order="wxyz")
-        pose2 = EndEffectorPose(translation=[2, 0, 0], rotation=[1,0,0,0],
-                               rotation_type="quat", rotation_order="wxyz")
-        relative = pose2 - pose1  # Transformation from pose1's frame to pose2's frame
-    """
-
-    pose_type = "end_effector"
-
-    def __init__(
-        self,
-        translation: list | np.ndarray | None = None,
-        rotation: list | np.ndarray | None = None,
-        rotation_type: str | None = None,
-        rotation_order: str | None = None,
-        homogeneous: np.ndarray | None = None,
-        degrees: bool = True,
-    ):
-        """
-        Initialize an end-effector pose.
-
-        Args:
-            translation: Translation vector [x, y, z]
-            rotation: Rotation in specified format
-            rotation_type: Type of rotation ("quat", "euler", "rotvec", "matrix", "rot6d")
-            rotation_order: Order/convention for the rotation type
-            homogeneous: Homogeneous transformation matrix (4, 4)
-                        If provided, overrides translation and rotation
-            degrees: For Euler angles, whether the input is in degrees (default True)
-        """
-        super().__init__()
-
-        # Cache for homogeneous matrix
-        self._homogeneous_cache: NDArray[np.float64] | None = None
-        self._cache_valid = False
-
-        # Handle homogeneous matrix input
-        if homogeneous is not None:
-            self._from_homogeneous(homogeneous)
-            return
-
-        # Store translation
-        self._translation = np.array(translation) if translation is not None else np.zeros(3)
-
-        # Store rotation as scipy Rotation object internally
-        if rotation is not None:
-            if rotation_type is None:
-                raise ValueError("rotation_type must be specified when rotation is provided")
-            self._set_rotation(rotation, rotation_type, rotation_order, degrees)
-        else:
-            self._rotation = Rotation.identity()
-
-    def _from_homogeneous(self, homogeneous: np.ndarray):
-        """Initialize from homogeneous transformation matrix"""
-        homogeneous = np.array(homogeneous)
-
-        # Extract translation (last column, first 3 rows)
-        self._translation = homogeneous[:3, 3]
-
-        # Extract rotation matrix (top-left 3x3)
-        rotation_matrix = homogeneous[:3, :3]
-
-        # Create Rotation object from matrix
-        self._rotation = Rotation.from_matrix(rotation_matrix)
-
-    @staticmethod
-    def _rot6d_to_matrix(rot6d: np.ndarray) -> np.ndarray:
-        """
-        Convert 6D rotation representation to rotation matrix.
-
-        Args:
-            rot6d: 6D rotation as (6,) array - first two rows of rotation matrix flattened
-
-        Returns:
-            Rotation matrix (3, 3)
-        """
-        rot6d = rot6d.reshape(2, 3)
-
-        # First two rows
-        row1 = rot6d[0]
-        row2 = rot6d[1]
-
-        # Normalize first row
-        row1 = row1 / np.linalg.norm(row1)
-
-        # Gram-Schmidt orthogonalization for second row
-        row2 = row2 - np.dot(row1, row2) * row1
-        row2 = row2 / np.linalg.norm(row2)
-
-        # Third row is cross product
-        row3 = np.cross(row1, row2)
-
-        # Construct rotation matrix
-        rotation_matrix = np.vstack([row1, row2, row3])
-
-        return rotation_matrix
-
-    @staticmethod
-    def _matrix_to_rot6d(rotation_matrix: np.ndarray) -> np.ndarray:
-        """
-        Convert rotation matrix to 6D rotation representation.
-
-        Args:
-            rotation_matrix: Rotation matrix (3, 3)
-
-        Returns:
-            6D rotation - (6,) array (first two rows flattened)
-        """
-        return rotation_matrix[:2, :].flatten()
-
-    def _set_rotation(
-        self,
-        rotation: list | np.ndarray,
-        rotation_type: str,
-        rotation_order: str | None = None,
-        degrees: bool = True,
-    ):
-        """Internal method to set rotation from various representations"""
-        rotation = np.array(rotation)
-        rot_type = RotationType(rotation_type.lower())
-
-        if rot_type == RotationType.QUAT:
-            quat_order = QuatOrder(rotation_order.lower()) if rotation_order else QuatOrder.WXYZ
-            if quat_order == QuatOrder.WXYZ:
-                # scipy uses xyzw order, so convert
-                quat_xyzw = np.array([rotation[1], rotation[2], rotation[3], rotation[0]])
-            else:
-                quat_xyzw = rotation
-            self._rotation = Rotation.from_quat(quat_xyzw)
-
-        elif rot_type == RotationType.EULER:
-            euler_order = EulerOrder(rotation_order.lower()) if rotation_order else EulerOrder.XYZ
-            self._rotation = Rotation.from_euler(euler_order.value, rotation, degrees=degrees)
-
-        elif rot_type == RotationType.ROTVEC:
-            self._rotation = Rotation.from_rotvec(rotation)
-
-        elif rot_type == RotationType.MATRIX:
-            self._rotation = Rotation.from_matrix(rotation)
-
-        elif rot_type == RotationType.ROT6D:
-            rotation_matrix = self._rot6d_to_matrix(rotation)
-            self._rotation = Rotation.from_matrix(rotation_matrix)
-
-        else:
-            raise ValueError(f"Unknown rotation type: {rotation_type}")
-
-        # Invalidate cache
-        self._cache_valid = False
-
-    @property
-    def translation(self) -> np.ndarray:
-        """
-        Get translation vector.
-
-        Returns:
-            Translation array - shape (3,)
-        """
-        return self._translation.copy()
-
-    @property
-    def quat_wxyz(self) -> np.ndarray:
-        """Get rotation as quaternion in wxyz order (w, x, y, z)"""
-        return self.to_rotation("quat", "wxyz")
-
-    @property
-    def quat_xyzw(self) -> np.ndarray:
-        """Get rotation as quaternion in xyzw order (x, y, z, w)"""
-        return self.to_rotation("quat", "xyzw")
-
-    @property
-    def euler_xyz(self) -> np.ndarray:
-        """Get rotation as Euler angles in xyz order (degrees)"""
-        return self.to_rotation("euler", "xyz")
-
-    @property
-    def rotvec(self) -> np.ndarray:
-        """Get rotation as rotation vector (axis-angle)"""
-        return self.to_rotation("rotvec")
-
-    @property
-    def rotation_matrix(self) -> np.ndarray:
-        """Get rotation as 3x3 rotation matrix"""
-        return self.to_rotation("matrix")
-
-    @property
-    def rot6d(self) -> np.ndarray:
-        """Get rotation as 6D representation (first two rows of rotation matrix)"""
-        return self.to_rotation("rot6d")
-
-    @property
-    def xyz_rot6d(self) -> np.ndarray:
-        """Get pose as concatenated translation and 6D rotation (9,)"""
-        return np.concatenate([self._translation, self.rot6d])
-
-    @property
-    def xyz_rotvec(self) -> np.ndarray:
-        """Get pose as concatenated translation and rotation vector (6,)"""
-        return np.concatenate([self._translation, self.rotvec])
-
-    @property
-    def homogeneous(self) -> np.ndarray:
-        """
-        Get homogeneous transformation matrix.
-
-        Returns:
-            Homogeneous matrix - shape (4, 4)
-        """
-        if not self._cache_valid:
-            self._homogeneous_cache = self._compute_homogeneous()
-            self._cache_valid = True
-        assert self._homogeneous_cache is not None
-        return self._homogeneous_cache.copy()
-
-    def _compute_homogeneous(self) -> np.ndarray:
-        """Compute homogeneous transformation matrix"""
-        H = np.eye(4)
-        H[:3, :3] = self._rotation.as_matrix()
-        H[:3, 3] = self._translation
-        return H
-
-    def to_rotation(
-        self,
-        rotation_type: str,
-        rotation_order: str | None = None,
-        degrees: bool = True,
-    ) -> np.ndarray:
-        """
-        Get rotation in specified representation.
-
-        Args:
-            rotation_type: Desired type ("quat", "euler", "rotvec", "matrix", "rot6d")
-            rotation_order: Order/convention for the rotation type
-            degrees: For Euler angles, return in degrees (default True)
-
-        Returns:
-            Rotation in requested format
-            - Shape (4,) for quat
-            - Shape (3,) for euler/rotvec
-            - Shape (6,) for rot6d
-            - Shape (3, 3) for matrix
-        """
-        rot_type = RotationType(rotation_type.lower())
-
-        if rot_type == RotationType.ROT6D:
-            rotation_matrix = self._rotation.as_matrix()
-            return self._matrix_to_rot6d(rotation_matrix)
-
-        if rot_type == RotationType.QUAT:
-            quat_order = QuatOrder(rotation_order.lower()) if rotation_order else QuatOrder.WXYZ
-            quat_xyzw = self._rotation.as_quat()
-            if quat_order == QuatOrder.WXYZ:
-                return np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])
-            else:
-                return quat_xyzw
-
-        elif rot_type == RotationType.EULER:
-            euler_order = EulerOrder(rotation_order.lower()) if rotation_order else EulerOrder.XYZ
-            return self._rotation.as_euler(euler_order.value, degrees=degrees)
-
-        elif rot_type == RotationType.ROTVEC:
-            return self._rotation.as_rotvec()
-
-        elif rot_type == RotationType.MATRIX:
-            return self._rotation.as_matrix()
-
-        else:
-            raise ValueError(f"Unknown rotation type: {rotation_type}")
-
-    def to_homogeneous(self) -> np.ndarray:
-        """
-        Convert pose to homogeneous transformation matrix.
-        (Alias for the homogeneous property)
-
-        Returns:
-            Homogeneous matrix - shape (4, 4)
-        """
-        return self.homogeneous
-
-    def set_rotation(
-        self,
-        rotation: list | np.ndarray,
-        rotation_type: str,
-        rotation_order: str | None = None,
-        degrees: bool = True,
-    ):
-        """
-        Set rotation from specified representation.
-
-        Args:
-            rotation: Rotation data
-            rotation_type: Type of rotation ("quat", "euler", "rotvec", "matrix", "rot6d")
-            rotation_order: Order/convention for the rotation type
-            degrees: For Euler angles, whether the input is in degrees (default True)
-        """
-        self._set_rotation(rotation, rotation_type, rotation_order, degrees)
-
-    def _compute_relative(self, other):  # type: ignore[override]
-        """
-        Compute relative transformation from other to self.
-
-        The result represents the transformation needed to go from other's frame to self's frame.
-        Mathematically: T_relative = T_other^{-1} * T_self
-
-        Args:
-            other: Reference end-effector pose
-
-        Returns:
-            EndEffectorPose representing the relative transformation
-        """
-        # Get homogeneous matrices
-        T_self = self.homogeneous
-        T_other = other.homogeneous
-
-        # Compute relative transformation: T_other^{-1} * T_self
-        T_relative = relative_transformation(T_other, T_self)
-
-        # Create new EndEffectorPose from relative transformation
-        return EndEffectorPose(homogeneous=T_relative)
-
-    @classmethod
-    def from_action_format(cls, data: np.ndarray, action_format: ActionFormat) -> "EndEffectorPose":
-        """
-        Create an EndEffectorPose from a flat array using the specified action format.
-
-        This is the inverse of the xyz_rot6d / xyz_rotvec / homogeneous properties.
-
-        Args:
-            data: Flat array whose layout depends on action_format.
-            action_format: One of ActionFormat.XYZ_ROT6D, XYZ_ROTVEC, or DEFAULT.
-
-        Returns:
-            EndEffectorPose instance.
-        """
-        if action_format == ActionFormat.XYZ_ROT6D:
-            return cls(translation=data[:3], rotation=data[3:], rotation_type="rot6d")
-        elif action_format == ActionFormat.XYZ_ROTVEC:
-            return cls(translation=data[:3], rotation=data[3:], rotation_type="rotvec")
-        elif action_format == ActionFormat.DEFAULT:
-            return cls(homogeneous=data.reshape(4, 4))
-        else:
-            raise ValueError(f"Unsupported ActionFormat: {action_format}")
-
-    def copy(self) -> "EndEffectorPose":
-        """
-        Create a deep copy of this end-effector pose.
-
-        Returns:
-            New EndEffectorPose instance with copied data
-        """
-        return EndEffectorPose(
-            translation=self._translation.copy(),
-            rotation=self._rotation.as_quat(),
-            rotation_type="quat",
-            rotation_order="xyzw",
-        )
-
-    def __repr__(self) -> str:
-        quat = self.to_rotation("quat", "wxyz")
-        return f"EndEffectorPose(translation={self.translation}, rotation_quat_wxyz={quat})"
-
-    def __eq__(self, other) -> bool:
-        if not isinstance(other, EndEffectorPose):
-            return False
-        return np.allclose(self._translation, other._translation) and np.allclose(
-            self._rotation.as_quat(), other._rotation.as_quat()
-        )
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py
deleted file mode 100644
index 14e2b0b3a37..00000000000
--- a/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py
+++ /dev/null
@@ -1,643 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Unified state and action processor for robotics."""
-
-from copy import deepcopy
-
-import numpy as np
-
-from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import (
-    ActionFormat,
-    ActionRepresentation,
-    ActionType,
-    ModalityConfig,
-)
-from vllm_omni.diffusion.models.gr00t.dataio.state_action.action_chunking import (
-    EndEffectorActionChunk,
-    JointActionChunk,
-)
-from vllm_omni.diffusion.models.gr00t.dataio.state_action.pose import EndEffectorPose, JointPose
-from vllm_omni.diffusion.models.gr00t.dataio.utils import (
-    apply_sin_cos_encoding,
-    nested_dict_to_numpy,
-    normalize_values_meanstd,
-    normalize_values_minmax,
-    parse_modality_configs,
-    unnormalize_values_meanstd,
-    unnormalize_values_minmax,
-)
-
-
-class StateActionProcessor:
-    """
-    Unified processor for robot state and action data.
-
-    Handles:
-    - State normalization (min/max, mean/std, sin/cos encoding)
-    - Action normalization
-    - Absolute <-> Relative action representation conversion
-    - Action processing with state dependency
-    """
-
-    def __init__(
-        self,
-        modality_configs: dict[str, dict[str, ModalityConfig]],
-        statistics: (dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None) = None,
-        use_percentiles: bool = False,
-        clip_outliers: bool = True,
-        apply_sincos_state_encoding: bool = False,
-        use_relative_action: bool = False,
-    ):
-        """
-        Initialize unified state and action processor.
-
-        Args:
-            modality_configs: Nested dict with structure:
-                {embodiment_tag: {modality: ModalityConfig}}
-                where modality in ["state", "action"]
-                Example: {"gr1": {"state": ModalityConfig(...), "action": ModalityConfig(...)}}
-            statistics: Optional nested dict with structure:
-                {embodiment_tag: {modality: {joint_group: {stat_type: values}}}}
-                where modality in ["state", "action", "relative_action"]
-                and stat_type in ["min", "max", "mean", "std", "q01", "q99"]
-                Example: {"gr1": {"state": {"left_arm": {"min": [...], "max": [...], ...}}}}
-            use_percentiles: Whether to use percentiles (q01/q99) instead of min/max
-            clip_outliers: Whether to clip normalized values to [-1, 1]
-            apply_sincos_state_encoding: Global flag to enable sin/cos encoding for states
-        """
-        self.modality_configs = parse_modality_configs(modality_configs)
-        self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {}
-        self.use_percentiles = use_percentiles
-        self.clip_outliers = clip_outliers
-        self.apply_sincos_state_encoding = apply_sincos_state_encoding
-        self.use_relative_action = use_relative_action
-
-        # Normalization parameters computed from statistics
-        self.norm_params: dict[str, dict[str, dict[str, dict[str, np.ndarray]]]] = {}
-        # Format: norm_params[embodiment_tag][modality][joint_group][stat_type]
-        # where stat_type in ["min", "max", "mean", "std", "dim"]
-
-        if statistics is not None:
-            self.set_statistics(statistics)
-
-    def set_statistics(
-        self,
-        statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]],
-        override: bool = False,
-    ) -> None:
-        """
-        Set dataset statistics for normalization.
-
-        Args:
-            statistics: Nested dict with structure:
-                {embodiment_tag: {modality: {joint_group: {stat_type: values}}}}
-        """
-        for key in statistics:
-            if key not in self.statistics or override:
-                self.statistics[key] = deepcopy(statistics[key])
-            else:
-                print(f"Embodiment tag {key} already in statistics, skipping updating")
-        self._compute_normalization_parameters()
-
-    def _compute_normalization_parameters(self) -> None:
-        """Compute and cache normalization parameters from statistics for all embodiments and modalities."""
-        for embodiment_tag in self.statistics:
-            self.norm_params[embodiment_tag] = {}
-
-            for modality in ["state", "action"]:
-                if modality not in self.statistics[embodiment_tag]:
-                    continue
-
-                self.norm_params[embodiment_tag][modality] = {}
-
-                for joint_group, stats in self.statistics[embodiment_tag][modality].items():
-                    if self.use_percentiles:
-                        min_vals = np.array(stats["q01"])
-                        max_vals = np.array(stats["q99"])
-                    else:
-                        min_vals = np.array(stats["min"])
-                        max_vals = np.array(stats["max"])
-
-                    mean_vals = np.array(stats["mean"])
-                    std_vals = np.array(stats["std"])
-
-                    # Compute range, ensuring it's not zero
-                    range_vals = max_vals - min_vals
-                    range_vals = np.maximum(range_vals, 1e-8)
-
-                    self.norm_params[embodiment_tag][modality][joint_group] = {
-                        "min": min_vals,
-                        "max": max_vals,
-                        "dim": np.array(range_vals.shape[0]),
-                        "mean": mean_vals,
-                        "std": std_vals,
-                    }
-
-            # Override absolute action stats with relative stats where specified
-            if "action" in self.modality_configs[embodiment_tag]:
-                modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys
-                action_configs = self.modality_configs[embodiment_tag]["action"].action_configs
-
-                if action_configs is not None:
-                    for key, action_config in zip(modality_keys, action_configs):
-                        if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action:
-                            if "relative_action" not in self.statistics[embodiment_tag]:
-                                raise ValueError(
-                                    f"Relative action statistics required for embodiment '{embodiment_tag}' "
-                                    f"but 'relative_action' not found in statistics"
-                                )
-                            if key not in self.statistics[embodiment_tag]["relative_action"]:
-                                raise ValueError(
-                                    f"Relative action statistics required for key '{key}' "
-                                    f"in embodiment '{embodiment_tag}' but not found"
-                                )
-                            action_dim = self.norm_params[embodiment_tag]["action"][key]["dim"]
-                            self.norm_params[embodiment_tag]["action"][key] = nested_dict_to_numpy(
-                                self.statistics[embodiment_tag]["relative_action"][key]
-                            )
-                            self.norm_params[embodiment_tag]["action"][key]["dim"] = action_dim
-
-    def apply_state(
-        self,
-        state: dict[str, np.ndarray],
-        embodiment_tag: str,
-    ) -> dict[str, np.ndarray]:
-        """
-        Apply state processing (normalization, encoding).
-
-        Args:
-            state: Dict mapping joint_group -> raw state values
-                Shape per group: (..., D) where D is state dimension
-            embodiment_tag: Embodiment identifier (e.g., "gr1")
-
-        Returns:
-            Dict mapping joint_group -> processed state values
-                - Sin/cos encoded groups: (..., 2*D)
-                - Other groups: (..., D)
-        """
-        normalized_values = {}
-        state = deepcopy(state)  # Avoid modifying input
-
-        # Get sin/cos embedding keys if enabled
-        sin_cos_keys = None
-        if self.apply_sincos_state_encoding:
-            state_config = self.modality_configs[embodiment_tag].get("state")
-            if state_config and hasattr(state_config, "sin_cos_embedding_keys"):
-                sin_cos_keys = state_config.sin_cos_embedding_keys
-
-        for joint_group in self.modality_configs[embodiment_tag]["state"].modality_keys:
-            if joint_group not in state:
-                raise KeyError(f"Joint group '{joint_group}' not found in state dict for embodiment '{embodiment_tag}'")
-
-            # Strategy 1: Sin/cos encoding (doubles dimension)
-            if sin_cos_keys and joint_group in sin_cos_keys:
-                normalized_values[joint_group] = apply_sin_cos_encoding(state[joint_group])
-
-            # Strategy 2: Mean/std normalization
-            elif (
-                hasattr(
-                    self.modality_configs[embodiment_tag]["state"],
-                    "mean_std_embedding_keys",
-                )
-                and self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys
-                and joint_group in self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys
-            ):
-                params = self.norm_params[embodiment_tag]["state"][joint_group]
-                normalized = normalize_values_meanstd(state[joint_group], params)
-                normalized_values[joint_group] = normalized
-
-            # Strategy 3: Min/max normalization to [-1, 1]
-            else:
-                params = self.norm_params[embodiment_tag]["state"][joint_group]
-                normalized = normalize_values_minmax(state[joint_group], params)
-
-                if self.clip_outliers:
-                    normalized = np.clip(normalized, -1.0, 1.0)
-
-                normalized_values[joint_group] = normalized
-
-        return normalized_values
-
-    def unapply_state(
-        self,
-        state: dict[str, np.ndarray],
-        embodiment_tag: str,
-    ) -> dict[str, np.ndarray]:
-        """
-        Reverse state processing (denormalization).
-
-        Args:
-            state: Dict mapping joint_group -> processed state values
-            embodiment_tag: Embodiment identifier
-
-        Returns:
-            Dict mapping joint_group -> raw state values
-
-        Raises:
-            ValueError: If attempting to reverse sin/cos encoding (not reversible)
-        """
-        unnormalized_values = {}
-
-        # Get sin/cos embedding keys if enabled
-        sin_cos_keys = None
-        if self.apply_sincos_state_encoding:
-            state_config = self.modality_configs[embodiment_tag].get("state")
-            if state_config and hasattr(state_config, "sin_cos_embedding_keys"):
-                sin_cos_keys = state_config.sin_cos_embedding_keys
-
-        for joint_group in self.modality_configs[embodiment_tag]["state"].modality_keys:
-            if joint_group not in state:
-                raise KeyError(f"Joint group '{joint_group}' not found in state dict for embodiment '{embodiment_tag}'")
-
-            # Sin/cos encoding is not reversible
-            if sin_cos_keys and joint_group in sin_cos_keys:
-                raise ValueError(
-                    f"Cannot unapply sin/cos encoding for joint group '{joint_group}' "
-                    f"in embodiment '{embodiment_tag}'. This transformation is not reversible."
-                )
-
-            # Reverse mean/std normalization
-            elif (
-                hasattr(
-                    self.modality_configs[embodiment_tag]["state"],
-                    "mean_std_embedding_keys",
-                )
-                and self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys
-                and joint_group in self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys
-            ):
-                params = self.norm_params[embodiment_tag]["state"][joint_group]
-                unnormalized = unnormalize_values_meanstd(state[joint_group], params)
-                unnormalized_values[joint_group] = unnormalized
-
-            # Reverse min/max normalization
-            else:
-                params = self.norm_params[embodiment_tag]["state"][joint_group]
-                unnormalized_values[joint_group] = unnormalize_values_minmax(state[joint_group], params)
-
-        return unnormalized_values
-
-    def apply_action(
-        self,
-        action: dict[str, np.ndarray],
-        embodiment_tag: str,
-        state: dict[str, np.ndarray] | None = None,
-    ) -> dict[str, np.ndarray]:
-        """
-        Apply action processing (absolute->relative conversion, normalization).
-
-        Processing order:
-        1. Convert absolute actions to relative (if configured)
-        2. Normalize actions
-
-        Args:
-            action: Dict mapping joint_group -> raw action values
-                Shape per group: (T, D) where T is action horizon, D is action dimension
-            embodiment_tag: Embodiment identifier
-            state: Optional dict mapping joint_group -> raw state values
-                Required if any action group uses ActionRepresentation.RELATIVE
-                Shape per group: (T_state, D) where last timestep is used as reference
-
-        Returns:
-            Dict mapping joint_group -> processed action values
-                Shape per group: (T, D)
-
-        Raises:
-            ValueError: If state is None but required for relative action conversion
-        """
-        action = deepcopy(action)  # Avoid modifying input
-
-        # Step 1: Convert absolute actions to relative (if needed)
-        modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys
-        action_configs = self.modality_configs[embodiment_tag]["action"].action_configs
-
-        if action_configs is not None:
-            for key, action_config in zip(modality_keys, action_configs):
-                if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action:
-                    if state is None:
-                        raise ValueError(
-                            f"State dict required for relative action processing of key '{key}' "
-                            f"in embodiment '{embodiment_tag}'"
-                        )
-
-                    # Determine which state key to use as reference
-                    state_key = action_config.state_key if action_config.state_key else key
-
-                    if state_key not in state:
-                        raise KeyError(
-                            f"Reference state key '{state_key}' not found in state dict "
-                            f"for embodiment '{embodiment_tag}'"
-                        )
-
-                    # Use last state as reference frame
-                    reference_state = state[state_key][-1]
-
-                    # Convert absolute to relative
-                    action[key] = self._convert_to_relative_action(
-                        action=action[key],
-                        reference_state=reference_state,
-                        action_type=action_config.type,
-                        action_format=action_config.format,
-                    )
-
-        # Step 2: Normalize actions
-        normalized_values = {}
-        for joint_group in modality_keys:
-            if joint_group not in action:
-                raise KeyError(
-                    f"Joint group '{joint_group}' not found in action dict for embodiment '{embodiment_tag}'"
-                )
-
-            params = self.norm_params[embodiment_tag]["action"][joint_group]
-            if (
-                self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys is not None
-                and joint_group in self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys
-            ):
-                normalized = normalize_values_meanstd(action[joint_group], params)
-            else:
-                normalized = normalize_values_minmax(action[joint_group], params)
-
-            if self.clip_outliers:
-                normalized = np.clip(normalized, -1.0, 1.0)
-
-            normalized_values[joint_group] = normalized
-
-        return normalized_values
-
-    def unapply_action(
-        self,
-        action: dict[str, np.ndarray],
-        embodiment_tag: str,
-        state: dict[str, np.ndarray] | None = None,
-    ) -> dict[str, np.ndarray]:
-        """
-        Reverse action processing (denormalization, relative->absolute conversion).
-
-        Processing order:
-        1. Denormalize actions
-        2. Convert relative actions to absolute (if configured)
-
-        Args:
-            action: Dict mapping joint_group -> processed action values
-                Shape per group: (T, D) or (B, T, D) for batched
-            embodiment_tag: Embodiment identifier
-            state: Optional dict mapping joint_group -> raw state values
-                Required if any action group uses ActionRepresentation.RELATIVE
-                Shape per group: (T_state, D) or (B, T_state, D) for batched
-
-        Returns:
-            Dict mapping joint_group -> raw absolute action values
-                Shape per group: (T, D) or (B, T, D) for batched
-
-        Raises:
-            ValueError: If state is None but required for relative->absolute conversion
-        """
-        # Step 1: Unnormalize actions
-        unnormalized_values = {}
-        modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys
-
-        for joint_group in modality_keys:
-            if joint_group not in action:
-                raise KeyError(
-                    f"Joint group '{joint_group}' not found in action dict for embodiment '{embodiment_tag}'"
-                )
-
-            params = self.norm_params[embodiment_tag]["action"][joint_group]
-            group_values = action[joint_group]
-
-            if (
-                self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys is not None
-                and joint_group in self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys
-            ):
-                unnormalized = unnormalize_values_meanstd(group_values, params)
-            else:
-                unnormalized = unnormalize_values_minmax(group_values, params)
-
-            unnormalized_values[joint_group] = unnormalized
-
-        # Step 2: Convert relative actions to absolute (if needed)
-        action_configs = self.modality_configs[embodiment_tag]["action"].action_configs
-
-        if action_configs is not None:
-            for key, action_config in zip(modality_keys, action_configs):
-                if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action:
-                    if state is None:
-                        raise ValueError(
-                            f"State dict required for relative->absolute conversion of key '{key}' "
-                            f"in embodiment '{embodiment_tag}'"
-                        )
-
-                    # Determine which state key to use as reference
-                    state_key = action_config.state_key if action_config.state_key else key
-
-                    if state_key not in state:
-                        raise KeyError(
-                            f"Reference state key '{state_key}' not found in state dict "
-                            f"for embodiment '{embodiment_tag}'"
-                        )
-
-                    relative_action = unnormalized_values[key]
-
-                    # Handle batched and unbatched cases
-                    is_batched = relative_action.ndim == 3
-                    if not is_batched:
-                        assert relative_action.ndim == 2
-                        reference_state = state[state_key]
-                        if reference_state.ndim == 2:
-                            reference_state = reference_state[None, :]
-                        relative_action = relative_action[None, :]
-                    else:
-                        reference_state = state[state_key]
-                        if reference_state.ndim == 2:
-                            reference_state = reference_state[None, :]
-
-                    # Convert batched relative actions to absolute
-                    absolute_actions = []
-                    for s, a in zip(reference_state, relative_action):
-                        # Use last timestep of state as reference
-                        absolute_action = self._convert_to_absolute_action(
-                            action=a,
-                            reference_state=s[-1],
-                            action_type=action_config.type,
-                            action_format=action_config.format,
-                        )
-                        absolute_actions.append(absolute_action)
-
-                    if is_batched:
-                        unnormalized_values[key] = np.stack(absolute_actions, axis=0)
-                    else:
-                        unnormalized_values[key] = absolute_actions[0]
-
-        return unnormalized_values
-
-    def apply(
-        self,
-        state: dict[str, np.ndarray],
-        action: dict[str, np.ndarray],
-        embodiment_tag: str,
-    ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
-        """
-        Apply both state and action processing together.
-
-        Convenience method that processes state and action in one call,
-        automatically passing raw state to action processor for relative conversion.
-
-        Args:
-            state: Dict mapping joint_group -> raw state values
-            action: Dict mapping joint_group -> raw action values
-            embodiment_tag: Embodiment identifier
-
-        Returns:
-            Tuple of (processed_state, processed_action)
-        """
-        processed_state = self.apply_state(state, embodiment_tag)
-        if action:
-            processed_action = self.apply_action(action, embodiment_tag, state=state)
-        else:
-            processed_action = {}
-        return processed_state, processed_action
-
-    def unapply(
-        self,
-        state: dict[str, np.ndarray],
-        action: dict[str, np.ndarray],
-        embodiment_tag: str,
-        raw_state: dict[str, np.ndarray] | None = None,
-    ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
-        """
-        Reverse both state and action processing together.
-
-        Args:
-            state: Dict mapping joint_group -> processed state values
-            action: Dict mapping joint_group -> processed action values
-            embodiment_tag: Embodiment identifier
-            raw_state: Optional dict of raw states for relative->absolute conversion
-                If None, will use unapplied state (but won't work for sin/cos encoded states)
-
-        Returns:
-            Tuple of (raw_state, raw_action)
-        """
-        # Unapply state first
-        try:
-            unapplied_state = self.unapply_state(state, embodiment_tag)
-        except ValueError as e:
-            if "sin/cos encoding" in str(e) and raw_state is None:
-                raise ValueError("Cannot unapply sin/cos encoded state. Please provide raw_state parameter.") from e
-            raise
-
-        # Use provided raw_state if available, otherwise use unapplied state
-        state_for_action = raw_state if raw_state is not None else unapplied_state
-
-        # Unapply action
-        unapplied_action = self.unapply_action(action, embodiment_tag, state=state_for_action)
-
-        return unapplied_state, unapplied_action
-
-    def get_state_dim(self, embodiment_tag: str, include_sincos_expansion: bool = False) -> int:
-        """
-        Get total state dimension after processing.
-
-        Args:
-            embodiment_tag: Embodiment identifier
-            include_sincos_expansion: If True, accounts for sin/cos encoding doubling dimensions
-
-        Returns:
-            Total state dimension across all joint groups
-        """
-        total_dim = 0
-        state_config = self.modality_configs[embodiment_tag]["state"]
-
-        # Get sin/cos embedding keys if enabled
-        sin_cos_keys = set()
-        if self.apply_sincos_state_encoding and hasattr(state_config, "sin_cos_embedding_keys"):
-            sin_cos_keys = set(state_config.sin_cos_embedding_keys)
-
-        for joint_group in state_config.modality_keys:
-            base_dim = self.norm_params[embodiment_tag]["state"][joint_group]["dim"].item()
-
-            # Sin/cos encoding doubles the dimension
-            if include_sincos_expansion and joint_group in sin_cos_keys:
-                total_dim += base_dim * 2
-            else:
-                total_dim += base_dim
-
-        return total_dim
-
-    def get_action_dim(self, embodiment_tag: str) -> int:
-        """
-        Get total action dimension.
-
-        Args:
-            embodiment_tag: Embodiment identifier
-
-        Returns:
-            Total action dimension across all joint groups
-        """
-        total_dim = 0
-        for joint_group in self.modality_configs[embodiment_tag]["action"].modality_keys:
-            total_dim += self.norm_params[embodiment_tag]["action"][joint_group]["dim"].item()
-        return total_dim
-
-    def _convert_to_relative_action(
-        self,
-        action: np.ndarray,
-        reference_state: np.ndarray,
-        action_type: ActionType,
-        action_format: ActionFormat,
-    ) -> np.ndarray:
-        """Convert absolute action to relative action using reference state."""
-        assert action.ndim == 2, f"Expected action shape (T, D), got {action.shape}"
-        assert reference_state.ndim == 1, f"Expected state shape (D,), got {reference_state.shape}"
-
-        if action_type == ActionType.EEF:
-            action_chunking = EndEffectorActionChunk.from_array(action, action_format)
-            reference_frame = EndEffectorPose.from_action_format(reference_state, action_format)
-
-        elif action_type == ActionType.NON_EEF:
-            action_chunking = JointActionChunk([JointPose(m) for m in action])
-            reference_frame = JointPose(reference_state)
-
-        else:
-            raise ValueError(f"Unknown ActionType: {action_type}")
-
-        relative_action_chunking = action_chunking.relative_chunking(reference_frame=reference_frame)
-        return relative_action_chunking.to(action_format)
-
-    def _convert_to_absolute_action(
-        self,
-        action: np.ndarray,
-        reference_state: np.ndarray,
-        action_type: ActionType,
-        action_format: ActionFormat,
-    ) -> np.ndarray:
-        """Convert relative action to absolute action using reference state."""
-        assert action.ndim == 2, f"Expected action shape (T, D), got {action.shape}"
-        assert reference_state.ndim == 1, f"Expected state shape (D,), got {reference_state.shape}"
-        assert reference_state.shape[0] == action.shape[1], (
-            f"State dim {reference_state.shape[0]} != action dim {action.shape[1]}"
-        )
-
-        if action_type == ActionType.EEF:
-            rel_action = EndEffectorActionChunk.from_array(action, action_format)
-            reference_frame = EndEffectorPose.from_action_format(reference_state, action_format)
-
-        elif action_type == ActionType.NON_EEF:
-            rel_action = JointActionChunk([JointPose(pose) for pose in action])
-            reference_frame = JointPose(reference_state)
-
-        else:
-            raise ValueError(f"Unknown ActionType: {action_type}")
-
-        abs_action = rel_action.to_absolute_chunking(reference_frame=reference_frame)
-        return abs_action.to(action_format)
-
-    def __str__(self) -> str:
-        return (
-            "StateActionProcessor("
-            f"modality_configs={self.modality_configs}, "
-            f"statistics={self.statistics}, "
-            f"use_percentiles={self.use_percentiles}, "
-            f"clip_outliers={self.clip_outliers}, "
-            f"apply_sincos_state_encoding={self.apply_sincos_state_encoding}, "
-            f"use_relative_action={self.use_relative_action})"
-        )
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/types.py b/vllm_omni/diffusion/models/gr00t/dataio/types.py
deleted file mode 100644
index d4948a1c234..00000000000
--- a/vllm_omni/diffusion/models/gr00t/dataio/types.py
+++ /dev/null
@@ -1,112 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from dataclasses import dataclass, field
-from enum import Enum
-from typing import Any
-
-import numpy as np
-
-from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag
-
-
-class MessageType(Enum):
-    START_OF_EPISODE = "start_of_episode"
-    END_OF_EPISODE = "end_of_episode"
-    EPISODE_STEP = "episode_step"
-    IMAGE = "image"
-    TEXT = "text"
-
-
-class ActionRepresentation(Enum):
-    RELATIVE = "relative"
-    DELTA = "delta"
-    ABSOLUTE = "absolute"
-
-
-class ActionType(Enum):
-    EEF = "eef"
-    NON_EEF = "non_eef"
-
-
-class ActionFormat(Enum):
-    DEFAULT = "default"
-    XYZ_ROT6D = "xyz+rot6d"
-    XYZ_ROTVEC = "xyz+rotvec"
-
-
-@dataclass
-class VLAStepData:
-    """
-    Represents a single step of VLA (Vision-Language-Action) data.
-
-    This is the core data structure returned by datasets, containing raw observation
-    and action data that will be processed by the SequenceVLAProcessor.
-    """
-
-    # Core data
-    images: dict[str, list[np.ndarray]]  # view_name -> list[np.ndarray] (for temporal stacking)
-    states: dict[str, np.ndarray]  # state_name -> np.ndarray (dim,) for single step or (horizon, dim) for trajectory
-    actions: dict[str, np.ndarray]  # action_name -> np.ndarray (horizon, dim) for action chunk
-    masks: dict[str, list[np.ndarray]] | None = None  # view_name -> list[np.ndarray] (H, W)
-    text: str | None = None  # Optional task description or instruction
-    embodiment: EmbodimentTag = EmbodimentTag.NEW_EMBODIMENT  # Optional embodiment tag for cross-embodiment training
-    is_demonstration: bool = (
-        False  # Whether the step is a demonstration. If True, no loss should be computed for this step.
-    )
-
-    # Flexible metadata that can be extended by users
-    metadata: dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class ActionConfig:
-    rep: ActionRepresentation
-    type: ActionType
-    format: ActionFormat
-    state_key: str | None = None
-
-
-@dataclass
-class ModalityConfig:
-    """Configuration for a modality defining how data should be sampled and loaded.
-
-    This class specifies which indices to sample relative to a base index and which
-    keys to load for a particular modality (e.g., video, state, action).
-    """
-
-    delta_indices: list[int]
-    """Delta indices to sample relative to the current index. The returned data will
-    correspond to the original data at a sampled base index + delta indices."""
-    modality_keys: list[str]
-    """The keys to load for the modality in the dataset."""
-    sin_cos_embedding_keys: list[str] | None = None
-    """Optional list of keys to apply sin/cos encoding. If None or empty, use
-    min/max normalization for all keys."""
-    mean_std_embedding_keys: list[str] | None = None
-    """Optional list of keys to apply mean/std normalization. If None or empty,
-    use min/max normalization for all keys."""
-    action_configs: list[ActionConfig] | None = None
-
-    def __post_init__(self):
-        """Validate fields and set default values."""
-        if self.delta_indices is None or not isinstance(self.delta_indices, list):
-            raise ValueError(f"delta_indices must be a non-None list, got {self.delta_indices!r}")
-        if self.modality_keys is None or not isinstance(self.modality_keys, list) or len(self.modality_keys) == 0:
-            raise ValueError(f"modality_keys must be a non-empty list, got {self.modality_keys!r}")
-        if self.action_configs is not None:
-            assert len(self.action_configs) == len(self.modality_keys), (
-                f"Number of action configs ({len(self.action_configs)}) must match "
-                f"number of modality keys ({len(self.modality_keys)})"
-            )
-            parsed_action_configs = []
-            for action_config in self.action_configs:
-                if isinstance(action_config, dict):
-                    action_config = ActionConfig(
-                        rep=ActionRepresentation[action_config["rep"]],
-                        type=ActionType[action_config["type"]],
-                        format=ActionFormat[action_config["format"]],
-                        state_key=action_config.get("state_key", None),
-                    )
-                parsed_action_configs.append(action_config)
-            self.action_configs = parsed_action_configs
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/utils.py b/vllm_omni/diffusion/models/gr00t/dataio/utils.py
deleted file mode 100644
index d3846d44228..00000000000
--- a/vllm_omni/diffusion/models/gr00t/dataio/utils.py
+++ /dev/null
@@ -1,187 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from dataclasses import asdict, is_dataclass
-from enum import Enum
-from typing import Any
-
-import numpy as np
-
-from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig
-
-
-def apply_sin_cos_encoding(values: np.ndarray) -> np.ndarray:
-    """Apply sin/cos encoding to values.
-
-    Args:
-        values: Array of shape (..., D) containing values to encode
-
-    Returns:
-        Array of shape (..., 2*D) with [sin, cos] concatenated
-
-    Note: This DOUBLES the dimension. For example:
-        Input:  [v₁, v₂, v₃] with shape (..., 3)
-        Output: [sin(v₁), sin(v₂), sin(v₃), cos(v₁), cos(v₂), cos(v₃)] with shape (..., 6)
-    """
-    sin_values = np.sin(values)
-    cos_values = np.cos(values)
-    # Concatenate sin and cos: [sin(v1), sin(v2), ..., cos(v1), cos(v2), ...]
-    return np.concatenate([sin_values, cos_values], axis=-1)
-
-
-def nested_dict_to_numpy(data: Any) -> Any:
-    """Recursively convert bottom-level list-of-lists to NumPy arrays.
-
-    Args:
-        data: A nested dictionary whose leaves are lists (or lists of lists).
-
-    Returns:
-        The same dictionary structure with leaf lists converted to ``np.ndarray``.
-    """
-    if isinstance(data, dict):
-        return {key: nested_dict_to_numpy(value) for key, value in data.items()}
-    elif isinstance(data, list):
-        # Convert lists to numpy arrays
-        # NumPy will handle both 1D and 2D cases appropriately
-        return np.array(data)
-    else:
-        return data
-
-
-def normalize_values_minmax(values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray:
-    """Min-max normalize ``values`` to ``[-1, 1]`` using ``params['min']`` and ``params['max']``.
-
-    Accepts 2D ``(T, D)`` or 3D ``(B, T, D)`` arrays. Features with ``min == max``
-    are emitted as 0.
-    """
-    min_vals = params["min"]
-    max_vals = params["max"]
-    normalized = np.zeros_like(values)
-
-    mask = ~np.isclose(max_vals, min_vals)
-
-    normalized[..., mask] = (values[..., mask] - min_vals[..., mask]) / (max_vals[..., mask] - min_vals[..., mask])
-    normalized[..., mask] = 2 * normalized[..., mask] - 1
-
-    return normalized
-
-
-def unnormalize_values_minmax(normalized_values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray:
-    """Inverse of :func:`normalize_values_minmax`.
-
-    Input values are clipped to ``[-1, 1]`` before being mapped back to
-    ``[params['min'], params['max']]``.
-    """
-
-    min_vals = params["min"]
-    max_vals = params["max"]
-    range_vals = max_vals - min_vals
-
-    # Unnormalize from [-1, 1]
-    unnormalized = (np.clip(normalized_values, -1.0, 1.0) + 1.0) / 2.0 * range_vals + min_vals
-    return unnormalized
-
-
-def normalize_values_meanstd(values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray:
-    """Z-score normalize ``values`` using ``params['mean']`` and ``params['std']``.
-
-    Features whose ``std == 0`` are passed through unchanged.
-    """
-    mean_vals = params["mean"]
-    std_vals = params["std"]
-
-    # Create mask for non-zero standard deviations
-    mask = std_vals != 0
-
-    # Initialize normalized array
-    normalized = np.zeros_like(values)
-
-    # Normalize only features with non-zero std
-    normalized[..., mask] = (values[..., mask] - mean_vals[..., mask]) / std_vals[..., mask]
-
-    # Keep original values for zero-std features
-    normalized[..., ~mask] = values[..., ~mask]
-
-    return normalized
-
-
-def unnormalize_values_meanstd(normalized_values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray:
-    """Inverse of :func:`normalize_values_meanstd` (``x * std + mean``).
-
-    Features whose ``std == 0`` are passed through unchanged.
-    """
-    mean_vals = params["mean"]
-    std_vals = params["std"]
-
-    # Create mask for non-zero standard deviations
-    mask = std_vals != 0
-
-    # Initialize unnormalized array
-    unnormalized = np.zeros_like(normalized_values)
-
-    # Unnormalize only features with non-zero std
-    unnormalized[..., mask] = normalized_values[..., mask] * std_vals[..., mask] + mean_vals[..., mask]
-
-    # Keep normalized values for zero-std features
-    unnormalized[..., ~mask] = normalized_values[..., ~mask]
-
-    return unnormalized
-
-
-def to_json_serializable(obj: Any) -> Any:
-    """
-    Recursively convert dataclasses and numpy arrays to JSON-serializable format.
-
-    Args:
-        obj: Object to convert (can be dataclass, numpy array, dict, list, etc.)
-
-    Returns:
-        JSON-serializable representation of the object
-    """
-    if is_dataclass(obj) and not isinstance(obj, type):
-        # Convert dataclass to dict, then recursively process the dict
-        return to_json_serializable(asdict(obj))
-    elif isinstance(obj, np.ndarray):
-        # Convert numpy array to list
-        return obj.tolist()
-    elif isinstance(obj, np.integer):
-        # Convert numpy integers to Python int
-        return int(obj)
-    elif isinstance(obj, np.floating):
-        # Convert numpy floats to Python float
-        return float(obj)
-    elif isinstance(obj, np.bool_):
-        # Convert numpy bool to Python bool
-        return bool(obj)
-    elif isinstance(obj, dict):
-        # Recursively process dictionary values
-        return {key: to_json_serializable(value) for key, value in obj.items()}
-    elif isinstance(obj, (list, tuple)):
-        # Recursively process list/tuple elements
-        return [to_json_serializable(item) for item in obj]
-    elif isinstance(obj, set):
-        # Convert set to list
-        return [to_json_serializable(item) for item in obj]
-    elif isinstance(obj, (str, int, float, bool, type(None))):
-        # Already JSON-serializable
-        return obj
-    elif isinstance(obj, Enum):
-        return obj.name
-    else:
-        # For other types, try to convert to string as fallback
-        # You might want to handle specific types differently
-        return str(obj)
-
-
-def parse_modality_configs(
-    modality_configs: dict[str, dict[str, ModalityConfig]],
-) -> dict[str, dict[str, ModalityConfig]]:
-    parsed_modality_configs = {}
-    for embodiment_tag, modality_config in modality_configs.items():
-        parsed_modality_configs[embodiment_tag] = {}
-        for modality, config in modality_config.items():
-            if isinstance(config, dict):
-                parsed_modality_configs[embodiment_tag][modality] = ModalityConfig(**config)
-            else:
-                parsed_modality_configs[embodiment_tag][modality] = config
-    return parsed_modality_configs
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/__init__.py b/vllm_omni/diffusion/models/gr00t/modeling/__init__.py
deleted file mode 100644
index 7aecb405ff9..00000000000
--- a/vllm_omni/diffusion/models/gr00t/modeling/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7
-from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7Processor
-
-__all__ = ["Gr00tN1d7", "Gr00tN1d7Processor"]
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py
deleted file mode 100644
index 3fcbce203f5..00000000000
--- a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py
+++ /dev/null
@@ -1,491 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import logging
-import os
-from typing import Any
-
-import torch
-from torch import nn
-from transformers import AutoConfig, AutoModel, PreTrainedModel
-from transformers.feature_extraction_utils import BatchFeature
-
-from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
-from vllm_omni.diffusion.models.gr00t.modeling.modules.dit import AlternateVLDiT, DiT, SelfAttentionTransformer
-from vllm_omni.diffusion.models.gr00t.modeling.modules.embodiment_conditioned_mlp import (
-    CategorySpecificMLP,
-    MultiEmbodimentActionEncoder,
-)
-from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7DataCollator
-
-logger = logging.getLogger(__name__)
-
-
-class Gr00tN1d7ActionHead(nn.Module):
-    """Action head component for flow matching diffusion policy."""
-
-    supports_gradient_checkpointing = True
-
-    def __init__(self, config: Gr00tN1d7Config):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.input_embedding_dim = config.input_embedding_dim
-
-        if config.use_alternate_vl_dit:
-            self.model = AlternateVLDiT(
-                **config.diffusion_model_cfg,
-                cross_attention_dim=config.backbone_embedding_dim,
-                attend_text_every_n_blocks=config.attend_text_every_n_blocks,
-            )
-            logger.info("Using AlternateVLDiT for diffusion model")
-        else:
-            self.model = DiT(
-                **config.diffusion_model_cfg,
-                cross_attention_dim=config.backbone_embedding_dim,
-            )
-            logger.info("Using DiT for diffusion model")
-        self.action_dim = config.max_action_dim
-        self.action_horizon = config.action_horizon
-        self.num_inference_timesteps = config.num_inference_timesteps
-
-        self.state_encoder = CategorySpecificMLP(
-            num_categories=config.max_num_embodiments,
-            input_dim=config.max_state_dim * config.state_history_length,
-            hidden_dim=self.hidden_size,
-            output_dim=self.input_embedding_dim,
-        )
-        self.action_encoder = MultiEmbodimentActionEncoder(
-            action_dim=self.action_dim,
-            hidden_size=self.input_embedding_dim,
-            num_embodiments=config.max_num_embodiments,
-        )
-        self.action_decoder = CategorySpecificMLP(
-            num_categories=config.max_num_embodiments,
-            input_dim=self.hidden_size,
-            hidden_dim=self.hidden_size,
-            output_dim=self.action_dim,
-        )
-
-        self.vlln = nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity()
-
-        vl_self_attention_cfg = config.vl_self_attention_cfg
-        if vl_self_attention_cfg and vl_self_attention_cfg.get("num_layers", 0) > 0:
-            self.vl_self_attention = SelfAttentionTransformer(**vl_self_attention_cfg)
-        else:
-            self.vl_self_attention = nn.Identity()
-
-        if config.add_pos_embed:
-            self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim)
-            nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)
-
-        self.num_timestep_buckets = config.num_timestep_buckets
-
-    def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature:
-        backbone_features = backbone_output["backbone_features"]
-        backbone_features = self.vlln(backbone_features)
-        backbone_features = self.vl_self_attention(backbone_features)
-        backbone_output["backbone_features"] = backbone_features
-        return backbone_output
-
-    def _encode_features(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature:
-        """
-        Encode features for the action head.
-
-        Args:
-            backbone_output: Output from the backbone model containing:
-                - backbone_features: [B, seq_len, backbone_embedding_dim]
-                - backbone_attention_mask: [B, seq_len]
-            action_input: Input containing:
-                - state: [B, state_history_length, max_state_dim]
-                - embodiment_id: [B] (embodiment IDs)
-
-        Returns:
-            BatchFeature containing:
-                - backbone_features: [B, seq_len, backbone_embedding_dim]
-                - state_features: [B, 1, input_embedding_dim]
-        """
-        backbone_output = self.process_backbone_output(backbone_output)
-
-        vl_embeds = backbone_output.backbone_features
-        embodiment_id = action_input.embodiment_id
-
-        state = action_input.state
-        current_T = state.shape[1]
-        assert current_T == self.config.state_history_length, "current_T != state_history_length"
-        # [B, state_history_length, max_state_dim] -> [B, 1, state_history_length * max_state_dim]
-        state = state.view(state.shape[0], 1, -1)
-
-        state_features = self.state_encoder(state, embodiment_id)
-
-        return BatchFeature(data={"backbone_features": vl_embeds, "state_features": state_features})
-
-    @torch.no_grad()
-    def get_action_with_features(
-        self,
-        backbone_features: torch.Tensor,
-        state_features: torch.Tensor,
-        embodiment_id: torch.Tensor,
-        backbone_output: BatchFeature,
-        action_input: BatchFeature,
-        options: dict[str, Any] | None = None,
-    ) -> BatchFeature:
-        """
-        Generate actions using the flow matching diffusion process.
-
-        Args:
-            backbone_features: [B, seq_len, backbone_embedding_dim]
-            state_features: [B, state_horizon, input_embedding_dim]
-            embodiment_id: [B] (embodiment IDs)
-            backbone_output: Output from the backbone model
-        """
-        vl_embeds = backbone_features
-
-        batch_size = vl_embeds.shape[0]
-        device = vl_embeds.device
-        _seed_env = os.environ.get("GR00T_NOISE_SEED")
-        if _seed_env is not None:
-            _gen = torch.Generator(device=device).manual_seed(int(_seed_env))
-            actions = torch.randn(
-                size=(batch_size, self.config.action_horizon, self.action_dim),
-                dtype=vl_embeds.dtype,
-                device=device,
-                generator=_gen,
-            )
-        else:
-            actions = torch.randn(
-                size=(batch_size, self.config.action_horizon, self.action_dim),
-                dtype=vl_embeds.dtype,
-                device=device,
-            )
-
-        dt = 1.0 / self.num_inference_timesteps
-        vel_strength = torch.ones_like(actions)
-
-        if "action" in action_input:
-            # If action in input when doing get action, it means we want to use RTC.
-            # action_horizon is the action horizon of the input action.
-            # rtc_overlap_steps is the number of steps to overlap with the previous action chunks.
-            # rtc_frozen_steps is the policy inference latency expressed as frozen action steps.
-            # rtc_ramp_rate is the rate of the ramp of denoising the actions.
-            assert options is not None, "options is not None"
-            assert "action_horizon" in options, "action_horizon is not in options"
-            assert "rtc_overlap_steps" in options, "rtc_overlap_steps is not in options"
-            assert "rtc_frozen_steps" in options, "rtc_frozen_steps is not in options"
-            assert "rtc_ramp_rate" in options, "rtc_ramp_rate is not in options"
-
-            action_horizon_before_padding = options["action_horizon"]
-
-            # Use previous action instead of pure noise to do inpainting
-            actions[:, : options["rtc_overlap_steps"], :] = action_input["action"][
-                :,
-                action_horizon_before_padding - options["rtc_overlap_steps"] : action_horizon_before_padding,
-                :,
-            ]
-            vel_strength[:, : options["rtc_frozen_steps"], :] = 0.0
-            # NOTE: use an exponential ramp strength to set the remaining unfrozen rtc_steps
-            intermediate_steps = options["rtc_overlap_steps"] - options["rtc_frozen_steps"]
-            # Create exponential ramp from 0 to 1 over intermediate steps
-            t = torch.linspace(0.0, 1.0, intermediate_steps + 2, device=device)
-            ramp = 1 - torch.exp(-options["rtc_ramp_rate"] * t)
-            ramp = ramp / ramp[-1].clamp_min(1e-8)  # normalize to [0,1]
-            ramp = ramp[1:-1]  # we will only take the middle part of the ramp, ignore the 0.0 and 1.0
-            # Apply ramp to the intermediate steps [batch, intermediate_steps, action_dim]
-            vel_strength[
-                :,
-                options["rtc_frozen_steps"] : options["rtc_overlap_steps"],
-                :,
-            ] = ramp[None, :, None].to(device)
-
-        # Run denoising steps.
-        for t in range(self.num_inference_timesteps):
-            t_discretized = t * self.num_timestep_buckets // self.num_inference_timesteps
-
-            # Embed noised action trajectory.
-            timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device)
-            action_features = self.action_encoder(actions, timesteps_tensor, embodiment_id)
-            # Add position embedding.
-            if self.config.add_pos_embed:
-                pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
-                pos_embs = self.position_embedding(pos_ids).unsqueeze(0)
-                action_features = action_features + pos_embs
-
-            # Join vision, language, state and action embedding along sequence dimension.
-            sa_embs = torch.cat((state_features, action_features), dim=1)
-
-            # Run model forward.
-            if self.config.use_alternate_vl_dit:
-                model_output = self.model(
-                    hidden_states=sa_embs,
-                    encoder_hidden_states=vl_embeds,
-                    timestep=timesteps_tensor,
-                    image_mask=backbone_output.image_mask,
-                    backbone_attention_mask=backbone_output.backbone_attention_mask,
-                )
-            else:
-                model_output = self.model(
-                    hidden_states=sa_embs,
-                    encoder_hidden_states=vl_embeds,
-                    timestep=timesteps_tensor,
-                )
-            pred = self.action_decoder(model_output, embodiment_id)
-
-            pred_velocity = pred[:, -self.action_horizon :]
-
-            # Update actions using euler integration.
-            actions = actions + dt * pred_velocity * vel_strength
-
-        return BatchFeature(
-            data={
-                "action_pred": actions,
-                "backbone_features": vl_embeds,
-                "state_features": state_features,
-            }
-        )
-
-    @torch.no_grad()
-    def get_action(
-        self,
-        backbone_output: BatchFeature,
-        action_input: BatchFeature,
-        options: dict[str, Any] | None = None,
-    ) -> BatchFeature:
-        """
-        Generate actions using the flow matching diffusion process.
-
-        Args:
-            backbone_output: Output from the backbone model containing:
-                - backbone_features: [B, seq_len, backbone_embedding_dim]
-                - backbone_attention_mask: [B, seq_len]
-            action_input: Input containing:
-                - state: [B, state_dim]
-                - embodiment_id: [B] (embodiment IDs)
-
-        Returns:
-            BatchFeature containing:
-                - action_pred: [B, action_horizon, action_dim] predicted actions
-        """
-        features = self._encode_features(backbone_output, action_input)
-        return self.get_action_with_features(
-            backbone_features=features.backbone_features,
-            state_features=features.state_features,
-            embodiment_id=action_input.embodiment_id,
-            backbone_output=backbone_output,
-            action_input=action_input,
-            options=options,
-        )
-
-    @property
-    def device(self):
-        return next(iter(self.parameters())).device
-
-    @property
-    def dtype(self):
-        return next(iter(self.parameters())).dtype
-
-    def prepare_input(self, batch: dict) -> BatchFeature:
-        """Prepare input batch for the action head."""
-        return BatchFeature(data=batch)
-
-
-def _load_qwen3_vl_cls() -> type:
-    """Import Qwen3-VL lazily so the missing-dep failure is clear and deferred."""
-    try:
-        from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration
-    except ImportError as exc:
-        raise ImportError(
-            "GR00T-N1.7 requires transformers>=4.57.1 for transformers.models.qwen3_vl "
-            "(the repo's declared floor is 4.56.0). Please upgrade transformers."
-        ) from exc
-    return Qwen3VLForConditionalGeneration
-
-
-class _Qwen3VLBackbone(nn.Module):
-    """GR00T adapter around the shared Qwen3-VL implementation."""
-
-    def __init__(
-        self,
-        model_name: str,
-        select_layer: int,
-        backbone_embedding_dim: int,
-        load_bf16: bool,
-        transformers_loading_kwargs: dict[str, Any] | None = None,
-    ):
-        super().__init__()
-        backbone_config = AutoConfig.from_pretrained(
-            model_name, **(transformers_loading_kwargs or {"trust_remote_code": True})
-        )
-        backbone_config.text_config._attn_implementation = "sdpa"
-        backbone_config.vision_config._attn_implementation = "sdpa"
-        self.model = _load_qwen3_vl_cls()(backbone_config).eval()
-        if load_bf16:
-            self.model.to(dtype=torch.bfloat16)
-
-        target_layers = select_layer if select_layer >= 0 else len(self.model.model.language_model.layers)
-        while len(self.model.model.language_model.layers) > target_layers:
-            self.model.model.language_model.layers.pop(-1)
-
-        # Drop the unused lm_head: feature-extractor only, and .to(dtype) above unties it (~0.58 GiB).
-        self.model.lm_head = None
-
-    def set_frozen_modules_to_eval_mode(self) -> None:
-        self.eval()
-
-    def prepare_input(self, batch: dict) -> BatchFeature:
-        return BatchFeature(data=batch)
-
-    def forward(self, vl_input: BatchFeature) -> BatchFeature:
-        self.set_frozen_modules_to_eval_mode()
-        keys_to_use = [
-            "input_ids",
-            "attention_mask",
-            "pixel_values",
-            "image_grid_thw",
-            "mm_token_type_ids",
-        ]
-        vl_input = {key: vl_input[key] for key in keys_to_use if key in vl_input}
-        # GR00T was trained against the **pre-final-RMSNorm** output of the last
-        # decoder layer (this was `outer.hidden_states[-1]` under HF 4.57.x semantics,
-        # where the @check_model_inputs tie-logic skipped overwriting because
-        # Qwen3VLCausalLMOutputWithPast has no `last_hidden_state` field).
-        # HF >=5.x changed the output_hidden_states tuple plumbing so that
-        # `hidden_states[-1]` is post-norm, which is a different tensor and breaks
-        # the trained action head. Capture the pre-norm tensor via a hook on the
-        # final RMSNorm regardless of transformers version.
-        _captured: list[torch.Tensor] = []
-
-        def _pre_norm_hook(_module, args, _out):
-            _captured.append(args[0])
-
-        norm = self.model.model.language_model.norm
-        handle = norm.register_forward_hook(_pre_norm_hook)
-        try:
-            self.model.model(**vl_input, return_dict=True)
-        finally:
-            handle.remove()
-        if not _captured:
-            raise RuntimeError("Failed to capture pre-norm hidden states from Qwen3-VL backbone")
-        backbone_features = _captured[-1]
-        image_mask = vl_input["input_ids"] == self.model.config.image_token_id
-        attention_mask = vl_input["attention_mask"] == 1
-        return BatchFeature(
-            data={
-                "backbone_features": backbone_features,
-                "backbone_attention_mask": attention_mask,
-                "image_mask": image_mask,
-            }
-        )
-
-
-def get_backbone_cls(config: Gr00tN1d7Config):
-    if "nvidia/Cosmos-Reason2" in config.model_name:
-        return _Qwen3VLBackbone
-    raise ValueError(f"Unsupported model name: {config.model_name}")
-
-
-class Gr00tN1d7(PreTrainedModel):
-    """Gr00tN1d7: VLA model with Cosmos-Reason2-2B (Qwen3-VL) backbone."""
-
-    config_class = Gr00tN1d7Config
-    supports_gradient_checkpointing = True
-    _tp_plan = {}
-
-    @property
-    def all_tied_weights_keys(self) -> dict[str, Any]:
-        return {}
-
-    def __init__(
-        self,
-        config: Gr00tN1d7Config,
-        transformers_loading_kwargs: dict | None = None,
-    ):
-        super().__init__(config)
-        if transformers_loading_kwargs is None:
-            transformers_loading_kwargs = {"trust_remote_code": True}
-        self.config = config
-
-        backbone_cls = get_backbone_cls(config)
-        self.backbone = backbone_cls(
-            model_name=config.model_name,
-            select_layer=config.select_layer,
-            backbone_embedding_dim=config.backbone_embedding_dim,
-            load_bf16=config.load_bf16,
-            transformers_loading_kwargs=transformers_loading_kwargs,
-        )
-
-        # Initialize action head
-        self.action_head = Gr00tN1d7ActionHead(config)
-
-        self.collator = Gr00tN1d7DataCollator(
-            model_name=config.model_name,
-            model_type=config.backbone_model_type,
-            transformers_loading_kwargs=transformers_loading_kwargs,
-        )
-
-    def prepare_input(self, inputs: dict) -> tuple[BatchFeature, BatchFeature]:
-        """Prepare inputs for backbone and action head."""
-
-        # NOTE -- currently the eval code doesn't use collator, so we need to add it here
-        if "vlm_content" in inputs:
-            # Fix for n_envs > 1: Process all environments' VLM content, not just the first
-            vlm_content_list = inputs["vlm_content"]
-            # Ensure vlm_content_list is always a list for consistent processing
-            if not isinstance(vlm_content_list, list):
-                vlm_content_list = [vlm_content_list]
-
-            # Process all VLM contents through the collator
-            prep = self.collator([{"vlm_content": vlm} for vlm in vlm_content_list])["inputs"]
-            inputs.pop("vlm_content")
-            inputs.update(prep)
-
-        backbone_inputs = self.backbone.prepare_input(inputs)
-        action_inputs = self.action_head.prepare_input(inputs)
-
-        backbone_inputs = backbone_inputs.to(device=self.device, dtype=self.dtype)
-        action_inputs = action_inputs.to(device=self.device, dtype=self.dtype)
-
-        return backbone_inputs, action_inputs
-
-    def forward(self, inputs: dict) -> BatchFeature:
-        """
-        Forward pass through the complete model.
-
-        Args:
-            inputs: Dictionary containing:
-                - Action inputs (state, action, embodiment_id, etc.)
-
-        Returns:
-            BatchFeature containing loss and other outputs
-        """
-        # Prepare inputs for backbone and action head
-        backbone_inputs, action_inputs = self.prepare_input(inputs)
-        backbone_outputs = self.backbone(backbone_inputs)
-        action_outputs = self.action_head(backbone_outputs, action_inputs)
-
-        return action_outputs
-
-    def get_action(self, inputs: dict, options: dict[str, Any] | None = None) -> BatchFeature:
-        """
-        Generate actions using the complete model.
-        """
-        # Prepare inputs for backbone and action head
-        backbone_inputs, action_inputs = self.prepare_input(inputs)
-
-        # Forward through backbone
-        backbone_outputs = self.backbone(backbone_inputs)
-        action_outputs = self.action_head.get_action(backbone_outputs, action_inputs, options)
-
-        return action_outputs
-
-    @property
-    def device(self):
-        return next(iter(self.parameters())).device
-
-    @property
-    def dtype(self):
-        return next(iter(self.parameters())).dtype
-
-
-# Register the model with HuggingFace
-AutoConfig.register("Gr00tN1d7", Gr00tN1d7Config)
-AutoModel.register(Gr00tN1d7Config, Gr00tN1d7)
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py
deleted file mode 100644
index 208f01a7cb5..00000000000
--- a/vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py
deleted file mode 100644
index 54fbe770d0d..00000000000
--- a/vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py
+++ /dev/null
@@ -1,344 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-import torch.nn.functional as F
-from diffusers import ConfigMixin, ModelMixin
-from diffusers.configuration_utils import register_to_config
-from diffusers.models.attention import Attention, FeedForward
-from diffusers.models.embeddings import SinusoidalPositionalEmbedding, TimestepEmbedding, Timesteps
-from torch import nn
-
-
-class AdaLayerNorm(nn.Module):
-    def __init__(
-        self,
-        embedding_dim: int,
-        norm_elementwise_affine: bool = False,
-        norm_eps: float = 1e-5,
-        chunk_dim: int = 0,
-    ):
-        super().__init__()
-        self.chunk_dim = chunk_dim
-        output_dim = embedding_dim * 2
-        self.silu = nn.SiLU()
-        self.linear = nn.Linear(embedding_dim, output_dim)
-        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        temb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        temb = self.linear(self.silu(temb))
-        scale, shift = temb.chunk(2, dim=1)
-        x = self.norm(x) * (1 + scale[:, None]) + shift[:, None]
-        return x
-
-
-class BasicTransformerBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        dropout=0.0,
-        cross_attention_dim: int | None = None,
-        activation_fn: str = "geglu",
-        attention_bias: bool = False,
-        upcast_attention: bool = False,
-        norm_elementwise_affine: bool = True,
-        norm_type: str = "layer_norm",
-        norm_eps: float = 1e-5,
-        final_dropout: bool = False,
-        attention_type: str = "default",
-        positional_embeddings: str | None = None,
-        num_positional_embeddings: int | None = None,
-        ff_inner_dim: int | None = None,
-        ff_bias: bool = True,
-        attention_out_bias: bool = True,
-    ):
-        super().__init__()
-        self.norm_type = norm_type
-
-        if positional_embeddings and (num_positional_embeddings is None):
-            raise ValueError(
-                "If `positional_embedding` type is defined, `num_position_embeddings` must also be defined."
-            )
-
-        if positional_embeddings == "sinusoidal":
-            self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)
-        else:
-            self.pos_embed = None
-
-        if norm_type == "ada_norm":
-            self.norm1 = AdaLayerNorm(dim)
-        else:
-            self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
-
-        self.attn1 = Attention(
-            query_dim=dim,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            dropout=dropout,
-            bias=attention_bias,
-            cross_attention_dim=cross_attention_dim,
-            upcast_attention=upcast_attention,
-            out_bias=attention_out_bias,
-        )
-
-        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
-        self.ff = FeedForward(
-            dim,
-            dropout=dropout,
-            activation_fn=activation_fn,
-            final_dropout=final_dropout,
-            inner_dim=ff_inner_dim,
-            bias=ff_bias,
-        )
-        self.final_dropout = nn.Dropout(dropout) if final_dropout else None
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        encoder_hidden_states: torch.Tensor | None = None,
-        encoder_attention_mask: torch.Tensor | None = None,
-        temb: torch.LongTensor | None = None,
-    ) -> torch.Tensor:
-        if self.norm_type == "ada_norm":
-            norm_hidden_states = self.norm1(hidden_states, temb)
-        else:
-            norm_hidden_states = self.norm1(hidden_states)
-
-        if self.pos_embed is not None:
-            norm_hidden_states = self.pos_embed(norm_hidden_states)
-
-        attn_output = self.attn1(
-            norm_hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            attention_mask=(encoder_attention_mask if encoder_hidden_states is not None else attention_mask),
-        )
-        if self.final_dropout:
-            attn_output = self.final_dropout(attn_output)
-
-        hidden_states = attn_output + hidden_states
-        if hidden_states.ndim == 4:
-            hidden_states = hidden_states.squeeze(1)
-
-        norm_hidden_states = self.norm3(hidden_states)
-        ff_output = self.ff(norm_hidden_states)
-
-        hidden_states = ff_output + hidden_states
-        if hidden_states.ndim == 4:
-            hidden_states = hidden_states.squeeze(1)
-        return hidden_states
-
-
-class TimestepEncoder(nn.Module):
-    def __init__(self, embedding_dim: int):
-        super().__init__()
-        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)
-        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
-
-    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
-        dtype = next(self.parameters()).dtype
-        return self.timestep_embedder(self.time_proj(timesteps).to(dtype))
-
-
-class DiT(ModelMixin, ConfigMixin):
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-        self,
-        num_attention_heads: int = 8,
-        attention_head_dim: int = 64,
-        output_dim: int = 26,
-        num_layers: int = 12,
-        dropout: float = 0.1,
-        attention_bias: bool = True,
-        activation_fn: str = "gelu-approximate",
-        num_embeds_ada_norm: int | None = 1000,
-        upcast_attention: bool = False,
-        norm_type: str = "ada_norm",
-        norm_elementwise_affine: bool = False,
-        norm_eps: float = 1e-5,
-        max_num_positional_embeddings: int = 512,
-        compute_dtype=torch.float32,
-        final_dropout: bool = True,
-        positional_embeddings: str | None = "sinusoidal",
-        interleave_self_attention=False,
-        cross_attention_dim: int | None = None,
-    ):
-        super().__init__()
-
-        self.inner_dim = num_attention_heads * attention_head_dim
-        self.gradient_checkpointing = False
-        self.timestep_encoder = TimestepEncoder(embedding_dim=self.inner_dim)
-
-        all_blocks = []
-        for idx in range(num_layers):
-            use_self_attn = idx % 2 == 1 and interleave_self_attention
-            curr_cross_attention_dim = cross_attention_dim if not use_self_attn else None
-            all_blocks.append(
-                BasicTransformerBlock(
-                    self.inner_dim,
-                    num_attention_heads,
-                    attention_head_dim,
-                    dropout=dropout,
-                    activation_fn=activation_fn,
-                    attention_bias=attention_bias,
-                    upcast_attention=upcast_attention,
-                    norm_type=norm_type,
-                    norm_elementwise_affine=norm_elementwise_affine,
-                    norm_eps=norm_eps,
-                    positional_embeddings=positional_embeddings,
-                    num_positional_embeddings=max_num_positional_embeddings,
-                    final_dropout=final_dropout,
-                    cross_attention_dim=curr_cross_attention_dim,
-                )
-            )
-        self.transformer_blocks = nn.ModuleList(all_blocks)
-        self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
-        self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
-        self.proj_out_2 = nn.Linear(self.inner_dim, output_dim)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.LongTensor | None = None,
-        encoder_attention_mask: torch.Tensor | None = None,
-        return_all_hidden_states: bool = False,
-    ):
-        temb = self.timestep_encoder(timestep)
-        hidden_states = hidden_states.contiguous()
-        encoder_hidden_states = encoder_hidden_states.contiguous()
-        all_hidden_states = [hidden_states]
-
-        for idx, block in enumerate(self.transformer_blocks):
-            if idx % 2 == 1 and self.config.interleave_self_attention:
-                hidden_states = block(hidden_states, temb=temb)
-            else:
-                hidden_states = block(
-                    hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=encoder_attention_mask,
-                    temb=temb,
-                )
-            all_hidden_states.append(hidden_states)
-
-        shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=1)
-        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
-        if return_all_hidden_states:
-            return self.proj_out_2(hidden_states), all_hidden_states
-        return self.proj_out_2(hidden_states)
-
-
-class AlternateVLDiT(DiT):
-    def __init__(self, *args, attend_text_every_n_blocks: int = 2, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.attend_text_every_n_blocks = attend_text_every_n_blocks
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.LongTensor | None = None,
-        encoder_attention_mask: torch.Tensor | None = None,
-        return_all_hidden_states: bool = False,
-        image_mask: torch.Tensor | None = None,
-        backbone_attention_mask: torch.Tensor | None = None,
-    ):
-        assert image_mask is not None, "Image mask is required"
-        assert self.config.interleave_self_attention, "Interleave self attention must be enabled"
-
-        temb = self.timestep_encoder(timestep)
-        hidden_states = hidden_states.contiguous()
-        encoder_hidden_states = encoder_hidden_states.contiguous()
-
-        image_attention_mask = image_mask & backbone_attention_mask
-        non_image_attention_mask = (~image_mask) & backbone_attention_mask
-
-        all_hidden_states = [hidden_states]
-        for idx, block in enumerate(self.transformer_blocks):
-            if idx % 2 == 1:
-                hidden_states = block(hidden_states, temb=temb)
-            else:
-                curr_encoder_attention_mask = (
-                    non_image_attention_mask
-                    if idx % (2 * self.attend_text_every_n_blocks) == 0
-                    else image_attention_mask
-                )
-                hidden_states = block(
-                    hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=curr_encoder_attention_mask,
-                    temb=temb,
-                )
-            all_hidden_states.append(hidden_states)
-
-        shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=1)
-        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
-        if return_all_hidden_states:
-            return self.proj_out_2(hidden_states), all_hidden_states
-        return self.proj_out_2(hidden_states)
-
-
-class SelfAttentionTransformer(ModelMixin, ConfigMixin):
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-        self,
-        num_attention_heads: int = 8,
-        attention_head_dim: int = 64,
-        output_dim: int = 26,
-        num_layers: int = 12,
-        dropout: float = 0.1,
-        attention_bias: bool = True,
-        activation_fn: str = "gelu-approximate",
-        num_embeds_ada_norm: int | None = 1000,
-        upcast_attention: bool = False,
-        max_num_positional_embeddings: int = 512,
-        compute_dtype=torch.float32,
-        final_dropout: bool = True,
-        positional_embeddings: str | None = "sinusoidal",
-        interleave_self_attention=False,
-    ):
-        super().__init__()
-
-        self.inner_dim = num_attention_heads * attention_head_dim
-        self.gradient_checkpointing = False
-        self.transformer_blocks = nn.ModuleList(
-            [
-                BasicTransformerBlock(
-                    self.inner_dim,
-                    num_attention_heads,
-                    attention_head_dim,
-                    dropout=dropout,
-                    activation_fn=activation_fn,
-                    attention_bias=attention_bias,
-                    upcast_attention=upcast_attention,
-                    positional_embeddings=positional_embeddings,
-                    num_positional_embeddings=max_num_positional_embeddings,
-                    final_dropout=final_dropout,
-                )
-                for _ in range(num_layers)
-            ]
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        return_all_hidden_states: bool = False,
-    ):
-        hidden_states = hidden_states.contiguous()
-        all_hidden_states = [hidden_states]
-        for block in self.transformer_blocks:
-            hidden_states = block(hidden_states)
-            all_hidden_states.append(hidden_states)
-        if return_all_hidden_states:
-            return hidden_states, all_hidden_states
-        return hidden_states
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py
deleted file mode 100644
index 81fcb4726f6..00000000000
--- a/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py
+++ /dev/null
@@ -1,178 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from vllm_omni.diffusion.utils.flow_matching import SinusoidalPositionalEncoding, swish
-
-
-class CategorySpecificLinear(nn.Module):
-    """Linear layer with category-specific weights and biases for multi-embodiment support."""
-
-    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int) -> None:
-        super().__init__()
-        self.num_categories = num_categories
-        # For each category, we have separate weights and biases.
-        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
-        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))
-
-    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
-        """Apply per-category linear projection.
-
-        Args:
-            x: ``[B, T, input_dim]`` input tensor.
-            cat_ids: ``[B]`` category / embodiment IDs.
-
-        Returns:
-            ``[B, T, hidden_dim]`` output tensor.
-        """
-        selected_W = self.W[cat_ids]
-        selected_b = self.b[cat_ids]
-        return torch.bmm(x, selected_W) + selected_b.unsqueeze(1)
-
-    def expand_action_dimension(
-        self,
-        old_action_dim: int,
-        new_action_dim: int,
-        expand_input: bool = False,
-        expand_output: bool = False,
-    ) -> None:
-        """Safely expand action dimension with explicit targeting.
-
-        Args:
-            old_action_dim: Original action dimension.
-            new_action_dim: New (larger) action dimension.
-            expand_input: Whether to expand input dimension (dim=1).
-            expand_output: Whether to expand output dimension (dim=2).
-        """
-        if new_action_dim <= old_action_dim:
-            raise ValueError(f"New action dim {new_action_dim} must be larger than old action dim {old_action_dim}")
-
-        # Expand input dimension (dim=1) only if explicitly requested AND dimensions match
-        if expand_input and self.W.shape[1] == old_action_dim:
-            repeat_times = new_action_dim // old_action_dim
-            remainder = new_action_dim % old_action_dim
-
-            new_W_parts = [self.W] * repeat_times
-            if remainder > 0:
-                new_W_parts.append(self.W[:, :remainder, :])
-
-            new_W = torch.cat(new_W_parts, dim=1)
-            self.W = nn.Parameter(new_W)
-
-        # Expand output dimension (dim=2) only if explicitly requested AND dimensions match
-        if expand_output and self.W.shape[2] == old_action_dim:
-            repeat_times = new_action_dim // old_action_dim
-            remainder = new_action_dim % old_action_dim
-
-            new_W_parts = [self.W] * repeat_times
-            if remainder > 0:
-                new_W_parts.append(self.W[:, :, :remainder])
-
-            new_W = torch.cat(new_W_parts, dim=2)
-            self.W = nn.Parameter(new_W)
-
-            # Expand bias for output dimension
-            if self.b.shape[1] == old_action_dim:
-                new_b_parts = [self.b] * repeat_times
-                if remainder > 0:
-                    new_b_parts.append(self.b[:, :remainder])
-
-                new_b = torch.cat(new_b_parts, dim=1)
-                self.b = nn.Parameter(new_b)
-
-
-class CategorySpecificMLP(nn.Module):
-    """Two-layer MLP with category-specific weights for multi-embodiment support."""
-
-    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int, output_dim: int) -> None:
-        super().__init__()
-        self.num_categories = num_categories
-        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
-        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)
-
-    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
-        """Apply two-layer per-category MLP.
-
-        Args:
-            x: ``[B, T, input_dim]`` input tensor.
-            cat_ids: ``[B]`` category / embodiment IDs.
-
-        Returns:
-            ``[B, T, output_dim]`` output tensor.
-        """
-        hidden = F.relu(self.layer1(x, cat_ids))
-        return self.layer2(hidden, cat_ids)
-
-    def expand_action_dimension(self, old_action_dim: int, new_action_dim: int) -> None:
-        """Expand action dimension by copying weights from existing dimensions.
-
-        Args:
-            old_action_dim: Original action dimension.
-            new_action_dim: New (larger) action dimension.
-        """
-        # self.layer1 does not take action_dim as input, so no expansion needed
-        self.layer2.expand_action_dimension(old_action_dim, new_action_dim, expand_input=False, expand_output=True)
-
-
-class MultiEmbodimentActionEncoder(nn.Module):
-    """Action encoder with multi-embodiment support and sinusoidal positional encoding."""
-
-    def __init__(self, action_dim: int, hidden_size: int, num_embodiments: int) -> None:
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.num_embodiments = num_embodiments
-
-        # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w}
-        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)  # (d -> w)
-        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)  # (2w -> w)
-        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)  # (w -> w)
-        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)
-
-    def forward(self, actions: torch.Tensor, timesteps: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
-        """Encode actions with embodiment-conditioned MLP and time embedding.
-
-        Args:
-            actions: ``[B, T, action_dim]`` action tensor.
-            timesteps: ``[B]`` timesteps (one scalar per batch item).
-            cat_ids: ``[B]`` category / embodiment IDs.
-
-        Returns:
-            ``[B, T, hidden_size]`` encoded action features.
-        """
-        B, T, _ = actions.shape
-
-        # 1) Expand each batch's single scalar time 'tau' across all T steps
-        #    so that shape => (B, T)
-        #    e.g. if timesteps is (B,), replicate across T
-        if timesteps.dim() == 1 and timesteps.shape[0] == B:
-            # shape (B,) => (B,T)
-            timesteps = timesteps.unsqueeze(1).expand(-1, T)
-        else:
-            raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.")
-
-        # 2) Standard action MLP step for shape => (B, T, w)
-        a_emb = self.W1(actions, cat_ids)
-
-        # 3) Get the sinusoidal encoding (B, T, w)
-        tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype)
-
-        # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish
-        x = torch.cat([a_emb, tau_emb], dim=-1)
-        x = swish(self.W2(x, cat_ids))
-
-        # 5) Finally W3 => (B, T, w)
-        x = self.W3(x, cat_ids)
-        return x
-
-    def expand_action_dimension(self, old_action_dim: int, new_action_dim: int) -> None:
-        """Expand action dimension by copying weights from existing dimensions.
-
-        Args:
-            old_action_dim: Original action dimension.
-            new_action_dim: New (larger) action dimension.
-        """
-        # Only W1 takes action_dim as input, so only expand its input dimension
-        self.W1.expand_action_dimension(old_action_dim, new_action_dim, expand_input=True, expand_output=False)
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
deleted file mode 100644
index ba128f4505b..00000000000
--- a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
+++ /dev/null
@@ -1,555 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import json
-import os
-import re
-import warnings
-from copy import deepcopy
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-import torchvision.transforms.v2 as transforms
-from PIL import Image
-from transformers import AutoProcessor, ProcessorMixin
-from transformers.feature_extraction_utils import BatchFeature
-from transformers.utils import cached_file
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig
-from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag
-from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor
-from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable
-
-
-class LetterBoxTransform:
-    """Pad image to square dimensions by adding black bars to the smaller side."""
-
-    def __call__(self, img: torch.Tensor) -> torch.Tensor:
-        import math
-
-        *leading_dims, c, h, w = img.shape
-        if h == w:
-            return img
-        max_dim = max(h, w)
-        pad_h, pad_w = max_dim - h, max_dim - w
-        pad_top, pad_left = pad_h // 2, pad_w // 2
-        pad_bottom, pad_right = pad_h - pad_top, pad_w - pad_left
-        if leading_dims:
-            batch_size = math.prod(leading_dims)
-            img_r = img.reshape(batch_size, c, h, w)
-            padded = transforms.functional.pad(img_r, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0)
-            return padded.reshape(leading_dims + [c, max_dim, max_dim])
-        return transforms.functional.pad(img, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0)
-
-
-def _build_eval_image_transform(
-    image_target_size: list[int],
-    image_crop_size: list[int],
-) -> transforms.Compose:
-    """Deterministic eval/inference image transform (letterbox → resize → centercrop → resize)."""
-    return transforms.Compose(
-        [
-            transforms.ToImage(),
-            LetterBoxTransform(),
-            transforms.Resize(size=image_target_size),
-            transforms.CenterCrop(size=image_crop_size),
-            transforms.Resize(size=image_target_size),
-        ]
-    )
-
-
-logger = init_logger(__name__)
-
-# Suppress protobuf deprecation warnings
-warnings.filterwarnings("ignore", category=DeprecationWarning, module="google.protobuf")
-
-EMBODIMENT_TAG_TO_PROJECTOR_INDEX = {
-    # Pretrain embodiment ids
-    "oxe_droid_relative_eef_relative_joint": 24,
-    "xdof_relative_eef_relative_joint": 27,
-    "xdof_relative_eef_relative_joint_subtask": 27,
-    "real_g1_relative_eef_relative_joints": 25,
-    "real_r1_pro_sharpa_relative_eef": 26,
-    "real_r1_pro_sharpa_relative_eef_human": 26,
-    "real_r1_pro_sharpa_relative_eef_maxinsights": 26,
-    "real_r1_pro_sharpa_relative_eef_mecka": 26,
-    # Posttrain embodiment ids
-    "unitree_g1_full_body_with_waist_height_nav_cmd": 25,
-    "unitree_g1_sonic": 11,
-    "simpler_env_google": 0,
-    "simpler_env_widowx": 1,
-    "libero_sim": 2,
-    "new_embodiment": 10,
-}
-
-QWEN3_VL_2B_PROCESSOR = "Qwen/Qwen3-VL-2B-Instruct"
-
-
-def build_processor(model_name: str, transformers_loading_kwargs: dict) -> ProcessorMixin:
-    try:
-        from transformers import Qwen3VLProcessor
-    except ImportError as exc:
-        raise ImportError(
-            "GR00T-N1.7 requires transformers>=4.57.1 for Qwen3VLProcessor "
-            "(the repo's declared floor is 4.56.0). Please upgrade transformers."
-        ) from exc
-    if model_name == "nvidia/Cosmos-Reason2-2B":
-        # Cosmos-Reason2-2B lacks a Qwen3VLProcessor; fall back to upstream artifacts.
-        logger.warning_once(
-            "Substituting processor from %s because %s does not ship one. "
-            "If you fine-tune Cosmos-Reason2-2B's tokenizer/image processor, "
-            "load the processor explicitly instead of relying on this fallback.",
-            QWEN3_VL_2B_PROCESSOR,
-            model_name,
-        )
-        model_name = QWEN3_VL_2B_PROCESSOR
-    return Qwen3VLProcessor.from_pretrained(model_name, **transformers_loading_kwargs)
-
-
-class Gr00tN1d7DataCollator:
-    def __init__(
-        self,
-        model_name: str,
-        model_type: str = "qwen",
-        transformers_loading_kwargs: dict | None = None,
-    ):
-        if transformers_loading_kwargs is None:
-            transformers_loading_kwargs = {}
-        self.processor = build_processor(model_name, transformers_loading_kwargs)
-        self.processor.tokenizer.padding_side = "left"
-        self.model_type = model_type
-        self.model_name = model_name
-
-    def __call__(self, features: list[dict[str, Any]]) -> BatchFeature:
-        batch = {}
-        keys = list(set().union(*(elem.keys() for elem in features)))
-
-        for key in keys:
-            values = [elem[key] for elem in features if key in elem]
-            if key == "vlm_content":
-                text_list = []
-                image_inputs = []
-                for v in values:
-                    curr_text_list = [v["text"]]
-
-                    text_list += curr_text_list
-                    curr_image_inputs = v["images"]
-                    image_inputs += curr_image_inputs
-
-                vlm_inputs = self.processor(
-                    text=text_list,
-                    images=image_inputs,
-                    return_tensors="pt",
-                    padding=True,
-                )
-                for k, v in vlm_inputs.items():
-                    batch[k] = v
-            elif key in (
-                "pixel_values",
-                "image_grid_thw",
-                "attention_mask",
-                "input_ids",
-            ):
-                raise Exception("Not implemented")
-            else:
-                batch[key] = torch.from_numpy(np.stack(values))
-        return BatchFeature(data={"inputs": batch})
-
-    def __str__(self):
-        return f"Gr00tN1d7DataCollator(model_name={self.model_name}, model_type={self.model_type})"
-
-
-class Gr00tN1d7Processor(ProcessorMixin):
-    data_collator_class = Gr00tN1d7DataCollator
-
-    def __init__(
-        self,
-        modality_configs: dict[str, dict[str, ModalityConfig]],
-        statistics: (dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None) = None,
-        use_percentiles: bool = False,
-        clip_outliers: bool = True,
-        image_crop_size: list[int] = None,
-        image_target_size: list[int] = None,
-        shortest_image_edge: int = 256,
-        crop_fraction: float = 0.95,
-        random_rotation_angle: int | None = None,
-        color_jitter_params: dict[str, float] | None = None,
-        formalize_language: bool = True,
-        model_name: str = "nvidia/Cosmos-Reason2-2B",
-        model_type: str = "qwen",
-        max_state_dim: int = 29,
-        max_action_dim: int = 29,
-        max_action_horizon: int = 50,
-        apply_sincos_state_encoding: bool = False,
-        use_relative_action: bool = False,
-        embodiment_id_mapping: dict[str, int] | None = None,
-        transformers_loading_kwargs: dict | None = None,
-        exclude_state: bool = False,
-        # Normalization
-        use_mean_std: bool = False,
-        **kwargs,  # absorb deprecated training-only keys from saved processor_config.json
-    ):
-        if transformers_loading_kwargs is None:
-            transformers_loading_kwargs = {"trust_remote_code": True}
-        if kwargs:
-            logger.debug("Gr00tN1d7Processor: ignoring unknown keys: %s", list(kwargs))
-        self.modality_configs = parse_modality_configs(modality_configs)
-
-        # Initialize StateActionProcessor for state/action normalization
-        self.state_action_processor = StateActionProcessor(
-            modality_configs=modality_configs,
-            statistics=statistics,
-            use_percentiles=use_percentiles,
-            clip_outliers=clip_outliers,
-            apply_sincos_state_encoding=apply_sincos_state_encoding,
-            use_relative_action=use_relative_action,
-        )
-
-        self.use_percentiles = use_percentiles
-        self.use_mean_std = use_mean_std
-        self.clip_outliers = clip_outliers
-        self.apply_sincos_state_encoding = apply_sincos_state_encoding
-        self.use_relative_action = use_relative_action
-
-        self.exclude_state = exclude_state
-
-        self.formalize_language = formalize_language
-        self.model_name = model_name
-        self.model_type = model_type
-
-        self.max_state_dim = max_state_dim
-        self.max_action_dim = max_action_dim
-        self.max_action_horizon = max_action_horizon
-
-        self.image_crop_size = image_crop_size
-        self.image_target_size = image_target_size
-        self.random_rotation_angle = random_rotation_angle
-        self.color_jitter_params = color_jitter_params
-        self.processor = build_processor(model_name, transformers_loading_kwargs)
-        self.processor.tokenizer.padding_side = "left"
-        self.embodiment_id_mapping = embodiment_id_mapping or EMBODIMENT_TAG_TO_PROJECTOR_INDEX
-        for k, v in EMBODIMENT_TAG_TO_PROJECTOR_INDEX.items():
-            if k not in self.embodiment_id_mapping:
-                self.embodiment_id_mapping[k] = v
-        self.shortest_image_edge = shortest_image_edge
-        self.crop_fraction = crop_fraction
-
-        self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {}
-
-        # Eval/inference image transform
-        self.eval_image_transform = _build_eval_image_transform(
-            image_target_size,
-            image_crop_size,
-        )
-        self._collator = self.data_collator_class(
-            model_name=model_name,
-            model_type=model_type,
-            transformers_loading_kwargs=transformers_loading_kwargs,
-        )
-
-    @property
-    def collator(self):
-        return self._collator
-
-    def set_statistics(
-        self,
-        statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]],
-        override: bool = False,
-    ) -> None:
-        """Set dataset statistics for normalization."""
-        for key in statistics:
-            if key not in self.statistics or override:
-                if override:
-                    logger.info("Overriding statistics for %s", key)
-                self.statistics[key] = deepcopy(statistics[key])
-            else:
-                logger.debug("Embodiment tag %s already in statistics, skipping update", key)
-
-        self.state_action_processor.set_statistics(statistics, override=override)
-
-        # Compute action dimensions for convenience
-        self.action_dim = {}
-        for embodiment_tag in self.state_action_processor.statistics:
-            self.action_dim[embodiment_tag] = self.state_action_processor.get_action_dim(embodiment_tag)
-
-    def decode_action(
-        self,
-        action: np.ndarray,
-        embodiment_tag: EmbodimentTag,
-        state: dict[str, np.ndarray] | None = None,
-    ):
-        """Undo action normalization and convert relative actions to absolute."""
-        # Split concatenated action into joint groups
-        out_dict = {}
-        start_idx = 0
-        joint_groups = self.modality_configs[embodiment_tag.value]["action"].modality_keys
-        action_horizon = len(self.modality_configs[embodiment_tag.value]["action"].delta_indices)
-        for key in joint_groups:
-            joint_dim = self.state_action_processor.norm_params[embodiment_tag.value]["action"][key]["dim"].item()
-            out_dict[key] = action[..., :action_horizon, start_idx : start_idx + joint_dim]
-            start_idx += joint_dim
-
-        # Use StateActionProcessor to unnormalize and convert to absolute
-        return self.state_action_processor.unapply_action(out_dict, embodiment_tag.value, state=state)
-
-    def _apply_vlm_processing(self, images: np.ndarray, language: str) -> BatchFeature:
-        pil_images = [Image.fromarray(np.transpose(v, (1, 2, 0))) for v in images]
-        conversation = [
-            {
-                "role": "user",
-                "content": [
-                    *[{"type": "image", "image": img} for img in pil_images],
-                    {"type": "text", "text": language},
-                ],
-            }
-        ]
-
-        text = self.processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False)
-        return {
-            "vlm_content": {
-                "text": text,
-                "images": pil_images,
-                "conversation": conversation,
-            }
-        }
-
-    def __call__(
-        self,
-        messages: list[dict[str, Any]],
-    ):
-        assert len(messages) == 1
-        content = messages[0]["content"]
-        embodiment_tag = content.embodiment
-        action_data = content.actions
-        state_data = content.states
-
-        norm_state_dict, normalized_actions = self.state_action_processor.apply(
-            state=state_data,
-            action=action_data,
-            embodiment_tag=embodiment_tag.value,
-        )
-
-        if normalized_actions:
-            action_keys = self.modality_configs[embodiment_tag.value]["action"].modality_keys
-            normalized_actions = torch.cat(
-                [torch.from_numpy(normalized_actions[key]) for key in action_keys],
-                dim=-1,
-            )  # (t, d)
-            action_dim = normalized_actions.shape[1]
-            normalized_actions = torch.cat(
-                [
-                    normalized_actions,
-                    torch.zeros(
-                        normalized_actions.shape[0],
-                        self.max_action_dim - normalized_actions.shape[1],
-                    ),
-                ],
-                dim=-1,
-            )  # (t, max_action_dim)
-            action_horizon = normalized_actions.shape[0]
-            assert action_horizon <= self.max_action_horizon, (
-                f"Action sequence length {action_horizon} exceeds max_action_horizon"
-                f" {self.max_action_horizon}. Increase model config action_horizon to"
-                f" >= {action_horizon}."
-            )
-            normalized_actions = torch.cat(
-                [
-                    normalized_actions,
-                    torch.zeros(
-                        self.max_action_horizon - normalized_actions.shape[0],
-                        self.max_action_dim,
-                    ),
-                ],
-                dim=0,
-            )  # (max_action_horizon, max_action_dim)
-            action_mask = torch.ones_like(normalized_actions)
-            action_mask[action_horizon:] = 0
-            action_mask[:, action_dim:] = 0
-        else:
-            normalized_actions = None
-            action_mask = None
-
-        state_keys = self.modality_configs[embodiment_tag.value]["state"].modality_keys
-        exclude_state = self.exclude_state or getattr(
-            self.modality_configs[embodiment_tag.value]["state"], "exclude_state", False
-        )
-        if exclude_state:
-            normalized_states = torch.cat(
-                [torch.from_numpy(np.zeros_like(state_data[key])) for key in state_keys], dim=-1
-            )
-        else:
-            normalized_states = torch.cat([torch.from_numpy(norm_state_dict[key]) for key in state_keys], dim=-1)
-        normalized_states = torch.cat(
-            [
-                normalized_states,
-                torch.zeros(
-                    normalized_states.shape[0],
-                    self.max_state_dim - normalized_states.shape[1],
-                ),
-            ],
-            dim=-1,
-        )
-
-        image_transform = self.eval_image_transform
-        image_keys = self.modality_configs[embodiment_tag.value]["video"].modality_keys
-
-        if self.formalize_language:
-            language = content.text.lower()
-            language = re.sub(r"[^\w\s]", "", language)
-        else:
-            language = content.text
-
-        vlm_inputs = self._get_vlm_inputs(
-            image_keys=image_keys,
-            images=content.images,
-            masks=content.masks,
-            image_transform=image_transform,
-            language=language,
-        )
-
-        transformed_inputs = {
-            "state": normalized_states.to(torch.get_default_dtype()),
-        }
-        if normalized_actions is not None:
-            transformed_inputs["action"] = normalized_actions.to(torch.get_default_dtype())
-        # Add VLM inputs
-        transformed_inputs.update(vlm_inputs)
-        if action_mask is not None:
-            transformed_inputs["action_mask"] = action_mask
-        transformed_inputs["embodiment_id"] = self.embodiment_id_mapping[embodiment_tag.value]
-        return transformed_inputs
-
-    def _get_vlm_inputs(
-        self,
-        image_keys: list[str],
-        images: list[Image.Image],
-        masks: dict[str, list[np.ndarray]] | None,
-        image_transform: transforms.Compose,
-        language: str,
-    ):
-        temporal_stacked_images = {}
-
-        if masks is not None:
-            raise ValueError("Mask-based transforms are not supported at inference.")
-        for view in image_keys:
-            assert view in images, f"{view} not in {images}"
-            temporal_stacked_images[view] = torch.stack([image_transform(img) for img in images[view]])  # (T, C, H, W)
-
-        for k, v in temporal_stacked_images.items():
-            assert isinstance(k, str), f"{k} is not a string"
-            assert isinstance(v, torch.Tensor), f"{v} is not a torch tensor"
-            assert v.ndim == 4, f"{v} is not a 4D tensor"
-            assert v.dtype == torch.uint8, f"{v} is not a uint8 tensor"
-            assert v.shape[1] == 3, f"{v} is not a 3 channel tensor"
-
-        stacked_images = (
-            torch.stack([temporal_stacked_images[view] for view in image_keys], dim=1).flatten(0, 1).numpy()
-        )  # (T*V, C, H, W), processor expects numpy array
-
-        vlm_inputs = self._apply_vlm_processing(stacked_images, language)
-        return vlm_inputs
-
-    def save_pretrained(self, save_directory: str | Path) -> list[Path]:
-        save_directory = Path(save_directory)
-        save_directory.mkdir(parents=True, exist_ok=True)
-        main_config_file = save_directory / "processor_config.json"
-        statistics_file = save_directory / "statistics.json"
-        embodiment_id_file = save_directory / "embodiment_id.json"
-
-        config = {
-            "processor_class": self.__class__.__name__,
-            "processor_kwargs": {
-                "modality_configs": to_json_serializable(self.modality_configs),
-                "image_crop_size": self.image_crop_size,
-                "image_target_size": self.image_target_size,
-                "random_rotation_angle": self.random_rotation_angle,
-                "color_jitter_params": self.color_jitter_params,
-                "shortest_image_edge": self.shortest_image_edge,
-                "crop_fraction": self.crop_fraction,
-                "model_name": self.model_name,
-                "model_type": self.model_type,
-                "formalize_language": self.formalize_language,
-                "max_state_dim": self.max_state_dim,
-                "max_action_dim": self.max_action_dim,
-                "max_action_horizon": self.max_action_horizon,
-                "use_percentiles": self.use_percentiles,
-                "use_mean_std": self.use_mean_std,
-                "clip_outliers": self.clip_outliers,
-                "apply_sincos_state_encoding": self.apply_sincos_state_encoding,
-                "use_relative_action": self.use_relative_action,
-                "exclude_state": self.exclude_state,
-            },
-        }
-        with open(main_config_file, "w") as f:
-            json.dump(config, f, indent=2)
-        with open(statistics_file, "w") as f:
-            json.dump(
-                to_json_serializable(self.state_action_processor.statistics),
-                f,
-                indent=2,
-            )
-        with open(embodiment_id_file, "w") as f:
-            json.dump(self.embodiment_id_mapping, f, indent=2)
-        return [main_config_file, statistics_file, embodiment_id_file]
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs):
-        transformers_loading_kwargs = kwargs.pop("transformers_loading_kwargs", {"trust_remote_code": True})
-        pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
-        config_file = pretrained_model_name_or_path / "processor_config.json"
-        statistics_file = pretrained_model_name_or_path / "statistics.json"
-        embodiment_id_file = pretrained_model_name_or_path / "embodiment_id.json"
-        is_local = os.path.isdir(pretrained_model_name_or_path)
-        if not is_local:
-            config_file = Path(cached_file(pretrained_model_name_or_path, "processor_config.json"))
-            statistics_file = Path(cached_file(pretrained_model_name_or_path, "statistics.json"))
-            embodiment_id_file = Path(cached_file(pretrained_model_name_or_path, "embodiment_id.json"))
-
-        with open(config_file) as f:
-            config = json.load(f)
-        with open(statistics_file) as f:
-            statistics = json.load(f)
-        if embodiment_id_file.exists():
-            with open(embodiment_id_file) as f:
-                embodiment_id_mapping = json.load(f)
-        else:
-            embodiment_id_mapping = None
-        processor_kwargs = config["processor_kwargs"]
-        processor_kwargs["statistics"] = statistics
-        processor_kwargs["embodiment_id_mapping"] = embodiment_id_mapping
-
-        # Backfill missing fields from older checkpoints.
-        processor_kwargs.setdefault("model_name", "nvidia/Cosmos-Reason2-2B")
-        processor_kwargs.setdefault("model_type", "qwen")
-        processor_kwargs.setdefault("clip_outliers", True)
-
-        # Directly override other processor kwargs
-        if kwargs:
-            # Override modality configs while keeping pretrained embodiment configs
-            modality_configs = kwargs.pop("modality_configs", {})
-            for embodiment_tag, modality_config in modality_configs.items():
-                processor_kwargs["modality_configs"][embodiment_tag] = modality_config
-            override_keys = [
-                "random_rotation_angle",
-                "color_jitter_params",
-                "use_relative_action",
-                "exclude_state",
-                "use_mean_std",
-                "model_name",
-                "model_type",
-                "max_action_horizon",
-                "max_state_dim",
-                "max_action_dim",
-            ]
-            for key in override_keys:
-                if key in kwargs:
-                    override = kwargs.pop(key)
-                    if override is not None:
-                        processor_kwargs[key] = override
-        return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs)
-
-
-AutoProcessor.register("Gr00tN1d7", Gr00tN1d7Processor)
diff --git a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py
deleted file mode 100644
index 98565a93faf..00000000000
--- a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py
+++ /dev/null
@@ -1,148 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from collections.abc import Iterable, Mapping
-from typing import Any
-
-import numpy as np
-import torch
-from torch import nn
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.models.gr00t.policy import Gr00tPolicy
-from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
-
-logger = init_logger(__name__)
-
-
-def _to_float32_action_dict(actions: Mapping[str, Any]) -> dict[str, np.ndarray]:
-    converted = {str(key): np.asarray(value, dtype=np.float32) for key, value in actions.items()}
-    if not converted:
-        raise RuntimeError("GR00T policy returned an empty action dict.")
-    return converted
-
-
-class Gr00tN1d7Pipeline(nn.Module):
-    """GR00T N1.7 policy pipeline backed by vLLM-Omni's local GR00T port.
-
-    vLLM-Omni owns the serving integration: OpenPI observations arrive through
-    `sampling_params.extra_args["robot_obs"]`, this pipeline runs GR00T policy
-    inference, and actions are returned through `DiffusionOutput.output["actions"]`.
-    """
-
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
-        super().__init__()
-        model_config = od_config.model_config
-        self.model_path = od_config.model
-        self.embodiment_tag = str(model_config.get("embodiment_tag") or "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT")
-        self.strict = bool(model_config.get("strict", True))
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-
-        logger.info("Loading GR00T N1.7 policy from %s with embodiment_tag=%s", self.model_path, self.embodiment_tag)
-        self.policy = Gr00tPolicy(
-            model_path=self.model_path,
-            embodiment_tag=self.embodiment_tag,
-            device=self.device,
-            strict=self.strict,
-        )
-        self._validate_policy_server_config(model_config.get("policy_server_config"))
-
-    def _validate_policy_server_config(self, psc: Mapping[str, Any] | None) -> None:
-        """Fail fast if the deploy handshake drifts from the loaded checkpoint.
-
-        ``policy_server_config`` is sent verbatim to the OpenPI client, so its
-        model/embodiment-specific values must match what the loaded policy actually
-        produces; otherwise the client is handed the wrong action contract.
-        """
-        if not isinstance(psc, Mapping):
-            return
-        action_config = self.policy.modality_configs["action"]
-        expected_horizon = len(action_config.delta_indices)
-        expected_keys = set(action_config.modality_keys)
-
-        if "action_horizon" in psc and psc["action_horizon"] != expected_horizon:
-            raise ValueError(
-                f"policy_server_config.action_horizon={psc['action_horizon']} != loaded model's "
-                f"action horizon {expected_horizon}."
-            )
-        if "action_keys" in psc and set(psc["action_keys"]) != expected_keys:
-            raise ValueError(
-                f"policy_server_config.action_keys={list(psc['action_keys'])} != loaded model's "
-                f"action keys {sorted(expected_keys)}."
-            )
-        psc_embodiment = psc.get("embodiment_tag")
-        if psc_embodiment is not None and psc_embodiment != self.embodiment_tag:
-            raise ValueError(
-                f"policy_server_config.embodiment_tag={psc_embodiment!r} != "
-                f"model_config.embodiment_tag={self.embodiment_tag!r}."
-            )
-
-    def reset(self) -> dict[str, Any]:
-        return self.policy.reset() or {}
-
-    @property
-    def weights_sources(self) -> tuple[Any, ...]:
-        return ()
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        consumed = list(weights)
-        if consumed:
-            raise RuntimeError(
-                f"Gr00tN1d7Pipeline.load_weights received {len(consumed)} weight tensors; "
-                "weights_sources=() should prevent this. GR00T weights are loaded directly by Gr00tPolicy."
-            )
-        return set()
-
-    def _dummy_actions(self) -> dict[str, np.ndarray]:
-        embodiment_value = self.policy.embodiment_tag.value
-        action_config = self.policy.modality_configs["action"]
-        horizon = len(action_config.delta_indices)
-        norm_params = self.policy.processor.state_action_processor.norm_params[embodiment_value]["action"]
-        actions = {}
-        for key in action_config.modality_keys:
-            dim = norm_params[key]["dim"]
-            dim = int(dim.item() if hasattr(dim, "item") else dim)
-            actions[key] = np.zeros((1, horizon, dim), dtype=np.float32)
-        return actions
-
-    @torch.inference_mode()
-    def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
-        del kwargs
-        extra_args = req.sampling_params.extra_args or {}
-        robot_obs = extra_args.get("robot_obs")
-        if robot_obs is None:
-            if req.request_id == DUMMY_DIFFUSION_REQUEST_ID:
-                return DiffusionOutput(output={"actions": self._dummy_actions()})
-            return DiffusionOutput(error="Gr00tN1d7Pipeline.forward expects sampling_params.extra_args['robot_obs'].")
-        if not isinstance(robot_obs, Mapping):
-            return DiffusionOutput(error=f"robot_obs must be a dict, got {type(robot_obs).__name__}.")
-
-        if extra_args.get("reset"):
-            self.reset()
-
-        policy_obs = _normalize_observation(robot_obs, language_key=self.policy.language_key)
-        result = self.policy.get_action(policy_obs)
-        actions = result[0] if isinstance(result, tuple) else result
-        if not isinstance(actions, Mapping):
-            return DiffusionOutput(error=f"GR00T policy returned {type(actions).__name__}; expected dict actions.")
-        # Return actions via output.output (like the DreamZero OpenPI policy) so the engine's
-        # empty-output guard passes.
-        return DiffusionOutput(output={"actions": _to_float32_action_dict(actions)})
-
-
-def _normalize_observation(robot_obs: Mapping[str, Any], *, language_key: str) -> dict[str, Any]:
-    obs: dict[str, Any] = {}
-    if "video" in robot_obs:
-        obs["video"] = robot_obs["video"]
-    elif "images" in robot_obs:
-        obs["video"] = robot_obs["images"]
-    if "state" in robot_obs:
-        obs["state"] = robot_obs["state"]
-    if "language" in robot_obs:
-        obs["language"] = robot_obs["language"]
-    else:
-        prompt = robot_obs.get("prompt")
-        if prompt is not None:
-            obs["language"] = {language_key: [[str(prompt)]]}
-    return obs
diff --git a/vllm_omni/diffusion/models/gr00t/policy.py b/vllm_omni/diffusion/models/gr00t/policy.py
deleted file mode 100644
index c16b344632f..00000000000
--- a/vllm_omni/diffusion/models/gr00t/policy.py
+++ /dev/null
@@ -1,434 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-from transformers import AutoModel, AutoProcessor
-
-from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import FINETUNE_ONLY_TAGS, POSTTRAIN_TAGS, EmbodimentTag
-from vllm_omni.diffusion.models.gr00t.dataio.types import MessageType, VLAStepData
-from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7
-from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7Processor
-
-
-def _rec_to_dtype(value: Any, dtype: torch.dtype) -> Any:
-    """Recursively convert floating-point tensors in nested collator output."""
-    if isinstance(value, torch.Tensor):
-        return value.to(dtype=dtype) if torch.is_floating_point(value) else value
-    if isinstance(value, dict) or hasattr(value, "items"):
-        return {key: _rec_to_dtype(item, dtype) for key, item in value.items()}
-    if isinstance(value, list):
-        return [_rec_to_dtype(item, dtype) for item in value]
-    return value
-
-
-class Gr00tPolicy:
-    """Core policy class for Gr00t model inference.
-
-    This policy handles the end-to-end inference pipeline:
-    1. Validates input observations
-    2. Processes observations with pretrained VLA processor
-    3. Runs model inference
-    4. Decodes and returns actions
-
-    The policy expects observations with specific modalities (video, state, language)
-    and returns actions in the format defined by the model's modality configuration.
-    """
-
-    def __init__(
-        self,
-        embodiment_tag: EmbodimentTag | str,
-        model_path: str,
-        *,
-        device: int | str,
-        strict: bool = True,
-    ):
-        """Initialize the Gr00t Policy.
-
-        Args:
-            embodiment_tag: The embodiment tag defining the robot/environment type.
-                Accepts an EmbodimentTag enum or a string (resolved case-insensitively).
-            model_path: Path to the pretrained model checkpoint directory
-            device: Device to run the model on (e.g., 'cuda:0', 0, 'cpu')
-            strict: Whether to enforce strict input validation (default: True)
-        """
-        self.strict = strict
-        if isinstance(embodiment_tag, str):
-            embodiment_tag = EmbodimentTag.resolve(embodiment_tag)
-        model_dir = Path(model_path)
-
-        # Load directly into bfloat16 to avoid a transient fp32 copy at load time.
-        self.model: Gr00tN1d7 = AutoModel.from_pretrained(model_dir, torch_dtype=torch.bfloat16)
-        self.model.eval()
-        self.model.to(device=device)
-
-        # Training saves processor files under a "processor/" subdirectory, but
-        # AutoProcessor expects them at the model root.  Fall back to the
-        # subdirectory when the root lacks a processor_config.json.
-        processor_dir = (
-            model_dir / "processor"
-            if (model_dir / "processor").is_dir() and not (model_dir / "processor_config.json").exists()
-            else model_dir
-        )
-        self.processor: Gr00tN1d7Processor = AutoProcessor.from_pretrained(processor_dir)
-
-        # Store embodiment-specific configurations
-        self.embodiment_tag = embodiment_tag
-        all_modality_configs = self.processor.modality_configs
-        if self.embodiment_tag.value not in all_modality_configs:
-            # Map raw checkpoint tag values to user-friendly enum names where possible.
-            supported_lines = []
-            for tag_value in sorted(all_modality_configs.keys()):
-                enum_name = EmbodimentTag.reverse_lookup(tag_value)
-                if enum_name != tag_value:
-                    supported_lines.append(f"  {enum_name:30s} (--embodiment-tag {enum_name})")
-                else:
-                    supported_lines.append(f"  {tag_value:30s} (internal, no public enum)")
-            supported_str = "\n".join(supported_lines)
-
-            hint = ""
-            if self.embodiment_tag in POSTTRAIN_TAGS:
-                hint = (
-                    f"\n\nHint: '{self.embodiment_tag.name}' is a posttrain tag that requires "
-                    f"a finetuned checkpoint, not the base model. "
-                    f"See the example READMEs for how to finetune and download checkpoints."
-                )
-            elif self.embodiment_tag in FINETUNE_ONLY_TAGS:
-                hint = (
-                    f"\n\nHint: '{self.embodiment_tag.name}' is for finetuning custom robots. "
-                    f"Use it with launch_finetune.py, not with the base model directly."
-                )
-
-            raise ValueError(
-                f"Embodiment tag '{self.embodiment_tag.name}' "
-                f"(value='{self.embodiment_tag.value}') is not supported "
-                f"by this checkpoint.\n\n"
-                f"Supported tags in this checkpoint:\n{supported_str}"
-                f"{hint}"
-            )
-        self.modality_configs = {
-            k: v for k, v in all_modality_configs[self.embodiment_tag.value].items() if k != "rl_info"
-        }
-        self.collate_fn = self.processor.collator
-
-        # Extract and validate language configuration
-        # Embodiments may define multiple language keys (e.g. paraphrases); use only the first at inference.
-        language_keys = self.modality_configs["language"].modality_keys
-        language_delta_indices = self.modality_configs["language"].delta_indices
-        assert len(language_keys) >= 1, "At least one language key is required"
-        assert len(language_delta_indices) == 1, "Only one language delta index is supported"
-        self.language_key = language_keys[0]
-
-    def _unbatch_observation(self, value: dict[str, Any]) -> list[dict[str, Any]]:
-        """Unbatch a batched observation into a list of single observations.
-
-        Args:
-            value: Batched observation with shape (B, ...) for each modality
-
-        Returns:
-            List of B observations, each with the batch dimension removed
-        """
-        unbatched_obs = []
-        # Infer batch size from the first video key
-        batch_size = value["video"][list(value["video"].keys())[0]].shape[0]
-
-        # Split each modality along the batch dimension
-        for i in range(batch_size):
-            unbatched_value = {
-                "video": {k: v[i] for k, v in value["video"].items()},
-                "state": {k: v[i] for k, v in value["state"].items()},
-                "language": {k: v[i] for k, v in value["language"].items()},
-            }
-            unbatched_obs.append(unbatched_value)
-        return unbatched_obs
-
-    def _to_vla_step_data(self, observation: dict[str, Any]) -> VLAStepData:
-        """Convert a single observation into a VLAStepData object for processing.
-
-        Args:
-            observation: Single observation dict with video, state, and language
-
-        Returns:
-            VLAStepData object ready for processor input
-        """
-        return VLAStepData(
-            images=observation["video"],
-            states=observation["state"],
-            actions={},  # No ground truth actions during inference
-            text=observation["language"][self.language_key][0],
-            embodiment=self.embodiment_tag,
-        )
-
-    def check_observation(self, observation: dict[str, Any]) -> None:
-        """Validate that the observation has the correct structure and types.
-
-        This method ensures that all required modalities are present and that their
-        data types, shapes, and dimensions match the model's expectations.
-
-        Expected observation structure:
-            - video: dict[str, np.ndarray[np.uint8, (B, T, H, W, C)]]
-                - B: batch size
-                - T: temporal horizon (number of frames)
-                - H, W: image height and width
-                - C: number of channels (must be 3 for RGB)
-            - state: dict[str, np.ndarray[np.float32, (B, T, D)]]
-                - B: batch size
-                - T: temporal horizon (number of state observations)
-                - D: state dimension
-            - language: dict[str, list[list[str]]]
-                - Shape: (B, T) where each element is a string
-                - T: temporal horizon (typically 1 for language)
-
-        Args:
-            observation: Dictionary containing video, state, and language modalities
-
-        Raises:
-            AssertionError: If any validation check fails
-        """
-        # Check that observation contains all required top-level modality keys
-        for modality in ["video", "state", "language"]:
-            assert modality in observation, f"Observation must contain a '{modality}' key"
-            assert isinstance(observation[modality], dict), (
-                f"Observation '{modality}' must be a dictionary. "
-                f"Got {type(observation[modality])}: {observation[modality]}"
-            )
-
-        # Track batch size across modalities to ensure consistency
-        bs = -1
-
-        for video_key in self.modality_configs["video"].modality_keys:
-            assert video_key in observation["video"], f"Video key '{video_key}' must be in observation"
-
-            # Set or verify batch size consistency across all video keys
-            if bs == -1:
-                bs = len(observation["video"][video_key])
-            else:
-                assert len(observation["video"][video_key]) == bs, (
-                    f"Video key '{video_key}' must have batch size {bs}. Got {len(observation['video'][video_key])}"
-                )
-
-            batched_video = observation["video"][video_key]
-
-            # Verify data type is numpy array
-            assert isinstance(batched_video, np.ndarray), (
-                f"Video key '{video_key}' must be a numpy array. Got {type(batched_video)}"
-            )
-
-            # Verify dtype is uint8 (standard for image data, range 0-255)
-            assert batched_video.dtype == np.uint8, (
-                f"Video key '{video_key}' must be a numpy array of type np.uint8. Got {batched_video.dtype}"
-            )
-
-            # Verify shape has 5 dimensions: (B, T, H, W, C)
-            assert batched_video.ndim == 5, (
-                f"Video key '{video_key}' must be a numpy array of shape (B, T, H, W, C), got {batched_video.shape}"
-            )
-
-            # Verify temporal dimension matches the expected horizon from config
-            assert batched_video.shape[1] == len(self.modality_configs["video"].delta_indices), (
-                f"Video key '{video_key}'s horizon must be "
-                f"{len(self.modality_configs['video'].delta_indices)}. Got {batched_video.shape[1]}"
-            )
-
-            # Verify channel dimension is 3 (RGB images)
-            assert batched_video.shape[-1] == 3, (
-                f"Video key '{video_key}'s channel 'C' must be 3. Got {batched_video.shape[-1]}"
-            )
-
-        for state_key in self.modality_configs["state"].modality_keys:
-            # Check that the expected state key exists in the observation
-            # Must happen before indexing; see video validation above.
-            assert state_key in observation["state"], f"State key '{state_key}' must be in observation"
-
-            # Set or verify batch size consistency across all state keys
-            if bs == -1:
-                bs = len(observation["state"][state_key])
-            else:
-                assert len(observation["state"][state_key]) == bs, (
-                    f"State key '{state_key}' must have batch size {bs}. Got {len(observation['state'][state_key])}"
-                )
-
-            batched_state = observation["state"][state_key]
-
-            # Verify data type is numpy array
-            assert isinstance(batched_state, np.ndarray), (
-                f"State key '{state_key}' must be a numpy array. Got {type(batched_state)}"
-            )
-
-            # Verify dtype is float32 (standard for continuous state values)
-            assert batched_state.dtype == np.float32, (
-                f"State key '{state_key}' must be a numpy array of type np.float32. Got {batched_state.dtype}"
-            )
-
-            # Verify shape has 3 dimensions: (B, T, D)
-            assert batched_state.ndim == 3, (
-                f"State key '{state_key}' must be a numpy array of shape (B, T, D), got {batched_state.shape}"
-            )
-
-            # Verify temporal dimension matches the expected horizon from config
-            assert batched_state.shape[1] == len(self.modality_configs["state"].delta_indices), (
-                f"State key '{state_key}'s horizon must be "
-                f"{len(self.modality_configs['state'].delta_indices)}. Got {batched_state.shape[1]}"
-            )
-
-        for language_key in self.modality_configs["language"].modality_keys:
-            # Check that the expected language key exists in the observation
-            # Must happen before indexing; see video validation above.
-            assert language_key in observation["language"], f"Language key '{language_key}' must be in observation"
-
-            # Set or verify batch size consistency (language uses len instead of .shape)
-            if bs == -1:
-                bs = len(observation["language"][language_key])
-            else:
-                assert len(observation["language"][language_key]) == bs, (
-                    f"Language key '{language_key}' must have batch size {bs}. "
-                    f"Got {len(observation['language'][language_key])}"
-                )
-
-            batched_language: list[list[str]] = observation["language"][language_key]
-
-            # Verify outer structure is a list (batch dimension)
-            assert isinstance(batched_language, list), (
-                f"Language key '{language_key}' must be a list. Got {type(batched_language)}"
-            )
-
-            # Validate each batch item
-            for batch_item in batched_language:
-                # Verify temporal dimension matches expected horizon
-                assert len(batch_item) == len(self.modality_configs["language"].delta_indices), (
-                    f"Language key '{language_key}'s horizon must be "
-                    f"{len(self.modality_configs['language'].delta_indices)}. Got {len(batched_language)}"
-                )
-
-                # Verify inner structure is also a list (temporal dimension)
-                assert isinstance(batch_item, list), f"Language batch item must be a list. Got {type(batch_item)}"
-
-                # Current implementation expects exactly one language instruction per timestep
-                assert len(batch_item) == 1, f"Language batch item must have exactly one item. Got {len(batch_item)}"
-
-                # Verify the instruction itself is a string
-                assert isinstance(batch_item[0], str), (
-                    f"Language batch item must be a string. Got {type(batch_item[0])}"
-                )
-
-    def _get_action(
-        self, observation: dict[str, Any], options: dict[str, Any] | None = None
-    ) -> tuple[dict[str, Any], dict[str, Any]]:
-        """Internal method to compute actions from observations.
-
-        Pipeline:
-        1. Unbatch observations into individual samples
-        2. Convert each to VLAStepData and process
-        3. Collate into model input batch
-        4. Run model inference
-        5. Decode and unnormalize actions
-
-        Args:
-            observation: Batched observation dictionary
-            options: Optional parameters (currently unused)
-
-        Returns:
-            Tuple of (actions_dict, info_dict)
-        """
-        # Step 1: Split batched observation into individual observations
-        unbatched_observations = self._unbatch_observation(observation)
-        processed_inputs = []
-
-        # Step 2: Process each observation through the VLA processor
-        states = []
-        for obs in unbatched_observations:
-            vla_step_data = self._to_vla_step_data(obs)
-            states.append(vla_step_data.states)  # dict[str, np.ndarray[np.float32, (T, D)]]
-            messages = [{"type": MessageType.EPISODE_STEP.value, "content": vla_step_data}]
-            processed_inputs.append(self.processor(messages))
-
-        # Step 3: Collate processed inputs into a single batch for model
-        collated_inputs = self.collate_fn(processed_inputs)
-        collated_inputs = _rec_to_dtype(collated_inputs, dtype=torch.bfloat16)
-
-        # Step 4: Run model inference to predict actions
-        with torch.inference_mode():
-            model_pred = self.model.get_action(**collated_inputs)
-        normalized_action = model_pred["action_pred"].float()
-
-        # Step 5: Decode actions from normalized space back to physical units
-        batched_states = {}
-        for k in self.modality_configs["state"].modality_keys:
-            batched_states[k] = np.stack([s[k] for s in states], axis=0)  # (B, T, D)
-        unnormalized_action = self.processor.decode_action(
-            normalized_action.cpu().numpy(), self.embodiment_tag, batched_states
-        )
-
-        # Cast all actions to float32 for consistency
-        casted_action = {key: value.astype(np.float32) for key, value in unnormalized_action.items()}
-        return casted_action, {}
-
-    def check_action(self, action: dict[str, Any]) -> None:
-        """Validate that the action has the correct structure and types.
-
-        This method ensures that all required action keys are present and that their
-        data types, shapes, and dimensions match the model's action space.
-
-        Expected action structure:
-            - action: dict[str, np.ndarray[np.float32, (B, T, D)]]
-                - B: batch size
-                - T: action horizon (number of future action steps)
-                - D: action dimension (e.g., joint positions, velocities, gripper state)
-
-        Args:
-            action: Dictionary containing action arrays for each action key
-
-        Raises:
-            AssertionError: If any validation check fails
-        """
-        # Validate each action key defined in the modality config
-        for action_key in self.modality_configs["action"].modality_keys:
-            # Check that the expected action key exists
-            assert action_key in action, f"Action key '{action_key}' must be in action"
-
-            action_arr = action[action_key]
-
-            # Verify data type is numpy array
-            assert isinstance(action_arr, np.ndarray), (
-                f"Action key '{action_key}' must be a numpy array. Got {type(action_arr)}"
-            )
-
-            # Verify dtype is float32 (standard for continuous actions)
-            assert action_arr.dtype == np.float32, (
-                f"Action key '{action_key}' must be a numpy array of type np.float32. Got {action_arr.dtype}"
-            )
-
-            # Verify shape has 3 dimensions: (B, T, D)
-            assert action_arr.ndim == 3, (
-                f"Action key '{action_key}' must be a numpy array of shape (B, T, D), got {action_arr.shape}"
-            )
-
-            # Verify action horizon matches the expected temporal dimension from config
-            assert action_arr.shape[1] == len(self.modality_configs["action"].delta_indices), (
-                f"Action key '{action_key}'s horizon must be "
-                f"{len(self.modality_configs['action'].delta_indices)}. Got {action_arr.shape[1]}"
-            )
-
-    def get_action(
-        self, observation: dict[str, Any], options: dict[str, Any] | None = None
-    ) -> tuple[dict[str, Any], dict[str, Any]]:
-        if self.strict:
-            self.check_observation(observation)
-        action, info = self._get_action(observation, options)
-        if self.strict:
-            self.check_action(action)
-        return action, info
-
-    def reset(self, options: dict[str, Any] | None = None) -> dict[str, Any]:
-        """Reset the policy to its initial state.
-
-        Args:
-            options: Dictionary containing the options for the reset
-
-        Returns:
-            Dictionary containing the info after resetting the policy
-        """
-        return {}
diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py
index 730c10f1c78..ed0a50fee38 100644
--- a/vllm_omni/diffusion/models/helios/helios_transformer.py
+++ b/vllm_omni/diffusion/models/helios/helios_transformer.py
@@ -11,7 +11,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import FP32LayerNorm
@@ -26,7 +25,6 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
@@ -604,13 +602,6 @@ class HeliosTransformer3DModel(nn.Module):
     guidance cross-attention, and chunked video generation support.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "blocks": ForwardPattern.Pattern_2,
-        },
-        has_separate_cfg=True,
-    )
-
     _repeated_blocks = ["HeliosTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["blocks"]
     packed_modules_mapping = {
diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py
new file mode 100644
index 00000000000..abc50eb8a22
--- /dev/null
+++ b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py
@@ -0,0 +1,1059 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+import json
+import math
+import os
+from collections.abc import Callable, Iterable
+from typing import Any
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import deprecate, logging
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    LlamaForCausalLM,
+    PreTrainedTokenizerFast,
+    T5EncoderModel,
+    T5Tokenizer,
+)
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+logger = logging.get_logger(__name__)
+
+
+def get_hidream_image_post_process_func(
+    od_config: OmniDiffusionConfig,
+):
+    model_name = od_config.model
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+    vae_config_path = os.path.join(model_path, "vae/config.json")
+    if not os.path.exists(vae_config_path):
+        raise FileNotFoundError(
+            f"VAE config not found at {vae_config_path}. "
+            "Please ensure the model path contains a valid VAE configuration."
+        )
+    with open(vae_config_path) as f:
+        vae_config = json.load(f)
+        vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+
+    def post_process_func(
+        images: torch.Tensor,
+    ):
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`list[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`list[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        timesteps (`torch.Tensor`): The timestep schedule from the scheduler.
+        num_inference_steps (`int`): The number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class HiDreamI1ImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, ProgressBarMixin):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+        llama_path = self.od_config.extras["auxiliary_text_encoder"]
+        if llama_path is None:
+            logger.warning(
+                f"auxiliary_text_encoder is not provided. "
+                f"Attempting to load default LLaMA model: {self.od_config.extras['default_llama_model_id']}"
+            )
+            llama_path = self.od_config.extras["default_llama_model_id"]
+        logger.info(f"Loading LLAMA model from: {llama_path}")
+
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+        ]
+
+        self.device = get_local_device()
+        model = od_config.model
+        # Check if model is a local path
+        local_files_only = os.path.exists(model)
+
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the in-repo component set before any from_pretrained
+        # (``text_encoder_4`` lives in a separate Llama repo and is unaffected).
+        hidream_subfolders = [
+            "scheduler",
+            "vae",
+            "text_encoder",
+            "tokenizer",
+            "text_encoder_2",
+            "tokenizer_2",
+            "text_encoder_3",
+            "tokenizer_3",
+        ]
+        prefetch_subfolders(model, hidream_subfolders, local_files_only=local_files_only)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKL.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
+        self.text_encoder = from_pretrained_with_prefetch(
+            CLIPTextModelWithProjection.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
+        )
+        self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+        self.text_encoder_2 = from_pretrained_with_prefetch(
+            CLIPTextModelWithProjection.from_pretrained,
+            model,
+            subfolder="text_encoder_2",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
+        )
+        self.tokenizer_2 = CLIPTokenizer.from_pretrained(
+            model, subfolder="tokenizer_2", local_files_only=local_files_only
+        )
+        self.text_encoder_3 = from_pretrained_with_prefetch(
+            T5EncoderModel.from_pretrained,
+            model,
+            subfolder="text_encoder_3",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
+        )
+        self.tokenizer_3 = T5Tokenizer.from_pretrained(
+            model, subfolder="tokenizer_3", local_files_only=local_files_only
+        )
+        self.text_encoder_4 = LlamaForCausalLM.from_pretrained(llama_path, output_hidden_states=True, dtype=dtype).to(
+            self.device
+        )
+        self.tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(llama_path, use_fast=False)
+        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HiDreamImageTransformer2DModel)
+        self.transformer = HiDreamImageTransformer2DModel(
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+        )
+
+        self.stage = None
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        # HiDreamImage latents are turned into 2x2 patches and packed.
+        # This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.default_sample_size = 128
+        if getattr(self, "tokenizer_4", None) is not None:
+            self.tokenizer_4.pad_token = self.tokenizer_4.eos_token
+
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        prompt_3,
+        prompt_4,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        negative_prompt_3=None,
+        negative_prompt_4=None,
+        prompt_embeds_t5=None,
+        prompt_embeds_llama3=None,
+        negative_prompt_embeds_t5=None,
+        negative_prompt_embeds_llama3=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        if prompt is not None and pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. "
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_2 is not None and pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_3 is not None and prompt_embeds_t5 is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_4 is not None and prompt_embeds_llama3 is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt is None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `pooled_prompt_embeds`."
+                "Cannot leave both `prompt` and `pooled_prompt_embeds` undefined."
+            )
+        elif prompt is None and prompt_embeds_t5 is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds_t5`."
+                "Cannot leave both `prompt` and `prompt_embeds_t5` undefined."
+            )
+        elif prompt is None and prompt_embeds_llama3 is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds_llama3`."
+                "Cannot leave both `prompt` and `prompt_embeds_llama3` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+        elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)):
+            raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}")
+        elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)):
+            raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}")
+
+        if negative_prompt is not None and negative_pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:"
+                f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:"
+                f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:"
+                f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:"
+                f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two."
+            )
+
+        if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None:
+            if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape:
+                raise ValueError(
+                    "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`"
+                    f" {negative_pooled_prompt_embeds.shape}."
+                )
+        if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None:
+            if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape:
+                raise ValueError(
+                    "`prompt_embeds_t5` and `negative_prompt_embeds_t5`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`"
+                    f" {negative_prompt_embeds_t5.shape}."
+                )
+        if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None:
+            if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape:
+                raise ValueError(
+                    "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`"
+                    f" {negative_prompt_embeds_llama3.shape}."
+                )
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: str | list[str] = None,
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or self.text_encoder_3.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        text_inputs = self.tokenizer_3(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, self.tokenizer_3.model_max_length),
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        attention_mask = text_inputs.attention_mask
+        untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_3.batch_decode(
+                untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_3(
+            text_input_ids.to(self.device), attention_mask=attention_mask.to(self.device)
+        )[0]
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        tokenizer,
+        text_encoder,
+        prompt: str | list[str],
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        model_max_length = self.tokenizer.model_max_length
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, model_max_length),
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = tokenizer.batch_decode(untruncated_ids[:, model_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {model_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True)
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds[0]
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+        return prompt_embeds
+
+    def _get_llama3_prompt_embeds(
+        self,
+        prompt: str | list[str] = None,
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or self.text_encoder_4.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        text_inputs = self.tokenizer_4(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, self.tokenizer_4.model_max_length),
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        attention_mask = text_inputs.attention_mask
+        untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_4.batch_decode(
+                untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}"
+            )
+
+        outputs = self.text_encoder_4(
+            text_input_ids.to(self.device),
+            attention_mask=attention_mask.to(self.device),
+            output_hidden_states=True,
+        )
+
+        prompt_embeds = outputs.hidden_states[1:]
+        prompt_embeds = torch.stack(prompt_embeds, dim=0)
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str] | None = None,
+        prompt_2: str | list[str] | None = None,
+        prompt_3: str | list[str] | None = None,
+        prompt_4: str | list[str] | None = None,
+        dtype: torch.dtype | None = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str | list[str] | None = None,
+        negative_prompt_2: str | list[str] | None = None,
+        negative_prompt_3: str | list[str] | None = None,
+        negative_prompt_4: str | list[str] | None = None,
+        prompt_embeds_t5: list[torch.FloatTensor] | None = None,
+        prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
+        pooled_prompt_embeds: torch.FloatTensor | None = None,
+        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+        max_sequence_length: int = 128,
+        lora_scale: float | None = None,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = pooled_prompt_embeds.shape[0]
+
+        if pooled_prompt_embeds is None:
+            pooled_prompt_embeds_1 = self._get_clip_prompt_embeds(
+                self.tokenizer, self.text_encoder, prompt, max_sequence_length, dtype
+            )
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+            if len(negative_prompt) > 1 and len(negative_prompt) != batch_size:
+                raise ValueError(f"negative_prompt must be of length 1 or {batch_size}")
+
+            negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds(
+                self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, dtype
+            )
+
+            if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1)
+
+        if pooled_prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            if len(prompt_2) > 1 and len(prompt_2) != batch_size:
+                raise ValueError(f"prompt_2 must be of length 1 or {batch_size}")
+
+            pooled_prompt_embeds_2 = self._get_clip_prompt_embeds(
+                self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, dtype
+            )
+
+            if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1:
+                pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1)
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+            negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+
+            if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size:
+                raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}")
+
+            negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds(
+                self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, dtype
+            )
+
+            if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1)
+
+        if pooled_prompt_embeds is None:
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1)
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1
+            )
+
+        if prompt_embeds_t5 is None:
+            prompt_3 = prompt_3 or prompt
+            prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
+
+            if len(prompt_3) > 1 and len(prompt_3) != batch_size:
+                raise ValueError(f"prompt_3 must be of length 1 or {batch_size}")
+
+            prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, dtype)
+
+            if prompt_embeds_t5.shape[0] == 1 and batch_size > 1:
+                prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds_t5 is None:
+            negative_prompt_3 = negative_prompt_3 or negative_prompt
+            negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
+
+            if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size:
+                raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}")
+
+            negative_prompt_embeds_t5 = self._get_t5_prompt_embeds(negative_prompt_3, max_sequence_length, dtype)
+
+            if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1:
+                negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1)
+
+        if prompt_embeds_llama3 is None:
+            prompt_4 = prompt_4 or prompt
+            prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4
+
+            if len(prompt_4) > 1 and len(prompt_4) != batch_size:
+                raise ValueError(f"prompt_4 must be of length 1 or {batch_size}")
+
+            prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, dtype)
+
+            if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1:
+                prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None:
+            negative_prompt_4 = negative_prompt_4 or negative_prompt
+            negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4
+
+            if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size:
+                raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}")
+
+            negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds(
+                negative_prompt_4, max_sequence_length, dtype
+            )
+
+            if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1:
+                negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+
+        # duplicate pooled_prompt_embeds for each generation per prompt
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt
+        bs_embed, seq_len, _ = prompt_embeds_t5.shape
+        if bs_embed == 1 and batch_size > 1:
+            prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1)
+        elif bs_embed > 1 and bs_embed != batch_size:
+            raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}")
+        prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt
+        _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape
+        if bs_embed == 1 and batch_size > 1:
+            prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+        elif bs_embed > 1 and bs_embed != batch_size:
+            raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}")
+        prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1)
+        prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim)
+
+        if do_classifier_free_guidance:
+            # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt
+            bs_embed, seq_len = negative_pooled_prompt_embeds.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}")
+            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt)
+            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+            # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt
+            bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}")
+            negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt
+            _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}")
+            negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1)
+            negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view(
+                -1, batch_size * num_images_per_prompt, seq_len, dim
+            )
+
+        return (
+            prompt_embeds_t5,
+            negative_prompt_embeds_t5,
+            prompt_embeds_llama3,
+            negative_prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        )
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        depr_message = (
+            f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this"
+            " method will be removed in a future version. Please use `pipe.vae.enable_slicing()`."
+        )
+        deprecate(
+            "enable_vae_slicing",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        depr_message = (
+            f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this"
+            " method will be removed in a future version. Please use `pipe.vae.disable_slicing()`."
+        )
+        deprecate(
+            "disable_vae_slicing",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        depr_message = (
+            f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}`"
+            " is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`."
+        )
+        deprecate(
+            "enable_vae_tiling",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        depr_message = (
+            f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is"
+            " deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`."
+        )
+        deprecate(
+            "disable_vae_tiling",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.disable_tiling()
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype)
+        else:
+            if latents.shape != shape:
+                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+            latents = latents.to(self.device)
+        return latents
+
+    def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len):
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            self.device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        return timesteps, num_inference_steps
+
+    def _extract_prompts(self, prompts):
+        """Extract prompt and negative_prompt from OmniPromptType list."""
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
+            negative_prompt = None
+        elif prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
+        else:
+            negative_prompt = None
+        return prompt, negative_prompt
+
+    def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool):
+        if get_classifier_free_guidance_world_size() == 1:
+            return True
+
+        if true_cfg_scale <= 1:
+            logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.")
+            return False
+
+        if not has_neg_prompt:
+            logger.warning(
+                "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings."
+            )
+            return False
+        return True
+
+    def diffuse(
+        self,
+        prompt_embeds_t5: torch.Tensor,
+        prompt_embeds_llama3: torch.Tensor,
+        pooled_prompt_embeds: torch.Tensor,
+        latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        do_true_cfg: bool,
+    ) -> torch.Tensor:
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timesteps=timestep,
+                    encoder_hidden_states_t5=prompt_embeds_t5,
+                    encoder_hidden_states_llama3=prompt_embeds_llama3,
+                    pooled_embeds=pooled_prompt_embeds,
+                    return_dict=False,
+                )[0]
+                noise_pred = -noise_pred
+
+                # TODO: Modify CFG guidance
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+
+                pbar.update()
+
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] = None,
+        prompt_2: str | list[str] | None = None,
+        prompt_3: str | list[str] | None = None,
+        prompt_4: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 5.0,
+        negative_prompt: str | list[str] | None = None,
+        negative_prompt_2: str | list[str] | None = None,
+        negative_prompt_3: str | list[str] | None = None,
+        negative_prompt_4: str | list[str] | None = None,
+        num_images_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.FloatTensor | None = None,
+        prompt_embeds_t5: torch.FloatTensor | None = None,
+        prompt_embeds_llama3: torch.FloatTensor | None = None,
+        negative_prompt_embeds_t5: torch.FloatTensor | None = None,
+        negative_prompt_embeds_llama3: torch.FloatTensor | None = None,
+        pooled_prompt_embeds: torch.FloatTensor | None = None,
+        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+        output_type: str | None = "pil",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end: Callable[[int, int], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 128,
+        **kwargs,
+    ):
+        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
+        prompt = extracted_prompt or prompt
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        generator = req.sampling_params.generator or generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or guidance_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
+
+        prompt_embeds = kwargs.get("prompt_embeds", None)
+        negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None)
+
+        if prompt_embeds is not None:
+            deprecation_message = (
+                "The `prompt_embeds` argument is deprecated."
+                " Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead."
+            )
+            deprecate("prompt_embeds", "0.35.0", deprecation_message)
+            prompt_embeds_t5 = prompt_embeds[0]
+            prompt_embeds_llama3 = prompt_embeds[1]
+
+        if negative_prompt_embeds is not None:
+            deprecation_message = (
+                "The `negative_prompt_embeds` argument is deprecated."
+                "Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead."
+            )
+            deprecate("negative_prompt_embeds", "0.35.0", deprecation_message)
+            negative_prompt_embeds_t5 = negative_prompt_embeds[0]
+            negative_prompt_embeds_llama3 = negative_prompt_embeds[1]
+
+        division = self.vae_scale_factor * 2
+        S_max = (self.default_sample_size * self.vae_scale_factor) ** 2
+        scale = S_max / (width * height)
+        scale = math.sqrt(scale)
+        width, height = int(width * scale // division * division), int(height * scale // division * division)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            prompt_3,
+            prompt_4,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            negative_prompt_4=negative_prompt_4,
+            prompt_embeds_t5=prompt_embeds_t5,
+            prompt_embeds_llama3=prompt_embeds_llama3,
+            negative_prompt_embeds_t5=negative_prompt_embeds_t5,
+            negative_prompt_embeds_llama3=negative_prompt_embeds_llama3,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        elif pooled_prompt_embeds is not None:
+            batch_size = pooled_prompt_embeds.shape[0]
+
+        # TODO: CFG guidance configuration
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
+
+        # 3. Encode prompt
+        lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None
+        (
+            prompt_embeds_t5,
+            negative_prompt_embeds_t5,
+            prompt_embeds_llama3,
+            negative_prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_3=prompt_3,
+            prompt_4=prompt_4,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            negative_prompt_4=negative_prompt_4,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds_t5=prompt_embeds_t5,
+            prompt_embeds_llama3=prompt_embeds_llama3,
+            negative_prompt_embeds_t5=negative_prompt_embeds_t5,
+            negative_prompt_embeds_llama3=negative_prompt_embeds_llama3,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0)
+            prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1)
+            pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            pooled_prompt_embeds.dtype,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
+        self._num_timesteps = len(timesteps)
+
+        # 6. Denoising loop
+        latents = self.diffuse(
+            prompt_embeds_t5,
+            prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            latents,
+            timesteps,
+            do_true_cfg,
+        )
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(
+            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
index e4c9e0b2a1e..5ce3098ad40 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
@@ -11,7 +11,6 @@
 import numpy as np
 import regex as re
 import torch
-from cache_dit import ForwardPattern
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
@@ -60,7 +59,6 @@
     AttentionMetadata,
 )
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
     get_classifier_free_guidance_rank,
@@ -2020,13 +2018,6 @@ def forward(self, x):
 
 
 class HunyuanImage3Model(nn.Module):
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "layers": ForwardPattern.Pattern_4,
-        },
-        check_forward_pattern=False,
-    )
-
     _sp_plan = {
         # Split custom_pos_emb tuple elements (cos, sin) at model forward input
         "pre_processor": {
diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
index 9224973dda8..9af9e848875 100644
--- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
+++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
@@ -6,7 +6,6 @@
 
 import torch
 import torch.nn.functional as F
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     CombinedTimestepTextProjEmbeddings,
     TimestepEmbedding,
@@ -23,7 +22,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size
@@ -593,13 +591,6 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
     tensor-parallel layers for the 54 main transformer blocks.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_0,
-        },
-        has_separate_cfg=True,
-    )
-
     _repeated_blocks = ["HunyuanVideo15TransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     packed_modules_mapping = {
diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
index 744c1806a5a..8f0ff446afd 100644
--- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
+++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
@@ -6,7 +6,6 @@
 
 import torch
 import torch.nn as nn
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
@@ -19,7 +18,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -583,14 +581,6 @@ class LongCatImageTransformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_1,
-            "single_transformer_blocks": ForwardPattern.Pattern_1,
-        },
-        has_separate_cfg=True,
-    )
-
     _repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
index 6072c04560a..dcf848390d4 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
@@ -23,7 +23,6 @@
 import torch.distributed
 import torch.nn as nn
 import torch.nn.functional as F
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings, PixArtAlphaTextProjection
 from diffusers.utils import (
     BaseOutput,
@@ -45,7 +44,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
@@ -1451,14 +1449,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
             The normalization layer to use.
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_0,
-        },
-        has_separate_cfg=True,
-        check_forward_pattern=False,
-    )
-
     _supports_gradient_checkpointing = True
     _skip_layerwise_casting_patterns = ["norm"]
     _repeated_blocks = ["LTX2VideoTransformerBlock"]
diff --git a/vllm_omni/diffusion/models/sd3/sd3_transformer.py b/vllm_omni/diffusion/models/sd3/sd3_transformer.py
index d8e03705112..89f06157758 100644
--- a/vllm_omni/diffusion/models/sd3/sd3_transformer.py
+++ b/vllm_omni/diffusion/models/sd3/sd3_transformer.py
@@ -5,7 +5,6 @@
 import torch.nn.functional as F
 
 # TODO replace this with vLLM implementation
-from cache_dit import ForwardPattern
 from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, SD35AdaLayerNormZeroX
@@ -20,7 +19,6 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 
 logger = init_logger(__name__)
@@ -388,12 +386,6 @@ class SD3Transformer2DModel(nn.Module):
     The Transformer model introduced in [Stable Diffusion 3](https://huggingface.co/papers/2403.03206).
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_1,
-        }
-    )
-
     _repeated_blocks = ["SD3TransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
 
diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
index c1ba258b28c..44446ae8381 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
@@ -534,8 +534,6 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
             self.llm_cfg,
             prefix="language_model",
         )
-        # We define this for Cache DiT compatibility
-        self.transformer = self.language_model.model
 
         # Vision model (understanding branch)
         self.vision_model = NEOVisionModel(self.vis_cfg)
diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
index 53e1a47c239..f9ff8aaf470 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
@@ -16,7 +16,6 @@
 
 import torch
 import torch.nn as nn
-from cache_dit import ForwardPattern
 from transformers.cache_utils import DynamicCache
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -31,7 +30,6 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, SensenovaCachedAdapter
 
 
 @dataclass
@@ -608,14 +606,6 @@ def forward(
 
 
 class SenseNovaU1Model(nn.Module):
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "layers": ForwardPattern.Pattern_3,
-        },
-        has_separate_cfg=True,
-        cached_adapter_cls=SensenovaCachedAdapter,
-    )
-
     def __init__(self, config, quant_config=None, prefix: str = ""):
         super().__init__()
         self.config = config
diff --git a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py b/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py
index c7bc9b2d4f6..83a669d92aa 100644
--- a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py
+++ b/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py
@@ -9,14 +9,12 @@
 
 import torch
 import torch.nn as nn
-from cache_dit import ForwardPattern
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.layers.fourier import GaussianFourierProjection
@@ -466,12 +464,6 @@ class StableAudioDiTModel(nn.Module):
     - Output: [B, out_channels, L]
     """
 
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "transformer_blocks": ForwardPattern.Pattern_3,
-        },
-    )
-
     _repeated_blocks = ["StableAudioDiTBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     _hsdp_shard_conditions = [is_transformer_block_module]
diff --git a/vllm_omni/diffusion/quantization/__init__.py b/vllm_omni/diffusion/quantization/__init__.py
deleted file mode 100644
index 208f01a7cb5..00000000000
--- a/vllm_omni/diffusion/quantization/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/diffusion/quantization/hsdp_fp8.py b/vllm_omni/diffusion/quantization/hsdp_fp8.py
deleted file mode 100644
index c2ddd68ae9c..00000000000
--- a/vllm_omni/diffusion/quantization/hsdp_fp8.py
+++ /dev/null
@@ -1,112 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""HSDP/FSDP2 compatibility for online FP8 quantization.
-
-vllm's ``Fp8OnlineLinearMethod.process_weights_after_loading`` ends with
-``layer.weight = qweight.t()`` so that the Cutlass FP8 GEMM kernel sees its
-B operand as column-major ``[K, N]`` (the TN layout required by Hopper FP8
-``wgmma`` instructions). The resulting tensor is a non-contiguous transpose
-view of a ``(out_features, in_features)`` row-major storage.
-
-FSDP2 ``fully_shard`` rejects non-contiguous parameters because dim-0 sharding
-on a column-major view cannot be a contiguous memcpy. The two requirements are
-fundamentally at odds with the same physical buffer.
-
-This module reconciles them by separating the views: keep the parameter as
-the underlying ``(out, in)`` row-major contiguous storage (FSDP-friendly),
-and inject the equivalent ``.t()`` at the GEMM call site so the Cutlass
-kernel still receives a column-major B (zero-copy stride flip).
-
-The transpose is injected inside ``ScaledMMLinearKernel._get_layer_params``,
-which is the single place where ``apply_weights`` reads the weight tensor.
-The rest of ``apply_weights`` -- including its ``output_shape = w.shape[1]``
-computation and the eventual ``apply_scaled_mm(B=w, ...)`` call -- then
-operates on a tensor whose shape ``(K, N)`` and stride ``(1, K)`` match
-what the upstream FP8 kernel expects. No other upstream code is overridden,
-which keeps us robust against future changes inside ``apply_weights``.
-"""
-
-from __future__ import annotations
-
-import types
-
-from torch import nn
-from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
-
-logger = init_logger(__name__)
-
-
-def _build_transposed_get_layer_params(original_bound_method):
-    """Return a ``_get_layer_params`` replacement that transposes ``w``.
-
-    ``original_bound_method`` is the kernel instance's existing
-    ``_get_layer_params`` (already bound to ``self``); we close over it and
-    apply ``.t()`` to the returned weight, leaving scales untouched.
-    """
-
-    def _get_layer_params(self, layer):
-        w, w_s, x_s, x_s_ub = original_bound_method(layer)
-        return w.t(), w_s, x_s, x_s_ub
-
-    return _get_layer_params
-
-
-def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int:
-    """Make online-FP8 linear layers in ``model`` FSDP2-compatible.
-
-    For every layer whose ``quant_method`` is an :class:`Fp8OnlineLinearMethod`
-    and whose weight is currently a non-contiguous transpose view, this
-    function:
-
-    1. Replaces ``layer.weight`` with the underlying ``(out, in)`` row-major
-       contiguous storage so FSDP2 ``fully_shard`` accepts it.
-    2. Patches the per-layer GEMM kernel's bound ``_get_layer_params`` method
-       to return a ``.t()`` view of the weight, so ``apply_weights`` and the
-       downstream ``apply_scaled_mm`` continue to see a column-major
-       ``(in, out)`` B with zero copies.
-
-    Layers whose weight is already contiguous (e.g. Marlin FP8, offline-
-    quantized checkpoints) or that use a different quant method are skipped.
-
-    Returns:
-        Number of layers rewritten.
-    """
-    n_patched = 0
-    patched_kernel_ids: set[int] = set()
-    for module in model.modules():
-        qm = getattr(module, "quant_method", None)
-        if not isinstance(qm, Fp8OnlineLinearMethod):
-            continue
-
-        weight = getattr(module, "weight", None)
-        if weight is None or weight.is_contiguous():
-            continue
-
-        # ``weight`` here is ``qweight.t()`` (a non-contiguous (in, out) view
-        # of an (out, in) row-major qweight storage). ``weight.t()`` recovers
-        # that storage; ``.contiguous()`` is a zero-copy alias when the source
-        # is already row-major contiguous, but defensively materializes if
-        # upstream ever stacks another view.
-        contig = weight.data.t().contiguous()
-        new_param = nn.Parameter(contig, requires_grad=False)
-        module.weight = new_param
-
-        kernel = qm.fp8_linear
-        # Each linear layer is expected to have its own kernel instance, but
-        # guard against shared instances to avoid stacking multiple ``.t()``
-        # patches on the same object (which would compose into identity).
-        if id(kernel) not in patched_kernel_ids:
-            original_get = kernel._get_layer_params
-            kernel._get_layer_params = types.MethodType(_build_transposed_get_layer_params(original_get), kernel)
-            patched_kernel_ids.add(id(kernel))
-
-        n_patched += 1
-
-    if n_patched:
-        logger.info(
-            "Rewrote %d FP8 linear layer(s) into FSDP2-compatible storage; "
-            "transpose to column-major B is now applied at GEMM time.",
-            n_patched,
-        )
-    return n_patched
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index 5ab84702452..996e74b32c9 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -156,11 +156,6 @@
         "pipeline_internvla_a1",
         "InternVLAA1Pipeline",
     ),
-    "Gr00tN1d7Pipeline": (
-        "gr00t",
-        "pipeline_gr00t",
-        "Gr00tN1d7Pipeline",
-    ),
     "LongCatImageEditPipeline": (
         "longcat_image",
         "pipeline_longcat_image_edit",
diff --git a/vllm_omni/diffusion/utils/flow_matching.py b/vllm_omni/diffusion/utils/flow_matching.py
deleted file mode 100644
index e81d405ec23..00000000000
--- a/vllm_omni/diffusion/utils/flow_matching.py
+++ /dev/null
@@ -1,28 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import math
-
-import torch
-from torch import nn
-
-
-def swish(x: torch.Tensor) -> torch.Tensor:
-    return x * torch.sigmoid(x)
-
-
-class SinusoidalPositionalEncoding(nn.Module):
-    def __init__(self, embedding_dim: int) -> None:
-        super().__init__()
-        self.embedding_dim = embedding_dim
-
-    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
-        timesteps = timesteps.float()
-        B, T = timesteps.shape
-        device = timesteps.device
-
-        half_dim = self.embedding_dim // 2
-        exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * (math.log(10000.0) / half_dim)
-        freqs = timesteps.unsqueeze(-1) * exponent.exp()  # (B, T, half_dim)
-        enc = torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1)  # (B, T, embedding_dim)
-        return enc
diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
index 288973637b7..278f1df40f5 100644
--- a/vllm_omni/diffusion/worker/diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py
@@ -168,17 +168,8 @@ def get_memory_context():
         # Apply torch.compile if not in eager mode
         if not self.od_config.enforce_eager:
             if current_omni_platform.supports_torch_inductor():
-                if hasattr(self.pipeline, "setup_compile"):
-                    try:
-                        self.pipeline.setup_compile()
-                    except Exception as exc:
-                        logger.warning(
-                            "Model runner: setup_compile() failed (%s); running without compile.",
-                            exc,
-                        )
-                else:
-                    self._compile_transformer("transformer")
-                    self._compile_transformer("transformer_2")
+                self._compile_transformer("transformer")
+                self._compile_transformer("transformer_2")
             else:
                 logger.warning(
                     "Model runner: Platform %s does not support torch inductor, skipping torch.compile.",
@@ -315,15 +306,8 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
                 # num_inference_steps at all (i.e., just resets state and clears
                 # stale residuals).
                 num_inference_steps = req.sampling_params.num_inference_steps
-                if num_inference_steps is None and self.od_config.cache_backend in (
-                    "tea_cache",
-                    "step_cache",
-                ):
-                    # TeaCache refresh ignores the value; step_cache refresh is a
-                    # no-op (per-chunk state resets in the denoise loop). DreamZero often
-                    # leaves sampling_params.num_inference_steps unset and uses the
-                    # pipeline default instead.
-                    num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0
+                if self.od_config.cache_backend == "tea_cache" and num_inference_steps is None:
+                    num_inference_steps = 0
 
                 if num_inference_steps is not None:
                     self.cache_backend.refresh(self.pipeline, num_inference_steps)
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index 62cddd4a6b5..0b4e877180a 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -161,12 +161,6 @@ def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDi
         return vllm_config
 
 
-def _get_cumem_allocator_class() -> type:
-    from vllm.device_allocator.cumem import CuMemAllocator
-
-    return CuMemAllocator
-
-
 def _resolve_ir_op_priority(od_config: OmniDiffusionConfig, vllm_config: VllmConfig) -> Any:
     ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config)
     ir_op_priority_func = get_diffusion_ir_op_priority_func(od_config)
@@ -461,7 +455,8 @@ def sleep(self, level: int = 1) -> bool:
         Args:
             level: Sleep level. Level 1 offloads weights, level 2 also saves buffers.
         """
-        CuMemAllocator = _get_cumem_allocator_class()
+        from vllm.device_allocator.cumem import CuMemAllocator
+
         allocator = CuMemAllocator.get_instance()
 
         usage_before = allocator.get_current_usage()
@@ -513,7 +508,8 @@ def wake_up(self, tags: list[str] | None = None) -> bool:
             tags: List of memory pool tags to re-activate (e.g., ["weights"]
                   to match Level 1 sleep). If None, all pools are re-activated.
         """
-        CuMemAllocator = _get_cumem_allocator_class()
+        from vllm.device_allocator.cumem import CuMemAllocator
+
         allocator = CuMemAllocator.get_instance()
         allocator.wake_up(tags)
         current_omni_platform.synchronize()
@@ -635,7 +631,8 @@ def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
         if is_sleep_enabled:
             current_omni_platform.synchronize()
             gc.collect()
-            CuMemAllocator = _get_cumem_allocator_class()
+            from vllm.device_allocator.cumem import CuMemAllocator
+
             allocator = CuMemAllocator.get_instance()
             if tag == "weights":
                 assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index ac6140291f1..721e58c7b7c 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -520,7 +520,7 @@ def _ensure_stage_replica_mm_uuids(
             return
 
         mm_data = prompt.get("multi_modal_data")
-        if not isinstance(mm_data, dict) or not mm_data:
+        if not isinstance(mm_data, Mapping) or not mm_data:
             return
 
         from vllm.multimodal.hasher import MultiModalHasher
@@ -842,14 +842,6 @@ def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | Non
                 "mag_max_skip_steps": 5,
                 "mag_retention_ratio": 0.1,
             }
-        if cache_backend in ("step_cache"):
-            return {
-                "step_cache_dit_enabled": True,
-                "velocity_sim_thresholds": [0.95, 0.93],
-                "velocity_skip_countdowns": [4, 2],
-                "step_cache_dit_min_history": 2,
-                "step_cache_dit_max_history": 2,
-            }
         return None
 
     @staticmethod
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 984706ae007..905ff4eded8 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -512,7 +512,7 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu
             "--cache-backend",
             type=str,
             default="none",
-            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'",
+            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache'",
         )
         omni_config_group.add_argument(
             "--cache-config",
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 35b9dc20fac..627a1070524 100755
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -137,7 +137,7 @@
     build_stage_sampling_params_list,
     get_default_sampling_params_list,
 )
-from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER, FileStorageHandle
+from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
 from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url, decode_input_reference
@@ -474,9 +474,6 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         await omni_init_app_state(engine_client, app.state, args)
 
-        # Start background processes
-        await STORAGE_MANAGER.start()
-
         # Conditionally register profiler endpoints based on stage YAML configs
         stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
         if _should_enable_profiler_endpoints(stage_configs):
@@ -2693,11 +2690,12 @@ def _video_error_from_exception(exc: Exception) -> VideoError:
     )
 
 
-async def _cleanup_video(video_id: str):
+def _cleanup_video(video_id: str, output_path: str | None):
     try:
-        await STORAGE_MANAGER.delete(video_id)
-    except Exception:
-        logger.warning("Failed to cleanup partial video file '%s'", video_id)
+        if output_path is not None:
+            os.remove(output_path)
+    except OSError:
+        logger.warning("Failed to cleanup partial video file '%s' for id=%s", output_path, video_id)
 
 
 async def _run_video_generation_job(
@@ -2716,6 +2714,7 @@ async def _run_video_generation_job(
 
     await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS})
     started_at = time.perf_counter()
+    output_path = None
     try:
         video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes(
             request,
@@ -2725,27 +2724,27 @@ async def _run_video_generation_job(
             reference_audio=reference_audio,
         )
 
-        save_context = await STORAGE_MANAGER.save(video_bytes, video_id)
-        logger.info("Video request %s persisted %s output file.", video_id, save_context.key)
-
-        updated_fields = {
-            "status": VideoGenerationStatus.COMPLETED,
-            "progress": 100,
-            "file_name": f"{video_id}.{job.file_extension}",
-            "completed_at": save_context.created_at,
-            "inference_time_s": time.perf_counter() - started_at,
-            "stage_durations": stage_durations,
-            "peak_memory_mb": peak_memory_mb,
-            "action": action,
-        }
-        if save_context.expires_at is not None:
-            updated_fields["expires_at"] = save_context.expires_at
+        file_name = f"{video_id}.{job.file_extension}"
+        output_path = await STORAGE_MANAGER.save(video_bytes, file_name)
+        logger.info("Video request %s persisted %s output file.", video_id, output_path)
 
-        await VIDEO_STORE.update_fields(video_id, updated_fields)
+        await VIDEO_STORE.update_fields(
+            video_id,
+            {
+                "status": VideoGenerationStatus.COMPLETED,
+                "progress": 100,
+                "file_name": file_name,
+                "completed_at": int(time.time()),
+                "inference_time_s": time.perf_counter() - started_at,
+                "stage_durations": stage_durations,
+                "peak_memory_mb": peak_memory_mb,
+                "action": action,
+            },
+        )
     except (EngineGenerateError, EngineDeadError) as exc:
         logger.exception("Video generation failed (engine error) for id=%s", video_id)
 
-        await _cleanup_video(video_id)
+        _cleanup_video(video_id, output_path)
         await VIDEO_STORE.update_fields(
             video_id,
             {
@@ -2765,7 +2764,7 @@ async def _run_video_generation_job(
     except Exception as exc:
         logger.exception("Video generation failed for id=%s", video_id)
 
-        await _cleanup_video(video_id)
+        _cleanup_video(video_id, output_path)
         await VIDEO_STORE.update_fields(
             video_id,
             {
@@ -2776,7 +2775,7 @@ async def _run_video_generation_job(
             },
         )
     except asyncio.CancelledError:
-        await _cleanup_video(video_id)
+        _cleanup_video(video_id, output_path)
         await VIDEO_STORE.pop(video_id)
         raise
     finally:
@@ -3164,7 +3163,7 @@ async def delete_video(video_id: str) -> VideoDeleteResponse:
     elif job.status is VideoGenerationStatus.FAILED:
         if job.file_name is not None:
             try:
-                await STORAGE_MANAGER.delete(video_id)
+                await STORAGE_MANAGER.delete(job.file_name)
             except Exception:
                 logger.warning("Failed to delete stored artifact for failed video job %s", video_id, exc_info=True)
 
@@ -3174,13 +3173,13 @@ async def delete_video(video_id: str) -> VideoDeleteResponse:
     if job.file_name is None:
         raise HTTPException(status_code=409, detail="Video output not yet available. Please try again later.")
 
-    await STORAGE_MANAGER.delete(video_id)
+    await STORAGE_MANAGER.delete(job.file_name)
     await VIDEO_STORE.pop(video_id)
     return VideoDeleteResponse(id=job.id, deleted=True)
 
 
 @router.get("/v1/videos/{video_id}/content")
-async def download_video(video_id: str) -> Response:
+async def download_video(video_id: str) -> FileResponse:
     """Download the generated file for a completed video job.
 
     Args:
@@ -3203,19 +3202,11 @@ async def download_video(video_id: str) -> Response:
     if not job.file_name:
         raise HTTPException(status_code=404, detail="Generation is still in-progress")
 
-    file_handle = await STORAGE_MANAGER.open(video_id)
-    if file_handle is None:
+    full_path = STORAGE_MANAGER.get_full_file_path(job.file_name)
+    if not os.path.exists(full_path):
         raise HTTPException(status_code=404, detail="Generated video file not found on disk")
 
-    file_name = job.file_name or f"{video_id}.{job.file_extension}"
-    if isinstance(file_handle, FileStorageHandle):
-        response = FileResponse(path=file_handle.path, media_type=job.media_type, filename=file_name)
-    else:
-        raise HTTPException(
-            status_code=500, detail=f"Server generated an unsupported file storage handle for file id {video_id}"
-        )
-
-    return response
+    return FileResponse(path=full_path, media_type=job.media_type, filename=job.file_name)
 
 
 @profiler_router.post("/start_profile")
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index f867bfadb45..1f7a6b79dfc 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/vllm_omni/entrypoints/openai/serving_speech.py
@@ -672,6 +672,7 @@ def _detect_tts_model_type(self) -> str | None:
         if model_stage in _OMNIVOICE_TTS_MODEL_STAGES:
             return "omnivoice"
         if model_stage in _COVO_AUDIO_MODEL_STAGES:
+            model_arch = getattr(self._tts_stage.engine_args, "model_arch", None)
             if model_arch and "CovoAudio" in model_arch:
                 return "covo_audio"
         if model_stage in _VOXCPM2_TTS_MODEL_STAGES:
diff --git a/vllm_omni/entrypoints/openai/storage.py b/vllm_omni/entrypoints/openai/storage.py
index a36688956e1..52aef8948f1 100644
--- a/vllm_omni/entrypoints/openai/storage.py
+++ b/vllm_omni/entrypoints/openai/storage.py
@@ -1,78 +1,32 @@
 import asyncio
-import contextlib
 import os
-import stat
-import time
 from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
 from tempfile import NamedTemporaryFile
-from typing import Generic, Literal, TypeVar
 
-from vllm.logger import init_logger
 
-from vllm_omni.config.server_settings import SERVER_SETTINGS_CONFIG, FileBackend
-
-logger = init_logger(__name__)
-
-
-@dataclass
-class SaveContext:
-    key: str
-    created_at: int
-    expires_at: int | None = None
-
-
-@dataclass(frozen=True)
-class BaseStorageHandle:
-    kind: str
-
-
-@dataclass(frozen=True)
-class FileStorageHandle(BaseStorageHandle):
-    path: str
-    kind: Literal["path"] = field(default="path", init=False)
-
-
-K = TypeVar("K", bound=BaseStorageHandle, covariant=True)
-
-
-class StorageBaseManager(Generic[K], ABC):
+class StorageBaseManager(ABC):
     @abstractmethod
-    async def save(self, *args, **kwargs) -> SaveContext:
+    async def save(self, *args, **kwargs) -> str:
         pass
 
     @abstractmethod
     async def delete(self, *args, **kwargs) -> bool:
         pass
 
-    async def start(self, *args, **kwargs) -> None:
-        pass
-
-    async def stop(self, *args, **kwargs) -> None:
-        pass
-
-    @abstractmethod
-    async def open(self, storage_key: str) -> K | None:
-        pass
-
 
-class LocalStorageManager(StorageBaseManager[FileStorageHandle]):
-    def __init__(self, storage_path: str, max_concurrency: int = 4):
-        self.storage_path = os.path.realpath(storage_path)
+class LocalStorageManager(StorageBaseManager):
+    def __init__(self, storage_path: str | None = None):
+        if storage_path is None:
+            storage_path = os.getenv("VLLM_OMNI_STORAGE_PATH", "/tmp/storage")
+        self.storage_path = storage_path
         os.makedirs(self.storage_path, exist_ok=True)
 
+        max_concurrency = int(os.getenv("VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "4"))
         self._io_semaphore = asyncio.Semaphore(max(1, max_concurrency))
 
-    async def open(self, storage_key: str) -> FileStorageHandle | None:
-        local_file = self.get_full_file_path(storage_key)
-        if not os.path.exists(local_file):
-            return None
-        return FileStorageHandle(path=local_file)
-
-    def _save_sync(self, data: bytes, file_name: str) -> SaveContext:
+    def _save_sync(self, data: bytes, file_name: str) -> str:
         filename = self.get_full_file_path(file_name)
         tmp_name: str | None = None
-
         try:
             with NamedTemporaryFile("wb", dir=self.storage_path, delete=False) as f:
                 tmp_name = f.name
@@ -80,8 +34,7 @@ def _save_sync(self, data: bytes, file_name: str) -> SaveContext:
                 f.flush()
                 os.fsync(f.fileno())
             os.replace(tmp_name, filename)
-            response = SaveContext(key=file_name, created_at=int(time.time()))
-            return response
+            return filename
         except Exception:
             if tmp_name is not None:
                 try:
@@ -90,7 +43,7 @@ def _save_sync(self, data: bytes, file_name: str) -> SaveContext:
                     pass
             raise
 
-    async def save(self, data: bytes, file_name: str) -> SaveContext:
+    async def save(self, data: bytes, file_name: str) -> str:
         async with self._io_semaphore:
             return await asyncio.to_thread(self._save_sync, data, file_name)
 
@@ -109,102 +62,9 @@ def exists(self, file_name: str) -> bool:
         return os.path.exists(self.get_full_file_path(file_name))
 
     def get_full_file_path(self, file_name: str) -> str:
-        full = os.path.realpath(os.path.join(self.storage_path, file_name))
-
-        if os.path.commonpath([self.storage_path, full]) != self.storage_path:
-            raise ValueError(f"Illegal storage key: {file_name!r}")
-
-        return full
-
-
-class LocalStorageTTLManager(LocalStorageManager):
-    def __init__(self, ttl_seconds: int, sweep_interval_seconds: int, *args, **kwargs):
-        if ttl_seconds <= 0:
-            raise ValueError("`ttl_seconds` must be greater than or equal to 1.")
-        if sweep_interval_seconds <= 0:
-            raise ValueError("`sweep_interval_seconds` must be greater than or equal to 1.")
-
-        self._ttl_seconds = ttl_seconds
-        self._sweep_interval_seconds = sweep_interval_seconds
-        self._sweeper_task: asyncio.Task[None] | None = None
-
-        super().__init__(*args, **kwargs)
-
-    async def save(self, data: bytes, file_name: str) -> SaveContext:
-        result = await super().save(data, file_name)
-        result.expires_at = result.created_at + self._ttl_seconds
-        return result
-
-    async def _sweep_once(self, cutoff: float) -> int:
-        expired: list[str] = []
-        for entry in os.scandir(self.storage_path):
-            try:
-                if not entry.is_file(follow_symlinks=False):
-                    continue
-                if entry.stat(follow_symlinks=False).st_mtime < cutoff:
-                    expired.append(entry.path)
-            except FileNotFoundError:
-                logger.debug("TTL sweep skipped %s; file removed during scan", entry.path)
-            except OSError:
-                logger.warning("TTL sweep failed to inspect %s", entry.path, exc_info=True)
-
-        deleted = 0
-        for path in expired:
-            try:
-                async with self._io_semaphore:
-                    st = await asyncio.to_thread(os.stat, path, follow_symlinks=False)
-                    if not stat.S_ISREG(st.st_mode):
-                        continue
-                    if st.st_mtime >= cutoff:
-                        continue
-                    await asyncio.to_thread(os.remove, path)
-                    deleted += 1
-            except FileNotFoundError:
-                logger.debug("TTL sweep skipped %s; file already removed", path)
-            except OSError:
-                logger.warning("TTL sweep failed to delete expired file %s", path, exc_info=True)
-
-        return deleted
-
-    async def _sweep_loop(self) -> None:
-        while True:
-            try:
-                cutoff = time.time() - self._ttl_seconds
-                await self._sweep_once(cutoff)
-            except Exception:
-                logger.exception("TTL sweep failed for storage path %s", self.storage_path)
-            await asyncio.sleep(self._sweep_interval_seconds)
-
-    async def start(self) -> None:
-        if self._sweeper_task is None or self._sweeper_task.done():
-            self._sweeper_task = asyncio.create_task(self._sweep_loop())
-
-    async def stop(self) -> None:
-        if self._sweeper_task is None:
-            return
-        self._sweeper_task.cancel()
-        with contextlib.suppress(asyncio.CancelledError):
-            await self._sweeper_task
-        self._sweeper_task = None
-
-
-def get_storage_manager(storage_config: FileBackend) -> StorageBaseManager[FileStorageHandle]:
-    if isinstance(storage_config, FileBackend):
-        if storage_config.file_ttl is not None and storage_config.ttl_sweep_interval is not None:
-            manager = LocalStorageTTLManager(
-                storage_path=storage_config.path,
-                max_concurrency=storage_config.file_concurrency,
-                ttl_seconds=storage_config.file_ttl,
-                sweep_interval_seconds=storage_config.ttl_sweep_interval,
-            )
-        else:
-            manager = LocalStorageManager(
-                storage_path=storage_config.path, max_concurrency=storage_config.file_concurrency
-            )
-    else:
-        raise ValueError("No supported storage managers")
-
-    return manager
+        return os.path.join(self.storage_path, file_name)
 
 
-STORAGE_MANAGER = get_storage_manager(SERVER_SETTINGS_CONFIG.storage)
+# Should implement TTL management on local file storage
+# and recovery on restart in the *STORE's
+STORAGE_MANAGER = LocalStorageManager()
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py
index ad753a91ed2..b6d3d3e9f85 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/base.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py
@@ -194,3 +194,6 @@ def extra_body_params(cls) -> frozenset[str]:
     "SpeechServingContext",
     "TTSModelAdapter",
 ]
+
+# Type alias kept for readability in adapter signatures.
+NormalizeFn = Callable[["OpenAICreateSpeechRequest"], None]
diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py
index b6526458363..69f722ac824 100644
--- a/vllm_omni/entrypoints/openpi/connection.py
+++ b/vllm_omni/entrypoints/openpi/connection.py
@@ -8,13 +8,7 @@
     Infer    -> client sends msgpack(obs), server sends msgpack(ndarray)
     Reset    -> client sends msgpack({endpoint:reset}), server sends msgpack(status)
 
-NumPy values use msgpack-numpy marker mappings. Outbound payloads follow the
-openpi-client format so official OpenPI clients can decode action responses:
-    ndarray -> {__ndarray__: true, dtype, shape, data}
-    scalar  -> {__npgeneric__: true, dtype, data}
-
-Inbound payloads accept both the openpi-client markers above and the legacy
-vLLM-native markers:
+NumPy values use the msgpack-numpy marker mapping:
     ndarray -> {nd: true, type, kind, shape, data}
     scalar  -> {nd: false, type, kind, data}
 """
@@ -47,16 +41,18 @@ def _pack_numpy(obj: Any) -> Any:
         if not obj.flags.c_contiguous:
             obj = np.ascontiguousarray(obj)
         return {
-            b"__ndarray__": True,
+            b"nd": True,
             b"data": obj.tobytes(),
-            b"dtype": obj.dtype.str,
+            b"type": obj.dtype.str,
+            b"kind": obj.dtype.kind,
             b"shape": obj.shape,
         }
     if isinstance(obj, np.generic):
         return {
-            b"__npgeneric__": True,
-            b"data": obj.item(),
-            b"dtype": obj.dtype.str,
+            b"nd": False,
+            b"data": obj.tobytes(),
+            b"type": obj.dtype.str,
+            b"kind": obj.dtype.kind,
         }
     raise TypeError(f"Unsupported type: {type(obj)!r}")
 
@@ -71,67 +67,27 @@ def _decode_marker_text(value: Any) -> str:
     return str(value)
 
 
-def _is_truthy_marker(value: Any) -> bool:
-    return value is True or value == 1
-
-
-def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any:
-    nd = _mapping_get(obj, "nd", _MISSING)
-    dtype = _mapping_get(obj, "type", _MISSING)
-    kind = _mapping_get(obj, "kind", _MISSING)
-    data = _mapping_get(obj, "data", _MISSING)
-    if nd is _MISSING or dtype is _MISSING or kind is _MISSING or data is _MISSING:
-        return _MISSING
-
-    dtype_obj = np.dtype(_decode_marker_text(dtype))
-    kind_text = _decode_marker_text(kind)
-    if dtype_obj.kind != kind_text:
-        raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}")
-    if dtype_obj.kind in ("V", "O", "c"):
-        raise ValueError(f"Unsupported dtype: {dtype_obj}")
-
-    array = np.frombuffer(data, dtype=dtype_obj).copy()
-    if nd:
-        shape = _mapping_get(obj, "shape", _MISSING)
-        if shape is _MISSING:
-            raise ValueError("NumPy ndarray marker is missing shape")
-        return array.reshape(tuple(shape))
-    return array[0]
-
-
-def _decode_openpi_numpy_marker(obj: dict[Any, Any]) -> Any:
-    if _is_truthy_marker(_mapping_get(obj, "__ndarray__", False)):
-        data = _mapping_get(obj, "data", _MISSING)
-        dtype = _mapping_get(obj, "dtype", _MISSING)
-        shape = _mapping_get(obj, "shape", _MISSING)
-        if data is _MISSING or dtype is _MISSING or shape is _MISSING:
-            raise ValueError("OpenPI ndarray marker is missing required fields")
-
-        dtype_obj = np.dtype(_decode_marker_text(dtype))
-        if dtype_obj.kind in ("V", "O", "c"):
-            raise ValueError(f"Unsupported dtype: {dtype_obj}")
-        array = np.frombuffer(data, dtype=dtype_obj).copy()
-        return array.reshape(tuple(shape))
-
-    if _is_truthy_marker(_mapping_get(obj, "__npgeneric__", False)):
-        data = _mapping_get(obj, "data", _MISSING)
-        dtype = _mapping_get(obj, "dtype", _MISSING)
-        if data is _MISSING or dtype is _MISSING:
-            raise ValueError("OpenPI scalar marker is missing required fields")
-        dtype_obj = np.dtype(_decode_marker_text(dtype))
-        if dtype_obj.kind in ("V", "O", "c"):
-            raise ValueError(f"Unsupported dtype: {dtype_obj}")
-        return dtype_obj.type(data)
-
-    return _MISSING
-
-
 def _unpack_numpy(obj: Any) -> Any:
     if isinstance(obj, dict):
-        for decoder in (_decode_vllm_numpy_marker, _decode_openpi_numpy_marker):
-            decoded = decoder(obj)
-            if decoded is not _MISSING:
-                return decoded
+        nd = _mapping_get(obj, "nd", _MISSING)
+        dtype = _mapping_get(obj, "type", _MISSING)
+        kind = _mapping_get(obj, "kind", _MISSING)
+        data = _mapping_get(obj, "data", _MISSING)
+        if nd is not _MISSING and dtype is not _MISSING and kind is not _MISSING and data is not _MISSING:
+            dtype_obj = np.dtype(_decode_marker_text(dtype))
+            kind_text = _decode_marker_text(kind)
+            if dtype_obj.kind != kind_text:
+                raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}")
+            if dtype_obj.kind in ("V", "O", "c"):
+                raise ValueError(f"Unsupported dtype: {dtype_obj}")
+
+            array = np.frombuffer(data, dtype=dtype_obj).copy()
+            if nd:
+                shape = _mapping_get(obj, "shape", _MISSING)
+                if shape is _MISSING:
+                    raise ValueError("NumPy ndarray marker is missing shape")
+                return array.reshape(tuple(shape))
+            return array[0]
         return {key: _unpack_numpy(value) for key, value in obj.items()}
     if isinstance(obj, list):
         return [_unpack_numpy(value) for value in obj]
diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
index d32ea62f284..1a387cdcc39 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
@@ -479,6 +479,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             # Keep additional information synchronized for async_chunk updates.
             self.enable_update_additional_information = True
 
+            # Expose streaming parameters
+            self.token_overlap_len = self.code2wav.token_overlap_len
+            self.mel_overlap_len = self.code2wav.mel_overlap_len
+            self.mel_window = self.code2wav.mel_window
+            self.mel_cache_len = self.code2wav.mel_cache_len
+            self.source_cache_len = self.code2wav.source_cache_len
+            self.speech_window = self.code2wav.speech_window
+            self._stream_audio_cache_by_req: dict[str, torch.Tensor] = {}
             self._stream_audio_cache_lock = Lock()
             self._stream_vocoder_cache_by_req: dict[str, dict[str, torch.Tensor]] = {}
         else:
@@ -503,6 +511,27 @@ def _create_llm_vllm_config(self, parent_config: VllmConfig) -> VllmConfig:
         # Use parent's cache config - critical for PagedAttention to work correctly
         return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"])
 
+    @staticmethod
+    def _cross_fade_audio(audio: torch.Tensor, prev_tail: torch.Tensor) -> torch.Tensor:
+        """Blend previous chunk tail into current chunk head using a Hamming window.
+
+        This mirrors upstream CosyVoice's `fade_in_out(...)` semantics:
+        update the current head in-place using a 2*overlap window, then
+        concatenate the unchanged remainder.
+        """
+        if audio.numel() == 0 or prev_tail.numel() == 0:
+            return audio
+        overlap = min(int(audio.numel()), int(prev_tail.numel()))
+        if overlap <= 0:
+            return audio
+        window = torch.hamming_window(2 * overlap, periodic=False, dtype=audio.dtype, device=audio.device)
+        fade_in = window[:overlap]
+        fade_out = window[overlap:]
+        blended = audio[:overlap] * fade_in + prev_tail[-overlap:].to(device=audio.device, dtype=audio.dtype) * fade_out
+        if overlap == int(audio.numel()):
+            return blended
+        return torch.cat([blended, audio[overlap:]], dim=0)
+
     def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_finished: bool) -> torch.Tensor:
         """Pass-through stitching for async_chunk.
 
@@ -510,9 +539,11 @@ def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_f
         Applying an additional waveform-domain fade/cache step introduces either
         duplicated overlap (if no tail trim) or duration shrink (if tail trim).
         """
-        if req_id is not None and stream_finished and hasattr(self, "_stream_vocoder_cache_by_req"):
+        if req_id is not None and stream_finished and hasattr(self, "_stream_audio_cache_by_req"):
             with self._stream_audio_cache_lock:
-                self._stream_vocoder_cache_by_req.pop(req_id, None)
+                self._stream_audio_cache_by_req.pop(req_id, None)
+                if hasattr(self, "_stream_vocoder_cache_by_req"):
+                    self._stream_vocoder_cache_by_req.pop(req_id, None)
         return audio
 
     @staticmethod
diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
index 49e8e515b25..cb3228c13a7 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+from contextlib import nullcontext
+
 import numpy as np
 import torch
 import torch.nn as nn
@@ -173,18 +175,19 @@ def _forward_mel(
         prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32)
         prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32)
 
-        feat, _ = self.flow_model.inference(
-            token=token,
-            token_len=token_len,
-            prompt_token=prompt_token,
-            prompt_token_len=prompt_token_len,
-            prompt_feat=prompt_feat,
-            prompt_feat_len=prompt_feat_len,
-            embedding=embedding,
-            streaming=streaming,
-            finalize=finalize,
-            n_timesteps=n_timesteps,
-        )
+        with nullcontext():
+            feat, _ = self.flow_model.inference(
+                token=token,
+                token_len=token_len,
+                prompt_token=prompt_token,
+                prompt_token_len=prompt_token_len,
+                prompt_feat=prompt_feat,
+                prompt_feat_len=prompt_feat_len,
+                embedding=embedding,
+                streaming=streaming,
+                finalize=finalize,
+                n_timesteps=n_timesteps,
+            )
 
         trim_mel = max(0, int(token_offset_tokens)) * int(self.token_mel_ratio)
         if trim_mel > 0:
@@ -192,6 +195,28 @@ def _forward_mel(
 
         return feat
 
+    @staticmethod
+    def _fade_speech(
+        speech: torch.Tensor,
+        prev_speech: torch.Tensor,
+    ) -> torch.Tensor:
+        """Blend previous speech tail into current speech head."""
+        if speech.numel() == 0 or prev_speech.numel() == 0:
+            return speech
+        overlap = min(int(speech.shape[-1]), int(prev_speech.shape[-1]))
+        if overlap <= 0:
+            return speech
+        window = torch.hamming_window(2 * overlap, periodic=False, dtype=speech.dtype, device=speech.device)
+        fade_in = window[:overlap].view(1, -1)
+        fade_out = window[overlap:].view(1, -1)
+        blended_head = (
+            speech[:, :overlap] * fade_in
+            + prev_speech[:, -overlap:].to(device=speech.device, dtype=speech.dtype) * fade_out
+        )
+        if overlap == int(speech.shape[-1]):
+            return blended_head
+        return torch.cat([blended_head, speech[:, overlap:]], dim=-1)
+
     @torch.inference_mode()
     def forward_streaming(
         self,
@@ -213,16 +238,17 @@ def forward_streaming(
         suffix. That preserves causal look-right handling without double
         trimming or duplicated overlap at chunk boundaries.
         """
-        feat = self._forward_mel(
-            token=token,
-            prompt_token=prompt_token,
-            prompt_feat=prompt_feat,
-            embedding=embedding,
-            n_timesteps=n_timesteps,
-            token_offset_tokens=token_offset_tokens,
-            streaming=True,
-            finalize=finalize,
-        )
+        with nullcontext():
+            feat = self._forward_mel(
+                token=token,
+                prompt_token=prompt_token,
+                prompt_feat=prompt_feat,
+                embedding=embedding,
+                n_timesteps=n_timesteps,
+                token_offset_tokens=token_offset_tokens,
+                streaming=True,
+                finalize=finalize,
+            )
         hift_weight = self.hift.m_source.l_linear.weight
         chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
 
@@ -242,7 +268,8 @@ def forward_streaming(
         if tts_mel.shape[-1] == 0:
             tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype)
         else:
-            tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
+            with nullcontext():
+                tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
 
         tts_speech = tts_speech.reshape(tts_speech.shape[0], -1)
         speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1])))
diff --git a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
index 32db8ef5b3b..ed35c93bd13 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
@@ -4,10 +4,8 @@
 
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
 Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
-  * ``sync_process_input_func`` (``text2flow_token_only``) runs when
-    ``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
-    from the emitted speech tokens. ``text2flow`` is the
-    ``custom_process_input_func`` variant.
+  * ``sync_process_input_func`` runs when ``deploy.async_chunk=false``:
+    stage 1 builds full-sequence flow input via ``text2flow``.
   * ``async_chunk_process_next_stage_input_func`` runs when
     ``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
     through the shared-memory connector.
diff --git a/vllm_omni/model_executor/models/cosyvoice3/utils.py b/vllm_omni/model_executor/models/cosyvoice3/utils.py
index d2c3c686574..5827e6bbbc2 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/utils.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/utils.py
@@ -13,6 +13,8 @@
 
 logger = logging.getLogger(__name__)
 
+IGNORE_ID = -1
+
 
 def dynamic_range_compression_torch(x, c=1, clip_val=1e-5):
     return torch.log(torch.clamp(x, min=clip_val) * c)
@@ -180,6 +182,26 @@ def log_mel_spectrogram(
     return log_spec
 
 
+def extract_speech_token(prompt_wav, speech_tokenizer_session, device):
+    speech = load_wav(prompt_wav, 16000)
+    assert speech.shape[1] / 16000 <= 30, "do not support extract speech token for audio longer than 30s"
+    feat = log_mel_spectrogram(speech, n_mels=128)
+    speech_token = (
+        speech_tokenizer_session.run(
+            None,
+            {
+                speech_tokenizer_session.get_inputs()[0].name: feat.detach().cpu().numpy(),
+                speech_tokenizer_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32),
+            },
+        )[0]
+        .flatten()
+        .tolist()
+    )
+    speech_token = torch.tensor([speech_token], dtype=torch.int32).to(device)
+    speech_token_len = torch.tensor([speech_token.shape[1]], dtype=torch.int32).to(device)
+    return speech_token, speech_token_len
+
+
 def extract_spk_embedding(prompt_wav, campplus_session, device):
     speech = load_wav(prompt_wav, 16000)
     feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000)
diff --git a/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py b/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py
index b9b06857d07..8153c2f385e 100644
--- a/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py
+++ b/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py
@@ -67,5 +67,6 @@ def __init__(self, **kwargs):
         self.inference = {
             "s_steps": 10,
             "cfg_alpha": 1.0,
+            "dynamic_the": False,
             "rescale_logits": False,
         }
diff --git a/vllm_omni/model_executor/models/covo_audio/token2wav.py b/vllm_omni/model_executor/models/covo_audio/token2wav.py
index 8a96e1f398b..4a6ce4b7e3a 100644
--- a/vllm_omni/model_executor/models/covo_audio/token2wav.py
+++ b/vllm_omni/model_executor/models/covo_audio/token2wav.py
@@ -598,7 +598,7 @@ def inference(
         latent_cond = torch.zeros(b, tgt_lens, self.target_dim).to(device)
         if prefix_target is not None:
             latent_cond[:, : prefix_target.size(1), :] = prefix_target
-        sample, _ = self.sample(
+        sample, trajectory = self.sample(
             tokens=token, audio=latent_cond, steps=s_steps, alpha=cfg_alpha, g_cond=ge, rescale_logits=rescale_logits
         )
         return sample
@@ -612,8 +612,9 @@ def solver(t, z):
                 output = torch.cat([tokens, audio, z], dim=-1)
                 return self.vectorfield_forward(inputs=output, times=t.unsqueeze(0), self_attn_mask=None, g_cond=g_cond)
             tokens_empty = torch.zeros(*audio.shape[:2], self.model_dim, device=tokens.device, dtype=tokens.dtype)
+            audio_empty = audio
             tokens_t = torch.cat([tokens_empty, tokens], dim=0)
-            audio_t = torch.cat([audio, audio], dim=0)
+            audio_t = torch.cat([audio_empty, audio], dim=0)
             audio_noizy_t = torch.cat([z, z], dim=0)
             t_t = torch.stack([t, t], dim=0)
             c = g_cond
diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
index 9627dc7ce13..c6e9064f133 100644
--- a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
+++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
@@ -348,16 +348,18 @@ def forward(
 
             if req_ids.numel() < 1:
                 continue
-            n = req_ids.numel()
-            if n % q != 0:
-                logger.warning(
-                    "DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.",
-                    n,
-                    q,
-                )
+            flat = req_ids
+            n = flat.numel()
+            if n == 0 or n % q != 0:
+                if n > 0:
+                    logger.warning(
+                        "DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.",
+                        n,
+                        q,
+                    )
                 continue
             frames = n // q
-            codes_qf = req_ids.reshape(q, frames)
+            codes_qf = flat.reshape(q, frames)
             parsed_ctx_frames[i] = ctx_frames
             parsed_total_frames[i] = frames
             valid_codes_qf.append(codes_qf)
diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
index eb056f16a10..46ceee49835 100644
--- a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
+++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
@@ -191,11 +191,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class _FastARDecoderLayer(nn.Module):
-    """Transformer decoder layer for Fast AR.
-
-    The batched ``forward`` runs SDPA over the full sequence and is cache-free,
-    while ``forward_one`` decodes a single step using a per-call KV cache.
-    """
+    """Transformer decoder layer for Fast AR (SDPA, no KV cache)."""
 
     def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None:
         super().__init__()
@@ -523,7 +519,7 @@ def forward(
         embed_buf[:bsz, 0, :] = projected
 
         # Position 1: embedding of semantic code.
-        code_embed = self.fast_embeddings(semantic_code)
+        code_embed = self.fast_embeddings(semantic_code.clamp(min=0))
         embed_buf[:bsz, 1, :] = code_embed
 
         use_sampling = do_sample and temperature > 0
@@ -546,7 +542,8 @@ def forward(
             logits = self.fast_output(self.fast_norm(hidden_out))
 
             # Residual codebooks (step >= 1) only have 1024 entries.
-            logits = logits[:, :residual_codebook_size]
+            if step >= 1:
+                logits = logits[:, :residual_codebook_size]
 
             if use_sampling:
                 scaled = logits * inv_temperature
diff --git a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py
index 3918a6f9462..8b8d8559ead 100644
--- a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py
+++ b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py
@@ -132,3 +132,13 @@ def estimate_fish_voice_clone_prompt_len_from_normalized(
         [0] * semantic_len,
     )
     return len(prompt_ids)
+
+
+def estimate_fish_voice_clone_prompt_len(tokenizer: Any, text: str, ref_text: str, semantic_len: int) -> int:
+    normalized_text, normalized_ref_text = normalize_fish_voice_clone_texts(text, ref_text)
+    return estimate_fish_voice_clone_prompt_len_from_normalized(
+        tokenizer,
+        normalized_text,
+        normalized_ref_text,
+        semantic_len,
+    )
diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py
index 666998e9dda..5cda6ec45f5 100644
--- a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py
+++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py
@@ -472,6 +472,10 @@ def __init__(
         self.num_heads = heads
         self.spkr_emb_adaln = spkr_emb_adaln
 
+        # Frame rates for block-causal mask token→mel conversion
+        self._input_frame_rate: float = 25.0
+        self._mel_framerate: float = 50.0
+
         # Speaker embedding dimension for AdaLN
         if spkr_emb_adaln:
             additional_dim = spkr_dim + 256 if use_wavlm_emb else spkr_dim
diff --git a/vllm_omni/model_executor/models/gr00t/__init__.py b/vllm_omni/model_executor/models/gr00t/__init__.py
deleted file mode 100644
index 208f01a7cb5..00000000000
--- a/vllm_omni/model_executor/models/gr00t/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/model_executor/models/gr00t/pipeline.py b/vllm_omni/model_executor/models/gr00t/pipeline.py
deleted file mode 100644
index 38a76f29769..00000000000
--- a/vllm_omni/model_executor/models/gr00t/pipeline.py
+++ /dev/null
@@ -1,26 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""GR00T N1.7 single-stage policy topology."""
-
-from vllm_omni.config.stage_config import (
-    PipelineConfig,
-    StageExecutionType,
-    StagePipelineConfig,
-)
-
-GR00T_N1D7_PIPELINE = PipelineConfig(
-    model_type="Gr00tN1d7",
-    model_arch="Gr00tN1d7Pipeline",
-    hf_architectures=("Gr00tN1d7",),
-    stages=(
-        StagePipelineConfig(
-            stage_id=0,
-            model_stage="diffusion",
-            execution_type=StageExecutionType.DIFFUSION,
-            input_sources=(),
-            final_output=True,
-            final_output_type="actions",
-            model_arch="Gr00tN1d7Pipeline",
-        ),
-    ),
-)
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py
index b6eef8dc4a5..dc5eb12b475 100644
--- a/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py
@@ -155,3 +155,15 @@ def __init__(
     @property
     def num_real_codes(self) -> int:
         return self.audio_stream_bos_id  # i.e. 1024 for the boson-ai checkpoint
+
+    # The full vocab size of each per-codebook output head.
+    @property
+    def codebook_output_size(self) -> int:
+        return self.codebook_size  # i.e. 1026 for the boson-ai checkpoint
+
+    # Canonical MusicGen-style delay pattern used by
+    # `HiggsAudioV2DelayPatternLogitsProcessor`: codebook k starts emitting
+    # real codes only after k frames. Kept as a function so subclasses can
+    # override if upstream ever publishes a different convention.
+    def default_delay_pattern(self) -> list[int]:
+        return list(range(self.num_codebooks))
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
index be08163e4d2..df636a4bbe2 100644
--- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
@@ -821,12 +821,12 @@ def forward(
         except Exception:
             self._last_step_query_start_loc = None
 
+        # Voice-clone ref-audio substitution: at the prompt-side audio
+        # placeholder span we overwrite the text-embed lookups with the
+        # multi-codebook embedding of the encoded reference clip carried in
+        # ``model_intermediate_buffer``. No-op for plain-text requests and
+        # for decode steps (which fall through to the audio-feedback path).
         if input_ids is not None and inputs_embeds is None:
-            # Voice-clone ref-audio substitution: at the prompt-side audio
-            # placeholder span we overwrite the text-embed lookups with the
-            # multi-codebook embedding of the encoded reference clip carried in
-            # ``model_intermediate_buffer``. No-op for plain-text requests and
-            # for decode steps (which fall through to the audio-feedback path).
             info_dicts = kwargs.get("model_intermediate_buffer")
             if info_dicts is None:
                 info_dicts = kwargs.get("runtime_additional_information")
@@ -836,14 +836,15 @@ def forward(
                 info_dicts,
             )
 
-            # Audio-feedback substitution: vLLM's AR runner calls forward()
-            # with ``inputs_embeds=None`` and lets the model embed its own
-            # input_ids inline. ``embed_input_ids`` is only invoked via the
-            # supports_mm_inputs/prompt_embeds branches, neither of which fires
-            # for the talker. So we apply the substitution directly here: at
-            # decode positions where input_ids == audio_token_id, swap the text
-            # embedding with embed_audio_codes(last audio frame stashed in
-            # ``self._audio_state[req_id]`` by :meth:`sample`).
+        # Audio-feedback substitution: vLLM's AR runner calls forward()
+        # with ``inputs_embeds=None`` and lets the model embed its own
+        # input_ids inline. ``embed_input_ids`` is only invoked via the
+        # supports_mm_inputs/prompt_embeds branches, neither of which fires
+        # for the talker. So we apply the substitution directly here: at
+        # decode positions where input_ids == audio_token_id, swap the text
+        # embedding with embed_audio_codes(last audio frame stashed in
+        # ``self._audio_state[req_id]`` by :meth:`sample`).
+        if input_ids is not None and inputs_embeds is None:
             hidden_states = self._maybe_apply_audio_feedback(hidden_states, input_ids)
 
         audio_mask = self.audio_token_mask(input_ids) if input_ids is not None else None
@@ -1505,7 +1506,7 @@ def _sample_audio_codes_upstream(
             sampled[~valid] = torch.argmax(logits[~valid].float(), dim=-1)
         return sampled
 
-    def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> None:
+    def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> int:
         """Mask non-audio tokens at audio-mode positions, in-place on ``logits``.
 
         Heuristic: per-request, find the last token seen so far (last of
@@ -1523,7 +1524,7 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -
         the stock sampler.
         """
         if logits is None or logits.ndim != 2:
-            return
+            return 0
         audio_bos = int(self.config.audio_bos_token_id)
         audio_id = int(self.config.audio_token_id)
         # (PyTorch align, corrected): upstream NEVER lets the model
@@ -1541,6 +1542,7 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -
         prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None)
         output_ids = getattr(sampling_metadata, "output_token_ids", None)
         num_rows = int(logits.shape[0])
+        biased = 0
         # Fallback "previous token" source: read each request's LAST input
         # position from ``self._last_step_input_ids`` using ``query_start_loc``
         # from the forward context. The previous version of this code took
@@ -1617,6 +1619,7 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -
                 mask[audio_eos_vocab] = row[audio_eos_vocab]
                 logits[i].copy_(mask)
                 slot_state["should_terminate"] = False
+                biased += 1
                 continue
             # Always force ``audio_token_id`` — upstream mirror.
             allowed: set[int] = {audio_id, *allowed_extra}
@@ -1626,7 +1629,9 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -
                 if 0 <= tok < row.shape[-1]:
                     mask[tok] = row[tok]
             logits[i].copy_(mask)
+            biased += 1
         self._last_first_audio_after_bos = first_after_bos
+        return biased
 
     has_postprocess: bool = True
 
diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py b/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py
new file mode 100644
index 00000000000..2a21aafcd39
--- /dev/null
+++ b/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py
@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bounded LRU cache for Higgs Audio v3 reference-audio codec encoding.
+
+Voice-clone requests that re-send the same ref_audio are otherwise re-encoded
+through the DAC encoder on every call. The cache keys by a stable content
+hash (size + sample rate + head/mid/tail sentinel bytes) so the same audio
+served via different URLs or repeated base64 payloads also hits warm.
+Tensors are cloned on get/put so callers and the cache cannot alias storage.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import threading
+from collections import OrderedDict
+from collections.abc import Callable
+
+import numpy as np
+import torch
+
+__all__ = ["cached_encode_reference_audio", "cache_stats", "clear_cache"]
+
+_DEFAULT_MAX_ENTRIES = 128
+_SENTINEL_BYTES = 4096
+
+_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
+_cache_lock = threading.Lock()
+_hits = 0
+_misses = 0
+
+
+def _content_hash(wav: np.ndarray, sr: int) -> str:
+    """Stable content hash. Head/mid/tail sentinel keeps cost O(1)."""
+    raw = wav.tobytes()
+    n = len(raw)
+    h = hashlib.sha256()
+    h.update(n.to_bytes(8, "little"))
+    h.update(int(sr).to_bytes(4, "little"))
+    if n == 0:
+        return h.hexdigest()
+    if n <= _SENTINEL_BYTES * 3:
+        h.update(raw)
+        return h.hexdigest()
+    h.update(raw[:_SENTINEL_BYTES])
+    mid = (n - _SENTINEL_BYTES) // 2
+    h.update(raw[mid : mid + _SENTINEL_BYTES])
+    h.update(raw[-_SENTINEL_BYTES:])
+    return h.hexdigest()
+
+
+def cached_encode_reference_audio(
+    wav: np.ndarray,
+    sr: int,
+    encode_fn: Callable[[np.ndarray, int], torch.Tensor],
+    max_entries: int = _DEFAULT_MAX_ENTRIES,
+) -> torch.Tensor:
+    """LRU-cached wrapper around ``encode_reference_audio``.
+
+    ``encode_fn`` is the actual codec encoder; passed in so this module
+    stays decoupled from the tokenizer import (avoids circular imports
+    during stage-input-processor discovery).
+
+    Returns a CPU tensor clone — never the cache's interior storage.
+    """
+    global _hits, _misses
+    key = _content_hash(wav, sr)
+    with _cache_lock:
+        cached = _cache.get(key)
+        if cached is not None:
+            _hits += 1
+            _cache.move_to_end(key)
+            return cached.clone()
+
+    # Miss: encode outside the lock so concurrent misses on different
+    # refs don't serialize on this single mutex.
+    codes = encode_fn(wav, sr)
+    codes_cpu = codes.detach().to("cpu").contiguous()
+    with _cache_lock:
+        if key in _cache:
+            # Concurrent miss winner already stored it; refresh LRU position
+            # and return that copy to avoid divergent tensors per request.
+            _cache.move_to_end(key)
+            _misses += 1  # we still paid the encode cost
+            return _cache[key].clone()
+        _cache[key] = codes_cpu
+        if len(_cache) > max_entries:
+            _cache.popitem(last=False)
+        _misses += 1
+    return codes_cpu.clone()
+
+
+def cache_stats() -> dict:
+    """Return (hits, misses, size, hit_rate) — for /metrics-style reporting."""
+    with _cache_lock:
+        total = _hits + _misses
+        return {
+            "hits": _hits,
+            "misses": _misses,
+            "size": len(_cache),
+            "hit_rate": (_hits / total) if total > 0 else 0.0,
+        }
+
+
+def clear_cache() -> None:
+    """Drop all cached entries. Used by tests."""
+    global _hits, _misses
+    with _cache_lock:
+        _cache.clear()
+        _hits = 0
+        _misses = 0
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/projectors.py b/vllm_omni/model_executor/models/ming_flash_omni/projectors.py
index 03527ba6327..42e53d1c635 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/projectors.py
+++ b/vllm_omni/model_executor/models/ming_flash_omni/projectors.py
@@ -108,6 +108,20 @@ def __init__(
             layers.append(nn.Linear(llm_dim, llm_dim))
         self.proj = nn.Sequential(*layers)
 
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Project audio features with temporal downsampling.
+
+        Args:
+            x: [B, T, audio_dim] audio encoder output (channel-last).
+
+        Returns:
+            [B, T', llm_dim] projected features (channel-last),
+            where T' = (T - ds_kernel_size + 2*(ds_kernel_size//2)) // ds_stride + 1.
+        """
+        # Conv1d expects [B, C, T], so transpose input
+        x = x.transpose(-1, -2)  # [B, audio_dim, T]
+        return self.proj(x)
+
     def forward_packed(
         self,
         packed: torch.Tensor,
@@ -141,6 +155,19 @@ def forward_packed(
         packed_proj = mlp(packed_proj)
         return packed_proj, proj_lens
 
+    def compute_output_length(self, input_length: torch.Tensor) -> torch.Tensor:
+        """Compute output sequence length after Conv1d downsampling.
+
+        Args:
+            input_length: Original mel spectrogram lengths.
+
+        Returns:
+            Output lengths after both convolutions.
+        """
+        length = (input_length - 3 + 2 * 1) // 2 + 1
+        length = (length - self.ds_kernel_size + 2 * (self.ds_kernel_size // 2)) // self.ds_stride + 1
+        return length
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py b/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py
index e0e8e73e669..436b92f1428 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py
+++ b/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py
@@ -15,6 +15,19 @@
 
 logger = init_logger(__name__)
 
+# Ported from
+# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/toolkit.py
+_TOKENIZE_PATTERN = re.compile(r"(?:[a-zA-Z]\.)+|[a-zA-Z]+(?:['\-][a-zA-Z]+)*|\d+(?:\.\d+)?|[\u4e00-\u9fff]|\s+|\S")
+
+
+def tokenize_mixed_text(text: str) -> list[str]:
+    return re.findall(_TOKENIZE_PATTERN, text)
+
+
+def tokenize_mixed_text_iterator(text: str):
+    for match in _TOKENIZE_PATTERN.finditer(text):
+        yield match.group(0)
+
 
 # Ported from
 # https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/text_segment_cut.py
@@ -205,6 +218,84 @@ def cut_text_by_semantic_length(
     return [f.replace(_DOT_PLACEHOLDER, ".") for f in result_fragments]
 
 
+# Streaming sentence boundary detection
+
+_RE_CJK = re.compile(r"[\u4e00-\u9fff]")
+_RE_DIGIT_LAST = re.compile(r"[0-9]")
+
+
+# Left for reference for now
+# Alternative to `cut_text_by_semantic_length`
+def detect_sentence_boundaries(
+    text: str,
+    max_length: int = 50,
+) -> list[str]:
+    """Accumulate tokens and flush at sentence boundaries.
+
+    Ported from the streaming sentence detection loop from the Ming repo
+    TTS branch, but operates on the full text since we have it available upfront.
+    """
+    sentences: list[str] = []
+    streaming_text: list[str] = []
+    count = 0
+
+    for ele in tokenize_mixed_text_iterator(text):
+        if len(ele) == 0:
+            continue
+
+        should_process = False
+        min_tokens = 12 if count == 0 else 8
+
+        if ele[-1] in "！？。，!?":
+            if len(streaming_text) >= min_tokens:
+                should_process = True
+            streaming_text.append(ele)
+
+        elif ele[-1] == ".":
+            if (
+                len(streaming_text) >= min_tokens
+                and streaming_text
+                and not _RE_DIGIT_LAST.search(streaming_text[-1][-1])
+            ):
+                should_process = True
+            streaming_text.append(ele)
+
+        elif ele[-1] == "\n":
+            if streaming_text:
+                joined = "".join(streaming_text)
+                if _RE_CJK.search(joined):
+                    if _RE_CJK.search(streaming_text[-1][-1]):
+                        ele = "，"
+                        streaming_text.append(ele)
+                else:
+                    if len(ele) > 1 and re.search(r"[a-zA-Z]", ele[-2]):
+                        ele = ele[:-1] + "."
+                    else:
+                        ele = ele[:-1]
+                    streaming_text.append(ele)
+
+            if len(streaming_text) >= min_tokens:
+                should_process = True
+        else:
+            streaming_text.append(ele)
+            continue
+
+        if should_process:
+            joined = "".join(streaming_text)
+            fragments = cut_text_by_semantic_length(joined, max_length)
+            sentences.extend(fragments)
+            streaming_text = []
+            count += 1
+
+    # Flush remaining
+    if streaming_text and re.search(r"[a-zA-Z\u4e00-\u9fff1-9]", "".join(streaming_text)):
+        joined = "".join(streaming_text)
+        fragments = cut_text_by_semantic_length(joined, max_length)
+        sentences.extend(fragments)
+
+    return sentences
+
+
 # number normalization for English. Ported from
 # https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/number_en.py
 
diff --git a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
index 8ef8a3a0e36..b42c4fdca62 100644
--- a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
+++ b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
@@ -30,6 +30,7 @@
     KEY_MAX_DECODE_STEPS,
     KEY_MIN_DECODE_STEPS,
     KEY_NEXT_EMBEDS,
+    KEY_PROMPT_LATENT_TAIL,
     KEY_PROMPT_LATENTS,
     KEY_REQUEST_ID,
     KEY_SIGMA,
@@ -183,6 +184,7 @@ def validate(self) -> None:
     "KEY_MAX_DECODE_STEPS",
     "KEY_MIN_DECODE_STEPS",
     "KEY_NEXT_EMBEDS",
+    "KEY_PROMPT_LATENT_TAIL",
     "KEY_PROMPT_LATENTS",
     "KEY_REQUEST_ID",
     "KEY_SIGMA",
diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py
index 4b3f527da6c..2f59e209401 100644
--- a/vllm_omni/model_executor/models/ming_tts/constants.py
+++ b/vllm_omni/model_executor/models/ming_tts/constants.py
@@ -55,6 +55,7 @@
 KEY_LAST_STOP_PROB = "ming_last_stop_prob"
 KEY_NEXT_EMBEDS = "ming_next_embeds"
 KEY_PROMPT_LATENTS = "ming_prompt_latents"
+KEY_PROMPT_LATENT_TAIL = "ming_prompt_latent_tail"
 KEY_SPEAKER_EMBEDDING = "ming_speaker_embedding"
 KEY_REQUEST_ID = "ming_request_id"
 KEY_CHUNK_ID = "ming_chunk_id"
diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts.py b/vllm_omni/model_executor/models/ming_tts/ming_tts.py
index 4140b403be5..667aded27d7 100644
--- a/vllm_omni/model_executor/models/ming_tts/ming_tts.py
+++ b/vllm_omni/model_executor/models/ming_tts/ming_tts.py
@@ -16,6 +16,7 @@
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 
 from .audio_prep import (
+    _coerce_prompt_latents,
     _find_audio_placeholder_positions,
     _find_speaker_placeholder_positions,
     _initial_history,
@@ -31,6 +32,7 @@
     KEY_MAX_DECODE_STEPS,
     KEY_MIN_DECODE_STEPS,
     KEY_NEXT_EMBEDS,
+    KEY_PROMPT_LATENT_TAIL,
     KEY_REQUEST_ID,
     KEY_SIGMA,
     KEY_SPEAKER_EMBEDDING,
@@ -173,6 +175,7 @@ def _prefill_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tenso
             dtype=torch.float32,
         )
         update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous()
+        update[KEY_PROMPT_LATENT_TAIL] = update[KEY_LATENT_HISTORY]
 
         speaker_embedding = info_dict.get(KEY_SPEAKER_EMBEDDING, info_dict.get("speaker_embedding"))
         speaker_embeddings = None
@@ -259,6 +262,7 @@ def _copy_runtime_controls(update: dict[str, Any], info_dict: dict[str, Any]) ->
 __all__ = [
     "MingTTSForConditionalGeneration",
     "_ModelSampleAdapter",
+    "_coerce_prompt_latents",
     "_find_audio_placeholder_positions",
     "_find_speaker_placeholder_positions",
     "_initial_history",
diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
index e7e59631f76..92daad967e0 100644
--- a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
+++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
@@ -3,6 +3,7 @@
 # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py
 from __future__ import annotations
 
+import warnings
 from collections.abc import Iterable
 from typing import Any
 
@@ -345,10 +346,10 @@ def _decode_one_step(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         params_dict = dict(self.named_parameters(remove_duplicate=False))
         loaded_params = AutoWeightsLoader(self).load_weights(weights, mapper=self.hf_to_vllm_mapper)
-        _warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.")
-        _warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.")
-        _warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.")
-        _warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.")
+        _warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.", fatal=True)
+        _warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.", fatal=True)
+        _warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.", fatal=True)
+        _warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.", fatal=True)
         return loaded_params
 
 
@@ -382,6 +383,7 @@ def _warn_missing_prefix(
     params_dict: dict[str, nn.Parameter],
     loaded_params: set[str],
     prefix: str,
+    fatal: bool = False,
 ) -> None:
     missing = {key for key in params_dict if key.startswith(prefix)} - loaded_params
     if not missing:
@@ -390,7 +392,9 @@ def _warn_missing_prefix(
         f"MingLLMModel: {len(missing)} {module_name} params not loaded "
         f"(prefix={prefix}). First few: {sorted(missing)[:5]}"
     )
-    raise RuntimeError(msg)
+    if fatal:
+        raise RuntimeError(msg)
+    warnings.warn(msg, stacklevel=3)
 
 
 __all__ = [
diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py
index 89c52485a39..a2cf75b07b1 100644
--- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py
+++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py
@@ -78,7 +78,7 @@ def num_quantizers(self) -> int:
 
     @property
     def codebook_size(self) -> int:
-        return int(self.quantizer_kwargs.get("codebook_size", 1024))
+        return int(self.quantizer_kwargs.get("codebook_size", 4096))
 
     @property
     def frame_rate(self) -> float:
diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
index d32af7c9020..e7e100a1893 100644
--- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
+++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
@@ -31,6 +31,8 @@
 
 logger = init_logger(__name__)
 
+_AUDIO_PAD_ID = 0  # padding id used before delay slot fires
+
 
 def _maybe_prefix(prefix: str, name: str) -> str:
     return f"{prefix}.{name}" if prefix else name
@@ -151,12 +153,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         # `make_omni_output`, which runs immediately before the sampler).
         self._batch_state: list[dict[str, Any]] | None = None
 
-        # Stacked weight caches built after load_weights(). Avoids a Python
-        # loop over n_vq heads/embeddings at every decode step.
-        # shapes: (n_vq, audio_vocab+1, hidden_size)
-        self._stacked_audio_head_w: torch.Tensor | None = None
-        self._stacked_audio_emb_w: torch.Tensor | None = None
-
     # ------------------------------------------------------------------
     # vLLM-Omni hooks
     # ------------------------------------------------------------------
@@ -398,15 +394,9 @@ def preprocess(
                     chunk = ref_codes[ref_offset:end_off]
                     if chunk.numel() > 0 and chunk.shape[0] == span_len:
                         codes = chunk.to(device=device, dtype=torch.long).clamp_(0, self.audio_vocab_size)
-                        if self._stacked_audio_emb_w is not None:
-                            codes_nq = codes.t()  # (n_vq, L)
-                            audio_embed = self._stacked_audio_emb_w[
-                                torch.arange(self.n_vq, device=device)[:, None], codes_nq
-                            ].sum(0)
-                        else:
-                            audio_embed = torch.zeros_like(embeds)
-                            for i, emb_layer in enumerate(self.audio_embeddings):
-                                audio_embed = audio_embed + emb_layer(codes[:, i])
+                        audio_embed = torch.zeros_like(embeds)
+                        for i, emb_layer in enumerate(self.audio_embeddings):
+                            audio_embed = audio_embed + emb_layer(codes[:, i])
                         embeds = embeds + audio_embed
                         last_ref_row = codes[-1].detach()
 
@@ -446,17 +436,10 @@ def preprocess(
         audio_codes_buf = (info_dict.get("audio_codes", {}) or {}).get("current")
         if isinstance(audio_codes_buf, torch.Tensor) and audio_codes_buf.numel() == self.n_vq:
             codes = audio_codes_buf.to(device=device, dtype=torch.long)
-            if self._stacked_audio_emb_w is not None:
-                # Single gather: stacked_emb[i, codes[i]] for each i → (n_vq, H) → sum → (1, H)
-                n_vq_idx = torch.arange(self.n_vq, device=device)
-                audio_embed = self._stacked_audio_emb_w[n_vq_idx, codes.clamp(0, self.audio_vocab_size)].sum(
-                    0, keepdim=True
-                )
-            else:
-                audio_embed = torch.zeros_like(text_embed)
-                for i, emb_layer in enumerate(self.audio_embeddings):
-                    code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1)
-                    audio_embed = audio_embed + emb_layer(code_i)
+            audio_embed = torch.zeros_like(text_embed)
+            for i, emb_layer in enumerate(self.audio_embeddings):
+                code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1)
+                audio_embed = audio_embed + emb_layer(code_i)
             combined = text_embed + audio_embed
         else:
             combined = text_embed
@@ -518,38 +501,17 @@ def _sample_audio_codes(
         if not bool(sampling_mask.any()):
             return codes
 
+        # Use deploy YAML defaults. The audio sampler isn't routed through
+        # vLLM's sampler so we apply the upstream's audio_top_k locally.
         audio_top_k = 25
         audio_temp = 1.7
 
-        if self._stacked_audio_head_w is not None:
-            # Single batched matmul: (n_vq, V+1, H) @ (H,) → (n_vq, V+1).
-            # Replaces n_vq serial nn.Linear calls.
-            all_logits = self._stacked_audio_head_w @ last_h.reshape(-1)  # (n_vq, V+1)
-            all_logits[:, -1] = float("-inf")  # invalid sentinel
-            all_logits[:, self.audio_pad_code] = float("-inf")
-
-            active_idx = sampling_mask.nonzero(as_tuple=True)[0]  # (n_active,)
-            active_logits = all_logits[active_idx]  # (n_active, V+1)
-            active_logits = active_logits / max(audio_temp, 1e-6)
-            if audio_top_k > 0 and audio_top_k < active_logits.shape[-1]:
-                top_vals, _ = torch.topk(active_logits, audio_top_k, dim=-1)
-                kth = top_vals[:, -1:].expand_as(active_logits)
-                active_logits = torch.where(
-                    active_logits < kth,
-                    torch.full_like(active_logits, float("-inf")),
-                    active_logits,
-                )
-            probs = torch.softmax(active_logits, dim=-1)
-            codes[active_idx] = torch.multinomial(probs, num_samples=1).squeeze(-1).long()
-            return codes
-
-        # Fallback: per-head loop (used before load_weights() completes).
         for i in range(n_vq):
             if not bool(sampling_mask[i]):
                 continue
             head = self.audio_heads[i]
             logits = head(last_h).reshape(-1)  # (V,)
-            logits[..., -1] = float("-inf")
+            logits[..., -1] = float("-inf")  # invalid sentinel
             logits[..., self.audio_pad_code] = float("-inf")
             sampled = self._sample_with_top_k(logits, audio_top_k, audio_temp)
             codes[i] = sampled.long()
@@ -591,19 +553,7 @@ def make_omni_output(
         # rows align with hidden rows, which align with info_dicts in order.
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
 
-        # Per-request (start, end) hidden-row spans from the runner. In mixed
-        # prefill+decode steps the per-request token counts differ, so an equal
-        # rows-per-request split would sample codes from the wrong rows; fall
-        # back to that split only when spans are unavailable.
-        spans = kwargs.get("request_token_spans")
-
-        # One accumulated-codes tensor per request, in batch order. The generic
-        # splitter (``to_payload_element``) routes ``per_req_codes[idx]`` to
-        # request ``idx``, so Stage 1 decodes each request's own codes instead
-        # of request 0's. Skipped requests get an empty placeholder so indices
-        # stay aligned (a shorter list would silently fall back to entry 0).
-        per_req_codes: list[torch.Tensor] = []
-        have_codes = False
+        audio_codes_list: list[torch.Tensor] = []
 
         if hidden.numel() > 0:
             num_rows = hidden.shape[0]
@@ -611,16 +561,10 @@ def make_omni_output(
 
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
-                    per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
-                if spans is not None and i < len(spans):
-                    row_start, row_end = spans[i]
-                    row_end = min(int(row_end), num_rows)
-                else:
-                    row_start = i * rows_per_req
-                    row_end = min(row_start + rows_per_req, num_rows)
+                row_start = i * rows_per_req
+                row_end = min(row_start + rows_per_req, num_rows)
                 if row_start >= row_end:
-                    per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
 
                 # Sample new audio codes from the last hidden state of this request.
@@ -638,20 +582,21 @@ def make_omni_output(
                     "current": new_codes,
                     "accumulated": updated_acc,
                 }
-                per_req_codes.append(updated_acc)
-                have_codes = True
+                audio_codes_list.append(updated_acc)
 
-        if not have_codes:
+        if not audio_codes_list:
             return OmniOutput(
                 text_hidden_states=hidden,
                 multimodal_outputs={},
             )
 
-        # Emit one accumulated-codes tensor per request as a list; the splitter
-        # picks ``element[idx]`` per request on the way to Stage 1.
+        # Forward the accumulated audio codes (T, NQ) of the first request to
+        # Stage 1.  Multi-request batching for the codec is not yet supported
+        # by this adapter; the existing code already takes the first stage
+        # output anyway.
         return OmniOutput(
             text_hidden_states=hidden,
-            multimodal_outputs={"codes": {"audio": per_req_codes}},
+            multimodal_outputs={"codes": {"audio": audio_codes_list[0]}},
         )
 
     # ------------------------------------------------------------------
@@ -716,14 +661,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         for n in backbone_loaded:
             loaded.add(f"model.{n}")
 
-        # Build stacked weight caches for vectorised decode ops.
-        self._stacked_audio_head_w = torch.stack(
-            [h.weight.detach() for h in self.audio_heads], dim=0
-        )  # (n_vq, audio_vocab+1, hidden_size)
-        self._stacked_audio_emb_w = torch.stack(
-            [e.weight.detach() for e in self.audio_embeddings], dim=0
-        )  # (n_vq, audio_vocab+1, hidden_size)
-
         return loaded
 
 
@@ -811,10 +748,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.logits_processor = LogitsProcessor(self.text_vocab_size)
         self._batch_state: list[dict[str, Any]] | None = None
 
-        # Stacked weight cache built after load_weights().
-        # shape: (n_vq, audio_vocab_size, hidden_size)
-        self._stacked_audio_emb_w: torch.Tensor | None = None
-
         self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
             ("audio_codes", "current"),
             ("audio_codes", "accumulated"),
@@ -904,11 +837,6 @@ def _build_input_embeds(
         embeds = self.embed_tokens[0](text_ids)
         if audio_codes is None:
             return embeds
-        if self._stacked_audio_emb_w is not None:
-            # Single gather: (n_vq, vocab, H)[arange[:,None], codes.T] → (n_vq, T, H) → sum → (T, H)
-            n_vq_idx = torch.arange(self.n_vq, device=audio_codes.device)
-            codes_t = audio_codes.t().clamp(0, self.audio_vocab_size - 1)  # (n_vq, T)
-            return embeds + self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t].sum(0)
         for i in range(self.n_vq):
             col = audio_codes[:, i].clamp_(0, self.audio_vocab_size - 1)
             embeds = embeds + self.embed_tokens[i + 1](col)
@@ -1005,27 +933,15 @@ def make_omni_output(
 
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
 
-        # See delay talker: real per-request row spans from the runner, with the
-        # equal-split as fallback when unavailable.
-        spans = kwargs.get("request_token_spans")
-        # Per-request accumulated codes in batch order, pre-filled with empty
-        # placeholders so every skip path (stopped / eos / bos / pad) keeps
-        # indices aligned for to_payload_element's per-request routing. A
-        # shorter list would silently fall back to request 0's codes.
-        per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts]
-        have_codes = False
+        audio_codes_list: list[torch.Tensor] = []
         if hidden.numel() > 0 and info_dicts:
             num_rows = hidden.shape[0]
             rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1))
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
                     continue
-                if spans is not None and i < len(spans):
-                    row_start, row_end = spans[i]
-                    row_end = min(int(row_end), num_rows)
-                else:
-                    row_start = i * rows_per_req
-                    row_end = min(row_start + rows_per_req, num_rows)
+                row_start = i * rows_per_req
+                row_end = min(row_start + rows_per_req, num_rows)
                 if row_start >= row_end:
                     continue
 
@@ -1056,7 +972,7 @@ def make_omni_output(
                     history_per_codebook=hist_per_cb,
                 ).squeeze(0)  # (n_vq,)
                 if int(state.get("step", 0)) < 5 or int(state.get("step", 0)) % 50 == 0:
-                    logger.debug(
+                    logger.info(
                         "[MossTTSRealtime make_omni] step=%d ch0=%d cursor=%d/%d",
                         int(state.get("step", 0)),
                         int(new_codes[0].item()),
@@ -1092,25 +1008,23 @@ def make_omni_output(
 
                 info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
                 state["step"] = int(state.get("step", 0)) + 1
-                per_req_codes[i] = updated_acc
-                have_codes = True
+                audio_codes_list.append(updated_acc)
 
-        if not have_codes:
+        if not audio_codes_list:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
 
-        # Emit one accumulated-codes tensor per request as a list; the splitter
-        # routes ``element[idx]`` to request idx, so each Stage-1 request decodes
-        # its own codes instead of request 0's.
-        # The realtime variant emits raw codes (no delay pattern), so we signal
-        # that to the chunk processor via a 1-D bool tensor in ``meta.finished``
-        # (already a tensor field on the schema). The processor reads its truth
-        # value as "skip de-delay" only for realtime; the delay path doesn't
-        # populate ``meta`` at all.
-        device = hidden.device
+        # Forward the first request's accumulated codes (multi-request
+        # batching for the codec is not yet wired up — same as delay).
+        # The realtime variant emits raw codes (no delay pattern), so we
+        # signal that to the chunk processor via a 1-D bool tensor in
+        # ``meta.finished`` (already a tensor field on the schema). The
+        # processor reads its truth value as "skip de-delay" only for
+        # realtime; the delay path doesn't populate ``meta`` at all.
+        device = audio_codes_list[0].device
         return OmniOutput(
             text_hidden_states=hidden,
             multimodal_outputs={
-                "codes": {"audio": per_req_codes},
+                "codes": {"audio": audio_codes_list[0]},
                 "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
             },
         )
@@ -1170,12 +1084,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         for n in backbone_loaded:
             loaded.add(f"model.{n}")
 
-        # Build stacked embedding cache for vectorised decode ops.
-        # embed_tokens[0] is the text embedding; [1..n_vq] are audio codebooks.
-        self._stacked_audio_emb_w = torch.stack(
-            [self.embed_tokens[i + 1].weight.detach() for i in range(self.n_vq)], dim=0
-        )  # (n_vq, audio_vocab_size, hidden_size)
-
         logger.info(
             "[MossTTSRealtime] loaded %d/%d params; skipped=%d (first 5: %s)",
             len(loaded),
diff --git a/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py
new file mode 100644
index 00000000000..a5a9f25acd3
--- /dev/null
+++ b/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration for MOSS-TTS-Nano in vLLM-Omni single-stage pipeline."""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class MossTTSNanoConfig(PretrainedConfig):
+    """Config for MOSS-TTS-Nano (OpenMOSS-Team/MOSS-TTS-Nano).
+
+    The model is a Global-Local GPT2-based AR LM (0.1B) paired with the
+    MOSS-Audio-Tokenizer-Nano codec.  Both components run in a single
+    vLLM-Omni generation stage.
+
+    Relevant fields from config.json:
+      gpt2_config  – backbone config (n_layer=12, n_embd=768, n_head=12, ...)
+      n_vq         – number of RVQ codebooks (16)
+      audio_vocab_size – per-codebook vocabulary size (1024)
+      audio_tokenizer_pretrained_name_or_path – HF hub path for codec model
+    """
+
+    model_type = "moss_tts_nano"
+
+    def __init__(self, **kwargs):
+        gpt2_cfg = kwargs.pop("gpt2_config", None) or {}
+        if hasattr(gpt2_cfg, "to_dict"):
+            gpt2_cfg = gpt2_cfg.to_dict()
+
+        super().__init__(**kwargs)
+
+        # --- GPT2 backbone parameters (exposed at top level for vLLM) ---
+        self.hidden_size = gpt2_cfg.get("n_embd", 768)
+        self.num_hidden_layers = gpt2_cfg.get("n_layer", 12)
+        self.num_attention_heads = gpt2_cfg.get("n_head", 12)
+        self.num_key_value_heads = self.num_attention_heads  # no GQA in GPT2
+        self.head_dim = self.hidden_size // self.num_attention_heads
+        self.vocab_size = gpt2_cfg.get("vocab_size", 16384)
+        self.max_position_embeddings = gpt2_cfg.get("n_positions", 32768)
+        self.intermediate_size = gpt2_cfg.get("n_inner", self.hidden_size * 4)
+
+        # --- Audio codec parameters ---
+        self.n_vq: int = getattr(self, "n_vq", 16)
+        self.audio_vocab_size: int = getattr(self, "audio_vocab_size", 1024)
+        self.audio_start_token_id: int = getattr(self, "audio_start_token_id", 6)
+        self.audio_end_token_id: int = getattr(self, "audio_end_token_id", 7)
+        self.audio_user_slot_token_id: int = getattr(self, "audio_user_slot_token_id", 8)
+        self.audio_assistant_slot_token_id: int = getattr(self, "audio_assistant_slot_token_id", 9)
+        self.audio_tokenizer_pretrained_name_or_path: str = getattr(
+            self,
+            "audio_tokenizer_pretrained_name_or_path",
+            "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
+        )
+        self.audio_tokenizer_sample_rate: int = getattr(self, "audio_tokenizer_sample_rate", 48000)
+
+        # vLLM requires speculative_config to be absent or None
+        self.speculative_config = None
+
+    def get_text_config(self, **kwargs):
+        """Return self so vLLM uses our top-level config."""
+        return self
diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py
index 70596098298..f52577f67aa 100644
--- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py
+++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py
@@ -172,6 +172,7 @@ class MossTTSNanoForGeneration(nn.Module):
     has_preprocess = False
     has_postprocess = False
     enable_update_additional_information = True
+    inject_omni_request_id_into_runtime_info = True
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         super().__init__()
diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py
index 43c496a1cb0..c5026086e23 100644
--- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py
+++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py
@@ -379,12 +379,13 @@ def _forward_generator(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutpu
         target_ids = torch.full((num_codebooks, target_len), mask_id, dtype=torch.long, device=device)
 
         # Conditional: [text] [ref_audio?] [target_mask]
-        cond_audio_start = text_ids.shape[1]
         if ref_audio_tokens is not None:
             ref_tokens = ref_audio_tokens.to(device)  # [8, T_ref]
             cond_ids = torch.cat([text_ids, ref_tokens, target_ids], dim=1)
+            cond_audio_start = text_ids.shape[1]
         else:
             cond_ids = torch.cat([text_ids, target_ids], dim=1)
+            cond_audio_start = text_ids.shape[1]
 
         cond_len = cond_ids.shape[1]
 
diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
index 7eacfb9bb92..fc5befeb959 100644
--- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
+++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
@@ -6,9 +6,8 @@
 Generates 8-codebook audio tokens from text via 32-step non-autoregressive
 iterative masked prediction with classifier-free guidance.
 
-Uses full bidirectional attention computed directly with PyTorch SDPA
-(torch.nn.functional.scaled_dot_product_attention); no auto-selected
-FlashAttention/SageAttention/DiffusionAttention backend is used.
+Uses vLLM-Omni's DiffusionAttention for optimized full (bidirectional) attention
+via FlashAttention/SageAttention/SDPA backends.
 """
 
 from __future__ import annotations
@@ -238,7 +237,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge
 
 
 # ---------------------------------------------------------------------------
-# Qwen3-style transformer blocks using PyTorch SDPA
+# Qwen3-style transformer blocks using DiffusionAttention
 # ---------------------------------------------------------------------------
 
 
@@ -257,7 +256,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class OmniVoiceAttention(nn.Module):
-    """Qwen3-style GQA attention using PyTorch SDPA (full bidirectional)."""
+    """Qwen3-style GQA attention using DiffusionAttention backend."""
 
     def __init__(self, config: OmniVoiceConfig):
         super().__init__()
@@ -316,17 +315,18 @@ def forward(
         v = v.permute(0, 2, 1, 3)
 
         # Convert [B, 1, S, S] bool mask to float mask for SDPA
-        # (0.0 where attend, -inf where masked)
         sdpa_mask = None
         if attention_mask is not None:
-            sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf"))
+            sdpa_mask = attention_mask.to(dtype=q.dtype)
+            sdpa_mask = sdpa_mask.masked_fill(~attention_mask, float("-inf"))
+            sdpa_mask = sdpa_mask.masked_fill(attention_mask, 0.0)
 
         out = F.scaled_dot_product_attention(
             q,
             k,
             v,
             attn_mask=sdpa_mask,
-            scale=self.scale,
+            scale=1.0 / math.sqrt(self.head_dim),
         )
 
         # Back to (batch, seq, heads * head_dim)
@@ -351,7 +351,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class OmniVoiceTransformerBlock(nn.Module):
-    """Single Qwen3 transformer block with PyTorch SDPA attention."""
+    """Single Qwen3 transformer block with DiffusionAttention."""
 
     def __init__(self, config: OmniVoiceConfig):
         super().__init__()
@@ -643,9 +643,10 @@ class OmniVoiceGenerator(nn.Module):
     - 32-step iterative unmasking with classifier-free guidance
 
     Optimizations:
-    - Full bidirectional attention via PyTorch SDPA (no auto-selected
-      FlashAttn/SageAttn/DiffusionAttention backend)
+    - DiffusionAttention (FlashAttn/SageAttn/SDPA auto-selected)
+    - TeaCache / Cache-DiT compatible (hook-based, non-intrusive)
     - regionally_compile() compatible for torch.compile on repeated blocks
+    - Sequence parallelism via SP hooks for multi-GPU
     """
 
     # For regionally_compile() support
diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py
index 28ca4f36289..8bb921a1eae 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py
@@ -27,7 +27,7 @@ class Qwen3TTSSpeakerEncoderConfig(PretrainedConfig):
     Args:
         mel_dim (`int`, *optional*, defaults to 128):
             The dimension of the input mel-spectrogram.
-        enc_dim (`int`, *optional*, defaults to 1024):
+        enc_dim (`int`, *optional*, defaults to 192):
             The dimension of the final speaker embedding.
         enc_channels (`list[int]`, *optional*, defaults to `[512, 512, 512, 512, 1536]`):
             A list of output channels for each TDNN/SERes2Net layer in the encoder.
@@ -80,11 +80,11 @@ class Qwen3TTSTalkerCodePredictorConfig(PretrainedConfig):
 
 
     Args:
-        vocab_size (`int`, *optional*, defaults to 2048):
+        vocab_size (`int`, *optional*, defaults to 151936):
             Vocabulary size of the Qwen3TTSTalkerCodePredictor model.
             Defines the number of different tokens that can be represented by the
             `inputs_ids` passed when calling [`Qwen3TTSTalkerCodePredictorModel`]
-        hidden_size (`int`, *optional*, defaults to 1024):
+        hidden_size (`int`, *optional*, defaults to 4096):
             Dimension of the hidden representations.
         intermediate_size (`int`, *optional*, defaults to 22016):
             Dimension of the MLP representations.
@@ -271,11 +271,11 @@ class Qwen3TTSTalkerConfig(PretrainedConfig):
 
 
     Args:
-        vocab_size (`int`, *optional*, defaults to 3072):
+        vocab_size (`int`, *optional*, defaults to 151936):
             Vocabulary size of the Qwen3TTSTalker model.
             Defines the number of different tokens that can be represented by the
             `inputs_ids` passed when calling [`Qwen3TTSTalkerModel`]
-        hidden_size (`int`, *optional*, defaults to 1024):
+        hidden_size (`int`, *optional*, defaults to 2048):
             Dimension of the hidden representations.
         intermediate_size (`int`, *optional*, defaults to 6144):
             Dimension of the MLP representations.
@@ -436,6 +436,7 @@ def __init__(
             self.rope_scaling["rope_type"] = self.rope_scaling["type"]
 
         if code_predictor_config is None:
+            code_predictor_config = {}
             self.code_predictor_config = Qwen3TTSTalkerCodePredictorConfig()
             logger.info("code_predictor_config is None. Initializing code_predictor model with default values")
         elif isinstance(code_predictor_config, Qwen3TTSTalkerCodePredictorConfig):
diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py
index 0d8f032949e..b0294b72583 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py
@@ -1320,7 +1320,9 @@ def _get_ref_audio() -> tuple[np.ndarray, int]:
             # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
             spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long)
             spk_embed = codec_embed(spk_tensor)
-            if spk_embed.ndim in (1, 2):
+            if spk_embed.ndim == 1:
+                spk_embed = spk_embed.view(1, 1, -1)
+            elif spk_embed.ndim == 2:
                 spk_embed = spk_embed.view(1, 1, -1)
             speaker_embed = spk_embed
             codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
@@ -1521,7 +1523,9 @@ def _first(x: object, default: object) -> object:
 
                 ref_code_len: int | None = None
                 if isinstance(ref_code, list):
-                    if ref_code:
+                    if ref_code and isinstance(ref_code[0], list):
+                        ref_code_len = len(ref_code)
+                    elif ref_code:
                         ref_code_len = len(ref_code)
                 elif hasattr(ref_code, "shape"):
                     try:
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
index c1bb7c049ad..64eb3d25320 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
@@ -286,6 +286,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             )
         self._codec_eos_token_id = int(getattr(self.talker_config, "codec_eos_token_id", -1))
 
+        self._eos_logit_bias: float = 0.0
+
         self.have_multimodal_outputs = True
         self.has_preprocess = True
         self.has_postprocess = True
@@ -515,6 +517,11 @@ def compute_logits(
         # Mask out invalid codec ids using the pre-built constant buffer.
         logits = logits.masked_fill(self._codec_disallowed_mask, float("-inf"))
 
+        if self._eos_logit_bias != 0.0:
+            eos_id = self._codec_eos_token_id
+            if 0 <= eos_id < logits.shape[-1]:
+                logits[:, eos_id] = logits[:, eos_id] + self._eos_logit_bias
+
         return logits
 
     # -------------------- Omni multimodal output plumbing --------------------
diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
index 447823abb11..6eb78eeae0d 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
@@ -912,6 +912,11 @@ def enable_cudagraph(
             self._cudagraph_wrapper.compile_shapes,
         )
 
+    def disable_cudagraph(self):
+        self._cudagraph_enabled = False
+        self._cudagraph_wrapper = None
+        logger.info("CUDA Graph disabled for decoder")
+
     def forward(self, codes):
         if codes.shape[1] != self.config.num_quantizers:
             raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py
index 4014014533d..5609abb394c 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py
@@ -33,6 +33,7 @@
 
 import random
 import typing as tp
+from math import ceil
 
 import numpy as np
 import torch
@@ -44,6 +45,10 @@
 logger = init_logger(__name__)
 
 
+def round_up_multiple(num, mult):
+    return ceil(num / mult) * mult
+
+
 def default(val: tp.Any, d: tp.Any) -> tp.Any:
     return val if val is not None else d
 
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2.py b/vllm_omni/model_executor/models/step_audio2/step_audio2.py
index 0a94ab9defb..82e5609ff40 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2.py
+++ b/vllm_omni/model_executor/models/step_audio2/step_audio2.py
@@ -26,17 +26,17 @@ class StepAudio2ForConditionalGeneration(nn.Module, SupportsMultiModal, Supports
     Step-Audio2 Main Controller
 
     Manages two-stage inference pipeline:
-    - Stage 0 (Thinker): Audio understanding and token generation
-    - Stage 1 (Token2Wav): Audio token to waveform synthesis
+    - Stage 1 (Thinker): Audio understanding and token generation
+    - Stage 2 (Token2Wav): Audio token to waveform synthesis
 
     Usage:
-        # Stage 0: Thinker
+        # Stage 1: Thinker
         model = StepAudio2ForConditionalGeneration(
             vllm_config=config,
             model_stage="thinker"
         )
 
-        # Stage 1: Token2Wav
+        # Stage 2: Token2Wav
         model = StepAudio2ForConditionalGeneration(
             vllm_config=config,
             model_stage="token2wav"
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py
index 2f2918a1be3..802cd63bef8 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py
+++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py
@@ -42,6 +42,39 @@ def audio_end(self) -> int:
         return self.audio_start + self.audio_vocab_size - 1
 
 
+@dataclass(frozen=True)
+class StepAudio2EncoderConfig:
+    """
+    Step-Audio2 audio encoder architecture configuration.
+
+    These values define the encoder structure and should match your model weights.
+    """
+
+    # Mel spectrogram parameters
+    n_mels: int = 128
+    """Number of mel frequency bins"""
+
+    # Encoder architecture
+    n_audio_ctx: int = 1500
+    """Audio context length (max sequence length)"""
+
+    n_audio_state: int = 512
+    """Audio encoder hidden state dimension"""
+
+    n_audio_head: int = 8
+    """Number of attention heads in encoder"""
+
+    n_audio_layer: int = 6
+    """Number of encoder layers"""
+
+    # Adapter parameters
+    kernel_size: int = 3
+    """Convolution kernel size for adapter"""
+
+    adapter_stride: int = 2
+    """Stride for adapter convolution (downsampling factor)"""
+
+
 @dataclass(frozen=True)
 class StepAudio2Token2WavConfig:
     """Token2Wav synthesis configuration."""
@@ -108,6 +141,7 @@ class StepAudio2ModelConfig:
 
 # Default configuration instances
 DEFAULT_TOKEN_CONFIG = StepAudio2TokenConfig()
+DEFAULT_ENCODER_CONFIG = StepAudio2EncoderConfig()
 DEFAULT_TOKEN2WAV_CONFIG = StepAudio2Token2WavConfig()
 DEFAULT_STREAM_CONFIG = StepAudio2StreamConfig()
 DEFAULT_MODEL_CONFIG = StepAudio2ModelConfig()
@@ -128,10 +162,12 @@ class StepAudio2ModelConfig:
 
 __all__ = [
     "StepAudio2TokenConfig",
+    "StepAudio2EncoderConfig",
     "StepAudio2Token2WavConfig",
     "StepAudio2StreamConfig",
     "StepAudio2ModelConfig",
     "DEFAULT_TOKEN_CONFIG",
+    "DEFAULT_ENCODER_CONFIG",
     "DEFAULT_TOKEN2WAV_CONFIG",
     "DEFAULT_STREAM_CONFIG",
     "DEFAULT_MODEL_CONFIG",
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py
index c92f2397813..7fd0bf4fbd7 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py
+++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py
@@ -609,7 +609,7 @@ class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, S
     Architecture:
         AudioEncoder (6 layers, 512 hidden) →
         Adaptor (512 → LLM dim) →
-        Qwen2 LLM (vocab=158720)
+        Qwen2 LLM (vocab=64012)
     """
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -868,3 +868,31 @@ def separate_tokens(
     def has_audio_output(token_ids: list[int], audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start) -> bool:
         """Check if generated tokens contain audio tokens"""
         return any(tid >= audio_start for tid in token_ids)
+
+
+class StepAudio2OutputProcessor:
+    """Helper class to process Step-Audio2 outputs"""
+
+    @staticmethod
+    def process_output(output_ids: torch.Tensor, tokenizer, remove_audio_padding: bool = True) -> dict:
+        """Process model output and separate text/audio tokens"""
+        if isinstance(output_ids, torch.Tensor):
+            output_ids = output_ids.squeeze().tolist()
+
+        # Separate tokens
+        text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(output_ids)
+
+        # Remove audio padding if requested
+        if remove_audio_padding and audio_tokens:
+            audio_tokens = [t for t in audio_tokens if t < DEFAULT_TOKEN_CONFIG.audio_eos]
+
+        # Decode text
+        text = tokenizer.decode(text_tokens, skip_special_tokens=False)
+
+        return {
+            "text": text,
+            "text_tokens": text_tokens,
+            "audio_tokens": audio_tokens,
+            "has_audio": len(audio_tokens) > 0,
+            "all_tokens": output_ids,
+        }
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py
index 86109ad9c7d..a59ba3fef19 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py
+++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py
@@ -94,6 +94,8 @@ def __init__(
         self.mel_cache_len = DEFAULT_STREAM_CONFIG.mel_cache_len
         self.source_cache_len = STREAM_SOURCE_CACHE_LEN
         self.speech_window: torch.Tensor | None = None  # created lazily on device
+        self.stream_cache: dict | None = None  # flow model causal cache
+        self.hift_cache_dict: dict = {}  # HiFT vocoder overlap cache
 
     def _ensure_models_loaded(self):
         """Lazy load models on first use"""
@@ -251,6 +253,125 @@ def forward(
         else:
             return wav
 
+    # ------------------------------------------------------------------
+    # Streaming (chunked) inference
+    # ------------------------------------------------------------------
+
+    def setup_stream(self, prompt_wav: str) -> None:
+        """Initialize flow + HiFT caches for chunked streaming inference.
+
+        Must be called once before the first ``stream_chunk()`` call for a
+        given utterance.
+        """
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        (
+            prompt_speech_tokens,
+            _,
+            spk_emb,
+            prompt_mels,
+            _,
+        ) = self.cache[prompt_wav]
+
+        # Lazily create the Hamming window on the correct device
+        if self.speech_window is None or self.speech_window.device != self.device:
+            self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to(
+                device=self.device, dtype=torch.float32
+            )
+
+        # Flow model: setup causal cache with prompt tokens + lookahead
+        pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len
+        self.stream_cache = self.flow.setup_cache(
+            torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]], dim=1),
+            prompt_mels,
+            spk_emb,
+            n_timesteps=self.n_timesteps,
+        )
+
+        # HiFT vocoder: empty overlap buffers
+        self.hift_cache_dict = {
+            "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device),
+            "source": torch.zeros(1, 1, 0, device=self.device),
+            "speech": torch.zeros(1, 0, device=self.device),
+        }
+
+    def stream_chunk(
+        self,
+        audio_tokens: list[int],
+        prompt_wav: str,
+        last_chunk: bool = False,
+    ) -> torch.Tensor:
+        """Process one chunk of audio tokens and return a waveform segment.
+
+        Args:
+            audio_tokens: Token IDs for this chunk (includes lookahead for
+                non-last chunks).
+            prompt_wav: Path to the prompt wav (for cached speaker embeddings).
+            last_chunk: Whether this is the final chunk of the utterance.
+
+        Returns:
+            1-D float tensor of audio samples (24 kHz).
+        """
+        if self.stream_cache is None:
+            raise ValueError("stream_cache not initialised – call setup_stream() first")
+
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
+
+        token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device)
+
+        with torch.amp.autocast(
+            str(self.device.type),
+            dtype=torch.float16 if self.float16 else torch.float32,
+        ):
+            chunk_mel, self.stream_cache = self.flow.inference_chunk(
+                token=token_tensor,
+                spk=spk_emb,
+                cache=self.stream_cache,
+                last_chunk=last_chunk,
+                n_timesteps=self.n_timesteps,
+            )
+
+        # Trim estimator attention cache to avoid unbounded growth
+        keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep
+        est_att = self.stream_cache.get("estimator_att_cache")
+        if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep):
+            self.stream_cache["estimator_att_cache"] = torch.cat(
+                [est_att[:, :, :, :, : prompt_mels.shape[1]], est_att[:, :, :, :, -keep:]],
+                dim=4,
+            )
+
+        # ---- HiFT vocoder with overlap-add ----
+        hift_cache_mel = self.hift_cache_dict["mel"]
+        hift_cache_source = self.hift_cache_dict["source"]
+        hift_cache_speech = self.hift_cache_dict["speech"]
+
+        mel = torch.cat([hift_cache_mel, chunk_mel], dim=2)
+        speech, source = self.hift(mel, hift_cache_source)
+
+        # Cross-fade overlap region for smooth concatenation
+        if hift_cache_speech.shape[-1] > 0:
+            speech = fade_in_out(speech, hift_cache_speech, self.speech_window)
+
+        # Update vocoder cache for the next chunk
+        self.hift_cache_dict = {
+            "mel": mel[..., -self.mel_cache_len :].clone().detach(),
+            "source": source[:, :, -self.source_cache_len :].clone().detach(),
+            "speech": speech[:, -self.source_cache_len :].clone().detach(),
+        }
+
+        # For non-last chunks, trim the tail that will be re-synthesised
+        if not last_chunk:
+            speech = speech[:, : -self.source_cache_len]
+
+        return speech.squeeze(0)  # [samples]
+
+    def reset_stream(self) -> None:
+        """Clear all streaming state after an utterance is complete."""
+        self.stream_cache = None
+        self.hift_cache_dict = {}
+
     # ------------------------------------------------------------------
     # Per-request streaming (does NOT mutate self – uses external state)
     # ------------------------------------------------------------------
diff --git a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py
index 13f1d0937d3..b87ec5aafef 100644
--- a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py
+++ b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py
@@ -65,6 +65,7 @@ def __init__(
         prefix: str = "",
     ) -> None:
         super().__init__()
+        self.layer_idx = layer_idx
         self.hidden_size = hidden_size
         self.num_heads = num_attention_heads
         self.head_dim = kv_channels if kv_channels else hidden_size // num_attention_heads
@@ -317,12 +318,18 @@ def precompute_fused_qkv(self) -> None:
                 ).detach()
 
     def compile_selective(self) -> list[str]:
-        """Compile the full model ``forward`` as one graph.
-
-        ``torch.compile`` is applied at the model level so Dynamo unrolls the
-        per-layer Python loop inside the graph. Graph breaks at PagedAttention
-        produce sub-graphs, but Dynamo memoises the whole trace once, so the
-        per-decode-step dispatch drops to just a few instead of one per layer.
+        """Compile the full model forward as one graph.
+
+        Earlier versions compiled ``layer.mlp`` + ``layer.self_attn.o_proj``
+        (PR #2690) and then the whole ``layer`` (perf/voxcpm2-streaming-vae).
+        Both still paid one Dynamo dispatch per layer per decode step.
+        V3 profiling showed 1,332 per-layer dispatches (~28 layers × ~47
+        decode steps) costing ~726 ms of CPU self-time for a long prompt.
+
+        Compiling ``forward`` at the model level lets Dynamo unroll the
+        28-layer Python loop inside the graph. Graph breaks at
+        PagedAttention produce sub-graphs but Dynamo memoises the whole
+        trace once, so the per-step dispatch drops from 28 to just a few.
         """
         if self._compiled_layers:
             return []
diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
index 9505bfbee38..3fcc683f3aa 100644
--- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
+++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
@@ -312,6 +312,7 @@ class _RequestState:
     decode_step_count: int = 0
     request_start_time: float = 0.0
     prefill_completed: bool = False
+    prefill_text: str = ""
     prompt_cache: dict | None = None
     prefill_masks: tuple | None = None
     is_stopping: bool = False
@@ -2509,6 +2510,7 @@ def preprocess(
                 token_ids = token_ids[1:]
 
             state = self._get_or_create_state(req_id)
+            state.prefill_text = ""
             state.decode_pad = None
             state.prefill_completed = False
             state.decode_step_count = 0
@@ -2570,26 +2572,29 @@ def preprocess(
                 )
 
             if state.prompt_cache is None and (ref_audio or (prompt_audio and prompt_text)):
-                try:
-                    state.prompt_cache = self._build_prompt_cache(
-                        ref_audio=ref_audio,
-                        prompt_audio=prompt_audio,
-                        prompt_text=prompt_text,
-                    )
-                    if (
-                        voice_name
-                        and state.prompt_cache is not None
-                        and state.prompt_cache.get("mode") == "reference"
-                        and "ref_audio_feat" in state.prompt_cache
-                    ):
-                        _key = self._speaker_cache.make_cache_key(
-                            voice_name, model_type="voxcpm2", created_at=_created_at
+                if state.prompt_cache is None:
+                    try:
+                        state.prompt_cache = self._build_prompt_cache(
+                            ref_audio=ref_audio,
+                            prompt_audio=prompt_audio,
+                            prompt_text=prompt_text,
                         )
-                        self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()})
-                        logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
-                except Exception as e:
-                    logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e)
-                    state.prompt_cache = None
+                        if (
+                            voice_name
+                            and state.prompt_cache is not None
+                            and state.prompt_cache.get("mode") == "reference"
+                            and "ref_audio_feat" in state.prompt_cache
+                        ):
+                            _key = self._speaker_cache.make_cache_key(
+                                voice_name, model_type="voxcpm2", created_at=_created_at
+                            )
+                            self._speaker_cache.put(
+                                _key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()}
+                            )
+                            logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
+                    except Exception as e:
+                        logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e)
+                        state.prompt_cache = None
 
             inputs = self._build_prefill_inputs(token_ids, dev, req_id)
             tts = self.tts
diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
index 8b6872cfeef..d12034c243b 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
+++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
@@ -244,6 +244,9 @@ def embed_multimodal(self, **kwargs):
         # Delegate to generation model for multimodal processing
         return self.model.embed_multimodal(**kwargs)
 
+    def last_index_of(self, list, value):
+        return len(list) - 1 - list[::-1].index(value)
+
     def forward(
         self,
         input_ids: torch.Tensor,
diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
index ce005936c63..2e50e34f32d 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
+++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
@@ -246,6 +246,7 @@ def __init__(
 
         self.n_local_heads: int = args.n_heads
         self.n_local_kv_heads: int = args.n_kv_heads
+        self.repeats = self.n_local_heads
         self.layer_id = layer_id
 
         self.head_dim = args.head_dim
@@ -275,6 +276,7 @@ def __init__(
             bias=args.use_biases,
         )
 
+        self.softmax_scale: float = self.args.head_dim**-0.5
         self.repeats = self.n_local_heads // self.n_local_kv_heads
 
     def _native_attention(
@@ -418,6 +420,13 @@ def __init__(
         self.acoustic_transformer_args = args
         assert isinstance(self.acoustic_transformer_args, AcousticTransformerArgs)
 
+        # currently assuming always 1 semantic codebook + N acoustic codebooks
+        self.num_non_acoustic_embeddings = 1
+        self.num_acoustic_codebooks = len(self.model_args.get_codebook_sizes()) - self.num_non_acoustic_embeddings
+
+        # flow matching utils
+        self.sigma = args.sigma
+
         # codebook sizes
         acoustic_codebook_sizes = self.model_args.get_codebook_sizes(
             pad_to_multiple=None, include_special_tokens=False
@@ -706,6 +715,7 @@ def __call__(
                 )
 
             text_tokens_for_audio = list[torch.Tensor]()
+            assert audios is not None
             audio_tokens_pt = list[torch.Tensor]()
             for audio_token_array in audio_tokens:
                 assert isinstance(audio_token_array, np.ndarray)
@@ -998,7 +1008,7 @@ def _parse_and_validate_audio_arrays(
 
         if audio_arrays is not None:
             if not isinstance(audio_arrays, (torch.Tensor, list)):
-                raise ValueError(f"Incorrect type of audio_arrays. Got type: {type(audio_arrays)}")
+                raise ValueError(f"Incorrect type of images. Got type: {type(audio_arrays)}")
             if isinstance(audio_arrays, torch.Tensor) and audio_arrays.dim() == 3:
                 audio_arrays = flatten_bn(audio_arrays)
             if isinstance(audio_arrays, torch.Tensor):
@@ -1043,11 +1053,24 @@ def fake_logits_for_audio_tokens(
         fake_logits[~is_eos, self.audio_tok_id] = 1.0
         return fake_logits
 
+    # TODO(chenyo): Remove this
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        text_logits = self.language_model.compute_logits(
+            hidden_states,
+        )
+        assert text_logits is not None
+        return text_logits
+
     def compute_mm_logits(
         self,
         hidden_states: torch.Tensor,
         cfg_alpha: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        audio_codes = None
+        mm_tokens = None
         audio_codes = self.acoustic_transformer(
             llm_hidden=hidden_states,
             cfg_alpha=cfg_alpha,
@@ -1065,6 +1088,7 @@ def compute_mm_logits(
         return fake_eos, mm_tokens
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # fmt: on
         remapping_rules = [
             (r"^acoustic_transformer\.(.*)$", r"\1"),  # noqa: E501
             (r"^audio_tokenizer\.(.*)$", r"\1"),  # noqa: E501
diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py
index 57a86a5e91c..fc753a58f37 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py
+++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py
@@ -652,7 +652,7 @@ def __init__(
     def layer_id(self) -> int:
         return self._layer_id
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
         r = self.attention.forward(self.attention_norm(x))
         if self.post_attention_norm is not None:
             r = self.post_attention_norm(r)
@@ -757,7 +757,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             # transformer
             layer_args = deepcopy(args)
             layer_args.attn_sliding_window_size = cur_window_size
-            assert layer_args.qk_norm, "qk_norm must be True for encoder"
+            assert layer_args.qk_norm, "qk_norm must be True for decoder"
             encoder_transformer = Transformer(
                 args=layer_args,
                 n_layers=n_layers,
diff --git a/vllm_omni/request.py b/vllm_omni/request.py
index b094f2d39fa..f936069be2d 100644
--- a/vllm_omni/request.py
+++ b/vllm_omni/request.py
@@ -30,16 +30,15 @@ class OmniRequest(Request):
 
     def __init__(
         self,
-        *args,
         prompt_embeds: PromptEmbedsPayload | torch.Tensor | None = None,
         # Optional external request ID for tracking
         external_req_id: str | None = None,
         additional_information: AdditionalInformationPayload | None = None,
+        *args,
         **kwargs,
     ):
-        if prompt_embeds is not None:
-            kwargs["prompt_embeds"] = self._maybe_decode_prompt_embeds(prompt_embeds)
-        super().__init__(*args, **kwargs)
+        prompt_embeds_tensor = self._maybe_decode_prompt_embeds(prompt_embeds)
+        super().__init__(prompt_embeds=prompt_embeds_tensor, *args, **kwargs)
         # Preserve serialized prompt embeddings payload (optional)
         self.prompt_embeds_payload: PromptEmbedsPayload | None = (
             prompt_embeds if isinstance(prompt_embeds, PromptEmbedsPayload) else None
diff --git a/vllm_omni/transformers_utils/processors/ming.py b/vllm_omni/transformers_utils/processors/ming.py
index b32f10d1204..685909029ad 100644
--- a/vllm_omni/transformers_utils/processors/ming.py
+++ b/vllm_omni/transformers_utils/processors/ming.py
@@ -143,6 +143,7 @@ def __call__(
         #      (conv1 has stride=1 and does not change T)
         #   2. AudioProjector Conv1d: kernel=3, stride=2, padding=1
         # Combined: T → ((T-1)//2 + 1 - 1)//2 + 1
+        # See also: AudioProjector.compute_output_length()
         encoder_feats_lengths = ((audio_feats_lengths - 3 + 2 * 1) // 2 + 1 - 3 + 2 * 1) // 2 + 1
         audio_feats = torch.cat(audio_feat_list, dim=0).unsqueeze(0)  # [1, T_total, n_mels]
 
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index cdf2286bfe4..349e1da0411 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -1313,20 +1313,6 @@ def _build_model_kwargs_extra(self) -> dict:
 
             traceback.print_exc()
 
-        # Per-request (start, end) hidden-row spans so make_omni_output can map
-        # flat hidden rows to the right request in mixed prefill+decode steps,
-        # instead of assuming an equal rows-per-request split (which samples the
-        # wrong rows whenever per-request token counts differ).
-        nstp = self._omni_num_scheduled_tokens_np
-        if nstp is not None and len(nstp) == len(self.input_batch.req_ids):
-            try:
-                model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(nstp)
-            except Exception as e:
-                # Visible on purpose: the fallback is the equal rows-per-request
-                # split, which can re-introduce the cross-request corruption this
-                # plumbing fixes — a silent failure here must not pass unnoticed.
-                logger.warning("[OMNI] Failed to compute request_token_spans: %s", e)
-
         if self._omni_query_start_loc_model_kwarg:
             try:
                 num_reqs = len(self.input_batch.req_ids)

From e78aa21fcd1bbfbb054fa097f981af78736eab3b Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 08:07:30 +0000
Subject: [PATCH 02/51] [BDE] Add engine-level KV cache helpers (adapter, chunk
 window, pool)

Phase-1 building blocks for BDE KV management, reusing vLLM's paged KV stack as
a library (verified against vLLM 0.23.0 on CPU):

- BDERequestAdapter: duck-types the Request surface KVCacheManager reads;
  advances per chunk (on_chunk_committed), not per denoise step.
- ChunkWindowSpec / ChunkWindowManager: chunk-granularity sliding window
  (sliding_window = window_chunks * chunk_size), registered via
  @register_kv_cache_spec; overrides only the token-skip math
  (chunk_window_skipped_tokens) for chunk-boundary eviction + sink protection.
- BDEKVConfig: disabled by default.
- build_kv_manager / compute_num_blocks: static sizing + metadata-only
  KVCacheManager construction.
- 15 tests, incl. a real allocate->free roundtrip driving KVCacheManager with
  the adapter (conformance against installed vLLM).

.gitignore: un-ignore vllm_omni/bde/kv_cache/ (the repo-wide kv_cache/ rule was
swallowing the source package), mirroring the existing !vllm_omni/diffusion/cache/.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .gitignore                             |   1 +
 tests/bde/test_kv_cache.py             | 174 +++++++++++++++++++++++++
 vllm_omni/bde/kv_cache/__init__.py     |  21 +++
 vllm_omni/bde/kv_cache/adapter.py      |  63 +++++++++
 vllm_omni/bde/kv_cache/chunk_window.py |  85 ++++++++++++
 vllm_omni/bde/kv_cache/config.py       |  35 +++++
 vllm_omni/bde/kv_cache/pool.py         |  68 ++++++++++
 7 files changed, 447 insertions(+)
 create mode 100644 tests/bde/test_kv_cache.py
 create mode 100644 vllm_omni/bde/kv_cache/__init__.py
 create mode 100644 vllm_omni/bde/kv_cache/adapter.py
 create mode 100644 vllm_omni/bde/kv_cache/chunk_window.py
 create mode 100644 vllm_omni/bde/kv_cache/config.py
 create mode 100644 vllm_omni/bde/kv_cache/pool.py

diff --git a/.gitignore b/.gitignore
index 40a6916b552..7a963cd68b6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -208,6 +208,7 @@ cache/
 .cache/
 diffusion_cache/
 kv_cache/
+!vllm_omni/bde/kv_cache/
 
 # Logs
 logs/
diff --git a/tests/bde/test_kv_cache.py b/tests/bde/test_kv_cache.py
new file mode 100644
index 00000000000..9cd16db92f8
--- /dev/null
+++ b/tests/bde/test_kv_cache.py
@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the BDE KV cache helpers (Phase 1, PR-2).
+
+Covers the request adapter, the chunk-window spec/manager (registration + the
+eviction policy), and the pool builder — exercised against the installed vLLM
+KV stack on CPU (block bookkeeping only, no GPU tensors).
+"""
+
+import pytest
+import torch
+
+from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.bde.kv_cache import (
+    BDEKVConfig,
+    BDERequestAdapter,
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    build_kv_manager,
+    compute_num_blocks,
+)
+from vllm_omni.bde.kv_cache.chunk_window import chunk_window_skipped_tokens
+
+BLOCK = 16
+
+
+def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+    )
+
+
+# --- ChunkWindowSpec registration -------------------------------------------
+
+
+def test_spec_registration_resolves_to_chunk_window_manager():
+    # Without explicit registration the MRO walk would fall back to the parent
+    # SlidingWindowManager; assert the subclass manager wins.
+    assert KVCacheSpecRegistry.get_manager_class(make_spec()) is ChunkWindowManager
+
+
+def test_spec_kind_is_sliding_window():
+    assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW
+
+
+def test_spec_rejects_inconsistent_window():
+    with pytest.raises(ValueError):
+        ChunkWindowSpec(
+            block_size=BLOCK,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float16,
+            sliding_window=99,  # != window_chunks * chunk_size
+            chunk_size=BLOCK,
+            window_chunks=2,
+        )
+
+
+# --- eviction policy (pure) -------------------------------------------------
+
+
+def test_sliding_replace_keeps_window():
+    # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
+    # the snap is to chunk boundaries.
+    skip = lambda n: chunk_window_skipped_tokens(
+        n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False
+    )
+    assert skip(32) == 0  # nothing past the window yet
+    assert skip(48) == 16  # one chunk fell out of the window
+    assert skip(64) == 32
+
+
+def test_sliding_replace_snaps_to_chunk_boundary():
+    # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
+    skip = chunk_window_skipped_tokens(
+        50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False
+    )
+    assert skip % 16 == 0 and skip == 16
+
+
+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
+
+
+def test_reset_at_boundary_drops_completed_past_sink():
+    skip = chunk_window_skipped_tokens(
+        48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True
+    )
+    # completed = 48; sink = 16 -> drop 32
+    assert skip == 32
+
+
+# --- BDEKVConfig ------------------------------------------------------------
+
+
+def test_kv_config_sliding_window_property():
+    assert BDEKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
+    assert BDEKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
+
+
+def test_kv_config_disabled_by_default():
+    assert BDEKVConfig().enable is False
+
+
+# --- BDERequestAdapter ------------------------------------------------------
+
+
+def test_adapter_advances_per_chunk_not_per_step():
+    a = BDERequestAdapter("r0", chunk_size=16)
+    assert a.num_computed_tokens == 0
+    assert a.num_tokens == 16  # in-flight chunk
+    a.on_chunk_committed()
+    assert a.num_computed_tokens == 16
+    assert a.num_tokens == 32
+
+
+def test_adapter_accounts_for_prefill_prefix():
+    a = BDERequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
+    assert a.num_computed_tokens == 4
+    assert a.num_prompt_tokens == 4
+    assert a.num_tokens == 20
+
+
+def test_adapter_status_is_vllm_enum():
+    assert isinstance(BDERequestAdapter("r0", chunk_size=16).status, RequestStatus)
+
+
+# --- pool / manager ---------------------------------------------------------
+
+
+def test_compute_num_blocks():
+    # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks.
+    assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
+
+
+def test_compute_num_blocks_validates():
+    with pytest.raises(ValueError):
+        compute_num_blocks(1 << 20, 0.5, 0)
+    with pytest.raises(ValueError):
+        compute_num_blocks(1 << 20, 1.5, 16 << 10)
+
+
+def test_build_manager_allocate_free_roundtrip():
+    """End-to-end: a BDERequestAdapter drives a real KVCacheManager.
+
+    This is the adapter conformance check — if the adapter were missing an
+    attribute the manager reads, allocate_slots/free would raise here.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
+    free_before = mgr.block_pool.get_num_free_blocks()
+
+    adapter = BDERequestAdapter("req-0", chunk_size=BLOCK)
+    blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
+    assert blocks is not None
+    assert mgr.block_pool.get_num_free_blocks() < free_before
+
+    mgr.free(adapter)
+    assert mgr.block_pool.get_num_free_blocks() == free_before
diff --git a/vllm_omni/bde/kv_cache/__init__.py b/vllm_omni/bde/kv_cache/__init__.py
new file mode 100644
index 00000000000..5a64851cccc
--- /dev/null
+++ b/vllm_omni/bde/kv_cache/__init__.py
@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+"""BDE engine-level KV cache helpers.
+
+Thin glue over vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` /
+``SlidingWindowManager``), used by the Block Diffusion Engine to manage KV for
+AR-diffusion models. See ``BDE_doc/dreamzero_kv_phase1_plan.md``.
+"""
+
+from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
+from vllm_omni.bde.kv_cache.chunk_window import ChunkWindowManager, ChunkWindowSpec
+from vllm_omni.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
+
+__all__ = [
+    "BDEKVConfig",
+    "BDERequestAdapter",
+    "ChunkWindowManager",
+    "ChunkWindowSpec",
+    "build_kv_manager",
+    "compute_num_blocks",
+]
diff --git a/vllm_omni/bde/kv_cache/adapter.py b/vllm_omni/bde/kv_cache/adapter.py
new file mode 100644
index 00000000000..2a639e9ce0f
--- /dev/null
+++ b/vllm_omni/bde/kv_cache/adapter.py
@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Adapt a BDE diffusion request to vLLM's ``KVCacheManager`` request surface."""
+
+from __future__ import annotations
+
+from vllm.v1.request import RequestStatus
+
+
+class BDERequestAdapter:
+    """Duck-types the subset of ``vllm.v1.request.Request`` that the
+    ``KVCacheManager`` reads (``allocate_slots`` / ``get_computed_blocks`` /
+    ``free`` and the coordinator they call into).
+
+    It is intentionally NOT a full ``Request``. The conformance test exercises a
+    real ``KVCacheManager`` against this adapter so the surface cannot silently
+    drift across vLLM versions.
+
+    A BDE request advances one *chunk* at a time: ``allocate_slots`` is called
+    once per chunk and ``num_computed_tokens`` advances only when a chunk is
+    committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk
+    reuse the same slots.
+    """
+
+    def __init__(
+        self,
+        request_id: str,
+        *,
+        chunk_size: int,
+        prefill_prefix_tokens: int = 0,
+    ) -> None:
+        self.request_id = request_id
+        self._chunk_size = chunk_size
+        self._prefill = prefill_prefix_tokens
+        self._completed_chunks = 0
+        # Filled only when cross-request prefix reuse is enabled (Phase 3).
+        self.block_hashes: list = []
+        self.skip_reading_prefix_cache = True
+        self.num_preemptions = 0
+        # vLLM watermark gate reads this; map the request lifecycle onto it.
+        self.status = RequestStatus.WAITING
+
+    @property
+    def num_computed_tokens(self) -> int:
+        """Persistent KV already materialized (committed chunks + prefill)."""
+        return self._prefill + self._completed_chunks * self._chunk_size
+
+    @property
+    def num_tokens(self) -> int:
+        """Total tokens once the in-flight chunk is committed."""
+        return self._prefill + (self._completed_chunks + 1) * self._chunk_size
+
+    @property
+    def num_prompt_tokens(self) -> int:
+        """The prefill prefix length (read by ``cache_blocks`` when caching)."""
+        return self._prefill
+
+    @property
+    def completed_chunks(self) -> int:
+        return self._completed_chunks
+
+    def on_chunk_committed(self) -> None:
+        """Advance by one chunk. Call once per chunk, not per denoise step."""
+        self._completed_chunks += 1
diff --git a/vllm_omni/bde/kv_cache/chunk_window.py b/vllm_omni/bde/kv_cache/chunk_window.py
new file mode 100644
index 00000000000..484c4782bda
--- /dev/null
+++ b/vllm_omni/bde/kv_cache/chunk_window.py
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Chunk-granularity sliding-window KV spec + manager for the BDE engine.
+
+A world-model keeps only the last ``W`` chunks of KV resident. That is a sliding
+window whose unit is a *chunk*, expressed as a ``SlidingWindowSpec`` subclass with
+``sliding_window = window_chunks * chunk_size`` plus a manager that evicts at
+chunk boundaries. Memory policy, refcounting, and ``null_block`` replacement stay
+in vLLM's ``BlockPool`` — only the token-skip math is overridden here.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
+from vllm.v1.kv_cache_interface import SlidingWindowSpec
+from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+
+
+def chunk_window_skipped_tokens(
+    num_computed_tokens: int,
+    *,
+    chunk_size: int,
+    sliding_window: int,
+    sink_chunks: int,
+    reset_at_boundary: bool,
+) -> int:
+    """Tokens outside the resident chunk window, snapped to a chunk boundary.
+
+    Pure function so the eviction policy is unit-testable without constructing a
+    manager. Two strategies:
+
+    - ``reset_at_boundary`` (DreamZero): at each chunk boundary everything past
+      the sink is dropped.
+    - otherwise (VGGT-style sliding replace): keep the last ``window`` tokens
+      (plus the sink); the skip count snaps down to a chunk boundary so a chunk
+      is never half-evicted.
+    """
+    sink = sink_chunks * chunk_size
+    if reset_at_boundary:
+        completed = (num_computed_tokens // chunk_size) * chunk_size
+        return max(0, completed - sink)
+    skipped = max(0, num_computed_tokens - sliding_window - sink)
+    return (skipped // chunk_size) * chunk_size
+
+
+class ChunkWindowManager(SlidingWindowManager):
+    """``SlidingWindowManager`` that evicts at chunk boundaries.
+
+    ``self.sliding_window`` is set by the base ``__init__``; the chunk fields are
+    read from ``self.kv_cache_spec`` (a :class:`ChunkWindowSpec`).
+    """
+
+    def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
+        spec = self.kv_cache_spec
+        return chunk_window_skipped_tokens(
+            num_computed_tokens,
+            chunk_size=spec.chunk_size,
+            sliding_window=self.sliding_window,
+            sink_chunks=spec.sink_chunks,
+            reset_at_boundary=spec.reset_at_boundary,
+        )
+
+
+# Register so KVCacheManager resolves ChunkWindowSpec to ChunkWindowManager.
+# Dispatch walks the spec's MRO, so without explicit registration the subclass
+# would silently fall back to the parent SlidingWindowManager (override ignored).
+# uniform_type_base_spec=None => its own KV cache group.
+@register_kv_cache_spec(manager_class=ChunkWindowManager, uniform_type_base_spec=None)
+@dataclass(frozen=True, kw_only=True)
+class ChunkWindowSpec(SlidingWindowSpec):
+    # sliding_window (inherited) MUST equal window_chunks * chunk_size.
+    chunk_size: int
+    window_chunks: int
+    sink_chunks: int = 0
+    reset_at_boundary: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.sliding_window != self.window_chunks * self.chunk_size:
+            raise ValueError(
+                "ChunkWindowSpec.sliding_window must equal "
+                f"window_chunks * chunk_size ({self.window_chunks} * "
+                f"{self.chunk_size}), got {self.sliding_window}"
+            )
diff --git a/vllm_omni/bde/kv_cache/config.py b/vllm_omni/bde/kv_cache/config.py
new file mode 100644
index 00000000000..84cb05780c1
--- /dev/null
+++ b/vllm_omni/bde/kv_cache/config.py
@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration for BDE engine-level KV cache management."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class BDEKVConfig:
+    """Settings for the BDE paged KV cache.
+
+    Disabled by default: when ``enable`` is False the BDE engine behaves exactly
+    like the base ``DiffusionEngine`` (no pool, no paged KV).
+    """
+
+    enable: bool = False
+    # Persistent KV tokens materialized per AR chunk. For DreamZero this is
+    # ``num_frame_per_block * frame_seqlen``.
+    chunk_size: int = 0
+    # Resident window in chunks. ``None`` means full attention (no eviction).
+    window_chunks: int | None = None
+    # Protected leading chunks (attention sink); never evicted.
+    sink_chunks: int = 0
+    # DreamZero-style window reset vs. VGGT-style sliding replace.
+    reset_at_boundary: bool = False
+    # Fraction of free device memory budgeted for the BDE KV pool.
+    gpu_memory_fraction: float = 0.1
+
+    @property
+    def sliding_window(self) -> int | None:
+        """Window size in tokens, or ``None`` for full attention."""
+        if self.window_chunks is None:
+            return None
+        return self.window_chunks * self.chunk_size
diff --git a/vllm_omni/bde/kv_cache/pool.py b/vllm_omni/bde/kv_cache/pool.py
new file mode 100644
index 00000000000..6619b8ecf0e
--- /dev/null
+++ b/vllm_omni/bde/kv_cache/pool.py
@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Construct a vLLM ``KVCacheManager`` for the BDE engine.
+
+Phase 1 uses static pool sizing (no profiling run): the worker passes the free
+device memory and a budget fraction. Construction is metadata-only — the backing
+KV tensors are allocated by the worker; here we only build the block bookkeeping
+so allocation / eviction / free flow through vLLM's ``KVCacheManager``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheSpec,
+    KVCacheTensor,
+)
+
+
+def compute_num_blocks(
+    available_bytes: int,
+    gpu_memory_fraction: float,
+    page_size_bytes: int,
+) -> int:
+    """Number of KV blocks that fit in ``fraction`` of the memory budget."""
+    if page_size_bytes <= 0:
+        raise ValueError(f"page_size_bytes must be positive, got {page_size_bytes}")
+    if not 0.0 < gpu_memory_fraction <= 1.0:
+        raise ValueError(f"gpu_memory_fraction must be in (0, 1], got {gpu_memory_fraction}")
+    budget = int(available_bytes * gpu_memory_fraction)
+    return max(0, budget // page_size_bytes)
+
+
+def build_kv_manager(
+    spec: KVCacheSpec,
+    layer_names: Sequence[str],
+    num_blocks: int,
+    max_model_len: int,
+    *,
+    enable_caching: bool = False,
+) -> KVCacheManager:
+    """Build a ``KVCacheManager`` with a single KV cache group for ``spec``.
+
+    Args:
+        spec: The KV cache spec for the group (e.g. a ``ChunkWindowSpec``).
+        layer_names: Attention layers sharing this group's block table.
+        num_blocks: Total physical blocks in the pool.
+        max_model_len: Upper bound on a request's sequence length.
+        enable_caching: Cross-request prefix caching (Phase 3); off in Phase 1.
+    """
+    layer_names = list(layer_names)
+    group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
+    tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)]
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=tensors,
+        kv_cache_groups=[group],
+    )
+    return KVCacheManager(
+        config,
+        max_model_len=max_model_len,
+        scheduler_block_size=spec.block_size,
+        hash_block_size=spec.block_size,
+        enable_caching=enable_caching,
+    )

From de280b51218134e4498ffaeb7e66ac7937331fae Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 09:03:26 +0000
Subject: [PATCH 03/51] [BDE] Test chunk-window eviction plateaus pool memory

Drive a real KVCacheManager over a 10-chunk rollout with BDERequestAdapter and
assert the pool free-block count plateaus once the window fills (out-of-window
blocks recycled via null_block) and the resident block set stays bounded by
window + current chunk. Guards the core bounded-memory promise of the chunk
window, which the pure-function / single-roundtrip tests did not cover.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_kv_cache.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/bde/test_kv_cache.py b/tests/bde/test_kv_cache.py
index 9cd16db92f8..ddb77974fab 100644
--- a/tests/bde/test_kv_cache.py
+++ b/tests/bde/test_kv_cache.py
@@ -172,3 +172,32 @@ def test_build_manager_allocate_free_roundtrip():
 
     mgr.free(adapter)
     assert mgr.block_pool.get_num_free_blocks() == free_before
+
+
+def test_window_eviction_plateaus_pool_memory():
+    """Multi-chunk rollout: out-of-window blocks are recycled, so pool memory
+    plateaus at ~window size instead of growing with rollout length.
+
+    This is the core promise of the chunk window — drive a real KVCacheManager
+    over many chunks and assert the free-block count stops decreasing once the
+    window is full, and the resident (non-null) blocks stay bounded.
+    """
+    spec = make_spec(chunk_size=BLOCK, window_chunks=2)  # window = 2 chunks
+    mgr = build_kv_manager(spec, ["l0"], num_blocks=64, max_model_len=4096)
+    null_id = mgr.block_pool.null_block.block_id
+    adapter = BDERequestAdapter("req", chunk_size=BLOCK)
+
+    free_after = []
+    distinct_real = []
+    for _ in range(10):
+        mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
+        ids = mgr.get_block_ids(adapter.request_id)[0]
+        distinct_real.append(len({b for b in ids if b != null_id}))
+        free_after.append(mgr.block_pool.get_num_free_blocks())
+        adapter.on_chunk_committed()
+
+    # Free-block count stops decreasing once the window is full (warmup = window).
+    assert free_after[-1] == free_after[spec.window_chunks]
+    assert free_after[-1] == free_after[-5]  # flat across the tail
+    # Resident non-null blocks bounded by window + the in-flight chunk.
+    assert max(distinct_real) <= spec.window_chunks + 1

From fe7600f2b9b9d766ea380592780ea1ed7f6dfcbc Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 09:46:53 +0000
Subject: [PATCH 04/51] [BDE] Add KV slot mapping, orchestrator, and model
 runner
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- slot_mapping.py: PagedAttention slot mapping for BDE's single-request layout
  (compute_slot_mapping / chunk_slot_mapping / resident_block_ids).
- manager.py: BDEKVCache — the engine-level KV orchestrator. Owns the vLLM
  KVCacheManager (one chunk-window group) + per-request adapter, and exposes the
  rollout ops (begin_request -> allocate_chunk / chunk_write_slots / commit_chunk
  / window_block_ids -> end_request). Lives in the runner (worker/GPU side).
- runner.py: BDEModelRunner(DiffusionModelRunner) — owns the BDEKVCache and
  brackets execute_model with the per-request KV lifecycle. Behavior-preserving
  when KV is disabled (the Phase-1 default).
- Tests: slot mapping + a full BDEKVCache lifecycle with eviction-plateau.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_kv_orchestration.py     |  66 +++++++++++++
 tests/bde/test_slot_mapping.py         |  78 +++++++++++++++
 vllm_omni/bde/kv_cache/__init__.py     |  10 ++
 vllm_omni/bde/kv_cache/manager.py      | 126 +++++++++++++++++++++++++
 vllm_omni/bde/kv_cache/slot_mapping.py |  71 ++++++++++++++
 vllm_omni/bde/runner.py                |  98 +++++++++++++++++++
 6 files changed, 449 insertions(+)
 create mode 100644 tests/bde/test_kv_orchestration.py
 create mode 100644 tests/bde/test_slot_mapping.py
 create mode 100644 vllm_omni/bde/kv_cache/manager.py
 create mode 100644 vllm_omni/bde/kv_cache/slot_mapping.py
 create mode 100644 vllm_omni/bde/runner.py

diff --git a/tests/bde/test_kv_orchestration.py b/tests/bde/test_kv_orchestration.py
new file mode 100644
index 00000000000..aef124af6b2
--- /dev/null
+++ b/tests/bde/test_kv_orchestration.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for BDEKVCache — the engine-level KV orchestration body (Phase 1)."""
+
+import pytest
+import torch
+
+from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
+
+DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16)
+
+
+def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
+    cfg = BDEKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+    return BDEKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
+
+
+def test_requires_enabled_config():
+    with pytest.raises(ValueError):
+        BDEKVCache(BDEKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_requires_bounded_window():
+    cfg = BDEKVConfig(enable=True, chunk_size=16, window_chunks=None)
+    with pytest.raises(ValueError):
+        BDEKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_full_request_lifecycle_and_eviction():
+    """begin -> per-chunk allocate/slots/commit over a long rollout -> free.
+
+    Exercises the orchestrator end-to-end and asserts the chunk window bounds
+    memory (pool plateaus) and frees cleanly.
+    """
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    adapter = kv.begin_request("req-0")
+    free_after = []
+    for k in range(10):
+        block_table = kv.allocate_chunk(adapter)
+        slots = kv.chunk_write_slots(adapter)
+        # Slots target real blocks in this chunk's table; length == chunk_size.
+        assert len(slots) == kv.spec.chunk_size
+        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.
+        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)
+
+    # Pool memory plateaus once the window is full (eviction recycles blocks).
+    assert free_after[-1] == free_after[kv.spec.window_chunks]
+
+    kv.end_request(adapter)
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_num_computed_advances_per_chunk():
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    a = kv.begin_request("r")
+    assert a.num_computed_tokens == 0
+    kv.allocate_chunk(a)
+    kv.commit_chunk(a)
+    assert a.num_computed_tokens == 16
+    kv.end_request(a)
diff --git a/tests/bde/test_slot_mapping.py b/tests/bde/test_slot_mapping.py
new file mode 100644
index 00000000000..4eabb4c6fb9
--- /dev/null
+++ b/tests/bde/test_slot_mapping.py
@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Slot-mapping / block-table tests for the BDE engine (Phase 1, Step 3)."""
+
+import pytest
+import torch
+
+from vllm_omni.bde.kv_cache import (
+    BDERequestAdapter,
+    ChunkWindowSpec,
+    build_kv_manager,
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    resident_block_ids,
+)
+
+BLOCK = 16
+
+
+def make_spec(*, window_chunks=2):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * BLOCK,
+        chunk_size=BLOCK,
+        window_chunks=window_chunks,
+    )
+
+
+def test_slot_mapping_matches_block_offsets():
+    # block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4).
+    slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4)
+    assert slots.tolist() == [20, 21, 8, 9, 36, 37]
+
+
+def test_slot_mapping_rejects_bad_block_size():
+    with pytest.raises(ValueError):
+        compute_slot_mapping([1], [0], block_size=0)
+
+
+def test_chunk_slot_mapping_targets_current_chunk():
+    # After 2 committed chunks (num_computed=32), the in-flight chunk maps to the
+    # block covering positions 32..47 — i.e. block table index 2.
+    block_ids = [7, 4, 3]  # block index 2 -> physical block 3
+    slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16)
+    # positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51
+    assert slots.tolist() == [48, 49, 50, 51]
+
+
+def test_resident_block_ids_excludes_null():
+    assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3]
+
+
+def test_blocktable_build_single_request():
+    """Integration: a real KVCacheManager block table -> a valid slot mapping.
+
+    Every slot for the in-flight chunk must land inside a real (non-null) block
+    that the manager actually allocated for the request.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
+    null_id = mgr.block_pool.null_block.block_id
+    adapter = BDERequestAdapter("req", chunk_size=BLOCK)
+
+    mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
+    block_ids = mgr.get_block_ids(adapter.request_id)[0]
+    real = set(resident_block_ids(block_ids, null_id))
+    assert real, "expected at least one real block after allocation"
+
+    slots = chunk_slot_mapping(
+        block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size
+    )
+    # Each slot resolves to a real allocated block, none to null_block.
+    blocks_used = {int(s) // spec.block_size for s in slots}
+    assert blocks_used <= real
+    assert null_id not in blocks_used
+    assert len(slots) == spec.chunk_size
diff --git a/vllm_omni/bde/kv_cache/__init__.py b/vllm_omni/bde/kv_cache/__init__.py
index 5a64851cccc..e3d6184cbf1 100644
--- a/vllm_omni/bde/kv_cache/__init__.py
+++ b/vllm_omni/bde/kv_cache/__init__.py
@@ -9,13 +9,23 @@
 from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
 from vllm_omni.bde.kv_cache.chunk_window import ChunkWindowManager, ChunkWindowSpec
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.bde.kv_cache.manager import BDEKVCache
 from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
+from vllm_omni.bde.kv_cache.slot_mapping import (
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    resident_block_ids,
+)
 
 __all__ = [
+    "BDEKVCache",
     "BDEKVConfig",
     "BDERequestAdapter",
     "ChunkWindowManager",
     "ChunkWindowSpec",
     "build_kv_manager",
+    "chunk_slot_mapping",
     "compute_num_blocks",
+    "compute_slot_mapping",
+    "resident_block_ids",
 ]
diff --git a/vllm_omni/bde/kv_cache/manager.py b/vllm_omni/bde/kv_cache/manager.py
new file mode 100644
index 00000000000..0b4000987da
--- /dev/null
+++ b/vllm_omni/bde/kv_cache/manager.py
@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""BDEKVCache — the engine-level KV cache orchestrator for one BDE model.
+
+This is the *body* of BDE's KV management: it owns a vLLM ``KVCacheManager`` (a
+single chunk-window group) and the per-request adapter lifecycle, and exposes the
+per-chunk operations a rollout needs — allocate, slot mapping, commit, window
+lookup, free. It lives in the model runner (worker / GPU side), co-located with
+the model and the KV tensors; the DreamZero pipeline calls these methods during a
+rollout. The main-process ``BDEEngine`` only selects the engine and is otherwise
+thin.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
+from vllm_omni.bde.kv_cache.chunk_window import ChunkWindowSpec
+from vllm_omni.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
+from vllm_omni.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
+
+
+class BDEKVCache:
+    """Owns the paged KV pool + per-request lifecycle for a BDE model.
+
+    Build once per loaded model (dimensions known); then per request:
+    ``begin_request`` → per chunk (``allocate_chunk`` → ``chunk_write_slots`` →
+    [model writes K/V] → ``commit_chunk``) → ``end_request``.
+    """
+
+    def __init__(
+        self,
+        config: BDEKVConfig,
+        *,
+        num_layers: int,
+        num_kv_heads: int,
+        head_size: int,
+        dtype: torch.dtype,
+        block_size: int,
+        max_model_len: int,
+        available_bytes: int,
+    ) -> None:
+        if not config.enable:
+            raise ValueError("BDEKVCache built with a disabled BDEKVConfig")
+        if config.window_chunks is None:
+            raise ValueError("Phase 1 requires a bounded window (window_chunks)")
+        if config.chunk_size <= 0:
+            raise ValueError("BDEKVConfig.chunk_size must be set (> 0)")
+
+        self.config = config
+        self.block_size = block_size
+        self.spec = ChunkWindowSpec(
+            block_size=block_size,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=dtype,
+            sliding_window=config.window_chunks * config.chunk_size,
+            chunk_size=config.chunk_size,
+            window_chunks=config.window_chunks,
+            sink_chunks=config.sink_chunks,
+            reset_at_boundary=config.reset_at_boundary,
+        )
+        # Each pool block spans all layers' K/V, so size against the per-layer
+        # page size times the layer count.
+        num_blocks = compute_num_blocks(
+            available_bytes,
+            config.gpu_memory_fraction,
+            self.spec.page_size_bytes * num_layers,
+        )
+        layer_names = [f"bde.layer.{i}" for i in range(num_layers)]
+        self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
+        self.num_blocks = num_blocks
+        self.null_block_id = self.manager.block_pool.null_block.block_id
+        self._adapters: dict[str, BDERequestAdapter] = {}
+
+    # -- request lifecycle ---------------------------------------------------
+
+    def begin_request(
+        self, request_id: str, *, prefill_prefix_tokens: int = 0
+    ) -> BDERequestAdapter:
+        adapter = BDERequestAdapter(
+            request_id,
+            chunk_size=self.spec.chunk_size,
+            prefill_prefix_tokens=prefill_prefix_tokens,
+        )
+        self._adapters[request_id] = adapter
+        return adapter
+
+    def end_request(self, adapter: BDERequestAdapter) -> None:
+        self.manager.free(adapter)
+        self._adapters.pop(adapter.request_id, None)
+
+    # -- per-chunk operations ------------------------------------------------
+
+    def allocate_chunk(self, adapter: BDERequestAdapter) -> list[int]:
+        """Allocate a chunk's blocks (evicting out-of-window blocks first).
+
+        Returns the request's full block table (incl. null_block placeholders).
+        """
+        blocks = self.manager.allocate_slots(adapter, num_new_tokens=self.spec.chunk_size)
+        if blocks is None:
+            # Phase 2 turns this into scheduler back-pressure; for the Phase-1
+            # single-request demo a full pool is a hard error.
+            raise RuntimeError("BDE KV pool exhausted while allocating a chunk")
+        return self.block_table(adapter)
+
+    def block_table(self, adapter: BDERequestAdapter) -> list[int]:
+        return list(self.manager.get_block_ids(adapter.request_id)[0])
+
+    def chunk_write_slots(self, adapter: BDERequestAdapter) -> torch.Tensor:
+        """Slot mapping for the in-flight chunk — the K/V write target."""
+        return chunk_slot_mapping(
+            self.block_table(adapter),
+            adapter.num_computed_tokens,
+            self.spec.chunk_size,
+            self.block_size,
+        )
+
+    def window_block_ids(self, adapter: BDERequestAdapter) -> list[int]:
+        """Resident (non-null) blocks the read path gathers the window from."""
+        return resident_block_ids(self.block_table(adapter), self.null_block_id)
+
+    def commit_chunk(self, adapter: BDERequestAdapter) -> None:
+        """Advance after the chunk's K/V is written. Once per chunk, not per step."""
+        adapter.on_chunk_committed()
diff --git a/vllm_omni/bde/kv_cache/slot_mapping.py b/vllm_omni/bde/kv_cache/slot_mapping.py
new file mode 100644
index 00000000000..6ca40138664
--- /dev/null
+++ b/vllm_omni/bde/kv_cache/slot_mapping.py
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Slot mapping + block-table access for the BDE engine.
+
+Phase 1 is single-request and owns its KV layout — it gathers the resident window
+into DreamZero's existing attention rather than calling vLLM's paged kernel — so a
+thin slot-mapping helper is used instead of vLLM's batch/Triton ``BlockTables``.
+The layout is the standard PagedAttention one:
+
+    slot(pos) = block_id(pos) * block_size + (pos % block_size)
+
+Write (commit) and read (gather) use the same formula, so the layout is
+self-consistent regardless of vLLM's internal kernel layout.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+
+def compute_slot_mapping(
+    block_ids: Sequence[int],
+    positions: torch.Tensor | Sequence[int],
+    block_size: int,
+) -> torch.Tensor:
+    """Map absolute token positions to physical KV-cache slots.
+
+    Args:
+        block_ids: physical block id per block index (the request's block table).
+        positions: absolute token positions to map (1-D).
+        block_size: tokens per block.
+
+    Returns:
+        ``LongTensor`` of physical slots, one per position.
+    """
+    if block_size <= 0:
+        raise ValueError(f"block_size must be positive, got {block_size}")
+    table = torch.as_tensor(block_ids, dtype=torch.long)
+    pos = torch.as_tensor(positions, dtype=torch.long)
+    block_index = torch.div(pos, block_size, rounding_mode="floor")
+    offset = pos % block_size
+    return table[block_index] * block_size + offset
+
+
+def chunk_slot_mapping(
+    block_ids: Sequence[int],
+    num_computed_tokens: int,
+    chunk_size: int,
+    block_size: int,
+) -> torch.Tensor:
+    """Slot mapping for the in-flight chunk's tokens (the commit write target).
+
+    The chunk occupies absolute positions
+    ``[num_computed_tokens, num_computed_tokens + chunk_size)``.
+    """
+    positions = torch.arange(
+        num_computed_tokens,
+        num_computed_tokens + chunk_size,
+        dtype=torch.long,
+    )
+    return compute_slot_mapping(block_ids, positions, block_size)
+
+
+def resident_block_ids(block_ids: Sequence[int], null_block_id: int) -> list[int]:
+    """Real (non-null) blocks currently resident, in table order.
+
+    These are the blocks the read path gathers the attention window from;
+    out-of-window positions are the shared ``null_block`` and are excluded.
+    """
+    return [int(b) for b in block_ids if int(b) != null_block_id]
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/bde/runner.py
new file mode 100644
index 00000000000..0cb2e71536b
--- /dev/null
+++ b/vllm_omni/bde/runner.py
@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+"""BDEModelRunner — the diffusion model runner for the BDE engine.
+
+Subclasses ``DiffusionModelRunner`` and owns the engine-level KV cache
+(:class:`BDEKVCache`). It brackets a request's rollout with the KV lifecycle
+(``begin_request`` / ``end_request``) and exposes the live ``BDEKVCache`` so the
+model pipeline (DreamZero) can drive the per-chunk allocate / slot-mapping /
+gather / commit operations during ``pipeline.forward``.
+
+When KV management is disabled (the Phase-1 default), the runner is
+behavior-preserving — it simply defers to the base ``DiffusionModelRunner``.
+"""
+
+from __future__ import annotations
+
+from vllm.logger import init_logger
+
+from vllm_omni.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.bde.kv_cache.manager import BDEKVCache
+from vllm_omni.diffusion.data import OmniDiffusionConfig, DiffusionOutput
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+
+logger = init_logger(__name__)
+
+
+def resolve_bde_kv_config(od_config: OmniDiffusionConfig) -> BDEKVConfig:
+    """Read the BDE KV config off the diffusion config (disabled if absent).
+
+    The field name is not yet fixed on ``OmniDiffusionConfig`` (it must avoid the
+    existing ``diffusion_kv_cache_*`` collision), so we look it up leniently and
+    default to a disabled config — keeping the runner behavior-preserving until
+    the field is wired.
+    """
+    raw = getattr(od_config, "bde_kv_config", None)
+    if raw is None:
+        raw = getattr(od_config, "kv_block_config", None)
+    if isinstance(raw, BDEKVConfig):
+        return raw
+    if isinstance(raw, dict):
+        return BDEKVConfig(**raw)
+    return BDEKVConfig()
+
+
+class BDEModelRunner(DiffusionModelRunner):
+    def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None:
+        super().__init__(vllm_config, od_config, device)
+        self.bde_kv_config = resolve_bde_kv_config(od_config)
+        # Built after the model is loaded (dimensions known); stays None while KV
+        # management is disabled.
+        self.kv_cache: BDEKVCache | None = None
+
+    def build_kv_cache(
+        self,
+        *,
+        num_layers: int,
+        num_kv_heads: int,
+        head_size: int,
+        dtype,
+        block_size: int,
+        max_model_len: int,
+        available_bytes: int,
+    ) -> None:
+        """Construct the BDEKVCache once the model's dimensions are known.
+
+        No-op when KV management is disabled.
+        """
+        if not self.bde_kv_config.enable:
+            return
+        self.kv_cache = BDEKVCache(
+            self.bde_kv_config,
+            num_layers=num_layers,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=dtype,
+            block_size=block_size,
+            max_model_len=max_model_len,
+            available_bytes=available_bytes,
+        )
+        logger.info(
+            "BDE KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s",
+            self.kv_cache.num_blocks,
+            self.bde_kv_config.chunk_size,
+            self.bde_kv_config.window_chunks,
+        )
+
+    def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        # KV disabled -> base behavior, unchanged.
+        if self.kv_cache is None:
+            return super().execute_model(req)
+
+        # Bracket the rollout: the pipeline drives per-chunk ops against
+        # self.kv_cache during pipeline.forward; we own begin/end here.
+        adapter = self.kv_cache.begin_request(req.request_id)
+        try:
+            return super().execute_model(req)
+        finally:
+            self.kv_cache.end_request(adapter)

From 112a8077af3d376b8118303b575db829569fbaac Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 09:47:07 +0000
Subject: [PATCH 05/51] [BDE] Wire BDE engine to BDEModelRunner via od_config
 override

BDEEngine.__init__ sets od_config.diffusion_model_runner_cls (to BDEModelRunner)
before the base __init__ builds the executor and spawns workers; the config is
propagated to each worker process. DiffusionWorker now reads this override and
falls back to the platform default when unset, so only BDE-routed models build
BDEModelRunner and every other model is unaffected.

worker_main hardcodes WorkerProc and both the worker/runner classes are
platform-selected, so a config-threaded runner override is the minimal, isolated
wiring (no empty BDEWorker/BDEExecutor subclasses needed). Adds the
diffusion_model_runner_cls config field and the gated worker hook.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_engine_routing.py              | 41 ++++++++++++++++++-
 vllm_omni/bde/engine.py                       | 36 ++++++++++++++--
 vllm_omni/diffusion/data.py                   |  5 +++
 .../diffusion/worker/diffusion_worker.py      |  9 +++-
 4 files changed, 85 insertions(+), 6 deletions(-)

diff --git a/tests/bde/test_engine_routing.py b/tests/bde/test_engine_routing.py
index a53eb758c78..b5550e3ee39 100644
--- a/tests/bde/test_engine_routing.py
+++ b/tests/bde/test_engine_routing.py
@@ -12,7 +12,13 @@
 
 import pytest
 
-from vllm_omni.bde.engine import BDEEngine
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+from vllm_omni.bde.engine import (
+    BDE_MODEL_RUNNER_CLS,
+    BDEEngine,
+    apply_bde_runner_default,
+)
 from vllm_omni.diffusion.data import (
     OmniDiffusionConfig,
     default_engine_backend_for_model,
@@ -94,3 +100,36 @@ def test_dreamzero_default_routes_to_bde_engine():
     # to BDEEngine through the same dispatcher used by make_engine.
     backend = default_engine_backend_for_model("DreamZeroPipeline")
     assert DiffusionEngine.resolve_engine_class(_cfg(backend)) is BDEEngine
+
+
+# --- BDE worker/runner wiring -----------------------------------------------
+
+
+def test_bde_model_runner_cls_resolves():
+    cls = resolve_obj_by_qualname(BDE_MODEL_RUNNER_CLS)
+    assert cls.__name__ == "BDEModelRunner"
+
+
+def test_apply_runner_default_sets_when_unset():
+    cfg = SimpleNamespace(diffusion_model_runner_cls=None)
+    apply_bde_runner_default(cfg)
+    assert cfg.diffusion_model_runner_cls == BDE_MODEL_RUNNER_CLS
+
+
+def test_apply_runner_default_respects_explicit_choice():
+    cfg = SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner")
+    apply_bde_runner_default(cfg)
+    assert cfg.diffusion_model_runner_cls == "my.custom.Runner"
+
+
+def test_config_runner_cls_field_defaults_none():
+    field = {f.name: f for f in fields(OmniDiffusionConfig)}["diffusion_model_runner_cls"]
+    assert field.default is None
+
+
+def test_worker_runner_selection_prefers_override():
+    # Mirrors the worker hook: an od_config override wins over the platform default.
+    over = SimpleNamespace(diffusion_model_runner_cls=BDE_MODEL_RUNNER_CLS)
+    assert (getattr(over, "diffusion_model_runner_cls", None) or "PLATFORM") == BDE_MODEL_RUNNER_CLS
+    unset = SimpleNamespace(diffusion_model_runner_cls=None)
+    assert (getattr(unset, "diffusion_model_runner_cls", None) or "PLATFORM") == "PLATFORM"
diff --git a/vllm_omni/bde/engine.py b/vllm_omni/bde/engine.py
index 69e0c1560c7..2df0acdd173 100644
--- a/vllm_omni/bde/engine.py
+++ b/vllm_omni/bde/engine.py
@@ -5,10 +5,27 @@
 
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.sched import SchedulerInterface
 
 logger = init_logger(__name__)
 
+#: Import path of the runner the BDE engine routes its workers to.
+BDE_MODEL_RUNNER_CLS = "vllm_omni.bde.runner.BDEModelRunner"
+
+
+def apply_bde_runner_default(od_config: OmniDiffusionConfig) -> None:
+    """Route the BDE engine's workers to ``BDEModelRunner``.
+
+    Sets ``od_config.diffusion_model_runner_cls`` in place, unless the caller
+    already chose a runner. The worker (in its own process) reads this off the
+    propagated ``od_config`` and builds the BDE runner instead of the platform
+    default, keeping the swap scoped to BDE-routed models.
+    """
+    if getattr(od_config, "diffusion_model_runner_cls", None) is None:
+        od_config.diffusion_model_runner_cls = BDE_MODEL_RUNNER_CLS
+
 
 class BDEEngine(DiffusionEngine):
     """AR-Diffusion engine with engine-level KV cache management.
@@ -24,7 +41,20 @@ class BDEEngine(DiffusionEngine):
     (resolved by :meth:`DiffusionEngine.make_engine`), so models that do not opt
     in keep using the base ``DiffusionEngine`` unchanged.
 
-    Phase 1 is a behavior-preserving subclass: it establishes the engine seam and
-    isolation boundary. The KV-cache lifecycle (pool construction, per-chunk
-    ``allocate_slots`` / ``free``) is added in later phases.
+    Architecture note: in the multiproc setup the KV cache lives in the worker /
+    runner process (GPU side), co-located with the model and KV tensors — so the
+    actual KV *body* is :class:`~vllm_omni.bde.kv_cache.manager.BDEKVCache`, owned
+    by :class:`~vllm_omni.bde.runner.BDEModelRunner`. ``BDEEngine`` itself is the
+    thin selection / injection seam; it wires the BDE executor → worker → runner
+    so DreamZero's rollout runs against the runner-owned KV cache.
     """
+
+    def __init__(
+        self,
+        od_config: OmniDiffusionConfig,
+        scheduler: SchedulerInterface | None = None,
+    ) -> None:
+        # Route this engine's workers to BDEModelRunner before the base __init__
+        # builds the executor and spawns workers (od_config is propagated to them).
+        apply_bde_runner_default(od_config)
+        super().__init__(od_config, scheduler=scheduler)
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 3c1323bbc42..1b713f2c062 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -554,6 +554,11 @@ class OmniDiffusionConfig:
     #                import path string.
     engine_backend: str = "default"
 
+    # Optional override for the diffusion model runner class (import path).
+    # The worker uses it instead of the platform default when set; the BDE engine
+    # sets it to BDEModelRunner. Unset -> platform default (existing behavior).
+    diffusion_model_runner_cls: str | None = None
+
     # HuggingFace specific parameters
     trust_remote_code: bool = False
     revision: str | None = None
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index 0b4e877180a..f2ad13d5076 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -203,8 +203,13 @@ def __init__(
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
-        # Create model runner using the platform-specified class
-        model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls()
+        # Create model runner. An od_config override (e.g. the BDE engine selecting
+        # BDEModelRunner) takes precedence; unset falls back to the platform default,
+        # so existing models are unaffected.
+        model_runner_cls_path = (
+            getattr(self.od_config, "diffusion_model_runner_cls", None)
+            or current_omni_platform.get_diffusion_model_runner_cls()
+        )
         model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path)
         self.model_runner = model_runner_cls(
             vllm_config=self.vllm_config,

From 99c97697ad62ca88bc92f18afbfa30dfb62353c6 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 09:47:17 +0000
Subject: [PATCH 06/51] [BDE] Add 2-GPU full-chain wiring test

Spawn a rank-0/rank-1 DiffusionWorker pair (world_size=2, CFG-parallel, NCCL
rendezvous) with the BDE runner override threaded through od_config, and assert
both workers build BDEModelRunner on their own GPU. skip_load_model avoids
needing model weights, so this verifies the engine->worker->runner wiring across
real process + GPU boundaries. Gated to >= 2 GPUs; run explicitly with
CUDA_VISIBLE_DEVICES=0,1.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_full_chain_gpu.py | 82 ++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)
 create mode 100644 tests/bde/test_full_chain_gpu.py

diff --git a/tests/bde/test_full_chain_gpu.py b/tests/bde/test_full_chain_gpu.py
new file mode 100644
index 00000000000..3e6b904c0b4
--- /dev/null
+++ b/tests/bde/test_full_chain_gpu.py
@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Full-chain BDE wiring test on real GPUs (GPU 0,1).
+
+Spawns a rank-0 / rank-1 ``DiffusionWorker`` pair (world_size=2, NCCL rendezvous)
+with the BDE runner override threaded through ``od_config``, and asserts both
+workers build ``BDEModelRunner`` on their own GPU. ``skip_load_model=True`` avoids
+needing DreamZero weights — this verifies the engine→worker→runner wiring across
+real process + GPU boundaries, not a full rollout.
+
+Run explicitly (needs >= 2 GPUs):
+    CUDA_VISIBLE_DEVICES=0,1 python -m pytest tests/bde/test_full_chain_gpu.py -v -s
+"""
+
+import multiprocessing as mp
+
+import pytest
+import torch
+
+requires_two_gpus = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.device_count() < 2,
+    reason="requires >= 2 GPUs",
+)
+
+RUNNER_OVERRIDE = "vllm_omni.bde.runner.BDEModelRunner"
+
+
+def _worker_entry(rank: int, od_config, q) -> None:
+    try:
+        from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+        worker = DiffusionWorker(
+            local_rank=rank, rank=rank, od_config=od_config, skip_load_model=True
+        )
+        q.put((rank, str(worker.device), type(worker.model_runner).__name__, None))
+    except Exception:  # pragma: no cover - reported to the parent
+        import traceback
+
+        q.put((rank, None, None, traceback.format_exc()))
+
+
+@requires_two_gpus
+def test_bde_full_chain_two_gpu():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    od_config = OmniDiffusionConfig.from_kwargs(
+        model="dummy",
+        model_class_name="DummyPipeline",
+        num_gpus=2,
+        # CFG-parallel across the 2 GPUs (the natural diffusion 2-way split).
+        parallel_config={"cfg_parallel_size": 2},
+        diffusion_load_format="dummy",
+        diffusion_model_runner_cls=RUNNER_OVERRIDE,
+    )
+    # Both ranks share this config (same master_port) so they rendezvous.
+    assert od_config.diffusion_model_runner_cls == RUNNER_OVERRIDE
+
+    ctx = mp.get_context("spawn")
+    q = ctx.Queue()
+    procs = [ctx.Process(target=_worker_entry, args=(rank, od_config, q)) for rank in (0, 1)]
+    for p in procs:
+        p.start()
+
+    results = {}
+    try:
+        for _ in range(2):
+            rank, device, runner, err = q.get(timeout=300)
+            results[rank] = (device, runner, err)
+    finally:
+        for p in procs:
+            p.join(timeout=60)
+            if p.is_alive():
+                p.terminate()
+
+    for rank in (0, 1):
+        assert rank in results, f"rank {rank} produced no result (hung?)"
+        device, runner, err = results[rank]
+        assert err is None, f"rank {rank} worker construction failed:\n{err}"
+        assert runner == "BDEModelRunner", f"rank {rank} built {runner}, expected BDEModelRunner"
+
+    # Each rank initialized its own GPU.
+    assert results[0][0] == "cuda:0"
+    assert results[1][0] == "cuda:1"

From d0cae958a9eefa8a1558da1df513882a98d5f312 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 15:06:09 +0000
Subject: [PATCH 07/51] [BDE] Add pool-backed KV write + window gather (Step 4)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Phase 1 keeps DreamZero's existing attention kernel; this only changes *where*
the past-window KV comes from (the pool) and *where* committed KV is stored.

- gather.py: allocate_kv_pool (per-layer flat pool indexed by slot_mapping),
  pool_write_chunk (scatter committed K/V into pool at slot_mapping),
  pool_gather_window (materialize resident window as contiguous (2,1,window,...)
  tensor — exact DreamZero kv_cache format), BDEPipelineMixin skeleton.
- BDEKVCache now owns the per-layer GPU pools (allocated on construction) and
  exposes write_chunk_kv(layer, new_k, new_v, adapter) + gather_window(layer,
  adapter) — both resolve the real block table through the orchestrator.
- 4 unit tests: write+gather roundtrip, multi-block gather, window trim, and
  an integration test driving full BDEKVCache (write->readback via the
  orchestrator's chunk_write_slots, gather_window shape check).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_gather.py           | 103 ++++++++++++++++++++++++
 vllm_omni/bde/kv_cache/__init__.py |  10 +++
 vllm_omni/bde/kv_cache/gather.py   | 125 +++++++++++++++++++++++++++++
 vllm_omni/bde/kv_cache/manager.py  |  49 +++++++++++
 4 files changed, 287 insertions(+)
 create mode 100644 tests/bde/test_gather.py
 create mode 100644 vllm_omni/bde/kv_cache/gather.py

diff --git a/tests/bde/test_gather.py b/tests/bde/test_gather.py
new file mode 100644
index 00000000000..ae4e1a7bc3f
--- /dev/null
+++ b/tests/bde/test_gather.py
@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for pool write + window gather (Step 4)."""
+
+import torch
+from vllm_omni.bde.kv_cache import (
+    BDEKVCache,
+    BDEKVConfig,
+    BDERequestAdapter,
+    allocate_kv_pool,
+    pool_gather_window,
+    pool_write_chunk,
+)
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def test_pool_write_gather_roundtrip():
+    """Write a single chunk to the pool then gather — K/V comes back identical."""
+    num_blocks = 4
+    kp, vp = allocate_kv_pool(num_blocks, BLOCK, 1, N_HEADS, HEAD_DIM, torch.float32, "cpu")
+    new_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
+    new_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
+    slot_mapping = torch.arange(num_blocks * BLOCK)  # identity mapping for the test
+    slot_mapping = slot_mapping[:BLOCK]
+
+    pool_write_chunk(kp[0], vp[0], new_k, new_v, slot_mapping)
+    window = pool_gather_window(kp[0], vp[0], window_block_ids=[0], block_size=BLOCK, max_attention_size=BLOCK)
+
+    assert window.shape == (2, 1, BLOCK, N_HEADS, HEAD_DIM)
+    assert torch.allclose(window[0, 0], new_k[0])
+    assert torch.allclose(window[1, 0], new_v[0])
+
+
+def test_pool_write_gather_multiple_blocks():
+    """Write 2 blocks (distinct positions) and gather both from the window."""
+    num_blocks = 8
+    kp, vp = allocate_kv_pool(num_blocks, BLOCK, 1, N_HEADS, HEAD_DIM, torch.float32, "cpu")
+
+    # Write block 0
+    k0 = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
+    v0 = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
+    s0 = torch.arange(BLOCK)  # block 0 starts at slot 0
+    pool_write_chunk(kp[0], vp[0], k0, v0, s0)
+
+    # Write block 2 (holes OK, block id doesn't need to be contiguous in slot space)
+    k2 = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
+    v2 = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
+    s2 = 2 * BLOCK + torch.arange(BLOCK)  # block 2
+    pool_write_chunk(kp[0], vp[0], k2, v2, s2)
+
+    # Gather blocks 0 and 2 — they appear in table order.
+    window = pool_gather_window(kp[0], vp[0], window_block_ids=[0, 2], block_size=BLOCK, max_attention_size=2 * BLOCK)
+    assert window.shape == (2, 1, 2 * BLOCK, N_HEADS, HEAD_DIM)
+    assert torch.allclose(window[0, 0, :BLOCK], k0[0])
+    assert torch.allclose(window[1, 0, :BLOCK], v0[0])
+    assert torch.allclose(window[0, 0, BLOCK:], k2[0])
+    assert torch.allclose(window[1, 0, BLOCK:], v2[0])
+
+
+def test_pool_write_gather_window_trim():
+    """Gather with max_attention_size smaller than the resident window trims to tail."""
+    num_blocks = 4
+    kp, vp = allocate_kv_pool(num_blocks, BLOCK, 1, N_HEADS, HEAD_DIM, torch.float32, "cpu")
+    k = torch.randn(1, 2 * BLOCK, N_HEADS, HEAD_DIM)
+    v = torch.randn(1, 2 * BLOCK, N_HEADS, HEAD_DIM)
+    slots = torch.cat([torch.arange(0 * BLOCK, 1 * BLOCK), torch.arange(1 * BLOCK, 2 * BLOCK)])
+    pool_write_chunk(kp[0], vp[0], k[:, :BLOCK], v[:, :BLOCK], slots[:BLOCK])  # block 0
+    pool_write_chunk(kp[0], vp[0], k[:, BLOCK:], v[:, BLOCK:], slots[BLOCK:])  # block 1
+    # max_attention = BLOCK — only last block.
+    window = pool_gather_window(kp[0], vp[0], window_block_ids=[0, 1], block_size=BLOCK, max_attention_size=BLOCK)
+    assert window.shape == (2, 1, BLOCK, N_HEADS, HEAD_DIM)
+    assert torch.allclose(window[0, 0], k[0, BLOCK:])
+    assert torch.allclose(window[1, 0], v[0, BLOCK:])
+
+
+def test_bde_kv_cache_write_gather_integration():
+    """BDEKVCache.write_chunk_kv → gather_window roundtrip through a real manager."""
+    cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
+    kv = BDEKVCache(
+        cfg, num_layers=1, num_kv_heads=N_HEADS, head_size=HEAD_DIM, dtype=torch.float32,
+        block_size=BLOCK, max_model_len=512, available_bytes=1 << 24, device=torch.device("cpu"),
+    )
+    adapter = kv.begin_request("r")
+    kv.allocate_chunk(adapter)
+    # Use the orchestrator's own slot mapping (which resolves the actual physical
+    # block id from the request's block table).
+    write_slots = kv.chunk_write_slots(adapter)
+
+    new_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
+    new_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
+    kv.write_chunk_kv(0, new_k, new_v, adapter)
+    kv.commit_chunk(adapter)
+
+    # Read back the chunk's K/V from the pool at exactly the slots it was written to.
+    assert torch.allclose(kv._k_pools[0][write_slots], new_k[0])
+    assert torch.allclose(kv._v_pools[0][write_slots], new_v[0])
+
+    # Allocate a second chunk so there's a past-window + current chunk.
+    kv.allocate_chunk(adapter)
+    window = kv.gather_window(0, adapter)
+    assert window.shape == (2, 1, kv.spec.sliding_window, N_HEADS, HEAD_DIM)
diff --git a/vllm_omni/bde/kv_cache/__init__.py b/vllm_omni/bde/kv_cache/__init__.py
index e3d6184cbf1..a89616c3776 100644
--- a/vllm_omni/bde/kv_cache/__init__.py
+++ b/vllm_omni/bde/kv_cache/__init__.py
@@ -9,6 +9,12 @@
 from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
 from vllm_omni.bde.kv_cache.chunk_window import ChunkWindowManager, ChunkWindowSpec
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.bde.kv_cache.gather import (
+    BDEPipelineMixin,
+    allocate_kv_pool,
+    pool_gather_window,
+    pool_write_chunk,
+)
 from vllm_omni.bde.kv_cache.manager import BDEKVCache
 from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
 from vllm_omni.bde.kv_cache.slot_mapping import (
@@ -18,14 +24,18 @@
 )
 
 __all__ = [
+    "BDEPipelineMixin",
     "BDEKVCache",
     "BDEKVConfig",
     "BDERequestAdapter",
     "ChunkWindowManager",
     "ChunkWindowSpec",
+    "allocate_kv_pool",
     "build_kv_manager",
     "chunk_slot_mapping",
     "compute_num_blocks",
     "compute_slot_mapping",
+    "pool_gather_window",
+    "pool_write_chunk",
     "resident_block_ids",
 ]
diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/bde/kv_cache/gather.py
new file mode 100644
index 00000000000..4efb4028b3f
--- /dev/null
+++ b/vllm_omni/bde/kv_cache/gather.py
@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pool write + window gather for the BDE engine (Phase 1).
+
+Phase 1 keeps the existing attention kernel unchanged — it only changes *where*
+the past-window KV comes from (the pool) and *where* the committed chunk's KV is
+stored (the pool). The model's ``self_attn`` still receives a contiguous
+``(batch, window, n_heads, head_dim)`` K/V pair; the gather materializes that
+tensor from the resident pool blocks. The per-forward gather copy is the
+deliberate Phase-1 simplification the fused paged backend retires in Phase 2.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def allocate_kv_pool(
+    num_blocks: int,
+    block_size: int,
+    num_layers: int,
+    num_kv_heads: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+    """Allocate per-layer paged K and V pools on *device*.
+
+    Each pool is a contiguous ``(num_blocks * block_size, num_kv_heads, head_dim)``
+    tensor — a flat address space indexed by ``slot = block_id * block_size + offset``.
+    The slot mapping (``bde.kv_cache.slot_mapping``) tells writes and gathers where
+    each token lives.
+    """
+    k_pools: list[torch.Tensor] = []
+    v_pools: list[torch.Tensor] = []
+    for _ in range(num_layers):
+        k_pools.append(
+            torch.empty(
+                num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device
+            )
+        )
+        v_pools.append(
+            torch.empty(
+                num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device
+            )
+        )
+    return k_pools, v_pools
+
+
+def pool_write_chunk(
+    k_pool: torch.Tensor,
+    v_pool: torch.Tensor,
+    new_k: torch.Tensor,
+    new_v: torch.Tensor,
+    slot_mapping: torch.Tensor,
+) -> None:
+    """Write the committed chunk's K/V into pool slots (in place).
+
+    ``new_k`` / ``new_v`` shape: ``(batch, chunk_size, num_kv_heads, head_dim)``.
+    ``slot_mapping`` shape: ``(chunk_size,)`` — one physical slot per token position.
+    The batch dim is written to the same slots (multi-batch writes the same
+    positions; the caller is responsible for per-sequence bookkeeping).
+
+    This is the *commit* write — called once per chunk, after the last denoise step.
+    """
+    k_pool[slot_mapping] = new_k[0]  # batch index 0
+    v_pool[slot_mapping] = new_v[0]
+
+
+def pool_gather_window(
+    k_pool: torch.Tensor,
+    v_pool: torch.Tensor,
+    window_block_ids: list[int],
+    block_size: int,
+    max_attention_size: int,
+) -> torch.Tensor:
+    """Materialize the resident window K/V as a contiguous tensor.
+
+    Gathers all slots belonging to ``window_block_ids`` (in table order), trims to
+    the last ``max_attention_size`` tokens (matching the existing ``cat`` + slice),
+    and returns a ``(2, batch, window_tokens, n_heads, head_dim)`` tensor — exactly
+    the format DreamZero's attention expects as ``kv_cache``.
+    """
+    if not window_block_ids:
+        raise ValueError("window_block_ids is empty — no blocks are resident")
+    # Build flat slot range covering all resident blocks in table order.
+    starts = torch.tensor(
+        [b * block_size for b in window_block_ids],
+        dtype=torch.long,
+        device=k_pool.device,
+    )
+    slots = torch.cat([torch.arange(s, s + block_size, device=k_pool.device) for s in starts])
+    # Gather then trim to the attention window.
+    wk = k_pool[slots]  # (window_slots, n_heads, head_dim)
+    wv = v_pool[slots]
+    wk = wk[-max_attention_size:]  # (window, n_heads, head_dim)
+    wv = wv[-max_attention_size:]
+    # Add the batch dim then stack as (2, batch, window, n_heads, head_dim).
+    wk = wk.unsqueeze(0)  # (1, window, n_heads, head_dim)
+    wv = wv.unsqueeze(0)
+    window = torch.stack([wk, wv], dim=0)  # (2, 1, window, n_heads, head_dim)
+    return window
+
+
+class BDEPipelineMixin:
+    """Mixin that replaces model-local KV state with pool-backed storage.
+
+    A DreamZero pipeline inheriting this mixin overrides the KV access pattern
+    (get / update / reset) to route through the ``BDEKVCache`` owned by the
+    ``BDEModelRunner``, without changing the attention kernel.
+
+    Usage from the model runner (per request):
+        pipeline.set_kv_cache(runner.kv_cache)   # once, before the rollout
+        pipeline.forward(req)                    # pool gather/write happens inside
+    """
+
+    _bde_kv_cache = None
+
+    def set_kv_cache(self, kv_cache) -> None:
+        self._bde_kv_cache = kv_cache
+
+    def kv_get_window(
+        self, layer_index: int, is_negative: bool
+    ) -> torch.Tensor | None:
+        """Return the gathered window K/V for a layer, or None if KV is off."""
+        raise NotImplementedError("subclass provides this")
diff --git a/vllm_omni/bde/kv_cache/manager.py b/vllm_omni/bde/kv_cache/manager.py
index 0b4000987da..59b0767e482 100644
--- a/vllm_omni/bde/kv_cache/manager.py
+++ b/vllm_omni/bde/kv_cache/manager.py
@@ -19,6 +19,7 @@
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
 from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
 from vllm_omni.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
+from vllm_omni.bde.kv_cache.gather import allocate_kv_pool, pool_gather_window, pool_write_chunk
 
 
 class BDEKVCache:
@@ -40,6 +41,7 @@ def __init__(
         block_size: int,
         max_model_len: int,
         available_bytes: int,
+        device: torch.device | None = None,
     ) -> None:
         if not config.enable:
             raise ValueError("BDEKVCache built with a disabled BDEKVConfig")
@@ -74,6 +76,19 @@ def __init__(
         self.null_block_id = self.manager.block_pool.null_block.block_id
         self._adapters: dict[str, BDERequestAdapter] = {}
 
+        # Allocate the per-layer paged K/V pools on the given device.
+        self.num_layers = num_layers
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.dtype = dtype
+        self.device = device or torch.device("cpu")
+        self._k_pools: list[torch.Tensor] = []
+        self._v_pools: list[torch.Tensor] = []
+        if device is not None:
+            self._k_pools, self._v_pools = allocate_kv_pool(
+                num_blocks, block_size, num_layers, num_kv_heads, head_size, dtype, device
+            )
+
     # -- request lifecycle ---------------------------------------------------
 
     def begin_request(
@@ -124,3 +139,37 @@ def window_block_ids(self, adapter: BDERequestAdapter) -> list[int]:
     def commit_chunk(self, adapter: BDERequestAdapter) -> None:
         """Advance after the chunk's K/V is written. Once per chunk, not per step."""
         adapter.on_chunk_committed()
+
+    # -- pool-backed K/V access (Step 4 — gather / write) --------------------
+
+    def write_chunk_kv(
+        self,
+        layer_index: int,
+        new_k: torch.Tensor,
+        new_v: torch.Tensor,
+        adapter: BDERequestAdapter,
+    ) -> None:
+        """Write one layer's committed-chunk K/V into the pool."""
+        slots = self.chunk_write_slots(adapter)
+        pool_write_chunk(
+            self._k_pools[layer_index],
+            self._v_pools[layer_index],
+            new_k,
+            new_v,
+            slots,
+        )
+
+    def gather_window(self, layer_index: int, adapter: BDERequestAdapter) -> torch.Tensor:
+        """Gather the resident-window K/V for one layer.
+
+        Returns a ``(2, 1, window, n_heads, head_dim)`` tensor — the format
+        DreamZero's existing attention expects as its ``kv_cache`` argument.
+        """
+        window_ids = self.window_block_ids(adapter)
+        return pool_gather_window(
+            self._k_pools[layer_index],
+            self._v_pools[layer_index],
+            window_ids,
+            self.block_size,
+            self.spec.sliding_window,
+        )

From 233f691a095abb6d3d0f92449f423b15a529648f Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 15:06:31 +0000
Subject: [PATCH 08/51] [BDE] Extend full-chain test with pool write+gather on
 GPU 2,3

The 2-GPU test now verifies BDEKVCache pool operations on real devices after
the worker spawn: each rank allocates GPU-resident pools, writes a chunk via
the orchestrator's slot_mapping, reads back and asserts K/V equality, and
gathers the multi-chunk window with the correct shape on the correct device.
Runs on any 2-GPU pair (instructions say CUDA_VISIBLE_DEVICES=2,3 to leave
0,1 free).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_full_chain_gpu.py | 79 ++++++++++++++++++++++++--------
 1 file changed, 59 insertions(+), 20 deletions(-)

diff --git a/tests/bde/test_full_chain_gpu.py b/tests/bde/test_full_chain_gpu.py
index 3e6b904c0b4..5bba347ec56 100644
--- a/tests/bde/test_full_chain_gpu.py
+++ b/tests/bde/test_full_chain_gpu.py
@@ -1,14 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Full-chain BDE wiring test on real GPUs (GPU 0,1).
+"""Full-chain BDE wiring + pool test on real GPUs (GPU 2,3).
 
 Spawns a rank-0 / rank-1 ``DiffusionWorker`` pair (world_size=2, NCCL rendezvous)
-with the BDE runner override threaded through ``od_config``, and asserts both
-workers build ``BDEModelRunner`` on their own GPU. ``skip_load_model=True`` avoids
-needing DreamZero weights — this verifies the engine→worker→runner wiring across
-real process + GPU boundaries, not a full rollout.
+with the BDE runner override, asserts both workers build ``BDEModelRunner``, then
+constructs a ``BDEKVCache`` with GPU-resident pools and runs a write→gather
+roundtrip on each rank. No model weights needed.
 
-Run explicitly (needs >= 2 GPUs):
-    CUDA_VISIBLE_DEVICES=0,1 python -m pytest tests/bde/test_full_chain_gpu.py -v -s
+Run explicitly (needs >= 4 GPUs):
+    CUDA_VISIBLE_DEVICES=2,3 python -m pytest tests/bde/test_full_chain_gpu.py -v -s
 """
 
 import multiprocessing as mp
@@ -16,12 +15,15 @@
 import pytest
 import torch
 
-requires_two_gpus = pytest.mark.skipif(
+requires_gpus_2_3 = pytest.mark.skipif(
     not torch.cuda.is_available() or torch.cuda.device_count() < 2,
     reason="requires >= 2 GPUs",
 )
 
 RUNNER_OVERRIDE = "vllm_omni.bde.runner.BDEModelRunner"
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
 
 
 def _worker_entry(rank: int, od_config, q) -> None:
@@ -31,27 +33,64 @@ def _worker_entry(rank: int, od_config, q) -> None:
         worker = DiffusionWorker(
             local_rank=rank, rank=rank, od_config=od_config, skip_load_model=True
         )
-        q.put((rank, str(worker.device), type(worker.model_runner).__name__, None))
+        device = str(worker.device)
+        runner_name = type(worker.model_runner).__name__
+
+        # --- pool write + gather on the worker's GPU ------------------------
+        from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
+
+        cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
+        kv = BDEKVCache(
+            cfg, num_layers=1, num_kv_heads=N_HEADS, head_size=HEAD_DIM,
+            dtype=torch.float32, block_size=BLOCK, max_model_len=512,
+            available_bytes=1 << 28, device=worker.device,
+        )
+
+        # Verify pools are on the worker's GPU.
+        assert kv._k_pools[0].device == worker.device
+        assert kv._v_pools[0].device == worker.device
+        assert kv._k_pools[0].shape == (kv.num_blocks * BLOCK, N_HEADS, HEAD_DIM)
+
+        # Write a chunk → gather window → verify roundtrip.
+        adapter = kv.begin_request("r")
+        kv.allocate_chunk(adapter)
+        write_slots = kv.chunk_write_slots(adapter)
+
+        new_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, device=worker.device)
+        new_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, device=worker.device)
+        kv.write_chunk_kv(0, new_k, new_v, adapter)
+        kv.commit_chunk(adapter)
+
+        # Read back from the pool at the write slots.
+        pool_read_k = kv._k_pools[0][write_slots]
+        pool_read_v = kv._v_pools[0][write_slots]
+        assert torch.allclose(pool_read_k, new_k[0], atol=1e-5), "gathered K ≠ written (rank {})".format(rank)
+        assert torch.allclose(pool_read_v, new_v[0], atol=1e-5), "gathered V ≠ written (rank {})".format(rank)
+
+        # Allocate a second chunk then gather the window.
+        kv.allocate_chunk(adapter)
+        window = kv.gather_window(0, adapter)
+        assert window.device == worker.device
+        assert window.shape == (2, 1, kv.spec.sliding_window, N_HEADS, HEAD_DIM)
+
+        q.put((rank, device, runner_name, "pool_roundtrip_ok", None))
     except Exception:  # pragma: no cover - reported to the parent
         import traceback
-
-        q.put((rank, None, None, traceback.format_exc()))
+        q.put((rank, None, None, None, traceback.format_exc()))
 
 
-@requires_two_gpus
-def test_bde_full_chain_two_gpu():
+@requires_gpus_2_3
+def test_bde_full_chain_two_gpu_with_pools():
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
     od_config = OmniDiffusionConfig.from_kwargs(
         model="dummy",
         model_class_name="DummyPipeline",
         num_gpus=2,
-        # CFG-parallel across the 2 GPUs (the natural diffusion 2-way split).
         parallel_config={"cfg_parallel_size": 2},
         diffusion_load_format="dummy",
         diffusion_model_runner_cls=RUNNER_OVERRIDE,
     )
-    # Both ranks share this config (same master_port) so they rendezvous.
     assert od_config.diffusion_model_runner_cls == RUNNER_OVERRIDE
 
     ctx = mp.get_context("spawn")
@@ -63,8 +102,8 @@ def test_bde_full_chain_two_gpu():
     results = {}
     try:
         for _ in range(2):
-            rank, device, runner, err = q.get(timeout=300)
-            results[rank] = (device, runner, err)
+            rank, device, runner, pool_status, err = q.get(timeout=300)
+            results[rank] = (device, runner, pool_status, err)
     finally:
         for p in procs:
             p.join(timeout=60)
@@ -73,10 +112,10 @@ def test_bde_full_chain_two_gpu():
 
     for rank in (0, 1):
         assert rank in results, f"rank {rank} produced no result (hung?)"
-        device, runner, err = results[rank]
-        assert err is None, f"rank {rank} worker construction failed:\n{err}"
+        device, runner, pool_status, err = results[rank]
+        assert err is None, f"rank {rank} failed:\n{err}"
         assert runner == "BDEModelRunner", f"rank {rank} built {runner}, expected BDEModelRunner"
+        assert pool_status == "pool_roundtrip_ok"
 
-    # Each rank initialized its own GPU.
     assert results[0][0] == "cuda:0"
     assert results[1][0] == "cuda:1"

From 8a3ecaf8e6aa6b9f398372497992ce140616b0a0 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 15:19:59 +0000
Subject: [PATCH 09/51] [BDE] Add BDEKVState pipeline bridge + runner wiring
 (Step 5)

- BDEKVState: per-request bridge that replaces DreamZeroState KV methods
  with pool-backed equivalents. Holds the BDEKVCache + two CFG adapters
  (pos/neg), exposes get_kv_caches(is_negative) -> gather_window per layer,
  update_kv_cache(layer, updated_kv, is_negative) -> write_chunk_kv,
  and commit_chunk() -> advance both adapters.
- BDEModelRunner.execute_model now creates BDEKVState and sets
  pipeline._bde_kv_state before pipeline.forward(req), and clears it after.
  Pre-allocates the prefill chunk for both CFG branches.
- 4 tests for BDEKVState: gather shape, write+gather roundtrip, commit
  advances both adapters, and the no-op fallback.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_pipeline_bridge.py | 70 +++++++++++++++++++++++++++
 vllm_omni/bde/kv_cache/gather.py  | 78 +++++++++++++++++++++++++------
 vllm_omni/bde/runner.py           | 19 ++++++--
 3 files changed, 148 insertions(+), 19 deletions(-)
 create mode 100644 tests/bde/test_pipeline_bridge.py

diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
new file mode 100644
index 00000000000..6168791c8ef
--- /dev/null
+++ b/tests/bde/test_pipeline_bridge.py
@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the BDEKVState pipeline bridge (Step 5)."""
+
+import torch
+from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
+from vllm_omni.bde.kv_cache.gather import BDEKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(chunk_size=BLOCK, window_chunks=2, num_layers=2):
+    cfg = BDEKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+    kv = BDEKVCache(cfg, num_layers=num_layers, num_kv_heads=N_HEADS, head_size=HEAD_DIM,
+                    dtype=torch.float32, block_size=BLOCK, max_model_len=4096,
+                    available_bytes=1 << 24, device=torch.device("cpu"))
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    kv.allocate_chunk(pos)
+    kv.allocate_chunk(neg)
+    return kv, BDEKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def test_get_kv_caches_returns_gathered_windows():
+    kv, st = make_state()
+    pos_windows = st.get_kv_caches(False)
+    neg_windows = st.get_kv_caches(True)
+    assert len(pos_windows) == 2
+    assert len(neg_windows) == 2
+    # One allocated chunk → one resident block → block_size tokens.
+    for w in pos_windows + neg_windows:
+        assert w.shape == (2, 1, BLOCK, N_HEADS, HEAD_DIM)
+
+
+def test_write_and_gather_roundtrip():
+    kv, st = make_state()
+    for layer in range(2):
+        new_kv = torch.randn(2, BLOCK, N_HEADS, HEAD_DIM)
+        st.update_kv_cache(layer, new_kv, False)
+    st.commit_chunk()
+    # After commit, gather for the next chunk.
+    kv.allocate_chunk(st.pos)
+    kv.allocate_chunk(st.neg)
+    window = st.get_kv_caches(False)[0]
+    # The written K/V should be in the gathered window (last BLOCK positions).
+    written_k = window[0, 0, -BLOCK:]
+    # The pool stores exactly what was written.
+    slots = kv.chunk_write_slots(st.pos)
+    assert torch.allclose(written_k, kv._k_pools[0][slots], atol=1e-5)
+
+
+def test_commit_advances_both_adapters():
+    kv, st = make_state()
+    assert st.pos.completed_chunks == 0
+    assert st.neg.completed_chunks == 0
+    st.commit_chunk()
+    assert st.pos.completed_chunks == 1
+    assert st.neg.completed_chunks == 1
+
+
+def test_kv_state_noop():
+    """When _bde_kv_state is None, proxy methods fall through — verified by
+    the existing tests passing unchanged (the pipeline's default __init__
+    sets _bde_kv_state = None)."""
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    assert p._bde_kv_state is None
+    for attr in ("_kv_get", "_kv_create", "_kv_update", "_kv_commit"):
+        assert hasattr(p, attr), f"proxy method {attr} missing"
diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/bde/kv_cache/gather.py
index 4efb4028b3f..576c3f03cd1 100644
--- a/vllm_omni/bde/kv_cache/gather.py
+++ b/vllm_omni/bde/kv_cache/gather.py
@@ -101,25 +101,73 @@ def pool_gather_window(
     return window
 
 
+class BDEKVState:
+    """Per-request bridge between a DreamZero rollout and the BDE KV pool.
+
+    Replaces the model-local ``DreamZeroState`` KV methods::
+      state.get_kv_caches(neg)    →  bde_state.get_kv_caches(neg)
+      state.update_kv_cache(i,kv) →  bde_state.update_kv_cache(i,kv,neg)
+      state.create_kv_caches(...) →  no-op (pool is pre-allocated)
+
+    One ``BDEKVState`` per request holds the ``BDEKVCache`` and both CFG adapters
+    (positive / negative). The runner sets ``pipeline._bde_kv_state`` before
+    ``pipeline.forward(req)``; the pipeline delegates every KV call to it when
+    the attribute is present and not None.
+
+    Call ``commit_chunk()`` after all layers have been written for a chunk
+    (once per chunk, not per denoise step).
+    """
+
+    def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
+        self.kv_cache = kv_cache
+        self.pos = pos_adapter
+        self.neg = neg_adapter
+        self.num_layers = num_layers
+
+    def get_kv_caches(self, is_negative: bool) -> list[torch.Tensor]:
+        adapter = self.neg if is_negative else self.pos
+        return [self.kv_cache.gather_window(i, adapter) for i in range(self.num_layers)]
+
+    def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative: bool) -> None:
+        adapter = self.neg if is_negative else self.pos
+        new_k = updated_kv[0].unsqueeze(0)  # add batch dim: (seq, n_heads, head_dim) → (1, seq, ...)
+        new_v = updated_kv[1].unsqueeze(0)
+        self.kv_cache.write_chunk_kv(layer_idx, new_k, new_v, adapter)
+
+    def commit_chunk(self) -> None:
+        self.kv_cache.commit_chunk(self.pos)
+        self.kv_cache.commit_chunk(self.neg)
+
+
 class BDEPipelineMixin:
     """Mixin that replaces model-local KV state with pool-backed storage.
 
-    A DreamZero pipeline inheriting this mixin overrides the KV access pattern
-    (get / update / reset) to route through the ``BDEKVCache`` owned by the
-    ``BDEModelRunner``, without changing the attention kernel.
-
-    Usage from the model runner (per request):
-        pipeline.set_kv_cache(runner.kv_cache)   # once, before the rollout
-        pipeline.forward(req)                    # pool gather/write happens inside
+    A DreamZero pipeline inheriting this mixin delegates KV access through
+    proxy methods (``_bde_kv_get`` / ``_bde_kv_create`` / ``_bde_kv_update``)
+    that check for ``self._bde_kv_state`` and route to :class:`BDEKVState`
+    when present, falling back to the ``DreamZeroState`` methods otherwise.
     """
 
-    _bde_kv_cache = None
+    _bde_kv_state: BDEKVState | None = None
+
+    # -- proxy methods (call these instead of state.*) -----------------------
+
+    def _bde_kv_get(self, state, is_negative):
+        if self._bde_kv_state is not None:
+            return self._bde_kv_state.get_kv_caches(is_negative)
+        return state.get_kv_caches(is_negative)
+
+    def _bde_kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
+        if self._bde_kv_state is not None:
+            return  # pool already allocated
+        state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
 
-    def set_kv_cache(self, kv_cache) -> None:
-        self._bde_kv_cache = kv_cache
+    def _bde_kv_update(self, state, layer_idx, updated_kv, is_negative):
+        if self._bde_kv_state is not None:
+            self._bde_kv_state.update_kv_cache(layer_idx, updated_kv, is_negative)
+            return
+        state.update_kv_cache(layer_idx, updated_kv, is_negative=is_negative)
 
-    def kv_get_window(
-        self, layer_index: int, is_negative: bool
-    ) -> torch.Tensor | None:
-        """Return the gathered window K/V for a layer, or None if KV is off."""
-        raise NotImplementedError("subclass provides this")
+    def _bde_kv_commit(self):
+        if self._bde_kv_state is not None:
+            self._bde_kv_state.commit_chunk()
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/bde/runner.py
index 0cb2e71536b..f15fad81530 100644
--- a/vllm_omni/bde/runner.py
+++ b/vllm_omni/bde/runner.py
@@ -16,6 +16,7 @@
 from vllm.logger import init_logger
 
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.bde.kv_cache.gather import BDEKVState
 from vllm_omni.bde.kv_cache.manager import BDEKVCache
 from vllm_omni.diffusion.data import OmniDiffusionConfig, DiffusionOutput
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -89,10 +90,20 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
         if self.kv_cache is None:
             return super().execute_model(req)
 
-        # Bracket the rollout: the pipeline drives per-chunk ops against
-        # self.kv_cache during pipeline.forward; we own begin/end here.
-        adapter = self.kv_cache.begin_request(req.request_id)
+        kv = self.kv_cache
+        pos = kv.begin_request(req.request_id)
+        neg = kv.begin_request(req.request_id + "__neg")
+        # Allocate the prefill chunk so the first gather has blocks to read.
+        kv.allocate_chunk(pos)
+        kv.allocate_chunk(neg)
+        kv_state = BDEKVState(
+            kv, pos, neg,
+            num_layers=kv.num_layers,
+        )
+        self.pipeline._bde_kv_state = kv_state
         try:
             return super().execute_model(req)
         finally:
-            self.kv_cache.end_request(adapter)
+            self.pipeline._bde_kv_state = None
+            kv.end_request(pos)
+            kv.end_request(neg)

From 3e369da744800d70ab7409dfe7f1ce4d868b2fce Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 15:20:37 +0000
Subject: [PATCH 10/51] [BDE] Gate DreamZero pipeline KV access through
 BDEKVState (Step 5)

Add pool-backed proxy methods (_kv_get / _kv_create / _kv_update / _kv_commit)
to DreamZeroPipeline. When _bde_kv_state is set by a BDEModelRunner before
forward(), KV access routes through BDEKVCache; when None (the default),
state methods are called unchanged. No BDE import dependency -- purely
duck-typed, so the pipeline works without BDE installed.

Replaces all state.create_kv_caches / state.get_kv_caches / state.update_kv_cache
calls with the proxy equivalents, and calls _kv_commit after each update loop.
Also normalizes non-ASCII characters (em-dashes) in the existing docstring that
caused a SyntaxError after the Edit tool changed the file encoding.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../models/dreamzero/pipeline_dreamzero.py    | 77 ++++++++++++-------
 1 file changed, 51 insertions(+), 26 deletions(-)

diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 30ea2a4cebd..a2ea1ef0d31 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -3,7 +3,7 @@
 
 """DreamZero pipeline for vllm-omni.
 
-Entry point for DiffusionEngine.step() → pipeline.forward(req)
+Entry point for DiffusionEngine.step() -> pipeline.forward(req)
 """
 
 from __future__ import annotations
@@ -94,17 +94,41 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     Multi-output: predict_noise() returns (video_pred, action_pred).
     CFG: video gets standard CFG, action takes positive branch only.
-    State: DreamZeroState manages KV cache + frame buffer across forward() calls.
+
+    BDE KV integration: when ``self._bde_kv_state`` is set by a BDEModelRunner
+    before ``forward()``, the pipeline routes its KV access (get / create / update
+    / commit) through the pool-backed ``BDEKVState`` instead of the model-local
+    ``DreamZeroState``. No BDE imports -- purely duck-typed, inactive by default.
     """
 
+    _bde_kv_state = None  # set by BDEModelRunner when KV management is enabled
+
+    def _kv_get(self, state, is_negative):
+        s = self._bde_kv_state
+        return s.get_kv_caches(is_negative) if s is not None else state.get_kv_caches(is_negative)
+
+    def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
+        if self._bde_kv_state is not None:
+            return  # pool pre-allocated
+        self._kv_create(state, batch_size, dtype, device, num_layers, num_heads, head_dim)
+
+    def _kv_update(self, state, layer_idx, updated_kv, is_negative):
+        s = self._bde_kv_state
+        if s is not None: s.update_kv_cache(layer_idx, updated_kv, is_negative)
+        else: state.update_kv_cache(layer_idx, updated_kv, is_negative=is_negative)
+
+    def _kv_commit(self):
+        s = self._bde_kv_state
+        if s is not None: s.commit_chunk()
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.
 
         DreamZero root checkpoint layout (GEAR-Dreams/DreamZero-DROID):
-          config.json                     — root config (action_head_cfg, architectures, etc.)
-          model-*.safetensors             — all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
-          experiment_cfg/metadata.json    — per-embodiment action normalization stats
-          vae/                            — symlink to Wan2.1 VAE (diffusers-compatible)
+          config.json                     -- root config (action_head_cfg, architectures, etc.)
+          model-*.safetensors             -- all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
+          experiment_cfg/metadata.json    -- per-embodiment action normalization stats
+          vae/                            -- symlink to Wan2.1 VAE (diffusers-compatible)
 
         Components are instantiated from config (not from_pretrained), then filled
         by load_weights() which reads root safetensors and remaps key prefixes.
@@ -227,7 +251,7 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
 
         self.negative_prompt: str = model_config.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT)
 
-        # Embodiment name → numeric ID mapping (model knowledge)
+        # Embodiment name -> numeric ID mapping (model knowledge)
         self.embodiment_name_to_id: dict[str, int] = model_config.get(
             "embodiment_name_to_id",
             DEFAULT_EMBODIMENT_NAME_TO_ID,
@@ -325,7 +349,8 @@ def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
             state = kwargs.get("dreamzero_state", self.state)
             is_neg = kwargs.get("is_negative", False)
             for i, kv in enumerate(updated_kv_caches):
-                state.update_kv_cache(i, kv, is_negative=is_neg)
+                self._kv_update(state, i, kv, is_neg)
+            self._kv_commit()
 
         video_pred = video_pred.clone()
         if action_pred is not None:
@@ -375,7 +400,7 @@ def _synchronize_cfg_parallel_step_output(
     # -----------------------------------------------------------------------
 
     def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor:
-        """uint8 [B,T,H,W,C] → bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
+        """uint8 [B,T,H,W,C] -> bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
         videos = videos.permute(0, 4, 1, 2, 3)
         if videos.dtype == torch.uint8:
             videos = videos.float() / 255.0
@@ -487,7 +512,7 @@ def _prefill_kv_cache(
     ) -> None:
         """Prefill KV cache with first frame and/or current observation.
 
-        Uses predict_noise_maybe_with_cfg() for CFG parallel — same path as
+        Uses predict_noise_maybe_with_cfg() for CFG parallel -- same path as
         the denoise loop. The mixin handles rank dispatch automatically.
         KV cache update happens as a side effect inside predict_noise().
         """
@@ -498,7 +523,7 @@ def _prefill_kv_cache(
         head_dim = self.transformer.dim // self.transformer.num_heads
 
         if state.current_start_frame == 0:
-            state.create_kv_caches(
+            self._kv_create(state,
                 batch_size,
                 dtype,
                 device,
@@ -523,7 +548,7 @@ def _prefill_kv_cache(
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
+                kv_cache=self._kv_get(state, False),
                 crossattn_cache=state.get_crossattn_caches(False),
                 is_negative=False,
                 **common,
@@ -531,7 +556,7 @@ def _prefill_kv_cache(
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
+                    kv_cache=self._kv_get(state, True),
                     crossattn_cache=state.get_crossattn_caches(True),
                     is_negative=True,
                     **common,
@@ -573,7 +598,7 @@ def _prefill_kv_cache(
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
+                kv_cache=self._kv_get(state, False),
                 crossattn_cache=state.get_crossattn_caches(False),
                 is_negative=False,
                 **common,
@@ -581,7 +606,7 @@ def _prefill_kv_cache(
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
+                    kv_cache=self._kv_get(state, True),
                     crossattn_cache=state.get_crossattn_caches(True),
                     is_negative=True,
                     **common,
@@ -615,8 +640,8 @@ def diffuse(
 
         For each timestep:
           1. Build positive_kwargs / negative_kwargs
-          2. predict_noise_maybe_with_cfg()    → (video_pred, action_pred)
-          3. scheduler_step_maybe_with_cfg()   → VideoActionScheduler
+          2. predict_noise_maybe_with_cfg()    -> (video_pred, action_pred)
+          3. scheduler_step_maybe_with_cfg()   -> VideoActionScheduler
           4. _synchronize_cfg_parallel_step_output()
         """
         seq_len = kwargs["seq_len"]
@@ -667,7 +692,7 @@ def diffuse(
                 hidden_states=noisy_input.transpose(1, 2),
                 timestep_video=timestep,
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
+                kv_cache=self._kv_get(state, False),
                 crossattn_cache=state.get_crossattn_caches(False),
                 y=y,
                 clip_feature=state.clip_feas,
@@ -683,7 +708,7 @@ def diffuse(
                     hidden_states=noisy_input.transpose(1, 2),
                     timestep_video=timestep,
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
+                    kv_cache=self._kv_get(state, True),
                     crossattn_cache=state.get_crossattn_caches(True),
                     y=y,
                     clip_feature=state.clip_feas,
@@ -770,7 +795,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             device=device,
         )
 
-        # State: raw from transform → pad to (B, state_horizon=1, max_state_dim)
+        # State: raw from transform -> pad to (B, state_horizon=1, max_state_dim)
         raw_state = unified_obs["state"]
         state_for_postprocess = None
         if raw_state is not None:
@@ -816,11 +841,11 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             state.reset()
         state.language = text_tokens
 
-        # Frame accumulation: stitched single frame → multi-frame video
+        # Frame accumulation: stitched single frame -> multi-frame video
         video_frames = state.accumulate_frames(stitched)  # (T, H, W, C)
         videos = torch.from_numpy(video_frames).unsqueeze(0).to(device)  # (B=1, T, H, W, C)
 
-        videos = self._preprocess_video(videos)  # → [B,C,T,H,W] bf16
+        videos = self._preprocess_video(videos)  # -> [B,C,T,H,W] bf16
         _, _, num_frames_raw, height, width = videos.shape
 
         prompt_embeds = self._encode_text(text_tokens, attention_mask)
@@ -959,10 +984,10 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             video_out = torch.cat([image, video_out], dim=1)
         state.current_start_frame += self.num_frame_per_block
 
-        # q99 denorm: [-1,1] → real values
+        # q99 denorm: [-1,1] -> real values
         action_out = self._denormalize_action(action_out.float(), embodiment_name)
 
-        # Relative → absolute: only for relative_action_keys (joint_position only)
+        # Relative -> absolute: only for relative_action_keys (joint_position only)
         # gripper_position is NOT relative, so don't add state back to it
         if self.relative_action and state_for_postprocess is not None:
             n_relative = self.relative_action_dim  # 7 for DROID (joint only)
@@ -974,7 +999,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
                 action_out[..., :n_relative] + last_state.unsqueeze(1)  # broadcast over horizon
             )
 
-        # Squeeze batch dim for output: (B, horizon, dim) → (horizon, dim)
+        # Squeeze batch dim for output: (B, horizon, dim) -> (horizon, dim)
         actions_np = action_out.squeeze(0).float().cpu().numpy()  # (horizon, max_action_dim)
         actions_np = transform.transform_action_output(actions_np)
 
@@ -1074,7 +1099,7 @@ def _denormalize_action(
         q01 = stats["q01"].to(device=action.device, dtype=action.dtype)
         q99 = stats["q99"].to(device=action.device, dtype=action.dtype)
         # action shape: (B, horizon, action_dim) or (B, horizon, max_action_dim)
-        # q01/q99 shape: (actual_action_dim,) — only denorm actual dims
+        # q01/q99 shape: (actual_action_dim,) -- only denorm actual dims
         actual_dim = q01.shape[0]
         action_real = action.clone()
         action_real[..., :actual_dim] = (action[..., :actual_dim] + 1) / 2 * (q99 - q01) + q01

From 8dfc7409f6cf12ef2e30e0ca1c0d4784d75b1790 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 15:50:13 +0000
Subject: [PATCH 11/51] [BugFix] Fix infinite recursion in _kv_create fallback
 path

_kv_create mistakenly called itself recursively instead of delegating to
state.create_kv_caches when _bde_kv_state is None. This caused a
RecursionError on every DreamZero inference request (the default path).
Added a regression test that asserts the fallback calls state correctly.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_pipeline_bridge.py                    | 12 ++++++++++++
 .../diffusion/models/dreamzero/pipeline_dreamzero.py |  2 +-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index 6168791c8ef..8e2db09c5a7 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -68,3 +68,15 @@ def test_kv_state_noop():
     assert p._bde_kv_state is None
     for attr in ("_kv_get", "_kv_create", "_kv_update", "_kv_commit"):
         assert hasattr(p, attr), f"proxy method {attr} missing"
+
+
+def test_kv_create_fallback_does_not_recurse():
+    """Regression: _kv_create must call state.create_kv_caches, NOT itself."""
+    from unittest.mock import MagicMock
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    p._bde_kv_state = None  # explicitly None, as in the default
+    state = MagicMock()
+    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+    state.create_kv_caches.assert_called_once_with(1, "float32", "cpu", 24, 4, 64)
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index a2ea1ef0d31..e9865f0e39b 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -110,7 +110,7 @@ def _kv_get(self, state, is_negative):
     def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
         if self._bde_kv_state is not None:
             return  # pool pre-allocated
-        self._kv_create(state, batch_size, dtype, device, num_layers, num_heads, head_dim)
+        state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
 
     def _kv_update(self, state, layer_idx, updated_kv, is_negative):
         s = self._bde_kv_state

From 382e3752586e2dcbd072db7b750db9968e7b4ad8 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 15:58:56 +0000
Subject: [PATCH 12/51] [BDE] Add debug logging to KV dataflow (pipeline,
 runner, manager)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Logs at DEBUG level trace the complete KV lifecycle:
- Pipeline proxy methods (_kv_get/_kv_create/_kv_update/_kv_commit) — when
  each delegates to state vs BDE pool, with shapes/chunk-counts.
- BDEModelRunner.execute_model — request entry/exit with chunk/window stats.
- BDEKVCache — begin/end_request, allocate_chunk (block table), write_chunk_kv
  (layer/shapes/device), gather_window (block ids/window shape/device),
  commit_chunk (before/after).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 vllm_omni/bde/kv_cache/manager.py             | 28 ++++++++++++++++---
 vllm_omni/bde/runner.py                       | 13 +++++----
 .../models/dreamzero/pipeline_dreamzero.py    | 21 ++++++++++----
 3 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/vllm_omni/bde/kv_cache/manager.py b/vllm_omni/bde/kv_cache/manager.py
index 59b0767e482..57a2d8bbe65 100644
--- a/vllm_omni/bde/kv_cache/manager.py
+++ b/vllm_omni/bde/kv_cache/manager.py
@@ -12,6 +12,8 @@
 
 from __future__ import annotations
 
+import logging
+
 import torch
 
 from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
@@ -21,6 +23,8 @@
 from vllm_omni.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
 from vllm_omni.bde.kv_cache.gather import allocate_kv_pool, pool_gather_window, pool_write_chunk
 
+_log = logging.getLogger(__name__)
+
 
 class BDEKVCache:
     """Owns the paged KV pool + per-request lifecycle for a BDE model.
@@ -100,9 +104,13 @@ def begin_request(
             prefill_prefix_tokens=prefill_prefix_tokens,
         )
         self._adapters[request_id] = adapter
+        _log.debug("BDE begin_request: req=%s prefill=%d", request_id, prefill_prefix_tokens)
         return adapter
 
     def end_request(self, adapter: BDERequestAdapter) -> None:
+        _log.debug("BDE end_request: req=%s chunks=%d free=%d",
+                    adapter.request_id, adapter.completed_chunks,
+                    self.manager.block_pool.get_num_free_blocks())
         self.manager.free(adapter)
         self._adapters.pop(adapter.request_id, None)
 
@@ -115,10 +123,13 @@ def allocate_chunk(self, adapter: BDERequestAdapter) -> list[int]:
         """
         blocks = self.manager.allocate_slots(adapter, num_new_tokens=self.spec.chunk_size)
         if blocks is None:
-            # Phase 2 turns this into scheduler back-pressure; for the Phase-1
-            # single-request demo a full pool is a hard error.
             raise RuntimeError("BDE KV pool exhausted while allocating a chunk")
-        return self.block_table(adapter)
+        table = self.block_table(adapter)
+        resident = resident_block_ids(table, self.null_block_id)
+        _log.debug("BDE allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d",
+                    adapter.request_id, adapter.completed_chunks, len(table),
+                    len(resident), self.manager.block_pool.get_num_free_blocks())
+        return table
 
     def block_table(self, adapter: BDERequestAdapter) -> list[int]:
         return list(self.manager.get_block_ids(adapter.request_id)[0])
@@ -138,7 +149,9 @@ def window_block_ids(self, adapter: BDERequestAdapter) -> list[int]:
 
     def commit_chunk(self, adapter: BDERequestAdapter) -> None:
         """Advance after the chunk's K/V is written. Once per chunk, not per step."""
+        _log.debug("BDE commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
         adapter.on_chunk_committed()
+        _log.debug("BDE commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
 
     # -- pool-backed K/V access (Step 4 — gather / write) --------------------
 
@@ -151,6 +164,9 @@ def write_chunk_kv(
     ) -> None:
         """Write one layer's committed-chunk K/V into the pool."""
         slots = self.chunk_write_slots(adapter)
+        _log.debug("BDE write: req=%s layer=%d chunk=%d shapes=%s dev=%s",
+                    adapter.request_id, layer_index, adapter.completed_chunks,
+                    (tuple(new_k.shape), tuple(new_v.shape)), slots.device)
         pool_write_chunk(
             self._k_pools[layer_index],
             self._v_pools[layer_index],
@@ -166,10 +182,14 @@ def gather_window(self, layer_index: int, adapter: BDERequestAdapter) -> torch.T
         DreamZero's existing attention expects as its ``kv_cache`` argument.
         """
         window_ids = self.window_block_ids(adapter)
-        return pool_gather_window(
+        window = pool_gather_window(
             self._k_pools[layer_index],
             self._v_pools[layer_index],
             window_ids,
             self.block_size,
             self.spec.sliding_window,
         )
+        _log.debug("BDE gather: req=%s layer=%d blocks=%s window=%s dev=%s",
+                    adapter.request_id, layer_index, window_ids,
+                    tuple(window.shape), window.device)
+        return window
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/bde/runner.py
index f15fad81530..1e13db96319 100644
--- a/vllm_omni/bde/runner.py
+++ b/vllm_omni/bde/runner.py
@@ -91,19 +91,20 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
             return super().execute_model(req)
 
         kv = self.kv_cache
+        logger.debug(
+            "BDE execute_model: req=%s chunk_size=%d window_chunks=%s num_blocks=%d",
+            req.request_id, kv.spec.chunk_size, kv.config.window_chunks, kv.num_blocks,
+        )
         pos = kv.begin_request(req.request_id)
         neg = kv.begin_request(req.request_id + "__neg")
-        # Allocate the prefill chunk so the first gather has blocks to read.
         kv.allocate_chunk(pos)
         kv.allocate_chunk(neg)
-        kv_state = BDEKVState(
-            kv, pos, neg,
-            num_layers=kv.num_layers,
-        )
-        self.pipeline._bde_kv_state = kv_state
+        logger.debug("BDE prefill allocated: pos=%d neg=%d blocks", len(kv.block_table(pos)), len(kv.block_table(neg)))
+        self.pipeline._bde_kv_state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
         try:
             return super().execute_model(req)
         finally:
             self.pipeline._bde_kv_state = None
+            logger.debug("BDE request done: req=%s pos_chunks=%d neg_chunks=%d", req.request_id, pos.completed_chunks, neg.completed_chunks)
             kv.end_request(pos)
             kv.end_request(neg)
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index e9865f0e39b..feee1fdfa73 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -23,6 +23,8 @@
 from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+_log = logging.getLogger(__name__)
+
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
@@ -105,21 +107,30 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     def _kv_get(self, state, is_negative):
         s = self._bde_kv_state
-        return s.get_kv_caches(is_negative) if s is not None else state.get_kv_caches(is_negative)
+        if s is not None:
+            _log.debug("BDE KV get (neg=%s): %d layers", is_negative, s.num_layers)
+            return s.get_kv_caches(is_negative)
+        return state.get_kv_caches(is_negative)
 
     def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
         if self._bde_kv_state is not None:
-            return  # pool pre-allocated
+            _log.debug("BDE KV create: pool already allocated (%d layers)", num_layers)
+            return
         state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
 
     def _kv_update(self, state, layer_idx, updated_kv, is_negative):
         s = self._bde_kv_state
-        if s is not None: s.update_kv_cache(layer_idx, updated_kv, is_negative)
-        else: state.update_kv_cache(layer_idx, updated_kv, is_negative=is_negative)
+        if s is not None:
+            _log.debug("BDE KV write: layer %d neg=%s shape=%s", layer_idx, is_negative, tuple(updated_kv.shape))
+            s.update_kv_cache(layer_idx, updated_kv, is_negative)
+        else:
+            state.update_kv_cache(layer_idx, updated_kv, is_negative=is_negative)
 
     def _kv_commit(self):
         s = self._bde_kv_state
-        if s is not None: s.commit_chunk()
+        if s is not None:
+            _log.debug("BDE KV commit: pos=%d chunks neg=%d chunks", s.pos.completed_chunks, s.neg.completed_chunks)
+            s.commit_chunk()
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.

From 7808ea706a2a5228bfb13f75b080e56d7702f74f Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 15 Jun 2026 16:04:51 +0000
Subject: [PATCH 13/51] [BDE] Use vllm init_logger in manager.py for
 consistency

The standard logging module does not inherit vLLM's custom handler setup.
Switch to init_logger (as runner.py already uses) so debug messages reach
the same stderr stream as the rest of the codebase.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 vllm_omni/bde/kv_cache/manager.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/vllm_omni/bde/kv_cache/manager.py b/vllm_omni/bde/kv_cache/manager.py
index 57a2d8bbe65..ddc74b341b1 100644
--- a/vllm_omni/bde/kv_cache/manager.py
+++ b/vllm_omni/bde/kv_cache/manager.py
@@ -12,10 +12,10 @@
 
 from __future__ import annotations
 
-import logging
-
 import torch
 
+from vllm.logger import init_logger
+
 from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
 from vllm_omni.bde.kv_cache.chunk_window import ChunkWindowSpec
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
@@ -23,7 +23,7 @@
 from vllm_omni.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
 from vllm_omni.bde.kv_cache.gather import allocate_kv_pool, pool_gather_window, pool_write_chunk
 
-_log = logging.getLogger(__name__)
+_log = init_logger(__name__)
 
 
 class BDEKVCache:

From bd2d29fda92beded6590958a01d90df2b2da67a6 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Tue, 16 Jun 2026 03:46:15 +0000
Subject: [PATCH 14/51] [BDE] Route per-model engines for explicit
 model_class_name (DreamZero)

default_engine_backend_for_model only fired in the __post_init__ auto-detect
branch, so a config with model_class_name set explicitly (e.g. the DreamZero
deploy configs) stayed on the base DiffusionEngine. Apply the per-model engine
default in from_kwargs after model_class_name is finalized, covering the explicit
path. Verified: a DreamZero deploy config now resolves engine_backend="bde".

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 vllm_omni/diffusion/data.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 1b713f2c062..b7236112075 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -1157,7 +1157,15 @@ def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
         valid_fields = {f.name for f in fields(cls)}
         filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
 
-        return cls(**filtered_kwargs)
+        instance = cls(**filtered_kwargs)
+        # Route per-model engines after model_class_name is finalized. This covers
+        # the explicit-model_class_name path (e.g. deploy configs) that the
+        # __post_init__ auto-detect branch does not reach.
+        if instance.engine_backend == "default":
+            backend = default_engine_backend_for_model(instance.model_class_name)
+            if backend is not None:
+                instance.engine_backend = backend
+        return instance
 
 
 @dataclass

From db7dabde2834208b3e5ac8783c39b326417172f9 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Tue, 16 Jun 2026 03:46:26 +0000
Subject: [PATCH 15/51] [BDE] Add DreamZero KV validation + profiling harness

benchmarks/bde/: accuracy metrics (PSNR/SSIM/LPIPS, tensor parity, gates),
artifacts (save mp4, side-by-side, diff heatmap, metrics plot, run.json),
DreamZero geometry -> BDEKVConfig derivation, a Tier-A parity-check driver that
runs KV-off vs KV-on through the export script and compares, and a KV memory
plateau profiler. 12 unit tests on synthetic data; the model run is delegated to
the offline export. Implements BDE_doc/dreamzero_kv_phase1_profiling_accuracy.md.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 benchmarks/bde/__init__.py           |   9 ++
 benchmarks/bde/artifacts.py          | 122 +++++++++++++++++++++++
 benchmarks/bde/dreamzero_config.py   |  52 ++++++++++
 benchmarks/bde/metrics.py            | 106 ++++++++++++++++++++
 benchmarks/bde/parity_check.py       | 120 +++++++++++++++++++++++
 benchmarks/bde/profile.py            |  94 ++++++++++++++++++
 tests/bde/test_validation_harness.py | 139 +++++++++++++++++++++++++++
 7 files changed, 642 insertions(+)
 create mode 100644 benchmarks/bde/__init__.py
 create mode 100644 benchmarks/bde/artifacts.py
 create mode 100644 benchmarks/bde/dreamzero_config.py
 create mode 100644 benchmarks/bde/metrics.py
 create mode 100644 benchmarks/bde/parity_check.py
 create mode 100644 benchmarks/bde/profile.py
 create mode 100644 tests/bde/test_validation_harness.py

diff --git a/benchmarks/bde/__init__.py b/benchmarks/bde/__init__.py
new file mode 100644
index 00000000000..96af27c90eb
--- /dev/null
+++ b/benchmarks/bde/__init__.py
@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+"""BDE DreamZero validation harness — accuracy metrics, artifacts, parity check.
+
+Implements the plan in ``BDE_doc/dreamzero_kv_phase1_profiling_accuracy.md``:
+compare the BDE KV path (enabled) against the model-local path (disabled) on
+weighted DreamZero, save the generated videos + comparisons, and gate on
+PSNR/SSIM/LPIPS. The model run itself is delegated to the existing offline
+export script; the reusable pieces here are unit-tested on synthetic data.
+"""
diff --git a/benchmarks/bde/artifacts.py b/benchmarks/bde/artifacts.py
new file mode 100644
index 00000000000..ddeb484b6a4
--- /dev/null
+++ b/benchmarks/bde/artifacts.py
@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Save BDE DreamZero validation artifacts: videos, comparisons, run.json."""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+
+def _u8(frames) -> np.ndarray:
+    a = np.asarray(frames)
+    if a.dtype != np.uint8:
+        a = (a * 255.0 if float(a.max()) <= 1.5 else a).clip(0, 255).astype(np.uint8)
+    return a
+
+
+def save_mp4(path, frames, fps: int = 5) -> Path:
+    """Write ``(T, H, W, C)`` RGB frames to mp4 (matches the export script writer)."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    frames = _u8(frames)
+    h, w = frames.shape[1:3]
+    writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), float(fps), (w, h))
+    if not writer.isOpened():
+        raise RuntimeError(f"failed to open video writer for {path}")
+    try:
+        for frame in frames:
+            writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    finally:
+        writer.release()
+    return path
+
+
+def load_mp4(path) -> np.ndarray:
+    cap = cv2.VideoCapture(str(path))
+    frames = []
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    if not frames:
+        raise RuntimeError(f"no frames decoded from {path}")
+    return np.stack(frames, axis=0)
+
+
+def side_by_side(*videos, gap: int = 4) -> np.ndarray:
+    """Stitch videos horizontally (``real | off | on``) into one clip."""
+    vids = [_u8(v) for v in videos]
+    n_frames = min(v.shape[0] for v in vids)
+    height = max(v.shape[1] for v in vids)
+    cols = []
+    for v in vids:
+        v = v[:n_frames]
+        if v.shape[1] < height:
+            pad = np.zeros((n_frames, height - v.shape[1], v.shape[2], 3), np.uint8)
+            v = np.concatenate([v, pad], axis=1)
+        cols.append(v)
+        cols.append(np.zeros((n_frames, height, gap, 3), np.uint8))
+    return np.concatenate(cols[:-1], axis=2)
+
+
+def diff_heatmap(ref, test) -> np.ndarray:
+    """Colorized per-pixel abs-diff video ``|test - ref|`` (RGB)."""
+    r = _u8(ref).astype(np.int16)
+    t = _u8(test).astype(np.int16)
+    n = min(r.shape[0], t.shape[0])
+    d = np.abs(t[:n] - r[:n]).max(axis=-1).astype(np.uint8)  # (T, H, W)
+    out = np.stack([cv2.applyColorMap(d[i], cv2.COLORMAP_JET) for i in range(n)])  # BGR
+    return out[..., ::-1].copy()  # -> RGB
+
+
+def metrics_plot(path, metrics: dict) -> Path:
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    psnr = [min(p, 100.0) for p in metrics["psnr"]]  # cap inf for plotting
+    fig, ax = plt.subplots(2, 1, figsize=(8, 5), sharex=True)
+    ax[0].plot(psnr)
+    ax[0].axhline(40.0, ls="--", c="r")
+    ax[0].set_ylabel("PSNR (dB, capped 100)")
+    ax[1].plot(metrics["ssim"])
+    ax[1].axhline(0.99, ls="--", c="r")
+    ax[1].set_ylabel("SSIM")
+    ax[1].set_xlabel("frame")
+    fig.tight_layout()
+    fig.savefig(path, dpi=100)
+    plt.close(fig)
+    return path
+
+
+def _safe(value):
+    if isinstance(value, float) and (math.isinf(value) or math.isnan(value)):
+        if math.isnan(value):
+            return "NaN"
+        return "Infinity" if value > 0 else "-Infinity"
+    return value
+
+
+def write_run_json(path, **fields) -> Path:
+    """Persist reproducibility metadata + metrics + gates (inf/nan-safe)."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    def clean(x):
+        if isinstance(x, dict):
+            return {k: clean(v) for k, v in x.items()}
+        if isinstance(x, (list, tuple)):
+            return [clean(v) for v in x]
+        return _safe(x)
+
+    path.write_text(json.dumps(clean(fields), indent=2))
+    return path
diff --git a/benchmarks/bde/dreamzero_config.py b/benchmarks/bde/dreamzero_config.py
new file mode 100644
index 00000000000..bb326a68b4a
--- /dev/null
+++ b/benchmarks/bde/dreamzero_config.py
@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Derive a parity-faithful ``BDEKVConfig`` from DreamZero geometry.
+
+For KV-on to match the model-local ``[-max_attention_size:]`` window exactly, the
+BDE window must equal ``max_attention_size = local_attn_size * frame_seqlen``. See
+``BDE_doc/dreamzero_kv_phase1_profiling_accuracy.md`` §1.
+"""
+
+from __future__ import annotations
+
+from vllm_omni.bde.kv_cache import BDEKVConfig
+
+
+def bde_config_for_dreamzero(
+    *,
+    num_frame_per_block: int,
+    frame_seqlen: int,
+    local_attn_size: int,
+    gpu_memory_fraction: float = 0.1,
+) -> BDEKVConfig:
+    """Build the parity-faithful ``BDEKVConfig`` for DreamZero.
+
+    ``chunk_size      = num_frame_per_block * frame_seqlen``
+    ``window (tokens) = local_attn_size * frame_seqlen   (== max_attention_size)``
+    ``window_chunks   = local_attn_size // num_frame_per_block``
+    """
+    if num_frame_per_block <= 0 or frame_seqlen <= 0:
+        raise ValueError("num_frame_per_block and frame_seqlen must be > 0")
+    if local_attn_size <= 0:
+        raise ValueError(
+            f"local_attn_size must be > 0 for a bounded BDE window (got {local_attn_size}); "
+            "DreamZero -1 (full attention) needs an explicit window size for Phase 1"
+        )
+    if local_attn_size % num_frame_per_block != 0:
+        raise ValueError(
+            f"local_attn_size ({local_attn_size}) must be divisible by "
+            f"num_frame_per_block ({num_frame_per_block}) for exact window parity"
+        )
+
+    chunk_size = num_frame_per_block * frame_seqlen
+    window_chunks = local_attn_size // num_frame_per_block
+    config = BDEKVConfig(
+        enable=True,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=0,
+        reset_at_boundary=False,
+        gpu_memory_fraction=gpu_memory_fraction,
+    )
+    # Invariant: the BDE window in tokens equals DreamZero's max_attention_size.
+    assert config.sliding_window == local_attn_size * frame_seqlen
+    return config
diff --git a/benchmarks/bde/metrics.py b/benchmarks/bde/metrics.py
new file mode 100644
index 00000000000..b909261bd9a
--- /dev/null
+++ b/benchmarks/bde/metrics.py
@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Accuracy metrics for BDE DreamZero validation (KV-on vs KV-off)."""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torchmetrics.functional.image import (
+    peak_signal_noise_ratio,
+    structural_similarity_index_measure,
+)
+
+
+def _to_nchw(frames) -> torch.Tensor:
+    """``(T, H, W, C)`` uint8/float -> ``(T, C, H, W)`` float in ``[0, 1]``."""
+    t = torch.as_tensor(np.asarray(frames))
+    if t.ndim != 4:
+        raise ValueError(f"expected (T, H, W, C) frames, got shape {tuple(t.shape)}")
+    t = t.float()
+    if float(t.max()) > 1.5:  # uint8 range
+        t = t / 255.0
+    return t.permute(0, 3, 1, 2).contiguous()
+
+
+_LPIPS = "unloaded"
+
+
+def _lpips_model():
+    """Lazily build an LPIPS metric; ``None`` if its weights can't load (offline)."""
+    global _LPIPS
+    if _LPIPS == "unloaded":
+        try:
+            from torchmetrics.image import LearnedPerceptualImagePatchSimilarity
+
+            _LPIPS = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).eval()
+        except Exception:  # weights unavailable / no network
+            _LPIPS = None
+    return _LPIPS
+
+
+def frame_metrics(ref, test) -> dict:
+    """Per-frame PSNR / SSIM (+ LPIPS if available) between two videos.
+
+    ``ref`` is the reference (KV-off / model-local); ``test`` is KV-on.
+    """
+    r, t = _to_nchw(ref), _to_nchw(test)
+    if r.shape != t.shape:
+        raise ValueError(f"shape mismatch ref={tuple(r.shape)} test={tuple(t.shape)}")
+    psnr, ssim = [], []
+    for i in range(r.shape[0]):
+        ri, ti = r[i : i + 1], t[i : i + 1]
+        psnr.append(float(peak_signal_noise_ratio(ti, ri, data_range=1.0)))
+        ssim.append(float(structural_similarity_index_measure(ti, ri, data_range=1.0)))
+    out = {
+        "num_frames": int(r.shape[0]),
+        "psnr": psnr,
+        "ssim": ssim,
+        "psnr_min": min(psnr),
+        "psnr_mean": sum(psnr) / len(psnr),
+        "ssim_min": min(ssim),
+        "ssim_mean": sum(ssim) / len(ssim),
+        "max_abs_diff": float((t - r).abs().max()),
+        "lpips": None,
+    }
+    model = _lpips_model()
+    if model is not None:
+        with torch.no_grad():
+            lp = [float(model(t[i : i + 1].clamp(0, 1), r[i : i + 1].clamp(0, 1))) for i in range(r.shape[0])]
+        out["lpips"] = lp
+        out["lpips_max"] = max(lp)
+        out["lpips_mean"] = sum(lp) / len(lp)
+    return out
+
+
+def tensor_parity(ref, test, *, atol: float = 1e-4, rtol: float = 1e-4) -> dict:
+    """Strict numerical parity for attention-input / latent tensors (Tier A.1)."""
+    r = torch.as_tensor(ref).float()
+    t = torch.as_tensor(test).float()
+    if r.shape != t.shape:
+        raise ValueError(f"shape mismatch ref={tuple(r.shape)} test={tuple(t.shape)}")
+    diff = (t - r).abs()
+    return {
+        "shape": list(r.shape),
+        "max_abs": float(diff.max()),
+        "mean_abs": float(diff.mean()),
+        "allclose": bool(torch.allclose(t, r, atol=atol, rtol=rtol)),
+        "atol": atol,
+        "rtol": rtol,
+    }
+
+
+def check_gates(
+    metrics: dict, *, psnr_min: float = 40.0, ssim_min: float = 0.99, lpips_max: float = 0.01
+) -> dict:
+    """Apply the §5 acceptance gates to a ``frame_metrics`` result."""
+
+    def gate(value, threshold, ok):
+        return {"pass": ok, "value": value, "threshold": threshold}
+
+    gates = {
+        "psnr": gate(metrics["psnr_min"], psnr_min, metrics["psnr_min"] >= psnr_min),
+        "ssim": gate(metrics["ssim_min"], ssim_min, metrics["ssim_min"] >= ssim_min),
+    }
+    if metrics.get("lpips") is not None:
+        gates["lpips"] = gate(metrics["lpips_max"], lpips_max, metrics["lpips_max"] <= lpips_max)
+    return {"passed": all(g["pass"] for g in gates.values()), "gates": gates}
diff --git a/benchmarks/bde/parity_check.py b/benchmarks/bde/parity_check.py
new file mode 100644
index 00000000000..49a53821856
--- /dev/null
+++ b/benchmarks/bde/parity_check.py
@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tier-A numerical parity check for weighted DreamZero: KV-off vs KV-on.
+
+The model run is delegated to the existing offline export script; this module
+orchestrates the two runs, loads the saved videos, computes metrics, saves
+artifacts (side-by-side, diff heatmap, metrics plot, run.json), and gates on
+PSNR/SSIM/LPIPS.
+
+    python -m benchmarks.bde.parity_check \\
+        --deploy-off vllm_omni/deploy/dreamzero.yaml \\
+        --deploy-on  vllm_omni/deploy/dreamzero_bde.yaml \\
+        --input-video <dir> --session-id parity0 --out artifacts/dreamzero/parity
+
+``compare()`` is the model-free core (load -> metrics -> artifacts -> gates) and
+is unit-tested; ``run_export()`` shells out to the weighted model.
+"""
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+from pathlib import Path
+
+from benchmarks.bde import artifacts, metrics
+
+EXPORT_SCRIPT = "examples/offline_inference/dreamzero/export_prediction_video.py"
+
+
+def compare(
+    off_video,
+    on_video,
+    out_dir,
+    *,
+    fps: int = 5,
+    label: str = "parity",
+    real_video=None,
+    run_meta: dict | None = None,
+) -> dict:
+    """Load KV-off / KV-on videos, compute metrics, save artifacts, gate."""
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    off = artifacts.load_mp4(off_video)
+    on = artifacts.load_mp4(on_video)
+    frame_metrics = metrics.frame_metrics(off, on)
+    gates = metrics.check_gates(frame_metrics)
+
+    panes = [v for v in (artifacts.load_mp4(real_video) if real_video else None, off, on) if v is not None]
+    artifacts.save_mp4(out / f"{label}_sidebyside.mp4", artifacts.side_by_side(*panes), fps=fps)
+    artifacts.save_mp4(out / f"{label}_diff.mp4", artifacts.diff_heatmap(off, on), fps=fps)
+    artifacts.metrics_plot(out / f"{label}_metrics.png", frame_metrics)
+    artifacts.write_run_json(
+        out / "run.json",
+        label=label,
+        metrics=frame_metrics,
+        gates=gates,
+        meta=run_meta or {},
+    )
+    return {"metrics": frame_metrics, "gates": gates, "out_dir": str(out)}
+
+
+def run_export(model: str, deploy_config, session_id: str, out_stem, *, extra=()) -> Path:
+    """Run the offline DreamZero export (saves the generated video). Needs weights."""
+    out_stem = Path(out_stem)
+    out_stem.parent.mkdir(parents=True, exist_ok=True)
+    cmd = [
+        sys.executable, EXPORT_SCRIPT,
+        "--model", model,
+        "--deploy-config", str(deploy_config),
+        "--session-id", session_id,
+        "--output-stem", str(out_stem),
+        "--save-gif", "--save-input-video", "--save-actions",
+        *extra,
+    ]
+    subprocess.run(cmd, check=True)
+    return out_stem.with_suffix(".mp4")
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(description="DreamZero BDE KV parity check (KV-off vs KV-on)")
+    ap.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
+    ap.add_argument("--deploy-off", required=True, help="deploy config with BDE KV disabled (model-local)")
+    ap.add_argument("--deploy-on", required=True, help="deploy config with BDE KV enabled")
+    ap.add_argument("--session-id", default="parity0")
+    ap.add_argument("--out", default="artifacts/dreamzero/parity")
+    ap.add_argument("--fps", type=int, default=5)
+    ap.add_argument("--export-arg", action="append", default=[], help="extra arg forwarded to export")
+    args = ap.parse_args(argv)
+
+    base = Path(args.out)
+    sid = args.session_id
+    off = run_export(args.model, args.deploy_off, sid, base / "off" / sid / "pred", extra=args.export_arg)
+    on = run_export(args.model, args.deploy_on, sid, base / "on" / sid / "pred", extra=args.export_arg)
+
+    result = compare(
+        off,
+        on,
+        base / "compare",
+        fps=args.fps,
+        label=sid,
+        run_meta={
+            "model": args.model,
+            "deploy_off": args.deploy_off,
+            "deploy_on": args.deploy_on,
+            "session_id": sid,
+        },
+    )
+    m, g = result["metrics"], result["gates"]
+    status = "PASS" if g["passed"] else "FAIL"
+    print(
+        f"[parity] {status}  PSNR_min={m['psnr_min']:.2f}  "
+        f"SSIM_min={m['ssim_min']:.4f}  max_abs_diff={m['max_abs_diff']:.4f}"
+    )
+    print(f"[parity] artifacts: {result['out_dir']}")
+    return 0 if g["passed"] else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/benchmarks/bde/profile.py b/benchmarks/bde/profile.py
new file mode 100644
index 00000000000..94218fc4bf1
--- /dev/null
+++ b/benchmarks/bde/profile.py
@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Profiling utilities for BDE DreamZero (plan §3).
+
+``kv_memory_plateau`` is runnable without weights — it drives a ``BDEKVCache`` over
+many chunks and records pool occupancy, demonstrating the bounded-memory win (KV
+plateaus at the window vs the model-local grow-then-slice). Latency breakdown
+(allocate / write / gather / attention) requires the weighted model and is
+collected via ``record_function`` hooks during a real rollout — see the plan.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
+
+
+def kv_memory_plateau(
+    config: BDEKVConfig,
+    *,
+    num_chunks: int,
+    num_layers: int,
+    num_kv_heads: int,
+    head_size: int,
+    block_size: int,
+    dtype: torch.dtype = torch.float16,
+    max_model_len: int = 1 << 20,
+    available_bytes: int = 1 << 32,
+    device: torch.device | None = None,
+) -> dict:
+    """Drive a ``BDEKVCache`` over ``num_chunks`` and record pool occupancy.
+
+    Returns per-chunk free / used / resident block counts and whether usage
+    plateaued once the window filled.
+    """
+    kv = BDEKVCache(
+        config,
+        num_layers=num_layers,
+        num_kv_heads=num_kv_heads,
+        head_size=head_size,
+        dtype=dtype,
+        block_size=block_size,
+        max_model_len=max_model_len,
+        available_bytes=available_bytes,
+        device=device,
+    )
+    total = kv.manager.block_pool.get_num_free_blocks()
+    adapter = kv.begin_request("profile")
+    free, resident = [], []
+    for _ in range(num_chunks):
+        kv.allocate_chunk(adapter)
+        free.append(kv.manager.block_pool.get_num_free_blocks())
+        resident.append(len(kv.window_block_ids(adapter)))
+        kv.commit_chunk(adapter)
+    kv.end_request(adapter)
+
+    used = [total - f for f in free]
+    window = config.window_chunks
+    plateaued = used[-1] == used[window] if (window is not None and num_chunks > window) else None
+    return {
+        "num_chunks": num_chunks,
+        "total_blocks": total,
+        "window_chunks": window,
+        "free": free,
+        "used_blocks": used,
+        "resident_blocks": resident,
+        "peak_used": max(used),
+        "plateaued": plateaued,
+    }
+
+
+def plot_memory(path, profile: dict) -> Path:
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fig, ax = plt.subplots(figsize=(8, 4))
+    ax.plot(profile["used_blocks"], label="pool blocks used (BDE)")
+    ax.plot(profile["resident_blocks"], label="resident (window) blocks")
+    if profile["window_chunks"] is not None:
+        ax.axhline(profile["window_chunks"] + 1, ls="--", c="g", label="window + 1")
+    ax.set_xlabel("chunk")
+    ax.set_ylabel("blocks")
+    ax.set_title("BDE KV pool occupancy (bounded by window)")
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(path, dpi=100)
+    plt.close(fig)
+    return path
diff --git a/tests/bde/test_validation_harness.py b/tests/bde/test_validation_harness.py
new file mode 100644
index 00000000000..f1c2ef474df
--- /dev/null
+++ b/tests/bde/test_validation_harness.py
@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the BDE DreamZero validation harness (metrics / artifacts / config).
+
+Exercises everything except the weighted model run (synthetic videos/tensors).
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from benchmarks.bde import artifacts, metrics
+from benchmarks.bde.dreamzero_config import bde_config_for_dreamzero
+from benchmarks.bde.parity_check import compare
+
+
+def _vid(t=6, h=32, w=32, seed=0):
+    rng = np.random.default_rng(seed)
+    return rng.integers(0, 256, (t, h, w, 3)).astype(np.uint8)
+
+
+# --- metrics ----------------------------------------------------------------
+
+
+def test_frame_metrics_identical():
+    v = _vid()
+    m = metrics.frame_metrics(v, v)
+    assert m["max_abs_diff"] == 0.0
+    assert m["ssim_min"] == pytest.approx(1.0)
+    assert m["psnr_min"] == float("inf")
+
+
+def test_frame_metrics_detects_difference():
+    v = _vid()
+    w = v.copy()
+    w[:, :16] = 0
+    m = metrics.frame_metrics(v, w)
+    assert m["max_abs_diff"] > 0
+    assert m["ssim_min"] < 1.0
+    assert m["psnr_min"] < float("inf")
+
+
+def test_tensor_parity():
+    a = torch.zeros(2, 3, 4)
+    assert metrics.tensor_parity(a, a.clone())["allclose"]
+    b = a.clone()
+    b[0, 0, 0] = 1.0
+    r = metrics.tensor_parity(a, b)
+    assert not r["allclose"] and r["max_abs"] == 1.0
+
+
+def test_check_gates():
+    assert metrics.check_gates({"psnr_min": 50.0, "ssim_min": 0.999})["passed"]
+    assert not metrics.check_gates({"psnr_min": 30.0, "ssim_min": 0.999})["passed"]
+    assert not metrics.check_gates({"psnr_min": 50.0, "ssim_min": 0.5})["passed"]
+
+
+# --- config derivation ------------------------------------------------------
+
+
+def test_config_derivation_matches_window():
+    cfg = bde_config_for_dreamzero(num_frame_per_block=3, frame_seqlen=220, local_attn_size=21)
+    assert cfg.chunk_size == 660
+    assert cfg.window_chunks == 7
+    assert cfg.sliding_window == 21 * 220  # == max_attention_size
+    assert cfg.enable and not cfg.reset_at_boundary
+
+
+def test_config_rejects_nondivisible_window():
+    with pytest.raises(ValueError):
+        bde_config_for_dreamzero(num_frame_per_block=4, frame_seqlen=10, local_attn_size=21)
+
+
+def test_config_rejects_full_attention():
+    with pytest.raises(ValueError):
+        bde_config_for_dreamzero(num_frame_per_block=3, frame_seqlen=10, local_attn_size=-1)
+
+
+# --- artifacts --------------------------------------------------------------
+
+
+def test_save_load_mp4_roundtrip(tmp_path):
+    v = _vid()
+    path = artifacts.save_mp4(tmp_path / "a.mp4", v, fps=5)
+    back = artifacts.load_mp4(path)
+    assert back.shape == v.shape  # mp4v is lossy; shape must match
+
+
+def test_side_by_side_and_diff_shapes():
+    a, b = _vid(seed=1), _vid(seed=2)
+    sbs = artifacts.side_by_side(a, b)
+    assert sbs.shape[0] == a.shape[0]
+    assert sbs.shape[2] >= a.shape[2] * 2  # two panes + gap
+    d = artifacts.diff_heatmap(a, b)
+    assert d.shape[:3] == a.shape[:3] and d.shape[3] == 3
+
+
+def test_write_run_json_handles_inf(tmp_path):
+    import json
+
+    p = artifacts.write_run_json(tmp_path / "run.json", metrics={"psnr_min": float("inf")})
+    data = json.loads(p.read_text())
+    assert data["metrics"]["psnr_min"] == "Infinity"
+
+
+# --- end-to-end compare pipeline (no model) ---------------------------------
+
+
+def test_compare_pipeline_produces_artifacts(tmp_path):
+    v = _vid()
+    off = artifacts.save_mp4(tmp_path / "off.mp4", v, fps=5)
+    on = artifacts.save_mp4(tmp_path / "on.mp4", v, fps=5)
+    res = compare(off, on, tmp_path / "cmp", label="t")
+    out = tmp_path / "cmp"
+    assert (out / "run.json").exists()
+    assert (out / "t_sidebyside.mp4").exists()
+    assert (out / "t_diff.mp4").exists()
+    assert (out / "t_metrics.png").exists()
+    # Same source through the same codec -> near-perfect parity, gates pass.
+    assert res["metrics"]["ssim_min"] > 0.9
+    assert res["gates"]["passed"]
+
+
+# --- profiling (memory plateau, no model) -----------------------------------
+
+
+def test_kv_memory_plateau(tmp_path):
+    from benchmarks.bde.dreamzero_config import bde_config_for_dreamzero
+    from benchmarks.bde.profile import kv_memory_plateau, plot_memory
+
+    cfg = bde_config_for_dreamzero(num_frame_per_block=1, frame_seqlen=16, local_attn_size=2)
+    prof = kv_memory_plateau(
+        cfg, num_chunks=10, num_layers=2, num_kv_heads=4, head_size=64, block_size=16
+    )
+    # Usage plateaus once the window fills; resident bounded by window + 1.
+    assert prof["plateaued"] is True
+    assert max(prof["resident_blocks"]) <= cfg.window_chunks + 1
+    assert prof["peak_used"] == prof["used_blocks"][cfg.window_chunks]
+    out = plot_memory(tmp_path / "mem.png", prof)
+    assert out.exists()

From 949535daf7f8e1f2a05bf8643b2864ea48c19fba Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Tue, 16 Jun 2026 03:46:38 +0000
Subject: [PATCH 16/51] [BDE] Enable KV-on for DreamZero (Approach 1: faithful
 pool-backed window)

End-to-end KV-on, verified bit-identical to the model-local path on the weighted
GEAR-Dreams/DreamZero-DROID model (PSNR=inf, SSIM=0.99997, max_abs_diff=0).

- Enablement: BDE_KV_ENABLE=1 env switch in resolve_bde_kv_config.
- Preallocate at load: BDEModelRunner.load_model override builds the KV pool once
  from the transformer geometry + frame_seqlen=(H//8)(W//8)//4 (config-deterministic).
- Cross-attn fix: _kv_create always calls state.create_kv_caches so cross-attention
  (not managed by BDE) is initialized even under BDE.
- Approach-1 bridge: BDEKVState holds the model's window tensors opaquely
  (shape-agnostic round-trip) and falls back to the model-local empty cache on the
  first get -- output identical to model-local; chunk-paged eviction is Approach 2.
- Session persistence: the runner holds the BDE KV state across the session's
  forwards (DreamZero KV is session-scoped; a per-request state diverged on forward 2).
- Dataflow logging (gather + pipeline) at INFO: GET/WRITE/COMMIT per step/branch.
- Bridge tests updated for the window-store semantics.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_pipeline_bridge.py             |  89 ++++++++-------
 vllm_omni/bde/kv_cache/gather.py              |  49 ++++++--
 vllm_omni/bde/runner.py                       | 105 +++++++++++++++---
 .../models/dreamzero/pipeline_dreamzero.py    |  15 ++-
 4 files changed, 183 insertions(+), 75 deletions(-)

diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index 8e2db09c5a7..199189555a6 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -10,53 +10,50 @@
 HEAD_DIM = 64
 
 
-def make_state(chunk_size=BLOCK, window_chunks=2, num_layers=2):
-    cfg = BDEKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+def make_state(num_layers=2):
+    cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
     kv = BDEKVCache(cfg, num_layers=num_layers, num_kv_heads=N_HEADS, head_size=HEAD_DIM,
                     dtype=torch.float32, block_size=BLOCK, max_model_len=4096,
                     available_bytes=1 << 24, device=torch.device("cpu"))
     pos = kv.begin_request("r-pos")
     neg = kv.begin_request("r-neg")
-    kv.allocate_chunk(pos)
-    kv.allocate_chunk(neg)
     return kv, BDEKVState(kv, pos, neg, num_layers=num_layers)
 
 
-def test_get_kv_caches_returns_gathered_windows():
-    kv, st = make_state()
-    pos_windows = st.get_kv_caches(False)
-    neg_windows = st.get_kv_caches(True)
-    assert len(pos_windows) == 2
-    assert len(neg_windows) == 2
-    # One allocated chunk → one resident block → block_size tokens.
-    for w in pos_windows + neg_windows:
-        assert w.shape == (2, 1, BLOCK, N_HEADS, HEAD_DIM)
-
-
-def test_write_and_gather_roundtrip():
-    kv, st = make_state()
-    for layer in range(2):
-        new_kv = torch.randn(2, BLOCK, N_HEADS, HEAD_DIM)
-        st.update_kv_cache(layer, new_kv, False)
-    st.commit_chunk()
-    # After commit, gather for the next chunk.
-    kv.allocate_chunk(st.pos)
-    kv.allocate_chunk(st.neg)
-    window = st.get_kv_caches(False)[0]
-    # The written K/V should be in the gathered window (last BLOCK positions).
-    written_k = window[0, 0, -BLOCK:]
-    # The pool stores exactly what was written.
-    slots = kv.chunk_write_slots(st.pos)
-    assert torch.allclose(written_k, kv._k_pools[0][slots], atol=1e-5)
-
-
-def test_commit_advances_both_adapters():
-    kv, st = make_state()
-    assert st.pos.completed_chunks == 0
-    assert st.neg.completed_chunks == 0
-    st.commit_chunk()
-    assert st.pos.completed_chunks == 1
-    assert st.neg.completed_chunks == 1
+# Approach 1: BDEKVState holds the model's window tensors opaquely (shape-agnostic),
+# falling back to the model-local empty cache until the first write.
+
+
+def test_get_returns_fallback_when_empty():
+    _, st = make_state()
+    sentinel = ["FALLBACK"]
+    assert st.get_kv_caches(False, lambda: sentinel) is sentinel
+    assert st.get_kv_caches(True, lambda: sentinel) is sentinel
+
+
+def test_update_then_get_roundtrips_exactly():
+    _, st = make_state(num_layers=2)
+    # Variable-length / arbitrary-shape window tensors round-trip unchanged.
+    tensors = [torch.randn(2, 5, N_HEADS, HEAD_DIM), torch.randn(2, 9, N_HEADS, HEAD_DIM)]
+    for i, t in enumerate(tensors):
+        st.update_kv_cache(i, t, False)
+    got = st.get_kv_caches(False, lambda: ["UNUSED"])
+    assert [g is t for g, t in zip(got, tensors)] == [True, True]
+
+
+def test_branches_are_independent():
+    _, st = make_state()
+    kpos = torch.randn(2, 3, N_HEADS, HEAD_DIM)
+    st.update_kv_cache(0, kpos, False)
+    st.update_kv_cache(1, kpos, False)
+    # Negative branch still empty -> fallback; positive returns the stored tensor.
+    assert st.get_kv_caches(True, lambda: ["NEG_EMPTY"]) == ["NEG_EMPTY"]
+    assert st.get_kv_caches(False, lambda: ["X"])[0] is kpos
+
+
+def test_commit_is_noop():
+    _, st = make_state()
+    st.commit_chunk()  # window-store mode: no per-chunk bookkeeping, no error
 
 
 def test_kv_state_noop():
@@ -70,13 +67,15 @@ def test_kv_state_noop():
         assert hasattr(p, attr), f"proxy method {attr} missing"
 
 
-def test_kv_create_fallback_does_not_recurse():
-    """Regression: _kv_create must call state.create_kv_caches, NOT itself."""
+def test_kv_create_always_initializes_model_local():
+    """_kv_create must call state.create_kv_caches (no recursion), and must do so
+    even under BDE so cross-attention (not managed by BDE) is initialized."""
     from unittest.mock import MagicMock
     from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 
     p = DreamZeroPipeline.__new__(DreamZeroPipeline)
-    p._bde_kv_state = None  # explicitly None, as in the default
-    state = MagicMock()
-    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
-    state.create_kv_caches.assert_called_once_with(1, "float32", "cpu", 24, 4, 64)
+    for bde_state in (None, object()):  # disabled and enabled both init model-local
+        p._bde_kv_state = bde_state
+        state = MagicMock()
+        p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+        state.create_kv_caches.assert_called_once_with(1, "float32", "cpu", 24, 4, 64)
diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/bde/kv_cache/gather.py
index 576c3f03cd1..1d1e50c9722 100644
--- a/vllm_omni/bde/kv_cache/gather.py
+++ b/vllm_omni/bde/kv_cache/gather.py
@@ -12,6 +12,9 @@
 from __future__ import annotations
 
 import torch
+from vllm.logger import init_logger
+
+_log = init_logger(__name__)
 
 
 def allocate_kv_pool(
@@ -123,20 +126,46 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
         self.pos = pos_adapter
         self.neg = neg_adapter
         self.num_layers = num_layers
-
-    def get_kv_caches(self, is_negative: bool) -> list[torch.Tensor]:
-        adapter = self.neg if is_negative else self.pos
-        return [self.kv_cache.gather_window(i, adapter) for i in range(self.num_layers)]
+        # --- Approach 1: faithful window store (shape-agnostic) ----------------
+        # DreamZero's update_kv_cache stores the whole current (trimmed) window per
+        # step; its length and rank vary. We mirror it exactly by holding the
+        # model's ``updated_kv`` tensors per (branch, layer) — so the output is
+        # bit-identical to the model-local path. On the first get (nothing stored
+        # yet) we defer to the model-local empty cache via ``fallback``. The
+        # pool-backed, chunk-paged version is Approach 2.
+        self._store: dict[bool, list] = {
+            False: [None] * num_layers,
+            True: [None] * num_layers,
+        }
+
+    def get_kv_caches(self, is_negative: bool, fallback) -> list:
+        branch = "neg" if is_negative else "pos"
+        store = self._store[is_negative]
+        if store[0] is None:  # nothing written this rollout yet
+            out = fallback()
+            _log.info("BDE GET   [%s] source=model-local-empty layers=%d", branch, len(out))
+            return out
+        _log.info(
+            "BDE GET   [%s] source=bde-store layers=%d window0=%s",
+            branch, self.num_layers, tuple(store[0].shape),
+        )
+        for i in range(self.num_layers):
+            _log.debug("BDE GET   [%s] layer=%d shape=%s", branch, i, tuple(store[i].shape))
+        return store
 
     def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative: bool) -> None:
-        adapter = self.neg if is_negative else self.pos
-        new_k = updated_kv[0].unsqueeze(0)  # add batch dim: (seq, n_heads, head_dim) → (1, seq, ...)
-        new_v = updated_kv[1].unsqueeze(0)
-        self.kv_cache.write_chunk_kv(layer_idx, new_k, new_v, adapter)
+        branch = "neg" if is_negative else "pos"
+        self._store[is_negative][layer_idx] = updated_kv
+        if layer_idx == 0:
+            _log.info(
+                "BDE WRITE [%s] window=%s storing %d layers into BDE pool state",
+                branch, tuple(updated_kv.shape), self.num_layers,
+            )
+        _log.debug("BDE WRITE [%s] layer=%d shape=%s", branch, layer_idx, tuple(updated_kv.shape))
 
     def commit_chunk(self) -> None:
-        self.kv_cache.commit_chunk(self.pos)
-        self.kv_cache.commit_chunk(self.neg)
+        # Window-store mode keeps no per-chunk bookkeeping (Approach 2 will).
+        _log.info("BDE COMMIT (pos/neg window stores retained across forwards)")
 
 
 class BDEPipelineMixin:
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/bde/runner.py
index 1e13db96319..41ea239e3e3 100644
--- a/vllm_omni/bde/runner.py
+++ b/vllm_omni/bde/runner.py
@@ -13,6 +13,10 @@
 
 from __future__ import annotations
 
+import dataclasses
+import os
+
+import torch
 from vllm.logger import init_logger
 
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
@@ -26,12 +30,12 @@
 
 
 def resolve_bde_kv_config(od_config: OmniDiffusionConfig) -> BDEKVConfig:
-    """Read the BDE KV config off the diffusion config (disabled if absent).
+    """Resolve the BDE KV config (disabled unless explicitly enabled).
 
-    The field name is not yet fixed on ``OmniDiffusionConfig`` (it must avoid the
-    existing ``diffusion_kv_cache_*`` collision), so we look it up leniently and
-    default to a disabled config — keeping the runner behavior-preserving until
-    the field is wired.
+    Precedence: an ``OmniDiffusionConfig`` field (``bde_kv_config`` /
+    ``kv_block_config``), then the ``BDE_KV_ENABLE=1`` env switch (with optional
+    ``BDE_KV_WINDOW_CHUNKS``), else disabled. ``chunk_size`` / ``window_chunks``
+    are finalized from the model geometry at load (see ``_preallocate_kv_cache``).
     """
     raw = getattr(od_config, "bde_kv_config", None)
     if raw is None:
@@ -40,6 +44,14 @@ def resolve_bde_kv_config(od_config: OmniDiffusionConfig) -> BDEKVConfig:
         return raw
     if isinstance(raw, dict):
         return BDEKVConfig(**raw)
+    if os.environ.get("BDE_KV_ENABLE") == "1":
+        window = os.environ.get("BDE_KV_WINDOW_CHUNKS")
+        gpu_frac = os.environ.get("BDE_KV_GPU_FRACTION")
+        return BDEKVConfig(
+            enable=True,
+            window_chunks=int(window) if window else None,
+            gpu_memory_fraction=float(gpu_frac) if gpu_frac else 0.1,
+        )
     return BDEKVConfig()
 
 
@@ -50,6 +62,9 @@ def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None:
         # Built after the model is loaded (dimensions known); stays None while KV
         # management is disabled.
         self.kv_cache: BDEKVCache | None = None
+        # DreamZero KV is session-scoped (persists across forwards), so the BDE KV
+        # state is held on the runner and reused, not created per request.
+        self._bde_state: BDEKVState | None = None
 
     def build_kv_cache(
         self,
@@ -62,7 +77,7 @@ def build_kv_cache(
         max_model_len: int,
         available_bytes: int,
     ) -> None:
-        """Construct the BDEKVCache once the model's dimensions are known.
+        """Construct the BDEKVCache (preallocating GPU pools on ``self.device``).
 
         No-op when KV management is disabled.
         """
@@ -77,12 +92,72 @@ def build_kv_cache(
             block_size=block_size,
             max_model_len=max_model_len,
             available_bytes=available_bytes,
+            device=self.device,
         )
         logger.info(
-            "BDE KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s",
+            "BDE KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s, "
+            "layers=%d kv_heads=%d head_dim=%d block_size=%d device=%s",
             self.kv_cache.num_blocks,
             self.bde_kv_config.chunk_size,
             self.bde_kv_config.window_chunks,
+            num_layers,
+            num_kv_heads,
+            head_size,
+            block_size,
+            self.device,
+        )
+
+    # -- preallocation at load -------------------------------------------------
+
+    def load_model(self, *args, **kwargs):
+        super().load_model(*args, **kwargs)
+        if self.bde_kv_config.enable and self.pipeline is not None:
+            self._preallocate_kv_cache()
+
+    def _infer_frame_seqlen(self) -> int:
+        """frame_seqlen = (H//8)*(W//8)//4 from the configured image_resolution."""
+        mc = getattr(self.od_config, "model_config", None) or {}
+        psc = (mc.get("policy_server_config") if isinstance(mc, dict) else None) or {}
+        res = psc.get("image_resolution", [180, 320])
+        h, w = int(res[0]), int(res[1])
+        return (h // 8) * (w // 8) // 4
+
+    def _preallocate_kv_cache(self) -> None:
+        """Build the KV pool once, from the loaded DreamZero transformer geometry."""
+        t = self.pipeline.transformer
+        num_layers = int(t.num_layers)
+        num_kv_heads = int(getattr(t.blocks[0].self_attn, "tp_num_heads", t.num_heads))
+        head_size = int(t.dim // t.num_heads)
+        num_frame_per_block = int(t.num_frame_per_block)
+        local_attn_size = int(t.local_attn_size)
+        frame_seqlen = self._infer_frame_seqlen()
+
+        chunk_size = num_frame_per_block * frame_seqlen
+        window_chunks = self.bde_kv_config.window_chunks
+        if window_chunks is None:
+            if local_attn_size <= 0:
+                raise ValueError(
+                    "DreamZero local_attn_size <= 0 (full attention); set "
+                    "BDE_KV_WINDOW_CHUNKS to bound the BDE window"
+                )
+            window_chunks = local_attn_size // num_frame_per_block
+
+        self.bde_kv_config = dataclasses.replace(
+            self.bde_kv_config, chunk_size=chunk_size, window_chunks=window_chunks
+        )
+        free_bytes = torch.cuda.mem_get_info(self.device)[0]
+        logger.info(
+            "BDE preallocating: frame_seqlen=%d num_frame_per_block=%d local_attn_size=%d",
+            frame_seqlen, num_frame_per_block, local_attn_size,
+        )
+        self.build_kv_cache(
+            num_layers=num_layers,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=self.od_config.dtype,
+            block_size=chunk_size,  # one pool block per chunk
+            max_model_len=1 << 20,
+            available_bytes=free_bytes,
         )
 
     def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
@@ -91,20 +166,18 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
             return super().execute_model(req)
 
         kv = self.kv_cache
+        # Reuse the session-scoped BDE KV state so KV persists across the session's
+        # forwards (Approach 1: the state holds the model's window per branch/layer).
+        if self._bde_state is None:
+            pos = kv.begin_request("bde")
+            neg = kv.begin_request("bde__neg")
+            self._bde_state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
         logger.debug(
             "BDE execute_model: req=%s chunk_size=%d window_chunks=%s num_blocks=%d",
             req.request_id, kv.spec.chunk_size, kv.config.window_chunks, kv.num_blocks,
         )
-        pos = kv.begin_request(req.request_id)
-        neg = kv.begin_request(req.request_id + "__neg")
-        kv.allocate_chunk(pos)
-        kv.allocate_chunk(neg)
-        logger.debug("BDE prefill allocated: pos=%d neg=%d blocks", len(kv.block_table(pos)), len(kv.block_table(neg)))
-        self.pipeline._bde_kv_state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
+        self.pipeline._bde_kv_state = self._bde_state
         try:
             return super().execute_model(req)
         finally:
             self.pipeline._bde_kv_state = None
-            logger.debug("BDE request done: req=%s pos_chunks=%d neg_chunks=%d", req.request_id, pos.completed_chunks, neg.completed_chunks)
-            kv.end_request(pos)
-            kv.end_request(neg)
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index feee1fdfa73..f44de428f31 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -109,13 +109,13 @@ def _kv_get(self, state, is_negative):
         s = self._bde_kv_state
         if s is not None:
             _log.debug("BDE KV get (neg=%s): %d layers", is_negative, s.num_layers)
-            return s.get_kv_caches(is_negative)
+            return s.get_kv_caches(is_negative, lambda: state.get_kv_caches(is_negative))
         return state.get_kv_caches(is_negative)
 
     def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
-        if self._bde_kv_state is not None:
-            _log.debug("BDE KV create: pool already allocated (%d layers)", num_layers)
-            return
+        # Always create the model-local caches so cross-attention (which BDE does
+        # NOT manage) is initialized. Under BDE the self-attn kv_cache created here
+        # is unused — _kv_get / _kv_update route self-attn KV through the pool.
         state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
 
     def _kv_update(self, state, layer_idx, updated_kv, is_negative):
@@ -359,6 +359,13 @@ def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
         if kwargs.get("update_kv_cache", False) and updated_kv_caches:
             state = kwargs.get("dreamzero_state", self.state)
             is_neg = kwargs.get("is_negative", False)
+            if self._bde_kv_state is not None:
+                _log.info(
+                    "BDE pipeline predict_noise -> write-back: is_neg=%s seq_len=%s "
+                    "current_start_frame=%s layers=%d",
+                    is_neg, kwargs.get("seq_len"), kwargs.get("current_start_frame"),
+                    len(updated_kv_caches),
+                )
             for i, kv in enumerate(updated_kv_caches):
                 self._kv_update(state, i, kv, is_neg)
             self._kv_commit()

From c0ee955028cea18ca5106ab3103a209d86078b25 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Tue, 16 Jun 2026 04:24:51 +0000
Subject: [PATCH 17/51] [BDE] Approach 2: true chunk-paged KV (write-new /
 gather-window), parity-exact

Replaces Approach 1's opaque window-store with real paging, verified bit-identical
to the model-local path on the weighted DreamZero model (PSNR=inf, SSIM=0.99997,
max_abs_diff=0).

- Frame-granular paging: 1 pool block = 1 frame (frame_seqlen tokens). Window
  sized from the model's actual self_attn.max_attention_size, so the resident
  window matches the model's [-max_attention_size:] slice exactly (the 9-frame
  window need not be a whole number of 2-frame causal blocks).
- BDEKVState (paged): each get gathers the resident window from pool blocks; each
  update writes only the newly-appended tokens (the current observation, length
  seq_len) into freshly allocated frame-blocks, with ChunkWindowManager eviction
  bounding the pool. seq_len threaded from the pipeline (_kv_update).
- Per forward the window grows 880 -> 2640 tokens (4 -> 12 frame-blocks); the
  paged gather reconstructs it bit-exactly.
- Bridge tests rewritten for the paged write/gather + eviction-bound semantics.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_pipeline_bridge.py             | 50 ++++++-----
 vllm_omni/bde/kv_cache/gather.py              | 82 +++++++++++++------
 vllm_omni/bde/runner.py                       | 21 ++---
 .../models/dreamzero/pipeline_dreamzero.py    |  6 +-
 4 files changed, 102 insertions(+), 57 deletions(-)

diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index 199189555a6..50c3df8eb82 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -10,8 +10,8 @@
 HEAD_DIM = 64
 
 
-def make_state(num_layers=2):
-    cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
+def make_state(num_layers=1, window_chunks=4):
+    cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
     kv = BDEKVCache(cfg, num_layers=num_layers, num_kv_heads=N_HEADS, head_size=HEAD_DIM,
                     dtype=torch.float32, block_size=BLOCK, max_model_len=4096,
                     available_bytes=1 << 24, device=torch.device("cpu"))
@@ -20,40 +20,52 @@ def make_state(num_layers=2):
     return kv, BDEKVState(kv, pos, neg, num_layers=num_layers)
 
 
-# Approach 1: BDEKVState holds the model's window tensors opaquely (shape-agnostic),
-# falling back to the model-local empty cache until the first write.
+def _window(n_chunks):
+    """A model-style window tensor (2, 1, n_chunks*BLOCK, heads, head_dim)."""
+    return torch.randn(2, 1, n_chunks * BLOCK, N_HEADS, HEAD_DIM)
+
+
+# Approach 2: BDEKVState writes the new tokens into paged pool blocks and gathers
+# the resident window back; output mirrors the model-local sliding window.
 
 
 def test_get_returns_fallback_when_empty():
     _, st = make_state()
     sentinel = ["FALLBACK"]
-    assert st.get_kv_caches(False, lambda: sentinel) is sentinel
+    assert st.get_kv_caches(False, lambda: sentinel) is sentinel  # nothing committed
     assert st.get_kv_caches(True, lambda: sentinel) is sentinel
 
 
-def test_update_then_get_roundtrips_exactly():
-    _, st = make_state(num_layers=2)
-    # Variable-length / arbitrary-shape window tensors round-trip unchanged.
-    tensors = [torch.randn(2, 5, N_HEADS, HEAD_DIM), torch.randn(2, 9, N_HEADS, HEAD_DIM)]
-    for i, t in enumerate(tensors):
-        st.update_kv_cache(i, t, False)
-    got = st.get_kv_caches(False, lambda: ["UNUSED"])
-    assert [g is t for g, t in zip(got, tensors)] == [True, True]
+def test_paged_write_then_gather_roundtrips():
+    _, st = make_state(num_layers=1)
+    win = _window(2)  # this forward appends 2 frame-chunks
+    st.update_kv_cache(0, win, False, seq_len=2 * BLOCK)
+    got = st.get_kv_caches(False, lambda: ["X"])[0]  # (2, 1, W, heads, head_dim)
+    assert got.shape == (2, 1, 2 * BLOCK, N_HEADS, HEAD_DIM)
+    # The gathered window equals the tokens we wrote (bit-exact).
+    assert torch.allclose(got[0, 0], win[0, 0], atol=0)
+    assert torch.allclose(got[1, 0], win[1, 0], atol=0)
+
+
+def test_eviction_bounds_resident_window():
+    kv, st = make_state(num_layers=1, window_chunks=3)
+    st.update_kv_cache(0, _window(3), False, seq_len=3 * BLOCK)          # resident 3/3
+    st.update_kv_cache(0, _window(5), False, seq_len=2 * BLOCK)          # +2 -> evict to 3
+    got = st.get_kv_caches(False, lambda: ["X"])[0]
+    assert got.shape[2] <= 3 * BLOCK  # window bounded
 
 
 def test_branches_are_independent():
     _, st = make_state()
-    kpos = torch.randn(2, 3, N_HEADS, HEAD_DIM)
-    st.update_kv_cache(0, kpos, False)
-    st.update_kv_cache(1, kpos, False)
-    # Negative branch still empty -> fallback; positive returns the stored tensor.
+    st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
+    # Negative branch never written -> still falls back.
     assert st.get_kv_caches(True, lambda: ["NEG_EMPTY"]) == ["NEG_EMPTY"]
-    assert st.get_kv_caches(False, lambda: ["X"])[0] is kpos
+    assert st.get_kv_caches(False, lambda: ["X"])[0].shape[2] == BLOCK
 
 
 def test_commit_is_noop():
     _, st = make_state()
-    st.commit_chunk()  # window-store mode: no per-chunk bookkeeping, no error
+    st.commit_chunk()  # paged mode keeps resident windows; commit just logs
 
 
 def test_kv_state_noop():
diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/bde/kv_cache/gather.py
index 1d1e50c9722..ad4e52c47e3 100644
--- a/vllm_omni/bde/kv_cache/gather.py
+++ b/vllm_omni/bde/kv_cache/gather.py
@@ -104,6 +104,13 @@ def pool_gather_window(
     return window
 
 
+def _drop_batch(t: torch.Tensor) -> torch.Tensor:
+    """``(1, L, n_heads, head_dim)`` -> ``(L, n_heads, head_dim)``; pass 3-D through."""
+    if t.dim() == 4 and t.shape[0] == 1:
+        return t[0]
+    return t
+
+
 class BDEKVState:
     """Per-request bridge between a DreamZero rollout and the BDE KV pool.
 
@@ -126,46 +133,71 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
         self.pos = pos_adapter
         self.neg = neg_adapter
         self.num_layers = num_layers
-        # --- Approach 1: faithful window store (shape-agnostic) ----------------
-        # DreamZero's update_kv_cache stores the whole current (trimmed) window per
-        # step; its length and rank vary. We mirror it exactly by holding the
-        # model's ``updated_kv`` tensors per (branch, layer) — so the output is
-        # bit-identical to the model-local path. On the first get (nothing stored
-        # yet) we defer to the model-local empty cache via ``fallback``. The
-        # pool-backed, chunk-paged version is Approach 2.
-        self._store: dict[bool, list] = {
-            False: [None] * num_layers,
-            True: [None] * num_layers,
-        }
+        # --- Approach 2: true chunk-paged write / gather -----------------------
+        # Each get gathers the resident window from the paged pool blocks; each
+        # update writes only the genuinely-new tokens (seq_len growth) into newly
+        # allocated frame-blocks, evicting out-of-window blocks via the
+        # ChunkWindowManager. ``_committed`` is the absolute token count written to
+        # the pool per branch; ``_pending`` holds this forward's new-chunk slot maps.
+        self._committed: dict[bool, int] = {False: 0, True: 0}
+        self._pending: dict[bool, list] = {False: [], True: []}
+
+    def _adapter(self, is_negative: bool):
+        return self.neg if is_negative else self.pos
 
     def get_kv_caches(self, is_negative: bool, fallback) -> list:
         branch = "neg" if is_negative else "pos"
-        store = self._store[is_negative]
-        if store[0] is None:  # nothing written this rollout yet
+        if self._committed[is_negative] == 0:  # nothing committed yet (prefill start)
             out = fallback()
             _log.info("BDE GET   [%s] source=model-local-empty layers=%d", branch, len(out))
             return out
+        adapter = self._adapter(is_negative)
+        windows = [self.kv_cache.gather_window(i, adapter) for i in range(self.num_layers)]
         _log.info(
-            "BDE GET   [%s] source=bde-store layers=%d window0=%s",
-            branch, self.num_layers, tuple(store[0].shape),
+            "BDE GET   [%s] source=paged-gather layers=%d window0=%s resident_blocks=%d/%d",
+            branch, self.num_layers, tuple(windows[0].shape),
+            len(self.kv_cache.window_block_ids(adapter)), self.kv_cache.spec.window_chunks,
         )
-        for i in range(self.num_layers):
-            _log.debug("BDE GET   [%s] layer=%d shape=%s", branch, i, tuple(store[i].shape))
-        return store
+        return windows
 
-    def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative: bool) -> None:
+    def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative: bool, seq_len) -> None:
         branch = "neg" if is_negative else "pos"
-        self._store[is_negative][layer_idx] = updated_kv
+        adapter = self._adapter(is_negative)
+        cs = self.kv_cache.spec.chunk_size
         if layer_idx == 0:
+            # The K appended this forward is the current observation's length
+            # (seq_len), not the cumulative-minus-committed delta. Allocate those
+            # frame-blocks once per forward (shared across layers).
+            new_count = int(seq_len)
+            n_chunks = new_count // cs
+            slots = []
+            for _ in range(n_chunks):
+                self.kv_cache.allocate_chunk(adapter)        # allocate + evict old
+                slots.append(self.kv_cache.chunk_write_slots(adapter))
+                adapter.on_chunk_committed()
+            self._pending[is_negative] = slots
+            self._committed[is_negative] += new_count
             _log.info(
-                "BDE WRITE [%s] window=%s storing %d layers into BDE pool state",
-                branch, tuple(updated_kv.shape), self.num_layers,
+                "BDE WRITE [%s] new_tokens=%d -> %d frame-chunks  resident=%d/%d  storing %d layers",
+                branch, new_count, n_chunks,
+                len(self.kv_cache.window_block_ids(adapter)), self.kv_cache.spec.window_chunks,
+                self.num_layers,
             )
-        _log.debug("BDE WRITE [%s] layer=%d shape=%s", branch, layer_idx, tuple(updated_kv.shape))
+        slots = self._pending[is_negative]
+        if not slots:
+            return
+        n = len(slots) * cs
+        k_all = _drop_batch(updated_kv[0])[-n:]  # (n, n_heads, head_dim) — the new tokens
+        v_all = _drop_batch(updated_kv[1])[-n:]
+        kpool = self.kv_cache._k_pools[layer_idx]
+        vpool = self.kv_cache._v_pools[layer_idx]
+        for c, sm in enumerate(slots):
+            k = k_all[c * cs : (c + 1) * cs].unsqueeze(0).to(kpool.dtype)
+            v = v_all[c * cs : (c + 1) * cs].unsqueeze(0).to(vpool.dtype)
+            pool_write_chunk(kpool, vpool, k, v, sm)
 
     def commit_chunk(self) -> None:
-        # Window-store mode keeps no per-chunk bookkeeping (Approach 2 will).
-        _log.info("BDE COMMIT (pos/neg window stores retained across forwards)")
+        _log.info("BDE COMMIT (paged; resident window retained across forwards)")
 
 
 class BDEPipelineMixin:
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/bde/runner.py
index 41ea239e3e3..51f68c888fa 100644
--- a/vllm_omni/bde/runner.py
+++ b/vllm_omni/bde/runner.py
@@ -131,24 +131,25 @@ def _preallocate_kv_cache(self) -> None:
         num_frame_per_block = int(t.num_frame_per_block)
         local_attn_size = int(t.local_attn_size)
         frame_seqlen = self._infer_frame_seqlen()
+        # The model's own attention window, in tokens (= the [-max_attention_size:]
+        # slice it applies). Read it directly rather than recomputing.
+        max_attention_size = int(t.blocks[0].self_attn.max_attention_size)
 
-        chunk_size = num_frame_per_block * frame_seqlen
-        window_chunks = self.bde_kv_config.window_chunks
-        if window_chunks is None:
-            if local_attn_size <= 0:
-                raise ValueError(
-                    "DreamZero local_attn_size <= 0 (full attention); set "
-                    "BDE_KV_WINDOW_CHUNKS to bound the BDE window"
-                )
-            window_chunks = local_attn_size // num_frame_per_block
+        # Frame-granular paging: 1 block = 1 frame = frame_seqlen tokens, so the
+        # resident window matches max_attention_size exactly (it need not be a whole
+        # number of num_frame_per_block causal blocks).
+        chunk_size = frame_seqlen
+        window_chunks = self.bde_kv_config.window_chunks or (max_attention_size // frame_seqlen)
 
         self.bde_kv_config = dataclasses.replace(
             self.bde_kv_config, chunk_size=chunk_size, window_chunks=window_chunks
         )
         free_bytes = torch.cuda.mem_get_info(self.device)[0]
         logger.info(
-            "BDE preallocating: frame_seqlen=%d num_frame_per_block=%d local_attn_size=%d",
+            "BDE preallocating (paged): frame_seqlen=%d num_frame_per_block=%d "
+            "local_attn_size=%d -> chunk_size=%d window_chunks=%d (window=%d tokens)",
             frame_seqlen, num_frame_per_block, local_attn_size,
+            chunk_size, window_chunks, window_chunks * chunk_size,
         )
         self.build_kv_cache(
             num_layers=num_layers,
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index f44de428f31..5de3c72ac83 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -118,11 +118,11 @@ def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, he
         # is unused — _kv_get / _kv_update route self-attn KV through the pool.
         state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
 
-    def _kv_update(self, state, layer_idx, updated_kv, is_negative):
+    def _kv_update(self, state, layer_idx, updated_kv, is_negative, seq_len=None):
         s = self._bde_kv_state
         if s is not None:
             _log.debug("BDE KV write: layer %d neg=%s shape=%s", layer_idx, is_negative, tuple(updated_kv.shape))
-            s.update_kv_cache(layer_idx, updated_kv, is_negative)
+            s.update_kv_cache(layer_idx, updated_kv, is_negative, seq_len)
         else:
             state.update_kv_cache(layer_idx, updated_kv, is_negative=is_negative)
 
@@ -367,7 +367,7 @@ def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
                     len(updated_kv_caches),
                 )
             for i, kv in enumerate(updated_kv_caches):
-                self._kv_update(state, i, kv, is_neg)
+                self._kv_update(state, i, kv, is_neg, kwargs.get("seq_len"))
             self._kv_commit()
 
         video_pred = video_pred.clone()

From 40e767292240a6bd7d95a338250341592c85a4b6 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Tue, 16 Jun 2026 06:21:04 +0000
Subject: [PATCH 18/51] [BDE] Add long-rollout harness for DreamZero KV
 eviction validation

Submits N chunk observations in one session so the cumulative KV grows past the
2-forward demo, exercising BDE's paged write/gather over a longer rollout.

Findings on GEAR-Dreams/DreamZero-DROID:
- BDE KV-on stays bit-identical to KV-off through a 4-forward rollout (33 frames,
  resident window grows 4 -> 28 blocks): PSNR=inf, SSIM=0.99997, max_abs_diff=0.
- The model's session reset (current_start_frame >= local_attn_size) fires at the
  window boundary, preempting the sliding-window eviction (max_attention_size=36
  frames). So DreamZero's KV is bounded by reset, not by KV eviction; BDE reset
  handling is the next gap for sessions that cross a reset.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../dreamzero/bde_long_rollout.py             | 54 +++++++++++++++++++
 1 file changed, 54 insertions(+)
 create mode 100644 examples/offline_inference/dreamzero/bde_long_rollout.py

diff --git a/examples/offline_inference/dreamzero/bde_long_rollout.py b/examples/offline_inference/dreamzero/bde_long_rollout.py
new file mode 100644
index 00000000000..9a234ee9c82
--- /dev/null
+++ b/examples/offline_inference/dreamzero/bde_long_rollout.py
@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Long DreamZero rollout in one session — to exercise BDE KV eviction.
+
+The stock export submits 2 forwards (prefill + 1 chunk = ~12 frames), which never
+reaches the model's attention window (max_attention_size). This harness submits
+``--num-chunks`` chunk observations in a single session so the cumulative KV grows
+past the window and the model (and BDE) must evict — letting us check that BDE's
+paged eviction stays parity-exact with the model-local path under real eviction.
+
+KV-off vs KV-on is toggled by the BDE_KV_ENABLE env var (same as the export):
+
+    BDE_KV_ENABLE=1 CUDA_VISIBLE_DEVICES=0 HF_HOME=/models \\
+      python examples/offline_inference/dreamzero/bde_long_rollout.py \\
+      --num-chunks 8 --output outputs/bde_parity/long_on.mp4
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import export_prediction_video as E  # noqa: E402
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
+    ap.add_argument("--deploy-config", type=Path, default=Path("vllm_omni/deploy/dreamzero.yaml"))
+    ap.add_argument("--video-dir", type=Path, default=Path("/models/dreamzero-assets"))
+    ap.add_argument("--prompt", default="A robot arm manipulates objects on a table.")
+    ap.add_argument("--session-id", default="long-rollout")
+    ap.add_argument("--num-chunks", type=int, default=8)
+    ap.add_argument("--output", type=Path, default=Path("outputs/bde_parity/long.mp4"))
+    args = ap.parse_args()
+
+    # base = [prefill_obs, chunk_obs]; extend the chunk into a long session.
+    _, base = E._build_observations(args.video_dir, prompt=args.prompt, session_id=args.session_id)
+    observations = [base[0]] + [base[1]] * args.num_chunks
+    print(f"[long-rollout] submitting {len(observations)} forwards in session {args.session_id!r}")
+
+    omni, outputs = E._run_generation(args.model, args.deploy_config, observations)
+    latents = torch.cat([E._extract_latents(o) for o in outputs], dim=2)
+    frames = E._decode_with_worker(omni, latents)
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    E._write_mp4(args.output, frames, fps=5)
+    print(f"SAVED_MP4={args.output}  frames={frames.shape[0]}")
+
+
+if __name__ == "__main__":
+    main()

From 0e7db6219b2cd9bcac7fe6d95c96374724744a20 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Tue, 16 Jun 2026 07:48:09 +0000
Subject: [PATCH 19/51] [BDE] Session-keyed KV state + model-reset handling
 (Phase 1 P0)

DreamZero KV is session-scoped and the model resets its sliding window at the
attention-window boundary (should_reset: current_start_frame >= local_attn_size).
The runner held a single persistent BDE state and ignored resets, so long /
multi-turn sessions broke after the window boundary.

- runner: key BDEKVState by session_id (dict) resolved from the request's
  extra_args, matching pipeline.forward; create/reuse per session.
- gather: BDEKVState.reset() frees resident pool blocks (end_request) and
  restarts both branch adapters (begin_request), clearing committed/pending.
- pipeline: _kv_reset() resets model-local state and the BDE pool together;
  wired into both reset sites (should_reset and explicit extra_args["reset"]).
- harness: bde_long_rollout.py --latents saves decoded latents for exact parity.
- test: test_reset_clears_session_window.

Verified: bde_long_rollout.py --num-chunks 8 (9 forwards, crosses 2 window
resets) runs error-free and is bit-exact vs KV-off (latents max_abs_diff=0,
PSNR=inf). tests/bde/ 70 passed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../dreamzero/bde_long_rollout.py             |  6 ++++
 tests/bde/test_pipeline_bridge.py             | 22 ++++++++++++++
 vllm_omni/bde/kv_cache/gather.py              | 19 ++++++++++++
 vllm_omni/bde/runner.py                       | 29 +++++++++++--------
 .../models/dreamzero/pipeline_dreamzero.py    | 16 ++++++++--
 5 files changed, 78 insertions(+), 14 deletions(-)

diff --git a/examples/offline_inference/dreamzero/bde_long_rollout.py b/examples/offline_inference/dreamzero/bde_long_rollout.py
index 9a234ee9c82..46b823fd901 100644
--- a/examples/offline_inference/dreamzero/bde_long_rollout.py
+++ b/examples/offline_inference/dreamzero/bde_long_rollout.py
@@ -34,6 +34,8 @@ def main() -> None:
     ap.add_argument("--session-id", default="long-rollout")
     ap.add_argument("--num-chunks", type=int, default=8)
     ap.add_argument("--output", type=Path, default=Path("outputs/bde_parity/long.mp4"))
+    ap.add_argument("--latents", type=Path, default=None,
+                    help="if set, torch.save the raw decoded latents here for exact parity checks")
     args = ap.parse_args()
 
     # base = [prefill_obs, chunk_obs]; extend the chunk into a long session.
@@ -43,6 +45,10 @@ def main() -> None:
 
     omni, outputs = E._run_generation(args.model, args.deploy_config, observations)
     latents = torch.cat([E._extract_latents(o) for o in outputs], dim=2)
+    if args.latents is not None:
+        args.latents.parent.mkdir(parents=True, exist_ok=True)
+        torch.save(latents.detach().cpu(), args.latents)
+        print(f"SAVED_LATENTS={args.latents}  shape={tuple(latents.shape)}")
     frames = E._decode_with_worker(omni, latents)
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index 50c3df8eb82..e4326bfef0e 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -63,6 +63,28 @@ def test_branches_are_independent():
     assert st.get_kv_caches(False, lambda: ["X"])[0].shape[2] == BLOCK
 
 
+def test_reset_clears_session_window():
+    """reset() drops the resident window so the next forward starts fresh —
+    mirrors DreamZero's window-boundary state.reset()."""
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    st.update_kv_cache(0, _window(2), False, seq_len=2 * BLOCK)
+    st.update_kv_cache(0, _window(1), True, seq_len=BLOCK)
+    assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK
+    free_before = kv.manager.block_pool.get_num_free_blocks()
+
+    st.reset()
+
+    # Both branches fall back to empty again, and pool blocks are freed.
+    assert st._committed == {False: 0, True: 0}
+    assert st._pending == {False: [], True: []}
+    assert st.get_kv_caches(False, lambda: ["EMPTY"]) == ["EMPTY"]
+    assert st.get_kv_caches(True, lambda: ["EMPTY"]) == ["EMPTY"]
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before
+    # Adapters are live again under the same ids — a post-reset write works.
+    st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
+    assert st.get_kv_caches(False, lambda: ["X"])[0].shape[2] == BLOCK
+
+
 def test_commit_is_noop():
     _, st = make_state()
     st.commit_chunk()  # paged mode keeps resident windows; commit just logs
diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/bde/kv_cache/gather.py
index ad4e52c47e3..c147861a0e0 100644
--- a/vllm_omni/bde/kv_cache/gather.py
+++ b/vllm_omni/bde/kv_cache/gather.py
@@ -199,6 +199,25 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
     def commit_chunk(self) -> None:
         _log.info("BDE COMMIT (paged; resident window retained across forwards)")
 
+    def reset(self) -> None:
+        """Drop this session's KV — mirrors the model-local ``state.reset()``.
+
+        DreamZero resets at the attention-window boundary (``should_reset``); the
+        model starts a fresh sliding window, so BDE must free the resident pool
+        blocks and start a fresh adapter for each branch. The ``BDEKVState``
+        object (and thus ``pipeline._bde_kv_state``) is preserved across the reset
+        so the runner's session mapping stays valid — only the pool-backed state
+        is recycled.
+        """
+        pos_id, neg_id = self.pos.request_id, self.neg.request_id
+        self.kv_cache.end_request(self.pos)        # free resident blocks
+        self.kv_cache.end_request(self.neg)
+        self.pos = self.kv_cache.begin_request(pos_id)   # fresh, empty window
+        self.neg = self.kv_cache.begin_request(neg_id)
+        self._committed = {False: 0, True: 0}
+        self._pending = {False: [], True: []}
+        _log.info("BDE RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)
+
 
 class BDEPipelineMixin:
     """Mixin that replaces model-local KV state with pool-backed storage.
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/bde/runner.py
index 51f68c888fa..b3182bc04d0 100644
--- a/vllm_omni/bde/runner.py
+++ b/vllm_omni/bde/runner.py
@@ -62,9 +62,9 @@ def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None:
         # Built after the model is loaded (dimensions known); stays None while KV
         # management is disabled.
         self.kv_cache: BDEKVCache | None = None
-        # DreamZero KV is session-scoped (persists across forwards), so the BDE KV
-        # state is held on the runner and reused, not created per request.
-        self._bde_state: BDEKVState | None = None
+        # DreamZero KV is session-scoped (persists across a session's forwards),
+        # so BDE KV state is keyed by session_id and reused, not created per request.
+        self._bde_states: dict[str, BDEKVState] = {}
 
     def build_kv_cache(
         self,
@@ -167,17 +167,22 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
             return super().execute_model(req)
 
         kv = self.kv_cache
-        # Reuse the session-scoped BDE KV state so KV persists across the session's
-        # forwards (Approach 1: the state holds the model's window per branch/layer).
-        if self._bde_state is None:
-            pos = kv.begin_request("bde")
-            neg = kv.begin_request("bde__neg")
-            self._bde_state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
+        # DreamZero KV is session-scoped (the model state persists across a
+        # session's forwards), so the BDE KV state is keyed by session_id and
+        # reused — matching how pipeline.forward resolves the model-local state.
+        extra_args = req.sampling_params.extra_args or {}
+        session_id = str(extra_args.get("session_id") or "default")
+        state = self._bde_states.get(session_id)
+        if state is None:
+            pos = kv.begin_request(f"bde__{session_id}")
+            neg = kv.begin_request(f"bde__{session_id}__neg")
+            state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
+            self._bde_states[session_id] = state
         logger.debug(
-            "BDE execute_model: req=%s chunk_size=%d window_chunks=%s num_blocks=%d",
-            req.request_id, kv.spec.chunk_size, kv.config.window_chunks, kv.num_blocks,
+            "BDE execute_model: req=%s session=%s chunk_size=%d window_chunks=%s num_blocks=%d",
+            req.request_id, session_id, kv.spec.chunk_size, kv.config.window_chunks, kv.num_blocks,
         )
-        self.pipeline._bde_kv_state = self._bde_state
+        self.pipeline._bde_kv_state = state
         try:
             return super().execute_model(req)
         finally:
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 5de3c72ac83..8218fcfb90e 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -132,6 +132,18 @@ def _kv_commit(self):
             _log.debug("BDE KV commit: pos=%d chunks neg=%d chunks", s.pos.completed_chunks, s.neg.completed_chunks)
             s.commit_chunk()
 
+    def _kv_reset(self, state):
+        """Reset model-local KV and, under BDE, the pooled session window too.
+
+        DreamZero resets its sliding window at the attention-window boundary; the
+        BDE pool must drop the same window so the next forward starts fresh and
+        stays parity-exact with the model-local path.
+        """
+        state.reset()
+        s = self._bde_kv_state
+        if s is not None:
+            s.reset()
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.
 
@@ -853,10 +865,10 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
         # Explicit reset from OpenPI serving is carried by `extra_args["reset"]`
         # on the next inference request after websocket reset/session switch.
         if extra_args.get("reset", False):
-            state.reset()
+            self._kv_reset(state)
         # Auto-reset based on model state (before accumulation)
         if state.should_reset(text_tokens, 0, self.transformer.local_attn_size):
-            state.reset()
+            self._kv_reset(state)
         state.language = text_tokens
 
         # Frame accumulation: stitched single frame -> multi-frame video

From 440f977db09fa047c9ec8dce9940432e2dffbbe1 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Wed, 17 Jun 2026 11:37:40 +0000
Subject: [PATCH 20/51] [BDE] Move cross-attention KV into the engine pool
 (Phase 1)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

DreamZero's cross-attn KV was model-local (DreamZeroState dicts), bypassing
BDE entirely. Move it into the engine pool so all persistent KV is under one
budget, and fix a pre-existing wiring gap that disarmed the cache.

Part A — transformer wiring gap (benefits the model-local path too):
- _forward_blocks now accepts and forwards crossattn_cache to each block;
  previously it was dropped, so every denoising step recomputed cross-attn
  k/v from context (40 projections × T steps wasted).

Part B — pool-backed cross-attn storage:
- causal_wan_model: project_kv() on both cross-attn classes.
- manager: BDEKVCache allocates per-layer (2, text_len, kv_heads, head_dim)
  cross-attn pools (dim-0 = [pos, neg]); write_cross_kv / read_cross_kv.
- gather: BDEKVState._cross_populated flags, get_cross_kv_caches(), reset
  clears the flags.
- runner: wire cross_attn_length (transformer.text_len) through build_kv_cache.
- pipeline: _kv_get_cross proxy + _kv_populate_cross (eager, once per session
  per branch, re-populated after reset); replace all 6 get_crossattn_caches
  call sites. Cross-attn k/v is computed from the full projected text (image
  tokens are recomputed in-attn, not cached).

Verified: bde_long_rollout.py --num-chunks 8 (9 forwards, crosses 2 resets)
bit-exact vs KV-off (latents max_abs_diff=0, PSNR=inf); cross-attn pool 800 MiB
(40 layers × 512 tok). tests/bde/ 70 passed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 vllm_omni/bde/kv_cache/gather.py              | 11 +++
 vllm_omni/bde/kv_cache/manager.py             | 76 +++++++++++++++++--
 vllm_omni/bde/runner.py                       | 11 ++-
 .../models/dreamzero/causal_wan_model.py      | 29 ++++++-
 .../models/dreamzero/pipeline_dreamzero.py    | 74 +++++++++++++++---
 5 files changed, 184 insertions(+), 17 deletions(-)

diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/bde/kv_cache/gather.py
index c147861a0e0..fb264fb593f 100644
--- a/vllm_omni/bde/kv_cache/gather.py
+++ b/vllm_omni/bde/kv_cache/gather.py
@@ -141,6 +141,10 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
         # the pool per branch; ``_pending`` holds this forward's new-chunk slot maps.
         self._committed: dict[bool, int] = {False: 0, True: 0}
         self._pending: dict[bool, list] = {False: [], True: []}
+        # Cross-attn KV is populated once after text encoding (eager), then
+        # read every denoising step.  Reset clears these flags so the next
+        # forward repopulates from the new text encoding.
+        self._cross_populated: dict[bool, bool] = {False: False, True: False}
 
     def _adapter(self, is_negative: bool):
         return self.neg if is_negative else self.pos
@@ -160,6 +164,12 @@ def get_kv_caches(self, is_negative: bool, fallback) -> list:
         )
         return windows
 
+    def get_cross_kv_caches(self, is_negative: bool, fallback) -> list[dict]:
+        """Return pool-backed cross-attn cache dicts, or fallback if not populated."""
+        if self._cross_populated[is_negative] and self.kv_cache.cross_attn_length > 0:
+            return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
+        return fallback()
+
     def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative: bool, seq_len) -> None:
         branch = "neg" if is_negative else "pos"
         adapter = self._adapter(is_negative)
@@ -216,6 +226,7 @@ def reset(self) -> None:
         self.neg = self.kv_cache.begin_request(neg_id)
         self._committed = {False: 0, True: 0}
         self._pending = {False: [], True: []}
+        self._cross_populated = {False: False, True: False}
         _log.info("BDE RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)
 
 
diff --git a/vllm_omni/bde/kv_cache/manager.py b/vllm_omni/bde/kv_cache/manager.py
index ddc74b341b1..ef7fcb43d52 100644
--- a/vllm_omni/bde/kv_cache/manager.py
+++ b/vllm_omni/bde/kv_cache/manager.py
@@ -45,6 +45,7 @@ def __init__(
         block_size: int,
         max_model_len: int,
         available_bytes: int,
+        cross_attn_length: int = 0,
         device: torch.device | None = None,
     ) -> None:
         if not config.enable:
@@ -56,6 +57,45 @@ def __init__(
 
         self.config = config
         self.block_size = block_size
+        self.num_layers = num_layers
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.dtype = dtype
+        self.cross_attn_length = cross_attn_length
+        self.device = device or torch.device("cpu")
+        self._adapters: dict[str, BDERequestAdapter] = {}
+
+        # -- cross-attention pool (static, one-time-fill) -----------------------
+        # Cross-attn KV is computed once from the text encoder and never changes.
+        # Allocate a per-layer contiguous tensor for each branch: (2, text_len,
+        # kv_heads, head_dim) where dim-0 = [pos, neg].  Allocated separately
+        # from the self-attn block pool (static, small — ~400 MiB for
+        # DreamZero I2V).  Both pools draw from the same GPU free memory but
+        # the cross-attn pool is sized directly rather than via the block pool
+        # budget (no eviction/chunk lifecycle needed).
+        self._cross_k: list[torch.Tensor] = []
+        self._cross_v: list[torch.Tensor] = []
+        if device is not None and cross_attn_length > 0:
+            cross_shape = (2, cross_attn_length, num_kv_heads, head_size)
+            bytes_per_element = dtype.itemsize
+            cross_bytes = (
+                2           # K + V
+                * 2         # pos + neg
+                * cross_attn_length
+                * num_kv_heads
+                * head_size
+                * bytes_per_element
+                * num_layers
+            )
+            for _ in range(num_layers):
+                self._cross_k.append(torch.empty(cross_shape, dtype=dtype, device=device))
+                self._cross_v.append(torch.empty(cross_shape, dtype=dtype, device=device))
+            _log.info(
+                "BDE cross-attn pool: %d layers × (%d tok × %d heads × %d) = %.1f MiB",
+                num_layers, cross_attn_length, num_kv_heads, head_size,
+                cross_bytes / (1024 * 1024),
+            )
+
         self.spec = ChunkWindowSpec(
             block_size=block_size,
             num_kv_heads=num_kv_heads,
@@ -78,14 +118,8 @@ def __init__(
         self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
         self.num_blocks = num_blocks
         self.null_block_id = self.manager.block_pool.null_block.block_id
-        self._adapters: dict[str, BDERequestAdapter] = {}
 
         # Allocate the per-layer paged K/V pools on the given device.
-        self.num_layers = num_layers
-        self.num_kv_heads = num_kv_heads
-        self.head_size = head_size
-        self.dtype = dtype
-        self.device = device or torch.device("cpu")
         self._k_pools: list[torch.Tensor] = []
         self._v_pools: list[torch.Tensor] = []
         if device is not None:
@@ -93,6 +127,36 @@ def __init__(
                 num_blocks, block_size, num_layers, num_kv_heads, head_size, dtype, device
             )
 
+    # -- cross-attention pool access -------------------------------------------
+    # Cross-attn KV is static once populated — write once (from text encoder),
+    # read many (every denoising step). Not managed through the paged block pool.
+
+    def write_cross_kv(
+        self, layer_idx: int, is_negative: bool,
+        k: torch.Tensor, v: torch.Tensor,
+    ) -> None:
+        """Write one layer's cross-attn K/V into the pool.
+
+        ``k`` / ``v``: ``(B, text_len, tp_num_heads, head_dim)``; only batch-0
+        is copied (B=1 for inference). The ``is_negative`` flag selects the
+        correct CFG branch slot.
+        """
+        branch = 1 if is_negative else 0
+        self._cross_k[layer_idx][branch].copy_(k[0])
+        self._cross_v[layer_idx][branch].copy_(v[0])
+
+    def read_cross_kv(self, layer_idx: int, is_negative: bool) -> dict:
+        """Return a pool-backed cross-attn cache dict for one layer.
+
+        The dict matches the ``{"is_init": True, "k": Tensor, "v": Tensor}``
+        convention the cross-attention module expects — it reads from the pool
+        slice rather than from the lazy-initialised model-local dict.
+        """
+        branch = 1 if is_negative else 0
+        k = self._cross_k[layer_idx][branch].unsqueeze(0)   # (1, L, heads, dim)
+        v = self._cross_v[layer_idx][branch].unsqueeze(0)
+        return {"is_init": True, "k": k, "v": v}
+
     # -- request lifecycle ---------------------------------------------------
 
     def begin_request(
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/bde/runner.py
index b3182bc04d0..7a7ffc39e4e 100644
--- a/vllm_omni/bde/runner.py
+++ b/vllm_omni/bde/runner.py
@@ -76,6 +76,7 @@ def build_kv_cache(
         block_size: int,
         max_model_len: int,
         available_bytes: int,
+        cross_attn_length: int = 0,
     ) -> None:
         """Construct the BDEKVCache (preallocating GPU pools on ``self.device``).
 
@@ -92,11 +93,12 @@ def build_kv_cache(
             block_size=block_size,
             max_model_len=max_model_len,
             available_bytes=available_bytes,
+            cross_attn_length=cross_attn_length,
             device=self.device,
         )
         logger.info(
             "BDE KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s, "
-            "layers=%d kv_heads=%d head_dim=%d block_size=%d device=%s",
+            "layers=%d kv_heads=%d head_dim=%d block_size=%d cross_attn_len=%d device=%s",
             self.kv_cache.num_blocks,
             self.bde_kv_config.chunk_size,
             self.bde_kv_config.window_chunks,
@@ -104,6 +106,7 @@ def build_kv_cache(
             num_kv_heads,
             head_size,
             block_size,
+            cross_attn_length,
             self.device,
         )
 
@@ -134,6 +137,11 @@ def _preallocate_kv_cache(self) -> None:
         # The model's own attention window, in tokens (= the [-max_attention_size:]
         # slice it applies). Read it directly rather than recomputing.
         max_attention_size = int(t.blocks[0].self_attn.max_attention_size)
+        # Cross-attn text length — static pool for the text-encoder output.
+        # The cached k/v spans the full text sequence (the image tokens
+        # prepended in _forward_blocks are stripped inside the cross-attn
+        # forward but are not part of the cached k/v).
+        cross_attn_length = int(getattr(t, "text_len", 0))
 
         # Frame-granular paging: 1 block = 1 frame = frame_seqlen tokens, so the
         # resident window matches max_attention_size exactly (it need not be a whole
@@ -159,6 +167,7 @@ def _preallocate_kv_cache(self) -> None:
             block_size=chunk_size,  # one pool block per chunk
             max_model_len=1 << 20,
             available_bytes=free_bytes,
+            cross_attn_length=cross_attn_length,
         )
 
     def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
index e187054e2f6..a6fbf1d004d 100644
--- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
+++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
@@ -290,6 +290,18 @@ def forward(
         x = self.o(x)
         return x
 
+    def project_kv(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Project encoder hidden states to cross-attn K/V (for eager caching).
+
+        Returns ``(k, v)`` where each has shape
+        ``(B, text_len, tp_num_heads, head_dim)`` — the same tensors the
+        ``forward`` method would compute on first call.
+        """
+        n, d = self.tp_num_heads, self.head_dim
+        k = self.norm_k(self.k(context)).unflatten(2, (n, d))
+        v = self.v(context).unflatten(2, (n, d))
+        return k, v
+
 
 class WanI2VCrossAttention(nn.Module):
     """Image-to-video cross-attention (splits first 257 image tokens).
@@ -360,6 +372,18 @@ def forward(
         x = self.o(x)
         return x
 
+    def project_kv(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Project encoder hidden states to cross-attn K/V for the text part only.
+
+        Image tokens (first 257) are *not* cached — they are always recomputed.
+        Returns ``(k, v)`` for the text tokens.
+        """
+        context = context[:, 257:]  # strip image tokens; cache text only
+        n, d = self.tp_num_heads, self.head_dim
+        k = self.norm_k(self.k(context)).unflatten(2, (n, d))
+        v = self.v(context).unflatten(2, (n, d))
+        return k, v
+
 
 WAN_CROSSATTENTION_CLASSES = {
     "t2v_cross_attn": WanT2VCrossAttention,
@@ -835,6 +859,7 @@ def _forward_blocks(
         state: torch.Tensor | None,
         kv_cache: list[torch.Tensor],
         current_start_frame: int,
+        crossattn_cache: list[dict] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
         x = x.flatten(start_dim=2).transpose(1, 2)
         B = x.shape[0]
@@ -884,6 +909,7 @@ def _forward_blocks(
                 context=context,
                 action_register_length=action_register_length,
                 kv_cache=kv_cache[block_index] if kv_cache else None,
+                crossattn_cache=crossattn_cache[block_index] if crossattn_cache else None,
                 current_start_frame=current_start_frame,
             )
             updated_kv_caches.append(updated_kv_cache)
@@ -907,8 +933,8 @@ def _forward_inference(
         context: torch.Tensor,
         seq_len: int,
         kv_cache: list[torch.Tensor],
-        crossattn_cache: list[torch.Tensor],
         current_start_frame: int,
+        crossattn_cache: list[dict] | None = None,
         y: torch.Tensor | None = None,
         clip_feature: torch.Tensor | None = None,
         action: torch.Tensor | None = None,
@@ -941,6 +967,7 @@ def _forward_inference(
             timestep_action=timestep_action,
             state=state,
             kv_cache=kv_cache,
+            crossattn_cache=crossattn_cache,
             current_start_frame=current_start_frame,
         )
 
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 8218fcfb90e..09e9ced8568 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -113,9 +113,11 @@ def _kv_get(self, state, is_negative):
         return state.get_kv_caches(is_negative)
 
     def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
-        # Always create the model-local caches so cross-attention (which BDE does
-        # NOT manage) is initialized. Under BDE the self-attn kv_cache created here
-        # is unused — _kv_get / _kv_update route self-attn KV through the pool.
+        # Always create the model-local caches for the non-BDE fallback path.
+        # Under BDE both self-attn and cross-attn KV are pool-backed — the
+        # state caches created here are unused when _bde_kv_state is set.
+        # Cross-attn is populated eagerly in _kv_populate_cross after text
+        # encoding; self-attn is routed through _kv_get / _kv_update.
         state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
 
     def _kv_update(self, state, layer_idx, updated_kv, is_negative, seq_len=None):
@@ -132,6 +134,52 @@ def _kv_commit(self):
             _log.debug("BDE KV commit: pos=%d chunks neg=%d chunks", s.pos.completed_chunks, s.neg.completed_chunks)
             s.commit_chunk()
 
+    def _kv_get_cross(self, state, is_negative):
+        """Cross-attn cache: pool-backed when BDE active, model-local otherwise."""
+        s = self._bde_kv_state
+        if s is not None:
+            return s.get_cross_kv_caches(
+                is_negative, lambda: state.get_crossattn_caches(is_negative)
+            )
+        return state.get_crossattn_caches(is_negative)
+
+    def _kv_populate_cross(self, context: torch.Tensor, is_negative: bool) -> None:
+        """Eagerly project cross-attn K/V for all layers into the BDE pool.
+
+        Called once per branch after text encoding (before the first denoising
+        step), and again after each window-boundary reset.  The attention
+        forward never writes to the cross-attn pool — it only reads.
+
+        The raw encoder output is projected through ``text_embedding`` first
+        so it matches the shape the cross-attn linear layers expect (the same
+        projection ``_forward_blocks`` applies internally).
+        """
+        s = self._bde_kv_state
+        if s is None:
+            return
+        # Only populate once per branch per session — reset clears this flag.
+        if s._cross_populated.get(is_negative, False):
+            return
+        projected = self.transformer.text_embedding(context)
+        for i, block in enumerate(self.transformer.blocks):
+            # The cross-attn forward caches text-only k/v — the image tokens
+            # (first 257, prepended inside _forward_blocks via img_emb) are
+            # always recomputed.  project_kv for I2V strips those image tokens,
+            # but we are feeding raw projected text (no image prepend), so
+            # compute k/v directly from the full projected text.
+            ca = block.cross_attn
+            n, d = ca.tp_num_heads, ca.head_dim
+            k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
+            v = ca.v(projected).unflatten(2, (n, d))
+            s.kv_cache.write_cross_kv(i, is_negative, k, v)
+        s._cross_populated[is_negative] = True
+        _log.info(
+            "BDE CROSS POPULATE [%s]: %d layers, context=%s",
+            "neg" if is_negative else "pos",
+            len(self.transformer.blocks),
+            tuple(context.shape),
+        )
+
     def _kv_reset(self, state):
         """Reset model-local KV and, under BDE, the pooled session window too.
 
@@ -579,7 +627,7 @@ def _prefill_kv_cache(
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
                 kv_cache=self._kv_get(state, False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
@@ -587,7 +635,7 @@ def _prefill_kv_cache(
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
                     kv_cache=self._kv_get(state, True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -629,7 +677,7 @@ def _prefill_kv_cache(
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
                 kv_cache=self._kv_get(state, False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
@@ -637,7 +685,7 @@ def _prefill_kv_cache(
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
                     kv_cache=self._kv_get(state, True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -723,7 +771,7 @@ def diffuse(
                 timestep_video=timestep,
                 encoder_hidden_states=prompt_embeds,
                 kv_cache=self._kv_get(state, False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                crossattn_cache=self._kv_get_cross(state, False),
                 y=y,
                 clip_feature=state.clip_feas,
                 action=noisy_input_action,
@@ -739,7 +787,7 @@ def diffuse(
                     timestep_video=timestep,
                     encoder_hidden_states=negative_prompt_embeds,
                     kv_cache=self._kv_get(state, True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     y=y,
                     clip_feature=state.clip_feas,
                     action=noisy_input_action,
@@ -895,6 +943,14 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
                 neg_inputs["attention_mask"].to(device),
             )
 
+        # -- Eager cross-attn population (BDE only) ----------------------------
+        # Project encoder hidden states to K/V once and store in the BDE pool
+        # so every denoising step reads without recomputing.  Re-populated after
+        # each window-boundary reset (_cross_populated clears → guard re-triggers).
+        self._kv_populate_cross(prompt_embeds, is_negative=False)
+        if negative_prompt_embeds is not None:
+            self._kv_populate_cross(negative_prompt_embeds, is_negative=True)
+
         # Extract first/last frame for CLIP + VAE encoding
         if num_frames_raw == 4 or num_frames_raw == 9:
             image = videos[:, :, -1:].transpose(1, 2)

From 9f1f9fe4d331be76cc36d65858a19b69c374d4c3 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Thu, 18 Jun 2026 08:40:36 +0000
Subject: [PATCH 21/51] [BDE] Add BDE-vs-baseline perf/precision harness for
 DreamZero
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Times each forward (per-chunk E2E), saves latents + decoded video + a timing
JSON. Works on both the BDE branch (BDE_KV_ENABLE=1) and a clean baseline
checkout — uses only export_prediction_video helpers, no BDE imports.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../dreamzero/bde_perf_compare.py             | 111 ++++++++++++++++++
 1 file changed, 111 insertions(+)
 create mode 100644 examples/offline_inference/dreamzero/bde_perf_compare.py

diff --git a/examples/offline_inference/dreamzero/bde_perf_compare.py b/examples/offline_inference/dreamzero/bde_perf_compare.py
new file mode 100644
index 00000000000..f307511fba5
--- /dev/null
+++ b/examples/offline_inference/dreamzero/bde_perf_compare.py
@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Timed DreamZero rollout for BDE-vs-main comparison.
+
+Runs one session of (prefill + N chunk) forwards, timing each `omni.generate`
+(the per-forward E2E), saving the decoded latents for a precision check, and
+dumping a timing JSON. Works on both the BDE branch (with BDE_KV_ENABLE=1) and a
+clean main checkout (baseline) — it only uses helpers from export_prediction_video,
+which exist on both.
+
+    BDE_KV_ENABLE=1 CUDA_VISIBLE_DEVICES=0 HF_HOME=/models \\
+      python examples/offline_inference/dreamzero/bde_perf_compare.py \\
+      --num-chunks 12 --tag bde \\
+      --latents outputs/bde_parity/perf_bde.pt --timing outputs/bde_parity/perf_bde.json
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import export_prediction_video as E  # noqa: E402
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
+    ap.add_argument("--deploy-config", type=Path, default=Path("vllm_omni/deploy/dreamzero.yaml"))
+    ap.add_argument("--video-dir", type=Path, default=Path("/models/dreamzero-assets"))
+    ap.add_argument("--prompt", default="A robot arm manipulates objects on a table.")
+    ap.add_argument("--session-id", default="perf-cmp")
+    ap.add_argument("--num-chunks", type=int, default=12)
+    ap.add_argument("--tag", default="run")
+    ap.add_argument("--latents", type=Path, default=None)
+    ap.add_argument("--video", type=Path, default=None)
+    ap.add_argument("--fps", type=int, default=5)
+    ap.add_argument("--timing", type=Path, default=None)
+    args = ap.parse_args()
+
+    _, base = E._build_observations(args.video_dir, prompt=args.prompt, session_id=args.session_id)
+    observations = [base[0]] + [base[1]] * args.num_chunks
+    n = len(observations)
+    print(f"[perf:{args.tag}] {n} forwards (1 prefill + {args.num_chunks} chunk) "
+          f"BDE_KV_ENABLE={os.environ.get('BDE_KV_ENABLE', '0')}")
+
+    t0 = time.perf_counter()
+    omni = E.Omni(
+        model=args.model,
+        deploy_config=str(args.deploy_config),
+        enforce_eager=True,
+        worker_extension_cls=E.WORKER_EXTENSION,
+    )
+    load_s = time.perf_counter() - t0
+    print(f"[perf:{args.tag}] engine load: {load_s:.2f}s")
+
+    outputs = []
+    per_forward = []
+    for index, obs in enumerate(observations):
+        sp = E.OmniDiffusionSamplingParams(
+            extra_args={"reset": index == 0, "session_id": obs["session_id"], "robot_obs": obs}
+        )
+        s = time.perf_counter()
+        result = omni.generate(obs["prompt"], sampling_params_list=[sp])  # blocking -> E2E
+        dt = time.perf_counter() - s
+        if not result:
+            raise RuntimeError(f"No output for forward {index}")
+        outputs.append(result[0])
+        per_forward.append(dt)
+        print(f"[perf:{args.tag}] forward {index:2d} ({'prefill' if index == 0 else 'chunk'}): {dt*1000:8.1f} ms")
+
+    prefill = per_forward[0]
+    steady = per_forward[2:] if len(per_forward) > 2 else per_forward[1:]  # drop prefill + 1 warmup
+    mean_steady = sum(steady) / len(steady)
+    summary = {
+        "tag": args.tag,
+        "bde_kv_enable": os.environ.get("BDE_KV_ENABLE", "0"),
+        "num_forwards": n,
+        "load_s": load_s,
+        "prefill_s": prefill,
+        "steady_mean_s": mean_steady,
+        "steady_min_s": min(steady),
+        "steady_max_s": max(steady),
+        "total_gen_s": sum(per_forward),
+        "per_forward_s": per_forward,
+    }
+    print(f"[perf:{args.tag}] prefill={prefill*1000:.1f}ms  "
+          f"steady_mean={mean_steady*1000:.1f}ms (n={len(steady)})  total={sum(per_forward):.2f}s")
+
+    # Decode once (after timing) so latents + the final video are both saved.
+    latents = torch.cat([E._extract_latents(o) for o in outputs], dim=2)
+    if args.latents is not None:
+        args.latents.parent.mkdir(parents=True, exist_ok=True)
+        torch.save(latents.detach().cpu(), args.latents)
+        print(f"[perf:{args.tag}] SAVED_LATENTS={args.latents} shape={tuple(latents.shape)}")
+    if args.video is not None:
+        frames = E._decode_with_worker(omni, latents)
+        args.video.parent.mkdir(parents=True, exist_ok=True)
+        E._write_mp4(args.video, frames, fps=args.fps)
+        print(f"[perf:{args.tag}] SAVED_MP4={args.video} frames={frames.shape[0]}")
+    if args.timing is not None:
+        args.timing.parent.mkdir(parents=True, exist_ok=True)
+        args.timing.write_text(json.dumps(summary, indent=2))
+        print(f"[perf:{args.tag}] SAVED_TIMING={args.timing}")
+
+
+if __name__ == "__main__":
+    main()

From 609736b9e972e4f2e39fa3aced90b79974396f2e Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Thu, 18 Jun 2026 09:59:13 +0000
Subject: [PATCH 22/51] [BDE] Memoize per-forward window gather (remove
 per-step redundancy)

The resident self-attn window is frozen across a forward's denoise steps (KV is
only written back after the last step), but _kv_get gathered it once per step
(16x/forward/branch), each time rebuilding the slot index with a per-layer Python
loop. This was the entire BDE-vs-baseline overhead (+15% / sawtooth scaling with
window size).

- BDEKVState memoizes the gathered window per branch (_gather_cache); invalidated
  on write-back (update_kv_cache layer 0) and on reset.
- gather_window_all_layers builds window_block_ids + a vectorized slot index ONCE
  per forward and shares it across all 40 layers (was rebuilt per layer).
- build_window_slots replaces the per-block Python arange/cat loop.

Result (H200, TP1/CFG1, eager, mean of 5 runs): gather calls 434 -> 26 (16.7x);
per-chunk E2E 5683 -> 4813 ms, now ~parity with the model-local baseline (4930 ms).
Bit-exact parity preserved (latents max_abs_diff=0, incl. across resets). Peak
memory unchanged (the cached window-set was already materialized every step).
tests/bde/ 72 passing (+2 memoization tests).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_pipeline_bridge.py | 28 ++++++++++++++++++
 vllm_omni/bde/kv_cache/gather.py  | 47 +++++++++++++++++++++++--------
 vllm_omni/bde/kv_cache/manager.py | 30 +++++++++++++++++++-
 3 files changed, 93 insertions(+), 12 deletions(-)

diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index e4326bfef0e..b321c61e651 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -47,6 +47,34 @@ def test_paged_write_then_gather_roundtrips():
     assert torch.allclose(got[1, 0], win[1, 0], atol=0)
 
 
+def test_gather_is_memoized_within_a_forward():
+    """The resident window is frozen across a forward's denoise steps, so repeated
+    get_kv_caches must reuse one gather; a write-back (layer 0) invalidates it."""
+    _, st = make_state(num_layers=2)
+    st.update_kv_cache(0, _window(2), False, seq_len=2 * BLOCK)  # commit one chunk
+    first = st.get_kv_caches(False, lambda: ["X"])
+    again = st.get_kv_caches(False, lambda: ["X"])
+    assert again is first  # cache HIT: identical object, no re-gather
+    # A new write-back grows the window -> next get must re-gather a fresh object.
+    st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
+    after = st.get_kv_caches(False, lambda: ["X"])
+    assert after is not first
+    # Negative branch is cached independently.
+    st.update_kv_cache(0, _window(1), True, seq_len=BLOCK)
+    neg1 = st.get_kv_caches(True, lambda: ["X"])
+    assert st.get_kv_caches(True, lambda: ["X"]) is neg1
+    assert neg1 is not after
+
+
+def test_reset_clears_gather_cache():
+    _, st = make_state(num_layers=1)
+    st.update_kv_cache(0, _window(2), False, seq_len=2 * BLOCK)
+    st.get_kv_caches(False, lambda: ["X"])
+    assert st._gather_cache[False] is not None
+    st.reset()
+    assert st._gather_cache == {False: None, True: None}
+
+
 def test_eviction_bounds_resident_window():
     kv, st = make_state(num_layers=1, window_chunks=3)
     st.update_kv_cache(0, _window(3), False, seq_len=3 * BLOCK)          # resident 3/3
diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/bde/kv_cache/gather.py
index fb264fb593f..e89f953a7f4 100644
--- a/vllm_omni/bde/kv_cache/gather.py
+++ b/vllm_omni/bde/kv_cache/gather.py
@@ -69,29 +69,41 @@ def pool_write_chunk(
     v_pool[slot_mapping] = new_v[0]
 
 
+def build_window_slots(
+    window_block_ids: list[int], block_size: int, device: torch.device
+) -> torch.Tensor:
+    """Flat slot index covering all resident blocks, in table order.
+
+    Vectorized (no Python loop): ``slot = block_id * block_size + offset`` for every
+    ``offset in [0, block_size)``. The window's blocks are identical across all layers
+    within a forward, so this is built once per forward and shared.
+    """
+    if not window_block_ids:
+        raise ValueError("window_block_ids is empty — no blocks are resident")
+    starts = torch.tensor(window_block_ids, dtype=torch.long, device=device) * block_size
+    offsets = torch.arange(block_size, device=device)
+    return (starts[:, None] + offsets[None, :]).reshape(-1)
+
+
 def pool_gather_window(
     k_pool: torch.Tensor,
     v_pool: torch.Tensor,
     window_block_ids: list[int],
     block_size: int,
     max_attention_size: int,
+    *,
+    slots: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Materialize the resident window K/V as a contiguous tensor.
 
     Gathers all slots belonging to ``window_block_ids`` (in table order), trims to
     the last ``max_attention_size`` tokens (matching the existing ``cat`` + slice),
     and returns a ``(2, batch, window_tokens, n_heads, head_dim)`` tensor — exactly
-    the format DreamZero's attention expects as ``kv_cache``.
+    the format DreamZero's attention expects as ``kv_cache``. Pass a precomputed
+    ``slots`` (from :func:`build_window_slots`) to share the index across layers.
     """
-    if not window_block_ids:
-        raise ValueError("window_block_ids is empty — no blocks are resident")
-    # Build flat slot range covering all resident blocks in table order.
-    starts = torch.tensor(
-        [b * block_size for b in window_block_ids],
-        dtype=torch.long,
-        device=k_pool.device,
-    )
-    slots = torch.cat([torch.arange(s, s + block_size, device=k_pool.device) for s in starts])
+    if slots is None:
+        slots = build_window_slots(window_block_ids, block_size, k_pool.device)
     # Gather then trim to the attention window.
     wk = k_pool[slots]  # (window_slots, n_heads, head_dim)
     wv = v_pool[slots]
@@ -141,6 +153,11 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
         # the pool per branch; ``_pending`` holds this forward's new-chunk slot maps.
         self._committed: dict[bool, int] = {False: 0, True: 0}
         self._pending: dict[bool, list] = {False: [], True: []}
+        # The resident window is frozen across a forward's denoise steps (KV is
+        # only written back after the last step), so the per-step gather is
+        # re-materializing the same tensors ~16x. Memoize the gathered window per
+        # branch; invalidate on write-back (window grows) and on reset.
+        self._gather_cache: dict[bool, list | None] = {False: None, True: None}
         # Cross-attn KV is populated once after text encoding (eager), then
         # read every denoising step.  Reset clears these flags so the next
         # forward repopulates from the new text encoding.
@@ -155,8 +172,12 @@ def get_kv_caches(self, is_negative: bool, fallback) -> list:
             out = fallback()
             _log.info("BDE GET   [%s] source=model-local-empty layers=%d", branch, len(out))
             return out
+        cached = self._gather_cache[is_negative]
+        if cached is not None:  # window unchanged this forward -> reuse (no re-gather)
+            return cached
         adapter = self._adapter(is_negative)
-        windows = [self.kv_cache.gather_window(i, adapter) for i in range(self.num_layers)]
+        windows = self.kv_cache.gather_window_all_layers(adapter)
+        self._gather_cache[is_negative] = windows
         _log.info(
             "BDE GET   [%s] source=paged-gather layers=%d window0=%s resident_blocks=%d/%d",
             branch, self.num_layers, tuple(windows[0].shape),
@@ -187,6 +208,9 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
                 adapter.on_chunk_committed()
             self._pending[is_negative] = slots
             self._committed[is_negative] += new_count
+            # The window just grew; the memoized gather is stale -> force re-gather
+            # on the next forward's first read of this branch.
+            self._gather_cache[is_negative] = None
             _log.info(
                 "BDE WRITE [%s] new_tokens=%d -> %d frame-chunks  resident=%d/%d  storing %d layers",
                 branch, new_count, n_chunks,
@@ -226,6 +250,7 @@ def reset(self) -> None:
         self.neg = self.kv_cache.begin_request(neg_id)
         self._committed = {False: 0, True: 0}
         self._pending = {False: [], True: []}
+        self._gather_cache = {False: None, True: None}
         self._cross_populated = {False: False, True: False}
         _log.info("BDE RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)
 
diff --git a/vllm_omni/bde/kv_cache/manager.py b/vllm_omni/bde/kv_cache/manager.py
index ef7fcb43d52..5c0163237df 100644
--- a/vllm_omni/bde/kv_cache/manager.py
+++ b/vllm_omni/bde/kv_cache/manager.py
@@ -21,7 +21,12 @@
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
 from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
 from vllm_omni.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
-from vllm_omni.bde.kv_cache.gather import allocate_kv_pool, pool_gather_window, pool_write_chunk
+from vllm_omni.bde.kv_cache.gather import (
+    allocate_kv_pool,
+    build_window_slots,
+    pool_gather_window,
+    pool_write_chunk,
+)
 
 _log = init_logger(__name__)
 
@@ -257,3 +262,26 @@ def gather_window(self, layer_index: int, adapter: BDERequestAdapter) -> torch.T
                     adapter.request_id, layer_index, window_ids,
                     tuple(window.shape), window.device)
         return window
+
+    def gather_window_all_layers(self, adapter: BDERequestAdapter) -> list[torch.Tensor]:
+        """Gather the resident-window K/V for every layer in one shot.
+
+        The window's block table is identical across layers within a forward, so
+        the block-id lookup and the flat slot index are computed **once** and shared
+        across all per-layer gathers — removing the per-layer Python work that the
+        single-layer ``gather_window`` repeats. Returns one
+        ``(2, 1, window, n_heads, head_dim)`` tensor per layer.
+        """
+        window_ids = self.window_block_ids(adapter)
+        slots = build_window_slots(window_ids, self.block_size, self._k_pools[0].device)
+        windows = [
+            pool_gather_window(
+                self._k_pools[i], self._v_pools[i], window_ids,
+                self.block_size, self.spec.sliding_window, slots=slots,
+            )
+            for i in range(self.num_layers)
+        ]
+        _log.debug("BDE gather-all: req=%s layers=%d blocks=%d window=%s",
+                    adapter.request_id, self.num_layers, len(window_ids),
+                    tuple(windows[0].shape))
+        return windows

From ee320939599127dfa66fa301337032d0d1834d39 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Thu, 18 Jun 2026 10:50:01 +0000
Subject: [PATCH 23/51] [BDE] Apply pre-commit lint/format fixes

ruff-check (E402 _log placement in pipeline_dreamzero, E731 lambda->def in
test_kv_cache, TID251 torch.cuda.device_count->torch.accelerator in
test_full_chain_gpu) + ruff-format across the BDE files. No behavior change;
tests/bde/ 72 passing.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 benchmarks/bde/metrics.py                     |  4 +-
 benchmarks/bde/parity_check.py                | 19 ++--
 .../dreamzero/bde_long_rollout.py             |  8 +-
 .../dreamzero/bde_perf_compare.py             | 14 +--
 tests/bde/test_engine_routing.py              |  1 -
 tests/bde/test_full_chain_gpu.py              | 23 +++--
 tests/bde/test_gather.py                      | 13 ++-
 tests/bde/test_kv_cache.py                    | 19 ++--
 tests/bde/test_pipeline_bridge.py             | 21 +++--
 tests/bde/test_slot_mapping.py                |  4 +-
 tests/bde/test_validation_harness.py          |  4 +-
 vllm_omni/bde/kv_cache/gather.py              | 36 ++++----
 vllm_omni/bde/kv_cache/manager.py             | 87 ++++++++++++-------
 vllm_omni/bde/runner.py                       | 20 +++--
 vllm_omni/diffusion/diffusion_engine.py       |  3 +-
 .../models/dreamzero/pipeline_dreamzero.py    | 17 ++--
 16 files changed, 173 insertions(+), 120 deletions(-)

diff --git a/benchmarks/bde/metrics.py b/benchmarks/bde/metrics.py
index b909261bd9a..2fe6a85e95c 100644
--- a/benchmarks/bde/metrics.py
+++ b/benchmarks/bde/metrics.py
@@ -89,9 +89,7 @@ def tensor_parity(ref, test, *, atol: float = 1e-4, rtol: float = 1e-4) -> dict:
     }
 
 
-def check_gates(
-    metrics: dict, *, psnr_min: float = 40.0, ssim_min: float = 0.99, lpips_max: float = 0.01
-) -> dict:
+def check_gates(metrics: dict, *, psnr_min: float = 40.0, ssim_min: float = 0.99, lpips_max: float = 0.01) -> dict:
     """Apply the §5 acceptance gates to a ``frame_metrics`` result."""
 
     def gate(value, threshold, ok):
diff --git a/benchmarks/bde/parity_check.py b/benchmarks/bde/parity_check.py
index 49a53821856..7ee8280f838 100644
--- a/benchmarks/bde/parity_check.py
+++ b/benchmarks/bde/parity_check.py
@@ -65,12 +65,19 @@ def run_export(model: str, deploy_config, session_id: str, out_stem, *, extra=()
     out_stem = Path(out_stem)
     out_stem.parent.mkdir(parents=True, exist_ok=True)
     cmd = [
-        sys.executable, EXPORT_SCRIPT,
-        "--model", model,
-        "--deploy-config", str(deploy_config),
-        "--session-id", session_id,
-        "--output-stem", str(out_stem),
-        "--save-gif", "--save-input-video", "--save-actions",
+        sys.executable,
+        EXPORT_SCRIPT,
+        "--model",
+        model,
+        "--deploy-config",
+        str(deploy_config),
+        "--session-id",
+        session_id,
+        "--output-stem",
+        str(out_stem),
+        "--save-gif",
+        "--save-input-video",
+        "--save-actions",
         *extra,
     ]
     subprocess.run(cmd, check=True)
diff --git a/examples/offline_inference/dreamzero/bde_long_rollout.py b/examples/offline_inference/dreamzero/bde_long_rollout.py
index 46b823fd901..2d81345efc3 100644
--- a/examples/offline_inference/dreamzero/bde_long_rollout.py
+++ b/examples/offline_inference/dreamzero/bde_long_rollout.py
@@ -34,8 +34,12 @@ def main() -> None:
     ap.add_argument("--session-id", default="long-rollout")
     ap.add_argument("--num-chunks", type=int, default=8)
     ap.add_argument("--output", type=Path, default=Path("outputs/bde_parity/long.mp4"))
-    ap.add_argument("--latents", type=Path, default=None,
-                    help="if set, torch.save the raw decoded latents here for exact parity checks")
+    ap.add_argument(
+        "--latents",
+        type=Path,
+        default=None,
+        help="if set, torch.save the raw decoded latents here for exact parity checks",
+    )
     args = ap.parse_args()
 
     # base = [prefill_obs, chunk_obs]; extend the chunk into a long session.
diff --git a/examples/offline_inference/dreamzero/bde_perf_compare.py b/examples/offline_inference/dreamzero/bde_perf_compare.py
index f307511fba5..274458d8039 100644
--- a/examples/offline_inference/dreamzero/bde_perf_compare.py
+++ b/examples/offline_inference/dreamzero/bde_perf_compare.py
@@ -44,8 +44,10 @@ def main() -> None:
     _, base = E._build_observations(args.video_dir, prompt=args.prompt, session_id=args.session_id)
     observations = [base[0]] + [base[1]] * args.num_chunks
     n = len(observations)
-    print(f"[perf:{args.tag}] {n} forwards (1 prefill + {args.num_chunks} chunk) "
-          f"BDE_KV_ENABLE={os.environ.get('BDE_KV_ENABLE', '0')}")
+    print(
+        f"[perf:{args.tag}] {n} forwards (1 prefill + {args.num_chunks} chunk) "
+        f"BDE_KV_ENABLE={os.environ.get('BDE_KV_ENABLE', '0')}"
+    )
 
     t0 = time.perf_counter()
     omni = E.Omni(
@@ -70,7 +72,7 @@ def main() -> None:
             raise RuntimeError(f"No output for forward {index}")
         outputs.append(result[0])
         per_forward.append(dt)
-        print(f"[perf:{args.tag}] forward {index:2d} ({'prefill' if index == 0 else 'chunk'}): {dt*1000:8.1f} ms")
+        print(f"[perf:{args.tag}] forward {index:2d} ({'prefill' if index == 0 else 'chunk'}): {dt * 1000:8.1f} ms")
 
     prefill = per_forward[0]
     steady = per_forward[2:] if len(per_forward) > 2 else per_forward[1:]  # drop prefill + 1 warmup
@@ -87,8 +89,10 @@ def main() -> None:
         "total_gen_s": sum(per_forward),
         "per_forward_s": per_forward,
     }
-    print(f"[perf:{args.tag}] prefill={prefill*1000:.1f}ms  "
-          f"steady_mean={mean_steady*1000:.1f}ms (n={len(steady)})  total={sum(per_forward):.2f}s")
+    print(
+        f"[perf:{args.tag}] prefill={prefill * 1000:.1f}ms  "
+        f"steady_mean={mean_steady * 1000:.1f}ms (n={len(steady)})  total={sum(per_forward):.2f}s"
+    )
 
     # Decode once (after timing) so latents + the final video are both saved.
     latents = torch.cat([E._extract_latents(o) for o in outputs], dim=2)
diff --git a/tests/bde/test_engine_routing.py b/tests/bde/test_engine_routing.py
index b5550e3ee39..bb301d3920f 100644
--- a/tests/bde/test_engine_routing.py
+++ b/tests/bde/test_engine_routing.py
@@ -11,7 +11,6 @@
 from types import SimpleNamespace
 
 import pytest
-
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
 from vllm_omni.bde.engine import (
diff --git a/tests/bde/test_full_chain_gpu.py b/tests/bde/test_full_chain_gpu.py
index 5bba347ec56..8081d7b49d4 100644
--- a/tests/bde/test_full_chain_gpu.py
+++ b/tests/bde/test_full_chain_gpu.py
@@ -16,7 +16,7 @@
 import torch
 
 requires_gpus_2_3 = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.device_count() < 2,
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 2,
     reason="requires >= 2 GPUs",
 )
 
@@ -30,9 +30,7 @@ def _worker_entry(rank: int, od_config, q) -> None:
     try:
         from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
-        worker = DiffusionWorker(
-            local_rank=rank, rank=rank, od_config=od_config, skip_load_model=True
-        )
+        worker = DiffusionWorker(local_rank=rank, rank=rank, od_config=od_config, skip_load_model=True)
         device = str(worker.device)
         runner_name = type(worker.model_runner).__name__
 
@@ -41,9 +39,15 @@ def _worker_entry(rank: int, od_config, q) -> None:
 
         cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
         kv = BDEKVCache(
-            cfg, num_layers=1, num_kv_heads=N_HEADS, head_size=HEAD_DIM,
-            dtype=torch.float32, block_size=BLOCK, max_model_len=512,
-            available_bytes=1 << 28, device=worker.device,
+            cfg,
+            num_layers=1,
+            num_kv_heads=N_HEADS,
+            head_size=HEAD_DIM,
+            dtype=torch.float32,
+            block_size=BLOCK,
+            max_model_len=512,
+            available_bytes=1 << 28,
+            device=worker.device,
         )
 
         # Verify pools are on the worker's GPU.
@@ -64,8 +68,8 @@ def _worker_entry(rank: int, od_config, q) -> None:
         # Read back from the pool at the write slots.
         pool_read_k = kv._k_pools[0][write_slots]
         pool_read_v = kv._v_pools[0][write_slots]
-        assert torch.allclose(pool_read_k, new_k[0], atol=1e-5), "gathered K ≠ written (rank {})".format(rank)
-        assert torch.allclose(pool_read_v, new_v[0], atol=1e-5), "gathered V ≠ written (rank {})".format(rank)
+        assert torch.allclose(pool_read_k, new_k[0], atol=1e-5), f"gathered K ≠ written (rank {rank})"
+        assert torch.allclose(pool_read_v, new_v[0], atol=1e-5), f"gathered V ≠ written (rank {rank})"
 
         # Allocate a second chunk then gather the window.
         kv.allocate_chunk(adapter)
@@ -76,6 +80,7 @@ def _worker_entry(rank: int, od_config, q) -> None:
         q.put((rank, device, runner_name, "pool_roundtrip_ok", None))
     except Exception:  # pragma: no cover - reported to the parent
         import traceback
+
         q.put((rank, None, None, None, traceback.format_exc()))
 
 
diff --git a/tests/bde/test_gather.py b/tests/bde/test_gather.py
index ae4e1a7bc3f..0cdd59b829f 100644
--- a/tests/bde/test_gather.py
+++ b/tests/bde/test_gather.py
@@ -2,10 +2,10 @@
 """Tests for pool write + window gather (Step 4)."""
 
 import torch
+
 from vllm_omni.bde.kv_cache import (
     BDEKVCache,
     BDEKVConfig,
-    BDERequestAdapter,
     allocate_kv_pool,
     pool_gather_window,
     pool_write_chunk,
@@ -79,8 +79,15 @@ def test_bde_kv_cache_write_gather_integration():
     """BDEKVCache.write_chunk_kv → gather_window roundtrip through a real manager."""
     cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
     kv = BDEKVCache(
-        cfg, num_layers=1, num_kv_heads=N_HEADS, head_size=HEAD_DIM, dtype=torch.float32,
-        block_size=BLOCK, max_model_len=512, available_bytes=1 << 24, device=torch.device("cpu"),
+        cfg,
+        num_layers=1,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=512,
+        available_bytes=1 << 24,
+        device=torch.device("cpu"),
     )
     adapter = kv.begin_request("r")
     kv.allocate_chunk(adapter)
diff --git a/tests/bde/test_kv_cache.py b/tests/bde/test_kv_cache.py
index ddb77974fab..8f77e3056f6 100644
--- a/tests/bde/test_kv_cache.py
+++ b/tests/bde/test_kv_cache.py
@@ -8,7 +8,6 @@
 
 import pytest
 import torch
-
 from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.request import RequestStatus
@@ -72,9 +71,9 @@ def test_spec_rejects_inconsistent_window():
 def test_sliding_replace_keeps_window():
     # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
     # the snap is to chunk boundaries.
-    skip = lambda n: chunk_window_skipped_tokens(
-        n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False
-    )
+    def skip(n):
+        return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+
     assert skip(32) == 0  # nothing past the window yet
     assert skip(48) == 16  # one chunk fell out of the window
     assert skip(64) == 32
@@ -82,25 +81,19 @@ def test_sliding_replace_keeps_window():
 
 def test_sliding_replace_snaps_to_chunk_boundary():
     # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
-    skip = chunk_window_skipped_tokens(
-        50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False
-    )
+    skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
     assert skip % 16 == 0 and skip == 16
 
 
 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
-    )
+    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
 
 
 def test_reset_at_boundary_drops_completed_past_sink():
-    skip = chunk_window_skipped_tokens(
-        48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True
-    )
+    skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True)
     # completed = 48; sink = 16 -> drop 32
     assert skip == 32
 
diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index b321c61e651..09ca24824e8 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -2,6 +2,7 @@
 """Tests for the BDEKVState pipeline bridge (Step 5)."""
 
 import torch
+
 from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
 from vllm_omni.bde.kv_cache.gather import BDEKVState
 
@@ -12,9 +13,17 @@
 
 def make_state(num_layers=1, window_chunks=4):
     cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
-    kv = BDEKVCache(cfg, num_layers=num_layers, num_kv_heads=N_HEADS, head_size=HEAD_DIM,
-                    dtype=torch.float32, block_size=BLOCK, max_model_len=4096,
-                    available_bytes=1 << 24, device=torch.device("cpu"))
+    kv = BDEKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 24,
+        device=torch.device("cpu"),
+    )
     pos = kv.begin_request("r-pos")
     neg = kv.begin_request("r-neg")
     return kv, BDEKVState(kv, pos, neg, num_layers=num_layers)
@@ -77,8 +86,8 @@ def test_reset_clears_gather_cache():
 
 def test_eviction_bounds_resident_window():
     kv, st = make_state(num_layers=1, window_chunks=3)
-    st.update_kv_cache(0, _window(3), False, seq_len=3 * BLOCK)          # resident 3/3
-    st.update_kv_cache(0, _window(5), False, seq_len=2 * BLOCK)          # +2 -> evict to 3
+    st.update_kv_cache(0, _window(3), False, seq_len=3 * BLOCK)  # resident 3/3
+    st.update_kv_cache(0, _window(5), False, seq_len=2 * BLOCK)  # +2 -> evict to 3
     got = st.get_kv_caches(False, lambda: ["X"])[0]
     assert got.shape[2] <= 3 * BLOCK  # window bounded
 
@@ -123,6 +132,7 @@ def test_kv_state_noop():
     the existing tests passing unchanged (the pipeline's default __init__
     sets _bde_kv_state = None)."""
     from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
     p = DreamZeroPipeline.__new__(DreamZeroPipeline)
     assert p._bde_kv_state is None
     for attr in ("_kv_get", "_kv_create", "_kv_update", "_kv_commit"):
@@ -133,6 +143,7 @@ def test_kv_create_always_initializes_model_local():
     """_kv_create must call state.create_kv_caches (no recursion), and must do so
     even under BDE so cross-attention (not managed by BDE) is initialized."""
     from unittest.mock import MagicMock
+
     from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 
     p = DreamZeroPipeline.__new__(DreamZeroPipeline)
diff --git a/tests/bde/test_slot_mapping.py b/tests/bde/test_slot_mapping.py
index 4eabb4c6fb9..89acecaf672 100644
--- a/tests/bde/test_slot_mapping.py
+++ b/tests/bde/test_slot_mapping.py
@@ -68,9 +68,7 @@ def test_blocktable_build_single_request():
     real = set(resident_block_ids(block_ids, null_id))
     assert real, "expected at least one real block after allocation"
 
-    slots = chunk_slot_mapping(
-        block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size
-    )
+    slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size)
     # Each slot resolves to a real allocated block, none to null_block.
     blocks_used = {int(s) // spec.block_size for s in slots}
     assert blocks_used <= real
diff --git a/tests/bde/test_validation_harness.py b/tests/bde/test_validation_harness.py
index f1c2ef474df..f4a15c7ad84 100644
--- a/tests/bde/test_validation_harness.py
+++ b/tests/bde/test_validation_harness.py
@@ -128,9 +128,7 @@ def test_kv_memory_plateau(tmp_path):
     from benchmarks.bde.profile import kv_memory_plateau, plot_memory
 
     cfg = bde_config_for_dreamzero(num_frame_per_block=1, frame_seqlen=16, local_attn_size=2)
-    prof = kv_memory_plateau(
-        cfg, num_chunks=10, num_layers=2, num_kv_heads=4, head_size=64, block_size=16
-    )
+    prof = kv_memory_plateau(cfg, num_chunks=10, num_layers=2, num_kv_heads=4, head_size=64, block_size=16)
     # Usage plateaus once the window fills; resident bounded by window + 1.
     assert prof["plateaued"] is True
     assert max(prof["resident_blocks"]) <= cfg.window_chunks + 1
diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/bde/kv_cache/gather.py
index e89f953a7f4..777ab0b9aab 100644
--- a/vllm_omni/bde/kv_cache/gather.py
+++ b/vllm_omni/bde/kv_cache/gather.py
@@ -36,16 +36,8 @@ def allocate_kv_pool(
     k_pools: list[torch.Tensor] = []
     v_pools: list[torch.Tensor] = []
     for _ in range(num_layers):
-        k_pools.append(
-            torch.empty(
-                num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device
-            )
-        )
-        v_pools.append(
-            torch.empty(
-                num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device
-            )
-        )
+        k_pools.append(torch.empty(num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device))
+        v_pools.append(torch.empty(num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device))
     return k_pools, v_pools
 
 
@@ -69,9 +61,7 @@ def pool_write_chunk(
     v_pool[slot_mapping] = new_v[0]
 
 
-def build_window_slots(
-    window_block_ids: list[int], block_size: int, device: torch.device
-) -> torch.Tensor:
+def build_window_slots(window_block_ids: list[int], block_size: int, device: torch.device) -> torch.Tensor:
     """Flat slot index covering all resident blocks, in table order.
 
     Vectorized (no Python loop): ``slot = block_id * block_size + offset`` for every
@@ -180,8 +170,11 @@ def get_kv_caches(self, is_negative: bool, fallback) -> list:
         self._gather_cache[is_negative] = windows
         _log.info(
             "BDE GET   [%s] source=paged-gather layers=%d window0=%s resident_blocks=%d/%d",
-            branch, self.num_layers, tuple(windows[0].shape),
-            len(self.kv_cache.window_block_ids(adapter)), self.kv_cache.spec.window_chunks,
+            branch,
+            self.num_layers,
+            tuple(windows[0].shape),
+            len(self.kv_cache.window_block_ids(adapter)),
+            self.kv_cache.spec.window_chunks,
         )
         return windows
 
@@ -203,7 +196,7 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
             n_chunks = new_count // cs
             slots = []
             for _ in range(n_chunks):
-                self.kv_cache.allocate_chunk(adapter)        # allocate + evict old
+                self.kv_cache.allocate_chunk(adapter)  # allocate + evict old
                 slots.append(self.kv_cache.chunk_write_slots(adapter))
                 adapter.on_chunk_committed()
             self._pending[is_negative] = slots
@@ -213,8 +206,11 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
             self._gather_cache[is_negative] = None
             _log.info(
                 "BDE WRITE [%s] new_tokens=%d -> %d frame-chunks  resident=%d/%d  storing %d layers",
-                branch, new_count, n_chunks,
-                len(self.kv_cache.window_block_ids(adapter)), self.kv_cache.spec.window_chunks,
+                branch,
+                new_count,
+                n_chunks,
+                len(self.kv_cache.window_block_ids(adapter)),
+                self.kv_cache.spec.window_chunks,
                 self.num_layers,
             )
         slots = self._pending[is_negative]
@@ -244,9 +240,9 @@ def reset(self) -> None:
         is recycled.
         """
         pos_id, neg_id = self.pos.request_id, self.neg.request_id
-        self.kv_cache.end_request(self.pos)        # free resident blocks
+        self.kv_cache.end_request(self.pos)  # free resident blocks
         self.kv_cache.end_request(self.neg)
-        self.pos = self.kv_cache.begin_request(pos_id)   # fresh, empty window
+        self.pos = self.kv_cache.begin_request(pos_id)  # fresh, empty window
         self.neg = self.kv_cache.begin_request(neg_id)
         self._committed = {False: 0, True: 0}
         self._pending = {False: [], True: []}
diff --git a/vllm_omni/bde/kv_cache/manager.py b/vllm_omni/bde/kv_cache/manager.py
index 5c0163237df..f1a02f98c69 100644
--- a/vllm_omni/bde/kv_cache/manager.py
+++ b/vllm_omni/bde/kv_cache/manager.py
@@ -13,20 +13,19 @@
 from __future__ import annotations
 
 import torch
-
 from vllm.logger import init_logger
 
 from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
 from vllm_omni.bde.kv_cache.chunk_window import ChunkWindowSpec
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
-from vllm_omni.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
 from vllm_omni.bde.kv_cache.gather import (
     allocate_kv_pool,
     build_window_slots,
     pool_gather_window,
     pool_write_chunk,
 )
+from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
+from vllm_omni.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
 
 _log = init_logger(__name__)
 
@@ -84,8 +83,8 @@ def __init__(
             cross_shape = (2, cross_attn_length, num_kv_heads, head_size)
             bytes_per_element = dtype.itemsize
             cross_bytes = (
-                2           # K + V
-                * 2         # pos + neg
+                2  # K + V
+                * 2  # pos + neg
                 * cross_attn_length
                 * num_kv_heads
                 * head_size
@@ -97,7 +96,10 @@ def __init__(
                 self._cross_v.append(torch.empty(cross_shape, dtype=dtype, device=device))
             _log.info(
                 "BDE cross-attn pool: %d layers × (%d tok × %d heads × %d) = %.1f MiB",
-                num_layers, cross_attn_length, num_kv_heads, head_size,
+                num_layers,
+                cross_attn_length,
+                num_kv_heads,
+                head_size,
                 cross_bytes / (1024 * 1024),
             )
 
@@ -137,8 +139,11 @@ def __init__(
     # read many (every denoising step). Not managed through the paged block pool.
 
     def write_cross_kv(
-        self, layer_idx: int, is_negative: bool,
-        k: torch.Tensor, v: torch.Tensor,
+        self,
+        layer_idx: int,
+        is_negative: bool,
+        k: torch.Tensor,
+        v: torch.Tensor,
     ) -> None:
         """Write one layer's cross-attn K/V into the pool.
 
@@ -158,15 +163,13 @@ def read_cross_kv(self, layer_idx: int, is_negative: bool) -> dict:
         slice rather than from the lazy-initialised model-local dict.
         """
         branch = 1 if is_negative else 0
-        k = self._cross_k[layer_idx][branch].unsqueeze(0)   # (1, L, heads, dim)
+        k = self._cross_k[layer_idx][branch].unsqueeze(0)  # (1, L, heads, dim)
         v = self._cross_v[layer_idx][branch].unsqueeze(0)
         return {"is_init": True, "k": k, "v": v}
 
     # -- request lifecycle ---------------------------------------------------
 
-    def begin_request(
-        self, request_id: str, *, prefill_prefix_tokens: int = 0
-    ) -> BDERequestAdapter:
+    def begin_request(self, request_id: str, *, prefill_prefix_tokens: int = 0) -> BDERequestAdapter:
         adapter = BDERequestAdapter(
             request_id,
             chunk_size=self.spec.chunk_size,
@@ -177,9 +180,12 @@ def begin_request(
         return adapter
 
     def end_request(self, adapter: BDERequestAdapter) -> None:
-        _log.debug("BDE end_request: req=%s chunks=%d free=%d",
-                    adapter.request_id, adapter.completed_chunks,
-                    self.manager.block_pool.get_num_free_blocks())
+        _log.debug(
+            "BDE end_request: req=%s chunks=%d free=%d",
+            adapter.request_id,
+            adapter.completed_chunks,
+            self.manager.block_pool.get_num_free_blocks(),
+        )
         self.manager.free(adapter)
         self._adapters.pop(adapter.request_id, None)
 
@@ -195,9 +201,14 @@ def allocate_chunk(self, adapter: BDERequestAdapter) -> list[int]:
             raise RuntimeError("BDE KV pool exhausted while allocating a chunk")
         table = self.block_table(adapter)
         resident = resident_block_ids(table, self.null_block_id)
-        _log.debug("BDE allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d",
-                    adapter.request_id, adapter.completed_chunks, len(table),
-                    len(resident), self.manager.block_pool.get_num_free_blocks())
+        _log.debug(
+            "BDE allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d",
+            adapter.request_id,
+            adapter.completed_chunks,
+            len(table),
+            len(resident),
+            self.manager.block_pool.get_num_free_blocks(),
+        )
         return table
 
     def block_table(self, adapter: BDERequestAdapter) -> list[int]:
@@ -233,9 +244,14 @@ def write_chunk_kv(
     ) -> None:
         """Write one layer's committed-chunk K/V into the pool."""
         slots = self.chunk_write_slots(adapter)
-        _log.debug("BDE write: req=%s layer=%d chunk=%d shapes=%s dev=%s",
-                    adapter.request_id, layer_index, adapter.completed_chunks,
-                    (tuple(new_k.shape), tuple(new_v.shape)), slots.device)
+        _log.debug(
+            "BDE write: req=%s layer=%d chunk=%d shapes=%s dev=%s",
+            adapter.request_id,
+            layer_index,
+            adapter.completed_chunks,
+            (tuple(new_k.shape), tuple(new_v.shape)),
+            slots.device,
+        )
         pool_write_chunk(
             self._k_pools[layer_index],
             self._v_pools[layer_index],
@@ -258,9 +274,14 @@ def gather_window(self, layer_index: int, adapter: BDERequestAdapter) -> torch.T
             self.block_size,
             self.spec.sliding_window,
         )
-        _log.debug("BDE gather: req=%s layer=%d blocks=%s window=%s dev=%s",
-                    adapter.request_id, layer_index, window_ids,
-                    tuple(window.shape), window.device)
+        _log.debug(
+            "BDE gather: req=%s layer=%d blocks=%s window=%s dev=%s",
+            adapter.request_id,
+            layer_index,
+            window_ids,
+            tuple(window.shape),
+            window.device,
+        )
         return window
 
     def gather_window_all_layers(self, adapter: BDERequestAdapter) -> list[torch.Tensor]:
@@ -276,12 +297,20 @@ def gather_window_all_layers(self, adapter: BDERequestAdapter) -> list[torch.Ten
         slots = build_window_slots(window_ids, self.block_size, self._k_pools[0].device)
         windows = [
             pool_gather_window(
-                self._k_pools[i], self._v_pools[i], window_ids,
-                self.block_size, self.spec.sliding_window, slots=slots,
+                self._k_pools[i],
+                self._v_pools[i],
+                window_ids,
+                self.block_size,
+                self.spec.sliding_window,
+                slots=slots,
             )
             for i in range(self.num_layers)
         ]
-        _log.debug("BDE gather-all: req=%s layers=%d blocks=%d window=%s",
-                    adapter.request_id, self.num_layers, len(window_ids),
-                    tuple(windows[0].shape))
+        _log.debug(
+            "BDE gather-all: req=%s layers=%d blocks=%d window=%s",
+            adapter.request_id,
+            self.num_layers,
+            len(window_ids),
+            tuple(windows[0].shape),
+        )
         return windows
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/bde/runner.py
index 7a7ffc39e4e..bf0f8445e89 100644
--- a/vllm_omni/bde/runner.py
+++ b/vllm_omni/bde/runner.py
@@ -22,7 +22,7 @@
 from vllm_omni.bde.kv_cache.config import BDEKVConfig
 from vllm_omni.bde.kv_cache.gather import BDEKVState
 from vllm_omni.bde.kv_cache.manager import BDEKVCache
-from vllm_omni.diffusion.data import OmniDiffusionConfig, DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 
@@ -149,15 +149,17 @@ def _preallocate_kv_cache(self) -> None:
         chunk_size = frame_seqlen
         window_chunks = self.bde_kv_config.window_chunks or (max_attention_size // frame_seqlen)
 
-        self.bde_kv_config = dataclasses.replace(
-            self.bde_kv_config, chunk_size=chunk_size, window_chunks=window_chunks
-        )
+        self.bde_kv_config = dataclasses.replace(self.bde_kv_config, chunk_size=chunk_size, window_chunks=window_chunks)
         free_bytes = torch.cuda.mem_get_info(self.device)[0]
         logger.info(
             "BDE preallocating (paged): frame_seqlen=%d num_frame_per_block=%d "
             "local_attn_size=%d -> chunk_size=%d window_chunks=%d (window=%d tokens)",
-            frame_seqlen, num_frame_per_block, local_attn_size,
-            chunk_size, window_chunks, window_chunks * chunk_size,
+            frame_seqlen,
+            num_frame_per_block,
+            local_attn_size,
+            chunk_size,
+            window_chunks,
+            window_chunks * chunk_size,
         )
         self.build_kv_cache(
             num_layers=num_layers,
@@ -189,7 +191,11 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
             self._bde_states[session_id] = state
         logger.debug(
             "BDE execute_model: req=%s session=%s chunk_size=%d window_chunks=%s num_blocks=%d",
-            req.request_id, session_id, kv.spec.chunk_size, kv.config.window_chunks, kv.num_blocks,
+            req.request_id,
+            session_id,
+            kv.spec.chunk_size,
+            kv.config.window_chunks,
+            kv.num_blocks,
         )
         self.pipeline._bde_kv_state = state
         try:
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 97a5b21069a..c68a75855d6 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -446,8 +446,7 @@ def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
                 engine_class = resolve_obj_by_qualname(backend)
             except (ImportError, ValueError) as e:
                 raise ValueError(
-                    f"Failed to load engine_backend '{backend}'. Ensure it is a valid "
-                    f"python path. Error: {e}"
+                    f"Failed to load engine_backend '{backend}'. Ensure it is a valid python path. Error: {e}"
                 ) from e
             if not issubclass(engine_class, DiffusionEngine):
                 raise TypeError(f"engine_backend must resolve to a DiffusionEngine subclass. Got {engine_class}.")
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 09e9ced8568..8294fdf1ef8 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -23,8 +23,6 @@
 from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
-_log = logging.getLogger(__name__)
-
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
@@ -54,6 +52,7 @@
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
+_log = logger
 MAX_DREAMZERO_SESSIONS = 64
 
 
@@ -138,9 +137,7 @@ def _kv_get_cross(self, state, is_negative):
         """Cross-attn cache: pool-backed when BDE active, model-local otherwise."""
         s = self._bde_kv_state
         if s is not None:
-            return s.get_cross_kv_caches(
-                is_negative, lambda: state.get_crossattn_caches(is_negative)
-            )
+            return s.get_cross_kv_caches(is_negative, lambda: state.get_crossattn_caches(is_negative))
         return state.get_crossattn_caches(is_negative)
 
     def _kv_populate_cross(self, context: torch.Tensor, is_negative: bool) -> None:
@@ -421,9 +418,10 @@ def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
             is_neg = kwargs.get("is_negative", False)
             if self._bde_kv_state is not None:
                 _log.info(
-                    "BDE pipeline predict_noise -> write-back: is_neg=%s seq_len=%s "
-                    "current_start_frame=%s layers=%d",
-                    is_neg, kwargs.get("seq_len"), kwargs.get("current_start_frame"),
+                    "BDE pipeline predict_noise -> write-back: is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
+                    is_neg,
+                    kwargs.get("seq_len"),
+                    kwargs.get("current_start_frame"),
                     len(updated_kv_caches),
                 )
             for i, kv in enumerate(updated_kv_caches):
@@ -601,7 +599,8 @@ def _prefill_kv_cache(
         head_dim = self.transformer.dim // self.transformer.num_heads
 
         if state.current_start_frame == 0:
-            self._kv_create(state,
+            self._kv_create(
+                state,
                 batch_size,
                 dtype,
                 device,

From bf11741d16978f4e1e7fc6a8b7f35e24a2b4a292 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Fri, 19 Jun 2026 08:00:30 +0000
Subject: [PATCH 24/51] [BDE] Phase 1 review fixes, move to experimental/,
 profiling tooling

Address the PR #4534 review findings, relocate BDE under a new
experimental package, and add memory/memoization profiling support.

Fixes (from the bde_phase1 review):
- P1 session-KV leak: bound BDEModelRunner._bde_states with an LRU at the
  same cap as DreamZeroPipeline._states and free pool blocks on eviction
  via a new BDEKVState.close(), so session-id churn can no longer leak the
  KV pool until exhaustion.
- P2 frame_seqlen source of truth: read it from the loaded transformer;
  the deploy-config inference is now only a drift cross-check (warns on
  mismatch). This surfaced a real 880-vs-220 mismatch on DreamZero-DROID.
- KV pool floor: guarantee 2*(window_chunks+1)+2 blocks so the memory
  fraction can't under-size the pool (fixes exhaustion exposed by the
  corrected 880-token frame granularity).
- P2 benchmark config helper: frame-granular (chunk_size=frame_seqlen,
  window_chunks=local_attn_size), matching the runtime.
- P3: remove the drifted, unused BDEPipelineMixin (its update path was
  missing the required seq_len argument).

Refactor:
- Move vllm_omni/bde -> vllm_omni/experimental/bde; update all imports,
  the engine resolver, and the dynamic runner-class path string.

Profiling:
- BDE_KV_NO_MEMO env toggle to A/B the per-forward gather memoization.
- gpu_mem_stats() worker RPC + peak-VRAM capture in bde_perf_compare.

Validated on H200 (DreamZero-DROID): tests/bde 73 passing; bit-exact
parity (max_abs_diff=0, PSNR=inf) for KV-on and for both memo/no-memo
vs the KV-off baseline. Three-way profile confirms the memoization win
(-9.25 GiB peak VRAM, -101 ms/chunk, sawtooth flattened).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 benchmarks/bde/dreamzero_config.py            | 22 +++----
 benchmarks/bde/profile.py                     |  2 +-
 .../dreamzero/bde_perf_compare.py             | 21 ++++++-
 tests/bde/test_engine_routing.py              |  4 +-
 tests/bde/test_full_chain_gpu.py              |  4 +-
 tests/bde/test_gather.py                      |  2 +-
 tests/bde/test_kv_cache.py                    |  4 +-
 tests/bde/test_kv_orchestration.py            | 29 ++++++++-
 tests/bde/test_pipeline_bridge.py             |  4 +-
 tests/bde/test_slot_mapping.py                |  2 +-
 tests/bde/test_validation_harness.py          | 19 ++++--
 vllm_omni/diffusion/diffusion_engine.py       |  2 +-
 .../models/dreamzero/video_export_worker.py   | 15 +++++
 vllm_omni/experimental/__init__.py            |  8 +++
 vllm_omni/{ => experimental}/bde/__init__.py  |  2 +-
 vllm_omni/{ => experimental}/bde/engine.py    |  6 +-
 .../bde/kv_cache/__init__.py                  | 16 +++--
 .../bde/kv_cache/adapter.py                   |  0
 .../bde/kv_cache/chunk_window.py              |  0
 .../{ => experimental}/bde/kv_cache/config.py |  0
 .../{ => experimental}/bde/kv_cache/gather.py | 60 +++++++------------
 .../bde/kv_cache/manager.py                   | 29 +++++++--
 .../{ => experimental}/bde/kv_cache/pool.py   |  0
 .../bde/kv_cache/slot_mapping.py              |  0
 vllm_omni/{ => experimental}/bde/runner.py    | 36 +++++++++--
 25 files changed, 196 insertions(+), 91 deletions(-)
 create mode 100644 vllm_omni/experimental/__init__.py
 rename vllm_omni/{ => experimental}/bde/__init__.py (86%)
 rename vllm_omni/{ => experimental}/bde/engine.py (90%)
 rename vllm_omni/{ => experimental}/bde/kv_cache/__init__.py (60%)
 rename vllm_omni/{ => experimental}/bde/kv_cache/adapter.py (100%)
 rename vllm_omni/{ => experimental}/bde/kv_cache/chunk_window.py (100%)
 rename vllm_omni/{ => experimental}/bde/kv_cache/config.py (100%)
 rename vllm_omni/{ => experimental}/bde/kv_cache/gather.py (87%)
 rename vllm_omni/{ => experimental}/bde/kv_cache/manager.py (89%)
 rename vllm_omni/{ => experimental}/bde/kv_cache/pool.py (100%)
 rename vllm_omni/{ => experimental}/bde/kv_cache/slot_mapping.py (100%)
 rename vllm_omni/{ => experimental}/bde/runner.py (80%)

diff --git a/benchmarks/bde/dreamzero_config.py b/benchmarks/bde/dreamzero_config.py
index bb326a68b4a..16721b25ac2 100644
--- a/benchmarks/bde/dreamzero_config.py
+++ b/benchmarks/bde/dreamzero_config.py
@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from vllm_omni.bde.kv_cache import BDEKVConfig
+from vllm_omni.experimental.bde.kv_cache import BDEKVConfig
 
 
 def bde_config_for_dreamzero(
@@ -20,9 +20,16 @@ def bde_config_for_dreamzero(
 ) -> BDEKVConfig:
     """Build the parity-faithful ``BDEKVConfig`` for DreamZero.
 
-    ``chunk_size      = num_frame_per_block * frame_seqlen``
+    Frame-granular paging (matches ``BDEModelRunner._preallocate_kv_cache``): one
+    pool block = one frame, so the resident window equals ``max_attention_size``
+    exactly without rounding to whole causal blocks.
+
+    ``chunk_size      = frame_seqlen``
     ``window (tokens) = local_attn_size * frame_seqlen   (== max_attention_size)``
-    ``window_chunks   = local_attn_size // num_frame_per_block``
+    ``window_chunks   = local_attn_size``
+
+    ``num_frame_per_block`` no longer governs paging (it sets the reset cadence); it
+    is kept as an argument for call-site compatibility.
     """
     if num_frame_per_block <= 0 or frame_seqlen <= 0:
         raise ValueError("num_frame_per_block and frame_seqlen must be > 0")
@@ -31,14 +38,9 @@ def bde_config_for_dreamzero(
             f"local_attn_size must be > 0 for a bounded BDE window (got {local_attn_size}); "
             "DreamZero -1 (full attention) needs an explicit window size for Phase 1"
         )
-    if local_attn_size % num_frame_per_block != 0:
-        raise ValueError(
-            f"local_attn_size ({local_attn_size}) must be divisible by "
-            f"num_frame_per_block ({num_frame_per_block}) for exact window parity"
-        )
 
-    chunk_size = num_frame_per_block * frame_seqlen
-    window_chunks = local_attn_size // num_frame_per_block
+    chunk_size = frame_seqlen
+    window_chunks = local_attn_size
     config = BDEKVConfig(
         enable=True,
         chunk_size=chunk_size,
diff --git a/benchmarks/bde/profile.py b/benchmarks/bde/profile.py
index 94218fc4bf1..50914b01242 100644
--- a/benchmarks/bde/profile.py
+++ b/benchmarks/bde/profile.py
@@ -14,7 +14,7 @@
 
 import torch
 
-from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
+from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
 
 
 def kv_memory_plateau(
diff --git a/examples/offline_inference/dreamzero/bde_perf_compare.py b/examples/offline_inference/dreamzero/bde_perf_compare.py
index 274458d8039..c71d5932986 100644
--- a/examples/offline_inference/dreamzero/bde_perf_compare.py
+++ b/examples/offline_inference/dreamzero/bde_perf_compare.py
@@ -77,9 +77,25 @@ def main() -> None:
     prefill = per_forward[0]
     steady = per_forward[2:] if len(per_forward) > 2 else per_forward[1:]  # drop prefill + 1 warmup
     mean_steady = sum(steady) / len(steady)
+
+    # Peak GPU memory over the whole rollout, queried from the worker process where
+    # the model + KV pools live. Monotonic high-water mark from process start.
+    peak_reserved_gib = None
+    peak_allocated_gib = None
+    try:
+        engine = getattr(omni.engine.stage_clients[0], "_engine", None)
+        if engine is not None:
+            stats = engine.executor.collective_rpc("gpu_mem_stats", unique_reply_rank=0, exec_all_ranks=True)
+            stats = stats[0] if isinstance(stats, (list, tuple)) else stats
+            peak_reserved_gib = stats["peak_reserved_gib"]
+            peak_allocated_gib = stats["peak_allocated_gib"]
+    except Exception as e:  # noqa: BLE001 — profiling extra must never fail the run
+        print(f"[perf:{args.tag}] WARN: could not read GPU mem stats: {e}")
+
     summary = {
         "tag": args.tag,
         "bde_kv_enable": os.environ.get("BDE_KV_ENABLE", "0"),
+        "bde_kv_no_memo": os.environ.get("BDE_KV_NO_MEMO", "0"),
         "num_forwards": n,
         "load_s": load_s,
         "prefill_s": prefill,
@@ -87,11 +103,14 @@ def main() -> None:
         "steady_min_s": min(steady),
         "steady_max_s": max(steady),
         "total_gen_s": sum(per_forward),
+        "peak_reserved_gib": peak_reserved_gib,
+        "peak_allocated_gib": peak_allocated_gib,
         "per_forward_s": per_forward,
     }
+    mem_str = f"  peak_reserved={peak_reserved_gib:.2f}GiB" if peak_reserved_gib is not None else ""
     print(
         f"[perf:{args.tag}] prefill={prefill * 1000:.1f}ms  "
-        f"steady_mean={mean_steady * 1000:.1f}ms (n={len(steady)})  total={sum(per_forward):.2f}s"
+        f"steady_mean={mean_steady * 1000:.1f}ms (n={len(steady)})  total={sum(per_forward):.2f}s{mem_str}"
     )
 
     # Decode once (after timing) so latents + the final video are both saved.
diff --git a/tests/bde/test_engine_routing.py b/tests/bde/test_engine_routing.py
index bb301d3920f..b5f52cf4efe 100644
--- a/tests/bde/test_engine_routing.py
+++ b/tests/bde/test_engine_routing.py
@@ -13,7 +13,7 @@
 import pytest
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
-from vllm_omni.bde.engine import (
+from vllm_omni.experimental.bde.engine import (
     BDE_MODEL_RUNNER_CLS,
     BDEEngine,
     apply_bde_runner_default,
@@ -51,7 +51,7 @@ def test_subclass_type_is_returned():
 
 
 def test_qualname_string_resolves():
-    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.bde.engine.BDEEngine"))
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.bde.engine.BDEEngine"))
     assert cls is BDEEngine
 
 
diff --git a/tests/bde/test_full_chain_gpu.py b/tests/bde/test_full_chain_gpu.py
index 8081d7b49d4..11e88c67ae5 100644
--- a/tests/bde/test_full_chain_gpu.py
+++ b/tests/bde/test_full_chain_gpu.py
@@ -20,7 +20,7 @@
     reason="requires >= 2 GPUs",
 )
 
-RUNNER_OVERRIDE = "vllm_omni.bde.runner.BDEModelRunner"
+RUNNER_OVERRIDE = "vllm_omni.experimental.bde.runner.BDEModelRunner"
 BLOCK = 16
 N_HEADS = 4
 HEAD_DIM = 64
@@ -35,7 +35,7 @@ def _worker_entry(rank: int, od_config, q) -> None:
         runner_name = type(worker.model_runner).__name__
 
         # --- pool write + gather on the worker's GPU ------------------------
-        from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
+        from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
 
         cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
         kv = BDEKVCache(
diff --git a/tests/bde/test_gather.py b/tests/bde/test_gather.py
index 0cdd59b829f..e61e66bd156 100644
--- a/tests/bde/test_gather.py
+++ b/tests/bde/test_gather.py
@@ -3,7 +3,7 @@
 
 import torch
 
-from vllm_omni.bde.kv_cache import (
+from vllm_omni.experimental.bde.kv_cache import (
     BDEKVCache,
     BDEKVConfig,
     allocate_kv_pool,
diff --git a/tests/bde/test_kv_cache.py b/tests/bde/test_kv_cache.py
index 8f77e3056f6..57f6c244908 100644
--- a/tests/bde/test_kv_cache.py
+++ b/tests/bde/test_kv_cache.py
@@ -12,7 +12,7 @@
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.request import RequestStatus
 
-from vllm_omni.bde.kv_cache import (
+from vllm_omni.experimental.bde.kv_cache import (
     BDEKVConfig,
     BDERequestAdapter,
     ChunkWindowManager,
@@ -20,7 +20,7 @@
     build_kv_manager,
     compute_num_blocks,
 )
-from vllm_omni.bde.kv_cache.chunk_window import chunk_window_skipped_tokens
+from vllm_omni.experimental.bde.kv_cache.chunk_window import chunk_window_skipped_tokens
 
 BLOCK = 16
 
diff --git a/tests/bde/test_kv_orchestration.py b/tests/bde/test_kv_orchestration.py
index aef124af6b2..266dc12566b 100644
--- a/tests/bde/test_kv_orchestration.py
+++ b/tests/bde/test_kv_orchestration.py
@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
+from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
 
 DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16)
 
@@ -64,3 +64,30 @@ def test_num_computed_advances_per_chunk():
     kv.commit_chunk(a)
     assert a.num_computed_tokens == 16
     kv.end_request(a)
+
+
+def test_state_close_frees_both_branch_blocks():
+    """BDEKVState.close() returns both CFG branches' pool blocks to the pool.
+
+    This is the primitive the runner's LRU eviction relies on: when a session is
+    evicted, close() must free the blocks both adapters hold, or session churn
+    leaks pool ownership (review P1).
+    """
+    from vllm_omni.experimental.bde.kv_cache.gather import BDEKVState
+
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    pos = kv.begin_request("bde__s")
+    neg = kv.begin_request("bde__s__neg")
+    state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
+    for _ in range(3):
+        for adapter in (pos, neg):
+            kv.allocate_chunk(adapter)
+            kv.commit_chunk(adapter)
+    # Both branches hold resident blocks now.
+    assert kv.manager.block_pool.get_num_free_blocks() < free_total
+
+    state.close()
+    # All blocks returned — no leak across an evicted session.
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index 09ca24824e8..ff58962ffe8 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -3,8 +3,8 @@
 
 import torch
 
-from vllm_omni.bde.kv_cache import BDEKVCache, BDEKVConfig
-from vllm_omni.bde.kv_cache.gather import BDEKVState
+from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
+from vllm_omni.experimental.bde.kv_cache.gather import BDEKVState
 
 BLOCK = 16
 N_HEADS = 4
diff --git a/tests/bde/test_slot_mapping.py b/tests/bde/test_slot_mapping.py
index 89acecaf672..31464b869c1 100644
--- a/tests/bde/test_slot_mapping.py
+++ b/tests/bde/test_slot_mapping.py
@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-from vllm_omni.bde.kv_cache import (
+from vllm_omni.experimental.bde.kv_cache import (
     BDERequestAdapter,
     ChunkWindowSpec,
     build_kv_manager,
diff --git a/tests/bde/test_validation_harness.py b/tests/bde/test_validation_harness.py
index f4a15c7ad84..2e778f915f2 100644
--- a/tests/bde/test_validation_harness.py
+++ b/tests/bde/test_validation_harness.py
@@ -58,16 +58,25 @@ def test_check_gates():
 
 
 def test_config_derivation_matches_window():
+    # Frame-granular: one pool block = one frame, so chunk_size == frame_seqlen and
+    # window_chunks == local_attn_size (matches BDEModelRunner._preallocate_kv_cache,
+    # which overwrites chunk_size to frame_seqlen but preserves window_chunks).
     cfg = bde_config_for_dreamzero(num_frame_per_block=3, frame_seqlen=220, local_attn_size=21)
-    assert cfg.chunk_size == 660
-    assert cfg.window_chunks == 7
+    assert cfg.chunk_size == 220
+    assert cfg.window_chunks == 21
     assert cfg.sliding_window == 21 * 220  # == max_attention_size
     assert cfg.enable and not cfg.reset_at_boundary
 
 
-def test_config_rejects_nondivisible_window():
-    with pytest.raises(ValueError):
-        bde_config_for_dreamzero(num_frame_per_block=4, frame_seqlen=10, local_attn_size=21)
+def test_config_window_is_frame_granular():
+    # Regression: frame granularity makes the window exact for any
+    # num_frame_per_block (no divisibility constraint), and window_chunks tracks
+    # local_attn_size rather than local_attn_size // num_frame_per_block — the old
+    # causal-block math the runtime would silently shrink to a too-small window.
+    cfg = bde_config_for_dreamzero(num_frame_per_block=4, frame_seqlen=10, local_attn_size=21)
+    assert cfg.chunk_size == 10
+    assert cfg.window_chunks == 21
+    assert cfg.sliding_window == 21 * 10
 
 
 def test_config_rejects_full_attention():
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index c68a75855d6..c4eeac40a64 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -438,7 +438,7 @@ def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
         if backend == "bde":
             # Lazy import to avoid a circular dependency (bde.engine imports this
             # module), mirroring how DiffusionExecutor.get_class imports backends.
-            from vllm_omni.bde.engine import BDEEngine
+            from vllm_omni.experimental.bde.engine import BDEEngine
 
             return BDEEngine
         if isinstance(backend, str):
diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
index 20c8cd8cdc7..dfab272d65e 100644
--- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
+++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
@@ -9,6 +9,21 @@
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
+    def gpu_mem_stats(self) -> dict:
+        """Peak GPU memory (GiB) for this worker's device, for profiling.
+
+        ``max_memory_reserved`` is the caching-allocator high-water mark (the best
+        in-process proxy for "peak VRAM"); ``max_memory_allocated`` is the live
+        tensor high-water mark. Both are monotonic from process start, so the
+        end-of-run values capture the whole rollout's peak.
+        """
+        dev = torch.cuda.current_device()
+        return {
+            "device": int(dev),
+            "peak_reserved_gib": torch.cuda.max_memory_reserved(dev) / (1024**3),
+            "peak_allocated_gib": torch.cuda.max_memory_allocated(dev) / (1024**3),
+        }
+
     def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor:
         if self.model_runner is None or self.model_runner.pipeline is None:
             raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
diff --git a/vllm_omni/experimental/__init__.py b/vllm_omni/experimental/__init__.py
new file mode 100644
index 00000000000..1be37e247b2
--- /dev/null
+++ b/vllm_omni/experimental/__init__.py
@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Experimental vLLM-Omni subsystems.
+
+Modules under ``vllm_omni.experimental`` are in active development and not yet
+covered by the stability guarantees of the rest of the package — their APIs,
+config surfaces, and defaults may change without notice. The Block Diffusion
+Engine (:mod:`vllm_omni.experimental.bde`) lives here.
+"""
diff --git a/vllm_omni/bde/__init__.py b/vllm_omni/experimental/bde/__init__.py
similarity index 86%
rename from vllm_omni/bde/__init__.py
rename to vllm_omni/experimental/bde/__init__.py
index 4dfcc40d2be..03c88964779 100644
--- a/vllm_omni/bde/__init__.py
+++ b/vllm_omni/experimental/bde/__init__.py
@@ -7,6 +7,6 @@
 ``OmniDiffusionConfig.engine_backend = "bde"``.
 """
 
-from vllm_omni.bde.engine import BDEEngine
+from vllm_omni.experimental.bde.engine import BDEEngine
 
 __all__ = ["BDEEngine"]
diff --git a/vllm_omni/bde/engine.py b/vllm_omni/experimental/bde/engine.py
similarity index 90%
rename from vllm_omni/bde/engine.py
rename to vllm_omni/experimental/bde/engine.py
index 2df0acdd173..e8914ee1a60 100644
--- a/vllm_omni/bde/engine.py
+++ b/vllm_omni/experimental/bde/engine.py
@@ -12,7 +12,7 @@
 logger = init_logger(__name__)
 
 #: Import path of the runner the BDE engine routes its workers to.
-BDE_MODEL_RUNNER_CLS = "vllm_omni.bde.runner.BDEModelRunner"
+BDE_MODEL_RUNNER_CLS = "vllm_omni.experimental.bde.runner.BDEModelRunner"
 
 
 def apply_bde_runner_default(od_config: OmniDiffusionConfig) -> None:
@@ -43,8 +43,8 @@ class BDEEngine(DiffusionEngine):
 
     Architecture note: in the multiproc setup the KV cache lives in the worker /
     runner process (GPU side), co-located with the model and KV tensors — so the
-    actual KV *body* is :class:`~vllm_omni.bde.kv_cache.manager.BDEKVCache`, owned
-    by :class:`~vllm_omni.bde.runner.BDEModelRunner`. ``BDEEngine`` itself is the
+    actual KV *body* is :class:`~vllm_omni.experimental.bde.kv_cache.manager.BDEKVCache`, owned
+    by :class:`~vllm_omni.experimental.bde.runner.BDEModelRunner`. ``BDEEngine`` itself is the
     thin selection / injection seam; it wires the BDE executor → worker → runner
     so DreamZero's rollout runs against the runner-owned KV cache.
     """
diff --git a/vllm_omni/bde/kv_cache/__init__.py b/vllm_omni/experimental/bde/kv_cache/__init__.py
similarity index 60%
rename from vllm_omni/bde/kv_cache/__init__.py
rename to vllm_omni/experimental/bde/kv_cache/__init__.py
index a89616c3776..040652ee758 100644
--- a/vllm_omni/bde/kv_cache/__init__.py
+++ b/vllm_omni/experimental/bde/kv_cache/__init__.py
@@ -6,25 +6,23 @@
 AR-diffusion models. See ``BDE_doc/dreamzero_kv_phase1_plan.md``.
 """
 
-from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
-from vllm_omni.bde.kv_cache.chunk_window import ChunkWindowManager, ChunkWindowSpec
-from vllm_omni.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.bde.kv_cache.gather import (
-    BDEPipelineMixin,
+from vllm_omni.experimental.bde.kv_cache.adapter import BDERequestAdapter
+from vllm_omni.experimental.bde.kv_cache.chunk_window import ChunkWindowManager, ChunkWindowSpec
+from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.experimental.bde.kv_cache.gather import (
     allocate_kv_pool,
     pool_gather_window,
     pool_write_chunk,
 )
-from vllm_omni.bde.kv_cache.manager import BDEKVCache
-from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
-from vllm_omni.bde.kv_cache.slot_mapping import (
+from vllm_omni.experimental.bde.kv_cache.manager import BDEKVCache
+from vllm_omni.experimental.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
+from vllm_omni.experimental.bde.kv_cache.slot_mapping import (
     chunk_slot_mapping,
     compute_slot_mapping,
     resident_block_ids,
 )
 
 __all__ = [
-    "BDEPipelineMixin",
     "BDEKVCache",
     "BDEKVConfig",
     "BDERequestAdapter",
diff --git a/vllm_omni/bde/kv_cache/adapter.py b/vllm_omni/experimental/bde/kv_cache/adapter.py
similarity index 100%
rename from vllm_omni/bde/kv_cache/adapter.py
rename to vllm_omni/experimental/bde/kv_cache/adapter.py
diff --git a/vllm_omni/bde/kv_cache/chunk_window.py b/vllm_omni/experimental/bde/kv_cache/chunk_window.py
similarity index 100%
rename from vllm_omni/bde/kv_cache/chunk_window.py
rename to vllm_omni/experimental/bde/kv_cache/chunk_window.py
diff --git a/vllm_omni/bde/kv_cache/config.py b/vllm_omni/experimental/bde/kv_cache/config.py
similarity index 100%
rename from vllm_omni/bde/kv_cache/config.py
rename to vllm_omni/experimental/bde/kv_cache/config.py
diff --git a/vllm_omni/bde/kv_cache/gather.py b/vllm_omni/experimental/bde/kv_cache/gather.py
similarity index 87%
rename from vllm_omni/bde/kv_cache/gather.py
rename to vllm_omni/experimental/bde/kv_cache/gather.py
index 777ab0b9aab..898f1acf0a0 100644
--- a/vllm_omni/bde/kv_cache/gather.py
+++ b/vllm_omni/experimental/bde/kv_cache/gather.py
@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+import os
+
 import torch
 from vllm.logger import init_logger
 
@@ -148,6 +150,10 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
         # re-materializing the same tensors ~16x. Memoize the gathered window per
         # branch; invalidate on write-back (window grows) and on reset.
         self._gather_cache: dict[bool, list | None] = {False: None, True: None}
+        # Profiling escape hatch: BDE_KV_NO_MEMO=1 disables the per-forward gather
+        # memoization so every denoise step re-gathers the window (the pre-memo
+        # behavior), for A/B measurement of the memoization win. Default: enabled.
+        self._memo_enabled = os.environ.get("BDE_KV_NO_MEMO") != "1"
         # Cross-attn KV is populated once after text encoding (eager), then
         # read every denoising step.  Reset clears these flags so the next
         # forward repopulates from the new text encoding.
@@ -162,12 +168,13 @@ def get_kv_caches(self, is_negative: bool, fallback) -> list:
             out = fallback()
             _log.info("BDE GET   [%s] source=model-local-empty layers=%d", branch, len(out))
             return out
-        cached = self._gather_cache[is_negative]
+        cached = self._gather_cache[is_negative] if self._memo_enabled else None
         if cached is not None:  # window unchanged this forward -> reuse (no re-gather)
             return cached
         adapter = self._adapter(is_negative)
         windows = self.kv_cache.gather_window_all_layers(adapter)
-        self._gather_cache[is_negative] = windows
+        if self._memo_enabled:
+            self._gather_cache[is_negative] = windows
         _log.info(
             "BDE GET   [%s] source=paged-gather layers=%d window0=%s resident_blocks=%d/%d",
             branch,
@@ -229,6 +236,18 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
     def commit_chunk(self) -> None:
         _log.info("BDE COMMIT (paged; resident window retained across forwards)")
 
+    def close(self) -> None:
+        """Final teardown — free both branches' resident pool blocks.
+
+        Unlike :meth:`reset`, this does **not** restart the adapters: the caller is
+        dropping this ``BDEKVState`` for good (session eviction / shutdown), so the
+        blocks return to the ``BlockPool`` and the dead adapters are discarded with
+        the state. Frees pool ownership that would otherwise leak when a session is
+        evicted from the runner's session map.
+        """
+        self.kv_cache.end_request(self.pos)
+        self.kv_cache.end_request(self.neg)
+
     def reset(self) -> None:
         """Drop this session's KV — mirrors the model-local ``state.reset()``.
 
@@ -240,8 +259,7 @@ def reset(self) -> None:
         is recycled.
         """
         pos_id, neg_id = self.pos.request_id, self.neg.request_id
-        self.kv_cache.end_request(self.pos)  # free resident blocks
-        self.kv_cache.end_request(self.neg)
+        self.close()  # free resident blocks for both branches
         self.pos = self.kv_cache.begin_request(pos_id)  # fresh, empty window
         self.neg = self.kv_cache.begin_request(neg_id)
         self._committed = {False: 0, True: 0}
@@ -249,37 +267,3 @@ def reset(self) -> None:
         self._gather_cache = {False: None, True: None}
         self._cross_populated = {False: False, True: False}
         _log.info("BDE RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)
-
-
-class BDEPipelineMixin:
-    """Mixin that replaces model-local KV state with pool-backed storage.
-
-    A DreamZero pipeline inheriting this mixin delegates KV access through
-    proxy methods (``_bde_kv_get`` / ``_bde_kv_create`` / ``_bde_kv_update``)
-    that check for ``self._bde_kv_state`` and route to :class:`BDEKVState`
-    when present, falling back to the ``DreamZeroState`` methods otherwise.
-    """
-
-    _bde_kv_state: BDEKVState | None = None
-
-    # -- proxy methods (call these instead of state.*) -----------------------
-
-    def _bde_kv_get(self, state, is_negative):
-        if self._bde_kv_state is not None:
-            return self._bde_kv_state.get_kv_caches(is_negative)
-        return state.get_kv_caches(is_negative)
-
-    def _bde_kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
-        if self._bde_kv_state is not None:
-            return  # pool already allocated
-        state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
-
-    def _bde_kv_update(self, state, layer_idx, updated_kv, is_negative):
-        if self._bde_kv_state is not None:
-            self._bde_kv_state.update_kv_cache(layer_idx, updated_kv, is_negative)
-            return
-        state.update_kv_cache(layer_idx, updated_kv, is_negative=is_negative)
-
-    def _bde_kv_commit(self):
-        if self._bde_kv_state is not None:
-            self._bde_kv_state.commit_chunk()
diff --git a/vllm_omni/bde/kv_cache/manager.py b/vllm_omni/experimental/bde/kv_cache/manager.py
similarity index 89%
rename from vllm_omni/bde/kv_cache/manager.py
rename to vllm_omni/experimental/bde/kv_cache/manager.py
index f1a02f98c69..b7badae3180 100644
--- a/vllm_omni/bde/kv_cache/manager.py
+++ b/vllm_omni/experimental/bde/kv_cache/manager.py
@@ -15,17 +15,17 @@
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.bde.kv_cache.adapter import BDERequestAdapter
-from vllm_omni.bde.kv_cache.chunk_window import ChunkWindowSpec
-from vllm_omni.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.bde.kv_cache.gather import (
+from vllm_omni.experimental.bde.kv_cache.adapter import BDERequestAdapter
+from vllm_omni.experimental.bde.kv_cache.chunk_window import ChunkWindowSpec
+from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.experimental.bde.kv_cache.gather import (
     allocate_kv_pool,
     build_window_slots,
     pool_gather_window,
     pool_write_chunk,
 )
-from vllm_omni.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
-from vllm_omni.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
+from vllm_omni.experimental.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
+from vllm_omni.experimental.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
 
 _log = init_logger(__name__)
 
@@ -121,6 +121,23 @@ def __init__(
             config.gpu_memory_fraction,
             self.spec.page_size_bytes * num_layers,
         )
+        # Floor: one forward needs the resident window plus the in-flight chunk for
+        # BOTH CFG branches (pos/neg are independent adapters allocating from the
+        # same pool), with a little eviction-transient headroom. The memory-fraction
+        # heuristic can under-size this once block_size grows — e.g. frame-granular
+        # paging at the true frame_seqlen makes each block larger and the pool
+        # fewer-blocks — so guarantee the minimum the rollout cannot run without,
+        # otherwise allocate_chunk hits an exhausted pool mid-forward.
+        min_blocks = 2 * (config.window_chunks + 1) + 2
+        if num_blocks < min_blocks:
+            _log.warning(
+                "BDE KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
+                "floor (2 CFG branches x (window_chunks=%d + 1) + 2 headroom)",
+                num_blocks,
+                min_blocks,
+                config.window_chunks,
+            )
+            num_blocks = min_blocks
         layer_names = [f"bde.layer.{i}" for i in range(num_layers)]
         self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
         self.num_blocks = num_blocks
diff --git a/vllm_omni/bde/kv_cache/pool.py b/vllm_omni/experimental/bde/kv_cache/pool.py
similarity index 100%
rename from vllm_omni/bde/kv_cache/pool.py
rename to vllm_omni/experimental/bde/kv_cache/pool.py
diff --git a/vllm_omni/bde/kv_cache/slot_mapping.py b/vllm_omni/experimental/bde/kv_cache/slot_mapping.py
similarity index 100%
rename from vllm_omni/bde/kv_cache/slot_mapping.py
rename to vllm_omni/experimental/bde/kv_cache/slot_mapping.py
diff --git a/vllm_omni/bde/runner.py b/vllm_omni/experimental/bde/runner.py
similarity index 80%
rename from vllm_omni/bde/runner.py
rename to vllm_omni/experimental/bde/runner.py
index bf0f8445e89..b68c8d98226 100644
--- a/vllm_omni/bde/runner.py
+++ b/vllm_omni/experimental/bde/runner.py
@@ -15,14 +15,16 @@
 
 import dataclasses
 import os
+from collections import OrderedDict
 
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.bde.kv_cache.gather import BDEKVState
-from vllm_omni.bde.kv_cache.manager import BDEKVCache
+from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.experimental.bde.kv_cache.gather import BDEKVState
+from vllm_omni.experimental.bde.kv_cache.manager import BDEKVCache
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import MAX_DREAMZERO_SESSIONS
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 
@@ -64,7 +66,12 @@ def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None:
         self.kv_cache: BDEKVCache | None = None
         # DreamZero KV is session-scoped (persists across a session's forwards),
         # so BDE KV state is keyed by session_id and reused, not created per request.
-        self._bde_states: dict[str, BDEKVState] = {}
+        # Bounded by the same LRU cap as DreamZeroPipeline._states: an evicted
+        # session's BDEKVState owns pool blocks (two CFG adapters), so without a
+        # bound, session-id churn in a long-running server would leak pool ownership
+        # until the KV pool is exhausted. Evicting frees those blocks (state.close()).
+        self._bde_states: OrderedDict[str, BDEKVState] = OrderedDict()
+        self._max_bde_states = MAX_DREAMZERO_SESSIONS
 
     def build_kv_cache(
         self,
@@ -133,7 +140,17 @@ def _preallocate_kv_cache(self) -> None:
         head_size = int(t.dim // t.num_heads)
         num_frame_per_block = int(t.num_frame_per_block)
         local_attn_size = int(t.local_attn_size)
-        frame_seqlen = self._infer_frame_seqlen()
+        # The loaded transformer is the source of truth for frame_seqlen (it derives
+        # max_attention_size from it). The deploy-config inference is only a
+        # cross-check — warn on drift rather than silently sizing chunks wrong.
+        frame_seqlen = int(getattr(t, "frame_seqlen", 0)) or self._infer_frame_seqlen()
+        inferred = self._infer_frame_seqlen()
+        if frame_seqlen != inferred:
+            logger.warning(
+                "BDE frame_seqlen mismatch: transformer=%d deploy-config=%d; using transformer",
+                frame_seqlen,
+                inferred,
+            )
         # The model's own attention window, in tokens (= the [-max_attention_size:]
         # slice it applies). Read it directly rather than recomputing.
         max_attention_size = int(t.blocks[0].self_attn.max_attention_size)
@@ -189,6 +206,15 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
             neg = kv.begin_request(f"bde__{session_id}__neg")
             state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
             self._bde_states[session_id] = state
+            # Evict the least-recently-used session(s) past the cap, freeing their
+            # pool blocks — mirrors DreamZeroPipeline._states so the BDE pool can't
+            # outlive the model-local state map it shadows.
+            while len(self._bde_states) > self._max_bde_states:
+                old_id, old_state = self._bde_states.popitem(last=False)
+                old_state.close()
+                logger.debug("BDE evicted session=%s (LRU); freed pool blocks", old_id)
+        # Track recency on every access (hit or miss) so the LRU order is correct.
+        self._bde_states.move_to_end(session_id)
         logger.debug(
             "BDE execute_model: req=%s session=%s chunk_size=%d window_chunks=%s num_blocks=%d",
             req.request_id,

From 522b5e578354c806934744ce1967b104204d0d72 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 20 Jun 2026 05:58:43 +0000
Subject: [PATCH 25/51] [BDE] Address PR #4534 review (project_kv, routing
 gate, commit_chunk docs)

Close inline review findings 2-5 from hsliuustc0106:
- Inline 2: delete the two unused project_kv methods in causal_wan_model.py
  (_kv_populate_cross inlines the identical projection); fix the stale comment.
- Inline 3: gate DreamZero->BDE engine/runner routing behind BDE_KV_ENABLE via
  resolve_default_engine_backend(), so non-BDE deploys stay on base
  DiffusionEngine and base is selectable again. +4 unit tests.
- Inline 4: docstrings clarifying the two commit_chunk paths (test-only vs the
  real advance inside BDEKVState.update_kv_cache).
- Inline 5: #4425 structured-config coordination is blocked (not in this tree);
  deferral marker left at the field defs + a note in the remaining-work doc.

Validated on H200 (DreamZero-DROID): tests/bde 77 passing; routing gate confirmed
end-to-end (KV-off -> base engine, BDE_KV_ENABLE=1 -> BDE); bit-exact parity
(max_abs_diff=0, PSNR=inf) preserved.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_engine_routing.py              | 42 ++++++++++++--
 vllm_omni/diffusion/data.py                   | 57 +++++++++++++++----
 .../models/dreamzero/causal_wan_model.py      | 24 --------
 .../models/dreamzero/pipeline_dreamzero.py    |  6 +-
 vllm_omni/experimental/bde/kv_cache/gather.py |  7 +++
 .../experimental/bde/kv_cache/manager.py      | 10 +++-
 6 files changed, 103 insertions(+), 43 deletions(-)

diff --git a/tests/bde/test_engine_routing.py b/tests/bde/test_engine_routing.py
index b5f52cf4efe..5dea9d4f184 100644
--- a/tests/bde/test_engine_routing.py
+++ b/tests/bde/test_engine_routing.py
@@ -13,16 +13,18 @@
 import pytest
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
-from vllm_omni.experimental.bde.engine import (
-    BDE_MODEL_RUNNER_CLS,
-    BDEEngine,
-    apply_bde_runner_default,
-)
 from vllm_omni.diffusion.data import (
     OmniDiffusionConfig,
+    _bde_routing_enabled,
     default_engine_backend_for_model,
+    resolve_default_engine_backend,
 )
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.bde.engine import (
+    BDE_MODEL_RUNNER_CLS,
+    BDEEngine,
+    apply_bde_runner_default,
+)
 
 
 def _cfg(backend):
@@ -101,6 +103,36 @@ def test_dreamzero_default_routes_to_bde_engine():
     assert DiffusionEngine.resolve_engine_class(_cfg(backend)) is BDEEngine
 
 
+# --- BDE routing is gated on BDE_KV_ENABLE (PR #4534 blast-radius review) ----
+
+
+def test_routing_enabled_reads_env(monkeypatch):
+    monkeypatch.setenv("BDE_KV_ENABLE", "1")
+    assert _bde_routing_enabled() is True
+    monkeypatch.delenv("BDE_KV_ENABLE", raising=False)
+    assert _bde_routing_enabled() is False
+
+
+def test_default_backend_applied_only_when_enabled(monkeypatch):
+    # DreamZero + unset ("default") backend: auto-route applies only with BDE on.
+    monkeypatch.setenv("BDE_KV_ENABLE", "1")
+    assert resolve_default_engine_backend("DreamZeroPipeline", "default") == "bde"
+    monkeypatch.delenv("BDE_KV_ENABLE", raising=False)
+    assert resolve_default_engine_backend("DreamZeroPipeline", "default") is None
+
+
+def test_explicit_backend_not_overridden_even_when_enabled(monkeypatch):
+    # An explicit engine_backend bypasses the auto-route gate entirely.
+    monkeypatch.setenv("BDE_KV_ENABLE", "1")
+    assert resolve_default_engine_backend("DreamZeroPipeline", "bde") is None
+    assert resolve_default_engine_backend("DreamZeroPipeline", "vllm_omni.experimental.bde.engine.BDEEngine") is None
+
+
+def test_non_bde_model_never_routed(monkeypatch):
+    monkeypatch.setenv("BDE_KV_ENABLE", "1")
+    assert resolve_default_engine_backend("WanS2VPipeline", "default") is None
+
+
 # --- BDE worker/runner wiring -----------------------------------------------
 
 
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index b7236112075..9f6b6df1c4b 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -495,12 +495,43 @@ def default_engine_backend_for_model(model_class_name: str | None) -> str | None
     on the Block Diffusion Engine with engine-level KV cache management; returns
     ``None`` to keep the base ``DiffusionEngine``. Extend
     ``_BDE_ENGINE_MODEL_CLASSES`` as more AR-DiT models migrate onto BDE.
+
+    NOTE: this is the *pure* "what would this model pick" lookup. Whether it is
+    actually applied is gated by :func:`_bde_routing_enabled` at the call sites, so
+    a deploy that hasn't opted into BDE stays on the base engine.
     """
     if model_class_name in _BDE_ENGINE_MODEL_CLASSES:
         return "bde"
     return None
 
 
+def _bde_routing_enabled() -> bool:
+    """Whether to auto-route a BDE-eligible model onto the BDE engine/runner.
+
+    Gated on the same switch as BDE KV (``BDE_KV_ENABLE``, read by
+    ``resolve_bde_kv_config``): a deploy that hasn't opted into BDE KV stays on the
+    base ``DiffusionEngine``/runner, keeping the (still-unvalidated for the
+    production CFG-parallel + CUDA-graph + compile topology) BDE runner off the
+    default path. An explicit ``engine_backend`` selection bypasses this gate
+    entirely (it only governs the ``"default"`` -> per-model auto-route). (PR #4534)
+    """
+    return os.environ.get("BDE_KV_ENABLE") == "1"
+
+
+def resolve_default_engine_backend(model_class_name: str | None, current_backend: str | None) -> str | None:
+    """Engine backend to apply for an auto-routed model, or ``None`` to leave as-is.
+
+    Applies the per-model default (:func:`default_engine_backend_for_model`) only
+    when ``current_backend`` is the unset ``"default"`` sentinel *and* BDE routing
+    is opted in (:func:`_bde_routing_enabled`). Returns ``None`` otherwise, so an
+    explicit backend choice and BDE-off deploys both stay untouched. Shared by the
+    two routing sites (``__post_init__`` auto-detect and ``from_kwargs``).
+    """
+    if current_backend != "default" or not _bde_routing_enabled():
+        return None
+    return default_engine_backend_for_model(model_class_name)
+
+
 @dataclass
 class OmniDiffusionConfig:
     # Model and path configuration (for convenience)
@@ -558,6 +589,11 @@ class OmniDiffusionConfig:
     # The worker uses it instead of the platform default when set; the BDE engine
     # sets it to BDEModelRunner. Unset -> platform default (existing behavior).
     diffusion_model_runner_cls: str | None = None
+    # DEFERRED (#4425 / RFC #4021 structured config): when rebasing onto the
+    # structured VllmOmniConfig, mirror engine_backend + diffusion_model_runner_cls
+    # into its DiffusionConfig (engine/runtime group) and parity test. enrich_config()
+    # there only copies fields present on DiffusionConfig, so without the mirror
+    # DreamZero's BDE routing silently drops on the structured-config path. (PR #4534)
 
     # HuggingFace specific parameters
     trust_remote_code: bool = False
@@ -1082,11 +1118,12 @@ def enrich_config(self) -> None:
                         self.set_tf_model_config(TransformerConfig())
                         self.update_multimodal_support()
                         # DreamZero is an AR-diffusion world model: route it to the
-                        # BDE engine (engine-level KV cache management) unless the
-                        # caller explicitly selected an engine backend.
-                        bde_default = default_engine_backend_for_model(self.model_class_name)
-                        if bde_default is not None and self.engine_backend == "default":
-                            self.engine_backend = bde_default
+                        # BDE engine (engine-level KV cache management) when BDE is
+                        # opted in (BDE_KV_ENABLE) and the caller did not explicitly
+                        # select an engine backend; otherwise stay on base.
+                        backend = resolve_default_engine_backend(self.model_class_name, self.engine_backend)
+                        if backend is not None:
+                            self.engine_backend = backend
                     else:
                         raise
                 elif architectures and len(architectures) == 1:
@@ -1160,11 +1197,11 @@ def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
         instance = cls(**filtered_kwargs)
         # Route per-model engines after model_class_name is finalized. This covers
         # the explicit-model_class_name path (e.g. deploy configs) that the
-        # __post_init__ auto-detect branch does not reach.
-        if instance.engine_backend == "default":
-            backend = default_engine_backend_for_model(instance.model_class_name)
-            if backend is not None:
-                instance.engine_backend = backend
+        # __post_init__ auto-detect branch does not reach. Gated on BDE_KV_ENABLE
+        # so non-BDE deploys stay on the base DiffusionEngine (PR #4534 review).
+        backend = resolve_default_engine_backend(instance.model_class_name, instance.engine_backend)
+        if backend is not None:
+            instance.engine_backend = backend
         return instance
 
 
diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
index a6fbf1d004d..06e541610f0 100644
--- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
+++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
@@ -290,18 +290,6 @@ def forward(
         x = self.o(x)
         return x
 
-    def project_kv(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        """Project encoder hidden states to cross-attn K/V (for eager caching).
-
-        Returns ``(k, v)`` where each has shape
-        ``(B, text_len, tp_num_heads, head_dim)`` — the same tensors the
-        ``forward`` method would compute on first call.
-        """
-        n, d = self.tp_num_heads, self.head_dim
-        k = self.norm_k(self.k(context)).unflatten(2, (n, d))
-        v = self.v(context).unflatten(2, (n, d))
-        return k, v
-
 
 class WanI2VCrossAttention(nn.Module):
     """Image-to-video cross-attention (splits first 257 image tokens).
@@ -372,18 +360,6 @@ def forward(
         x = self.o(x)
         return x
 
-    def project_kv(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        """Project encoder hidden states to cross-attn K/V for the text part only.
-
-        Image tokens (first 257) are *not* cached — they are always recomputed.
-        Returns ``(k, v)`` for the text tokens.
-        """
-        context = context[:, 257:]  # strip image tokens; cache text only
-        n, d = self.tp_num_heads, self.head_dim
-        k = self.norm_k(self.k(context)).unflatten(2, (n, d))
-        v = self.v(context).unflatten(2, (n, d))
-        return k, v
-
 
 WAN_CROSSATTENTION_CLASSES = {
     "t2v_cross_attn": WanT2VCrossAttention,
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 8294fdf1ef8..5d19ed7ce44 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -161,9 +161,9 @@ def _kv_populate_cross(self, context: torch.Tensor, is_negative: bool) -> None:
         for i, block in enumerate(self.transformer.blocks):
             # The cross-attn forward caches text-only k/v — the image tokens
             # (first 257, prepended inside _forward_blocks via img_emb) are
-            # always recomputed.  project_kv for I2V strips those image tokens,
-            # but we are feeding raw projected text (no image prepend), so
-            # compute k/v directly from the full projected text.
+            # always recomputed.  Here we feed raw projected text (no image
+            # prepend), so we compute k/v directly from the full projected text,
+            # mirroring the cross-attn forward's own k/v projection.
             ca = block.cross_attn
             n, d = ca.tp_num_heads, ca.head_dim
             k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
diff --git a/vllm_omni/experimental/bde/kv_cache/gather.py b/vllm_omni/experimental/bde/kv_cache/gather.py
index 898f1acf0a0..09a5e163314 100644
--- a/vllm_omni/experimental/bde/kv_cache/gather.py
+++ b/vllm_omni/experimental/bde/kv_cache/gather.py
@@ -234,6 +234,13 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
             pool_write_chunk(kpool, vpool, k, v, sm)
 
     def commit_chunk(self) -> None:
+        """No-op for DreamZero (logs only).
+
+        The per-chunk advance already happened in :meth:`update_kv_cache`, which
+        calls ``adapter.on_chunk_committed()`` at layer 0 while allocating each
+        frame-chunk. This method exists so the pipeline's ``_kv_commit`` bridge has
+        a symmetric call; the resident window is retained across forwards.
+        """
         _log.info("BDE COMMIT (paged; resident window retained across forwards)")
 
     def close(self) -> None:
diff --git a/vllm_omni/experimental/bde/kv_cache/manager.py b/vllm_omni/experimental/bde/kv_cache/manager.py
index b7badae3180..58c58de5d9a 100644
--- a/vllm_omni/experimental/bde/kv_cache/manager.py
+++ b/vllm_omni/experimental/bde/kv_cache/manager.py
@@ -245,7 +245,15 @@ def window_block_ids(self, adapter: BDERequestAdapter) -> list[int]:
         return resident_block_ids(self.block_table(adapter), self.null_block_id)
 
     def commit_chunk(self, adapter: BDERequestAdapter) -> None:
-        """Advance after the chunk's K/V is written. Once per chunk, not per step."""
+        """Advance the adapter by one chunk after its K/V is written.
+
+        This is the standalone per-chunk advance used by the unit lifecycle and the
+        ``tests/bde`` suite. In the DreamZero production path the advance happens
+        elsewhere: :meth:`BDEKVState.update_kv_cache` (gather.py) calls
+        ``adapter.on_chunk_committed()`` directly as it allocates each frame-chunk,
+        and the pipeline's ``_kv_commit`` -> :meth:`BDEKVState.commit_chunk` is only
+        a no-op log. Call once per chunk, not per denoise step.
+        """
         _log.debug("BDE commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
         adapter.on_chunk_committed()
         _log.debug("BDE commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)

From 01bc4171df2a6ee9799a6994eb61eb10b6931641 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 20 Jun 2026 06:52:41 +0000
Subject: [PATCH 26/51] Merge origin/main into bde-kv-phase1 (resolve DreamZero
 conflicts)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Bring PR #4534 up to date with upstream main (DreamZero CUDA graph /
torch.compile / DiT step-cache #4213, tp & cross-attn cache fix #4154,
GR00T, stepcache, etc.). Conflict resolution:

- pipeline_dreamzero.py:
  * diffuse loop: adopt upstream's run_dit (DiT step-cache) wrapper, keep
    BDE's _kv_get / _kv_get_cross proxies for all KV reads.
  * reset: adopt upstream's reset_reason (session/inference) model; route
    both session and inference (window-boundary) resets through _kv_reset
    so the BDE pool resets too (parity). _kv_reset gains clear_video_latents.
  * take upstream's state.append_video_latents(video_out).
  * cross-attn: BDE now uses the model-local cross-attn cache (upstream #4154
    added k_img/v_img image-token caching that BDE's text-only pool lacked);
    _kv_get_cross returns model-local, _kv_populate_cross is a no-op. BDE's
    self-attn paged pool is unchanged.
- causal_wan_model.py: silent auto-merge duplicated the crossattn_cache arg in
  _forward_blocks (upstream #4154 and our wiring-gap fix are the same change) —
  de-duplicated to one param.
- video_export_worker.py: keep both gpu_mem_stats (ours) and
  _latents_to_uint8_frames (upstream); switch gpu_mem_stats off the now-banned
  torch.cuda.* to torch.accelerator.* (TID251).
- bde_perf_compare.py: update to upstream's new _build_observations signature.

Validated on H200 (DreamZero-DROID): tests/bde 77 passing; whole package
compiles; KV-off (base engine) and BDE_KV_ENABLE=1 both run on the merged code
and stay bit-exact (max_abs_diff=0, PSNR=inf).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 README.md                                     |    1 +
 docs/models/supported_models.md               |    6 +-
 docs/serving/speech_api.md                    |    7 +
 docs/serving/videos_api.md                    |    7 +-
 .../examples/online_serving/image_to_video.md |   18 +-
 .../examples/online_serving/text_to_video.md  |   18 +-
 .../dreamzero/bde_perf_compare.py             |    9 +-
 .../dreamzero/client_schedule.py              |  129 +
 .../dreamzero/export_prediction_video.py      |  136 +-
 .../online_serving/image_to_video/README.md   |   11 +-
 .../online_serving/text_to_video/README.md    |   11 +-
 recipes/NVIDIA/GR00T-N1.7.md                  |  101 +
 recipes/README.md                             |   11 +-
 recipes/mistralai/Voxtral-TTS.md              |   42 +
 requirements/common.txt                       |    2 +
 tests/config/test_server_settings.py          |   84 +
 tests/diffusion/cache/test_cache_backends.py  |  115 +-
 tests/diffusion/cache/test_cache_dit.py       |   47 +-
 tests/diffusion/cache/test_stepcache.py       |  116 +
 tests/diffusion/models/gr00t/__init__.py      |    0
 tests/diffusion/models/gr00t/test_pipeline.py |  141 +
 .../models/ltx2/test_ltx2_3_pipeline.py       |   17 +-
 tests/diffusion/quantization/test_hsdp_fp8.py |  147 +
 tests/diffusion/test_diffusion_worker.py      |   25 +-
 tests/dreamzero/test_crossattn_cache.py       |  114 +
 tests/dreamzero/test_fused_qk_rms_norm.py     |  112 +
 tests/dreamzero/test_qkv_fusion.py            |  117 +
 tests/e2e/offline_inference/test_moss_tts.py  |  333 +--
 .../test_moss_tts_realtime.py                 |  290 ++
 .../test_dreamzero_expansion.py               |    2 -
 tests/e2e/online_serving/test_gr00t_openpi.py |  163 ++
 .../openai_api/test_openpi_connection.py      |   41 +
 .../openai_api/test_storage_manager.py        |   94 +
 .../openai_api/test_video_server.py           |   51 +-
 tests/gr00t/__init__.py                       |    2 +
 tests/gr00t/openpi_client_helper.py           |  183 ++
 .../models/moss_tts/test_per_request_codes.py |  176 ++
 tests/test_omni_request.py                    |   89 +
 vllm_omni/config/pipeline_registry.py         |    4 +
 vllm_omni/config/server_settings.py           |   58 +
 vllm_omni/deploy/Gr00tN1d7.yaml               |   41 +
 vllm_omni/deploy/dreamzero.yaml               |    7 +-
 vllm_omni/deploy/dreamzero_tp1_cfg2.yaml      |    7 +-
 vllm_omni/diffusion/cache/base.py             |    2 +
 .../diffusion/cache/cache_dit_backend.py      | 2407 +++++------------
 vllm_omni/diffusion/cache/selector.py         |   13 +-
 .../diffusion/cache/stepcache/__init__.py     |   66 +
 .../diffusion/cache/stepcache/backend.py      |  129 +
 vllm_omni/diffusion/cache/stepcache/config.py |   51 +
 vllm_omni/diffusion/cache/stepcache/state.py  |   51 +
 vllm_omni/diffusion/data.py                   |   13 +
 vllm_omni/diffusion/distributed/hsdp.py       |    7 +-
 .../model_loader/diffusers_loader.py          |   27 +-
 .../models/bagel/bagel_transformer.py         |    9 +
 .../models/cosmos3/transformer_cosmos3.py     |   12 +
 .../diffusion/models/dreamid_omni/fusion.py   |    9 +
 .../models/dreamzero/causal_wan_model.py      |  124 +-
 .../models/dreamzero/pipeline_dreamzero.py    |  359 ++-
 .../models/dreamzero/state_dreamzero.py       |  111 +-
 .../models/dreamzero/video_export_worker.py   |   31 +-
 .../dreamzero/wan_vae_feat_cache_patch.py     |   92 +
 .../ernie_image/ernie_image_transformer.py    |    8 +
 .../diffusion/models/flux/flux_transformer.py |    9 +
 .../models/flux2/flux2_transformer.py         |    9 +
 .../flux2_klein/flux2_klein_transformer.py    |    9 +
 .../models/glm_image/glm_image_transformer.py |   10 +
 vllm_omni/diffusion/models/gr00t/__init__.py  |    6 +
 .../models/gr00t/configs/__init__.py          |    8 +
 .../gr00t/configs/embodiment_configs.py       |   12 +
 .../models/gr00t/configs/gr00t_n1d7.py        |  109 +
 .../diffusion/models/gr00t/dataio/__init__.py |    2 +
 .../models/gr00t/dataio/embodiment_tags.py    |  191 ++
 .../gr00t/dataio/state_action/__init__.py     |    2 +
 .../dataio/state_action/action_chunking.py    |  654 +++++
 .../models/gr00t/dataio/state_action/pose.py  |  709 +++++
 .../state_action/state_action_processor.py    |  643 +++++
 .../diffusion/models/gr00t/dataio/types.py    |  112 +
 .../diffusion/models/gr00t/dataio/utils.py    |  187 ++
 .../models/gr00t/modeling/__init__.py         |    7 +
 .../models/gr00t/modeling/gr00t_n1d7.py       |  491 ++++
 .../models/gr00t/modeling/modules/__init__.py |    2 +
 .../models/gr00t/modeling/modules/dit.py      |  344 +++
 .../modules/embodiment_conditioned_mlp.py     |  178 ++
 .../gr00t/modeling/processing_gr00t_n1d7.py   |  555 ++++
 .../diffusion/models/gr00t/pipeline_gr00t.py  |  148 +
 vllm_omni/diffusion/models/gr00t/policy.py    |  434 +++
 .../models/helios/helios_transformer.py       |    9 +
 .../pipeline_hidream_i1_image.py              | 1059 --------
 .../hunyuan_image3_transformer.py             |    9 +
 .../hunyuan_video_15_transformer.py           |    9 +
 .../longcat_image_transformer.py              |   10 +
 .../diffusion/models/ltx2/ltx2_transformer.py |   10 +
 .../diffusion/models/sd3/sd3_transformer.py   |    8 +
 .../sensenova_u1/pipeline_sensenova_u1.py     |    2 +
 .../sensenova_u1/sensenova_u1_transformer.py  |   10 +
 .../stable_audio/stable_audio_transformer.py  |    8 +
 vllm_omni/diffusion/quantization/__init__.py  |    2 +
 vllm_omni/diffusion/quantization/hsdp_fp8.py  |  112 +
 vllm_omni/diffusion/registry.py               |    5 +
 vllm_omni/diffusion/utils/flow_matching.py    |   28 +
 .../worker/diffusion_model_runner.py          |   24 +-
 .../diffusion/worker/diffusion_worker.py      |   15 +-
 vllm_omni/engine/async_omni_engine.py         |   10 +-
 vllm_omni/entrypoints/cli/serve.py            |    2 +-
 vllm_omni/entrypoints/openai/api_server.py    |   73 +-
 .../entrypoints/openai/serving_speech.py      |    1 -
 vllm_omni/entrypoints/openai/storage.py       |  170 +-
 .../entrypoints/openai/tts_adapters/base.py   |    3 -
 vllm_omni/entrypoints/openpi/connection.py    |   98 +-
 .../models/cosyvoice3/cosyvoice3.py           |   35 +-
 .../models/cosyvoice3/cosyvoice3_code2wav.py  |   73 +-
 .../models/cosyvoice3/pipeline.py             |    6 +-
 .../model_executor/models/cosyvoice3/utils.py |   22 -
 .../models/covo_audio/config_covo_audio.py    |    1 -
 .../models/covo_audio/token2wav.py            |    5 +-
 .../fish_speech/fish_speech_dac_decoder.py    |   18 +-
 .../models/fish_speech/fish_speech_fast_ar.py |   11 +-
 .../models/fish_speech/prompt_utils.py        |   10 -
 .../models/glm_tts/glm_tts_dit.py             |    4 -
 .../model_executor/models/gr00t/__init__.py   |    2 +
 .../model_executor/models/gr00t/pipeline.py   |   26 +
 .../configuration_higgs_audio_v2.py           |   12 -
 .../higgs_audio_v2/higgs_audio_v2_talker.py   |   35 +-
 .../models/higgs_audio_v3/ref_audio_cache.py  |  111 -
 .../models/ming_flash_omni/projectors.py      |   27 -
 .../models/ming_flash_omni/text_processing.py |   91 -
 .../models/ming_tts/config_ming_tts.py        |    2 -
 .../models/ming_tts/constants.py              |    1 -
 .../models/ming_tts/ming_tts.py               |    4 -
 .../models/ming_tts/ming_tts_llm.py           |   14 +-
 .../models/moss_tts/audio_tokenizer.py        |    2 +-
 .../moss_tts/modeling_moss_tts_talker.py      |  166 +-
 .../configuration_moss_tts_nano.py            |   60 -
 .../moss_tts_nano/modeling_moss_tts_nano.py   |    1 -
 .../models/omnivoice/omnivoice.py             |    3 +-
 .../models/omnivoice/omnivoice_generator.py   |   23 +-
 .../qwen3_tts/configuration_qwen3_tts.py      |   11 +-
 .../models/qwen3_tts/prompt_embeds_builder.py |    8 +-
 .../models/qwen3_tts/qwen3_tts_talker.py      |    7 -
 .../modeling_qwen3_tts_tokenizer_v2.py        |    5 -
 .../qwen3_tts/tokenizer_25hz/vq/core_vq.py    |    5 -
 .../models/step_audio2/step_audio2.py         |    8 +-
 .../step_audio2/step_audio2_constants.py      |   36 -
 .../models/step_audio2/step_audio2_thinker.py |   30 +-
 .../step_audio2/step_audio2_token2wav.py      |  121 -
 .../models/voxcpm2/minicpm4_paged.py          |   19 +-
 .../models/voxcpm2/voxcpm2_talker.py          |   43 +-
 .../models/voxtral_tts/voxtral_tts.py         |    3 -
 .../voxtral_tts_audio_generation.py           |   26 +-
 .../voxtral_tts_audio_tokenizer.py            |    4 +-
 vllm_omni/request.py                          |    7 +-
 .../transformers_utils/processors/ming.py     |    1 -
 vllm_omni/worker/gpu_model_runner.py          |   14 +
 153 files changed, 10076 insertions(+), 4151 deletions(-)
 create mode 100644 examples/offline_inference/dreamzero/client_schedule.py
 create mode 100644 recipes/NVIDIA/GR00T-N1.7.md
 create mode 100644 tests/config/test_server_settings.py
 create mode 100644 tests/diffusion/cache/test_stepcache.py
 create mode 100644 tests/diffusion/models/gr00t/__init__.py
 create mode 100644 tests/diffusion/models/gr00t/test_pipeline.py
 create mode 100644 tests/diffusion/quantization/test_hsdp_fp8.py
 create mode 100644 tests/dreamzero/test_crossattn_cache.py
 create mode 100644 tests/dreamzero/test_fused_qk_rms_norm.py
 create mode 100644 tests/dreamzero/test_qkv_fusion.py
 create mode 100644 tests/e2e/offline_inference/test_moss_tts_realtime.py
 create mode 100644 tests/e2e/online_serving/test_gr00t_openpi.py
 create mode 100644 tests/entrypoints/openai_api/test_storage_manager.py
 create mode 100644 tests/gr00t/__init__.py
 create mode 100644 tests/gr00t/openpi_client_helper.py
 create mode 100644 tests/model_executor/models/moss_tts/test_per_request_codes.py
 create mode 100644 tests/test_omni_request.py
 create mode 100644 vllm_omni/config/server_settings.py
 create mode 100644 vllm_omni/deploy/Gr00tN1d7.yaml
 create mode 100644 vllm_omni/diffusion/cache/stepcache/__init__.py
 create mode 100644 vllm_omni/diffusion/cache/stepcache/backend.py
 create mode 100644 vllm_omni/diffusion/cache/stepcache/config.py
 create mode 100644 vllm_omni/diffusion/cache/stepcache/state.py
 create mode 100644 vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/__init__.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/configs/__init__.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/dataio/__init__.py
 create mode 100755 vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/dataio/types.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/dataio/utils.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/modeling/__init__.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py
 create mode 100644 vllm_omni/diffusion/models/gr00t/policy.py
 delete mode 100644 vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py
 create mode 100644 vllm_omni/diffusion/quantization/__init__.py
 create mode 100644 vllm_omni/diffusion/quantization/hsdp_fp8.py
 create mode 100644 vllm_omni/diffusion/utils/flow_matching.py
 create mode 100644 vllm_omni/model_executor/models/gr00t/__init__.py
 create mode 100644 vllm_omni/model_executor/models/gr00t/pipeline.py
 delete mode 100644 vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py
 delete mode 100644 vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py

diff --git a/README.md b/README.md
index cf1658b4574..e300b6b8d1e 100644
--- a/README.md
+++ b/README.md
@@ -68,6 +68,7 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
 - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
 - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
+- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
 
 ## Contributing
 
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 119f3bf683c..c224c1a6a2f 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -6,6 +6,8 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
 
 If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
 
+For deployment recipes (TTS and multimodal), see [recipes.vllm.ai](https://recipes.vllm.ai) and the in-repo [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes) directory.
+
 ## List of Supported Models
 
 <style>
@@ -22,6 +24,7 @@ th {
 | `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
 | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
+| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | |
 | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -63,13 +66,14 @@ th {
 | `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
 | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
 | `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
 | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
 | `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
 | `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
+| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
 | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
 | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md
index 935deba4f2d..b70c282979f 100644
--- a/docs/serving/speech_api.md
+++ b/docs/serving/speech_api.md
@@ -9,6 +9,13 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
 
 See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano.
 
+!!! tip "Deployment recipes"
+    TTS deployment recipes are published at
+    [recipes.vllm.ai](https://recipes.vllm.ai) (e.g.
+    [Qwen3-TTS](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice),
+    [Higgs-Audio v3](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b)).
+    The in-repo runbooks live under [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes).
+
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
 ## Quick Start
diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md
index 4d156bdf0f0..95b5dbe7734 100644
--- a/docs/serving/videos_api.md
+++ b/docs/serving/videos_api.md
@@ -195,13 +195,16 @@ curl -X POST http://localhost:8091/v1/videos/sync \
 
 ## Storage
 
-Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are
+Set `VLLM_OMNI_SERVER_STORAGE__PATH` to control where asynchronous video outputs are
 stored:
 
 ```bash
-export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
 ```
 
+> `VLLM_OMNI_STORAGE_PATH` is deprecated and will be removed in a future release;
+> use `VLLM_OMNI_SERVER_STORAGE__PATH` instead.
+
 ## Model-Specific Examples
 
 For complete text-to-video, image-to-video, and model-specific video-to-video
diff --git a/docs/user_guide/examples/online_serving/image_to_video.md b/docs/user_guide/examples/online_serving/image_to_video.md
index 781f0c2a5ed..2dbc57633b8 100644
--- a/docs/user_guide/examples/online_serving/image_to_video.md
+++ b/docs/user_guide/examples/online_serving/image_to_video.md
@@ -84,16 +84,26 @@ curl -X POST http://localhost:8091/v1/videos/sync \
 Generated video files are stored on local disk by the async video API.
 Local file storage behavior can be controlled via the following environment variables:
 
-- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
-- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files
+  (default: `/tmp/storage`)
+- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete/open
+  operations (default: `4`)
+- `VLLM_OMNI_SERVER_STORAGE__FILE_TTL`: optional TTL for generated files in
+  seconds
+- `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL`: optional sweep interval in
+  seconds for enforcing file TTL; defaults to `300` when TTL is set
 
 Example:
 
 ```bash
-export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
-export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
+export VLLM_OMNI_SERVER_STORAGE__FILE_TTL=86400
+export VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL=300
 ```
 
+See also: [Configuration Options](../../../configuration/README.md)
+
 ## API Calls
 
 ### Method 1: Using curl
diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md
index b918aac19d0..d4d520d5e76 100644
--- a/docs/user_guide/examples/online_serving/text_to_video.md
+++ b/docs/user_guide/examples/online_serving/text_to_video.md
@@ -90,16 +90,26 @@ curl -X POST http://localhost:8091/v1/videos/sync \
 Generated video files are stored on local disk by the async video API.
 Local file storage behavior can be controlled via the following environment variables:
 
-- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
-- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files
+  (default: `/tmp/storage`)
+- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete/open
+  operations (default: `4`)
+- `VLLM_OMNI_SERVER_STORAGE__FILE_TTL`: optional TTL for generated files in
+  seconds
+- `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL`: optional sweep interval in
+  seconds for enforcing file TTL; defaults to `300` when TTL is set
 
 Example:
 
 ```bash
-export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
-export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
+export VLLM_OMNI_SERVER_STORAGE__FILE_TTL=86400
+export VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL=300
 ```
 
+See also: [Configuration Options](../../../configuration/README.md)
+
 ## API Calls
 
 ### Method 1: Using curl
diff --git a/examples/offline_inference/dreamzero/bde_perf_compare.py b/examples/offline_inference/dreamzero/bde_perf_compare.py
index c71d5932986..ab4725eb086 100644
--- a/examples/offline_inference/dreamzero/bde_perf_compare.py
+++ b/examples/offline_inference/dreamzero/bde_perf_compare.py
@@ -41,8 +41,13 @@ def main() -> None:
     ap.add_argument("--timing", type=Path, default=None)
     args = ap.parse_args()
 
-    _, base = E._build_observations(args.video_dir, prompt=args.prompt, session_id=args.session_id)
-    observations = [base[0]] + [base[1]] * args.num_chunks
+    _, observations = E._build_observations(
+        args.video_dir,
+        prompt=args.prompt,
+        session_id=args.session_id,
+        num_chunks=args.num_chunks,
+        repeat_chunk_observations=True,
+    )
     n = len(observations)
     print(
         f"[perf:{args.tag}] {n} forwards (1 prefill + {args.num_chunks} chunk) "
diff --git a/examples/offline_inference/dreamzero/client_schedule.py b/examples/offline_inference/dreamzero/client_schedule.py
new file mode 100644
index 00000000000..e4dbad55a03
--- /dev/null
+++ b/examples/offline_inference/dreamzero/client_schedule.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero AR client frame schedule aligned with upstream ``test_client_AR.py``."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+
+# Frame schedule constants (matching upstream debug_inference.py / test_client_AR.py).
+ACTION_HORIZON = 24
+DEFAULT_ACTION_DIM = 8
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+DEFAULT_NUM_AR_CHUNKS = 15
+
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+
+def build_frame_schedule(
+    total_frames: int,
+    num_chunks: int,
+    *,
+    logger: logging.Logger | None = None,
+) -> list[list[int]]:
+    """Build the frame index schedule for multi-frame chunks.
+
+    Args:
+        total_frames: Number of frames available in each camera video.
+        num_chunks: Number of 4-frame chunks to schedule after the initial frame.
+
+    Returns:
+        A list of frame-index lists. Each inner list has four indices. The
+        returned list may be shorter than ``num_chunks`` when the videos run
+        out of frames (upstream stops early instead of erroring).
+    """
+    chunks: list[list[int]] = []
+    current_frame = 23
+    for _ in range(num_chunks):
+        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            if logger is not None:
+                logger.info(
+                    "Frame %s >= %s, stopping at %s chunks",
+                    indices[-1],
+                    total_frames,
+                    len(chunks),
+                )
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    *,
+    prompt: str,
+    session_id: str,
+) -> dict[str, Any]:
+    """Build an observation dict from real video frames.
+
+    For one frame each image key is ``(H, W, 3)``. For four frames each key is
+    ``(4, H, W, 3)``.
+    """
+    obs: dict[str, Any] = {}
+    for camera_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]
+        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def build_ar_observations(
+    camera_frames: dict[str, np.ndarray],
+    *,
+    prompt: str,
+    session_id: str,
+    num_chunks: int = DEFAULT_NUM_AR_CHUNKS,
+    repeat_chunk_observations: bool = False,
+    logger: logging.Logger | None = None,
+) -> list[dict[str, Any]]:
+    """Build the AR observation sequence used by upstream ``test_client_AR.py``.
+
+    Step 0 sends frame ``[0]``. Each subsequent step sends one 4-frame chunk.
+
+    ``num_chunks`` counts only the 4-frame chunks after the initial frame, not
+    the total number of inferences.
+    """
+    if num_chunks < 0:
+        raise ValueError("num_chunks must be non-negative")
+
+    total_frames = min(frames.shape[0] for frames in camera_frames.values())
+    observations = [
+        make_obs_from_video(
+            camera_frames,
+            [0],
+            prompt=prompt,
+            session_id=session_id,
+        )
+    ]
+
+    chunk_schedule = build_frame_schedule(total_frames, num_chunks, logger=logger)
+    if repeat_chunk_observations and chunk_schedule and len(chunk_schedule) < num_chunks:
+        while len(chunk_schedule) < num_chunks:
+            chunk_schedule.append(chunk_schedule[-1])
+
+    for frame_indices in chunk_schedule:
+        observations.append(
+            make_obs_from_video(
+                camera_frames,
+                frame_indices,
+                prompt=prompt,
+                session_id=session_id,
+            )
+        )
+    return observations
diff --git a/examples/offline_inference/dreamzero/export_prediction_video.py b/examples/offline_inference/dreamzero/export_prediction_video.py
index 82f56ff26e5..4c560418c2b 100644
--- a/examples/offline_inference/dreamzero/export_prediction_video.py
+++ b/examples/offline_inference/dreamzero/export_prediction_video.py
@@ -11,6 +11,11 @@
 import cv2
 import numpy as np
 import torch
+from client_schedule import (
+    CAMERA_FILES,
+    DEFAULT_NUM_AR_CHUNKS,
+    build_ar_observations,
+)
 from PIL import Image
 
 from vllm_omni import Omni
@@ -20,13 +25,6 @@
 WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
 ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
 DEFAULT_SESSION_PREFIX = "dreamzero-export"
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-ACTION_HORIZON = 24
-CAMERA_FILES = {
-    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
-    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
-    "observation/wrist_image_left": "wrist_image_left.mp4",
-}
 
 
 def _parse_args() -> argparse.Namespace:
@@ -49,6 +47,17 @@ def _parse_args() -> argparse.Namespace:
         default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
     )
     parser.add_argument("--session-id", default=None)
+    parser.add_argument(
+        "--num-chunks",
+        type=int,
+        default=DEFAULT_NUM_AR_CHUNKS,
+        help="Number of 4-frame chunks to send after the initial frame (default: 15).",
+    )
+    parser.add_argument(
+        "--repeat-chunk-observations",
+        action="store_true",
+        help="Repeat the last valid chunk observation when bundled assets run out of frames.",
+    )
     parser.add_argument("--save-input-video", action="store_true")
     parser.add_argument("--save-gif", action="store_true")
     parser.add_argument("--save-actions", action="store_true")
@@ -84,57 +93,25 @@ def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
     return camera_frames
 
 
-def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
-    chunks: list[list[int]] = []
-    current_frame = 23
-    for _ in range(num_chunks):
-        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
-        if indices[-1] >= total_frames:
-            break
-        chunks.append(indices)
-        current_frame += ACTION_HORIZON
-    return chunks
-
-
-def _make_obs_from_video(
-    camera_frames: dict[str, np.ndarray],
-    frame_indices: list[int],
-    *,
+def _build_observations(
+    video_dir: Path,
     prompt: str,
     session_id: str,
-) -> dict:
-    obs: dict = {}
-    for camera_key, all_frames in camera_frames.items():
-        selected = all_frames[frame_indices]
-        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
-
-    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
-    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
-    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-    obs["prompt"] = prompt
-    obs["session_id"] = session_id
-    return obs
-
-
-def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
+    *,
+    num_chunks: int,
+    repeat_chunk_observations: bool,
+) -> tuple[dict[str, np.ndarray], list[dict]]:
     camera_frames = _load_camera_frames(video_dir)
-    total_frames = min(frames.shape[0] for frames in camera_frames.values())
-    chunks = _build_frame_schedule(total_frames, 1)
-    observations = [
-        _make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
-    ]
-    if chunks:
-        observations.append(
-            _make_obs_from_video(
-                camera_frames,
-                chunks[0],
-                prompt=prompt,
-                session_id=session_id,
-            )
-        )
+    observations = build_ar_observations(
+        camera_frames,
+        prompt=prompt,
+        session_id=session_id,
+        num_chunks=num_chunks,
+        repeat_chunk_observations=repeat_chunk_observations,
+    )
     if len(observations) < 2:
-        raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
-    return camera_frames, observations[:2]
+        raise RuntimeError("Need at least two DreamZero observations to export a prediction video.")
+    return camera_frames, observations
 
 
 def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
@@ -229,6 +206,26 @@ def _run_generation(
     return omni, outputs
 
 
+def _decode_accumulated_with_worker(omni: Omni, session_id: str) -> np.ndarray:
+    """Decode server-side accumulated latents when the pipeline supports it."""
+    stage_client = omni.engine.stage_clients[0]
+    engine = getattr(stage_client, "_engine", None)
+    if engine is None:
+        raise RuntimeError("DreamZero export requires inline diffusion stage access.")
+
+    decoded = engine.executor.collective_rpc(
+        "decode_accumulated_video_latents_to_uint8",
+        args=(session_id,),
+        unique_reply_rank=0,
+        exec_all_ranks=True,
+    )
+    if isinstance(decoded, torch.Tensor):
+        decoded = decoded.numpy()
+    if not isinstance(decoded, np.ndarray):
+        raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
+    return decoded
+
+
 def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
     stage_client = omni.engine.stage_clients[0]
     engine = getattr(stage_client, "_engine", None)
@@ -248,6 +245,15 @@ def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
     return decoded
 
 
+def _decode_video(omni: Omni, session_id: str, outputs: list[OmniRequestOutput]) -> np.ndarray:
+    """Decode full AR rollout by concatenating per-step latents from all outputs."""
+    latent_steps = [_extract_latents(output) for output in outputs]
+    latent_t = [step.shape[2] for step in latent_steps]
+    print(f"Decoding {len(latent_steps)} latent steps, latent_t={latent_t}, total_t={sum(latent_t)}")
+    full_latents = torch.cat(latent_steps, dim=2)
+    return _decode_with_worker(omni, full_latents)
+
+
 def main() -> None:
     args = _parse_args()
     session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
@@ -256,6 +262,12 @@ def main() -> None:
         video_dir=args.video_dir,
         prompt=args.prompt,
         session_id=session_id,
+        num_chunks=args.num_chunks,
+        repeat_chunk_observations=args.repeat_chunk_observations,
+    )
+    print(
+        f"Running {len(observations)} inferences "
+        f"(1 initial + {len(observations) - 1} chunks; requested {args.num_chunks} post-initial chunks)"
     )
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
@@ -273,17 +285,19 @@ def main() -> None:
             deploy_config_path=args.deploy_config,
             observations=observations,
         )
-        latent_steps = [_extract_latents(output) for output in outputs]
-        full_latents = torch.cat(latent_steps, dim=2)
-        frames = _decode_with_worker(omni, full_latents)
+        frames = _decode_video(omni, session_id, outputs)
     finally:
         if omni is not None:
             omni.close()
 
+    expected_frames = 8 * len(observations) + 1
     mp4_path = args.output_dir / f"{args.output_stem}.mp4"
 
     _write_mp4(mp4_path, frames, fps=args.fps)
-    print(f"SAVED_MP4={mp4_path}")
+    print(
+        f"SAVED_MP4={mp4_path} frames={len(frames)} "
+        f"expected≈{expected_frames} duration≈{len(frames) / max(args.fps, 1):.1f}s"
+    )
 
     if args.save_gif:
         gif_path = args.output_dir / f"{args.output_stem}.gif"
@@ -294,8 +308,10 @@ def main() -> None:
         npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
         np.savez(
             npz_path,
-            step0=np.asarray(outputs[0].multimodal_output.get("actions")),
-            step1=np.asarray(outputs[1].multimodal_output.get("actions")),
+            **{
+                f"step{index}": np.asarray(output.multimodal_output.get("actions"))
+                for index, output in enumerate(outputs)
+            },
         )
         print(f"SAVED_ACTIONS={npz_path}")
 
diff --git a/examples/online_serving/image_to_video/README.md b/examples/online_serving/image_to_video/README.md
index 285eeb27983..f18cea1c8ce 100644
--- a/examples/online_serving/image_to_video/README.md
+++ b/examples/online_serving/image_to_video/README.md
@@ -120,14 +120,17 @@ Use `/v1/videos/sync` if you want to write the MP4 directly to a file. `POST /v1
 Generated video files are stored on local disk by the async video API.
 Local file storage behavior can be controlled via the following environment variables:
 
-- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
-- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+`VLLM_OMNI_STORAGE_PATH` and `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` are deprecated and will be
+removed in a future release; use the names above instead.
 
 Example:
 
 ```bash
-export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
-export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
 ```
 
 ## API Calls
diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md
index c01e0602ff9..35b7f028490 100644
--- a/examples/online_serving/text_to_video/README.md
+++ b/examples/online_serving/text_to_video/README.md
@@ -87,14 +87,17 @@ curl -X POST http://localhost:8091/v1/videos/sync \
 Generated video files are stored on local disk by the async video API.
 Local file storage behavior can be controlled via the following environment variables:
 
-- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
-- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+`VLLM_OMNI_STORAGE_PATH` and `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` are deprecated and will be
+removed in a future release; use the names above instead.
 
 Example:
 
 ```bash
-export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
-export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
 ```
 
 ## API Calls
diff --git a/recipes/NVIDIA/GR00T-N1.7.md b/recipes/NVIDIA/GR00T-N1.7.md
new file mode 100644
index 00000000000..91f676a524a
--- /dev/null
+++ b/recipes/NVIDIA/GR00T-N1.7.md
@@ -0,0 +1,101 @@
+# GR00T-N1.7
+
+> NVIDIA Isaac GR00T-N1.7-3B robot VLA policy served over the OpenPI WebSocket protocol
+
+## Summary
+
+- Vendor: NVIDIA
+- Model: `nvidia/GR00T-N1.7-3B`
+- Task: Vision-Language-Action (VLA) inference for robot manipulation
+- Mode: Online serving via OpenPI WebSocket endpoint
+- Maintainer: timzsu
+
+## When to use this recipe
+
+Use this recipe when you need to serve GR00T-N1.7 as a real-time robot policy
+over the OpenPI WebSocket API. It configures the DROID embodiment
+(`OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT`) and exposes the standard DROID action
+keys (`eef_9d`, `gripper_position`, `joint_position`) with action horizon 40.
+
+## References
+
+- Upstream model: <https://huggingface.co/nvidia/GR00T-N1.7-3B>
+- Upstream codebase: <https://github.com/NVIDIA/Isaac-GR00T>
+- OpenPI client library: <https://github.com/Physical-Intelligence/openpi>
+- Pipeline: `vllm_omni.diffusion.models.gr00t.pipeline_gr00t.Gr00tN1d7Pipeline`
+- Deploy config: [`vllm_omni/deploy/Gr00tN1d7.yaml`](../../vllm_omni/deploy/Gr00tN1d7.yaml)
+- E2E test: [`tests/e2e/online_serving/test_gr00t_openpi.py`](../../tests/e2e/online_serving/test_gr00t_openpi.py)
+
+## Environment
+
+- OS: Linux
+- Python: 3.11+
+- Driver / runtime: NVIDIA CUDA
+- Hardware: 1 NVIDIA GPU. The upstream model card lists 16 GB+ VRAM for inference (e.g. RTX 4090, L40, H100); in practice this serving path uses ~6 GiB peak (bf16, TP=1, `max_num_seqs: 1`).
+- vLLM-Omni version or commit: use versions from your current checkout
+
+## Start server
+
+From repository root:
+
+```bash
+vllm serve nvidia/GR00T-N1.7-3B \
+  --omni \
+  --host 127.0.0.1 \
+  --port 8000 \
+  --served-model-name gr00t-n1d7 \
+  --stage-configs-path vllm_omni/deploy/Gr00tN1d7.yaml
+```
+
+Notes:
+
+- Only `max_num_seqs: 1` is supported (configured in the deploy YAML). The
+  policy is markovian (`state_history_length=1`, `reset()` returns `{}`) — the
+  reason for the cap is that `Gr00tPolicy` does its own per-sample
+  (un)batching inside `get_action` and is not integrated with vLLM's
+  continuous batching path.
+- This pipeline is a thin wrapper around the upstream HF policy
+  (`AutoModel.from_pretrained` + `enforce_eager`, `tensor_parallel_size=1`,
+  pipeline `load_weights` is a no-op). The standard diffusion accelerators
+  (SP, CFG, TeaCache, VAE tiling) do not transfer to a flow-matching action
+  policy at batch size 1, so a native-kernel port is intentionally out of
+  scope for this recipe — the value is the OpenPI serving integration, not
+  kernel-level acceleration.
+- The WebSocket endpoint is `ws://127.0.0.1:8000/v1/realtime/robot/openpi`.
+  The server handshake message (first frame after connect) is a msgpack-encoded
+  dict with `action_horizon`, `action_keys`, `embodiment_tag`, and
+  `needs_session_id`.
+
+## Verification
+
+```python
+from tests.gr00t.openpi_client_helper import run_policy_session, validate_session_result
+validate_session_result(run_policy_session(host="127.0.0.1", port=8000))
+```
+
+Or run the e2e test suite:
+
+```bash
+python -m pytest tests/e2e/online_serving/test_gr00t_openpi.py -v
+```
+
+The test sends a synthetic two-frame DROID observation and checks:
+
+- GR00T metadata contract: `image_resolution`, `action_horizon`, `action_keys`, `embodiment_tag`
+- Action shapes: `eef_9d (1,40,9)`, `gripper_position (1,40,1)`, `joint_position (1,40,7)`
+- All action values are finite float32
+- Reset response is `"reset successful"`
+
+## Notes
+
+- **Do not change the model-specific `policy_server_config` values.** `action_horizon`,
+  `action_keys`, and `supported_embodiments` are fixed by the GR00T-N1.7 checkpoint and
+  are validated against the loaded policy at startup (the server refuses to start on a
+  mismatch). Only `image_resolution` and `needs_session_id` are deployment knobs.
+- To switch embodiment, edit `embodiment_tag` under both `model_config` and
+  `policy_server_config` in `vllm_omni/deploy/Gr00tN1d7.yaml`. Supported values:
+  `OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT` (default), `XDOF`, `XDOF_SUBTASK`,
+  `REAL_G1`, `REAL_R1_PRO_SHARPA`, `LIBERO_PANDA`, `SIMPLER_ENV_GOOGLE`,
+  `SIMPLER_ENV_WIDOWX`.
+- GR00T weights are loaded directly by `Gr00tPolicy` via `AutoModel.from_pretrained`;
+  the pipeline's `load_weights` is intentionally a no-op.
diff --git a/recipes/README.md b/recipes/README.md
index 1f5b5be5765..e3ee2bc5639 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -12,6 +12,13 @@ that is aligned with
 that external repository as a reference for structure rather than the place to
 add files for this repo. Use one Markdown file per model family by default.
 
+> **Note:** The vLLM-Omni TTS recipes (12 models across 8 providers) are also
+> published on the rendered recipes site at
+> [recipes.vllm.ai](https://recipes.vllm.ai) (see
+> [vllm-project/recipes#554](https://github.com/vllm-project/recipes/pull/554)).
+> Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`,
+> e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b).
+
 Example layout:
 
 ```text
@@ -30,11 +37,12 @@ recipes/
 | [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
 | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
 | [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
+| [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB |
 | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
-| [`inclusionAI/Ming-omni-tts-0.5B.md`](./inclusionAI/Ming-omni-tts-0.5B.md) | Offline + online dense Ming TTS/audio generation | 1x AMD `gfx942` (ROCm 7.2) |
+| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
@@ -42,6 +50,7 @@ recipes/
 | [`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 |
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
 | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
+| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
diff --git a/recipes/mistralai/Voxtral-TTS.md b/recipes/mistralai/Voxtral-TTS.md
index 595d5c29e5f..aae48a6895a 100644
--- a/recipes/mistralai/Voxtral-TTS.md
+++ b/recipes/mistralai/Voxtral-TTS.md
@@ -105,3 +105,45 @@ voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../exam
 ) or adding a `ref_audio` request parameter will fail
 because the public checkpoint does not provide the encoder weights needed to
 convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.`
+
+### 1x NVIDIA RTX A6000 48GB
+
+#### Environment
+
+- OS: Ubuntu 22.04
+- Python: 3.12
+- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
+- GPU: NVIDIA RTX A6000, 48 GB
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
+- Required package: `pip install mistral-common`
+
+#### Command
+
+```bash
+vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
+```
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "mistralai/Voxtral-4B-TTS-2603",
+    "input": "Artificial intelligence is transforming the way we interact with technology.",
+    "voice": "neutral_female",
+    "response_format": "wav"
+  }' -o voxtral_test.wav
+```
+
+Supported voice presets: `ar_male`, `casual_female`, `casual_male`, `cheerful_female`,
+`de_female`, `de_male`, `es_female`, `es_male`, `fr_female`, `fr_male`, `hi_female`,
+`hi_male`, `it_female`, `it_male`, `neutral_female`, `neutral_male`, `nl_female`,
+`nl_male`, `pt_female`, `pt_male`.
+
+#### Notes
+
+- Memory usage: Peak VRAM ~40.5 GiB during inference.
+- Generation time: ~8 seconds for ~200 words of English text.
+- Audio quality is high; mono 24 kHz output.
diff --git a/requirements/common.txt b/requirements/common.txt
index 469563d3ab0..6d9dd63ba40 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -28,3 +28,5 @@ aenum==3.1.16
 pyzmq>=25.0.0
 janus>=1.0.0
 msgpack>=1.0.0
+pydantic>=2.1.0
+pydantic-settings>=2.0.0
diff --git a/tests/config/test_server_settings.py b/tests/config/test_server_settings.py
new file mode 100644
index 00000000000..53c15524ce6
--- /dev/null
+++ b/tests/config/test_server_settings.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for ServerSettings env-var parameterization."""
+
+import pytest
+
+from vllm_omni.config.server_settings import ServerSettings
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+ENV_VARS = [
+    "VLLM_OMNI_SERVER_STORAGE__PATH",
+    "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY",
+    "VLLM_OMNI_SERVER_STORAGE__FILE_TTL",
+    "VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL",
+    "VLLM_OMNI_STORAGE_PATH",
+    "VLLM_OMNI_STORAGE_MAX_CONCURRENCY",
+]
+
+
+@pytest.fixture(autouse=True)
+def clean_env(monkeypatch):
+    for name in ENV_VARS:
+        monkeypatch.delenv(name, raising=False)
+
+
+def test_defaults_when_no_env_set():
+    storage = ServerSettings().storage
+
+    assert storage.path == "/tmp/storage"
+    assert storage.file_concurrency == 4
+    assert storage.file_ttl is None
+    assert storage.ttl_sweep_interval is None
+
+
+def test_nested_env_vars_populate_storage(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__PATH", "/var/storage")
+    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", "8")
+    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_TTL", "120")
+    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL", "30")
+
+    storage = ServerSettings().storage
+
+    assert storage.path == "/var/storage"
+    assert storage.file_concurrency == 8
+    assert storage.file_ttl == 120
+    assert storage.ttl_sweep_interval == 30
+
+
+def test_ttl_sweep_interval_defaults_when_ttl_set(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_TTL", "60")
+
+    storage = ServerSettings().storage
+
+    assert storage.file_ttl == 60
+    assert storage.ttl_sweep_interval == 300
+
+
+def test_deprecated_path_alias_populates_and_warns(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_STORAGE_PATH", "/legacy/path")
+
+    with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_PATH"):
+        storage = ServerSettings().storage
+
+    assert storage.path == "/legacy/path"
+
+
+def test_deprecated_concurrency_alias_populates_and_warns(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "16")
+
+    with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_MAX_CONCURRENCY"):
+        storage = ServerSettings().storage
+
+    assert storage.file_concurrency == 16
+
+
+def test_new_env_var_wins_over_deprecated_alias(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__PATH", "/new/path")
+    monkeypatch.setenv("VLLM_OMNI_STORAGE_PATH", "/legacy/path")
+
+    with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_PATH"):
+        storage = ServerSettings().storage
+
+    assert storage.path == "/new/path"
diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py
index 021473ca5cc..f5981b92748 100644
--- a/tests/diffusion/cache/test_cache_backends.py
+++ b/tests/diffusion/cache/test_cache_backends.py
@@ -13,10 +13,12 @@
 
 from unittest.mock import Mock, patch
 
+import cache_dit
 import pytest
+from cache_dit import ForwardPattern
 
 from vllm_omni.diffusion.cache.cache_dit_backend import (
-    CUSTOM_DIT_ENABLERS,
+    CacheDiTAdapterConfig,
     CacheDiTBackend,
 )
 from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
@@ -45,13 +47,20 @@ def test_init_with_config_object(self):
         assert backend.config.Fn_compute_blocks == 4
         assert backend.enabled is False
 
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-    def test_enable_single_transformer(self, mock_cache_dit):
+    def test_enable_single_transformer(self, mock_cache_dit, mock_block_adapter):
         """Test enabling cache-dit on single-transformer pipeline."""
         # Mock pipeline
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DiTPipeline"
         mock_transformer = Mock()
+        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "layers": ForwardPattern.Pattern_0,
+            },
+        )
+
         mock_pipeline.transformer = mock_transformer
 
         # Mock cache_dit functions
@@ -66,14 +75,20 @@ def test_enable_single_transformer(self, mock_cache_dit):
         assert backend._refresh_func is not None
         mock_cache_dit.enable_cache.assert_called_once()
 
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-    def test_refresh(self, mock_cache_dit):
+    def test_refresh(self, mock_cache_dit, mock_block_adapter):
         """Test refreshing cache context with SCM mask policy updates when num_inference_steps changes."""
         # Mock pipeline
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DiTPipeline"
         mock_transformer = Mock()
         mock_pipeline.transformer = mock_transformer
+        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "layers": ForwardPattern.Pattern_0,
+            },
+        )
 
         # Mock cache_dit functions
         mock_cache_dit.enable_cache = Mock()
@@ -124,23 +139,24 @@ def test_refresh(self, mock_cache_dit):
         assert "cache_config" in call_args[1]
         assert backend._last_num_inference_steps == 100
 
-    def test_hunyuan_custom_enabler_registered(self):
-        """Test HunyuanImage3 custom cache-dit enabler is registered."""
-        assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS
-
-    def test_dreamid_custom_enabler_registered(self):
-        """Test DreamIDOmni custom cache-dit enabler is registered."""
-        assert "DreamIDOmniPipeline" in CUSTOM_DIT_ENABLERS
-
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
-        """Test HunyuanImage3 custom enabler uses pipeline.model for cache enable/refresh."""
+        """Test HunyuanImage3 uses pipeline.transformer for cache enable/refresh.
+
+        NOTE: HunyuanImage3 no longer has a custom enabler, so this tests against the generic path.
+        """
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
         mock_pipeline.model = Mock()
         mock_pipeline.model.layers = Mock()
-
+        mock_pipeline.model._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "layers": ForwardPattern.Pattern_4,
+            },
+        )
+        # NOTE: pipe.transformer is pipe.model in HunyuanImage3Pipelines
+        mock_pipeline.transformer = mock_pipeline.model
         mock_cache_dit.enable_cache = Mock()
         mock_cache_dit.refresh_context = Mock()
 
@@ -152,9 +168,9 @@ def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mo
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.model
-        assert adapter_kwargs["blocks"] is mock_pipeline.model.layers
-        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_4
-        assert len(adapter_kwargs["params_modifiers"]) == 1
+        assert len(adapter_kwargs["blocks"]) == 1
+        assert adapter_kwargs["blocks"][0] == mock_pipeline.model.layers
+        assert adapter_kwargs["forward_pattern"][0] == ForwardPattern.Pattern_4
         mock_cache_dit.enable_cache.assert_called_once()
 
         backend.refresh(mock_pipeline, num_inference_steps=12)
@@ -163,24 +179,23 @@ def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mo
         assert call_args[0][0] is mock_pipeline.model
         assert call_args[1]["num_inference_steps"] == 12
 
-    def test_enable_hunyuan_pipeline_requires_model_layers(self):
-        """Test HunyuanImage3 enabler fails with a formatted pipeline class name."""
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
-        mock_pipeline.model = Mock(spec=[])
-
-        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
-
-        with pytest.raises(ValueError, match="HunyuanImage3Pipeline"):
-            backend.enable(mock_pipeline)
-
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
+        """Test DreamID uses pipeline.transformer for cache enable/refresh.
+
+        NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
+        """
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
         mock_pipeline.transformer = Mock()
         mock_pipeline.transformer.fused_blocks = Mock()
+        mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "fused_blocks": ForwardPattern.Pattern_0,
+            },
+            has_separate_cfg=True,
+        )
 
         mock_cache_dit.enable_cache = Mock()
         mock_cache_dit.refresh_context = Mock()
@@ -193,10 +208,10 @@ def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_bl
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.transformer
-        assert adapter_kwargs["blocks"] is mock_pipeline.transformer.fused_blocks
-        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_0
+        assert len(adapter_kwargs["blocks"]) == 1
+        assert adapter_kwargs["blocks"][0] == mock_pipeline.transformer.fused_blocks
+        assert adapter_kwargs["forward_pattern"][0] == ForwardPattern.Pattern_0
         assert adapter_kwargs["has_separate_cfg"] is True
-        assert len(adapter_kwargs["params_modifiers"]) == 1
         mock_cache_dit.enable_cache.assert_called_once()
 
         backend.refresh(mock_pipeline, num_inference_steps=12)
@@ -205,15 +220,43 @@ def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_bl
         assert call_args[0][0] is mock_pipeline.transformer
         assert call_args[1]["num_inference_steps"] == 12
 
-    def test_enable_dreamid_pipeline_requires_fused_blocks(self):
+    @pytest.mark.parametrize("num_inference_steps", [1, 7])
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    def test_refresh_scm_bypassed_for_unsupported_step_counts(
+        self, mock_cache_dit, mock_block_adapter, num_inference_steps
+    ):
+        """Ensure SCM is bypassed when num_inference_steps < 8 and not in (4, 6),
+        because cache_dit.steps_mask() raises for unsupported step count.
+        For these cases, we fall back to the non-SCM path to avoid crashing.
+        """
         mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
-        mock_pipeline.transformer = Mock(spec=[])
+        mock_pipeline.__class__.__name__ = "DiTPipeline"
+        mock_transformer = Mock()
+        mock_pipeline.transformer = mock_transformer
+        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "layers": ForwardPattern.Pattern_0,
+            },
+        )
 
-        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        mock_cache_dit.enable_cache = Mock()
+        mock_cache_dit.refresh_context = Mock()
+        mock_cache_dit.steps_mask = cache_dit.steps_mask
+
+        # Create a cache config with an scm policy & enable it
+        config = DiffusionCacheConfig(scm_steps_mask_policy="fast")
+        backend = CacheDiTBackend(config)
+        backend.enable(mock_pipeline)
 
-        with pytest.raises(ValueError, match="DreamIDOmniPipeline"):
-            backend.enable(mock_pipeline)
+        backend.refresh(mock_pipeline, num_inference_steps=num_inference_steps)
+
+        mock_cache_dit.refresh_context.assert_called_once()
+        call_args = mock_cache_dit.refresh_context.call_args
+        # Ensure that we properly guard, i.e., cache config is filtered
+        assert call_args[0][0] == mock_transformer
+        assert call_args[1]["num_inference_steps"] == num_inference_steps
+        assert "cache_config" not in call_args[1]
 
 
 class TestTeaCacheBackend:
diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py
index 146084a508c..0a05aeb761f 100644
--- a/tests/diffusion/cache/test_cache_dit.py
+++ b/tests/diffusion/cache/test_cache_dit.py
@@ -5,6 +5,7 @@
 Model specific tests for CacheDiT enablement.
 """
 
+import sys
 from unittest.mock import Mock, patch
 
 import pytest
@@ -12,36 +13,49 @@
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
 
 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
+from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
 from vllm_omni.platforms import current_omni_platform
 
+# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
+# the class even though the dependency may not be set up correctly; this is ok for these
+# tests because we just inspect it and never initialize the model.
+for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
+    sys.modules.setdefault(mod, Mock())
+# isort: split
+from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel  # noqa: E402
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-SEPARATE_CFG_ENABLERS = [
-    cd_backend.enable_cache_for_dreamid_omni,
-    cd_backend.enable_cache_for_ltx2,
-    cd_backend.enable_cache_for_helios,
-    cd_backend.enable_cache_for_wan22,
-    cd_backend.enable_cache_for_longcat_image,
-    cd_backend.enable_cache_for_cosmos3,
+SEPARATE_CFG_TRANSFORMERS = [
+    DreamIdOmniModel,
+    LTX2VideoTransformer3DModel,
+    HeliosTransformer3DModel,
+    LongCatImageTransformer2DModel,
+    Cosmos3VFMTransformer,
 ]
 
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
 
 
-@pytest.mark.parametrize("enabler", SEPARATE_CFG_ENABLERS)
+@pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS)
+def test_cache_dit_configs_have_separate_cfg(transformer_model):
+    """Check models with separate CFG set has_separate_cfg=True in their configs."""
+    assert hasattr(transformer_model, "_cache_dit_adapter_config")
+    assert isinstance(transformer_model._cache_dit_adapter_config, CacheDiTAdapterConfig)
+    assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
+
+
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler):
-    """Ensure that custom enablers for models with separate CFG pass
-    the param through to cache_dit correctly.
-
-    Regression test for: https://github.com/vllm-project/vllm-omni/pull/2860
-    """
+def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
+    """Ensure that Wan22, which has a custom enabler, setts custom CFG correctly."""
     mock_pipeline = Mock()
-    enabler(mock_pipeline, SAMPLE_CACHE_CONFIG)
+    cd_backend.enable_cache_for_wan22(mock_pipeline, SAMPLE_CACHE_CONFIG)
 
     mock_cache_dit.enable_cache.assert_called_once()
     adapter_kwargs = mock_block_adapter.call_args.kwargs
@@ -55,6 +69,7 @@ def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
     mock_pipeline = Mock()
     gen_layers = object()
     mock_pipeline.transformer.gen_layers = gen_layers
+    mock_pipeline.transformer._cache_dit_adapter_config = Cosmos3VFMTransformer._cache_dit_adapter_config
 
     cd_backend.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG)
 
diff --git a/tests/diffusion/cache/test_stepcache.py b/tests/diffusion/cache/test_stepcache.py
new file mode 100644
index 00000000000..1fc3bbb5922
--- /dev/null
+++ b/tests/diffusion/cache/test_stepcache.py
@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for StepCacheBackend and StepCacheState."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.cache.stepcache import (
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheBackend,
+    StepCacheConfig,
+    StepCacheState,
+    is_stepcache_active,
+)
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_prev_predictions(sim: float, dim: int = 8) -> list[tuple[torch.Tensor]]:
+    v_prev = torch.randn(1, dim)
+    if sim >= 0.99:
+        v_last = v_prev.clone()
+    else:
+        v_last = torch.randn(1, dim)
+    return [(v_prev,), (v_last,)]
+
+
+class TestStepCacheConfig:
+    def test_from_diffusion_cache_config_respects_step_cache_dit_enabled(self):
+        enabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=True))
+        disabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=False))
+        assert enabled.enabled is True
+        assert disabled.enabled is False
+
+
+class TestStepCacheState:
+    def test_disabled_always_runs(self):
+        state = StepCacheState(StepCacheConfig(enabled=False))
+        assert state.should_run_step(_make_prev_predictions(0.99)) is True
+
+    def test_insufficient_history_always_runs(self):
+        state = StepCacheState(StepCacheConfig(min_history_steps=2))
+        assert state.should_run_step([(torch.randn(1, 4),)]) is True
+
+    def test_high_similarity_skips_with_countdown(self):
+        state = StepCacheState(
+            StepCacheConfig(
+                sim_thresholds=(0.95,),
+                skip_countdowns=(2,),
+            )
+        )
+        preds = _make_prev_predictions(0.99)
+        assert state.should_run_step(preds) is False
+        assert state.skip_countdown == 2
+        assert state.should_run_step(preds) is False
+        assert state.should_run_step(preds) is True
+
+    def test_low_similarity_runs(self):
+        state = StepCacheState(StepCacheConfig())
+        preds = _make_prev_predictions(0.5)
+        assert state.should_run_step(preds) is True
+
+    def test_trim_history(self):
+        state = StepCacheState(StepCacheConfig(max_history=2))
+        history: list[tuple[torch.Tensor]] = [(torch.randn(1, 2),) for _ in range(4)]
+        state.trim_history(history)
+        assert len(history) == 2
+
+
+class TestStepCacheBackend:
+    def test_enable_dreamzero_pipeline(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamZeroPipeline"
+
+        backend = StepCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        cfg = getattr(mock_pipeline, STEP_CACHE_CONFIG_ATTR)
+        assert isinstance(cfg, StepCacheConfig)
+        assert is_stepcache_active(mock_pipeline)
+        state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR)
+        assert isinstance(state, StepCacheState)
+
+    def test_refresh_resets_state(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamZeroPipeline"
+        backend = StepCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+        state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR)
+        state.skip_countdown = 3
+        backend.refresh(mock_pipeline, num_inference_steps=16, verbose=False)
+        assert state.skip_countdown == 0
+
+    def test_enable_unsupported_pipeline_raises(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "QwenImagePipeline"
+        backend = StepCacheBackend(DiffusionCacheConfig())
+
+        with pytest.raises(ValueError, match="step_cache backend does not support"):
+            backend.enable(mock_pipeline)
+
+    def test_selector_returns_backend(self):
+        backend = get_cache_backend("step_cache", {"velocity_sim_thresholds": [0.9]})
+        assert isinstance(backend, StepCacheBackend)
+        assert backend.config.velocity_sim_thresholds == [0.9]
+
+    def test_legacy_backend_aliases(self):
+        assert isinstance(get_cache_backend("stepcache", {}), StepCacheBackend)
+        assert isinstance(get_cache_backend("step_cache_dit", {}), StepCacheBackend)
diff --git a/tests/diffusion/models/gr00t/__init__.py b/tests/diffusion/models/gr00t/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/diffusion/models/gr00t/test_pipeline.py b/tests/diffusion/models/gr00t/test_pipeline.py
new file mode 100644
index 00000000000..72118c3f93d
--- /dev/null
+++ b/tests/diffusion/models/gr00t/test_pipeline.py
@@ -0,0 +1,141 @@
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from vllm_omni.diffusion.models.gr00t import pipeline_gr00t
+from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeGr00tPolicy:
+    instances = []
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+        self.reset_calls = 0
+        self.seen_obs = None
+        self.embodiment_tag = SimpleNamespace(value="fake_embodiment")
+        self.language_key = "annotation.language.language_instruction"
+        self.modality_configs = {
+            "action": SimpleNamespace(
+                delta_indices=[0, 1],
+                modality_keys=["arm", "gripper"],
+            )
+        }
+        self.processor = SimpleNamespace(
+            state_action_processor=SimpleNamespace(
+                norm_params={
+                    "fake_embodiment": {
+                        "action": {
+                            "arm": {"dim": np.array(2)},
+                            "gripper": {"dim": np.array(1)},
+                        }
+                    }
+                }
+            )
+        )
+        FakeGr00tPolicy.instances.append(self)
+
+    def get_action(self, obs):
+        self.seen_obs = obs
+        return {
+            "arm": np.array([[[1.0, 2.0]]], dtype=np.float64),
+            "gripper": [[[3.0]]],
+        }, {"latency_ms": 1.0}
+
+    def reset(self):
+        self.reset_calls += 1
+        return {"reset": True}
+
+
+@pytest.fixture(autouse=True)
+def fake_gr00t_policy(monkeypatch):
+    FakeGr00tPolicy.instances.clear()
+    monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", FakeGr00tPolicy)
+
+
+def _pipeline():
+    od_config = SimpleNamespace(
+        model="nvidia/GR00T-N1.7-3B",
+        model_config={
+            "embodiment_tag": "LIBERO_PANDA",
+            "strict": False,
+        },
+        custom_pipeline_args={},
+    )
+    return Gr00tN1d7Pipeline(od_config=od_config)
+
+
+def test_pipeline_initializes_local_policy():
+    pipeline = _pipeline()
+
+    policy = FakeGr00tPolicy.instances[0]
+    assert policy.kwargs["model_path"] == "nvidia/GR00T-N1.7-3B"
+    assert policy.kwargs["embodiment_tag"] == "LIBERO_PANDA"
+    assert policy.kwargs["strict"] is False
+    assert pipeline.weights_sources == ()
+    assert pipeline.load_weights(iter(())) == set()
+
+
+def test_forward_returns_dict_actions_in_output():
+    pipeline = _pipeline()
+    req = OmniDiffusionRequest(
+        prompts=["pick"],
+        request_id="req",
+        sampling_params=OmniDiffusionSamplingParams(
+            extra_args={
+                "robot_obs": {
+                    "images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)},
+                    "state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)},
+                    "prompt": "pick the cube",
+                    "session_id": "session-a",
+                },
+                "reset": True,
+            }
+        ),
+    )
+
+    output = pipeline.forward(req)
+
+    assert output.error is None
+    actions = output.output["actions"]
+    assert set(actions) == {"arm", "gripper"}
+    assert actions["arm"].dtype == np.float32
+    np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32))
+    policy = FakeGr00tPolicy.instances[0]
+    assert "video" in policy.seen_obs
+    assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
+    assert "images" not in policy.seen_obs
+    assert "prompt" not in policy.seen_obs
+    assert "session_id" not in policy.seen_obs
+    assert policy.reset_calls == 1
+
+
+def test_dummy_warmup_returns_shape_correct_zero_actions():
+    pipeline = _pipeline()
+    req = OmniDiffusionRequest(
+        prompts=["dummy run"],
+        request_id="dummy_req_id",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+    )
+
+    output = pipeline.forward(req)
+
+    assert output.error is None
+    actions = output.output["actions"]
+    assert set(actions) == {"arm", "gripper"}
+    assert actions["arm"].shape == (1, 2, 2)
+    assert actions["gripper"].shape == (1, 2, 1)
+    assert not actions["arm"].any()
+    assert FakeGr00tPolicy.instances[0].seen_obs is None
+
+
+def test_reset_delegates_to_policy():
+    pipeline = _pipeline()
+
+    assert pipeline.reset() == {"reset": True}
+    assert FakeGr00tPolicy.instances[0].reset_calls == 1
diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
index bf6a7563623..9f38f097acd 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
@@ -692,16 +692,27 @@ def test_post_process_funcs_registered(self):
             assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS"
             assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func"
 
-    def test_cache_dit_enablers_registered(self):
-        """Pipeline variants must be registered in CUSTOM_DIT_ENABLERS."""
+    def test_cache_dit_for_ltx2_does_not_have_custom_enablers_registered(self):
+        """Pipeline variants are *not* registered in CUSTOM_DIT_ENABLERS."""
         from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
 
+        # NOTE: We used to have custom enablers for this model, but refactored to handle
+        # it more generically. Now we only need to ensure it has git cache adapter config.
         expected = [
             "LTX23Pipeline",
             "LTX23ImageToVideoPipeline",
         ]
         for name in expected:
-            assert name in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS"
+            assert name not in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS"
+
+    def test_ltx2_transformer_has_dit_cache_config(self):
+        """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
+        from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+        from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+
+        adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
+        assert isinstance(adapter_config, CacheDiTAdapterConfig)
+        assert adapter_config.has_separate_cfg
 
 
 class TestVocoderSampleRateDetection:
diff --git a/tests/diffusion/quantization/test_hsdp_fp8.py b/tests/diffusion/quantization/test_hsdp_fp8.py
new file mode 100644
index 00000000000..97279acf2f5
--- /dev/null
+++ b/tests/diffusion/quantization/test_hsdp_fp8.py
@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HSDP/FSDP2 compatibility with online FP8 quantization."""
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.quantization.hsdp_fp8 import (
+    _build_transposed_get_layer_params,
+    prepare_fp8_layers_for_fsdp,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+# --- shared test doubles ---
+
+
+class _ToyKernel:
+    """Minimal kernel that mirrors ``FP8ScaledMMLinearKernel``."""
+
+    def __init__(self):
+        self.layer_param_names = ("weight", "weight_scale", "input_scale", "input_scale_ub")
+
+    def _get_layer_params(self, layer):
+        w, w_s, x_s, x_s_ub = self.layer_param_names
+        return (
+            getattr(layer, w),
+            getattr(layer, w_s),
+            getattr(layer, x_s, None),
+            getattr(layer, x_s_ub, None),
+        )
+
+
+# --- helper to build a toy module simulating online-FP8 post-load state ---
+
+
+def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32):
+    """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
+    ``quant_method`` is an ``Fp8OnlineLinearMethod`` with a ``_ToyKernel``.
+    """
+    from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
+
+    # Create the quant-method instance without calling __init__ (avoids
+    # heavyweight upstream config dependencies).  isinstance() still works.
+    qm = object.__new__(Fp8OnlineLinearMethod)
+    qm.fp8_linear = _ToyKernel()
+
+    # randn does not support float8 on CPU; zeros does and the test only
+    # cares about shape/stride, not the actual weight values.
+    qweight = torch.zeros(out_features, in_features, dtype=torch.float8_e4m3fn)
+    weight = nn.Parameter(qweight.t())  # (in, out) non-contiguous column-major view
+
+    module = nn.Module()
+    module.quant_method = qm
+    module.weight = weight
+    module.weight_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
+    # Set input_scale / input_scale_ub to verify _get_layer_params
+    # passes them through unchanged.
+    module.input_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
+    module.input_scale_ub = nn.Parameter(torch.ones(1, dtype=torch.float32))
+
+    return module
+
+
+# --- tests ---
+
+
+def test_transposed_get_layer_params():
+    import types
+
+    kernel = _ToyKernel()
+
+    # original bound method + patch + re-bind, matching prepare_fp8_layers_for_fsdp
+    original_bound = kernel._get_layer_params
+    patched_func = _build_transposed_get_layer_params(original_bound)
+    kernel._get_layer_params = types.MethodType(patched_func, kernel)
+
+    module = _make_fp8_toy_module(16, 32)
+    # Reproduce the storage rewrite from prepare_fp8_layers_for_fsdp:
+    # weight was qweight.t()  →  .t() recovers qweight (row-major contiguous).
+    module.weight = nn.Parameter(module.weight.data.t(), requires_grad=False)
+    assert module.weight.is_contiguous(), "qweight should be row-major contiguous"
+
+    w, w_s, x_s, x_s_ub = kernel._get_layer_params(module)
+
+    # patched method transposes row-major → column-major for Cutlass
+    assert w.shape == (32, 16)
+    assert not w.is_contiguous()
+    # scales are passed through unchanged
+    assert w_s is module.weight_scale
+    assert x_s is module.input_scale
+    assert x_s_ub is module.input_scale_ub
+
+
+def test_rewrites_non_contiguous_weight():
+    module = _make_fp8_toy_module(16, 32)
+    assert not module.weight.is_contiguous()
+
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+    assert n_rewritten_layers == 1
+    assert module.weight.is_contiguous()
+    assert module.weight.shape == (16, 32)
+
+
+def test_patched_get_layer_params_returns_transposed_view():
+    module = _make_fp8_toy_module(16, 32)
+    prepare_fp8_layers_for_fsdp(module)
+
+    kernel = module.quant_method.fp8_linear
+    w, w_s, x_s, x_s_ub = kernel._get_layer_params(module)
+
+    # Cutlass expects column-major B: (in, out), stride (1, in)
+    assert w.shape == (32, 16)
+    assert not w.is_contiguous()
+    assert w.stride() == (1, 32)
+
+    # scales pass through unchanged
+    assert w_s is module.weight_scale
+    assert x_s is module.input_scale
+    assert x_s_ub is module.input_scale_ub
+
+
+def test_contiguous_weights_skipped():
+    module = _make_fp8_toy_module(16, 32)
+    # make weight contiguous before calling
+    module.weight = nn.Parameter(module.weight.data.t().contiguous())
+    assert module.weight.is_contiguous()
+
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+    assert n_rewritten_layers == 0
+
+
+def test_kernel_patched_only_once_per_instance():
+    module = _make_fp8_toy_module(16, 32)
+    kernel = module.quant_method.fp8_linear
+    original = kernel._get_layer_params
+
+    prepare_fp8_layers_for_fsdp(module)
+    first_patched = kernel._get_layer_params
+
+    # Second call with same kernel id must not stack patches
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+    assert n_rewritten_layers == 0
+    assert kernel._get_layer_params is first_patched
+    assert kernel._get_layer_params is not original
diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py
index 8ad6c273416..fb45cd41d0d 100644
--- a/tests/diffusion/test_diffusion_worker.py
+++ b/tests/diffusion/test_diffusion_worker.py
@@ -23,6 +23,15 @@
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
 
+def patch_cumem_allocator(mocker: MockerFixture):
+    mock_allocator_class = mocker.Mock()
+    mocker.patch(
+        "vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class",
+        return_value=mock_allocator_class,
+    )
+    return mock_allocator_class
+
+
 @pytest.fixture
 def mock_od_config(mocker: MockerFixture):
     """Create a mock OmniDiffusionConfig."""
@@ -118,7 +127,7 @@ def setup_allocator(self, mocker: MockerFixture):
         """
         Unified interception of Allocators, and provision of default security values.
         """
-        self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        self.mock_allocator_class = patch_cumem_allocator(mocker)
         self.mock_allocator = mocker.Mock()
         self.mock_allocator_class.get_instance.return_value = self.mock_allocator
         self.mock_allocator.get_current_usage.return_value = 4 * 1024**3
@@ -126,7 +135,7 @@ def setup_allocator(self, mocker: MockerFixture):
 
     def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
-        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mock_allocator_class = patch_cumem_allocator(mocker)
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
         mock_platform.get_device_total_memory.return_value = 80 * 1024**3
@@ -158,7 +167,7 @@ def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
 
     def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 2 (offload all, save buffers)."""
-        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mock_allocator_class = patch_cumem_allocator(mocker)
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3]
         mock_platform.get_device_total_memory.return_value = 80 * 1024**3
@@ -201,7 +210,7 @@ def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
 
     def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
         """Test that sleep validates memory was actually freed."""
-        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mock_allocator_class = patch_cumem_allocator(mocker)
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_platform.get_free_memory.return_value = 10 * 1024**3
         mock_platform.get_device_total_memory.return_value = 80 * 1024**3
@@ -226,7 +235,7 @@ def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_wor
     def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep uses device-scoped fallback when NVML is unavailable."""
 
-        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mock_allocator_class = patch_cumem_allocator(mocker)
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
         mock_get_process_memory.side_effect = [None, None]
@@ -252,7 +261,7 @@ class TestDiffusionWorkerWakeUp:
 
     def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up without saved buffers (level 1 sleep)."""
-        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mock_allocator_class = patch_cumem_allocator(mocker)
 
         # Setup allocator mock
         mock_allocator = mocker.Mock()
@@ -272,7 +281,7 @@ def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):
 
     def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up with saved buffers (level 2 sleep)."""
-        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mock_allocator_class = patch_cumem_allocator(mocker)
 
         # Setup allocator mock
         mock_allocator = mocker.Mock()
@@ -317,7 +326,7 @@ def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
 
     def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up only restores buffers that were saved."""
-        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mock_allocator_class = patch_cumem_allocator(mocker)
 
         # Setup allocator mock
         mock_allocator = mocker.Mock()
diff --git a/tests/dreamzero/test_crossattn_cache.py b/tests/dreamzero/test_crossattn_cache.py
new file mode 100644
index 00000000000..94361c090a0
--- /dev/null
+++ b/tests/dreamzero/test_crossattn_cache.py
@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests that DreamZero's cross-attention KV cache is actually init'd and reused.
+
+DreamZero's I2V cross-attention caches the text k/v and image k_img/v_img in a
+per-layer ``crossattn_cache`` dict: they are computed on the first forward of a
+session (``is_init`` False -> True) and reused on every later step, because the
+text/image context is session-invariant. Only the query (which depends on the
+per-step hidden state) is recomputed.
+
+These tests lock in that ``WanI2VCrossAttention.forward`` initialises the cache
+on the first call and reuses it (not recompute) on later calls, and recomputes
+every call when no cache is provided.
+
+CPU-only: ``WanI2VCrossAttention.__init__`` needs a tensor-parallel group (it
+builds ColumnParallel/RowParallel linears), so we bypass ``__init__`` and stub
+the projections/attention to count calls and exercise only the caching logic.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.causal_wan_model import WanI2VCrossAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_CLIP_LEN = 257  # image tokens the I2V cross-attn splits off the front of context
+
+
+def _make_cross_attn(n_heads: int = 2, head_dim: int = 4) -> WanI2VCrossAttention:
+    """A WanI2VCrossAttention whose submodules are call-counting stubs.
+
+    ``__init__`` is bypassed (it needs a TP group); we only wire what
+    ``forward`` touches. Each projection returns a [B, S, n*head_dim] tensor so
+    the subsequent ``unflatten(2, (n, head_dim))`` works, and is a MagicMock so
+    we can assert how many times it was invoked.
+    """
+    inner = n_heads * head_dim
+    module = WanI2VCrossAttention.__new__(WanI2VCrossAttention)
+    module.tp_num_heads = n_heads
+    module.head_dim = head_dim
+
+    def _proj() -> MagicMock:
+        return MagicMock(side_effect=lambda t: torch.zeros(t.shape[0], t.shape[1], inner))
+
+    module.q = _proj()
+    module.k = _proj()
+    module.v = _proj()
+    module.k_img = _proj()
+    module.v_img = _proj()
+    module.norm_q = lambda t: t
+    module.norm_k = lambda t: t
+    module.norm_k_img = lambda t: t
+    # attn returns the query unchanged ([B, S, n, d]); flatten(2) downstream is happy.
+    module.attn = lambda q, k, v: q
+    module.o = lambda t: t
+    return module
+
+
+def _fresh_cache() -> dict:
+    return {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None}
+
+
+def _inputs(inner: int = 8):
+    x = torch.randn(1, 3, inner)  # B=1, query_len=3
+    context = torch.randn(1, _CLIP_LEN + 5, inner)  # 257 image tokens + 5 text tokens
+    return x, context
+
+
+def test_cache_inits_on_first_call_and_reuses_on_second() -> None:
+    module = _make_cross_attn()
+    cache = _fresh_cache()
+    x, context = _inputs()
+
+    # First call: cache is initialised and every k/v projection runs once.
+    module.forward(x, context, crossattn_cache=cache)
+    assert cache["is_init"] is True
+    for key in ("k", "v", "k_img", "v_img"):
+        assert cache[key] is not None, f"{key} should be cached after first call"
+    assert module.k.call_count == 1
+    assert module.v.call_count == 1
+    assert module.k_img.call_count == 1
+    assert module.v_img.call_count == 1
+
+    cached = {key: cache[key] for key in ("k", "v", "k_img", "v_img")}
+
+    # Second call: k/v/k_img/v_img are read from cache, NOT recomputed.
+    module.forward(x, context, crossattn_cache=cache)
+    assert module.k.call_count == 1
+    assert module.v.call_count == 1
+    assert module.k_img.call_count == 1
+    assert module.v_img.call_count == 1
+    for key, tensor in cached.items():
+        assert cache[key] is tensor, f"{key} should be the same cached object"
+
+    # The query always depends on the per-step hidden state, so it is recomputed.
+    assert module.q.call_count == 2
+
+
+def test_no_cache_recomputes_every_call() -> None:
+    module = _make_cross_attn()
+    x, context = _inputs()
+
+    module.forward(x, context, crossattn_cache=None)
+    module.forward(x, context, crossattn_cache=None)
+
+    # Without a cache, the image/text k/v are recomputed on every call.
+    assert module.k_img.call_count == 2
+    assert module.v_img.call_count == 2
+    assert module.k.call_count == 2
+    assert module.v.call_count == 2
diff --git a/tests/dreamzero/test_fused_qk_rms_norm.py b/tests/dreamzero/test_fused_qk_rms_norm.py
new file mode 100644
index 00000000000..dbbb439400b
--- /dev/null
+++ b/tests/dreamzero/test_fused_qk_rms_norm.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for DreamZero's fused q/k qk-norm TP all-reduce.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the TP world size and the collective are mocked, so a single process
+simulates one rank (all-reduce is identity). This exercises the cat/split and
+count-scaling logic that a real multi-rank run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.causal_wan_model import (
+    DistributedRMSNorm,
+    fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.dreamzero.causal_wan_model"
+
+
+def _make_norms(dim: int) -> tuple[DistributedRMSNorm, DistributedRMSNorm]:
+    """Two norms with distinct random weights so a swapped q/k slice is caught."""
+    torch.manual_seed(0)
+    norm_q = DistributedRMSNorm(dim)
+    norm_k = DistributedRMSNorm(dim)
+    norm_q.weight.data.normal_()
+    norm_k.weight.data.normal_()
+    return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+    """Single-rank all-reduce: the reduction over one rank is the identity."""
+    return tensor
+
+
+def test_fused_matches_separate_tp1() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_matches_separate_tp_gt1() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    # tp_size=2 exercises the reduce branch and the `global_count *= tp_size`
+    # path in both the per-norm forward and the fused helper.
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce),
+    ):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_single_all_reduce_over_packed_pair() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch(
+            f"{_MODULE}.tensor_model_parallel_all_reduce",
+            side_effect=_identity_all_reduce,
+        ) as mock_ar,
+    ):
+        fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    # Exactly one collective, over a tensor whose last dim packs the q+k pair.
+    assert mock_ar.call_count == 1
+    packed = mock_ar.call_args[0][0]
+    assert packed.shape[-1] == 2
+
+
+def test_fused_identity_fallback_when_qk_norm_disabled() -> None:
+    # qk_norm=False makes the norms nn.Identity; the helper must fall back to
+    # plain passthrough without touching the collective.
+    norm_q, norm_k = torch.nn.Identity(), torch.nn.Identity()
+    q = torch.randn(2, 5, 64)
+    k = torch.randn(2, 5, 64)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce") as mock_ar:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, k, rtol=0, atol=0)
+    mock_ar.assert_not_called()
diff --git a/tests/dreamzero/test_qkv_fusion.py b/tests/dreamzero/test_qkv_fusion.py
new file mode 100644
index 00000000000..16b78693f69
--- /dev/null
+++ b/tests/dreamzero/test_qkv_fusion.py
@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence test for DreamZero self-attn QKV fusion.
+
+CausalWanSelfAttention now projects q/k/v with a single QKVParallelLinear and
+splits the result, instead of three separate ColumnParallelLinear. The DreamZero
+checkpoint still ships separate q/k/v weights, so load_weights routes each into
+the packed `qkv` param with its shard id. This test proves the chain
+"load separate q/k/v via shard id -> single GEMM -> split" reproduces the three
+independent projections, given identical weights.
+
+CPU-only: a single-rank (world_size=1) tensor-parallel group is initialized and
+the GEMM dispatcher is forced to the default (CPU-compatible) path, following
+tests/diffusion/layers/test_adalayernorm.py.
+"""
+
+import os
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Minimal world_size=1 TP group required to build QKVParallelLinear."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29507")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for CPU test tensors."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+def test_fused_qkv_load_and_split_matches_separate_projections():
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    torch.manual_seed(0)
+    dim, head_dim, num_heads = 16, 4, 4  # inner_dim == dim, no GQA
+
+    qkv = QKVParallelLinear(
+        hidden_size=dim,
+        head_size=head_dim,
+        total_num_heads=num_heads,
+        bias=True,
+    )
+
+    # Distinct random weight + bias per projection, loaded the exact way
+    # DreamZero.load_weights does: param.weight_loader(param, tensor, shard_id).
+    refs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+    for shard_id in ("q", "k", "v"):
+        w = torch.randn(dim, dim)
+        b = torch.randn(dim)
+        refs[shard_id] = (w, b)
+        qkv.weight.weight_loader(qkv.weight, w, shard_id)
+        qkv.bias.weight_loader(qkv.bias, b, shard_id)
+
+    x = torch.randn(2, 3, dim)
+    fused, _ = qkv(x)
+    q, k, v = fused.split([dim, dim, dim], dim=-1)
+
+    for got, shard_id in zip((q, k, v), ("q", "k", "v"), strict=True):
+        w, b = refs[shard_id]
+        ref = torch.nn.functional.linear(x, w, b)
+        torch.testing.assert_close(got, ref)
+
+
+def test_fused_qkv_shard_ids_are_not_transposed():
+    # Guard against q/k/v being routed to the wrong slice: give each projection a
+    # constant, distinct weight so a swapped shard id would surface immediately.
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    dim, head_dim, num_heads = 16, 4, 4
+    qkv = QKVParallelLinear(
+        hidden_size=dim,
+        head_size=head_dim,
+        total_num_heads=num_heads,
+        bias=True,
+    )
+
+    scale = {"q": 1.0, "k": 2.0, "v": 3.0}
+    for shard_id, s in scale.items():
+        qkv.weight.weight_loader(qkv.weight, torch.full((dim, dim), s), shard_id)
+        qkv.bias.weight_loader(qkv.bias, torch.zeros(dim), shard_id)
+
+    x = torch.ones(1, dim)
+    fused, _ = qkv(x)
+    q, k, v = fused.split([dim, dim, dim], dim=-1)
+
+    # F.linear(ones, full(s)) == s * dim per element.
+    torch.testing.assert_close(q, torch.full((1, dim), scale["q"] * dim))
+    torch.testing.assert_close(k, torch.full((1, dim), scale["k"] * dim))
+    torch.testing.assert_close(v, torch.full((1, dim), scale["v"] * dim))
diff --git a/tests/e2e/offline_inference/test_moss_tts.py b/tests/e2e/offline_inference/test_moss_tts.py
index 505db1af718..19a8ac6d6d7 100644
--- a/tests/e2e/offline_inference/test_moss_tts.py
+++ b/tests/e2e/offline_inference/test_moss_tts.py
@@ -1,77 +1,97 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""E2E offline inference tests for MOSS-TTS two-stage pipeline.
-
-Tests the MossTTSDelayModel path using MOSS-VoiceGenerator (1.7B) — the
-smallest MossTTSDelayModel variant — so the test runs on a single A10G or L4
-without requiring an 80 GB GPU.
-
-A second test class covers MOSS-TTS-Realtime (1.7B, MossTTSRealtime) to
-exercise the local-transformer generation path.
-
-Both classes use a module-scoped engine so the model is loaded once per test
-module execution.  Do NOT add a second engine fixture in this file — that
-triggers mid-module teardown races (see skill invariant I4).
+"""E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B).
+
+Uses the standard omni_runner + pytestmark pattern (one module-scoped engine
+per file, skill invariant I4).  The realtime variant lives in the sibling file
+test_moss_tts_realtime.py.
+
+MOSS-VoiceGenerator synthesizes speech from a text *instruction* that describes
+the desired voice (e.g. "a warm female voice with an American accent").  It does
+NOT accept a reference audio clip.  The request format requires running
+AutoProcessor.from_pretrained once per call to encode the (text, instruction)
+pair into the (prompt_token_ids, codes.ref) grid that the MossTTSDelayModel
+talker expects — same as examples/offline_inference/text_to_speech/moss_tts/
+end2end.py:_build_unified_codes.
+
+No determinism test: VoiceGenerator produces variable-length output even with
+a fixed seed; waveform length reproducibility is not guaranteed.
 """
 
 from __future__ import annotations
 
+import gc
 import os
-import urllib.request
 
 import pytest
 import torch
 from vllm import SamplingParams
 
 from tests.helpers.mark import hardware_test
-from vllm_omni import Omni
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 # ---------------------------------------------------------------------------
-# Shared constants
+# Constants
 # ---------------------------------------------------------------------------
 
-# Selected by the nightly "TTS · Function Test" (-m "full_model and L4 and tts").
-# Without the tts + run-level marks these tests are collected but never run.
-# TODO: Fix this test
-# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
-
-SAMPLE_RATE = 24_000  # All MOSS-TTS full variants output 24 kHz
+MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator"
+SAMPLE_RATE = 24_000
 
-REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
-
-_DEFAULT_SAMPLING = SamplingParams(
+# Stage 0 = AR talker; max_tokens=512 keeps tests fast (~20 s of audio).
+# Stage 1 = codec decoder; greedy, large context to hold all codec tokens.
+_STAGE0_PARAMS = SamplingParams(
     temperature=1.7,
     top_p=0.8,
     top_k=25,
-    max_tokens=512,  # short for tests (~20 s at 24 kHz, ~12.5 tokens/s)
+    max_tokens=512,
+    seed=42,
+    detokenize=False,
+)
+_STAGE1_PARAMS = SamplingParams(
+    temperature=0.0,
+    top_p=1.0,
+    top_k=-1,
+    max_tokens=65536,
     seed=42,
     detokenize=False,
 )
+_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS]
 
 # ---------------------------------------------------------------------------
-# Session fixtures
+# Deploy config
 # ---------------------------------------------------------------------------
 
 
-@pytest.fixture(scope="session")
-def ref_audio_path(tmp_path_factory) -> str:
-    """Download the upstream reference clip once per session.
+def _get_test_config() -> str:
+    """Derive a CI-friendly config from moss_voice_generator.yaml.
 
-    Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments.
+    Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom
+    for Stage 1 (0.12) on a shared L4/A10G.  max_num_seqs=1 keeps per-test
+    peak memory predictable.
     """
-    cache_dir = tmp_path_factory.mktemp("moss_tts_ref")
-    target = cache_dir / "zh_1.wav"
-    try:
-        with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
-            target.write_bytes(resp.read())
-    except Exception as exc:
-        msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}"
-        if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
-            pytest.skip(msg)
-        pytest.fail(msg)
-    if not target.exists() or target.stat().st_size == 0:
-        pytest.fail(f"Reference clip empty after download: {target}")
-    return str(target)
+    return modify_stage_config(
+        get_deploy_config_path("moss_voice_generator.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "max_num_seqs": 1,
+                    "gpu_memory_utilization": 0.45,
+                },
+            },
+        },
+    )
+
+
+# ---------------------------------------------------------------------------
+# pytestmark — one engine for the whole module
+# ---------------------------------------------------------------------------
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True),
+]
 
 
 # ---------------------------------------------------------------------------
@@ -79,147 +99,142 @@ def ref_audio_path(tmp_path_factory) -> str:
 # ---------------------------------------------------------------------------
 
 
-def _build_request(text: str, ref_audio_path: str, seed: int = 42) -> dict:
+def _build_request(text: str, instruction: str) -> dict:
+    """Build a VoiceGenerator request using AutoProcessor.
+
+    Calls proc.build_user_message + proc() to produce the unified-codes grid
+    (L, 1+n_vq) and splits it into prompt_token_ids + codes.ref, exactly as
+    end2end.py:_build_unified_codes does for the voice_generator variant.
+    The processor (including its CPU-resident audio_tokenizer) is freed before
+    returning so it does not compete with the running engine.
+
+    Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments where
+    the HF snapshot is not cached.
+    """
+    from transformers import AutoProcessor
+
+    try:
+        proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True)
+    except Exception as exc:
+        msg = f"Cannot load AutoProcessor for {MODEL}: {exc}"
+        if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
+            pytest.skip(msg)
+        pytest.fail(msg)
+
+    user_msg = proc.build_user_message(text=text, instruction=instruction)
+    batch = proc(conversations=[[user_msg]], mode="generation")
+    unified = batch["input_ids"][0]  # (L, 1+n_vq)
+    text_ids = unified[:, 0].tolist()
+    audio_codes = unified[:, 1:].contiguous().to(torch.int64)  # (L, n_vq)
+    del proc
+    gc.collect()
+
     return {
-        "prompt": "<|im_start|>assistant\n",
-        "additional_information": {
-            "task_type": ["voice_clone"],
-            "text": [text],
-            "mode": ["voice_clone"],
-            "prompt_audio_path": [ref_audio_path],
-            "seed": [seed],
-        },
+        "prompt_token_ids": text_ids,
+        "additional_information": {"codes": {"ref": audio_codes}},
     }
 
 
-def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
-    """Run one request; return (waveform_cpu, sample_rate)."""
-    for stage_outputs in omni.generate(request, _DEFAULT_SAMPLING):
-        for req_output in stage_outputs.request_output:
-            mm = req_output.outputs[0].multimodal_output
-            assert mm is not None, "Expected non-None multimodal_output"
-            # Never use ``a or b`` on tensor values — a multi-element tensor
-            # raises on truthiness. Check ``is None`` explicitly and handle the
-            # pre-consolidation list form.
-            audio = mm.get("audio")
-            if audio is None:
-                audio = mm.get("model_outputs")
-            sr = mm.get("sr")
-            assert audio is not None, "Expected 'audio' or 'model_outputs' in multimodal_output"
-            if isinstance(audio, list):
-                audio = torch.cat(
-                    [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
-                    dim=0,
-                )
-            assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}"
-            return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
-    raise AssertionError("No stage outputs received")
+def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]:
+    """Run one request and return (waveform_cpu, sample_rate).
+
+    Iterates all OmniRequestOutputs from generate() via the top-level
+    multimodal_output property (which aggregates from completion outputs).
+    AR-stage outputs have empty multimodal_output and are silently skipped.
+    With async_chunk=True and FINAL_ONLY, the codec stage consolidates its
+    chunks into a single tensor before yielding — but we handle list fallback
+    for robustness.
+    """
+    chunks: list[torch.Tensor] = []
+    sr_final = SAMPLE_RATE
+    for out in omni_runner.omni.generate(request, _SAMPLING):
+        mm = out.multimodal_output
+        if not mm:
+            continue
+        audio = mm.get("audio")
+        if audio is None:
+            audio = mm.get("model_outputs")
+        if audio is None:
+            continue
+        sr = mm.get("sr")
+        if sr is not None:
+            sr_final = int(sr.item() if hasattr(sr, "item") else sr)
+        if isinstance(audio, list):
+            audio = torch.cat(
+                [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
+                dim=0,
+            )
+        if isinstance(audio, torch.Tensor) and audio.numel() > 0:
+            chunks.append(audio.reshape(-1).cpu())
+    if not chunks:
+        raise AssertionError("No audio output received from generate()")
+    return torch.cat(chunks, dim=0), sr_final
 
 
 # ---------------------------------------------------------------------------
-# MossTTSDelayModel — MOSS-VoiceGenerator 1.7B
-# (tests the delay-pattern generation path on a small GPU)
+# Tests
 # ---------------------------------------------------------------------------
 
-_DELAY_MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator"
-
-
-@pytest.fixture(scope="module")
-def delay_engine():
-    return Omni(model=_DELAY_MODEL, stage_init_timeout=300)
 
-
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_delay_english(delay_engine, ref_audio_path):
-    """MossTTSDelayModel: English voice_clone produces non-empty 24 kHz audio."""
-    req = _build_request("Hello, this is a MOSS-TTS voice cloning test.", ref_audio_path)
-    audio, sr = _collect_audio(delay_engine, req)
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_moss_tts_english(omni_runner: OmniRunner) -> None:
+    """VoiceGenerator: English instruction produces non-empty 24 kHz audio."""
+    req = _build_request(
+        "Hello, this is a MOSS voice design test.",
+        "a warm female voice with an American accent",
+    )
+    audio, sr = _collect_audio(omni_runner, req)
 
     assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
     assert audio.numel() > 0, "Audio tensor is empty"
     assert not torch.all(audio == 0), "Audio is silence"
 
 
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_delay_chinese(delay_engine, ref_audio_path):
-    """MossTTSDelayModel: Chinese input produces non-empty audio."""
-    req = _build_request("你好，这是语音合成测试。", ref_audio_path)
-    audio, sr = _collect_audio(delay_engine, req)
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_moss_tts_chinese(omni_runner: OmniRunner) -> None:
+    """VoiceGenerator: Chinese input produces non-empty audio."""
+    req = _build_request(
+        "你好，这是语音合成测试。",
+        "一个清晰温暖的女声",
+    )
+    audio, sr = _collect_audio(omni_runner, req)
 
     assert sr == SAMPLE_RATE
     assert audio.numel() > 0
     assert not torch.all(audio == 0)
 
 
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path):
-    """MossTTSDelayModel: same seed yields identical waveforms."""
-    req = _build_request("Reproducibility check.", ref_audio_path, seed=99)
-    audio1, _ = _collect_audio(delay_engine, req)
-    audio2, _ = _collect_audio(delay_engine, req)
-
-    assert audio1.shape == audio2.shape, "Shapes differ across identical seeds"
-    assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms differ across identical seeds"
-
-
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_delay_batch(delay_engine, ref_audio_path):
-    """MossTTSDelayModel: batch of two requests each returns non-empty audio."""
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_moss_tts_batch(omni_runner: OmniRunner) -> None:
+    """VoiceGenerator: batch of two requests each returns non-empty audio."""
     requests = [
-        _build_request("First sentence.", ref_audio_path),
-        _build_request("Second sentence.", ref_audio_path),
+        _build_request("First sentence.", "a warm female voice"),
+        _build_request("Second sentence.", "a young male voice"),
     ]
     results: list[torch.Tensor] = []
-    for stage_outputs in delay_engine.generate(requests, [_DEFAULT_SAMPLING] * 2):
-        for req_output in stage_outputs.request_output:
-            mm = req_output.outputs[0].multimodal_output
-            assert mm is not None
-            audio = mm.get("audio")
-            if audio is None:
-                audio = mm.get("model_outputs")
-            assert audio is not None
-            if isinstance(audio, list):
-                audio = torch.cat(
-                    [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
-                    dim=0,
-                )
+    for out in omni_runner.omni.generate(requests, _SAMPLING):
+        mm = out.multimodal_output
+        if not mm:
+            continue
+        audio = mm.get("audio")
+        if audio is None:
+            audio = mm.get("model_outputs")
+        if audio is None:
+            continue
+        if isinstance(audio, list):
+            audio = torch.cat(
+                [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
+                dim=0,
+            )
+        if isinstance(audio, torch.Tensor) and audio.numel() > 0:
             results.append(audio.reshape(-1).cpu())
 
-    assert len(results) == 2, f"Expected 2 outputs, got {len(results)}"
+    # With async_chunk streaming + FINAL_ONLY, each request produces one
+    # consolidated output. The >= 2 guard handles backward-compatibility if
+    # consolidation behavior changes.
+    assert len(results) >= 2, f"Expected at least 2 audio outputs, got {len(results)}"
     for i, audio in enumerate(results):
-        assert audio.numel() > 0, f"Audio[{i}] is empty"
-
-
-# ---------------------------------------------------------------------------
-# MossTTSRealtime — MOSS-TTS-Realtime 1.7B
-# (tests the local-transformer generation path)
-# ---------------------------------------------------------------------------
-
-_REALTIME_MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
-
-
-@pytest.fixture(scope="module")
-def realtime_engine():
-    return Omni(model=_REALTIME_MODEL, stage_init_timeout=300)
-
-
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_realtime_english(realtime_engine, ref_audio_path):
-    """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
-    req = _build_request("This is a real-time TTS streaming test.", ref_audio_path)
-    audio, sr = _collect_audio(realtime_engine, req)
-
-    assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
-    assert audio.numel() > 0, "Audio tensor is empty"
-    assert not torch.all(audio == 0), "Audio is silence"
-
-
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_realtime_deterministic(realtime_engine, ref_audio_path):
-    """MossTTSRealtime: same seed yields identical waveforms."""
-    req = _build_request("Determinism check for realtime model.", ref_audio_path, seed=7)
-    audio1, _ = _collect_audio(realtime_engine, req)
-    audio2, _ = _collect_audio(realtime_engine, req)
-
-    assert audio1.shape == audio2.shape
-    assert torch.allclose(audio1, audio2, atol=1e-4)
+        assert audio.numel() > 0, f"Audio chunk[{i}] is empty"
diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py
new file mode 100644
index 00000000000..540c8750edb
--- /dev/null
+++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py
@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B).
+
+Uses the standard omni_runner + pytestmark pattern (one module-scoped engine
+per file, skill invariant I4).  The delay-model variant lives in the sibling
+file test_moss_tts.py.
+
+The request format mirrors end2end.py:_build_realtime_prompt exactly:
+  1. snapshot_download to locate the MossTTSRealtimeProcessor module.
+  2. MOSS-Audio-Tokenizer (separate HF repo) to encode the reference clip.
+  3. Build the (L, 17) int grid; split into prompt_token_ids + codes.ref +
+     ids.all (remaining text tokens fed one-per-step during decode).
+
+Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments.
+
+No determinism test: MossTTSRealtime uses a local depth transformer that is
+sensitive to async scheduling timing; waveform reproducibility is not
+guaranteed across back-to-back calls.
+"""
+
+from __future__ import annotations
+
+import gc
+import importlib.util
+import os
+import sys
+import urllib.request
+from pathlib import Path
+
+import numpy as np
+import pytest
+import soundfile as sf
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
+SAMPLE_RATE = 24_000
+
+REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
+
+_STAGE0_PARAMS = SamplingParams(
+    temperature=1.0,
+    top_p=0.95,
+    top_k=50,
+    max_tokens=512,
+    seed=42,
+    detokenize=False,
+)
+_STAGE1_PARAMS = SamplingParams(
+    temperature=0.0,
+    top_p=1.0,
+    top_k=-1,
+    max_tokens=65536,
+    seed=42,
+    detokenize=False,
+)
+_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS]
+
+# ---------------------------------------------------------------------------
+# Deploy config
+# ---------------------------------------------------------------------------
+
+
+def _get_test_config() -> str:
+    """Derive a CI-friendly config from moss_tts_realtime.yaml.
+
+    Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom
+    for Stage 1 (0.12) on a shared L4/A10G.  max_num_seqs=1 keeps per-test
+    peak memory predictable.
+    """
+    return modify_stage_config(
+        get_deploy_config_path("moss_tts_realtime.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "max_num_seqs": 1,
+                    "gpu_memory_utilization": 0.45,
+                },
+            },
+        },
+    )
+
+
+# ---------------------------------------------------------------------------
+# pytestmark — one engine for the whole module
+# ---------------------------------------------------------------------------
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True),
+]
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def ref_audio_path(tmp_path_factory) -> str:
+    """Download the upstream reference clip once per test module.
+
+    Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments.
+    """
+    target = tmp_path_factory.mktemp("moss_tts_realtime_ref") / "zh_1.wav"
+    try:
+        with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
+            target.write_bytes(resp.read())
+    except Exception as exc:
+        msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}"
+        if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
+            pytest.skip(msg)
+        pytest.fail(msg)
+    if not target.exists() or target.stat().st_size == 0:
+        pytest.fail(f"Reference clip empty after download: {target}")
+    return str(target)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
+    """Load and resample reference audio to (1, T) float32 at target_sr."""
+    data, sr = sf.read(path, always_2d=True)  # (T, C)
+    wav = torch.from_numpy(data.T.astype("float32"))  # (C, T)
+    if wav.shape[0] > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+    if sr != target_sr:
+        import torchaudio
+
+        wav = torchaudio.functional.resample(wav, sr, target_sr)
+    return wav  # (1, T)
+
+
+def _build_request(ref_audio_path: str, text: str) -> dict:
+    """Build a MOSS-TTS-Realtime request.
+
+    Mirrors end2end.py:_build_realtime_prompt exactly:
+      1. Locate MossTTSRealtimeProcessor in the HF snapshot.
+      2. Encode ref audio via MOSS-Audio-Tokenizer (separate repo).
+      3. Build the (L, 17) int grid and return prompt_token_ids +
+         additional_information: {codes.ref, ids.all}.
+
+    Frees the codec model before returning to avoid competing with the
+    running vllm engine for GPU/CPU memory.
+    """
+    from huggingface_hub import snapshot_download
+    from transformers import AutoModel, AutoTokenizer
+
+    # Step 1: locate the realtime processor module in the snapshot.
+    try:
+        snap_dir = Path(snapshot_download(repo_id=MODEL))
+    except Exception as exc:
+        msg = f"Cannot locate snapshot for {MODEL}: {exc}"
+        if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
+            pytest.skip(msg)
+        pytest.fail(msg)
+
+    proc_module_path = snap_dir / "processing_mossttsrealtime.py"
+    if not proc_module_path.exists():
+        pytest.fail(f"Realtime processor module missing at {proc_module_path}")
+
+    mod_name = "_moss_tts_realtime_proc"
+    if mod_name not in sys.modules:
+        spec = importlib.util.spec_from_file_location(mod_name, proc_module_path)
+        mod = importlib.util.module_from_spec(spec)  # type: ignore[arg-type]
+        sys.modules[mod_name] = mod
+        spec.loader.exec_module(mod)  # type: ignore[union-attr]
+    else:
+        mod = sys.modules[mod_name]
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
+    realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
+
+    # Step 2: encode the reference audio clip via MOSS-Audio-Tokenizer.
+    try:
+        codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).to("cpu").eval()
+    except Exception as exc:
+        msg = f"Cannot load MOSS-Audio-Tokenizer: {exc}"
+        if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
+            pytest.skip(msg)
+        pytest.fail(msg)
+
+    ref_wav = _load_ref_audio(ref_audio_path)
+    wav_2d = ref_wav if ref_wav.dim() == 2 else ref_wav.unsqueeze(0)
+    with torch.no_grad():
+        enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
+    codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
+    audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy()  # (T, 16)
+    del codec
+    gc.collect()
+
+    # Step 3: build the (L, 17) prefill grid.
+    system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens)
+
+    assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
+    assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
+    assistant_grid[:, 0] = assistant_header_ids
+
+    text_ids_only = tokenizer.encode(text, add_special_tokens=False)
+    if not text_ids_only:
+        pytest.fail(f"Empty text after tokenization: {text!r}")
+
+    PREFILL_MAX_TEXT = 12
+    cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
+    text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
+    text_grid[:, 0] = text_ids_only[:cur_len]
+    text_grid[-1, 1] = 1025  # audio_bos at the last prefilled text token
+
+    grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
+    prompt_token_ids = grid[:, 0].tolist()
+    audio_codes_tensor = torch.from_numpy(grid[:, 1:].astype(np.int64).copy())  # (L, 16)
+    remaining_text_ids = list(text_ids_only[cur_len:])
+
+    additional_info: dict = {"codes": {"ref": audio_codes_tensor}}
+    if remaining_text_ids:
+        additional_info["ids"] = {"all": remaining_text_ids}
+
+    return {
+        "prompt_token_ids": prompt_token_ids,
+        "additional_information": additional_info,
+    }
+
+
+def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]:
+    """Run one request and return (waveform_cpu, sample_rate)."""
+    chunks: list[torch.Tensor] = []
+    sr_final = SAMPLE_RATE
+    for out in omni_runner.omni.generate(request, _SAMPLING):
+        mm = out.multimodal_output
+        if not mm:
+            continue
+        audio = mm.get("audio")
+        if audio is None:
+            audio = mm.get("model_outputs")
+        if audio is None:
+            continue
+        sr = mm.get("sr")
+        if sr is not None:
+            sr_final = int(sr.item() if hasattr(sr, "item") else sr)
+        if isinstance(audio, list):
+            audio = torch.cat(
+                [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
+                dim=0,
+            )
+        if isinstance(audio, torch.Tensor) and audio.numel() > 0:
+            chunks.append(audio.reshape(-1).cpu())
+    if not chunks:
+        raise AssertionError("No audio output received from generate()")
+    return torch.cat(chunks, dim=0), sr_final
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None:
+    """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
+    req = _build_request(ref_audio_path, "This is a real-time TTS streaming test.")
+    audio, sr = _collect_audio(omni_runner, req)
+
+    assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
+    assert audio.numel() > 0, "Audio tensor is empty"
+    assert not torch.all(audio == 0), "Audio is silence"
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None:
+    """MossTTSRealtime: Chinese input produces non-empty audio."""
+    req = _build_request(ref_audio_path, "你好，这是语音合成测试。")
+    audio, sr = _collect_audio(omni_runner, req)
+
+    assert sr == SAMPLE_RATE
+    assert audio.numel() > 0
+    assert not torch.all(audio == 0)
diff --git a/tests/e2e/online_serving/test_dreamzero_expansion.py b/tests/e2e/online_serving/test_dreamzero_expansion.py
index bc321589132..79429c1235a 100644
--- a/tests/e2e/online_serving/test_dreamzero_expansion.py
+++ b/tests/e2e/online_serving/test_dreamzero_expansion.py
@@ -28,8 +28,6 @@
         server_args=[
             "--deploy-config",
             "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml",
-            "--enforce-eager",
-            "--disable-log-stats",
         ],
         env_dict={
             "ATTENTION_BACKEND": "torch",
diff --git a/tests/e2e/online_serving/test_gr00t_openpi.py b/tests/e2e/online_serving/test_gr00t_openpi.py
new file mode 100644
index 00000000000..40b8293d64f
--- /dev/null
+++ b/tests/e2e/online_serving/test_gr00t_openpi.py
@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end online serving test for GR00T N1.7 through the OpenPI robot endpoint."""
+
+import os
+
+import numpy as np
+import pytest
+
+from tests.gr00t import openpi_client_helper as openpi_client
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = "nvidia/GR00T-N1.7-3B"
+
+pytest.importorskip("websockets")
+pytest.importorskip("openpi_client.msgpack_numpy")
+
+test_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("Gr00tN1d7.yaml"),
+            server_args=["--disable-log-stats"],
+            env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1", "GR00T_NOISE_SEED": "42"},
+            init_timeout=1200,
+            stage_init_timeout=900,
+        ),
+        id="gr00t-n1d7-openpi",
+    )
+]
+
+# Reference values captured from the Isaac-GR00T ZMQ server (nvidia/GR00T-N1.7-3B,
+# OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT embodiment) using build_droid_observation()
+# inputs: zero images (256×256), identity eef_9d, zero gripper/joint, "pick up the object".
+# Outputs are bit-reproducible across resets and across runs (max_diff=0.0).
+_REF_EEF_9D = np.array(
+    [
+        [
+            0.014888550154864788,
+            -0.00039258378092199564,
+            -0.013574761338531971,
+            0.9999837875366211,
+            0.005678884219378233,
+            -0.00034708858584053814,
+            -0.005677700508385897,
+            0.9999783635139465,
+            0.0033210734836757183,
+        ],
+        [
+            0.020188070833683014,
+            -0.0003105341165792197,
+            -0.02232760563492775,
+            0.9997346997261047,
+            0.02301345206797123,
+            0.000983047066256404,
+            -0.02302766777575016,
+            0.9995648860931396,
+            0.018431292846798897,
+        ],
+        [
+            -0.007266733795404434,
+            -0.05537768080830574,
+            0.03667901083827019,
+            0.992686927318573,
+            0.1206541359424591,
+            0.003906540106981993,
+            -0.12024091929197311,
+            0.9853788614273071,
+            0.12070894986391068,
+        ],
+    ],
+    dtype=np.float32,
+)  # rows = step 0, step 4, step 39
+
+_REF_GRIPPER = np.array([[0.0], [0.0078125], [0.939453125]], dtype=np.float32)  # steps 0, 4, 39
+
+_REF_JOINT = np.array(
+    [
+        [
+            -0.0010484338272362947,
+            0.0014262489276006818,
+            -0.003565810853615403,
+            -3.846167237497866e-05,
+            -0.0002604846959002316,
+            0.008521700277924538,
+            -0.006872728932648897,
+        ],
+        [
+            -0.009435676969587803,
+            0.0021475711837410927,
+            -0.0031688229646533728,
+            -1.8328893929719925e-05,
+            0.0005945992306806147,
+            0.019159257411956787,
+            0.0009468861389905214,
+        ],
+        [
+            -0.02944088727235794,
+            -0.08419207483530045,
+            -0.0251418836414814,
+            0.00540524534881115,
+            0.04752273112535477,
+            -0.012884500436484814,
+            0.024298785254359245,
+        ],
+    ],
+    dtype=np.float32,
+)  # rows = step 0, step 4, step 39
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_gr00t_n1d7_openpi_online(omni_server) -> None:
+    result = openpi_client.run_policy_session(
+        host=omni_server.host,
+        port=omni_server.port,
+        session_id="gr00t-online-e2e",
+    )
+    openpi_client.validate_session_result(result)
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_gr00t_n1d7_openpi_precision(omni_server) -> None:
+    """Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs)."""
+    client = openpi_client.OpenPIWebsocketClient(host=omni_server.host, port=omni_server.port)
+    try:
+        obs = openpi_client.build_droid_observation(session_id="gr00t-precision-e2e")
+        client.reset({})
+        actions = client.infer(obs)
+    finally:
+        client.close()
+
+    steps = [0, 4, 39]
+    np.testing.assert_allclose(
+        actions["eef_9d"][0, steps, :],
+        _REF_EEF_9D,
+        atol=1e-2,
+        rtol=0.0,
+        err_msg="eef_9d action mismatch vs Isaac-GR00T reference",
+    )
+    np.testing.assert_allclose(
+        actions["gripper_position"][0, steps, :],
+        _REF_GRIPPER,
+        atol=1e-2,
+        rtol=0.0,
+        err_msg="gripper_position action mismatch vs Isaac-GR00T reference",
+    )
+    np.testing.assert_allclose(
+        actions["joint_position"][0, steps, :],
+        _REF_JOINT,
+        atol=1e-2,
+        rtol=0.0,
+        err_msg="joint_position action mismatch vs Isaac-GR00T reference",
+    )
diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py
index 83d17a28cd3..398ae649415 100644
--- a/tests/entrypoints/openai_api/test_openpi_connection.py
+++ b/tests/entrypoints/openai_api/test_openpi_connection.py
@@ -88,6 +88,47 @@ def test_unpack_accepts_msgpack_numpy_marker_dicts():
     np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32))
 
 
+def test_unpack_accepts_openpi_client_ndarray_markers():
+    image = np.arange(6, dtype=np.uint8).reshape(2, 3)
+    payload = {
+        "observation/exterior_image_0_left": {
+            b"__ndarray__": True,
+            b"dtype": image.dtype.str,
+            b"shape": image.shape,
+            b"data": image.tobytes(),
+        }
+    }
+
+    decoded = openpi_connection._unpack_numpy(payload)
+
+    np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image)
+    assert decoded["observation/exterior_image_0_left"].dtype == np.uint8
+
+
+def test_unpack_openpi_client_packed_observation():
+    openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy")
+    image = np.zeros((180, 320, 3), dtype=np.uint8)
+    obs = {
+        "observation/exterior_image_0_left": image,
+        "observation/joint_position": np.zeros(7, dtype=np.float32),
+    }
+
+    decoded = openpi_connection._unpack(openpi_msgpack.packb(obs))
+
+    np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image)
+    np.testing.assert_allclose(decoded["observation/joint_position"], obs["observation/joint_position"])
+
+
+def test_pack_uses_openpi_client_ndarray_markers():
+    openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy")
+    actions = np.asarray([[1.0, 2.0]], dtype=np.float32)
+
+    packed = openpi_connection._pack(actions)
+    decoded = openpi_msgpack.unpackb(packed)
+
+    np.testing.assert_allclose(decoded, actions)
+
+
 def test_unpack_leaves_user_dict_without_numpy_kind_marker_unchanged():
     payload = {
         "metadata": {
diff --git a/tests/entrypoints/openai_api/test_storage_manager.py b/tests/entrypoints/openai_api/test_storage_manager.py
new file mode 100644
index 00000000000..30e50576672
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_storage_manager.py
@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for the OpenAI video storage managers.
+"""
+
+import asyncio
+import os
+import time
+
+import pytest
+
+from vllm_omni.entrypoints.openai import storage as storage_module
+from vllm_omni.entrypoints.openai.storage import FileStorageHandle, LocalStorageManager, LocalStorageTTLManager
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_local_storage_open_returns_file_handle_for_saved_key(tmp_path):
+    storage = LocalStorageManager(storage_path=str(tmp_path / "storage"))
+    save_context = asyncio.run(storage.save(b"video-bytes", "video-123"))
+
+    handle = asyncio.run(storage.open(save_context.key))
+
+    assert isinstance(handle, FileStorageHandle)
+    assert handle.path == os.path.join(storage.storage_path, "video-123")
+    with open(handle.path, "rb") as file:
+        assert file.read() == b"video-bytes"
+
+
+def test_local_storage_open_returns_none_for_missing_key(tmp_path):
+    storage = LocalStorageManager(storage_path=str(tmp_path / "storage"))
+
+    handle = asyncio.run(storage.open("missing-video"))
+
+    assert handle is None
+
+
+def test_local_storage_ttl_save_sets_expiration_metadata(tmp_path, monkeypatch):
+    storage = LocalStorageTTLManager(
+        storage_path=str(tmp_path / "storage"),
+        max_concurrency=1,
+        ttl_seconds=60,
+        sweep_interval_seconds=300,
+    )
+    monkeypatch.setattr(storage_module.time, "time", lambda: 1_700_000_000)
+
+    save_context = asyncio.run(storage.save(b"video-bytes", "video-ttl"))
+
+    assert save_context.key == "video-ttl"
+    assert save_context.created_at == 1_700_000_000
+    assert save_context.expires_at == 1_700_000_060
+
+
+def test_local_storage_ttl_sweeper_removes_expired_file(tmp_path):
+    storage = LocalStorageTTLManager(
+        storage_path=str(tmp_path / "storage"),
+        max_concurrency=1,
+        ttl_seconds=1,
+        sweep_interval_seconds=60,
+    )
+
+    async def setup_file() -> tuple[str, str]:
+        save_context = await storage.save(b"video-bytes", "video-expired")
+        file_path = storage.get_full_file_path(save_context.key)
+        return save_context.key, file_path
+
+    storage_key, file_path = asyncio.run(setup_file())
+    expired_mtime = time.time() - 10
+    os.utime(file_path, (expired_mtime, expired_mtime))
+    assert os.path.exists(file_path)
+
+    deleted = asyncio.run(storage._sweep_once(time.time() - 1))
+
+    assert deleted == 1
+    assert not os.path.exists(file_path)
+    assert asyncio.run(storage.open(storage_key)) is None
+
+
+def test_local_storage_ttl_sweeper_keeps_files_when_path_missing(tmp_path):
+    storage = LocalStorageTTLManager(
+        storage_path=str(tmp_path / "storage"),
+        max_concurrency=1,
+        ttl_seconds=1,
+        sweep_interval_seconds=60,
+    )
+
+    missing_path = storage.get_full_file_path("video-missing")
+    assert not os.path.exists(missing_path)
+
+    deleted = asyncio.run(storage._sweep_once(time.time() - 1))
+
+    assert deleted == 0
+    assert not os.path.exists(missing_path)
diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py
index b54afd54a88..dbf76639167 100644
--- a/tests/entrypoints/openai_api/test_video_server.py
+++ b/tests/entrypoints/openai_api/test_video_server.py
@@ -1020,6 +1020,31 @@ def test_invalid_lora_returns_400(test_client):
     assert "lora object" in failed["error"]["message"].lower()
 
 
+def test_failed_generation_awaits_storage_cleanup(test_client, isolated_video_backends, mocker: MockerFixture):
+    """Regression (merge seam): when async generation raises, the failure handler
+    must ``await _cleanup_video(video_id)`` (single-arg, async) and still record
+    FAILED. Upstream carried a sync ``_cleanup_video(video_id, output_path)`` whose
+    stale call in the generic handler sat outside the conflict markers; against the
+    PR's async storage manager that raised NameError before the FAILED update,
+    wedging the job in IN_PROGRESS and orphaning the artifact."""
+    _store, _tasks, storage = isolated_video_backends
+    delete_spy = mocker.spy(storage, "delete")
+    mocker.patch.object(
+        OmniOpenAIServingVideo,
+        "generate_video_bytes",
+        side_effect=RuntimeError("GPU exploded"),
+    )
+
+    response = test_client.post("/v1/videos", data={"prompt": "will fail"})
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+
+    failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value)
+    assert failed["error"]["code"] == 500
+    assert "GPU exploded" in failed["error"]["message"]
+    delete_spy.assert_called_once_with(video_id)
+
+
 def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture):
     mocker.patch.object(
         OmniOpenAIServingVideo,
@@ -1178,7 +1203,7 @@ def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: Moc
     final = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
     file_name = final["file_name"]
     assert file_name is not None
-    file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, file_name)
+    file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, video_id)
     assert os.path.exists(file_path)
 
     delete_resp = test_client.delete(f"/v1/videos/{video_id}")
@@ -1189,6 +1214,30 @@ def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: Moc
     assert not os.path.exists(file_path)
 
 
+def test_download_completed_job_uses_storage_open_and_download_name(test_client, mocker: MockerFixture):
+    video_bytes = b"stored-video-data"
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=video_bytes,
+    )
+    create_resp = test_client.post("/v1/videos", data={"prompt": "Download this video"})
+    assert create_resp.status_code == 200
+    video_id = create_resp.json()["id"]
+
+    final = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    file_name = final["file_name"]
+    assert file_name == f"{video_id}.mp4"
+
+    storage_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, video_id)
+    assert os.path.exists(storage_path)
+
+    response = test_client.get(f"/v1/videos/{video_id}/content")
+    assert response.status_code == 200
+    assert response.content == video_bytes
+    assert response.headers["content-type"] == "video/mp4"
+    assert file_name in response.headers["content-disposition"]
+
+
 def test_delete_in_progress_job_cancels_task_and_removes_metadata(test_client):
     handler = BlockingVideoHandler()
     test_client.app.state.openai_serving_video = handler
diff --git a/tests/gr00t/__init__.py b/tests/gr00t/__init__.py
new file mode 100644
index 00000000000..208f01a7cb5
--- /dev/null
+++ b/tests/gr00t/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/tests/gr00t/openpi_client_helper.py b/tests/gr00t/openpi_client_helper.py
new file mode 100644
index 00000000000..e5943ca3307
--- /dev/null
+++ b/tests/gr00t/openpi_client_helper.py
@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+try:
+    import websockets.sync.client as websockets_client
+except ImportError:  # pragma: no cover - optional e2e dependency
+    websockets_client = None
+
+try:
+    from openpi_client import msgpack_numpy
+except ImportError:  # pragma: no cover - optional e2e dependency
+    msgpack_numpy = None
+
+PING_INTERVAL_SECS = 300
+PING_TIMEOUT_SECS = 3600
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8000
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_SESSION_ID = "gr00t-smoke"
+ACTION_KEYS = {"eef_9d", "gripper_position", "joint_position"}
+LANGUAGE_KEY = "annotation.language.language_instruction"
+
+
+def _identity_eef_9d_state() -> np.ndarray:
+    state = np.zeros((1, 1, 9), dtype=np.float32)
+    state[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
+    return state
+
+
+def require_dependencies() -> None:
+    missing = []
+    if websockets_client is None:
+        missing.append("websockets")
+    if msgpack_numpy is None:
+        missing.append("openpi-client")
+    if missing:
+        raise ModuleNotFoundError(f"GR00T OpenPI test dependencies are missing: {', '.join(missing)}")
+
+
+@dataclass(frozen=True)
+class Gr00tServerMetadata:
+    action_horizon: int
+    action_keys: set[str]
+    embodiment_tag: str
+    needs_session_id: bool
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]):
+        required_keys = ("action_horizon", "action_keys", "embodiment_tag", "needs_session_id")
+        missing_keys = [key for key in required_keys if key not in payload]
+        if missing_keys:
+            raise ValueError(f"Missing GR00T metadata keys: {missing_keys}")
+
+        return cls(
+            action_horizon=int(payload["action_horizon"]),
+            action_keys={str(key) for key in payload["action_keys"]},
+            embodiment_tag=str(payload["embodiment_tag"]),
+            needs_session_id=bool(payload["needs_session_id"]),
+        )
+
+
+class OpenPIWebsocketClient:
+    def __init__(
+        self,
+        *,
+        host: str = DEFAULT_HOST,
+        port: int = DEFAULT_PORT,
+        path: str = DEFAULT_PATH,
+    ) -> None:
+        require_dependencies()
+        self._uri = f"ws://{host}:{port}{path}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._connect()
+
+    def _connect(self):
+        conn = websockets_client.connect(
+            self._uri,
+            compression=None,
+            max_size=None,
+            ping_interval=PING_INTERVAL_SECS,
+            ping_timeout=PING_TIMEOUT_SECS,
+        )
+        metadata = msgpack_numpy.unpackb(conn.recv())
+        if not isinstance(metadata, dict):
+            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
+        return conn, metadata
+
+    def get_server_metadata(self) -> dict[str, Any]:
+        return dict(self._server_metadata)
+
+    def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]:
+        payload = dict(obs)
+        payload["endpoint"] = "infer"
+        self._ws.send(self._packer.pack(payload))
+        response = msgpack_numpy.unpackb(self._ws.recv())
+        if isinstance(response, dict) and response.get("type") == "error":
+            raise RuntimeError(f"Inference failed: {response['message']}")
+        if not isinstance(response, dict):
+            raise RuntimeError(f"Expected dict actions from GR00T OpenPI endpoint, got {type(response)!r}")
+        return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
+
+    def reset(self, reset_info: dict[str, Any] | None = None) -> str:
+        payload = dict(reset_info or {})
+        payload["endpoint"] = "reset"
+        self._ws.send(self._packer.pack(payload))
+        response = msgpack_numpy.unpackb(self._ws.recv())
+        if not isinstance(response, dict) or response.get("status") != "reset successful":
+            raise RuntimeError(f"Unexpected reset response: {response!r}")
+        return str(response["status"])
+
+    def close(self) -> None:
+        self._ws.close()
+
+
+def build_droid_observation(*, session_id: str = DEFAULT_SESSION_ID) -> dict[str, Any]:
+    return {
+        "session_id": session_id,
+        "video": {
+            "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
+            "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
+        },
+        "state": {
+            "eef_9d": _identity_eef_9d_state(),
+            "gripper_position": np.zeros((1, 1, 1), dtype=np.float32),
+            "joint_position": np.zeros((1, 1, 7), dtype=np.float32),
+        },
+        "language": {LANGUAGE_KEY: [["pick up the object"]]},
+    }
+
+
+def run_policy_session(
+    *,
+    host: str = DEFAULT_HOST,
+    port: int = DEFAULT_PORT,
+    path: str = DEFAULT_PATH,
+    session_id: str = DEFAULT_SESSION_ID,
+) -> dict[str, Any]:
+    client = OpenPIWebsocketClient(host=host, port=port, path=path)
+    try:
+        metadata = client.get_server_metadata()
+        actions = client.infer(build_droid_observation(session_id=session_id))
+        reset_status = client.reset({})
+        return {
+            "metadata": metadata,
+            "actions": actions,
+            "reset_status": reset_status,
+            "session_id": session_id,
+        }
+    finally:
+        client.close()
+
+
+def validate_session_result(result: dict[str, Any]) -> None:
+    metadata = Gr00tServerMetadata.from_dict(result["metadata"])
+    if metadata.embodiment_tag != "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT":
+        raise AssertionError(f"Unexpected embodiment_tag: {metadata.embodiment_tag}")
+    if not metadata.needs_session_id:
+        raise AssertionError("GR00T test expects needs_session_id metadata")
+    if metadata.action_keys != ACTION_KEYS:
+        raise AssertionError(f"Unexpected action keys: {metadata.action_keys}")
+    if result["reset_status"] != "reset successful":
+        raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}")
+
+    actions = result["actions"]
+    if set(actions) != ACTION_KEYS:
+        raise AssertionError(f"Unexpected action keys: {set(actions)}")
+    expected_shapes = {
+        "eef_9d": (1, metadata.action_horizon, 9),
+        "gripper_position": (1, metadata.action_horizon, 1),
+        "joint_position": (1, metadata.action_horizon, 7),
+    }
+    for key, expected_shape in expected_shapes.items():
+        if actions[key].shape != expected_shape:
+            raise AssertionError(f"Action {key} shape mismatch: expected {expected_shape}, got {actions[key].shape}")
+        if actions[key].dtype != np.float32:
+            raise AssertionError(f"Action {key} dtype mismatch: expected float32, got {actions[key].dtype}")
+        if not np.isfinite(actions[key]).all():
+            raise AssertionError(f"Action {key} contains non-finite values")
diff --git a/tests/model_executor/models/moss_tts/test_per_request_codes.py b/tests/model_executor/models/moss_tts/test_per_request_codes.py
new file mode 100644
index 00000000000..eabb4fdc02f
--- /dev/null
+++ b/tests/model_executor/models/moss_tts/test_per_request_codes.py
@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MOSS-TTS per-request code routing in the talker.
+
+Guards the cross-request audio-corruption class (RFC #4316 Issue 1, bug report
+#4355): under continuous batching every Stage-1 request must decode its *own*
+codes, not request 0's. The talker emits ``codes["audio"]`` as a batch-ordered
+**list** so the generic splitter (``to_payload_element``) routes
+``element[idx]`` to request ``idx``; the list must stay index-aligned with the
+batch (placeholder for skipped requests) or a shorter list silently falls back
+to entry 0 — the original corruption. The hidden-row mapping must also use the
+runner's real per-request token spans, not an equal split, or mixed
+prefill+decode steps sample each request's codes from the wrong rows.
+
+CPU-only and weight-free: ``make_omni_output`` is driven directly on a bare
+talker with ``_sample_audio_codes`` mocked to echo the hidden row it was given,
+so an assertion on the output value proves which row each request sampled from.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+N_VQ = 4
+
+
+@functools.lru_cache(maxsize=1)
+def _delay_talker_cls():
+    """Defer talker import (pulls vLLM model_executor) until first use."""
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSDelayTalkerForGeneration,
+    )
+
+    return MossTTSDelayTalkerForGeneration
+
+
+def _make_bare_delay_talker():
+    """A talker whose only behaviour is the routing/span logic under test.
+
+    ``_sample_audio_codes`` is replaced with a stub returning an (N_VQ,) tensor
+    filled with the value of the hidden row it received. With a ``hidden`` where
+    row r holds the scalar r, the emitted codes reveal which row each request
+    sampled.
+    """
+    cls = _delay_talker_cls()
+    talker = cls.__new__(cls)
+    talker._batch_state = None
+    talker.n_vq = N_VQ
+
+    def _stub_sample(last_h, state):
+        row_value = int(last_h.reshape(-1)[0].item())
+        return torch.full((N_VQ,), row_value, dtype=torch.long)
+
+    talker._sample_audio_codes = _stub_sample
+    return talker
+
+
+def _row_indexed_hidden(num_rows: int) -> torch.Tensor:
+    """(num_rows, 1) tensor where row r holds the scalar r."""
+    return torch.arange(num_rows, dtype=torch.float32).reshape(num_rows, 1)
+
+
+def _audio_payload(result):
+    return result.multimodal_outputs["codes"]["audio"]
+
+
+class TestPerRequestRouting:
+    def test_mixed_prefill_decode_routes_by_span(self):
+        """req 0 prefills 3 rows, req 1 decodes 1 row in the same step.
+
+        With real spans, req 0 samples its last row (2) and req 1 its last row
+        (3); the payload is a batch-aligned list, element[idx] per request.
+        """
+        talker = _make_bare_delay_talker()
+        hidden = _row_indexed_hidden(4)
+        info_dicts = [{}, {}]
+        spans = [(0, 3), (3, 4)]
+
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
+        audio = _audio_payload(result)
+
+        assert isinstance(audio, list)
+        assert len(audio) == 2  # one entry per request, batch-aligned
+        assert audio[0].shape == (1, N_VQ)
+        assert torch.all(audio[0] == 2)  # req 0: last row of span [0,3)
+        assert torch.all(audio[1] == 3)  # req 1: last row of span [3,4)
+        # No cross-request leakage.
+        assert not torch.equal(audio[0], audio[1])
+
+    def test_no_request_samples_outside_its_span(self):
+        talker = _make_bare_delay_talker()
+        spans = [(0, 2), (2, 5), (5, 6)]
+        hidden = _row_indexed_hidden(6)
+        info_dicts = [{}, {}, {}]
+
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
+        audio = _audio_payload(result)
+
+        assert len(audio) == 3
+        sampled = []
+        for i, (start, end) in enumerate(spans):
+            value = int(audio[i].reshape(-1)[0].item())
+            assert start <= value < end, f"req {i} sampled row {value} outside [{start},{end})"
+            sampled.append(value)
+        assert len(set(sampled)) == len(sampled)  # distinct rows, no sharing
+
+    def test_skipped_request_keeps_list_aligned(self):
+        """An empty-span request gets an empty placeholder at its own index, so
+        later requests are not shifted onto request 0's codes."""
+        talker = _make_bare_delay_talker()
+        # req 1 has an empty span (start == end) and must not displace req 2.
+        spans = [(0, 1), (1, 1), (1, 2)]
+        hidden = _row_indexed_hidden(2)
+        info_dicts = [{}, {}, {}]
+
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
+        audio = _audio_payload(result)
+
+        assert len(audio) == 3
+        assert torch.all(audio[0] == 0)  # req 0: row 0
+        assert audio[1].numel() == 0  # req 1: empty placeholder, aligned
+        assert torch.all(audio[2] == 1)  # req 2: row 1, not shifted
+
+    def test_accumulates_into_per_request_history(self):
+        talker = _make_bare_delay_talker()
+        hidden = _row_indexed_hidden(2)
+        prior = torch.full((1, N_VQ), 7, dtype=torch.long)
+        info_dicts = [{"audio_codes": {"accumulated": prior}}, {}]
+        spans = [(0, 1), (1, 2)]
+
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
+        audio = _audio_payload(result)
+
+        assert audio[0].shape == (2, N_VQ)
+        assert torch.all(audio[0][0] == 7)  # prior history kept
+        assert torch.all(audio[0][1] == 0)  # this step's row 0 appended
+        assert torch.all(audio[1] == 1)
+
+    def test_runs_without_spans_backward_compat(self):
+        """Missing request_token_spans must not crash (equal-split fallback)
+        and still yields a batch-aligned list."""
+        talker = _make_bare_delay_talker()
+        hidden = _row_indexed_hidden(2)
+        info_dicts = [{}, {}]
+
+        result = talker.make_omni_output(hidden, runtime_additional_information=info_dicts)
+        audio = _audio_payload(result)
+        assert isinstance(audio, list)
+        assert len(audio) == 2
+
+    def test_no_codes_returns_empty_payload(self):
+        """Zero hidden rows → no codes → empty multimodal payload (no list)."""
+        talker = _make_bare_delay_talker()
+        hidden = torch.zeros((0, 1), dtype=torch.float32)
+        result = talker.make_omni_output(hidden, runtime_additional_information=[{}], request_token_spans=[(0, 0)])
+        assert result.multimodal_outputs == {}
diff --git a/tests/test_omni_request.py b/tests/test_omni_request.py
new file mode 100644
index 00000000000..09f893a0399
--- /dev/null
+++ b/tests/test_omni_request.py
@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for OmniRequest substitutability with the base Request.
+
+vllm-omni rebinds ``vllm.v1.request.Request`` to ``OmniRequest`` at import
+time (see ``vllm_omni/patch.py``). vLLM core constructs ``Request`` positionally
+in some paths (notably ``vllm/v1/worker/gpu/warmup.py::warmup_kernels`` for V2
+model-runner architectures like Qwen3ForCausalLM). The omni-specific params
+must therefore be keyword-only so positional construction stays Liskov-
+substitutable with the base class — including base-style calls that pass
+``prompt_embeds`` (itself a positional-capable base param) positionally.
+"""
+
+import inspect
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.request import Request
+
+from vllm_omni.engine import PromptEmbedsPayload
+from vllm_omni.request import OmniRequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_omni_params_are_keyword_only():
+    """The three omni params must be keyword-only after ``*args``.
+
+    Guards against re-introducing the import-time rebind bug where omni params
+    came first positionally and broke positional ``Request(...)`` construction.
+    """
+    params = inspect.signature(OmniRequest.__init__).parameters
+    for name in ("prompt_embeds", "external_req_id", "additional_information"):
+        assert params[name].kind is inspect.Parameter.KEYWORD_ONLY, name
+
+
+def test_positional_construction_matches_base_request():
+    """Reproduces the warmup_kernels call: Request(id, tokens, sp, pp)."""
+    req = OmniRequest("req-0", [1, 2, 3], SamplingParams(), None)
+
+    assert isinstance(req, Request)
+    assert req.request_id == "req-0"
+    assert req.prompt_token_ids == [1, 2, 3]
+    # Omni params default cleanly when constructed positionally.
+    assert req.external_req_id is None
+    assert req.additional_information is None
+    assert req.prompt_embeds_payload is None
+
+
+def test_positional_prompt_embeds_does_not_collide():
+    """Base-style positional call that includes ``prompt_embeds``.
+
+    ``prompt_embeds`` is a positional-capable base param (after
+    ``arrival_time``). A base-style call must not trip
+    ``got multiple values for argument 'prompt_embeds'`` — the subclass must
+    not also inject its keyword override when the value arrived positionally.
+    Positional order: request_id, prompt_token_ids, sampling_params,
+    pooling_params, client_index, arrival_time, prompt_embeds.
+    """
+    embeds = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    req = OmniRequest("req-pe", [1, 2], SamplingParams(), None, 0, None, embeds)
+
+    assert isinstance(req, Request)
+    assert torch.equal(req.prompt_embeds, embeds)
+    # Positional tensor is not a serialized payload.
+    assert req.prompt_embeds_payload is None
+    assert req.external_req_id is None
+
+
+def test_keyword_omni_params_round_trip():
+    """Keyword omni params are preserved; serialized embeds are decoded."""
+    arr = np.arange(6, dtype=np.float32).reshape(2, 3)
+    payload = PromptEmbedsPayload(data=arr.tobytes(), shape=[2, 3], dtype="float32")
+
+    req = OmniRequest(
+        request_id="req-1",
+        prompt_token_ids=[7, 8],
+        sampling_params=SamplingParams(),
+        pooling_params=None,
+        prompt_embeds=payload,
+        external_req_id="ext-1",
+    )
+
+    assert req.external_req_id == "ext-1"
+    # The serialized payload is retained and decoded into a tensor on the base.
+    assert req.prompt_embeds_payload is payload
+    assert torch.equal(req.prompt_embeds, torch.from_numpy(arr))
diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py
index 106b2bf12b8..032bcbc1017 100644
--- a/vllm_omni/config/pipeline_registry.py
+++ b/vllm_omni/config/pipeline_registry.py
@@ -77,6 +77,10 @@
         "vllm_omni.model_executor.models.glm_image.pipeline",
         "GLM_IMAGE_PIPELINE",
     ),
+    "Gr00tN1d7": (
+        "vllm_omni.model_executor.models.gr00t.pipeline",
+        "GR00T_N1D7_PIPELINE",
+    ),
     "hunyuan_image_3_moe": (
         "vllm_omni.model_executor.models.hunyuan_image3.pipeline",
         "HUNYUAN_IMAGE3_PIPELINE",
diff --git a/vllm_omni/config/server_settings.py b/vllm_omni/config/server_settings.py
new file mode 100644
index 00000000000..bb63712ba10
--- /dev/null
+++ b/vllm_omni/config/server_settings.py
@@ -0,0 +1,58 @@
+import os
+import warnings
+from typing import Literal
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+def _resolve_deprecated_env(old_name: str, new_name: str, value):
+    old = os.environ.get(old_name)
+    if old is None:
+        return value
+    warnings.warn(
+        f"{old_name} is deprecated and will be removed in a future release; use {new_name} instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return value if new_name in os.environ else old
+
+
+class FileBackend(BaseModel):
+    type: Literal["file"] = "file"
+    path: str = Field(default="/tmp/storage", validate_default=True, description="Local path to store completed files.")
+    file_concurrency: int = Field(
+        default=4, validate_default=True, description="Maximum number of file operations permitted at a time"
+    )
+    file_ttl: int | None = Field(
+        default=None, description="Optional TTL (in seconds) configuration settings for locally stored files."
+    )
+    ttl_sweep_interval: int | None = Field(
+        default=None, description="Optional frequency (in seconds) to enforce file TTLs."
+    )
+
+    @field_validator("path", mode="before")
+    @classmethod
+    def _migrate_path(cls, value):
+        return _resolve_deprecated_env("VLLM_OMNI_STORAGE_PATH", "VLLM_OMNI_SERVER_STORAGE__PATH", value)
+
+    @field_validator("file_concurrency", mode="before")
+    @classmethod
+    def _migrate_file_concurrency(cls, value):
+        return _resolve_deprecated_env(
+            "VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", value
+        )
+
+    @model_validator(mode="after")
+    def set_default_ttl_sweep_interval(self) -> "FileBackend":
+        if self.file_ttl is not None and self.ttl_sweep_interval is None:
+            self.ttl_sweep_interval = 300
+        return self
+
+
+class ServerSettings(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="VLLM_OMNI_SERVER_", env_nested_delimiter="__")
+    storage: FileBackend = Field(default_factory=FileBackend)
+
+
+SERVER_SETTINGS_CONFIG = ServerSettings()
diff --git a/vllm_omni/deploy/Gr00tN1d7.yaml b/vllm_omni/deploy/Gr00tN1d7.yaml
new file mode 100644
index 00000000000..0e90d8536b7
--- /dev/null
+++ b/vllm_omni/deploy/Gr00tN1d7.yaml
@@ -0,0 +1,41 @@
+# GR00T N1.7 deploy: single diffusion stage.
+#
+# Topology is declared in vllm_omni/model_executor/models/gr00t/pipeline.py.
+# This default uses one GPU with TP=1 and CFG parallel disabled.
+
+pipeline: Gr00tN1d7
+# Required: async_chunk defaults to True, which the deploy validator rejects for a
+# single-stage pipeline (no next-stage processor). Must be false here.
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 1
+    enforce_eager: true
+    model_class_name: Gr00tN1d7Pipeline
+    model_config:
+      embodiment_tag: OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT
+      # OpenPI client handshake — model/embodiment-specific values for GR00T-N1.7.
+      # Sent verbatim to the client (the pipeline never reads these); they describe the
+      # action space / inputs the client must produce for this checkpoint + embodiment.
+      policy_server_config:
+        image_resolution: [180, 320]
+        needs_session_id: true
+        embodiment_tag: OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT
+        action_horizon: 40
+        action_keys:
+          - eef_9d
+          - gripper_position
+          - joint_position
+        supported_embodiments:
+          - OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT
+          - XDOF
+          - XDOF_SUBTASK
+          - REAL_G1
+          - REAL_R1_PRO_SHARPA
+          - REAL_R1_PRO_SHARPA_HUMAN
+          - REAL_R1_PRO_SHARPA_MAXINSIGHTS
+          - REAL_R1_PRO_SHARPA_MECKA
+          - LIBERO_PANDA
+          - SIMPLER_ENV_GOOGLE
+          - SIMPLER_ENV_WIDOWX
diff --git a/vllm_omni/deploy/dreamzero.yaml b/vllm_omni/deploy/dreamzero.yaml
index e77afad8409..079101ba0e0 100644
--- a/vllm_omni/deploy/dreamzero.yaml
+++ b/vllm_omni/deploy/dreamzero.yaml
@@ -12,7 +12,7 @@ stages:
   - stage_id: 0
     devices: "0"
     max_num_seqs: 1
-    enforce_eager: true
+    enforce_eager: false
     model_class_name: DreamZeroPipeline
     model_config:
       default_robot_embodiment: roboarena
@@ -23,3 +23,8 @@ stages:
         needs_stereo_camera: false
         needs_session_id: true
         action_space: joint_position
+    cache_backend: step_cache
+    cache_config:
+      step_cache_dit_enabled: true
+      velocity_sim_thresholds: [0.95, 0.93]
+      velocity_skip_countdowns: [4, 2]
diff --git a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
index ed9988c92a3..a453d43f3bf 100644
--- a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
+++ b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
@@ -9,7 +9,7 @@ stages:
   - stage_id: 0
     devices: "0,1"
     max_num_seqs: 1
-    enforce_eager: true
+    enforce_eager: false
     model_class_name: DreamZeroPipeline
     parallel_config:
       tensor_parallel_size: 1
@@ -23,3 +23,8 @@ stages:
         needs_stereo_camera: false
         needs_session_id: true
         action_space: joint_position
+    cache_backend: step_cache
+    cache_config:
+      step_cache_dit_enabled: true
+      velocity_sim_thresholds: [0.95, 0.93]
+      velocity_skip_countdowns: [4, 2]
diff --git a/vllm_omni/diffusion/cache/base.py b/vllm_omni/diffusion/cache/base.py
index d7807622497..aafbf0ba7a3 100644
--- a/vllm_omni/diffusion/cache/base.py
+++ b/vllm_omni/diffusion/cache/base.py
@@ -13,6 +13,8 @@
    the cache-dit library. Inherits from CacheBackend. Used via cache_backend="cache_dit".
 2. TeaCacheBackend: Hook-based backend for TeaCache acceleration. Inherits from
    CacheBackend. Used via cache_backend="tea_cache".
+3. StepCacheBackend: Velocity cosine step-skipping for DreamZero. Inherits from
+   CacheBackend. Used via cache_backend="step_cache".
 
 All backends implement the same interface:
 - enable(pipeline): Enable cache on the pipeline
diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py
index ff1eea6aeaa..5add71daa61 100644
--- a/vllm_omni/diffusion/cache/cache_dit_backend.py
+++ b/vllm_omni/diffusion/cache/cache_dit_backend.py
@@ -10,11 +10,19 @@
 import functools
 from collections.abc import Callable
 from contextlib import ExitStack
-from typing import Any, Optional
+from dataclasses import dataclass
+from typing import Any, Optional, TypeAlias
 
 import cache_dit
 import torch
-from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifier, TaylorSeerCalibratorConfig
+from cache_dit import (
+    BlockAdapter,
+    CalibratorConfig,
+    DBCacheConfig,
+    ForwardPattern,
+    ParamsModifier,
+    TaylorSeerCalibratorConfig,
+)
 from cache_dit.caching.block_adapters import FakeDiffusionPipeline
 from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
@@ -25,10 +33,29 @@
 
 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig
-from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline
 
 logger = init_logger(__name__)
 
+RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None]
+
+
+@dataclass
+class CacheDiTAdapterConfig:
+    """Config for creating a Cache DiT's block adapter; to enable CacheDiT,
+    most models just need to define an instance of this class as a class
+    var in the DiT.
+    """
+
+    block_forward_patterns: dict[str, ForwardPattern]
+    has_separate_cfg: bool = False
+    cached_adapter_cls: type[CachedAdapter] | None = None
+    check_forward_pattern: bool = True
+
+
+# Registry of custom cache-dit enablers for specific models
+# Maps pipeline names to their cache-dit enablement functions
+CUSTOM_DIT_ENABLERS: dict[str, Callable] = {}
+
 
 # Small helper to centralize cache-dit summaries.
 def cache_summary(pipeline: Any, details: bool = True) -> None:
@@ -37,14 +64,60 @@ def cache_summary(pipeline: Any, details: bool = True) -> None:
         cache_dit.summary(pipeline.transformer_2, details=details)
 
 
-# Registry of custom cache-dit enablers for specific models
-# Maps pipeline names to their cache-dit enablement functions
-# Models in this registry require custom handling (e.g., dual-transformer architectures)
-# Will be populated after function definitions
-CUSTOM_DIT_ENABLERS: dict[str, Callable] = {}
+def default_get_pipeline_transformer(pipeline: Any) -> Any:
+    return pipeline.transformer
+
+
+def build_cache_context_refresh(
+    cache_config: DiffusionCacheConfig,
+    get_pipeline_transformer: Callable[[Any], Any] = default_get_pipeline_transformer,
+) -> RefreshCacheContextFunc:
+    """Build the cache context refresh func for a single Transformer."""
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
+
+        Args:
+            pipeline: The pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        transformer = get_pipeline_transformer(pipeline)
+
+        # Bypass SCM for step counts that don't support predefined masks (e.g., vLLM's 1-step dummy run)
+        scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6)
+
+        if cache_config.scm_steps_mask_policy is None or not scm_supported_steps:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
+def _resolve_calibrator_config(cache_config: DiffusionCacheConfig) -> CalibratorConfig | None:
+    """Resolves the Calibrator subconfig For DiT cache."""
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+    # In the future, more calibrators will likely be added to DiT Cache,
+    # e.g., focal; handle them generically here.
+    return calibrator
 
 
-def _build_db_cache_config(cache_config: Any) -> DBCacheConfig:
+def _build_db_cache_config(cache_config: DiffusionCacheConfig) -> DBCacheConfig:
     """Build DBCacheConfig with optional SCM (Step Computation Masking) support.
 
     Args:
@@ -53,9 +126,8 @@ def _build_db_cache_config(cache_config: Any) -> DBCacheConfig:
     Returns:
         DBCacheConfig instance with SCM support if configured.
     """
-
     return DBCacheConfig(
-        # we will refresh the context when gets num_inference_steps in the first inference request
+        # we will refresh the context when we get num_inference_steps in the first inference request
         num_inference_steps=None,
         Fn_compute_blocks=cache_config.Fn_compute_blocks,
         Bn_compute_blocks=cache_config.Bn_compute_blocks,
@@ -68,7 +140,100 @@ def _build_db_cache_config(cache_config: Any) -> DBCacheConfig:
     )
 
 
-def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+def enable_cache_for_dit(
+    pipeline: Any,
+    cache_config: Any,
+    block_adapter: BlockAdapter | None = None,
+    adapter_cls: type[CachedAdapter] | None = None,
+) -> RefreshCacheContextFunc:
+    """Enable cache-dit for regular single-transformer DiT models.
+
+    Args:
+        pipeline: The diffusion pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+        block_adapter: Custom block adapters for specific model architectures.
+        adapter_cls: Custom cached adapter class for specific model architectures.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    # Build DBCacheConfig with optional SCM support
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    # Build calibrator config if TaylorSeer is enabled
+    calibrator_config = _resolve_calibrator_config(cache_config)
+
+    logger.info(
+        f"Enabling cache-dit on transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    # Enable cache-dit on the transformer
+    transformer = default_get_pipeline_transformer(pipeline)
+
+    # If we have a custom cached adapter subclass, call apply directly
+    if adapter_cls is not None:
+        adapter_cls.apply(
+            transformer if block_adapter is None else block_adapter,
+            cache_config=db_cache_config,
+            calibrator_config=calibrator_config,
+        )
+    else:
+        # Otherwise, call enable cache, which will call CachedAdapter.apply for us
+        cache_dit.enable_cache(
+            transformer if block_adapter is None else block_adapter,
+            cache_config=db_cache_config,
+            calibrator_config=calibrator_config,
+        )
+
+    return build_cache_context_refresh(cache_config)
+
+
+### Complex / custom enablers for DiT cache
+# NOTE (Alex): This case is rare; you should only really need to do this if you have a dual transformer
+# architecture, since it hasn't been handled generically yet, or if the model class has unique attributes
+# that it sets during Cache DiT enablement.
+#
+# For the vast majority of models, you should only need to add a _cache_dit_adapter_config attribute
+# to the Transformer class, which controls the forward pattern, whether or not we have separate CFG,
+# and so on.
+
+
+# from https://github.com/vipshop/cache-dit/pull/542
+def _split_wan22_inference_steps(pipeline, num_inference_steps: int) -> tuple[int, int]:
+    """Split inference steps into high-noise and low-noise steps for Wan2.2.
+
+    This is an internal helper function specific to Wan2.2's dual-transformer
+    architecture that uses boundary_ratio to determine the split point.
+
+    Args:
+        num_inference_steps: Total number of inference steps.
+
+    Returns:
+        A tuple of (num_high_noise_steps, num_low_noise_steps).
+    """
+    if pipeline.boundary_ratio is not None:
+        boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps
+    else:
+        boundary_timestep = None
+
+    # Set timesteps to calculate the split
+    device = next(pipeline.transformer.parameters()).device
+    pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
+
+    timesteps = pipeline.scheduler.timesteps
+    num_high_noise_steps = 0  # high-noise steps for transformer
+    for t in timesteps:
+        if boundary_timestep is None or t >= boundary_timestep:
+            num_high_noise_steps += 1
+    # low-noise steps for transformer_2
+    num_low_noise_steps = num_inference_steps - num_high_noise_steps
+    return num_high_noise_steps, num_low_noise_steps
+
+
+def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc:
     """Enable cache-dit for Wan2.2 single or dual-transformer architecture.
 
     Wan2.2 can use single or dual transformers (transformer and transformer_2) that need
@@ -81,10 +246,11 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
     Returns:
         A refresh function that can be called to update cache context with new num_inference_steps.
     """
+    # Build DBCacheConfig with optional SCM support
+    db_cache_config = _build_db_cache_config(cache_config)
 
     if getattr(pipeline, "transformer_2", None) is None:
         logger.info("transformer_2 not found, enabling cache-dit for single transformer mode")
-        db_cache_config = _build_db_cache_config(cache_config)
         cache_dit.enable_cache(
             BlockAdapter(
                 transformer=pipeline.transformer,
@@ -97,27 +263,7 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
             ),
             cache_config=db_cache_config,
         )
-
-        def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-            """Refresh cache context for single transformer."""
-            if cache_config.scm_steps_mask_policy is None:
-                cache_dit.refresh_context(
-                    pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
-                )
-            else:
-                cache_dit.refresh_context(
-                    pipeline.transformer,
-                    cache_config=DBCacheConfig().reset(
-                        num_inference_steps=num_inference_steps,
-                        steps_computation_mask=cache_dit.steps_mask(
-                            mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps
-                        ),
-                        steps_computation_policy=cache_config.scm_steps_policy,
-                    ),
-                    verbose=verbose,
-                )
-
-        return refresh_cache_context
+        return build_cache_context_refresh(cache_config)
 
     cache_dit.enable_cache(
         BlockAdapter(
@@ -150,47 +296,11 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool
             ],
             has_separate_cfg=True,
         ),
-        cache_config=DBCacheConfig(
-            Fn_compute_blocks=cache_config.Fn_compute_blocks,
-            Bn_compute_blocks=cache_config.Bn_compute_blocks,
-            max_warmup_steps=cache_config.max_warmup_steps,
-            max_cached_steps=cache_config.max_cached_steps,
-            max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
-            residual_diff_threshold=cache_config.residual_diff_threshold,
-            num_inference_steps=None,
-        ),
+        cache_config=db_cache_config,
     )
 
-    # from https://github.com/vipshop/cache-dit/pull/542
-    def _split_inference_steps(num_inference_steps: int) -> tuple[int, int]:
-        """Split inference steps into high-noise and low-noise steps for Wan2.2.
-
-        This is an internal helper function specific to Wan2.2's dual-transformer
-        architecture that uses boundary_ratio to determine the split point.
-
-        Args:
-            num_inference_steps: Total number of inference steps.
-
-        Returns:
-            A tuple of (num_high_noise_steps, num_low_noise_steps).
-        """
-        if pipeline.boundary_ratio is not None:
-            boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps
-        else:
-            boundary_timestep = None
-
-        # Set timesteps to calculate the split
-        device = next(pipeline.transformer.parameters()).device
-        pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
-
-        timesteps = pipeline.scheduler.timesteps
-        num_high_noise_steps = 0  # high-noise steps for transformer
-        for t in timesteps:
-            if boundary_timestep is None or t >= boundary_timestep:
-                num_high_noise_steps += 1
-        # low-noise steps for transformer_2
-        num_low_noise_steps = num_inference_steps - num_high_noise_steps
-        return num_high_noise_steps, num_low_noise_steps
+    refresh_trans_one = build_cache_context_refresh(cache_config)
+    refresh_trans_two = build_cache_context_refresh(cache_config, lambda pipeline: pipeline.transformer_2)
 
     def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
         """Refresh cache context for both transformers with new num_inference_steps.
@@ -199,50 +309,14 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool
             pipeline: The Wan2.2 pipeline instance.
             num_inference_steps: New number of inference steps.
         """
-
-        num_high_noise_steps, num_low_noise_steps = _split_inference_steps(num_inference_steps)
-        # Refresh context for high-noise transformer
-        if cache_config.scm_steps_mask_policy is None:
-            # cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_high_noise_steps, verbose=verbose)
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                num_inference_steps=num_high_noise_steps,
-                verbose=verbose,
-            )
-            cache_dit.refresh_context(
-                pipeline.transformer_2,
-                num_inference_steps=num_low_noise_steps,
-                verbose=verbose,
-            )
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_high_noise_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_high_noise_steps
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-            cache_dit.refresh_context(
-                pipeline.transformer_2,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_low_noise_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_low_noise_steps
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
+        num_high_noise_steps, num_low_noise_steps = _split_wan22_inference_steps(pipeline, num_inference_steps)
+        refresh_trans_one(pipeline, num_high_noise_steps, verbose)
+        refresh_trans_two(pipeline, num_low_noise_steps, verbose)
 
     return refresh_cache_context
 
 
-def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc:
     """Enable cache-dit for Wan2.2 S2V.
 
     S2V uses a single transformer, but unlike the other Wan2.2 variants its
@@ -316,1635 +390,550 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool
     return refresh_cache_context
 
 
-def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for LongCatImage pipeline.
-
-    Args:
-        pipeline: The LongCatImage pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
+class BagelCachedContextManager(CachedContextManager):
+    """
+    Custom CachedContextManager for Bagel that safely handles NaiveCache objects
+    (mapped to encoder_hidden_states) by skipping tensor operations on them.
     """
-    # Build DBCacheConfig for transformer
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    # Build ParamsModifier for transformer
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator,
-    )
-
-    logger.info(
-        f"Enabling cache-dit on LongCatImage transformer with BlockAdapter: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
 
-    # Enable cache-dit using BlockAdapter for transformer
-    cache_dit.enable_cache(
-        (
-            BlockAdapter(
-                transformer=pipeline.transformer,
-                blocks=[
-                    pipeline.transformer.transformer_blocks,
-                    pipeline.transformer.single_transformer_blocks,
-                ],
-                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
-                params_modifiers=[modifier],
-                has_separate_cfg=True,
-            )
-        ),
-        cache_config=db_cache_config,
-    )
+    @torch.compiler.disable
+    def apply_cache(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        prefix: str = "Bn",
+        encoder_prefix: str = "Bn_encoder",
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # Allow Bn and Fn prefix to be used for residual cache.
+        if "Bn" in prefix:
+            hidden_states_prev = self.get_Bn_buffer(prefix)
+        else:
+            hidden_states_prev = self.get_Fn_buffer(prefix)
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
+        assert hidden_states_prev is not None, f"{prefix}_buffer must be set before"
 
-        Args:
-            pipeline: The LongCatImage pipeline instance.
-            num_inference_steps: New number of inference steps.
-        """
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        if self.is_cache_residual():
+            hidden_states = hidden_states_prev + hidden_states
         else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
+            # If cache is not residual, we use the hidden states directly
+            hidden_states = hidden_states_prev
 
-    return refresh_cache_context
+        hidden_states = hidden_states.contiguous()
 
+        if encoder_hidden_states is not None:
+            if "Bn" in encoder_prefix:
+                encoder_hidden_states_prev = self.get_Bn_encoder_buffer(encoder_prefix)
+            else:
+                encoder_hidden_states_prev = self.get_Fn_encoder_buffer(encoder_prefix)
 
-def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for Flux.1-dev pipeline.
+            if encoder_hidden_states_prev is not None:
+                if self.is_encoder_cache_residual():
+                    # FIX: Check if encoder_hidden_states is a tensor before adding
+                    if isinstance(encoder_hidden_states, torch.Tensor) and isinstance(
+                        encoder_hidden_states_prev, torch.Tensor
+                    ):
+                        encoder_hidden_states = encoder_hidden_states_prev + encoder_hidden_states
+                else:
+                    # If encoder cache is not residual, we use the encoder hidden states directly
+                    encoder_hidden_states = encoder_hidden_states_prev
 
-    Args:
-        pipeline: The Flux pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-    Returns:
-        A refresh function that can be called with a new ``num_inference_steps``
-        to update the cache context for the pipeline.
-    """
-    # Build DBCacheConfig for transformer
-    db_cache_config = _build_db_cache_config(cache_config)
+            # FIX: Check if encoder_hidden_states is a tensor before calling contiguous
+            if isinstance(encoder_hidden_states, torch.Tensor):
+                encoder_hidden_states = encoder_hidden_states.contiguous()
 
-    calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+        return hidden_states, encoder_hidden_states
 
-    # Build ParamsModifier for transformer
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator,
-    )
 
-    logger.info(
-        f"Enabling cache-dit on Flux transformer with BlockAdapter: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
+class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2):
+    """
+    Custom CachedBlocks for Bagel that safely handles NaiveCache objects
+    by adding isinstance checks in call_Mn_blocks and compute_or_prune.
+    """
 
-    # Enable cache-dit using BlockAdapter for transformer
-    cache_dit.enable_cache(
-        (
-            BlockAdapter(
-                transformer=pipeline.transformer,
-                blocks=[
-                    pipeline.transformer.transformer_blocks,
-                    pipeline.transformer.single_transformer_blocks,
-                ],
-                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
-                params_modifiers=[modifier],
-            )
-        ),
-        cache_config=db_cache_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
-
-        Args:
-            pipeline: The Flux pipeline instance.
-            num_inference_steps: New number of inference steps.
-        """
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
+    def call_Mn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
+        for block in self._Mn_blocks():
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
             )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
 
-    return refresh_cache_context
+        # compute hidden_states residual
+        hidden_states = hidden_states.contiguous()
 
+        hidden_states_residual = hidden_states - original_hidden_states
 
-def enable_cache_for_flux2_klein(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for FLUX.2-klein-4B pipeline.
+        if (
+            encoder_hidden_states is not None
+            and original_encoder_hidden_states is not None
+            and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
+        ):
+            encoder_hidden_states = encoder_hidden_states.contiguous()
+            encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
+        else:
+            encoder_hidden_states_residual = None
 
-    Args:
-        pipeline: The FLUX.2-klein-4B pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-    Returns:
-        A refresh function that can be called with a new ``num_inference_steps``
-        to update the cache context for the pipeline.
-    """
-    # Build DBCacheConfig for transformer
-    db_cache_config = _build_db_cache_config(cache_config)
-    # The Fn_compute_blocks = 2 override is the most important decision here,
-    # and the rationale (quality degradation at Fn=1) only lives in flux2_klein.
-    db_cache_config.Fn_compute_blocks = 2
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            hidden_states_residual,
+            encoder_hidden_states_residual,
+        )
 
-    calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+    def compute_or_prune(
+        self,
+        block_id: int,  # Block index in the transformer blocks
+        # Below are the inputs to the block
+        block,  # The transformer block to be executed
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        # NOTE: Although Bagel likely won't use pruning, implementing safe version just in case.
+        # Copy-pasted from original but adding checks.
 
-    # Build ParamsModifier for transformer
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator,
-    )
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
 
-    logger.info(
-        f"Enabling cache-dit on Flux2-Klein transformer with BlockAdapter: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
+        can_use_prune = self._maybe_prune(
+            block_id,
+            hidden_states,
+            prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
+        )
 
-    # Enable cache-dit using BlockAdapter for transformer
-    cache_dit.enable_cache(
-        (
-            BlockAdapter(
-                transformer=pipeline.transformer,
-                blocks=[
-                    pipeline.transformer.transformer_blocks,
-                    pipeline.transformer.single_transformer_blocks,
-                ],
-                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2],
-                params_modifiers=[modifier],
+        torch._dynamo.graph_break()
+        if can_use_prune:
+            self.context_manager.add_pruned_step()
+            hidden_states, encoder_hidden_states = self.context_manager.apply_prune(
+                hidden_states,
+                encoder_hidden_states,
+                prefix=(
+                    f"{self.cache_prefix}_{block_id}_Bn_residual"
+                    if self.context_manager.is_cache_residual()
+                    else f"{self.cache_prefix}_Bn_hidden_states"
+                ),
+                encoder_prefix=(
+                    f"{self.cache_prefix}_{block_id}_Bn_encoder_residual"
+                    if self.context_manager.is_encoder_cache_residual()
+                    else f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states"
+                ),
             )
-        ),
-        cache_config=db_cache_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
-
-        Args:
-            pipeline: The FLUX.2-klein-4B pipeline instance.
-            num_inference_steps: New number of inference steps.
-        """
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+            torch._dynamo.graph_break()
         else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    Fn_compute_blocks=db_cache_config.Fn_compute_blocks,
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
+            # Normal steps: Compute the block and cache the residuals.
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
             )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
+            if not self._skip_prune(block_id):
+                hidden_states = hidden_states.contiguous()
+                hidden_states_residual = hidden_states - original_hidden_states
 
-    return refresh_cache_context
-
-
-def enable_cache_for_stable_audio_open(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for Stable Audio Open pipeline.
-
-    Args:
-        pipeline: The StableAudioPipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-
-    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
-    """
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+                if (
+                    encoder_hidden_states is not None
+                    and original_encoder_hidden_states is not None
+                    and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
+                ):
+                    encoder_hidden_states = encoder_hidden_states.contiguous()
+                    encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
+                else:
+                    encoder_hidden_states_residual = None
 
-    # StableAudio is officially registered in CacheDiT as Pattern_3:
-    # https://github.com/vipshop/cache-dit/blob/69e82bd1/src/cache_dit/caching/block_adapters/__init__.py#L562
-    #
-    # Pattern_3 is required because StableAudioDiT uses cross-attention
-    # with static encoder_hidden_states that do not change inside the
-    # transformer block loop.
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=pipeline.transformer,
-            blocks=pipeline.transformer.transformer_blocks,
-            forward_pattern=ForwardPattern.Pattern_3,
-            params_modifiers=[
-                ParamsModifier(
-                    cache_config=db_cache_config,
-                    calibrator_config=calibrator_config,
+                self.context_manager.set_Fn_buffer(
+                    original_hidden_states,
+                    prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
                 )
-            ],
-        ),
-        cache_config=db_cache_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
+                if self.context_manager.is_cache_residual():
+                    self.context_manager.set_Bn_buffer(
+                        hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_{block_id}_Bn_residual",
+                    )
+                else:
+                    self.context_manager.set_Bn_buffer(
+                        hidden_states,
+                        prefix=f"{self.cache_prefix}_{block_id}_Bn_hidden_states",
+                    )
+                if encoder_hidden_states_residual is not None:
+                    if self.context_manager.is_encoder_cache_residual():
+                        self.context_manager.set_Bn_encoder_buffer(
+                            encoder_hidden_states_residual,
+                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_residual",
+                        )
+                    else:
+                        self.context_manager.set_Bn_encoder_buffer(
+                            encoder_hidden_states_residual,
+                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states",
+                        )
+            torch._dynamo.graph_break()
 
-        Args:
-            pipeline: The StableAudioPipeline instance.
-            num_inference_steps: New number of inference steps.
-            verbose: Whether to log refresh operations.
-        """
-        # Bypass SCM for step counts that don't support predefined masks (e.g., vLLM's 1-step dummy run)
-        scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6)
+        return hidden_states, encoder_hidden_states
 
-        if cache_config.scm_steps_mask_policy is None or not scm_supported_steps:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            updated_scm_config = DBCacheConfig().reset(
-                num_inference_steps=num_inference_steps,
-                steps_computation_mask=cache_dit.steps_mask(
-                    mask_policy=cache_config.scm_steps_mask_policy,
-                    total_steps=num_inference_steps,
-                ),
-                steps_computation_policy=cache_config.scm_steps_policy,
-            )
 
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=updated_scm_config,
-                verbose=verbose,
-            )
+class Wan22S2VCachedBlocks(CachedBlocks_Pattern_3_4_5):
+    """CacheDiT blocks wrapper that preserves S2V per-layer audio injection."""
 
-    return refresh_cache_context
+    def _run_block(self, block_id: int, block: torch.nn.Module, hidden_states: torch.Tensor, *args, **kwargs):
+        hidden_states = block(hidden_states, *args, **kwargs)
+        hidden_states, new_encoder_hidden_states = self._process_block_outputs(hidden_states)
+        original_after_transformer_block = getattr(
+            self.transformer,
+            "_cache_dit_original_after_transformer_block",
+            getattr(self.transformer, "after_transformer_block", None),
+        )
+        if original_after_transformer_block is not None:
+            hidden_states = original_after_transformer_block(block_id, hidden_states)
+        return hidden_states, new_encoder_hidden_states
 
+    def call_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        for block_id, block in enumerate(self.transformer_blocks):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+        return hidden_states, new_encoder_hidden_states
 
-def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for StableDiffusion3Pipeline.
-
-    Args:
-        pipeline: The StableDiffusion3 pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-    """
-    # Build DBCacheConfig for transformer
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    # Build ParamsModifier for transformer
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator,
-    )
-
-    logger.info(
-        f"Enabling cache-dit on StableDiffusion3 transformer with BlockAdapter: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    # Enable cache-dit using BlockAdapter for transformer
-    cache_dit.enable_cache(
-        (
-            BlockAdapter(
-                transformer=pipeline.transformer,
-                blocks=pipeline.transformer.transformer_blocks,
-                forward_pattern=ForwardPattern.Pattern_1,
-                params_modifiers=[modifier],
-            )
-        ),
-        cache_config=db_cache_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
-
-        Args:
-            pipeline: The LongCatImage pipeline instance.
-            num_inference_steps: New number of inference steps.
-        """
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-    return refresh_cache_context
-
-
-def enable_cache_for_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for LTX2 pipelines (audio-video transformer blocks)."""
-    transformer = get_transformer_from_pipeline(pipeline)
-
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    blocks = transformer.transformer_blocks
-
-    logger.info(
-        f"Enabling cache-dit on LTX2 transformer: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=transformer,
-            blocks=blocks,
-            # LTX2 blocks return (hidden_states, audio_hidden_states)
-            forward_pattern=ForwardPattern.Pattern_0,
-            # Treat audio_hidden_states as encoder_hidden_states in Pattern_0
-            check_forward_pattern=False,
-            has_separate_cfg=True,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        transformer = get_transformer_from_pipeline(pipeline)
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-    return refresh_cache_context
-
-
-def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for regular single-transformer DiT models.
-
-    Args:
-        pipeline: The diffusion pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-
-    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
-    """
-    # Build DBCacheConfig with optional SCM support
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    # Build calibrator config if TaylorSeer is enabled
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    logger.info(
-        f"Enabling cache-dit on transformer: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    # Enable cache-dit on the transformer
-    transformer = get_transformer_from_pipeline(pipeline)
-    cache_dit.enable_cache(
-        transformer,
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
-
-        Args:
-            pipeline: The diffusion pipeline instance.
-            num_inference_steps: New number of inference steps.
-        """
-        transformer = get_transformer_from_pipeline(pipeline)
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-    return refresh_cache_context
-
-
-def enable_cache_for_hunyuan_image3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for HunyuanImage3 pipeline.
-
-    HunyuanImage3 stores its main transformer stack at ``pipeline.model`` with
-    decoder blocks in ``pipeline.model.layers``.
-    """
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    transformer = getattr(pipeline, "model", None)
-    if transformer is None or not hasattr(transformer, "layers"):
-        raise ValueError(
-            f"HunyuanImage3 cache-dit enabler expects pipeline.model.layers, got pipeline={pipeline.__class__.__name__}"
-        )
-
-    logger.info(
-        f"Enabling cache-dit on HunyuanImage3 model: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-    # HunyuanImage3 decoder layers are single-stream `hidden_states` blocks
-    # (`HunyuanImage3DecoderLayer.forward(hidden_states, ...)`) while caller
-    # still reads outputs as tuple (`layer_outputs[0]`). Pattern_4 matches this:
-    # hidden-only input with hidden-first tuple output.
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=transformer,
-            blocks=transformer.layers,
-            forward_pattern=ForwardPattern.Pattern_4,
-            params_modifiers=[modifier],
-            check_forward_pattern=False,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        transformer = getattr(pipeline, "model", None)
-        if transformer is None:
-            raise ValueError("HunyuanImage3 cache-dit refresh expects pipeline.model to exist.")
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-    return refresh_cache_context
-
-
-class BagelCachedContextManager(CachedContextManager):
-    """
-    Custom CachedContextManager for Bagel that safely handles NaiveCache objects
-    (mapped to encoder_hidden_states) by skipping tensor operations on them.
-    """
-
-    @torch.compiler.disable
-    def apply_cache(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
-        prefix: str = "Bn",
-        encoder_prefix: str = "Bn_encoder",
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        # Allow Bn and Fn prefix to be used for residual cache.
-        if "Bn" in prefix:
-            hidden_states_prev = self.get_Bn_buffer(prefix)
-        else:
-            hidden_states_prev = self.get_Fn_buffer(prefix)
-
-        assert hidden_states_prev is not None, f"{prefix}_buffer must be set before"
-
-        if self.is_cache_residual():
-            hidden_states = hidden_states_prev + hidden_states
-        else:
-            # If cache is not residual, we use the hidden states directly
-            hidden_states = hidden_states_prev
-
-        hidden_states = hidden_states.contiguous()
-
-        if encoder_hidden_states is not None:
-            if "Bn" in encoder_prefix:
-                encoder_hidden_states_prev = self.get_Bn_encoder_buffer(encoder_prefix)
-            else:
-                encoder_hidden_states_prev = self.get_Fn_encoder_buffer(encoder_prefix)
-
-            if encoder_hidden_states_prev is not None:
-                if self.is_encoder_cache_residual():
-                    # FIX: Check if encoder_hidden_states is a tensor before adding
-                    if isinstance(encoder_hidden_states, torch.Tensor) and isinstance(
-                        encoder_hidden_states_prev, torch.Tensor
-                    ):
-                        encoder_hidden_states = encoder_hidden_states_prev + encoder_hidden_states
-                else:
-                    # If encoder cache is not residual, we use the encoder hidden states directly
-                    encoder_hidden_states = encoder_hidden_states_prev
-
-            # FIX: Check if encoder_hidden_states is a tensor before calling contiguous
-            if isinstance(encoder_hidden_states, torch.Tensor):
-                encoder_hidden_states = encoder_hidden_states.contiguous()
-
-        return hidden_states, encoder_hidden_states
-
-
-class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2):
-    """
-    Custom CachedBlocks for Bagel that safely handles NaiveCache objects
-    by adding isinstance checks in call_Mn_blocks and compute_or_prune.
-    """
-
-    def call_Mn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
-        for block in self._Mn_blocks():
-            hidden_states = block(
-                hidden_states,
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
-            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
-
-        # compute hidden_states residual
-        hidden_states = hidden_states.contiguous()
-
-        hidden_states_residual = hidden_states - original_hidden_states
-
-        if (
-            encoder_hidden_states is not None
-            and original_encoder_hidden_states is not None
-            and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
-        ):
-            encoder_hidden_states = encoder_hidden_states.contiguous()
-            encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
-        else:
-            encoder_hidden_states_residual = None
-
-        return (
-            hidden_states,
-            encoder_hidden_states,
-            hidden_states_residual,
-            encoder_hidden_states_residual,
-        )
-
-    def compute_or_prune(
-        self,
-        block_id: int,  # Block index in the transformer blocks
-        # Below are the inputs to the block
-        block,  # The transformer block to be executed
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        # NOTE: Although Bagel likely won't use pruning, implementing safe version just in case.
-        # Copy-pasted from original but adding checks.
-
-        original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
-
-        can_use_prune = self._maybe_prune(
-            block_id,
-            hidden_states,
-            prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
-        )
-
-        torch._dynamo.graph_break()
-        if can_use_prune:
-            self.context_manager.add_pruned_step()
-            hidden_states, encoder_hidden_states = self.context_manager.apply_prune(
-                hidden_states,
-                encoder_hidden_states,
-                prefix=(
-                    f"{self.cache_prefix}_{block_id}_Bn_residual"
-                    if self.context_manager.is_cache_residual()
-                    else f"{self.cache_prefix}_Bn_hidden_states"
-                ),
-                encoder_prefix=(
-                    f"{self.cache_prefix}_{block_id}_Bn_encoder_residual"
-                    if self.context_manager.is_encoder_cache_residual()
-                    else f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states"
-                ),
-            )
-            torch._dynamo.graph_break()
-        else:
-            # Normal steps: Compute the block and cache the residuals.
-            hidden_states = block(
-                hidden_states,
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
-            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
-            if not self._skip_prune(block_id):
-                hidden_states = hidden_states.contiguous()
-                hidden_states_residual = hidden_states - original_hidden_states
-
-                if (
-                    encoder_hidden_states is not None
-                    and original_encoder_hidden_states is not None
-                    and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
-                ):
-                    encoder_hidden_states = encoder_hidden_states.contiguous()
-                    encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
-                else:
-                    encoder_hidden_states_residual = None
-
-                self.context_manager.set_Fn_buffer(
-                    original_hidden_states,
-                    prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
-                )
-                if self.context_manager.is_cache_residual():
-                    self.context_manager.set_Bn_buffer(
-                        hidden_states_residual,
-                        prefix=f"{self.cache_prefix}_{block_id}_Bn_residual",
-                    )
-                else:
-                    self.context_manager.set_Bn_buffer(
-                        hidden_states,
-                        prefix=f"{self.cache_prefix}_{block_id}_Bn_hidden_states",
-                    )
-                if encoder_hidden_states_residual is not None:
-                    if self.context_manager.is_encoder_cache_residual():
-                        self.context_manager.set_Bn_encoder_buffer(
-                            encoder_hidden_states_residual,
-                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_residual",
-                        )
-                    else:
-                        self.context_manager.set_Bn_encoder_buffer(
-                            encoder_hidden_states_residual,
-                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states",
-                        )
-            torch._dynamo.graph_break()
-
-        return hidden_states, encoder_hidden_states
-
-
-class Wan22S2VCachedBlocks(CachedBlocks_Pattern_3_4_5):
-    """CacheDiT blocks wrapper that preserves S2V per-layer audio injection."""
-
-    def _run_block(self, block_id: int, block: torch.nn.Module, hidden_states: torch.Tensor, *args, **kwargs):
-        hidden_states = block(hidden_states, *args, **kwargs)
-        hidden_states, new_encoder_hidden_states = self._process_block_outputs(hidden_states)
-        original_after_transformer_block = getattr(
-            self.transformer,
-            "_cache_dit_original_after_transformer_block",
-            getattr(self.transformer, "after_transformer_block", None),
-        )
-        if original_after_transformer_block is not None:
-            hidden_states = original_after_transformer_block(block_id, hidden_states)
-        return hidden_states, new_encoder_hidden_states
-
-    def call_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        new_encoder_hidden_states = None
-        for block_id, block in enumerate(self.transformer_blocks):
-            hidden_states, new_encoder_hidden_states = self._run_block(
-                block_id,
-                block,
-                hidden_states,
-                *args,
-                **kwargs,
-            )
-        return hidden_states, new_encoder_hidden_states
-
-    def call_Fn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        new_encoder_hidden_states = None
-        for block_id, block in enumerate(self._Fn_blocks()):
-            hidden_states, new_encoder_hidden_states = self._run_block(
-                block_id,
-                block,
-                hidden_states,
-                *args,
-                **kwargs,
-            )
-        return hidden_states, new_encoder_hidden_states
-
-    def call_Mn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        original_hidden_states = hidden_states
-        new_encoder_hidden_states = None
-        start_idx = self.context_manager.Fn_compute_blocks()
-        for block_id, block in enumerate(self._Mn_blocks(), start=start_idx):
-            hidden_states, new_encoder_hidden_states = self._run_block(
-                block_id,
-                block,
-                hidden_states,
-                *args,
-                **kwargs,
-            )
-
-        hidden_states = hidden_states.contiguous()
-        hidden_states_residual = hidden_states - original_hidden_states.to(hidden_states.device)
-
-        return (
-            hidden_states,
-            new_encoder_hidden_states,
-            hidden_states_residual,
-        )
-
-    def call_Bn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        new_encoder_hidden_states = None
-        if self.context_manager.Bn_compute_blocks() == 0:
-            return hidden_states, new_encoder_hidden_states
-
-        start_idx = len(self.transformer_blocks) - self.context_manager.Bn_compute_blocks()
-        for block_id, block in enumerate(self._Bn_blocks(), start=start_idx):
-            hidden_states, new_encoder_hidden_states = self._run_block(
-                block_id,
-                block,
-                hidden_states,
-                *args,
-                **kwargs,
-            )
-
-        return hidden_states, new_encoder_hidden_states
-
-
-class Wan22S2VCachedAdapter(CachedAdapter):
-    """CacheDiT adapter that uses Wan22S2VCachedBlocks for S2V audio injection.
-
-    Only overrides collect_unified_blocks to use Wan22S2VCachedBlocks (which
-    calls after_transformer_block per-layer internally). The base class
-    mock_transformer handles the forward wrapping — after_transformer_block is
-    permanently replaced with a no-op in enable_cache_for_wan22_s2v() to prevent
-    double injection.
-    """
-
-    @classmethod
-    def collect_unified_blocks(
-        cls,
-        block_adapter: BlockAdapter,
-        contexts_kwargs: list[dict],
-    ) -> list[dict[str, torch.nn.ModuleList]]:
-        BlockAdapter.assert_normalized(block_adapter)
-
-        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
-        assert hasattr(block_adapter.pipe, "_context_manager")
-
-        for i in range(len(block_adapter.transformer)):
-            unified_blocks_bind_context = {}
-            for j in range(len(block_adapter.blocks[i])):
-                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
-                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
-                    [
-                        Wan22S2VCachedBlocks(
-                            block_adapter.blocks[i][j],
-                            transformer=block_adapter.transformer[i],
-                            forward_pattern=block_adapter.forward_pattern[i][j],
-                            check_forward_pattern=block_adapter.check_forward_pattern,
-                            check_num_outputs=block_adapter.check_num_outputs,
-                            cache_prefix=block_adapter.blocks_name[i][j],
-                            cache_context=block_adapter.unique_blocks_name[i][j],
-                            context_manager=block_adapter.pipe._context_manager,
-                            cache_type=cache_config.cache_type,
-                        )
-                    ]
-                )
-
-            total_cached_blocks.append(unified_blocks_bind_context)
-
-        return total_cached_blocks
-
-
-class BagelCachedAdapter(CachedAdapter):
-    """
-    Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks.
-    """
-
-    @classmethod
-    def create_context(
-        cls,
-        block_adapter: BlockAdapter,
-        **context_kwargs,
-    ) -> tuple[list[str], list[dict[str, Any]]]:
-        # Override to use BagelCachedContextManager
-
-        BlockAdapter.assert_normalized(block_adapter)
-
-        if BlockAdapter.is_cached(block_adapter.pipe):
-            return block_adapter.pipe
-
-        # Check context_kwargs
-        context_kwargs = cls.check_context_kwargs(block_adapter, **context_kwargs)
-
-        # Each Pipeline should have it's own context manager instance.
-        cache_config: BasicCacheConfig = context_kwargs.get("cache_config", None)
-        assert cache_config is not None, "cache_config can not be None."
-
-        # Apply cache on pipeline: wrap cache context
-        pipe_cls_name = block_adapter.pipe.__class__.__name__
-
-        # USE CUSTOM CONTEXT MANAGER
-        context_manager = BagelCachedContextManager(
-            name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
-            persistent_context=isinstance(block_adapter.pipe, FakeDiffusionPipeline),
-        )
-
-        flatten_contexts, contexts_kwargs = cls.modify_context_params(block_adapter, **context_kwargs)
-
-        block_adapter.pipe._context_manager = context_manager  # instance level
-
-        if not context_manager.persistent_context:
-            original_call = block_adapter.pipe.__class__.__call__
-
-            @functools.wraps(original_call)
-            def new_call(self, *args, **kwargs):
-                with ExitStack() as stack:
-                    # cache context will be reset for each pipe inference
-                    for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
-                        stack.enter_context(
-                            context_manager.enter_context(
-                                context_manager.reset_context(
-                                    context_name,
-                                    **context_kwargs,
-                                ),
-                            )
-                        )
-                    outputs = original_call(self, *args, **kwargs)
-                    cls.apply_stats_hooks(block_adapter)
-                    return outputs
-
-            block_adapter.pipe.__class__.__call__ = new_call
-            block_adapter.pipe.__class__._original_call = original_call
-
-        else:
-            # Init persistent cache context for transformer
-            for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
-                context_manager.reset_context(
-                    context_name,
-                    **context_kwargs,
-                )
-
-        block_adapter.pipe.__class__._is_cached = True
-
-        cls.apply_params_hooks(block_adapter, contexts_kwargs)
-
-        return flatten_contexts, contexts_kwargs
-
-    @classmethod
-    def collect_unified_blocks(
-        cls,
-        block_adapter: BlockAdapter,
-        contexts_kwargs: list[dict],
-    ) -> list[dict[str, torch.nn.ModuleList]]:
-        # Override to use BagelCachedBlocks
-
-        BlockAdapter.assert_normalized(block_adapter)
-
-        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
-        assert hasattr(block_adapter.pipe, "_context_manager")
-        # Skipping isinstance check for ContextManager._supported_managers to avoid import issues
-
-        for i in range(len(block_adapter.transformer)):
-            unified_blocks_bind_context = {}
-            for j in range(len(block_adapter.blocks[i])):
-                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
-
-                # Directly instantiate BagelCachedBlocks
-                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
-                    [
-                        BagelCachedBlocks(
-                            # 0. Transformer blocks configuration
-                            block_adapter.blocks[i][j],
-                            transformer=block_adapter.transformer[i],
-                            forward_pattern=block_adapter.forward_pattern[i][j],
-                            check_forward_pattern=block_adapter.check_forward_pattern,
-                            check_num_outputs=block_adapter.check_num_outputs,
-                            # 1. Cache/Prune context configuration
-                            cache_prefix=block_adapter.blocks_name[i][j],
-                            cache_context=block_adapter.unique_blocks_name[i][j],
-                            context_manager=block_adapter.pipe._context_manager,
-                            cache_type=cache_config.cache_type,
-                        )
-                    ]
-                )
-
-            total_cached_blocks.append(unified_blocks_bind_context)
-
-        return total_cached_blocks
-
-
-def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for Bagel model (via OmniDiffusion pipeline).
-
-    Args:
-        pipeline: The OmniDiffusion pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-
-    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
-    """
-    # Build DBCacheConfig
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    # Build calibrator config if TaylorSeer is enabled
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    # Access the transformer: BagelPipeline -> Qwen2MoTForCausalLM -> Qwen2MoTModel
-    # BagelPipeline has self.language_model which is Qwen2MoTForCausalLM
-    # Qwen2MoTForCausalLM has self.model which is Qwen2MoTModel
-    transformer = pipeline.language_model.model
-
-    logger.info(
-        f"Enabling cache-dit on Bagel transformer: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    # Enable cache-dit on the transformer
-    # Pattern_0 corresponds to (hidden_states, encoder_hidden_states) input, output
-    # Custom adapter for Bagel to handle NaiveCache correctly
-    # from vllm_omni.diffusion.cache.bagel_cache_adapter import BagelCachedAdapter # No longer needed
-    BagelCachedAdapter.apply(
-        BlockAdapter(
-            transformer=transformer,
-            blocks=transformer.layers,
-            forward_pattern=ForwardPattern.Pattern_0,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        transformer = pipeline.language_model.model
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-    return refresh_cache_context
-
-
-class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5):
-    """
-    Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden
-    states during denoising.
-    """
-
-    @classmethod
-    def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool:
-        if kwargs.get("cache_dit_skip", False):
-            return False
-
-        # Prefix/text forwards either omit image_gen_indicators or update the
-        # DynamicCache. Denoising forwards are gen-only and use update_cache=False.
-        if kwargs.get("update_cache", True):
-            return False
-
-        exist_gen = kwargs.get("exist_gen")
-        exist_und = kwargs.get("exist_und")
-        if exist_gen is None or exist_und is None:
-            image_gen_indicators = kwargs.get("image_gen_indicators")
-            if image_gen_indicators is None:
-                return False
-            exist_gen = image_gen_indicators.any().item()
-            exist_und = (~image_gen_indicators).any().item()
-
-        return exist_gen and not exist_und
-
-    @staticmethod
-    def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
-        kwargs = dict(kwargs)
-        kwargs.pop("cache_dit_skip", None)
-        return kwargs
-
-    def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
-        block_kwargs = self._strip_cache_only_kwargs(kwargs)
-        if not self._is_denoising_call(kwargs):
-            hidden_states, new_encoder_hidden_states = self.call_blocks(
+    def call_Fn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        for block_id, block in enumerate(self._Fn_blocks()):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
                 hidden_states,
                 *args,
-                **block_kwargs,
-            )
-            return self._process_forward_outputs(hidden_states, new_encoder_hidden_states)
-
-        return super().forward(hidden_states, *args, **block_kwargs)
-
-
-class SensenovaCachedAdapter(CachedAdapter):
-    """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks."""
-
-    @classmethod
-    def collect_unified_blocks(
-        cls,
-        block_adapter: BlockAdapter,
-        contexts_kwargs: list[dict],
-    ) -> list[dict[str, torch.nn.ModuleList]]:
-        BlockAdapter.assert_normalized(block_adapter)
-
-        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
-        assert hasattr(block_adapter.pipe, "_context_manager")
-
-        for i in range(len(block_adapter.transformer)):
-            unified_blocks_bind_context = {}
-            for j in range(len(block_adapter.blocks[i])):
-                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
-                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
-                    [
-                        SensenovaCachedBlocks(
-                            # 0. Transformer blocks configuration
-                            block_adapter.blocks[i][j],
-                            transformer=block_adapter.transformer[i],
-                            forward_pattern=block_adapter.forward_pattern[i][j],
-                            check_forward_pattern=block_adapter.check_forward_pattern,
-                            check_num_outputs=block_adapter.check_num_outputs,
-                            # 1. Cache/Prune context configuration
-                            cache_prefix=block_adapter.blocks_name[i][j],
-                            cache_context=block_adapter.unique_blocks_name[i][j],
-                            context_manager=block_adapter.pipe._context_manager,
-                            cache_type=cache_config.cache_type,
-                        )
-                    ]
-                )
-
-            total_cached_blocks.append(unified_blocks_bind_context)
-
-        return total_cached_blocks
-
-
-def enable_cache_for_sensenova_u1(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for SenseNova-U1 model.
-
-    Args:
-        pipeline: The SenseNova-U1 pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-
-    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
-    """
-    transformer = getattr(getattr(pipeline, "language_model", None), "model", None)
-    if transformer is None or not hasattr(transformer, "layers"):
-        raise ValueError("SenseNovaU1Pipeline cache-dit expects pipeline.language_model.model.layers.")
-
-    # Build DBCacheConfig
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
-    logger.info(
-        f"Enabling cache-dit on SenseNova-U1 decoder layers: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    SensenovaCachedAdapter.apply(
-        BlockAdapter(
-            transformer=transformer,
-            blocks=transformer.layers,
-            forward_pattern=ForwardPattern.Pattern_3,
-            params_modifiers=[modifier],
-            check_forward_pattern=True,
-            has_separate_cfg=True,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        transformer = pipeline.language_model.model
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
+                **kwargs,
             )
+        return hidden_states, new_encoder_hidden_states
 
-    return refresh_cache_context
-
-
-def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for Flux.2-dev pipeline.
-
-    Args:
-        pipeline: The Flux2 pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-    Returns:
-        A refresh function that can be called with a new ``num_inference_steps``
-        to update the cache context for the pipeline.
-    """
-    # Build DBCacheConfig for transformer
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    # Build ParamsModifier for transformer
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator,
-    )
-
-    logger.info(
-        f"Enabling cache-dit on Flux transformer with BlockAdapter: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    # Enable cache-dit using BlockAdapter for transformer
-    cache_dit.enable_cache(
-        (
-            BlockAdapter(
-                transformer=pipeline.transformer,
-                blocks=[
-                    pipeline.transformer.transformer_blocks,
-                    pipeline.transformer.single_transformer_blocks,
-                ],
-                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2],
-                params_modifiers=[modifier],
+    def call_Mn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        new_encoder_hidden_states = None
+        start_idx = self.context_manager.Fn_compute_blocks()
+        for block_id, block in enumerate(self._Mn_blocks(), start=start_idx):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
             )
-        ),
-        cache_config=db_cache_config,
-    )
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
+        hidden_states = hidden_states.contiguous()
+        hidden_states_residual = hidden_states - original_hidden_states.to(hidden_states.device)
 
-        Args:
-            pipeline: The Flux2 pipeline instance.
-            num_inference_steps: New number of inference steps.
-        """
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
+        return (
+            hidden_states,
+            new_encoder_hidden_states,
+            hidden_states_residual,
+        )
+
+    def call_Bn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        if self.context_manager.Bn_compute_blocks() == 0:
+            return hidden_states, new_encoder_hidden_states
+
+        start_idx = len(self.transformer_blocks) - self.context_manager.Bn_compute_blocks()
+        for block_id, block in enumerate(self._Bn_blocks(), start=start_idx):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
             )
 
-    return refresh_cache_context
+        return hidden_states, new_encoder_hidden_states
 
 
-def enable_cache_for_glm_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for GlmImage pipeline.
+class Wan22S2VCachedAdapter(CachedAdapter):
+    """CacheDiT adapter that uses Wan22S2VCachedBlocks for S2V audio injection.
 
-    Args:
-        pipeline: The GlmImage pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-    Returns:
-        A refresh function that can be called with a new ``num_inference_steps``
-        to update the cache context for the pipeline.
+    Only overrides collect_unified_blocks to use Wan22S2VCachedBlocks (which
+    calls after_transformer_block per-layer internally). The base class
+    mock_transformer handles the forward wrapping — after_transformer_block is
+    permanently replaced with a no-op in enable_cache_for_wan22_s2v() to prevent
+    double injection.
     """
-    # Build DBCacheConfig for transformer
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
 
-    # Build ParamsModifier for transformer
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator,
-    )
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        BlockAdapter.assert_normalized(block_adapter)
 
-    logger.info(
-        f"Enabling cache-dit on GlmImage transformer with BlockAdapter: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
 
-    # Enable cache-dit using BlockAdapter for transformer
-    # Note: We don't use patch_functor here because it's designed for diffusers' GlmImage,
-    # and our vllm-omni implementation has a different forward signature.
-    # We use ForwardPattern.Pattern_0 because our block returns (hidden_states, encoder_hidden_states)
-    cache_dit.enable_cache(
-        (
-            BlockAdapter(
-                transformer=pipeline.transformer,
-                blocks=pipeline.transformer.transformer_blocks,
-                forward_pattern=ForwardPattern.Pattern_0,
-                params_modifiers=[modifier],
-                patch_functor=None,
-                has_separate_cfg=True,
-            )
-        ),
-        cache_config=db_cache_config,
-    )
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        Wan22S2VCachedBlocks(
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        """Refresh cache context for the transformer with new num_inference_steps.
+            total_cached_blocks.append(unified_blocks_bind_context)
 
-        Args:
-            pipeline: The GlmImage pipeline instance.
-            num_inference_steps: New number of inference steps.
-        """
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
+        return total_cached_blocks
 
-    return refresh_cache_context
 
+class BagelCachedAdapter(CachedAdapter):
+    """
+    Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks.
+    """
 
-def enable_cache_for_dreamid_omni(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for DreamID-Omni fused pipeline.
+    @classmethod
+    def create_context(
+        cls,
+        block_adapter: BlockAdapter,
+        **context_kwargs,
+    ) -> tuple[list[str], list[dict[str, Any]]]:
+        # Override to use BagelCachedContextManager
 
-    Args:
-        pipeline: The DreamIDOmni pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
+        BlockAdapter.assert_normalized(block_adapter)
 
-    Returns:
-        A refresh function that can be called with a new ``num_inference_steps``
-        to update the cache context for the pipeline.
-    """
-    transformer = getattr(pipeline, "transformer", None)
-    fused_blocks = getattr(transformer, "fused_blocks", None)
-    if transformer is None or fused_blocks is None:
-        raise ValueError("DreamIDOmniPipeline cache-dit expects pipeline.transformer.fused_blocks.")
+        if BlockAdapter.is_cached(block_adapter.pipe):
+            return block_adapter.pipe
 
-    db_cache_config = _build_db_cache_config(cache_config)
+        # Check context_kwargs
+        context_kwargs = cls.check_context_kwargs(block_adapter, **context_kwargs)
 
-    calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+        # Each Pipeline should have it's own context manager instance.
+        cache_config: BasicCacheConfig = context_kwargs.get("cache_config", None)
+        assert cache_config is not None, "cache_config can not be None."
 
-    modifier = ParamsModifier(
-        cache_config=db_cache_config,
-        calibrator_config=calibrator,
-    )
+        # Apply cache on pipeline: wrap cache context
+        pipe_cls_name = block_adapter.pipe.__class__.__name__
 
-    logger.info(
-        f"Enabling cache-dit on DreamID-Omni fused blocks: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
+        # USE CUSTOM CONTEXT MANAGER
+        context_manager = BagelCachedContextManager(
+            name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
+            persistent_context=isinstance(block_adapter.pipe, FakeDiffusionPipeline),
+        )
 
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=transformer,
-            blocks=fused_blocks,
-            forward_pattern=ForwardPattern.Pattern_0,
-            params_modifiers=[modifier],
-            has_separate_cfg=True,
-        ),
-        cache_config=db_cache_config,
-    )
+        flatten_contexts, contexts_kwargs = cls.modify_context_params(block_adapter, **context_kwargs)
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        transformer = pipeline.transformer
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
+        block_adapter.pipe._context_manager = context_manager  # instance level
 
-    return refresh_cache_context
+        if not context_manager.persistent_context:
+            original_call = block_adapter.pipe.__class__.__call__
 
+            @functools.wraps(original_call)
+            def new_call(self, *args, **kwargs):
+                with ExitStack() as stack:
+                    # cache context will be reset for each pipe inference
+                    for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
+                        stack.enter_context(
+                            context_manager.enter_context(
+                                context_manager.reset_context(
+                                    context_name,
+                                    **context_kwargs,
+                                ),
+                            )
+                        )
+                    outputs = original_call(self, *args, **kwargs)
+                    cls.apply_stats_hooks(block_adapter)
+                    return outputs
 
-def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for ERNIE-Image pipeline.
+            block_adapter.pipe.__class__.__call__ = new_call
+            block_adapter.pipe.__class__._original_call = original_call
 
-    ERNIE-Image blocks have signature:
-        forward(x, rotary_pos_emb, temb, attention_mask) -> x
+        else:
+            # Init persistent cache context for transformer
+            for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
+                context_manager.reset_context(
+                    context_name,
+                    **context_kwargs,
+                )
 
-    Where x is hidden_states (concatenated image + text tokens).
-    This matches Pattern_3 which expects:
-        - Input: hidden_states only
-        - Output: hidden_states only
+        block_adapter.pipe.__class__._is_cached = True
 
-    Args:
-        pipeline: The ERNIE-Image pipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
-    """
-    db_cache_config = _build_db_cache_config(cache_config)
+        cls.apply_params_hooks(block_adapter, contexts_kwargs)
 
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+        return flatten_contexts, contexts_kwargs
 
-    logger.info(
-        f"Enabling cache-dit on ERNIE-Image transformer: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        # Override to use BagelCachedBlocks
 
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=pipeline.transformer,
-            blocks=pipeline.transformer.layers,
-            forward_pattern=ForwardPattern.Pattern_3,
-            check_forward_pattern=False,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
+        BlockAdapter.assert_normalized(block_adapter)
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+        # Skipping isinstance check for ContextManager._supported_managers to avoid import issues
 
-    return refresh_cache_context
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
 
+                # Directly instantiate BagelCachedBlocks
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        BagelCachedBlocks(
+                            # 0. Transformer blocks configuration
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            # 1. Cache/Prune context configuration
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
 
-def enable_cache_for_helios(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for Helios pipeline.
+            total_cached_blocks.append(unified_blocks_bind_context)
 
-    Helios extends Wan2.2 with multi-term memory patches and guidance cross-attention.
-    Its transformer blocks have the same single-output signature as Wan2.2 blocks
-    (returns hidden_states only), so we use ForwardPattern.Pattern_2.
+        return total_cached_blocks
 
-    Args:
-        pipeline: The HeliosPipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
 
-    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
+class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5):
+    """
+    Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden
+    states during denoising.
     """
-    db_cache_config = _build_db_cache_config(cache_config)
 
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+    @classmethod
+    def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool:
+        if kwargs.get("cache_dit_skip", False):
+            return False
 
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=pipeline.transformer,
-            blocks=[pipeline.transformer.blocks],
-            forward_pattern=[ForwardPattern.Pattern_2],
-            params_modifiers=[
-                ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
-            ],
-            has_separate_cfg=True,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
+        # Prefix/text forwards either omit image_gen_indicators or update the
+        # DynamicCache. Denoising forwards are gen-only and use update_cache=False.
+        if kwargs.get("update_cache", True):
+            return False
 
-    logger.info(
-        f"Enabling cache-dit on Helios transformer: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
+        exist_gen = kwargs.get("exist_gen")
+        exist_und = kwargs.get("exist_und")
+        if exist_gen is None or exist_und is None:
+            image_gen_indicators = kwargs.get("image_gen_indicators")
+            if image_gen_indicators is None:
+                return False
+            exist_gen = image_gen_indicators.any().item()
+            exist_und = (~image_gen_indicators).any().item()
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
+        return exist_gen and not exist_und
 
-    return refresh_cache_context
+    @staticmethod
+    def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+        kwargs = dict(kwargs)
+        kwargs.pop("cache_dit_skip", None)
+        return kwargs
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
+        block_kwargs = self._strip_cache_only_kwargs(kwargs)
+        if not self._is_denoising_call(kwargs):
+            hidden_states, new_encoder_hidden_states = self.call_blocks(
+                hidden_states,
+                *args,
+                **block_kwargs,
+            )
+            return self._process_forward_outputs(hidden_states, new_encoder_hidden_states)
 
+        return super().forward(hidden_states, *args, **block_kwargs)
 
-def enable_cache_for_hunyuan_video_15(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for HunyuanVideo 1.5 pipeline.
 
-    HunyuanVideo 1.5 uses a single transformer with has_separate_cfg=True
-    (separate conditional/unconditional forward passes). The _sp_plan scatter
-    is applied at the transformer input boundary (empty-string key), so CacheDiT
-    sees already-sharded hidden_states throughout transformer_blocks.
-    """
-    db_cache_config = _build_db_cache_config(cache_config)
+class SensenovaCachedAdapter(CachedAdapter):
+    """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks."""
 
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        BlockAdapter.assert_normalized(block_adapter)
 
-    logger.info(
-        f"Enabling cache-dit on HunyuanVideo15 transformer: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
 
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=pipeline.transformer,
-            blocks=pipeline.transformer.transformer_blocks,
-            forward_pattern=ForwardPattern.Pattern_0,
-            check_forward_pattern=True,
-            has_separate_cfg=True,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        SensenovaCachedBlocks(
+                            # 0. Transformer blocks configuration
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            # 1. Cache/Prune context configuration
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
+            total_cached_blocks.append(unified_blocks_bind_context)
 
-    return refresh_cache_context
+        return total_cached_blocks
 
 
-def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc:
     """Enable cache-dit for Cosmos3.
 
     Cosmos3 has a dual-pathway architecture (UND + GEN) but only the GEN
@@ -1959,41 +948,6 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int]
     Returns:
         A refresh function that can be called to update cache context with new num_inference_steps.
     """
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    logger.info(
-        f"Enabling cache-dit on Cosmos3 gen_layers: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=pipeline.transformer,
-            blocks=[pipeline.transformer.gen_layers],
-            # Cosmos3 GEN blocks return only hidden_states.  Per-layer UND K/V
-            # conditioning uses the transformer's cache-dit fallback path.
-            forward_pattern=[ForwardPattern.Pattern_3],
-            params_modifiers=[
-                ParamsModifier(
-                    cache_config=db_cache_config,
-                    calibrator_config=calibrator_config,
-                ),
-            ],
-            check_forward_pattern=False,
-            has_separate_cfg=True,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
     # The T2I denoising loop skips the unconditional forward outside the
     # guidance interval as a speed optimization. cache-dit distinguishes the
     # conditional vs unconditional passes purely by transformer-forward parity
@@ -2001,25 +955,8 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int]
     # step accounting. Still do both cond/uncond CFG steps when cache-dit is active.
     # CFG is instead neutralized via scale=1.0 outside the interval.
     pipeline._cache_dit_requires_paired_cfg = True
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy,
-                        total_steps=num_inference_steps,
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-    return refresh_cache_context
+    block_adapter = CacheDiTBackend.maybe_build_block_adapter(pipeline)
+    return enable_cache_for_dit(pipeline, cache_config, block_adapter)
 
 
 # Register custom cache-dit enablers after function definitions
@@ -2027,29 +964,8 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool
     {
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
+        "Wan22TI2VPipeline": enable_cache_for_wan22,
         "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
-        "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3,
-        "FluxPipeline": enable_cache_for_flux,
-        "Flux2KleinPipeline": enable_cache_for_flux2_klein,
-        "LongCatImagePipeline": enable_cache_for_longcat_image,
-        "LongCatImageEditPipeline": enable_cache_for_longcat_image,
-        "StableAudioPipeline": enable_cache_for_stable_audio_open,
-        "StableDiffusion3Pipeline": enable_cache_for_sd3,
-        "LTX2Pipeline": enable_cache_for_ltx2,
-        "LTX2ImageToVideoPipeline": enable_cache_for_ltx2,
-        "LTX2TwoStagesPipeline": enable_cache_for_ltx2,
-        "LTX2ImageToVideoTwoStagesPipeline": enable_cache_for_ltx2,
-        "LTX23Pipeline": enable_cache_for_ltx2,
-        "LTX23ImageToVideoPipeline": enable_cache_for_ltx2,
-        "BagelPipeline": enable_cache_for_bagel,
-        "SenseNovaU1Pipeline": enable_cache_for_sensenova_u1,
-        "GlmImagePipeline": enable_cache_for_glm_image,
-        "Flux2Pipeline": enable_cache_for_flux2,
-        "DreamIDOmniPipeline": enable_cache_for_dreamid_omni,
-        "ErnieImagePipeline": enable_cache_for_ernie_image,
-        "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15,
-        "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15,
-        "HeliosPipeline": enable_cache_for_helios,
         "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
     }
 )
@@ -2092,6 +1008,43 @@ def __init__(self, cache_config: Any = None):
         self._refresh_func: Callable[[Any, int, bool], None] | None = None
         self._last_num_inference_steps: int | None = None
 
+    @staticmethod
+    def maybe_build_block_adapter(pipeline) -> BlockAdapter | None:
+        """If a module defines `_cache_dit_adapter_config`, build the corresponding
+        block adapter.
+        """
+        transformer = default_get_pipeline_transformer(pipeline)
+
+        adapter_cfg: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None)
+        if adapter_cfg is None:
+            return None
+
+        block_attrs, forward_pattern = zip(*(adapter_cfg.block_forward_patterns).items())
+        missing_attrs = [block_attr for block_attr in block_attrs if not hasattr(transformer, block_attr)]
+
+        if missing_attrs:
+            logger.warning("Missing Cache DiT block attributes: %s", missing_attrs)
+
+        block_adapter = BlockAdapter(
+            transformer=transformer,
+            blocks=[getattr(transformer, block_attr) for block_attr in block_attrs],
+            forward_pattern=list(forward_pattern),
+            has_separate_cfg=adapter_cfg.has_separate_cfg,
+            check_forward_pattern=adapter_cfg.check_forward_pattern,
+        )
+        return block_adapter
+
+    @staticmethod
+    def maybe_get_cached_adapter_cls(pipeline) -> type[CachedAdapter] | None:
+        """If a module has a custom cached adapter type registered, e.g., SenseNova, retrieve it
+        from the transformer's CacheDiTAdapterConfig."""
+        transformer = default_get_pipeline_transformer(pipeline)
+
+        adapter_cfg: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None)
+        if adapter_cfg is None:
+            return None
+        return adapter_cfg.cached_adapter_cls
+
     def enable(self, pipeline: Any) -> None:
         """Enable cache-dit on the pipeline if configured.
 
@@ -2110,8 +1063,12 @@ def enable(self, pipeline: Any) -> None:
             logger.info(f"Using custom cache-dit enabler for model: {pipeline_name}")
             self._refresh_func = CUSTOM_DIT_ENABLERS[pipeline_name](pipeline, self.config)
         else:
-            # For regular single-transformer models
-            self._refresh_func = enable_cache_for_dit(pipeline, self.config)
+            # Common case; either the model doesn't explicitly support dit cache yet,
+            # Or it defines its _cache_dit_adapter_config, which describes how we should
+            # create its block adapter.
+            block_adapter = self.maybe_build_block_adapter(pipeline)
+            adapter_cls = self.maybe_get_cached_adapter_cls(pipeline)
+            self._refresh_func = enable_cache_for_dit(pipeline, self.config, block_adapter, adapter_cls)
 
         self.enabled = True
         logger.info(f"Cache-dit enabled successfully on {pipeline_name}")
diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py
index e6c2e70e318..3b002abe0bf 100644
--- a/vllm_omni/diffusion/cache/selector.py
+++ b/vllm_omni/diffusion/cache/selector.py
@@ -3,6 +3,7 @@
 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
 from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
+from vllm_omni.diffusion.cache.stepcache import StepCacheBackend
 from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
@@ -14,14 +15,15 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
     - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface
     - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface
     - mag_cache: Uses MagCacheBackend with enable()/refresh() interface
+    - step_cache: Uses StepCacheBackend for DreamZero velocity step skip
 
     Args:
-        cache_backend: Cache backend name ("cache_dit", "tea_cache", "mag_cache", or None).
+        cache_backend: Cache backend name ("cache_dit", "tea_cache",
+            "mag_cache", "step_cache", or None).
         cache_config: Cache configuration (dict or DiffusionCacheConfig instance).
 
     Returns:
-        Cache backend instance (CacheDiTBackend, TeaCacheBackend, or MagCacheBackend)
-        if cache_backend is set, None otherwise.
+        Cache backend instance if cache_backend is set, None otherwise.
 
     Raises:
         ValueError: If cache_backend is unsupported.
@@ -38,7 +40,10 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
         return TeaCacheBackend(cache_config)
     elif cache_backend == "mag_cache":
         return MagCacheBackend(cache_config)
+    elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"):
+        return StepCacheBackend(cache_config)
     else:
         raise ValueError(
-            f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache', 'mag_cache'"
+            f"Unsupported cache backend: {cache_backend}. "
+            "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'"
         )
diff --git a/vllm_omni/diffusion/cache/stepcache/__init__.py b/vllm_omni/diffusion/cache/stepcache/__init__.py
new file mode 100644
index 00000000000..a7e1852b0d4
--- /dev/null
+++ b/vllm_omni/diffusion/cache/stepcache/__init__.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+StepCache: velocity-based step skipping for DreamZero-style DiT denoising.
+
+Unlike block-level Cache-DiT, this backend skips entire DiT forwards during the
+pipeline denoise loop when successive velocity predictions are highly aligned.
+
+Usage:
+    from vllm_omni import Omni
+
+    omni = Omni(
+        model="...",
+        cache_backend="step_cache",
+    )
+
+    # Alternative: environment variable
+    # export DIFFUSION_CACHE_BACKEND=step_cache
+"""
+
+from vllm_omni.diffusion.cache.stepcache.backend import (
+    CUSTOM_STEP_CACHE_DIT_ENABLERS,
+    CUSTOM_STEPCACHE_ENABLERS,
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_DIT_CONFIG_ATTR,
+    STEP_CACHE_DIT_STATE_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheBackend,
+    StepCacheDiTBackend,
+    StepCacheDiTConfig,
+    StepCacheDiTState,
+    enable_dreamzero_step_cache_dit,
+    enable_dreamzero_stepcache,
+    get_step_cache_dit_config,
+    get_step_cache_dit_state,
+    get_stepcache_config,
+    get_stepcache_state,
+    is_step_cache_dit_active,
+    is_stepcache_active,
+)
+from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
+from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
+
+__all__ = [
+    "CUSTOM_STEP_CACHE_DIT_ENABLERS",
+    "CUSTOM_STEPCACHE_ENABLERS",
+    "STEP_CACHE_CONFIG_ATTR",
+    "STEP_CACHE_DIT_CONFIG_ATTR",
+    "STEP_CACHE_DIT_STATE_ATTR",
+    "STEP_CACHE_STATE_ATTR",
+    "StepCacheBackend",
+    "StepCacheConfig",
+    "StepCacheDiTBackend",
+    "StepCacheDiTConfig",
+    "StepCacheDiTState",
+    "StepCacheState",
+    "enable_dreamzero_step_cache_dit",
+    "enable_dreamzero_stepcache",
+    "get_step_cache_dit_config",
+    "get_step_cache_dit_state",
+    "get_stepcache_config",
+    "get_stepcache_state",
+    "is_step_cache_dit_active",
+    "is_stepcache_active",
+]
diff --git a/vllm_omni/diffusion/cache/stepcache/backend.py b/vllm_omni/diffusion/cache/stepcache/backend.py
new file mode 100644
index 00000000000..517d2369a3c
--- /dev/null
+++ b/vllm_omni/diffusion/cache/stepcache/backend.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+StepCache backend implementation.
+
+This module provides the stepcache backend that implements the CacheBackend
+interface for DreamZero-style velocity-based step skipping in the denoise loop.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheBackend
+from vllm_omni.diffusion.cache.stepcache.config import (
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheConfig,
+)
+from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+logger = init_logger(__name__)
+
+# Backward-compatible aliases for older imports and pipeline attribute names.
+STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR
+STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR
+StepCacheDiTConfig = StepCacheConfig
+StepCacheDiTState = StepCacheState
+
+
+def is_stepcache_active(pipeline: Any) -> bool:
+    """Return True when stepcache is enabled on *pipeline*."""
+    cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None)
+    return isinstance(cfg, StepCacheConfig) and cfg.enabled
+
+
+def get_stepcache_config(pipeline: Any) -> StepCacheConfig | None:
+    cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None)
+    return cfg if isinstance(cfg, StepCacheConfig) else None
+
+
+def get_stepcache_state(pipeline: Any) -> StepCacheState | None:
+    state = getattr(pipeline, STEP_CACHE_STATE_ATTR, None)
+    return state if isinstance(state, StepCacheState) else None
+
+
+def is_step_cache_dit_active(pipeline: Any) -> bool:
+    return is_stepcache_active(pipeline)
+
+
+def get_step_cache_dit_config(pipeline: Any) -> StepCacheConfig | None:
+    return get_stepcache_config(pipeline)
+
+
+def get_step_cache_dit_state(pipeline: Any) -> StepCacheState | None:
+    return get_stepcache_state(pipeline)
+
+
+def enable_dreamzero_stepcache(pipeline: Any, config: DiffusionCacheConfig) -> None:
+    """Enable stepcache for DreamZeroPipeline."""
+    cache_config = StepCacheConfig.from_diffusion_cache_config(config)
+    setattr(pipeline, STEP_CACHE_CONFIG_ATTR, cache_config)
+    setattr(pipeline, STEP_CACHE_STATE_ATTR, StepCacheState(cache_config))
+
+    logger.info(
+        "stepcache enabled on DreamZeroPipeline (enabled=%s, thresholds=%s, countdowns=%s)",
+        cache_config.enabled,
+        cache_config.sim_thresholds,
+        cache_config.skip_countdowns,
+    )
+
+
+enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache
+
+CUSTOM_STEPCACHE_ENABLERS = {
+    "DreamZeroPipeline": enable_dreamzero_stepcache,
+}
+
+CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS
+
+
+class StepCacheBackend(CacheBackend):
+    """
+    Velocity cosine step-skipping cache backend for DreamZero.
+
+    Attaches :class:`StepCacheConfig` and :class:`StepCacheState` to supported
+    pipelines. The denoise loop calls :meth:`StepCacheState.should_run_step` to
+    decide whether to run ``predict_noise``.
+
+    Example:
+        >>> from vllm_omni.diffusion.data import DiffusionCacheConfig
+        >>> backend = StepCacheBackend(DiffusionCacheConfig())
+        >>> backend.enable(pipeline)
+        >>> backend.refresh(pipeline, num_inference_steps=16)
+    """
+
+    def enable(self, pipeline: Any) -> None:
+        """Enable stepcache on a supported pipeline."""
+        pipeline_type = pipeline.__class__.__name__
+
+        if pipeline_type in CUSTOM_STEPCACHE_ENABLERS:
+            logger.info("Using custom stepcache enabler for model: %s", pipeline_type)
+            CUSTOM_STEPCACHE_ENABLERS[pipeline_type](pipeline, self.config)
+        else:
+            raise ValueError(
+                f"step_cache backend does not support {pipeline_type}. Supported: {sorted(CUSTOM_STEPCACHE_ENABLERS)}"
+            )
+
+        self.enabled = True
+
+    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh stepcache state for a new generation."""
+        state = get_stepcache_state(pipeline)
+        if state is not None:
+            state.reset()
+            if verbose:
+                logger.debug(
+                    "stepcache state refreshed (num_inference_steps=%d)",
+                    num_inference_steps,
+                )
+        elif verbose and is_stepcache_active(pipeline):
+            logger.warning("stepcache config active but state not found on pipeline")
+
+
+StepCacheDiTBackend = StepCacheBackend
diff --git a/vllm_omni/diffusion/cache/stepcache/config.py b/vllm_omni/diffusion/cache/stepcache/config.py
new file mode 100644
index 00000000000..00a7984ae7d
--- /dev/null
+++ b/vllm_omni/diffusion/cache/stepcache/config.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Configuration for step-level DiT velocity caching."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+# Pipeline attributes set by StepCacheBackend.enable().
+STEP_CACHE_CONFIG_ATTR = "_stepcache_config"
+STEP_CACHE_STATE_ATTR = "_stepcache_state"
+
+
+@dataclass(frozen=True)
+class StepCacheConfig:
+    """Runtime config for velocity-based step skipping.
+
+    Skips entire DiT forwards when successive video velocity predictions are
+    highly aligned (cosine similarity above configured thresholds).
+
+    Reference: DreamZero paper DiT Caching / dreamzero.git ``should_run_model``.
+    """
+
+    enabled: bool = True
+    min_history_steps: int = 2
+    max_history: int = 2
+    sim_thresholds: tuple[float, ...] = (0.95, 0.93)
+    skip_countdowns: tuple[int, ...] = (4, 2)
+
+    @classmethod
+    def from_diffusion_cache_config(cls, config: DiffusionCacheConfig) -> StepCacheConfig:
+        enabled = config.step_cache_dit_enabled
+
+        thresholds = tuple(config.velocity_sim_thresholds)
+        countdowns = tuple(config.velocity_skip_countdowns)
+        if len(thresholds) != len(countdowns):
+            raise ValueError(
+                "velocity_sim_thresholds and velocity_skip_countdowns must have the same length; "
+                f"got {len(thresholds)} and {len(countdowns)}."
+            )
+
+        return cls(
+            enabled=enabled,
+            min_history_steps=config.step_cache_dit_min_history,
+            max_history=config.step_cache_dit_max_history,
+            sim_thresholds=thresholds,
+            skip_countdowns=countdowns,
+        )
diff --git a/vllm_omni/diffusion/cache/stepcache/state.py b/vllm_omni/diffusion/cache/stepcache/state.py
new file mode 100644
index 00000000000..db786d4f210
--- /dev/null
+++ b/vllm_omni/diffusion/cache/stepcache/state.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Mutable state for step-level DiT velocity caching."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
+
+
+class StepCacheState:
+    """Per-generation mutable state (skip countdown)."""
+
+    def __init__(self, config: StepCacheConfig) -> None:
+        self.config = config
+        self.skip_countdown = 0
+
+    def reset(self) -> None:
+        self.skip_countdown = 0
+
+    def should_run_step(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> bool:
+        """Return True when the DiT forward should execute at this scheduler step."""
+        if not self.config.enabled:
+            return True
+
+        if len(prev_predictions) < self.config.min_history_steps:
+            return True
+
+        if self.skip_countdown > 1:
+            self.skip_countdown -= 1
+            return False
+        if self.skip_countdown == 1:
+            self.skip_countdown = 0
+            return True
+
+        v_last = prev_predictions[-1][0].flatten(1).float()
+        v_prev = prev_predictions[-2][0].flatten(1).float()
+        sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean()
+
+        for threshold, countdown in zip(self.config.sim_thresholds, self.config.skip_countdowns):
+            if sim > threshold:
+                self.skip_countdown = countdown
+                return False
+
+        return True
+
+    def trim_history(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> None:
+        while len(prev_predictions) > self.config.max_history:
+            prev_predictions.pop(0)
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 9f6b6df1c4b..ec0f45fb9ae 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -302,6 +302,8 @@ class DiffusionCacheConfig:
                     scm_steps_mask_policy, scm_steps_policy
         - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio,
                     mag_ratios, mag_calibrate
+        - step_cache: step_cache_dit_enabled, velocity_sim_thresholds,
+                          velocity_skip_countdowns, step_cache_dit_min_history
 
     Example:
         >>> # From dict (user-facing API) - partial config uses defaults for missing keys
@@ -375,6 +377,13 @@ class DiffusionCacheConfig:
     # "repeat" refreshes every force_refresh_step_hint steps.
     force_refresh_step_policy: str = "once"
 
+    # step_cache parameters [step_cache only] — DreamZero velocity schedule
+    step_cache_dit_enabled: bool = True
+    velocity_sim_thresholds: list[float] = field(default_factory=lambda: [0.95, 0.93])
+    velocity_skip_countdowns: list[int] = field(default_factory=lambda: [4, 2])
+    step_cache_dit_min_history: int = 2
+    step_cache_dit_max_history: int = 2
+
     # Additional parameters that may be passed but not explicitly defined
     _extra_params: dict[str, Any] = field(default_factory=dict, repr=False)
 
@@ -1110,6 +1119,10 @@ def enrich_config(self) -> None:
                         self.model_class_name = "WanS2VPipeline"
                     self.tf_model_config = TransformerConfig()
                     self.update_multimodal_support()
+                elif model_type == "Gr00tN1d7" or "Gr00tN1d7" in architectures:
+                    self.model_class_name = "Gr00tN1d7Pipeline"
+                    self.set_tf_model_config(TransformerConfig())
+                    self.update_multimodal_support()
                 elif model_type == "vla":
                     from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 
diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py
index 42675c7c095..3a056d6d45b 100644
--- a/vllm_omni/diffusion/distributed/hsdp.py
+++ b/vllm_omni/diffusion/distributed/hsdp.py
@@ -140,8 +140,13 @@ def apply_hsdp_to_model(
         fs_rank,
     )
 
+    # When the model contains FP8 parameters (online quantization), let FSDP
+    # keep the original storage dtype on all-gather instead of casting to
+    # hsdp_config.param_dtype (typically bfloat16). FP8 GEMM kernels expect
+    # FP8 inputs; an implicit FP8 -> bf16 cast would silently break them.
+    has_fp8_params = any(p.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for p in model.parameters())
     mp_policy = MixedPrecisionPolicy(
-        param_dtype=hsdp_config.param_dtype,
+        param_dtype=None if has_fp8_params else hsdp_config.param_dtype,
         reduce_dtype=hsdp_config.reduce_dtype,
         output_dtype=hsdp_config.output_dtype,
         cast_forward_inputs=False,
diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py
index b2339843ed8..1aa40aae5cf 100644
--- a/vllm_omni/diffusion/model_loader/diffusers_loader.py
+++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py
@@ -602,10 +602,17 @@ def _init_from_load_format(
             model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
             with set_current_diffusion_config(self.od_config):
                 model = model_cls(od_config=self.od_config)
-            if not is_hsdp and target_device.type != "cpu":
+            # HSDP normally defers GPU placement to apply_hsdp_to_model to keep peak
+            # load-time memory on CPU. Online quantization (e.g. fp8) runs CUDA-only
+            # kernels inside load_weights via the layerwise loader, so when a quant
+            # config is set we initialize on the accelerator like the non-HSDP path;
+            # apply_hsdp_to_model shards GPU-resident params equally well.
+            hsdp_defer_to_cpu = is_hsdp and self.quant_config is None
+            if not hsdp_defer_to_cpu and target_device.type != "cpu":
                 model.to(target_device)
         else:
-            device_ctx = target_device if not is_hsdp else contextlib.nullcontext()
+            hsdp_defer_to_cpu = is_hsdp and self.quant_config is None
+            device_ctx = contextlib.nullcontext() if hsdp_defer_to_cpu else target_device
             with device_ctx:
                 if load_format == "default":
                     model = initialize_model(self.od_config)
@@ -646,6 +653,22 @@ def _load_model_with_hsdp(
         model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
         self.load_weights(model)
 
+        # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights
+        # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel
+        # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params.
+        # Rewrite affected layers in-place to row-major contiguous storage and
+        # shift the .t() to GEMM-call time. Layers using other quant methods or
+        # already-contiguous weights are left untouched.
+        if self.quant_config is not None:
+            from vllm_omni.diffusion.quantization.hsdp_fp8 import (
+                prepare_fp8_layers_for_fsdp,
+            )
+
+            for _attr in ("transformer", "transformer_2"):
+                _trans = getattr(model, _attr, None)
+                if _trans is not None:
+                    prepare_fp8_layers_for_fsdp(_trans)
+
         # Collect all transformers to shard (some models have transformer_2 for MoE)
         transformers_to_shard = []
         transformer = getattr(model, "transformer", None)
diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py
index bc3eef49d68..e6802596762 100644
--- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py
+++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py
@@ -13,6 +13,7 @@
 import numpy as np
 import torch
 import torch.distributed as dist
+from cache_dit import ForwardPattern
 from torch import nn
 from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
 from transformers.models.qwen2.modeling_qwen2 import (
@@ -36,6 +37,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata as DiffusionAttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention
+from vllm_omni.diffusion.cache.cache_dit_backend import BagelCachedAdapter, CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
@@ -772,6 +774,13 @@ def forward(
 
 
 class Qwen2MoTModel(Qwen2PreTrainedModel):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "layers": ForwardPattern.Pattern_0,
+        },
+        cached_adapter_cls=BagelCachedAdapter,
+    )
+
     _layerwise_offload_blocks_attrs = ["layers"]
 
     @staticmethod
diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
index 59086684628..4543d692270 100644
--- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
@@ -18,6 +18,7 @@
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
@@ -30,6 +31,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
@@ -964,6 +966,16 @@ class Cosmos3VFMTransformer(nn.Module):
     ``Attention`` layer (with Ulysses all-to-all) or plain SDPA accordingly.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        # Cosmos3 GEN blocks return only hidden_states.  Per-layer UND K/V
+        # conditioning uses the transformer's cache-dit fallback path.
+        block_forward_patterns={
+            "gen_layers": ForwardPattern.Pattern_3,
+        },
+        has_separate_cfg=True,
+        check_forward_pattern=False,
+    )
+
     _repeated_blocks = ["Cosmos3GenDecoderLayer"]
 
     _layerwise_offload_blocks_attrs = ["gen_layers"]
diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py
index 42393542712..acf24dd2f64 100644
--- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py
+++ b/vllm_omni/diffusion/models/dreamid_omni/fusion.py
@@ -3,11 +3,13 @@
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -251,6 +253,13 @@ def forward(
 
 
 class FusionModel(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "fused_blocks": ForwardPattern.Pattern_0,
+        },
+        has_separate_cfg=True,
+    )
+
     _layerwise_offload_blocks_attrs = ["fused_blocks"]
 
     packed_modules_mapping = {
diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
index 06e541610f0..f058ba2bd3a 100644
--- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
+++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
@@ -21,11 +21,12 @@
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
-    get_tp_group,
+    tensor_model_parallel_all_reduce,
 )
 from vllm.model_executor.layers.conv import Conv3dLayer
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
+    QKVParallelLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.utils import set_weight_attrs
@@ -174,22 +175,82 @@ def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> Non
             raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.")
         param.data.copy_(shard)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        tp_size = get_tensor_model_parallel_world_size()
+    def _local_sum_sq(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
+        """Per-rank float32 activation, local sum-of-squares, and local width."""
         x_float = x.float()
         local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
-        local_count = x.shape[-1]
+        return x_float, local_sum_sq, x.shape[-1]
+
+    def _scale(
+        self,
+        x_float: torch.Tensor,
+        global_sum_sq: torch.Tensor,
+        global_count: int,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply the (already reduced) RMS and the per-rank weight shard."""
+        mean_sq = global_sum_sq / global_count
+        return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+        x_float, local_sum_sq, local_count = self._local_sum_sq(x)
 
         if tp_size > 1:
-            global_sum_sq = local_sum_sq.clone()
-            torch.distributed.all_reduce(global_sum_sq, group=get_tp_group().device_group)
+            # Use vLLM's collective (custom all-reduce / symmetric-mem fast path
+            # for small tensors) instead of raw torch.distributed.all_reduce, and
+            # take the return value so the custom-AR path (which may return a new
+            # buffer) is handled correctly. No .clone() needed.
+            global_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
             global_count = local_count * tp_size
         else:
             global_sum_sq = local_sum_sq
             global_count = local_count
 
-        mean_sq = global_sum_sq / global_count
-        return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
+        return self._scale(x_float, global_sum_sq, global_count, x)
+
+
+def fused_qk_rms_norm(
+    norm_q: nn.Module,
+    norm_k: nn.Module,
+    q: torch.Tensor,
+    k: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply q/k :class:`DistributedRMSNorm` with a SINGLE fused TP all-reduce.
+
+    WHY: self-attention norms q and k every step; on TP each norm issues its own
+    tiny all-reduce of the per-token sum-of-squares. For DreamZero's decode-like
+    per-step forward these latency-bound micro-collectives dominate, so we pack
+    both sum-of-squares into one tensor and reduce once (2 collectives → 1).
+
+    NUMERICALLY IDENTICAL to ``norm_q(q), norm_k(k)``: all-reduce is elementwise,
+    so packing along the last dim reduces each slice independently with the same
+    fp32 accumulation. Requires q and k to share the same shape (true for
+    self-attention — both come from the same hidden states). Falls back to
+    independent application when either norm is not a DistributedRMSNorm
+    (e.g. nn.Identity when qk_norm=False).
+    """
+    if not (isinstance(norm_q, DistributedRMSNorm) and isinstance(norm_k, DistributedRMSNorm)):
+        return norm_q(q), norm_k(k)
+
+    # The fused path reduces one packed sum-of-squares and reuses q's token width
+    # as the RMS count for BOTH q and k, so q and k must share the same shape.
+    # (Self-attention guarantees this: q and k are projected from the same x.)
+    assert q.shape == k.shape, "fused_qk_rms_norm requires q and k to have the same shape."
+
+    tp_size = get_tensor_model_parallel_world_size()
+    q_float, q_sum_sq, count = norm_q._local_sum_sq(q)
+    k_float, k_sum_sq, _ = norm_k._local_sum_sq(k)
+
+    if tp_size > 1:
+        packed = torch.cat([q_sum_sq, k_sum_sq], dim=-1)
+        packed = tensor_model_parallel_all_reduce(packed)
+        q_sum_sq, k_sum_sq = packed[..., 0:1], packed[..., 1:2]
+        count = count * tp_size
+
+    q_out = norm_q._scale(q_float, q_sum_sq, count, q)
+    k_out = norm_k._scale(k_float, k_sum_sq, count, k)
+    return q_out, k_out
 
 
 # ── Projections ─────────────────────────────────────────────────────
@@ -337,22 +398,35 @@ def forward(
         context = context[:, 257:]
         n, d = self.tp_num_heads, self.head_dim
         q = self.norm_q(self.q(x)).unflatten(2, (n, d))
+        # context (text) and context_img (clip features) are constant within a
+        # session, so k/v and k_img/v_img are cached on first call and reused.
+        # context_img == img_emb(clip feature) is session-invariant (state.clip_feas
+        # is set once at current_start_frame==0 and only cleared on reset, which also
+        # rebuilds crossattn_cache). Caching k_img/v_img removes a per-step
+        # norm_k_img all-reduce plus the k_img/v_img projection GEMMs in steady state.
+        # Only q depends on the per-step x and is always recomputed.
         if crossattn_cache is not None:
             if not crossattn_cache["is_init"]:
                 crossattn_cache["is_init"] = True
                 k = self.norm_k(self.k(context)).unflatten(2, (n, d))
                 v = self.v(context).unflatten(2, (n, d))
+                k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
+                v_img = self.v_img(context_img).unflatten(2, (n, d))
                 crossattn_cache["k"] = k
                 crossattn_cache["v"] = v
+                crossattn_cache["k_img"] = k_img
+                crossattn_cache["v_img"] = v_img
             else:
                 k = crossattn_cache["k"]
                 v = crossattn_cache["v"]
+                k_img = crossattn_cache["k_img"]
+                v_img = crossattn_cache["v_img"]
         else:
             k = self.norm_k(self.k(context)).unflatten(2, (n, d))
             v = self.v(context).unflatten(2, (n, d))
+            k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
+            v_img = self.v_img(context_img).unflatten(2, (n, d))
         x = self.attn(q, k, v)
-        k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
-        v_img = self.v_img(context_img).unflatten(2, (n, d))
         img_x = self.attn(q, k_img, v_img)
         x = x.flatten(2)
         img_x = img_x.flatten(2)
@@ -403,9 +477,21 @@ def __init__(
         self.num_action_per_block = num_action_per_block
         self.num_state_per_block = num_state_per_block
         self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen
-        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
-        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
-        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        # Fused QKV projection: q/k/v all come from x, so a single column-parallel
+        # GEMM replaces three (fewer launches / better GEMM util in the decode-like
+        # per-step forward). No GQA here: total_num_kv_heads defaults to num_heads.
+        self.qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            bias=True,
+        )
+        # The forward splits qkv into three equal q/k/v shards, which is only
+        # correct without GQA. total_num_kv_heads defaults to total_num_heads
+        # here; fail loud if a future checkpoint introduces GQA so the equal
+        # split does not silently misalign k/v.
+        if self.qkv.total_num_kv_heads != self.qkv.total_num_heads:
+            raise ValueError("Self-attn QKV fusion requires no GQA (total_num_kv_heads == total_num_heads).")
         self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
         self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
@@ -431,9 +517,15 @@ def forward(
         """Inference-only forward (KV cache path)."""
         n, d = self.tp_num_heads, self.head_dim
 
-        q = self.norm_q(self.q(x)).unflatten(2, (n, d))
-        k = self.norm_k(self.k(x)).unflatten(2, (n, d))
-        v = self.v(x).unflatten(2, (n, d))
+        # Single fused QKV GEMM, then split into the per-rank q/k/v shards.
+        qkv, _ = self.qkv(x)
+        qk_size = self.tp_num_heads * self.head_dim
+        q, k, v = qkv.split([qk_size, qk_size, qk_size], dim=-1)
+        # Fuse q/k qk-norm into a single TP all-reduce (both come from x, same shape).
+        q, k = fused_qk_rms_norm(self.norm_q, self.norm_k, q, k)
+        q = q.unflatten(2, (n, d))
+        k = k.unflatten(2, (n, d))
+        v = v.unflatten(2, (n, d))
 
         updated_kv_cache: torch.Tensor | None = None
 
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 5d19ed7ce44..c9372a8187d 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -23,6 +23,10 @@
 from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from vllm_omni.diffusion.cache.stepcache import (
+    get_stepcache_state,
+    is_stepcache_active,
+)
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
@@ -56,10 +60,6 @@
 MAX_DREAMZERO_SESSIONS = 64
 
 
-# ---------------------------------------------------------------------------
-# ---------------------------------------------------------------------------
-
-
 class VideoActionScheduler:
     """Wraps video + action schedulers into single .step() interface."""
 
@@ -134,57 +134,36 @@ def _kv_commit(self):
             s.commit_chunk()
 
     def _kv_get_cross(self, state, is_negative):
-        """Cross-attn cache: pool-backed when BDE active, model-local otherwise."""
-        s = self._bde_kv_state
-        if s is not None:
-            return s.get_cross_kv_caches(is_negative, lambda: state.get_crossattn_caches(is_negative))
+        """Cross-attn cache: always model-local (even under BDE).
+
+        Upstream #4154 extended the I2V cross-attn cache to also store image-token
+        ``k_img``/``v_img``; BDE's cross-attn pool only stores text ``k``/``v``, so a
+        pool-backed dict makes the I2V forward ``KeyError`` on ``'k_img'``. BDE
+        manages *self*-attn KV via the paged pool; cross-attn uses the model's own
+        cache (now image-token-caching), which BDE's reset path already clears.
+        """
         return state.get_crossattn_caches(is_negative)
 
     def _kv_populate_cross(self, context: torch.Tensor, is_negative: bool) -> None:
-        """Eagerly project cross-attn K/V for all layers into the BDE pool.
-
-        Called once per branch after text encoding (before the first denoising
-        step), and again after each window-boundary reset.  The attention
-        forward never writes to the cross-attn pool — it only reads.
+        """No-op: cross-attn is model-local (see :meth:`_kv_get_cross`).
 
-        The raw encoder output is projected through ``text_embedding`` first
-        so it matches the shape the cross-attn linear layers expect (the same
-        projection ``_forward_blocks`` applies internally).
+        Kept as a stable hook. The BDE cross-attn pool is unused after upstream
+        #4154 added ``k_img``/``v_img`` caching, so there is nothing to eagerly
+        populate; the model's own cross-attn cache handles (and now caches) the
+        image-token projections.
         """
-        s = self._bde_kv_state
-        if s is None:
-            return
-        # Only populate once per branch per session — reset clears this flag.
-        if s._cross_populated.get(is_negative, False):
-            return
-        projected = self.transformer.text_embedding(context)
-        for i, block in enumerate(self.transformer.blocks):
-            # The cross-attn forward caches text-only k/v — the image tokens
-            # (first 257, prepended inside _forward_blocks via img_emb) are
-            # always recomputed.  Here we feed raw projected text (no image
-            # prepend), so we compute k/v directly from the full projected text,
-            # mirroring the cross-attn forward's own k/v projection.
-            ca = block.cross_attn
-            n, d = ca.tp_num_heads, ca.head_dim
-            k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
-            v = ca.v(projected).unflatten(2, (n, d))
-            s.kv_cache.write_cross_kv(i, is_negative, k, v)
-        s._cross_populated[is_negative] = True
-        _log.info(
-            "BDE CROSS POPULATE [%s]: %d layers, context=%s",
-            "neg" if is_negative else "pos",
-            len(self.transformer.blocks),
-            tuple(context.shape),
-        )
+        return
 
-    def _kv_reset(self, state):
+    def _kv_reset(self, state, *, clear_video_latents: bool = True):
         """Reset model-local KV and, under BDE, the pooled session window too.
 
         DreamZero resets its sliding window at the attention-window boundary; the
         BDE pool must drop the same window so the next forward starts fresh and
-        stays parity-exact with the model-local path.
+        stays parity-exact with the model-local path. ``clear_video_latents=False``
+        mirrors ``state.reset_inference_state()`` (window-boundary "inference" reset),
+        which keeps the accumulated video latents for export.
         """
-        state.reset()
+        state.reset(clear_video_latents=clear_video_latents)
         s = self._bde_kv_state
         if s is not None:
             s.reset()
@@ -297,6 +276,9 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         self._max_session_states = MAX_DREAMZERO_SESSIONS
         self.state = self._get_or_create_state("default")
 
+        # DiT step cache is configured by StepCacheBackend
+        # (cache_backend="step_cache") via pipeline._stepcache_config.
+
         # Keep runtime inference settings separate from the training-time config.
         self.num_inference_steps: int = model_config.get(
             "num_inference_steps",
@@ -398,6 +380,27 @@ def _load_repo_json(model_path: str, relative_path: str, local_files_only: bool)
 
     def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
         """Call CausalWanModel, return (video_pred, action_pred)."""
+        video_pred, action_pred = self._predict_noise_eager(kwargs)
+
+        if is_stepcache_active(self):
+            video_pred = video_pred.clone()
+            if action_pred is not None:
+                action_pred = action_pred.clone()
+
+        if action_pred is None:
+            batch_size = kwargs["hidden_states"].shape[0]
+            action_pred = torch.empty(
+                batch_size,
+                0,
+                self.transformer.action_dim,
+                device=video_pred.device,
+                dtype=video_pred.dtype,
+            )
+        return (video_pred, action_pred)
+
+    def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Eager DiT forward; also handles KV-cache write-back on prefill."""
+        self._cudagraph_mark_step_begin()
         video_pred, action_pred, updated_kv_caches = self.transformer(
             x=kwargs["hidden_states"],
             timestep=kwargs["timestep_video"],
@@ -428,19 +431,118 @@ def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
                 self._kv_update(state, i, kv, is_neg, kwargs.get("seq_len"))
             self._kv_commit()
 
-        video_pred = video_pred.clone()
-        if action_pred is not None:
-            action_pred = action_pred.clone()
-        else:
-            batch_size = kwargs["hidden_states"].shape[0]
-            action_pred = torch.empty(
-                batch_size,
-                0,
-                self.transformer.action_dim,
-                device=video_pred.device,
-                dtype=video_pred.dtype,
-            )  # CFG-parallel-safe dummy action pred
-        return (video_pred, action_pred)
+        return video_pred, action_pred
+
+    # -----------------------------------------------------------------------
+    # torch.compile setup (paper D.2 extended: encoders + VAE + DiT blocks)
+    # -----------------------------------------------------------------------
+
+    def setup_compile(self) -> None:
+        """Compile DreamZero encoders, VAE, and per-block DiT for inference.
+
+        Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False``
+        on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and
+        compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop).
+        DiT blocks use per-block ``fullgraph=True``.
+        """
+        if not torch.cuda.is_available():
+            logger.info("DreamZero setup_compile skipped: CUDA not available.")
+            return
+
+        from vllm_omni.diffusion.models.dreamzero.wan_vae_feat_cache_patch import (
+            apply_wan_vae_feat_cache_tensor_patch,
+        )
+
+        apply_wan_vae_feat_cache_tensor_patch()
+
+        compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
+        # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
+        dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
+
+        logger.info(
+            "DreamZero: torch.compile text/image/VAE encode + per-block DiT (encoders reduce-overhead, DiT default)."
+        )
+
+        try:
+            self.text_encoder.forward = torch.compile(self.text_encoder.forward, **compile_ro)
+        except Exception as exc:
+            logger.warning("DreamZero: text_encoder compile failed (%s); skipping.", exc)
+
+        try:
+            self.image_encoder.model.visual.forward = torch.compile(
+                self.image_encoder.model.visual.forward,
+                **compile_ro,
+            )
+        except Exception as exc:
+            logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
+
+        try:
+            self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
+        except Exception as exc:
+            logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
+
+        compiled_blocks = 0
+        for block in self.transformer.blocks:
+            try:
+                block.forward = torch.compile(block.forward, **dit_compile)
+                compiled_blocks += 1
+            except Exception as exc:
+                logger.warning(
+                    "DreamZero: transformer block %d compile failed (%s); leaving remaining eager.",
+                    compiled_blocks,
+                    exc,
+                )
+                break
+        if compiled_blocks:
+            logger.info(
+                "DreamZero: compiled %d/%d transformer blocks.",
+                compiled_blocks,
+                len(self.transformer.blocks),
+            )
+
+        self.warmup_compile()
+
+    def warmup_compile(self) -> None:
+        """Warm up compiled text/image/VAE paths before timed inference."""
+        if not torch.cuda.is_available():
+            return
+
+        device = next(self.text_encoder.parameters()).device
+        with torch.inference_mode():
+            try:
+                text_tokens = torch.zeros(1, 16, dtype=torch.long, device=device)
+                attention_mask = torch.ones_like(text_tokens)
+                self._encode_text(text_tokens, attention_mask)
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (text_encoder) skipped: %s", exc)
+
+            try:
+                image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
+                self._encode_image(image, self.num_frames, 180, 320)
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
+
+            try:
+                latent_h, latent_w = 180 // 8, 320 // 8
+                dummy_latent = torch.zeros(
+                    1,
+                    16,
+                    self.num_frame_per_block,
+                    latent_h * 2,
+                    latent_w * 2,
+                    dtype=self.vae.dtype,
+                    device=device,
+                )
+                mean = self.vae_latents_mean.to(device=device, dtype=self.vae.dtype)
+                inv_std = self.vae_latents_inv_std.to(device=device, dtype=self.vae.dtype)
+                dummy_latent_denorm = dummy_latent / inv_std + mean
+                self._cudagraph_mark_step_begin()
+                self.vae.decode(dummy_latent_denorm, return_dict=False)
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (vae decode) skipped: %s", exc)
+
+        torch.accelerator.synchronize(device)
+        logger.info("DreamZero compile warmup finished (text / image / vae decode).")
 
     def combine_cfg_noise(
         self,
@@ -489,12 +591,20 @@ def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor:
             videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4)
         return videos.to(dtype=torch.bfloat16)
 
+    @staticmethod
+    def _cudagraph_mark_step_begin() -> None:
+        try:
+            torch.compiler.cudagraph_mark_step_begin()
+        except Exception:
+            pass
+
     # -----------------------------------------------------------------------
     # Text encoding
     # -----------------------------------------------------------------------
 
     def _encode_text(self, text_tokens: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
         """Encode text prompt via UMT5."""
+        self._cudagraph_mark_step_begin()
         seq_lens = attention_mask.gt(0).sum(dim=1).long()
         prompt_emb = self.text_encoder(
             text_tokens,
@@ -523,6 +633,7 @@ def _encode_image(
         batch_size = image.shape[0]
 
         with torch.amp.autocast(dtype=torch.bfloat16, device_type=device.type):
+            self._cudagraph_mark_step_begin()
             clip_context = self.image_encoder.encode_image(image)
 
             msk = torch.ones(batch_size, num_frames, height // 8, width // 8, device=device)
@@ -570,8 +681,23 @@ def decode_video_latents(self, video_latents: torch.Tensor) -> torch.Tensor:
         inv_std = self.vae_latents_inv_std.to(device=vae_device, dtype=vae_dtype)
         latents = latents / inv_std + mean
         with torch.no_grad():
+            self._cudagraph_mark_step_begin()
             return self.vae.decode(latents, return_dict=False)[0]
 
+    def decode_accumulated_video_latents(self, session_id: str | None = None) -> torch.Tensor:
+        """Decode all AR-chunk latents accumulated for ``session_id``."""
+        state = self._get_or_create_state(session_id)
+        latents = state.get_concatenated_video_latents()
+        if latents is None:
+            session_key = str(session_id or "default")
+            raise RuntimeError(f"No accumulated video latents for session {session_key!r}.")
+        return self.decode_video_latents(latents)
+
+    def clear_accumulated_video_latents(self, session_id: str | None = None) -> None:
+        """Clear accumulated video latents for ``session_id`` without resetting KV state."""
+        state = self._get_or_create_state(session_id)
+        state.clear_video_latents()
+
     # -----------------------------------------------------------------------
     # KV cache prefill
     # -----------------------------------------------------------------------
@@ -737,6 +863,12 @@ def diffuse(
 
         noisy_input = video_latents
         noisy_input_action = action_latents
+        # ---- step_cache (StepCacheBackend / dreamzero.git) ----
+        _cached_flow_pred: torch.Tensor | None = None
+        _cached_flow_pred_action: torch.Tensor | None = None
+        _prev_predictions: list[tuple[torch.Tensor]] = []
+        _step_cache = get_stepcache_state(self) if is_stepcache_active(self) else None
+
         for index in range(len(timesteps_video)):
             video_timestep = timesteps_video[index]
             action_timestep = timesteps_action[index]
@@ -765,46 +897,58 @@ def diffuse(
             else:
                 y = state.ys[:, :, -self.num_frame_per_block :]
 
-            positive_kwargs = dict(
-                hidden_states=noisy_input.transpose(1, 2),
-                timestep_video=timestep,
-                encoder_hidden_states=prompt_embeds,
-                kv_cache=self._kv_get(state, False),
-                crossattn_cache=self._kv_get_cross(state, False),
-                y=y,
-                clip_feature=state.clip_feas,
-                action=noisy_input_action,
-                timestep_action=timestep_action,
-                is_negative=False,
-                **common_kwargs,
-            )
-
-            # Negative (uncond) kwargs
-            if do_true_cfg and negative_prompt_embeds is not None:
-                negative_kwargs = dict(
+            run_dit = _step_cache is None or _step_cache.should_run_step(_prev_predictions)
+            if run_dit:
+                positive_kwargs = dict(
                     hidden_states=noisy_input.transpose(1, 2),
                     timestep_video=timestep,
-                    encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=self._kv_get(state, True),
-                    crossattn_cache=self._kv_get_cross(state, True),
+                    encoder_hidden_states=prompt_embeds,
+                    kv_cache=self._kv_get(state, False),
+                    crossattn_cache=self._kv_get_cross(state, False),
                     y=y,
                     clip_feature=state.clip_feas,
                     action=noisy_input_action,
                     timestep_action=timestep_action,
-                    is_negative=True,
+                    is_negative=False,
                     **common_kwargs,
                 )
-            else:
-                negative_kwargs = None
 
-            noise_pred = self.predict_noise_maybe_with_cfg(
-                positive_kwargs=positive_kwargs,
-                negative_kwargs=negative_kwargs,
-                do_true_cfg=do_true_cfg,
-                true_cfg_scale=self.cfg_scale,
-                cfg_normalize=False,
-            )
-            flow_pred, flow_pred_action = noise_pred
+                if do_true_cfg and negative_prompt_embeds is not None:
+                    negative_kwargs = dict(
+                        hidden_states=noisy_input.transpose(1, 2),
+                        timestep_video=timestep,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        kv_cache=self._kv_get(state, True),
+                        crossattn_cache=self._kv_get_cross(state, True),
+                        y=y,
+                        clip_feature=state.clip_feas,
+                        action=noisy_input_action,
+                        timestep_action=timestep_action,
+                        is_negative=True,
+                        **common_kwargs,
+                    )
+                else:
+                    negative_kwargs = None
+
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                    do_true_cfg=do_true_cfg,
+                    true_cfg_scale=self.cfg_scale,
+                    cfg_normalize=False,
+                )
+                flow_pred, flow_pred_action = noise_pred
+                _cached_flow_pred = flow_pred
+                _cached_flow_pred_action = flow_pred_action
+
+                _prev_predictions.append((flow_pred,))
+                if _step_cache is not None:
+                    _step_cache.trim_history(_prev_predictions)
+            else:
+                # Reuse previous prediction (DiT step-skipping cache).
+                assert _cached_flow_pred is not None, "DiT cache: no cached prediction available for step skip."
+                flow_pred = _cached_flow_pred
+                flow_pred_action = _cached_flow_pred_action
 
             latents = (noisy_input, noisy_input_action)
             t = (video_timestep, action_timestep)
@@ -913,9 +1057,17 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
         # on the next inference request after websocket reset/session switch.
         if extra_args.get("reset", False):
             self._kv_reset(state)
-        # Auto-reset based on model state (before accumulation)
-        if state.should_reset(text_tokens, 0, self.transformer.local_attn_size):
-            self._kv_reset(state)
+        else:
+            # Auto-reset based on model state (before accumulation). Both a "session"
+            # reset and an "inference" (window-boundary) reset clear the model-local
+            # KV window, so route both through _kv_reset to drop the BDE pool window
+            # too and stay parity-exact; "inference" keeps the accumulated video
+            # latents (reset_inference_state == reset(clear_video_latents=False)).
+            reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
+            if reset_reason == "session":
+                self._kv_reset(state)
+            elif reset_reason == "inference":
+                self._kv_reset(state, clear_video_latents=False)
         state.language = text_tokens
 
         # Frame accumulation: stitched single frame -> multi-frame video
@@ -1069,7 +1221,9 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             video_out = torch.cat([image, video_out], dim=1)
         state.current_start_frame += self.num_frame_per_block
 
-        # q99 denorm: [-1,1] -> real values
+        state.append_video_latents(video_out)
+
+        # q99 denorm: [-1,1] → real values
         action_out = self._denormalize_action(action_out.float(), embodiment_name)
 
         # Relative -> absolute: only for relative_action_keys (joint_position only)
@@ -1214,10 +1368,27 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     .replace("img_emb.proj.3.", "img_emb.fc2.")
                     .replace("img_emb.proj.4.", "img_emb.norm2.")
                 )
+
+                # Self-attn q/k/v are fused into a single QKVParallelLinear; route
+                # each separate checkpoint weight/bias to the packed `qkv` param
+                # with its shard id. cross_attn keeps separate q/k/v (q from x,
+                # k/v from context — not fusible), so it is left untouched here.
+                qkv_shard_id: str | None = None
+                for shard_id in ("q", "k", "v"):
+                    needle = f".self_attn.{shard_id}."
+                    if needle in new_name:
+                        new_name = new_name.replace(needle, ".self_attn.qkv.")
+                        qkv_shard_id = shard_id
+                        break
+
                 if new_name in params:
                     param = params[new_name]
-                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                    weight_loader(param, tensor)
+                    if qkv_shard_id is not None:
+                        # QKVParallelLinear.weight_loader needs the shard id.
+                        param.weight_loader(param, tensor, qkv_shard_id)
+                    else:
+                        weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                        weight_loader(param, tensor)
                     loaded.add(new_name)
                 elif new_name in buffers:
                     buffers[new_name].data.copy_(tensor)
diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
index 0972d1b5a2c..5815d93d4b1 100644
--- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
@@ -23,7 +23,8 @@ class DreamZeroState:
     Lifecycle:
         - Created once in DreamZeroPipeline.__init__()
         - Mutated every forward() call (frame append, KV cache grow)
-        - reset() on new session / language change / local_attn_size exceeded
+        - reset() on new session / language change
+        - reset_inference_state() on local_attn_size exceeded (KV only)
     """
 
     def __init__(self) -> None:
@@ -65,50 +66,106 @@ def accumulate_frames(self, stitched: np.ndarray) -> np.ndarray:
         self.call_count += 1
         return np.stack(frames, axis=0)  # (T, H, W, C)
 
+    def append_video_latents(self, video_out: torch.Tensor) -> None:
+        """Append one AR chunk of normalized video latents for later decode.
+
+        Args:
+            video_out: ``(B, T, C, H, W)`` tensor returned by the denoise loop
+                (before the ``transpose(1, 2)`` stored in ``DiffusionOutput``).
+        """
+        if video_out.dim() != 5:
+            raise ValueError(f"Expected 5D video_out, got shape {tuple(video_out.shape)}")
+        # Upstream ``torch.cat(self.video_across_time, dim=2)`` uses (B, C, T, H, W).
+        chunk = video_out.transpose(1, 2).detach().cpu()
+        self.video_latents_across_time.append(chunk)
+        logger.info(
+            "append_video_latents: chunk_shape=%s total_chunks=%d total_latent_t=%d",
+            tuple(chunk.shape),
+            len(self.video_latents_across_time),
+            int(sum(c.shape[2] for c in self.video_latents_across_time)),
+        )
+
+    def get_concatenated_video_latents(self) -> torch.Tensor | None:
+        """Return all accumulated chunks concatenated along the time dimension."""
+        if not self.video_latents_across_time:
+            return None
+        if len(self.video_latents_across_time) == 1:
+            return self.video_latents_across_time[0]
+        return torch.cat(self.video_latents_across_time, dim=2)
+
+    def clear_video_latents(self) -> None:
+        """Drop accumulated video latents without resetting KV/frame state."""
+        self.video_latents_across_time = []
+
     # ------------------------------------------------------------------
     # Reset / should_reset
     # ------------------------------------------------------------------
 
-    def reset(self) -> None:
-        """Clear all state."""
-        self.stitched_buffer: deque[np.ndarray] = deque(maxlen=FRAMES_PER_CHUNK)
-        self.call_count: int = 0
+    def reset(self, *, clear_video_latents: bool = True) -> None:
+        """Clear session state.
 
-        # KV cache once robot-policy diffusion supports that integration.
-        self.kv_cache: list[torch.Tensor] | None = None
-        self.kv_cache_neg: list[torch.Tensor] | None = None
-        self.crossattn_cache: list[dict[str, bool | torch.Tensor | None]] | None = None
-        self.crossattn_cache_neg: list[dict[str, bool | torch.Tensor | None]] | None = None
-        self.current_start_frame: int = 0
-
-        self.clip_feas: torch.Tensor | None = None
-        self.ys: torch.Tensor | None = None
-        self.language: torch.Tensor | None = None
+        Args:
+            clear_video_latents: When ``False``, keep ``video_latents_across_time``
+                so offline/online export can decode the full AR rollout.
+        """
+        saved_video_latents = [] if clear_video_latents else list(self.video_latents_across_time)
+        self.stitched_buffer = deque(maxlen=FRAMES_PER_CHUNK)
+        self.call_count = 0
+        # Normalized VAE latents per AR chunk, each (B, C, T, H, W) on CPU.
+        # Concatenated along T before decode (matches upstream video_across_time).
+        self.video_latents_across_time = saved_video_latents
 
-    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
-        """Determine if state should be reset before this forward()."""
+        # KV cache once robot-policy diffusion supports that integration.
+        self.kv_cache = None
+        self.kv_cache_neg = None
+        self.crossattn_cache = None
+        self.crossattn_cache_neg = None
+        self.current_start_frame = 0
+
+        self.clip_feas = None
+        self.ys = None
+        if clear_video_latents:
+            self.language = None
+
+    def reset_inference_state(self) -> None:
+        """Reset KV/frame state after local attention rolls without dropping video latents."""
+        self.reset(clear_video_latents=False)
+
+    def reset_reason(
+        self,
+        text_tokens: torch.Tensor | None,
+        num_video_frames: int,
+        local_attn_size: int,
+    ) -> str | None:
+        """Return why state should reset before the next forward(), if any."""
         if self.language is None:
             logger.info("language is None, resetting")
-            return True
+            return "session"
 
         if text_tokens is not None and not torch.equal(self.language, text_tokens):
             logger.info("language changed, resetting")
-            return True
+            return "session"
 
         # NOTE: after accumulate_frames, num_video_frames is the accumulated T
         # (1 for first call, 4 for subsequent). Only reset on true single-frame
         # which happens when the stitched_buffer was cleared externally.
         if num_video_frames == 1 and self.call_count > 1:
             logger.info("single frame input after first call, resetting")
-            return True
+            return "session"
 
         if local_attn_size != -1 and self.current_start_frame >= local_attn_size:
             logger.info(
-                "current_start_frame %d >= local_attn_size %d, resetting", self.current_start_frame, local_attn_size
+                "current_start_frame %d >= local_attn_size %d, resetting inference state",
+                self.current_start_frame,
+                local_attn_size,
             )
-            return True
+            return "inference"
+
+        return None
 
-        return False
+    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
+        """Determine if state should be reset before this forward()."""
+        return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None
 
     # ------------------------------------------------------------------
     # KV cache management
@@ -131,8 +188,12 @@ def create_kv_caches(
             torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
         ]
 
-        self.crossattn_cache = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
-        self.crossattn_cache_neg = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
+        self.crossattn_cache = [
+            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
+        ]
+        self.crossattn_cache_neg = [
+            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
+        ]
 
     def update_kv_cache(
         self,
diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
index dfab272d65e..530a3fbec89 100644
--- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
+++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
@@ -17,20 +17,37 @@ def gpu_mem_stats(self) -> dict:
         tensor high-water mark. Both are monotonic from process start, so the
         end-of-run values capture the whole rollout's peak.
         """
-        dev = torch.cuda.current_device()
+        dev = torch.accelerator.current_device_index()
         return {
             "device": int(dev),
-            "peak_reserved_gib": torch.cuda.max_memory_reserved(dev) / (1024**3),
-            "peak_allocated_gib": torch.cuda.max_memory_allocated(dev) / (1024**3),
+            "peak_reserved_gib": torch.accelerator.max_memory_reserved(dev) / (1024**3),
+            "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3),
         }
 
+    @staticmethod
+    def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
+        decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
+        decoded = decoded.clamp(-1, 1) * 0.5 + 0.5
+        return (decoded * 255.0).round().to(torch.uint8).cpu()
+
     def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor:
         if self.model_runner is None or self.model_runner.pipeline is None:
             raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
 
         with torch.inference_mode():
             decoded = self.model_runner.pipeline.decode_video_latents(video_latents)
-            decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
-            decoded = decoded.clamp(-1, 1) * 0.5 + 0.5
-            decoded = (decoded * 255.0).round().to(torch.uint8).cpu()
-        return decoded
+            return self._latents_to_uint8_frames(decoded)
+
+    def decode_accumulated_video_latents_to_uint8(self, session_id: str = "default") -> torch.Tensor:
+        """Decode all AR-chunk latents accumulated on the server for one session."""
+        if self.model_runner is None or self.model_runner.pipeline is None:
+            raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
+
+        with torch.inference_mode():
+            decoded = self.model_runner.pipeline.decode_accumulated_video_latents(session_id)
+            return self._latents_to_uint8_frames(decoded)
+
+    def clear_accumulated_video_latents(self, session_id: str = "default") -> None:
+        if self.model_runner is None or self.model_runner.pipeline is None:
+            raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
+        self.model_runner.pipeline.clear_accumulated_video_latents(session_id)
diff --git a/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py
new file mode 100644
index 00000000000..57d9b866c56
--- /dev/null
+++ b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py
@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Make diffusers Wan VAE feat_cache compile-friendly.
+
+diffusers ``WanResample`` stores the string sentinel ``"Rep"`` in feat_cache slots.
+``tensor == "Rep"`` returns a non-Tensor and breaks ``torch.compile(fullgraph=True)``.
+We replace it with a scalar int8 tensor marker (value 0).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_PATCHED = False
+
+
+def _is_feat_rep(entry: Any) -> bool:
+    if not isinstance(entry, torch.Tensor):
+        return False
+    return entry.numel() == 1 and entry.dtype == torch.int8
+
+
+def _make_feat_rep(x: torch.Tensor) -> torch.Tensor:
+    return torch.zeros(1, device=x.device, dtype=torch.int8)
+
+
+def _wan_resample_forward_tensorized(self, x, feat_cache=None, feat_idx=[0]):
+    """``WanResample.forward`` with tensor feat_cache sentinel (no ``"Rep"`` string)."""
+    from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T
+
+    b, c, t, h, w = x.size()
+    if self.mode == "upsample3d":
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            if feat_cache[idx] is None:
+                feat_cache[idx] = _make_feat_rep(x)
+                feat_idx[0] += 1
+            else:
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and not _is_feat_rep(feat_cache[idx]):
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+                        dim=2,
+                    )
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and _is_feat_rep(feat_cache[idx]):
+                    cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+                if _is_feat_rep(feat_cache[idx]):
+                    x = self.time_conv(x)
+                else:
+                    x = self.time_conv(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+
+                x = x.reshape(b, 2, c, t, h, w)
+                x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                x = x.reshape(b, c, t * 2, h, w)
+    t = x.shape[2]
+    x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+    x = self.resample(x)
+    x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
+
+    if self.mode == "downsample3d":
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            if feat_cache[idx] is None:
+                feat_cache[idx] = x.clone()
+                feat_idx[0] += 1
+            else:
+                cache_x = x[:, :, -1:, :, :].clone()
+                x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+    return x
+
+
+def apply_wan_vae_feat_cache_tensor_patch() -> None:
+    """Monkeypatch diffusers WanResample once per process."""
+    global _PATCHED
+    if _PATCHED:
+        return
+
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample
+
+    WanResample.forward = _wan_resample_forward_tensorized
+    _PATCHED = True
+    logger.info("DreamZero: applied Wan VAE feat_cache tensor sentinel patch.")
diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py
index aa9b0d8b79a..8398cb7aba2 100644
--- a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py
+++ b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py
@@ -8,6 +8,7 @@
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -21,6 +22,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -459,6 +461,12 @@ def forward(self, x: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
 
 
 class ErnieImageTransformer2DModel(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "layers": ForwardPattern.Pattern_3,
+        },
+        check_forward_pattern=False,
+    )
     _repeated_blocks = ["ErnieImageSharedAdaLNBlock"]
     _layerwise_offload_blocks_attrs = ["layers"]
 
diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py
index 88d8126c8ad..9d89368bec0 100644
--- a/vllm_omni/diffusion/models/flux/flux_transformer.py
+++ b/vllm_omni/diffusion/models/flux/flux_transformer.py
@@ -6,6 +6,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     CombinedTimestepGuidanceTextProjEmbeddings,
     CombinedTimestepTextProjEmbeddings,
@@ -25,6 +26,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -524,6 +526,13 @@ class FluxTransformer2DModel(nn.Module):
             The dimensions to use for the rotary positional embeddings.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+            "single_transformer_blocks": ForwardPattern.Pattern_1,
+        }
+    )
+
     # the small and frequently-repeated block(s) of a model
     # -- typically a transformer layer
     # used for torch compile optimizations
diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py
index 0c6896c131d..5a65db257da 100644
--- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py
+++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py
@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Any
 
 import torch
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     TimestepEmbedding,
     Timesteps,
@@ -25,6 +26,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -755,6 +757,13 @@ class Flux2Transformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+            "single_transformer_blocks": ForwardPattern.Pattern_2,
+        }
+    )
+
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
     _sp_plan = {
         "": {
diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
index a708a8ece6a..cf2c9f78a19 100644
--- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
+++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
@@ -21,6 +21,7 @@
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     TimestepEmbedding,
     Timesteps,
@@ -40,6 +41,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -768,6 +770,13 @@ class Flux2Transformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+            "single_transformer_blocks": ForwardPattern.Pattern_2,
+        }
+    )
+
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
index 56a7066e390..c933d4e7f31 100644
--- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
+++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
@@ -9,6 +9,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.transformers.transformer_glm_image import GlmImageCombinedTimestepSizeEmbeddings
 from vllm.logger import init_logger
@@ -20,6 +21,8 @@
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
@@ -935,6 +938,13 @@ class GlmImageTransformer2DModel(CachedTransformer):
             `od_config.tf_model_config`.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_0,
+        },
+        has_separate_cfg=True,
+    )
+
     _repeated_blocks = ["GlmImageTransformerBlock"]
     # SP plan using GlmImagePrepare module for sharding hidden_states and RoPE together.
     # Similar to Qwen-Image's ImageRopePrepare, this creates a module boundary where
diff --git a/vllm_omni/diffusion/models/gr00t/__init__.py b/vllm_omni/diffusion/models/gr00t/__init__.py
new file mode 100644
index 00000000000..a10202a40bc
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
+
+__all__ = ["Gr00tN1d7Pipeline"]
diff --git a/vllm_omni/diffusion/models/gr00t/configs/__init__.py b/vllm_omni/diffusion/models/gr00t/configs/__init__.py
new file mode 100644
index 00000000000..9648b859751
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/configs/__init__.py
@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+MODEL_CONFIG_TYPES: dict[str, type] = {}
+
+
+def register_model_config(shortname: str, configtype: type) -> None:
+    MODEL_CONFIG_TYPES[shortname] = configtype
diff --git a/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py b/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py
new file mode 100644
index 00000000000..7d11a651cf6
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py
@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.models.gr00t.dataio.types import (
+    ActionConfig,
+    ActionFormat,
+    ActionRepresentation,
+    ActionType,
+    ModalityConfig,
+)
+
+__all__ = ["ActionConfig", "ActionFormat", "ActionRepresentation", "ActionType", "ModalityConfig"]
diff --git a/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py
new file mode 100644
index 00000000000..d7d4527fb93
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from transformers import PretrainedConfig
+
+from . import register_model_config
+
+
+def _default_diffusion_model_cfg() -> dict:
+    return {
+        "positional_embeddings": None,
+        "num_layers": 16,
+        "num_attention_heads": 32,
+        "attention_head_dim": 48,
+        "norm_type": "ada_norm",
+        "output_dim": 1024,
+        "interleave_self_attention": True,
+    }
+
+
+class Gr00tN1d7Config(PretrainedConfig):
+    model_type = "Gr00tN1d7"
+
+    def __init__(
+        self,
+        model_dtype: str = "bfloat16",
+        model_name: str = "nvidia/Cosmos-Reason2-2B",
+        backbone_model_type: str = "qwen",
+        model_revision: str | None = None,
+        backbone_embedding_dim: int = 2048,
+        select_layer: int = 12,
+        reproject_vision: bool = False,
+        use_flash_attention: bool = True,
+        load_bf16: bool = False,
+        image_crop_size: tuple | None = (230, 230),
+        image_target_size: tuple | None = (256, 256),
+        shortest_image_edge: int | None = None,
+        crop_fraction: float | None = None,
+        random_rotation_angle: int | None = None,
+        color_jitter_params: dict | None = None,
+        formalize_language: bool = True,
+        apply_sincos_state_encoding: bool = False,
+        use_percentiles: bool = True,
+        use_relative_action: bool = False,
+        max_state_dim: int = 132,
+        max_action_dim: int = 132,
+        action_horizon: int = 40,
+        hidden_size: int = 1024,
+        input_embedding_dim: int = 1536,
+        state_history_length: int = 1,
+        add_pos_embed: bool = True,
+        attn_dropout: float = 0.2,
+        use_vlln: bool = True,
+        max_seq_len: int = 1024,
+        use_alternate_vl_dit: bool = True,
+        attend_text_every_n_blocks: int = 2,
+        diffusion_model_cfg: dict | None = None,
+        vl_self_attention_cfg: dict | None = None,
+        num_inference_timesteps: int = 4,
+        num_timestep_buckets: int = 1000,
+        exclude_state: bool = False,
+        use_mean_std: bool = False,
+        max_num_embodiments: int = 32,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_dtype = model_dtype
+        self.model_name = model_name
+        self.backbone_model_type = backbone_model_type
+        self.model_revision = model_revision
+        self.backbone_embedding_dim = backbone_embedding_dim
+        self.select_layer = select_layer
+        self.reproject_vision = reproject_vision
+        self.use_flash_attention = use_flash_attention
+        self.load_bf16 = load_bf16
+        self.image_crop_size = image_crop_size
+        self.image_target_size = image_target_size
+        self.shortest_image_edge = shortest_image_edge
+        self.crop_fraction = crop_fraction
+        self.random_rotation_angle = random_rotation_angle
+        self.color_jitter_params = color_jitter_params
+        self.formalize_language = formalize_language
+        self.apply_sincos_state_encoding = apply_sincos_state_encoding
+        self.use_percentiles = use_percentiles
+        self.use_relative_action = use_relative_action
+        self.max_state_dim = max_state_dim
+        self.max_action_dim = max_action_dim
+        self.action_horizon = action_horizon
+        self.hidden_size = hidden_size
+        self.input_embedding_dim = input_embedding_dim
+        self.state_history_length = state_history_length
+        self.add_pos_embed = add_pos_embed
+        self.attn_dropout = attn_dropout
+        self.use_vlln = use_vlln
+        self.max_seq_len = max_seq_len
+        self.use_alternate_vl_dit = use_alternate_vl_dit
+        self.attend_text_every_n_blocks = attend_text_every_n_blocks
+        self.diffusion_model_cfg = (
+            diffusion_model_cfg if diffusion_model_cfg is not None else _default_diffusion_model_cfg()
+        )
+        self.vl_self_attention_cfg = vl_self_attention_cfg
+        self.num_inference_timesteps = num_inference_timesteps
+        self.num_timestep_buckets = num_timestep_buckets
+        self.exclude_state = exclude_state
+        self.use_mean_std = use_mean_std
+        self.max_num_embodiments = max_num_embodiments
+
+
+register_model_config("Gr00tN1d7", Gr00tN1d7Config)
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/__init__.py b/vllm_omni/diffusion/models/gr00t/dataio/__init__.py
new file mode 100644
index 00000000000..208f01a7cb5
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/dataio/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py b/vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py
new file mode 100755
index 00000000000..3ed0c9b0360
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py
@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from enum import Enum
+
+"""
+Embodiment tags are used to identify the robot embodiment in the data.
+
+Naming convention:
+<dataset>_<robot_name>
+
+If using multiple datasets, e.g. sim GR1 and real GR1, we can drop the dataset name and use only the robot name.
+"""
+
+
+class EmbodimentTag(Enum):
+    """Embodiment tags supported by the GR00T N1.7 checkpoint.
+
+    Pretrain tags (baked into the base model nvidia/GR00T-N1.7-3B, inference-ready):
+    - OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT -> "oxe_droid_relative_eef_relative_joint"
+    - XDOF                                  -> "xdof_relative_eef_relative_joint"
+    - XDOF_SUBTASK                          -> "xdof_relative_eef_relative_joint_subtask"
+    - REAL_G1                               -> "real_g1_relative_eef_relative_joints"
+    - REAL_R1_PRO_SHARPA                    -> "real_r1_pro_sharpa_relative_eef"
+    - REAL_R1_PRO_SHARPA_HUMAN              -> "real_r1_pro_sharpa_relative_eef_human"
+    - REAL_R1_PRO_SHARPA_MAXINSIGHTS        -> "real_r1_pro_sharpa_relative_eef_maxinsights"
+    - REAL_R1_PRO_SHARPA_MECKA              -> "real_r1_pro_sharpa_relative_eef_mecka"
+
+    Pre-registered posttrain tags (require finetuned checkpoint):
+    - UNITREE_G1           -> "unitree_g1_full_body_with_waist_height_nav_cmd"
+    - UNITREE_G1_SONIC     -> "unitree_g1_sonic"
+    - SIMPLER_ENV_GOOGLE   -> "simpler_env_google"
+    - SIMPLER_ENV_WIDOWX   -> "simpler_env_widowx"
+    - LIBERO_PANDA         -> "libero_sim"
+
+    Finetuning tag (for custom robots):
+    - NEW_EMBODIMENT       -> "new_embodiment"
+
+    Use ``EmbodimentTag.resolve(s)`` to look up a tag by name or value,
+    case-insensitively.
+    """
+
+    OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT = "oxe_droid_relative_eef_relative_joint"
+    """
+    The Open-X-Embodiment DROID robot with relative EEF and relative joint position actions.
+    """
+
+    XDOF = "xdof_relative_eef_relative_joint"
+    """
+    The generic X-DOF robot with relative EEF and relative joint position actions.
+    """
+
+    XDOF_SUBTASK = "xdof_relative_eef_relative_joint_subtask"
+    """
+    The generic X-DOF robot (subtask variant).
+    """
+
+    REAL_G1 = "real_g1_relative_eef_relative_joints"
+    """
+    Real-world Unitree G1 with relative EEF and relative joint actions.
+    """
+
+    REAL_R1_PRO_SHARPA = "real_r1_pro_sharpa_relative_eef"
+    """
+    Real-world R1 Pro Sharpa with relative EEF actions.
+    """
+
+    REAL_R1_PRO_SHARPA_HUMAN = "real_r1_pro_sharpa_relative_eef_human"
+    """
+    Real-world R1 Pro Sharpa with relative EEF actions (human teleop data).
+    """
+
+    REAL_R1_PRO_SHARPA_MAXINSIGHTS = "real_r1_pro_sharpa_relative_eef_maxinsights"
+    """
+    Real-world R1 Pro Sharpa with relative EEF actions (MaxInsights data, single-cam).
+    """
+
+    REAL_R1_PRO_SHARPA_MECKA = "real_r1_pro_sharpa_relative_eef_mecka"
+    """
+    Real-world R1 Pro Sharpa with relative EEF actions (Mecka data, single-cam).
+    """
+
+    UNITREE_G1 = "unitree_g1_full_body_with_waist_height_nav_cmd"
+    """
+    The Unitree G1 robot (sim, full-body with waist height and nav commands).
+    """
+
+    UNITREE_G1_SONIC = "unitree_g1_sonic"
+    """
+    The Unitree G1 robot with SONIC whole-body controller. VLA action space is SONIC latents.
+    """
+
+    SIMPLER_ENV_GOOGLE = "simpler_env_google"
+    """
+    The SimplerEnv Google robot.
+    """
+
+    SIMPLER_ENV_WIDOWX = "simpler_env_widowx"
+    """
+    The SimplerEnv WidowX robot.
+    """
+
+    LIBERO_PANDA = "libero_sim"
+    """
+    The LIBERO Panda robot (used for LIBERO-Goal, LIBERO-Object, LIBERO-Spatial, LIBERO-10).
+    """
+
+    # New embodiment during post-training
+    NEW_EMBODIMENT = "new_embodiment"
+    """
+    Any new embodiment.
+    """
+
+    @classmethod
+    def resolve(cls, tag: "str | EmbodimentTag") -> "EmbodimentTag":
+        """Resolve a string to an EmbodimentTag, case-insensitively.
+
+        Matches by enum **name** first (e.g. ``"xdof"`` -> ``XDOF``), then by
+        enum **value** (e.g. ``"xdof_relative_eef_relative_joint"`` -> ``XDOF``).
+
+        Raises:
+            ValueError: If *tag* does not match any known embodiment.
+        """
+        if isinstance(tag, cls):
+            return tag
+        key = tag.strip()
+        key_lower = key.lower()
+        # Match by enum name (case-insensitive)
+        for member in cls:
+            if member.name.lower() == key_lower:
+                return member
+        # Match by enum value (case-insensitive)
+        for member in cls:
+            if member.value.lower() == key_lower:
+                return member
+
+        def _fmt(tags):
+            return "\n".join(f"    {m.name:40s} -> {m.value}" for m in tags)
+
+        msg = (
+            f"Unknown embodiment tag: {tag!r}\n\n"
+            f"  Base model tags (work with nvidia/GR00T-N1.7-3B):\n"
+            f"{_fmt(PRETRAIN_TAGS)}\n\n"
+            f"  Posttrain tags (require a finetuned checkpoint):\n"
+            f"{_fmt(POSTTRAIN_TAGS)}\n\n"
+            f"  Finetuning-only tags (for custom robots):\n"
+            f"{_fmt(FINETUNE_ONLY_TAGS)}"
+        )
+        raise ValueError(msg)
+
+    @classmethod
+    def reverse_lookup(cls, value: str) -> "str":
+        """Map a tag value string back to its enum name, or return the value as-is."""
+        for member in cls:
+            if member.value == value:
+                return member.name
+        return value
+
+
+# Module-level tag category sets (cannot be Enum class attributes).
+PRETRAIN_TAGS: frozenset[EmbodimentTag] = frozenset(
+    {
+        EmbodimentTag.OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT,
+        EmbodimentTag.XDOF,
+        EmbodimentTag.XDOF_SUBTASK,
+        EmbodimentTag.REAL_G1,
+        EmbodimentTag.REAL_R1_PRO_SHARPA,
+        EmbodimentTag.REAL_R1_PRO_SHARPA_HUMAN,
+        EmbodimentTag.REAL_R1_PRO_SHARPA_MAXINSIGHTS,
+        EmbodimentTag.REAL_R1_PRO_SHARPA_MECKA,
+    }
+)
+"""Tags baked into the base model (nvidia/GR00T-N1.7-3B) — usable without finetuning."""
+
+POSTTRAIN_TAGS: frozenset[EmbodimentTag] = frozenset(
+    {
+        EmbodimentTag.UNITREE_G1,
+        EmbodimentTag.UNITREE_G1_SONIC,
+        EmbodimentTag.SIMPLER_ENV_GOOGLE,
+        EmbodimentTag.SIMPLER_ENV_WIDOWX,
+        EmbodimentTag.LIBERO_PANDA,
+    }
+)
+"""Tags that require a finetuned checkpoint."""
+
+FINETUNE_ONLY_TAGS: frozenset[EmbodimentTag] = frozenset(
+    {
+        EmbodimentTag.NEW_EMBODIMENT,
+    }
+)
+"""Tags for custom robots (finetuning only, not in any shipped checkpoint)."""
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py
new file mode 100644
index 00000000000..208f01a7cb5
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py
new file mode 100644
index 00000000000..4128d6cae6c
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py
@@ -0,0 +1,654 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Sequence
+from typing import Generic, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy import interpolate
+from scipy.spatial.transform import Rotation, Slerp
+
+from vllm_omni.diffusion.models.gr00t.dataio.state_action.pose import EndEffectorPose, JointPose, Pose
+from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat
+
+PoseType = TypeVar("PoseType", bound=Pose)
+
+
+class ActionChunk(Generic[PoseType]):
+    """
+    Abstract base class for robot action chunking.
+
+    This class provides common functionality for different action chunking types
+    including relative and delta action chunking computation with optional reference frames,
+    interpolation, and format conversion.
+    """
+
+    def __init__(
+        self,
+        poses: Sequence[PoseType],
+        times: Sequence[float] | NDArray[np.float64] | None = None,
+    ):
+        """
+        Initialize action chunking from a list of poses.
+
+        Args:
+            poses: Sequence of Pose objects
+            times: Optional sequence of timestamps for each pose. If None, assumes
+                   uniform spacing starting from 0 with step 1.0
+
+        Raises:
+            ValueError: If action chunking is empty or times length doesn't match poses
+        """
+        if not poses:
+            raise ValueError("ActionChunk must contain at least one pose")
+
+        self._poses: list[PoseType] = list(poses)
+
+        # Set up times
+        if times is None:
+            self._times = np.arange(len(poses), dtype=np.float64)
+        else:
+            if len(times) != len(poses):
+                raise ValueError("Number of times must match number of poses")
+            self._times = np.array(times, dtype=np.float64)
+
+    @property
+    def poses(self) -> list[PoseType]:
+        """Get the list of poses"""
+        return self._poses.copy()
+
+    @property
+    def times(self) -> NDArray[np.float64]:
+        """Get the timestamps"""
+        return self._times.copy()
+
+    @property
+    def num_poses(self) -> int:
+        """Get the number of poses in the action chunking"""
+        return len(self._poses)
+
+    def relative_chunking(self, reference_frame: PoseType | None = None) -> "ActionChunk[PoseType]":
+        """
+        Compute the relative action chunking with respect to a reference frame.
+
+        If reference_frame is None, uses the first pose in the action chunking as reference.
+        All poses are transformed to be relative to the reference frame.
+
+        Args:
+            reference_frame: Optional reference pose. If None, uses first pose.
+
+        Returns:
+            A new ActionChunk of the same type where all poses are relative to the reference frame.
+        """
+        if not self._poses:
+            return self.__class__([], times=[])
+
+        # Use the first pose as the reference if one is not provided.
+        ref_pose = reference_frame if reference_frame is not None else self._poses[0]
+
+        # Use the polymorphic subtraction defined in the Pose subclasses.
+        # The subtraction returns the same type as the operands
+        relative_poses: list[PoseType] = [pose - ref_pose for pose in self._poses]  # type: ignore[misc]
+
+        # Return a new instance of the same action chunking class
+        # (e.g., JointActionChunk or EndEffectorActionChunk)
+        return self.__class__(relative_poses, times=self.times)
+
+    def delta_chunking(self, reference_frame: PoseType | None = None) -> "ActionChunk[PoseType]":
+        """
+        Compute the delta action chunking where each pose represents the relative
+        transformation from the previous frame.
+
+        If reference_frame is provided, it is treated as the first frame, and the
+        first delta will be from reference_frame to the first pose in the action chunking.
+        Otherwise, the first pose in the delta action chunking will be the identity/zero transformation.
+
+        Args:
+            reference_frame: Optional reference pose to use as the first frame.
+
+        Returns:
+            A new ActionChunk of the same type where each pose is relative to the previous pose.
+        """
+        if not self._poses:
+            return self.__class__([], times=[])
+
+        delta_poses: list[PoseType] = []
+
+        # Determine the initial reference for the very first pose.
+        # If a reference_frame is given, the first delta is pose[0] - reference_frame.
+        # If not, the first delta is pose[0] - pose[0], resulting in an identity/zero pose.
+        prev_pose = reference_frame if reference_frame is not None else self._poses[0]
+
+        for current_pose in self._poses:
+            delta: PoseType = current_pose - prev_pose  # type: ignore[assignment]
+            delta_poses.append(delta)
+            prev_pose = current_pose  # Update the reference for the next step
+
+        return self.__class__(delta_poses, times=self.times.tolist())
+
+    def to_absolute_chunking(self, reference_frame: PoseType) -> "ActionChunk[PoseType]":
+        """
+        Convert a relative action chunking to an absolute action chunking by applying
+        the relative poses on top of a reference frame.
+
+        This is the inverse operation of relative_chunking(). Each relative pose
+        is composed with the reference frame to produce absolute poses.
+
+        Args:
+            reference_frame: The reference pose to apply the relative action chunking on top of.
+
+        Returns:
+            A new ActionChunk of the same type with absolute poses.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses
+        """
+        raise NotImplementedError("Subclasses must implement to_absolute_chunking")
+
+    def interpolate(
+        self,
+        num_points: int | None = None,
+        times: NDArray[np.float64] | None = None,
+    ) -> "ActionChunk":
+        """
+        Interpolate the action chunking to generate intermediate poses.
+        Must be implemented by subclasses.
+
+        Args:
+            num_points: Number of evenly-spaced points to generate
+            times: Specific timestamps at which to interpolate
+
+        Returns:
+            A new ActionChunk with interpolated poses
+        """
+        raise NotImplementedError("Subclasses must implement interpolate")
+
+    def to(self, action_format: ActionFormat) -> NDArray[np.float64]:
+        """
+        Convert action chunking to the specified action format.
+        Must be implemented by subclasses.
+
+        Args:
+            action_format: The desired output format
+
+        Returns:
+            Array in the requested format
+
+        Raises:
+            NotImplementedError: If not implemented by subclass
+        """
+        raise NotImplementedError("Subclasses must implement to method")
+
+    def __len__(self) -> int:
+        """Return the number of poses in the action chunking"""
+        return len(self._poses)
+
+    def __getitem__(self, index: int) -> PoseType:
+        """Get a pose by index"""
+        return self._poses[index]
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(num_poses={len(self._poses)}, "
+            f"time_range=[{self._times[0]:.2f}, {self._times[-1]:.2f}])"
+        )
+
+
+class JointActionChunk(ActionChunk[JointPose]):
+    """
+    Represents action chunking in joint space as a sequence of joint configurations.
+
+    Examples:
+        # Create a joint action chunking
+        joint_poses = [
+            JointPose([0.0, 0.0, 0.0]),
+            JointPose([0.5, 0.5, 0.5]),
+            JointPose([1.0, 1.0, 1.0]),
+        ]
+        action_chunking = JointActionChunk(joint_poses)
+
+        # Get relative trajectory (all poses relative to first pose)
+        relative_traj = action_chunking.relative_chunking()
+
+        # Get relative trajectory with custom reference
+        reference = JointPose([0.1, 0.1, 0.1])
+        relative_traj = action_chunking.relative_chunking(reference_frame=reference)
+
+        # Get delta trajectory (incremental changes)
+        delta_traj = action_chunking.delta_chunking()
+
+        # Convert relative trajectory back to absolute
+        reference = JointPose([0.1, 0.1, 0.1])
+        absolute_traj = relative_traj.to_absolute_chunking(reference_frame=reference)
+
+        # Interpolate trajectory
+        interpolated = action_chunking.interpolate(num_points=10)
+
+        # Convert to desired format
+        from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat
+        array_data = action_chunking.to(ActionFormat.DEFAULT)  # Returns joint array
+    """
+
+    def __init__(
+        self,
+        poses: Sequence[JointPose],
+        times: Sequence[float] | NDArray[np.float64] | None = None,
+    ):
+        """
+        Initialize a joint trajectory from a list of joint poses.
+
+        Args:
+            poses: Sequence of JointPose objects
+            times: Optional sequence of timestamps for each pose
+
+        Raises:
+            TypeError: If poses are not all JointPose objects
+        """
+        # Validate all poses are JointPose
+        if not all(isinstance(p, JointPose) for p in poses):
+            raise TypeError("All poses must be JointPose objects for JointActionChunk")
+
+        super().__init__(poses, times)
+
+    def interpolate(
+        self,
+        num_points: int | None = None,
+        times: NDArray[np.float64] | None = None,
+    ) -> "JointActionChunk":
+        """
+        Interpolate the joint action chunking to generate intermediate configurations.
+
+        Uses linear interpolation for joint values.
+
+        Args:
+            num_points: Number of evenly-spaced points to generate (including endpoints).
+                       Only used if times is None.
+            times: Specific timestamps at which to interpolate. If provided,
+                  num_points is ignored.
+
+        Returns:
+            A new JointActionChunk with interpolated poses
+
+        Raises:
+            ValueError: If neither num_points nor times is provided, or if
+                       interpolation times are outside the trajectory range
+        """
+        if num_points is None and times is None:
+            raise ValueError("Must provide either num_points or times")
+
+        if len(self._poses) < 2:
+            raise ValueError("Need at least 2 poses for interpolation")
+
+        # Prepare data: extract joint values
+        timestamps = self._times.copy()
+        joint_values = np.array([pose.joints for pose in self._poses])  # (N, num_joints)
+
+        # Find and remove non-monotonic timestamps
+        drop_indices = [idx for idx in range(1, len(timestamps)) if timestamps[idx] <= timestamps[idx - 1]]
+
+        if drop_indices:
+            for idx in drop_indices:
+                print(
+                    f"Dropping timestamp pair - Previous: {timestamps[idx - 1]}, "
+                    f"Current: {timestamps[idx]} at index {idx}"
+                )
+            timestamps = np.delete(timestamps, drop_indices)
+            joint_values = np.delete(joint_values, drop_indices, axis=0)
+
+        # Check if we still have enough poses after cleanup
+        if len(timestamps) < 2:
+            raise ValueError("Need at least 2 poses with monotonic timestamps for interpolation")
+
+        # Create interpolator
+        joint_interp = interpolate.interp1d(timestamps, joint_values, kind="linear", axis=0)
+
+        # Generate interpolation times if not provided
+        if times is None:
+            assert num_points is not None  # Type narrowing for type checker
+            interp_times = np.linspace(timestamps[0], timestamps[-1], num_points)
+        else:
+            interp_times = np.array(times, dtype=np.float64)
+
+        # Check that interpolation times are within bounds
+        if np.any(interp_times < timestamps[0]) or np.any(interp_times > timestamps[-1]):
+            raise ValueError(f"Interpolation times must be within [{timestamps[0]}, {timestamps[-1]}]")
+
+        # Interpolate joint values
+        interp_joint_values = joint_interp(interp_times)
+
+        # Create interpolated poses
+        joint_names = self._poses[0].joint_names
+        interpolated_poses = [
+            JointPose(joints=interp_joint_values[i], joint_names=joint_names) for i in range(len(interp_times))
+        ]
+
+        return JointActionChunk(interpolated_poses, times=interp_times)
+
+    def to_array(self) -> NDArray[np.float64]:
+        """
+        Convert trajectory to array of joint values.
+
+        Returns:
+            Array with shape (N, num_joints) where N is the number of poses
+        """
+        return np.array([pose.joints for pose in self._poses])
+
+    def to_absolute_chunking(self, reference_frame: JointPose) -> "JointActionChunk":
+        """
+        Convert a relative joint action chunking to an absolute action chunking by adding
+        the relative joint positions to the reference frame.
+
+        This is the inverse operation of relative_chunking(). Each relative joint
+        configuration is added to the reference frame to produce absolute joint positions.
+
+        Args:
+            reference_frame: The reference joint pose to apply the relative trajectory on top of.
+
+        Returns:
+            A new JointActionChunk with absolute joint positions.
+
+        Raises:
+            ValueError: If joint dimensions don't match
+        """
+        if not self._poses:
+            return JointActionChunk([], times=[])
+
+        if len(self._poses[0].joints) != len(reference_frame.joints):
+            raise ValueError(
+                f"Cannot apply relative trajectory: "
+                f"joint dimensions don't match ({len(self._poses[0].joints)} vs {len(reference_frame.joints)})"
+            )
+
+        # Add each relative pose to the reference frame
+        absolute_poses: list[JointPose] = []
+        for relative_pose in self._poses:
+            absolute_joints = reference_frame.joints + relative_pose.joints
+            absolute_pose = JointPose(joints=absolute_joints, joint_names=reference_frame.joint_names)
+            absolute_poses.append(absolute_pose)
+
+        return JointActionChunk(absolute_poses, times=self.times)
+
+    def to(self, action_format: ActionFormat) -> NDArray[np.float64]:
+        """
+        Convert trajectory to the desired format.
+
+        Args:
+            action_format: The desired output format
+
+        Returns:
+            Array in the requested format
+
+        Raises:
+            ValueError: If the action format is not supported for joint trajectories
+        """
+        if action_format == ActionFormat.DEFAULT:
+            return self.to_array()
+        else:
+            raise ValueError(
+                f"ActionFormat {action_format} is not supported for JointActionChunk. "
+                f"Only {ActionFormat.DEFAULT} is supported."
+            )
+
+
+class EndEffectorActionChunk(ActionChunk[EndEffectorPose]):
+    """
+    Represents action chunking in Cartesian space as a sequence of end-effector poses.
+
+    Examples:
+        # Create an end-effector action chunking
+        ee_poses = [
+            EndEffectorPose(translation=[0, 0, 0], rotation=[1, 0, 0, 0],
+                          rotation_type="quat", rotation_order="wxyz"),
+            EndEffectorPose(translation=[1, 0, 0], rotation=[0.707, 0, 0, 0.707],
+                          rotation_type="quat", rotation_order="wxyz"),
+            EndEffectorPose(translation=[2, 0, 0], rotation=[0, 0, 0, 1],
+                          rotation_type="quat", rotation_order="wxyz"),
+        ]
+        action_chunking = EndEffectorActionChunk(ee_poses)
+
+        # Get relative trajectory (all poses relative to first pose)
+        relative_traj = action_chunking.relative_chunking()
+
+        # Get relative trajectory with custom reference frame
+        reference = EndEffectorPose(translation=[0.5, 0, 0], rotation=[1,0,0,0],
+                                   rotation_type="quat", rotation_order="wxyz")
+        relative_traj = action_chunking.relative_chunking(reference_frame=reference)
+
+        # Get delta trajectory
+        delta_traj = action_chunking.delta_chunking()
+
+        # Convert relative trajectory back to absolute
+        reference = EndEffectorPose(translation=[0.5, 0, 0], rotation=[1,0,0,0],
+                                   rotation_type="quat", rotation_order="wxyz")
+        absolute_traj = relative_traj.to_absolute_chunking(reference_frame=reference)
+
+        # Interpolate trajectory
+        interpolated = action_chunking.interpolate(num_points=10)
+
+        # Convert to desired format
+        from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat
+        homo_matrices = action_chunking.to(ActionFormat.DEFAULT)      # (N, 4, 4) homogeneous matrices
+        xyz_rot6d = action_chunking.to(ActionFormat.XYZ_ROT6D)        # (N, 9) xyz + rot6d
+        xyz_rotvec = action_chunking.to(ActionFormat.XYZ_ROTVEC)      # (N, 6) xyz + rotvec
+    """
+
+    def __init__(
+        self,
+        poses: Sequence[EndEffectorPose],
+        times: Sequence[float] | NDArray[np.float64] | None = None,
+    ):
+        """
+        Initialize an end-effector trajectory from a list of end-effector poses.
+
+        Args:
+            poses: Sequence of EndEffectorPose objects
+            times: Optional sequence of timestamps for each pose
+
+        Raises:
+            TypeError: If poses are not all EndEffectorPose objects
+        """
+        # Validate all poses are EndEffectorPose
+        if not all(isinstance(p, EndEffectorPose) for p in poses):
+            raise TypeError("All poses must be EndEffectorPose objects for EndEffectorActionChunk")
+
+        super().__init__(poses, times)
+
+    @classmethod
+    def from_array(cls, data: np.ndarray, action_format: ActionFormat) -> "EndEffectorActionChunk":
+        """
+        Create an EndEffectorActionChunk from a 2-D array using the specified action format.
+
+        This is the inverse of ``.to(action_format)``.
+
+        Args:
+            data: Array of shape (N, D) where D depends on the action_format.
+            action_format: The format that describes the layout of each row.
+
+        Returns:
+            EndEffectorActionChunk with N poses.
+        """
+        poses = [EndEffectorPose.from_action_format(row, action_format) for row in data]
+        return cls(poses)
+
+    def interpolate(
+        self,
+        num_points: int | None = None,
+        times: NDArray[np.float64] | None = None,
+    ) -> "EndEffectorActionChunk":
+        """
+        Interpolate the action chunking to generate intermediate poses.
+
+        Uses linear interpolation for translation and SLERP (Spherical Linear
+        Interpolation) for rotation.
+
+        Args:
+            num_points: Number of evenly-spaced points to generate (including endpoints).
+                       Only used if times is None.
+            times: Specific timestamps at which to interpolate. If provided,
+                  num_points is ignored.
+
+        Returns:
+            A new EndEffectorActionChunk with interpolated poses
+
+        Raises:
+            ValueError: If neither num_points nor times is provided, or if
+                       interpolation times are outside the trajectory range
+        """
+        if num_points is None and times is None:
+            raise ValueError("Must provide either num_points or times")
+
+        if len(self._poses) < 2:
+            raise ValueError("Need at least 2 poses for interpolation")
+
+        # Prepare data: extract positions and rotations
+        timestamps = self._times.copy()
+        homogeneous_matrices = np.array([pose.homogeneous for pose in self._poses])
+        positions = homogeneous_matrices[:, :3, 3]
+        rotations = Rotation.from_matrix(homogeneous_matrices[:, :3, :3])
+
+        # Find indices where timestamps are not monotonically increasing
+        drop_indices = [idx for idx in range(1, len(timestamps)) if timestamps[idx] <= timestamps[idx - 1]]
+
+        # Remove the problematic timestamps and corresponding data
+        if drop_indices:
+            for idx in drop_indices:
+                print(
+                    f"Dropping timestamp pair - Previous: {timestamps[idx - 1]}, "
+                    f"Current: {timestamps[idx]} at index {idx}"
+                )
+            timestamps = np.delete(timestamps, drop_indices)
+            positions = np.delete(positions, drop_indices, axis=0)
+            rotations = Rotation.from_matrix(np.delete(homogeneous_matrices[:, :3, :3], drop_indices, axis=0))
+
+        # Check if we still have enough poses after cleanup
+        if len(timestamps) < 2:
+            raise ValueError("Need at least 2 poses with monotonic timestamps for interpolation")
+
+        # Create interpolators
+        pos_interp = interpolate.interp1d(timestamps, positions, kind="linear", axis=0)
+        rot_interp = Slerp(timestamps, rotations)
+
+        # Generate interpolation times if not provided
+        if times is None:
+            assert num_points is not None  # Type narrowing for type checker
+            interp_times = np.linspace(timestamps[0], timestamps[-1], num_points)
+        else:
+            interp_times = np.array(times, dtype=np.float64)
+
+        # Check that interpolation times are within bounds
+        if np.any(interp_times < timestamps[0]) or np.any(interp_times > timestamps[-1]):
+            raise ValueError(f"Interpolation times must be within [{timestamps[0]}, {timestamps[-1]}]")
+
+        # Interpolate positions and rotations
+        interp_positions = pos_interp(interp_times)
+        interp_rotations = rot_interp(interp_times)
+
+        # Create interpolated poses
+        interpolated_poses = []
+        for i in range(len(interp_times)):
+            pose = EndEffectorPose(
+                translation=interp_positions[i],
+                rotation=interp_rotations[i].as_matrix(),
+                rotation_type="matrix",
+            )
+            interpolated_poses.append(pose)
+
+        return EndEffectorActionChunk(interpolated_poses, times=interp_times)
+
+    def to_homogeneous_matrices(self) -> NDArray[np.float64]:
+        """
+        Convert trajectory to array of homogeneous transformation matrices.
+
+        Returns:
+            Array of homogeneous matrices with shape (N, 4, 4) where N is the number of poses
+        """
+        return np.array([pose.homogeneous for pose in self._poses])
+
+    def to_translation_rot6d(self) -> NDArray[np.float64]:
+        """
+        Convert trajectory to array of translations and 6D rotations.
+
+        Returns:
+            Array with shape (N, 9) - 3 for xyz + 6 for rot6d
+        """
+        translations = np.array([pose.translation for pose in self._poses])  # (N, 3)
+        rotations = np.array([pose.rot6d for pose in self._poses])  # (N, 6)
+
+        # Concatenate translation and rotation
+        xyz_rot6d = np.concatenate([translations, rotations], axis=1)  # (N, 9)
+
+        return xyz_rot6d
+
+    def to_translation_rotvec(self) -> NDArray[np.float64]:
+        """
+        Convert trajectory to array of translations and rotation vectors.
+
+        Returns:
+            Array with shape (N, 6) - 3 for xyz + 3 for rotvec
+        """
+        translations = np.array([pose.translation for pose in self._poses])  # (N, 3)
+        rotations = np.array([pose.rotvec for pose in self._poses])  # (N, 3)
+
+        # Concatenate translation and rotation
+        xyz_rotvec = np.concatenate([translations, rotations], axis=1)  # (N, 6)
+
+        return xyz_rotvec
+
+    def to_absolute_chunking(self, reference_frame: EndEffectorPose) -> "EndEffectorActionChunk":
+        """
+        Convert a relative end-effector action chunking to an absolute action chunking by
+        composing each relative transformation with the reference frame.
+
+        This is the inverse operation of relative_chunking(). Each relative pose
+        represents a transformation that is applied on top of the reference frame
+        to produce absolute poses.
+
+        Args:
+            reference_frame: The reference end-effector pose to apply the relative trajectory on top of.
+
+        Returns:
+            A new EndEffectorActionChunk with absolute poses.
+        """
+        if not self._poses:
+            return EndEffectorActionChunk([], times=[])
+
+        # Get reference frame as homogeneous matrix
+        T_ref = reference_frame.homogeneous
+
+        # Compose each relative transformation with the reference frame
+        absolute_poses: list[EndEffectorPose] = []
+        for relative_pose in self._poses:
+            # Get relative transformation as homogeneous matrix
+            T_relative = relative_pose.homogeneous
+
+            # Compose transformations: T_absolute = T_ref @ T_relative
+            T_absolute = T_ref @ T_relative
+
+            # Create absolute pose from composed transformation
+            absolute_pose = EndEffectorPose(homogeneous=T_absolute)
+            absolute_poses.append(absolute_pose)
+
+        return EndEffectorActionChunk(absolute_poses, times=self.times)
+
+    def to(self, action_format: ActionFormat) -> NDArray[np.float64]:
+        """
+        Convert trajectory to the desired format.
+
+        Args:
+            action_format: The desired output format
+
+        Returns:
+            Array in the requested format
+
+        Raises:
+            ValueError: If the action format is not supported
+        """
+        if action_format == ActionFormat.DEFAULT:
+            return self.to_homogeneous_matrices()
+        elif action_format == ActionFormat.XYZ_ROT6D:
+            return self.to_translation_rot6d()
+        elif action_format == ActionFormat.XYZ_ROTVEC:
+            return self.to_translation_rotvec()
+        else:
+            raise ValueError(f"Unsupported action format: {action_format}")
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py
new file mode 100644
index 00000000000..0e0849f17b4
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py
@@ -0,0 +1,709 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from enum import Enum
+from typing import TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.spatial.transform import Rotation
+
+from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat
+
+# TypeVar for self-type preservation in Pose operations
+PoseT = TypeVar("PoseT", bound="Pose")
+
+
+def invert_transformation(transform: NDArray[np.float64]) -> NDArray[np.float64]:
+    """
+    Invert a homogeneous transformation matrix.
+
+    Args:
+        transform: A 4x4 homogeneous transformation matrix
+
+    Returns:
+        The inverse of the transformation matrix (4x4)
+    """
+    R = transform[:3, :3]  # Extract the rotation matrix
+    t = transform[:3, 3]  # Extract the translation vector
+
+    # Inverse of the rotation matrix is its transpose (since it's orthogonal)
+    R_inv = R.T
+
+    # Inverse of the translation is -R_inv * t
+    t_inv = -R_inv @ t
+
+    # Construct the inverse transformation matrix
+    T_inv = np.eye(4)
+    T_inv[:3, :3] = R_inv
+    T_inv[:3, 3] = t_inv
+
+    return T_inv
+
+
+def relative_transformation(
+    base_transform: NDArray[np.float64], target_transform: NDArray[np.float64]
+) -> NDArray[np.float64]:
+    """
+    Compute the relative transformation between two poses.
+
+    Args:
+        base_transform: Initial 4x4 homogeneous transformation matrix
+        target_transform: Current 4x4 homogeneous transformation matrix
+
+    Returns:
+        The relative transformation matrix (4x4) from base_transform to target_transform
+    """
+    # Relative transformation is base_transform^{-1} * target_transform
+    T_relative = invert_transformation(base_transform) @ target_transform
+    return T_relative
+
+
+class RotationType(Enum):
+    """Supported rotation representation types"""
+
+    QUAT = "quat"
+    EULER = "euler"
+    ROTVEC = "rotvec"
+    MATRIX = "matrix"
+    ROT6D = "rot6d"
+
+
+class EulerOrder(Enum):
+    """Common Euler angle conventions"""
+
+    XYZ = "xyz"
+    ZYX = "zyx"
+    XZY = "xzy"
+    YXZ = "yxz"
+    YZX = "yzx"
+    ZXY = "zxy"
+
+
+class QuatOrder(Enum):
+    """Quaternion ordering conventions"""
+
+    WXYZ = "wxyz"  # scalar-first (w, x, y, z)
+    XYZW = "xyzw"  # scalar-last (x, y, z, w)
+
+
+class Pose:
+    """
+    Abstract base class for robot poses.
+
+    This class provides common functionality for different pose representations
+    including relative pose computation via the subtraction operator.
+    """
+
+    pose_type: str
+
+    def __sub__(self: PoseT, other: PoseT) -> PoseT:
+        """
+        Compute relative transformation between two poses.
+
+        For EndEffectorPose: Computes the relative transformation from other to self.
+        Result represents the transformation needed to go from other's frame to self's frame.
+
+        For JointPose: Computes the joint-space difference (self - other).
+
+        Args:
+            other: The reference pose to compute relative transformation from
+
+        Returns:
+            Relative pose (same type as self)
+
+        Raises:
+            TypeError: If poses are not of the same type
+
+        Examples:
+            # End-effector poses
+            pose1 = EndEffectorPose(translation=[1, 0, 0], rotation=[1,0,0,0],
+                                   rotation_type="quat", rotation_order="wxyz")
+            pose2 = EndEffectorPose(translation=[2, 0, 0], rotation=[1,0,0,0],
+                                   rotation_type="quat", rotation_order="wxyz")
+            relative = pose2 - pose1  # Transformation from pose1 to pose2
+
+            # Joint poses
+            joint1 = JointPose([0.0, 0.5, 1.0])
+            joint2 = JointPose([0.1, 0.6, 1.2])
+            joint_diff = joint2 - joint1  # Joint differences: [0.1, 0.1, 0.2]
+        """
+        if type(self) is not type(other):
+            raise TypeError(
+                f"Cannot compute relative transformation between different pose types: "
+                f"{type(self).__name__} and {type(other).__name__}"
+            )
+
+        return self._compute_relative(other)
+
+    def _compute_relative(self: PoseT, other: PoseT) -> PoseT:
+        """
+        Internal method to compute relative transformation.
+        Must be implemented by subclasses.
+
+        Args:
+            other: The reference pose
+
+        Returns:
+            Relative pose
+        """
+        raise NotImplementedError("Subclasses must implement _compute_relative")
+
+    def copy(self: PoseT) -> PoseT:
+        """
+        Create a deep copy of this pose.
+        Must be implemented by subclasses.
+
+        Returns:
+            New Pose instance with copied data
+        """
+        raise NotImplementedError("Subclasses must implement copy")
+
+
+class JointPose(Pose):
+    """
+    Represents a robot configuration in joint space.
+
+    This class stores joint angles/positions for a robot manipulator.
+    Unlike end-effector poses, joint poses represent the configuration
+    of all joints in the kinematic chain.
+
+    Examples:
+        # Create a 6-DOF joint configuration
+        joint_pose = JointPose(
+            joints=[0.0, -np.pi/4, np.pi/2, 0.0, np.pi/4, 0.0],
+            joint_names=["shoulder_pan", "shoulder_lift", "elbow",
+                        "wrist_1", "wrist_2", "wrist_3"]
+        )
+
+        # Create with default joint names
+        joint_pose = JointPose(joints=[0.0, 0.5, 1.0])
+
+        # Get as dictionary
+        joint_dict = joint_pose.to_dict()  # {"joint_0": 0.0, ...}
+
+        # Access individual joints
+        first_joint = joint_pose.joints[0]
+        num_joints = joint_pose.num_joints
+
+        # Compute relative joint displacement
+        joint1 = JointPose([0.0, 0.5, 1.0])
+        joint2 = JointPose([0.1, 0.6, 1.2])
+        relative = joint2 - joint1  # [0.1, 0.1, 0.2]
+    """
+
+    pose_type = "joint"
+
+    def __init__(
+        self,
+        joints: list | np.ndarray,
+        joint_names: list | None = None,
+    ):
+        """
+        Initialize a joint pose.
+
+        Args:
+            joints: Joint angles/positions as array-like of shape (n,)
+            joint_names: Optional list of names for each joint. If None,
+                        defaults to ["joint_0", "joint_1", ...]
+        """
+        super().__init__()
+        self.joints = np.array(joints, dtype=np.float64)
+
+        # Set defaults and validate joint_names
+        if joint_names is None:
+            self.joint_names = [f"joint_{i}" for i in range(len(self.joints))]
+        else:
+            if len(joint_names) != len(self.joints):
+                raise ValueError(
+                    f"Number of joint names ({len(joint_names)}) must match number of joints ({len(self.joints)})"
+                )
+            self.joint_names = joint_names
+
+    @property
+    def num_joints(self) -> int:
+        """
+        Get the number of joints.
+
+        Returns:
+            Number of joints in the configuration
+        """
+        return len(self.joints)
+
+    def to_dict(self) -> dict:
+        """
+        Convert joint configuration to dictionary.
+
+        Returns:
+            Dictionary mapping joint names to joint values
+        """
+        return dict(zip(self.joint_names, self.joints))
+
+    def _compute_relative(self, other):  # type: ignore[override]
+        """
+        Compute relative joint displacement.
+
+        Args:
+            other: Reference joint pose
+
+        Returns:
+            JointPose representing the joint-space difference (self - other)
+
+        Raises:
+            ValueError: If joint dimensions don't match
+        """
+        if len(self.joints) != len(other.joints):
+            raise ValueError(
+                f"Cannot compute relative joint pose: "
+                f"joint dimensions don't match ({len(self.joints)} vs {len(other.joints)})"
+            )
+
+        relative_joints = self.joints - other.joints
+        return JointPose(joints=relative_joints, joint_names=self.joint_names)
+
+    def copy(self) -> "JointPose":
+        """
+        Create a deep copy of this joint pose.
+
+        Returns:
+            New JointPose instance with copied data
+        """
+        return JointPose(
+            joints=self.joints.copy(),
+            joint_names=self.joint_names.copy(),
+        )
+
+    def __repr__(self) -> str:
+        if len(self.joints) <= 6:
+            joints_str = np.array2string(self.joints, precision=4, suppress_small=True)
+        else:
+            joints_str = f"[{self.joints[0]:.4f}, ..., {self.joints[-1]:.4f}] ({len(self.joints)} joints)"
+
+        return f"JointPose(joints={joints_str})"
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, JointPose):
+            return False
+        return np.allclose(self.joints, other.joints) and self.joint_names == other.joint_names
+
+    def __getitem__(self, index) -> float | NDArray[np.float64]:
+        """Allow indexing: joint_pose[0] returns first joint value"""
+        return self.joints[index]
+
+    def __len__(self) -> int:
+        """Allow len(): len(joint_pose) returns number of joints"""
+        return len(self.joints)
+
+
+class EndEffectorPose(Pose):
+    """
+    Represents a single end-effector pose with translation and rotation components.
+
+    This class handles Cartesian space representations of robot end-effector poses,
+    supporting multiple rotation representations (quaternions, Euler angles, rotation
+    vectors, rotation matrices, etc.).
+
+    Examples:
+        # Create with quaternion (wxyz order)
+        pose = EndEffectorPose(
+            translation=[1.0, 2.0, 3.0],
+            rotation=[1.0, 0.0, 0.0, 0.0],
+            rotation_type="quat",
+            rotation_order="wxyz"
+        )
+
+        # Create with Euler angles (degrees by default)
+        pose = EndEffectorPose(
+            translation=[1, 2, 3],
+            rotation=[0, 0, 90],
+            rotation_type="euler",
+            rotation_order="xyz"
+        )
+
+        # Create with Euler angles in radians
+        pose = EndEffectorPose(
+            translation=[1, 2, 3],
+            rotation=[0, 0, np.pi/2],
+            rotation_type="euler",
+            rotation_order="xyz",
+            degrees=False
+        )
+
+        # Create from homogeneous matrix
+        H = np.eye(4)
+        H[:3, 3] = [1, 2, 3]
+        pose = EndEffectorPose(homogeneous=H)
+
+        # Convert between representations
+        quat_wxyz = pose.to_rotation("quat", "wxyz")
+        euler_zyx = pose.to_rotation("euler", "zyx")
+        rot6d = pose.to_rotation("rot6d")
+
+        # Compute relative transformation
+        pose1 = EndEffectorPose(translation=[1, 0, 0], rotation=[1,0,0,0],
+                               rotation_type="quat", rotation_order="wxyz")
+        pose2 = EndEffectorPose(translation=[2, 0, 0], rotation=[1,0,0,0],
+                               rotation_type="quat", rotation_order="wxyz")
+        relative = pose2 - pose1  # Transformation from pose1's frame to pose2's frame
+    """
+
+    pose_type = "end_effector"
+
+    def __init__(
+        self,
+        translation: list | np.ndarray | None = None,
+        rotation: list | np.ndarray | None = None,
+        rotation_type: str | None = None,
+        rotation_order: str | None = None,
+        homogeneous: np.ndarray | None = None,
+        degrees: bool = True,
+    ):
+        """
+        Initialize an end-effector pose.
+
+        Args:
+            translation: Translation vector [x, y, z]
+            rotation: Rotation in specified format
+            rotation_type: Type of rotation ("quat", "euler", "rotvec", "matrix", "rot6d")
+            rotation_order: Order/convention for the rotation type
+            homogeneous: Homogeneous transformation matrix (4, 4)
+                        If provided, overrides translation and rotation
+            degrees: For Euler angles, whether the input is in degrees (default True)
+        """
+        super().__init__()
+
+        # Cache for homogeneous matrix
+        self._homogeneous_cache: NDArray[np.float64] | None = None
+        self._cache_valid = False
+
+        # Handle homogeneous matrix input
+        if homogeneous is not None:
+            self._from_homogeneous(homogeneous)
+            return
+
+        # Store translation
+        self._translation = np.array(translation) if translation is not None else np.zeros(3)
+
+        # Store rotation as scipy Rotation object internally
+        if rotation is not None:
+            if rotation_type is None:
+                raise ValueError("rotation_type must be specified when rotation is provided")
+            self._set_rotation(rotation, rotation_type, rotation_order, degrees)
+        else:
+            self._rotation = Rotation.identity()
+
+    def _from_homogeneous(self, homogeneous: np.ndarray):
+        """Initialize from homogeneous transformation matrix"""
+        homogeneous = np.array(homogeneous)
+
+        # Extract translation (last column, first 3 rows)
+        self._translation = homogeneous[:3, 3]
+
+        # Extract rotation matrix (top-left 3x3)
+        rotation_matrix = homogeneous[:3, :3]
+
+        # Create Rotation object from matrix
+        self._rotation = Rotation.from_matrix(rotation_matrix)
+
+    @staticmethod
+    def _rot6d_to_matrix(rot6d: np.ndarray) -> np.ndarray:
+        """
+        Convert 6D rotation representation to rotation matrix.
+
+        Args:
+            rot6d: 6D rotation as (6,) array - first two rows of rotation matrix flattened
+
+        Returns:
+            Rotation matrix (3, 3)
+        """
+        rot6d = rot6d.reshape(2, 3)
+
+        # First two rows
+        row1 = rot6d[0]
+        row2 = rot6d[1]
+
+        # Normalize first row
+        row1 = row1 / np.linalg.norm(row1)
+
+        # Gram-Schmidt orthogonalization for second row
+        row2 = row2 - np.dot(row1, row2) * row1
+        row2 = row2 / np.linalg.norm(row2)
+
+        # Third row is cross product
+        row3 = np.cross(row1, row2)
+
+        # Construct rotation matrix
+        rotation_matrix = np.vstack([row1, row2, row3])
+
+        return rotation_matrix
+
+    @staticmethod
+    def _matrix_to_rot6d(rotation_matrix: np.ndarray) -> np.ndarray:
+        """
+        Convert rotation matrix to 6D rotation representation.
+
+        Args:
+            rotation_matrix: Rotation matrix (3, 3)
+
+        Returns:
+            6D rotation - (6,) array (first two rows flattened)
+        """
+        return rotation_matrix[:2, :].flatten()
+
+    def _set_rotation(
+        self,
+        rotation: list | np.ndarray,
+        rotation_type: str,
+        rotation_order: str | None = None,
+        degrees: bool = True,
+    ):
+        """Internal method to set rotation from various representations"""
+        rotation = np.array(rotation)
+        rot_type = RotationType(rotation_type.lower())
+
+        if rot_type == RotationType.QUAT:
+            quat_order = QuatOrder(rotation_order.lower()) if rotation_order else QuatOrder.WXYZ
+            if quat_order == QuatOrder.WXYZ:
+                # scipy uses xyzw order, so convert
+                quat_xyzw = np.array([rotation[1], rotation[2], rotation[3], rotation[0]])
+            else:
+                quat_xyzw = rotation
+            self._rotation = Rotation.from_quat(quat_xyzw)
+
+        elif rot_type == RotationType.EULER:
+            euler_order = EulerOrder(rotation_order.lower()) if rotation_order else EulerOrder.XYZ
+            self._rotation = Rotation.from_euler(euler_order.value, rotation, degrees=degrees)
+
+        elif rot_type == RotationType.ROTVEC:
+            self._rotation = Rotation.from_rotvec(rotation)
+
+        elif rot_type == RotationType.MATRIX:
+            self._rotation = Rotation.from_matrix(rotation)
+
+        elif rot_type == RotationType.ROT6D:
+            rotation_matrix = self._rot6d_to_matrix(rotation)
+            self._rotation = Rotation.from_matrix(rotation_matrix)
+
+        else:
+            raise ValueError(f"Unknown rotation type: {rotation_type}")
+
+        # Invalidate cache
+        self._cache_valid = False
+
+    @property
+    def translation(self) -> np.ndarray:
+        """
+        Get translation vector.
+
+        Returns:
+            Translation array - shape (3,)
+        """
+        return self._translation.copy()
+
+    @property
+    def quat_wxyz(self) -> np.ndarray:
+        """Get rotation as quaternion in wxyz order (w, x, y, z)"""
+        return self.to_rotation("quat", "wxyz")
+
+    @property
+    def quat_xyzw(self) -> np.ndarray:
+        """Get rotation as quaternion in xyzw order (x, y, z, w)"""
+        return self.to_rotation("quat", "xyzw")
+
+    @property
+    def euler_xyz(self) -> np.ndarray:
+        """Get rotation as Euler angles in xyz order (degrees)"""
+        return self.to_rotation("euler", "xyz")
+
+    @property
+    def rotvec(self) -> np.ndarray:
+        """Get rotation as rotation vector (axis-angle)"""
+        return self.to_rotation("rotvec")
+
+    @property
+    def rotation_matrix(self) -> np.ndarray:
+        """Get rotation as 3x3 rotation matrix"""
+        return self.to_rotation("matrix")
+
+    @property
+    def rot6d(self) -> np.ndarray:
+        """Get rotation as 6D representation (first two rows of rotation matrix)"""
+        return self.to_rotation("rot6d")
+
+    @property
+    def xyz_rot6d(self) -> np.ndarray:
+        """Get pose as concatenated translation and 6D rotation (9,)"""
+        return np.concatenate([self._translation, self.rot6d])
+
+    @property
+    def xyz_rotvec(self) -> np.ndarray:
+        """Get pose as concatenated translation and rotation vector (6,)"""
+        return np.concatenate([self._translation, self.rotvec])
+
+    @property
+    def homogeneous(self) -> np.ndarray:
+        """
+        Get homogeneous transformation matrix.
+
+        Returns:
+            Homogeneous matrix - shape (4, 4)
+        """
+        if not self._cache_valid:
+            self._homogeneous_cache = self._compute_homogeneous()
+            self._cache_valid = True
+        assert self._homogeneous_cache is not None
+        return self._homogeneous_cache.copy()
+
+    def _compute_homogeneous(self) -> np.ndarray:
+        """Compute homogeneous transformation matrix"""
+        H = np.eye(4)
+        H[:3, :3] = self._rotation.as_matrix()
+        H[:3, 3] = self._translation
+        return H
+
+    def to_rotation(
+        self,
+        rotation_type: str,
+        rotation_order: str | None = None,
+        degrees: bool = True,
+    ) -> np.ndarray:
+        """
+        Get rotation in specified representation.
+
+        Args:
+            rotation_type: Desired type ("quat", "euler", "rotvec", "matrix", "rot6d")
+            rotation_order: Order/convention for the rotation type
+            degrees: For Euler angles, return in degrees (default True)
+
+        Returns:
+            Rotation in requested format
+            - Shape (4,) for quat
+            - Shape (3,) for euler/rotvec
+            - Shape (6,) for rot6d
+            - Shape (3, 3) for matrix
+        """
+        rot_type = RotationType(rotation_type.lower())
+
+        if rot_type == RotationType.ROT6D:
+            rotation_matrix = self._rotation.as_matrix()
+            return self._matrix_to_rot6d(rotation_matrix)
+
+        if rot_type == RotationType.QUAT:
+            quat_order = QuatOrder(rotation_order.lower()) if rotation_order else QuatOrder.WXYZ
+            quat_xyzw = self._rotation.as_quat()
+            if quat_order == QuatOrder.WXYZ:
+                return np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])
+            else:
+                return quat_xyzw
+
+        elif rot_type == RotationType.EULER:
+            euler_order = EulerOrder(rotation_order.lower()) if rotation_order else EulerOrder.XYZ
+            return self._rotation.as_euler(euler_order.value, degrees=degrees)
+
+        elif rot_type == RotationType.ROTVEC:
+            return self._rotation.as_rotvec()
+
+        elif rot_type == RotationType.MATRIX:
+            return self._rotation.as_matrix()
+
+        else:
+            raise ValueError(f"Unknown rotation type: {rotation_type}")
+
+    def to_homogeneous(self) -> np.ndarray:
+        """
+        Convert pose to homogeneous transformation matrix.
+        (Alias for the homogeneous property)
+
+        Returns:
+            Homogeneous matrix - shape (4, 4)
+        """
+        return self.homogeneous
+
+    def set_rotation(
+        self,
+        rotation: list | np.ndarray,
+        rotation_type: str,
+        rotation_order: str | None = None,
+        degrees: bool = True,
+    ):
+        """
+        Set rotation from specified representation.
+
+        Args:
+            rotation: Rotation data
+            rotation_type: Type of rotation ("quat", "euler", "rotvec", "matrix", "rot6d")
+            rotation_order: Order/convention for the rotation type
+            degrees: For Euler angles, whether the input is in degrees (default True)
+        """
+        self._set_rotation(rotation, rotation_type, rotation_order, degrees)
+
+    def _compute_relative(self, other):  # type: ignore[override]
+        """
+        Compute relative transformation from other to self.
+
+        The result represents the transformation needed to go from other's frame to self's frame.
+        Mathematically: T_relative = T_other^{-1} * T_self
+
+        Args:
+            other: Reference end-effector pose
+
+        Returns:
+            EndEffectorPose representing the relative transformation
+        """
+        # Get homogeneous matrices
+        T_self = self.homogeneous
+        T_other = other.homogeneous
+
+        # Compute relative transformation: T_other^{-1} * T_self
+        T_relative = relative_transformation(T_other, T_self)
+
+        # Create new EndEffectorPose from relative transformation
+        return EndEffectorPose(homogeneous=T_relative)
+
+    @classmethod
+    def from_action_format(cls, data: np.ndarray, action_format: ActionFormat) -> "EndEffectorPose":
+        """
+        Create an EndEffectorPose from a flat array using the specified action format.
+
+        This is the inverse of the xyz_rot6d / xyz_rotvec / homogeneous properties.
+
+        Args:
+            data: Flat array whose layout depends on action_format.
+            action_format: One of ActionFormat.XYZ_ROT6D, XYZ_ROTVEC, or DEFAULT.
+
+        Returns:
+            EndEffectorPose instance.
+        """
+        if action_format == ActionFormat.XYZ_ROT6D:
+            return cls(translation=data[:3], rotation=data[3:], rotation_type="rot6d")
+        elif action_format == ActionFormat.XYZ_ROTVEC:
+            return cls(translation=data[:3], rotation=data[3:], rotation_type="rotvec")
+        elif action_format == ActionFormat.DEFAULT:
+            return cls(homogeneous=data.reshape(4, 4))
+        else:
+            raise ValueError(f"Unsupported ActionFormat: {action_format}")
+
+    def copy(self) -> "EndEffectorPose":
+        """
+        Create a deep copy of this end-effector pose.
+
+        Returns:
+            New EndEffectorPose instance with copied data
+        """
+        return EndEffectorPose(
+            translation=self._translation.copy(),
+            rotation=self._rotation.as_quat(),
+            rotation_type="quat",
+            rotation_order="xyzw",
+        )
+
+    def __repr__(self) -> str:
+        quat = self.to_rotation("quat", "wxyz")
+        return f"EndEffectorPose(translation={self.translation}, rotation_quat_wxyz={quat})"
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, EndEffectorPose):
+            return False
+        return np.allclose(self._translation, other._translation) and np.allclose(
+            self._rotation.as_quat(), other._rotation.as_quat()
+        )
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py
new file mode 100644
index 00000000000..14e2b0b3a37
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py
@@ -0,0 +1,643 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unified state and action processor for robotics."""
+
+from copy import deepcopy
+
+import numpy as np
+
+from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import (
+    ActionFormat,
+    ActionRepresentation,
+    ActionType,
+    ModalityConfig,
+)
+from vllm_omni.diffusion.models.gr00t.dataio.state_action.action_chunking import (
+    EndEffectorActionChunk,
+    JointActionChunk,
+)
+from vllm_omni.diffusion.models.gr00t.dataio.state_action.pose import EndEffectorPose, JointPose
+from vllm_omni.diffusion.models.gr00t.dataio.utils import (
+    apply_sin_cos_encoding,
+    nested_dict_to_numpy,
+    normalize_values_meanstd,
+    normalize_values_minmax,
+    parse_modality_configs,
+    unnormalize_values_meanstd,
+    unnormalize_values_minmax,
+)
+
+
+class StateActionProcessor:
+    """
+    Unified processor for robot state and action data.
+
+    Handles:
+    - State normalization (min/max, mean/std, sin/cos encoding)
+    - Action normalization
+    - Absolute <-> Relative action representation conversion
+    - Action processing with state dependency
+    """
+
+    def __init__(
+        self,
+        modality_configs: dict[str, dict[str, ModalityConfig]],
+        statistics: (dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None) = None,
+        use_percentiles: bool = False,
+        clip_outliers: bool = True,
+        apply_sincos_state_encoding: bool = False,
+        use_relative_action: bool = False,
+    ):
+        """
+        Initialize unified state and action processor.
+
+        Args:
+            modality_configs: Nested dict with structure:
+                {embodiment_tag: {modality: ModalityConfig}}
+                where modality in ["state", "action"]
+                Example: {"gr1": {"state": ModalityConfig(...), "action": ModalityConfig(...)}}
+            statistics: Optional nested dict with structure:
+                {embodiment_tag: {modality: {joint_group: {stat_type: values}}}}
+                where modality in ["state", "action", "relative_action"]
+                and stat_type in ["min", "max", "mean", "std", "q01", "q99"]
+                Example: {"gr1": {"state": {"left_arm": {"min": [...], "max": [...], ...}}}}
+            use_percentiles: Whether to use percentiles (q01/q99) instead of min/max
+            clip_outliers: Whether to clip normalized values to [-1, 1]
+            apply_sincos_state_encoding: Global flag to enable sin/cos encoding for states
+        """
+        self.modality_configs = parse_modality_configs(modality_configs)
+        self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {}
+        self.use_percentiles = use_percentiles
+        self.clip_outliers = clip_outliers
+        self.apply_sincos_state_encoding = apply_sincos_state_encoding
+        self.use_relative_action = use_relative_action
+
+        # Normalization parameters computed from statistics
+        self.norm_params: dict[str, dict[str, dict[str, dict[str, np.ndarray]]]] = {}
+        # Format: norm_params[embodiment_tag][modality][joint_group][stat_type]
+        # where stat_type in ["min", "max", "mean", "std", "dim"]
+
+        if statistics is not None:
+            self.set_statistics(statistics)
+
+    def set_statistics(
+        self,
+        statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]],
+        override: bool = False,
+    ) -> None:
+        """
+        Set dataset statistics for normalization.
+
+        Args:
+            statistics: Nested dict with structure:
+                {embodiment_tag: {modality: {joint_group: {stat_type: values}}}}
+        """
+        for key in statistics:
+            if key not in self.statistics or override:
+                self.statistics[key] = deepcopy(statistics[key])
+            else:
+                print(f"Embodiment tag {key} already in statistics, skipping updating")
+        self._compute_normalization_parameters()
+
+    def _compute_normalization_parameters(self) -> None:
+        """Compute and cache normalization parameters from statistics for all embodiments and modalities."""
+        for embodiment_tag in self.statistics:
+            self.norm_params[embodiment_tag] = {}
+
+            for modality in ["state", "action"]:
+                if modality not in self.statistics[embodiment_tag]:
+                    continue
+
+                self.norm_params[embodiment_tag][modality] = {}
+
+                for joint_group, stats in self.statistics[embodiment_tag][modality].items():
+                    if self.use_percentiles:
+                        min_vals = np.array(stats["q01"])
+                        max_vals = np.array(stats["q99"])
+                    else:
+                        min_vals = np.array(stats["min"])
+                        max_vals = np.array(stats["max"])
+
+                    mean_vals = np.array(stats["mean"])
+                    std_vals = np.array(stats["std"])
+
+                    # Compute range, ensuring it's not zero
+                    range_vals = max_vals - min_vals
+                    range_vals = np.maximum(range_vals, 1e-8)
+
+                    self.norm_params[embodiment_tag][modality][joint_group] = {
+                        "min": min_vals,
+                        "max": max_vals,
+                        "dim": np.array(range_vals.shape[0]),
+                        "mean": mean_vals,
+                        "std": std_vals,
+                    }
+
+            # Override absolute action stats with relative stats where specified
+            if "action" in self.modality_configs[embodiment_tag]:
+                modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys
+                action_configs = self.modality_configs[embodiment_tag]["action"].action_configs
+
+                if action_configs is not None:
+                    for key, action_config in zip(modality_keys, action_configs):
+                        if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action:
+                            if "relative_action" not in self.statistics[embodiment_tag]:
+                                raise ValueError(
+                                    f"Relative action statistics required for embodiment '{embodiment_tag}' "
+                                    f"but 'relative_action' not found in statistics"
+                                )
+                            if key not in self.statistics[embodiment_tag]["relative_action"]:
+                                raise ValueError(
+                                    f"Relative action statistics required for key '{key}' "
+                                    f"in embodiment '{embodiment_tag}' but not found"
+                                )
+                            action_dim = self.norm_params[embodiment_tag]["action"][key]["dim"]
+                            self.norm_params[embodiment_tag]["action"][key] = nested_dict_to_numpy(
+                                self.statistics[embodiment_tag]["relative_action"][key]
+                            )
+                            self.norm_params[embodiment_tag]["action"][key]["dim"] = action_dim
+
+    def apply_state(
+        self,
+        state: dict[str, np.ndarray],
+        embodiment_tag: str,
+    ) -> dict[str, np.ndarray]:
+        """
+        Apply state processing (normalization, encoding).
+
+        Args:
+            state: Dict mapping joint_group -> raw state values
+                Shape per group: (..., D) where D is state dimension
+            embodiment_tag: Embodiment identifier (e.g., "gr1")
+
+        Returns:
+            Dict mapping joint_group -> processed state values
+                - Sin/cos encoded groups: (..., 2*D)
+                - Other groups: (..., D)
+        """
+        normalized_values = {}
+        state = deepcopy(state)  # Avoid modifying input
+
+        # Get sin/cos embedding keys if enabled
+        sin_cos_keys = None
+        if self.apply_sincos_state_encoding:
+            state_config = self.modality_configs[embodiment_tag].get("state")
+            if state_config and hasattr(state_config, "sin_cos_embedding_keys"):
+                sin_cos_keys = state_config.sin_cos_embedding_keys
+
+        for joint_group in self.modality_configs[embodiment_tag]["state"].modality_keys:
+            if joint_group not in state:
+                raise KeyError(f"Joint group '{joint_group}' not found in state dict for embodiment '{embodiment_tag}'")
+
+            # Strategy 1: Sin/cos encoding (doubles dimension)
+            if sin_cos_keys and joint_group in sin_cos_keys:
+                normalized_values[joint_group] = apply_sin_cos_encoding(state[joint_group])
+
+            # Strategy 2: Mean/std normalization
+            elif (
+                hasattr(
+                    self.modality_configs[embodiment_tag]["state"],
+                    "mean_std_embedding_keys",
+                )
+                and self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys
+                and joint_group in self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys
+            ):
+                params = self.norm_params[embodiment_tag]["state"][joint_group]
+                normalized = normalize_values_meanstd(state[joint_group], params)
+                normalized_values[joint_group] = normalized
+
+            # Strategy 3: Min/max normalization to [-1, 1]
+            else:
+                params = self.norm_params[embodiment_tag]["state"][joint_group]
+                normalized = normalize_values_minmax(state[joint_group], params)
+
+                if self.clip_outliers:
+                    normalized = np.clip(normalized, -1.0, 1.0)
+
+                normalized_values[joint_group] = normalized
+
+        return normalized_values
+
+    def unapply_state(
+        self,
+        state: dict[str, np.ndarray],
+        embodiment_tag: str,
+    ) -> dict[str, np.ndarray]:
+        """
+        Reverse state processing (denormalization).
+
+        Args:
+            state: Dict mapping joint_group -> processed state values
+            embodiment_tag: Embodiment identifier
+
+        Returns:
+            Dict mapping joint_group -> raw state values
+
+        Raises:
+            ValueError: If attempting to reverse sin/cos encoding (not reversible)
+        """
+        unnormalized_values = {}
+
+        # Get sin/cos embedding keys if enabled
+        sin_cos_keys = None
+        if self.apply_sincos_state_encoding:
+            state_config = self.modality_configs[embodiment_tag].get("state")
+            if state_config and hasattr(state_config, "sin_cos_embedding_keys"):
+                sin_cos_keys = state_config.sin_cos_embedding_keys
+
+        for joint_group in self.modality_configs[embodiment_tag]["state"].modality_keys:
+            if joint_group not in state:
+                raise KeyError(f"Joint group '{joint_group}' not found in state dict for embodiment '{embodiment_tag}'")
+
+            # Sin/cos encoding is not reversible
+            if sin_cos_keys and joint_group in sin_cos_keys:
+                raise ValueError(
+                    f"Cannot unapply sin/cos encoding for joint group '{joint_group}' "
+                    f"in embodiment '{embodiment_tag}'. This transformation is not reversible."
+                )
+
+            # Reverse mean/std normalization
+            elif (
+                hasattr(
+                    self.modality_configs[embodiment_tag]["state"],
+                    "mean_std_embedding_keys",
+                )
+                and self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys
+                and joint_group in self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys
+            ):
+                params = self.norm_params[embodiment_tag]["state"][joint_group]
+                unnormalized = unnormalize_values_meanstd(state[joint_group], params)
+                unnormalized_values[joint_group] = unnormalized
+
+            # Reverse min/max normalization
+            else:
+                params = self.norm_params[embodiment_tag]["state"][joint_group]
+                unnormalized_values[joint_group] = unnormalize_values_minmax(state[joint_group], params)
+
+        return unnormalized_values
+
+    def apply_action(
+        self,
+        action: dict[str, np.ndarray],
+        embodiment_tag: str,
+        state: dict[str, np.ndarray] | None = None,
+    ) -> dict[str, np.ndarray]:
+        """
+        Apply action processing (absolute->relative conversion, normalization).
+
+        Processing order:
+        1. Convert absolute actions to relative (if configured)
+        2. Normalize actions
+
+        Args:
+            action: Dict mapping joint_group -> raw action values
+                Shape per group: (T, D) where T is action horizon, D is action dimension
+            embodiment_tag: Embodiment identifier
+            state: Optional dict mapping joint_group -> raw state values
+                Required if any action group uses ActionRepresentation.RELATIVE
+                Shape per group: (T_state, D) where last timestep is used as reference
+
+        Returns:
+            Dict mapping joint_group -> processed action values
+                Shape per group: (T, D)
+
+        Raises:
+            ValueError: If state is None but required for relative action conversion
+        """
+        action = deepcopy(action)  # Avoid modifying input
+
+        # Step 1: Convert absolute actions to relative (if needed)
+        modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys
+        action_configs = self.modality_configs[embodiment_tag]["action"].action_configs
+
+        if action_configs is not None:
+            for key, action_config in zip(modality_keys, action_configs):
+                if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action:
+                    if state is None:
+                        raise ValueError(
+                            f"State dict required for relative action processing of key '{key}' "
+                            f"in embodiment '{embodiment_tag}'"
+                        )
+
+                    # Determine which state key to use as reference
+                    state_key = action_config.state_key if action_config.state_key else key
+
+                    if state_key not in state:
+                        raise KeyError(
+                            f"Reference state key '{state_key}' not found in state dict "
+                            f"for embodiment '{embodiment_tag}'"
+                        )
+
+                    # Use last state as reference frame
+                    reference_state = state[state_key][-1]
+
+                    # Convert absolute to relative
+                    action[key] = self._convert_to_relative_action(
+                        action=action[key],
+                        reference_state=reference_state,
+                        action_type=action_config.type,
+                        action_format=action_config.format,
+                    )
+
+        # Step 2: Normalize actions
+        normalized_values = {}
+        for joint_group in modality_keys:
+            if joint_group not in action:
+                raise KeyError(
+                    f"Joint group '{joint_group}' not found in action dict for embodiment '{embodiment_tag}'"
+                )
+
+            params = self.norm_params[embodiment_tag]["action"][joint_group]
+            if (
+                self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys is not None
+                and joint_group in self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys
+            ):
+                normalized = normalize_values_meanstd(action[joint_group], params)
+            else:
+                normalized = normalize_values_minmax(action[joint_group], params)
+
+            if self.clip_outliers:
+                normalized = np.clip(normalized, -1.0, 1.0)
+
+            normalized_values[joint_group] = normalized
+
+        return normalized_values
+
+    def unapply_action(
+        self,
+        action: dict[str, np.ndarray],
+        embodiment_tag: str,
+        state: dict[str, np.ndarray] | None = None,
+    ) -> dict[str, np.ndarray]:
+        """
+        Reverse action processing (denormalization, relative->absolute conversion).
+
+        Processing order:
+        1. Denormalize actions
+        2. Convert relative actions to absolute (if configured)
+
+        Args:
+            action: Dict mapping joint_group -> processed action values
+                Shape per group: (T, D) or (B, T, D) for batched
+            embodiment_tag: Embodiment identifier
+            state: Optional dict mapping joint_group -> raw state values
+                Required if any action group uses ActionRepresentation.RELATIVE
+                Shape per group: (T_state, D) or (B, T_state, D) for batched
+
+        Returns:
+            Dict mapping joint_group -> raw absolute action values
+                Shape per group: (T, D) or (B, T, D) for batched
+
+        Raises:
+            ValueError: If state is None but required for relative->absolute conversion
+        """
+        # Step 1: Unnormalize actions
+        unnormalized_values = {}
+        modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys
+
+        for joint_group in modality_keys:
+            if joint_group not in action:
+                raise KeyError(
+                    f"Joint group '{joint_group}' not found in action dict for embodiment '{embodiment_tag}'"
+                )
+
+            params = self.norm_params[embodiment_tag]["action"][joint_group]
+            group_values = action[joint_group]
+
+            if (
+                self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys is not None
+                and joint_group in self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys
+            ):
+                unnormalized = unnormalize_values_meanstd(group_values, params)
+            else:
+                unnormalized = unnormalize_values_minmax(group_values, params)
+
+            unnormalized_values[joint_group] = unnormalized
+
+        # Step 2: Convert relative actions to absolute (if needed)
+        action_configs = self.modality_configs[embodiment_tag]["action"].action_configs
+
+        if action_configs is not None:
+            for key, action_config in zip(modality_keys, action_configs):
+                if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action:
+                    if state is None:
+                        raise ValueError(
+                            f"State dict required for relative->absolute conversion of key '{key}' "
+                            f"in embodiment '{embodiment_tag}'"
+                        )
+
+                    # Determine which state key to use as reference
+                    state_key = action_config.state_key if action_config.state_key else key
+
+                    if state_key not in state:
+                        raise KeyError(
+                            f"Reference state key '{state_key}' not found in state dict "
+                            f"for embodiment '{embodiment_tag}'"
+                        )
+
+                    relative_action = unnormalized_values[key]
+
+                    # Handle batched and unbatched cases
+                    is_batched = relative_action.ndim == 3
+                    if not is_batched:
+                        assert relative_action.ndim == 2
+                        reference_state = state[state_key]
+                        if reference_state.ndim == 2:
+                            reference_state = reference_state[None, :]
+                        relative_action = relative_action[None, :]
+                    else:
+                        reference_state = state[state_key]
+                        if reference_state.ndim == 2:
+                            reference_state = reference_state[None, :]
+
+                    # Convert batched relative actions to absolute
+                    absolute_actions = []
+                    for s, a in zip(reference_state, relative_action):
+                        # Use last timestep of state as reference
+                        absolute_action = self._convert_to_absolute_action(
+                            action=a,
+                            reference_state=s[-1],
+                            action_type=action_config.type,
+                            action_format=action_config.format,
+                        )
+                        absolute_actions.append(absolute_action)
+
+                    if is_batched:
+                        unnormalized_values[key] = np.stack(absolute_actions, axis=0)
+                    else:
+                        unnormalized_values[key] = absolute_actions[0]
+
+        return unnormalized_values
+
+    def apply(
+        self,
+        state: dict[str, np.ndarray],
+        action: dict[str, np.ndarray],
+        embodiment_tag: str,
+    ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
+        """
+        Apply both state and action processing together.
+
+        Convenience method that processes state and action in one call,
+        automatically passing raw state to action processor for relative conversion.
+
+        Args:
+            state: Dict mapping joint_group -> raw state values
+            action: Dict mapping joint_group -> raw action values
+            embodiment_tag: Embodiment identifier
+
+        Returns:
+            Tuple of (processed_state, processed_action)
+        """
+        processed_state = self.apply_state(state, embodiment_tag)
+        if action:
+            processed_action = self.apply_action(action, embodiment_tag, state=state)
+        else:
+            processed_action = {}
+        return processed_state, processed_action
+
+    def unapply(
+        self,
+        state: dict[str, np.ndarray],
+        action: dict[str, np.ndarray],
+        embodiment_tag: str,
+        raw_state: dict[str, np.ndarray] | None = None,
+    ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
+        """
+        Reverse both state and action processing together.
+
+        Args:
+            state: Dict mapping joint_group -> processed state values
+            action: Dict mapping joint_group -> processed action values
+            embodiment_tag: Embodiment identifier
+            raw_state: Optional dict of raw states for relative->absolute conversion
+                If None, will use unapplied state (but won't work for sin/cos encoded states)
+
+        Returns:
+            Tuple of (raw_state, raw_action)
+        """
+        # Unapply state first
+        try:
+            unapplied_state = self.unapply_state(state, embodiment_tag)
+        except ValueError as e:
+            if "sin/cos encoding" in str(e) and raw_state is None:
+                raise ValueError("Cannot unapply sin/cos encoded state. Please provide raw_state parameter.") from e
+            raise
+
+        # Use provided raw_state if available, otherwise use unapplied state
+        state_for_action = raw_state if raw_state is not None else unapplied_state
+
+        # Unapply action
+        unapplied_action = self.unapply_action(action, embodiment_tag, state=state_for_action)
+
+        return unapplied_state, unapplied_action
+
+    def get_state_dim(self, embodiment_tag: str, include_sincos_expansion: bool = False) -> int:
+        """
+        Get total state dimension after processing.
+
+        Args:
+            embodiment_tag: Embodiment identifier
+            include_sincos_expansion: If True, accounts for sin/cos encoding doubling dimensions
+
+        Returns:
+            Total state dimension across all joint groups
+        """
+        total_dim = 0
+        state_config = self.modality_configs[embodiment_tag]["state"]
+
+        # Get sin/cos embedding keys if enabled
+        sin_cos_keys = set()
+        if self.apply_sincos_state_encoding and hasattr(state_config, "sin_cos_embedding_keys"):
+            sin_cos_keys = set(state_config.sin_cos_embedding_keys)
+
+        for joint_group in state_config.modality_keys:
+            base_dim = self.norm_params[embodiment_tag]["state"][joint_group]["dim"].item()
+
+            # Sin/cos encoding doubles the dimension
+            if include_sincos_expansion and joint_group in sin_cos_keys:
+                total_dim += base_dim * 2
+            else:
+                total_dim += base_dim
+
+        return total_dim
+
+    def get_action_dim(self, embodiment_tag: str) -> int:
+        """
+        Get total action dimension.
+
+        Args:
+            embodiment_tag: Embodiment identifier
+
+        Returns:
+            Total action dimension across all joint groups
+        """
+        total_dim = 0
+        for joint_group in self.modality_configs[embodiment_tag]["action"].modality_keys:
+            total_dim += self.norm_params[embodiment_tag]["action"][joint_group]["dim"].item()
+        return total_dim
+
+    def _convert_to_relative_action(
+        self,
+        action: np.ndarray,
+        reference_state: np.ndarray,
+        action_type: ActionType,
+        action_format: ActionFormat,
+    ) -> np.ndarray:
+        """Convert absolute action to relative action using reference state."""
+        assert action.ndim == 2, f"Expected action shape (T, D), got {action.shape}"
+        assert reference_state.ndim == 1, f"Expected state shape (D,), got {reference_state.shape}"
+
+        if action_type == ActionType.EEF:
+            action_chunking = EndEffectorActionChunk.from_array(action, action_format)
+            reference_frame = EndEffectorPose.from_action_format(reference_state, action_format)
+
+        elif action_type == ActionType.NON_EEF:
+            action_chunking = JointActionChunk([JointPose(m) for m in action])
+            reference_frame = JointPose(reference_state)
+
+        else:
+            raise ValueError(f"Unknown ActionType: {action_type}")
+
+        relative_action_chunking = action_chunking.relative_chunking(reference_frame=reference_frame)
+        return relative_action_chunking.to(action_format)
+
+    def _convert_to_absolute_action(
+        self,
+        action: np.ndarray,
+        reference_state: np.ndarray,
+        action_type: ActionType,
+        action_format: ActionFormat,
+    ) -> np.ndarray:
+        """Convert relative action to absolute action using reference state."""
+        assert action.ndim == 2, f"Expected action shape (T, D), got {action.shape}"
+        assert reference_state.ndim == 1, f"Expected state shape (D,), got {reference_state.shape}"
+        assert reference_state.shape[0] == action.shape[1], (
+            f"State dim {reference_state.shape[0]} != action dim {action.shape[1]}"
+        )
+
+        if action_type == ActionType.EEF:
+            rel_action = EndEffectorActionChunk.from_array(action, action_format)
+            reference_frame = EndEffectorPose.from_action_format(reference_state, action_format)
+
+        elif action_type == ActionType.NON_EEF:
+            rel_action = JointActionChunk([JointPose(pose) for pose in action])
+            reference_frame = JointPose(reference_state)
+
+        else:
+            raise ValueError(f"Unknown ActionType: {action_type}")
+
+        abs_action = rel_action.to_absolute_chunking(reference_frame=reference_frame)
+        return abs_action.to(action_format)
+
+    def __str__(self) -> str:
+        return (
+            "StateActionProcessor("
+            f"modality_configs={self.modality_configs}, "
+            f"statistics={self.statistics}, "
+            f"use_percentiles={self.use_percentiles}, "
+            f"clip_outliers={self.clip_outliers}, "
+            f"apply_sincos_state_encoding={self.apply_sincos_state_encoding}, "
+            f"use_relative_action={self.use_relative_action})"
+        )
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/types.py b/vllm_omni/diffusion/models/gr00t/dataio/types.py
new file mode 100644
index 00000000000..d4948a1c234
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/dataio/types.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+import numpy as np
+
+from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag
+
+
+class MessageType(Enum):
+    START_OF_EPISODE = "start_of_episode"
+    END_OF_EPISODE = "end_of_episode"
+    EPISODE_STEP = "episode_step"
+    IMAGE = "image"
+    TEXT = "text"
+
+
+class ActionRepresentation(Enum):
+    RELATIVE = "relative"
+    DELTA = "delta"
+    ABSOLUTE = "absolute"
+
+
+class ActionType(Enum):
+    EEF = "eef"
+    NON_EEF = "non_eef"
+
+
+class ActionFormat(Enum):
+    DEFAULT = "default"
+    XYZ_ROT6D = "xyz+rot6d"
+    XYZ_ROTVEC = "xyz+rotvec"
+
+
+@dataclass
+class VLAStepData:
+    """
+    Represents a single step of VLA (Vision-Language-Action) data.
+
+    This is the core data structure returned by datasets, containing raw observation
+    and action data that will be processed by the SequenceVLAProcessor.
+    """
+
+    # Core data
+    images: dict[str, list[np.ndarray]]  # view_name -> list[np.ndarray] (for temporal stacking)
+    states: dict[str, np.ndarray]  # state_name -> np.ndarray (dim,) for single step or (horizon, dim) for trajectory
+    actions: dict[str, np.ndarray]  # action_name -> np.ndarray (horizon, dim) for action chunk
+    masks: dict[str, list[np.ndarray]] | None = None  # view_name -> list[np.ndarray] (H, W)
+    text: str | None = None  # Optional task description or instruction
+    embodiment: EmbodimentTag = EmbodimentTag.NEW_EMBODIMENT  # Optional embodiment tag for cross-embodiment training
+    is_demonstration: bool = (
+        False  # Whether the step is a demonstration. If True, no loss should be computed for this step.
+    )
+
+    # Flexible metadata that can be extended by users
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ActionConfig:
+    rep: ActionRepresentation
+    type: ActionType
+    format: ActionFormat
+    state_key: str | None = None
+
+
+@dataclass
+class ModalityConfig:
+    """Configuration for a modality defining how data should be sampled and loaded.
+
+    This class specifies which indices to sample relative to a base index and which
+    keys to load for a particular modality (e.g., video, state, action).
+    """
+
+    delta_indices: list[int]
+    """Delta indices to sample relative to the current index. The returned data will
+    correspond to the original data at a sampled base index + delta indices."""
+    modality_keys: list[str]
+    """The keys to load for the modality in the dataset."""
+    sin_cos_embedding_keys: list[str] | None = None
+    """Optional list of keys to apply sin/cos encoding. If None or empty, use
+    min/max normalization for all keys."""
+    mean_std_embedding_keys: list[str] | None = None
+    """Optional list of keys to apply mean/std normalization. If None or empty,
+    use min/max normalization for all keys."""
+    action_configs: list[ActionConfig] | None = None
+
+    def __post_init__(self):
+        """Validate fields and set default values."""
+        if self.delta_indices is None or not isinstance(self.delta_indices, list):
+            raise ValueError(f"delta_indices must be a non-None list, got {self.delta_indices!r}")
+        if self.modality_keys is None or not isinstance(self.modality_keys, list) or len(self.modality_keys) == 0:
+            raise ValueError(f"modality_keys must be a non-empty list, got {self.modality_keys!r}")
+        if self.action_configs is not None:
+            assert len(self.action_configs) == len(self.modality_keys), (
+                f"Number of action configs ({len(self.action_configs)}) must match "
+                f"number of modality keys ({len(self.modality_keys)})"
+            )
+            parsed_action_configs = []
+            for action_config in self.action_configs:
+                if isinstance(action_config, dict):
+                    action_config = ActionConfig(
+                        rep=ActionRepresentation[action_config["rep"]],
+                        type=ActionType[action_config["type"]],
+                        format=ActionFormat[action_config["format"]],
+                        state_key=action_config.get("state_key", None),
+                    )
+                parsed_action_configs.append(action_config)
+            self.action_configs = parsed_action_configs
diff --git a/vllm_omni/diffusion/models/gr00t/dataio/utils.py b/vllm_omni/diffusion/models/gr00t/dataio/utils.py
new file mode 100644
index 00000000000..d3846d44228
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/dataio/utils.py
@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import asdict, is_dataclass
+from enum import Enum
+from typing import Any
+
+import numpy as np
+
+from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig
+
+
+def apply_sin_cos_encoding(values: np.ndarray) -> np.ndarray:
+    """Apply sin/cos encoding to values.
+
+    Args:
+        values: Array of shape (..., D) containing values to encode
+
+    Returns:
+        Array of shape (..., 2*D) with [sin, cos] concatenated
+
+    Note: This DOUBLES the dimension. For example:
+        Input:  [v₁, v₂, v₃] with shape (..., 3)
+        Output: [sin(v₁), sin(v₂), sin(v₃), cos(v₁), cos(v₂), cos(v₃)] with shape (..., 6)
+    """
+    sin_values = np.sin(values)
+    cos_values = np.cos(values)
+    # Concatenate sin and cos: [sin(v1), sin(v2), ..., cos(v1), cos(v2), ...]
+    return np.concatenate([sin_values, cos_values], axis=-1)
+
+
+def nested_dict_to_numpy(data: Any) -> Any:
+    """Recursively convert bottom-level list-of-lists to NumPy arrays.
+
+    Args:
+        data: A nested dictionary whose leaves are lists (or lists of lists).
+
+    Returns:
+        The same dictionary structure with leaf lists converted to ``np.ndarray``.
+    """
+    if isinstance(data, dict):
+        return {key: nested_dict_to_numpy(value) for key, value in data.items()}
+    elif isinstance(data, list):
+        # Convert lists to numpy arrays
+        # NumPy will handle both 1D and 2D cases appropriately
+        return np.array(data)
+    else:
+        return data
+
+
+def normalize_values_minmax(values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray:
+    """Min-max normalize ``values`` to ``[-1, 1]`` using ``params['min']`` and ``params['max']``.
+
+    Accepts 2D ``(T, D)`` or 3D ``(B, T, D)`` arrays. Features with ``min == max``
+    are emitted as 0.
+    """
+    min_vals = params["min"]
+    max_vals = params["max"]
+    normalized = np.zeros_like(values)
+
+    mask = ~np.isclose(max_vals, min_vals)
+
+    normalized[..., mask] = (values[..., mask] - min_vals[..., mask]) / (max_vals[..., mask] - min_vals[..., mask])
+    normalized[..., mask] = 2 * normalized[..., mask] - 1
+
+    return normalized
+
+
+def unnormalize_values_minmax(normalized_values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray:
+    """Inverse of :func:`normalize_values_minmax`.
+
+    Input values are clipped to ``[-1, 1]`` before being mapped back to
+    ``[params['min'], params['max']]``.
+    """
+
+    min_vals = params["min"]
+    max_vals = params["max"]
+    range_vals = max_vals - min_vals
+
+    # Unnormalize from [-1, 1]
+    unnormalized = (np.clip(normalized_values, -1.0, 1.0) + 1.0) / 2.0 * range_vals + min_vals
+    return unnormalized
+
+
+def normalize_values_meanstd(values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray:
+    """Z-score normalize ``values`` using ``params['mean']`` and ``params['std']``.
+
+    Features whose ``std == 0`` are passed through unchanged.
+    """
+    mean_vals = params["mean"]
+    std_vals = params["std"]
+
+    # Create mask for non-zero standard deviations
+    mask = std_vals != 0
+
+    # Initialize normalized array
+    normalized = np.zeros_like(values)
+
+    # Normalize only features with non-zero std
+    normalized[..., mask] = (values[..., mask] - mean_vals[..., mask]) / std_vals[..., mask]
+
+    # Keep original values for zero-std features
+    normalized[..., ~mask] = values[..., ~mask]
+
+    return normalized
+
+
+def unnormalize_values_meanstd(normalized_values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray:
+    """Inverse of :func:`normalize_values_meanstd` (``x * std + mean``).
+
+    Features whose ``std == 0`` are passed through unchanged.
+    """
+    mean_vals = params["mean"]
+    std_vals = params["std"]
+
+    # Create mask for non-zero standard deviations
+    mask = std_vals != 0
+
+    # Initialize unnormalized array
+    unnormalized = np.zeros_like(normalized_values)
+
+    # Unnormalize only features with non-zero std
+    unnormalized[..., mask] = normalized_values[..., mask] * std_vals[..., mask] + mean_vals[..., mask]
+
+    # Keep normalized values for zero-std features
+    unnormalized[..., ~mask] = normalized_values[..., ~mask]
+
+    return unnormalized
+
+
+def to_json_serializable(obj: Any) -> Any:
+    """
+    Recursively convert dataclasses and numpy arrays to JSON-serializable format.
+
+    Args:
+        obj: Object to convert (can be dataclass, numpy array, dict, list, etc.)
+
+    Returns:
+        JSON-serializable representation of the object
+    """
+    if is_dataclass(obj) and not isinstance(obj, type):
+        # Convert dataclass to dict, then recursively process the dict
+        return to_json_serializable(asdict(obj))
+    elif isinstance(obj, np.ndarray):
+        # Convert numpy array to list
+        return obj.tolist()
+    elif isinstance(obj, np.integer):
+        # Convert numpy integers to Python int
+        return int(obj)
+    elif isinstance(obj, np.floating):
+        # Convert numpy floats to Python float
+        return float(obj)
+    elif isinstance(obj, np.bool_):
+        # Convert numpy bool to Python bool
+        return bool(obj)
+    elif isinstance(obj, dict):
+        # Recursively process dictionary values
+        return {key: to_json_serializable(value) for key, value in obj.items()}
+    elif isinstance(obj, (list, tuple)):
+        # Recursively process list/tuple elements
+        return [to_json_serializable(item) for item in obj]
+    elif isinstance(obj, set):
+        # Convert set to list
+        return [to_json_serializable(item) for item in obj]
+    elif isinstance(obj, (str, int, float, bool, type(None))):
+        # Already JSON-serializable
+        return obj
+    elif isinstance(obj, Enum):
+        return obj.name
+    else:
+        # For other types, try to convert to string as fallback
+        # You might want to handle specific types differently
+        return str(obj)
+
+
+def parse_modality_configs(
+    modality_configs: dict[str, dict[str, ModalityConfig]],
+) -> dict[str, dict[str, ModalityConfig]]:
+    parsed_modality_configs = {}
+    for embodiment_tag, modality_config in modality_configs.items():
+        parsed_modality_configs[embodiment_tag] = {}
+        for modality, config in modality_config.items():
+            if isinstance(config, dict):
+                parsed_modality_configs[embodiment_tag][modality] = ModalityConfig(**config)
+            else:
+                parsed_modality_configs[embodiment_tag][modality] = config
+    return parsed_modality_configs
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/__init__.py b/vllm_omni/diffusion/models/gr00t/modeling/__init__.py
new file mode 100644
index 00000000000..7aecb405ff9
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/modeling/__init__.py
@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7
+from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7Processor
+
+__all__ = ["Gr00tN1d7", "Gr00tN1d7Processor"]
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py
new file mode 100644
index 00000000000..3fcbce203f5
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py
@@ -0,0 +1,491 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+import os
+from typing import Any
+
+import torch
+from torch import nn
+from transformers import AutoConfig, AutoModel, PreTrainedModel
+from transformers.feature_extraction_utils import BatchFeature
+
+from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
+from vllm_omni.diffusion.models.gr00t.modeling.modules.dit import AlternateVLDiT, DiT, SelfAttentionTransformer
+from vllm_omni.diffusion.models.gr00t.modeling.modules.embodiment_conditioned_mlp import (
+    CategorySpecificMLP,
+    MultiEmbodimentActionEncoder,
+)
+from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7DataCollator
+
+logger = logging.getLogger(__name__)
+
+
+class Gr00tN1d7ActionHead(nn.Module):
+    """Action head component for flow matching diffusion policy."""
+
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: Gr00tN1d7Config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.input_embedding_dim = config.input_embedding_dim
+
+        if config.use_alternate_vl_dit:
+            self.model = AlternateVLDiT(
+                **config.diffusion_model_cfg,
+                cross_attention_dim=config.backbone_embedding_dim,
+                attend_text_every_n_blocks=config.attend_text_every_n_blocks,
+            )
+            logger.info("Using AlternateVLDiT for diffusion model")
+        else:
+            self.model = DiT(
+                **config.diffusion_model_cfg,
+                cross_attention_dim=config.backbone_embedding_dim,
+            )
+            logger.info("Using DiT for diffusion model")
+        self.action_dim = config.max_action_dim
+        self.action_horizon = config.action_horizon
+        self.num_inference_timesteps = config.num_inference_timesteps
+
+        self.state_encoder = CategorySpecificMLP(
+            num_categories=config.max_num_embodiments,
+            input_dim=config.max_state_dim * config.state_history_length,
+            hidden_dim=self.hidden_size,
+            output_dim=self.input_embedding_dim,
+        )
+        self.action_encoder = MultiEmbodimentActionEncoder(
+            action_dim=self.action_dim,
+            hidden_size=self.input_embedding_dim,
+            num_embodiments=config.max_num_embodiments,
+        )
+        self.action_decoder = CategorySpecificMLP(
+            num_categories=config.max_num_embodiments,
+            input_dim=self.hidden_size,
+            hidden_dim=self.hidden_size,
+            output_dim=self.action_dim,
+        )
+
+        self.vlln = nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity()
+
+        vl_self_attention_cfg = config.vl_self_attention_cfg
+        if vl_self_attention_cfg and vl_self_attention_cfg.get("num_layers", 0) > 0:
+            self.vl_self_attention = SelfAttentionTransformer(**vl_self_attention_cfg)
+        else:
+            self.vl_self_attention = nn.Identity()
+
+        if config.add_pos_embed:
+            self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim)
+            nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)
+
+        self.num_timestep_buckets = config.num_timestep_buckets
+
+    def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature:
+        backbone_features = backbone_output["backbone_features"]
+        backbone_features = self.vlln(backbone_features)
+        backbone_features = self.vl_self_attention(backbone_features)
+        backbone_output["backbone_features"] = backbone_features
+        return backbone_output
+
+    def _encode_features(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature:
+        """
+        Encode features for the action head.
+
+        Args:
+            backbone_output: Output from the backbone model containing:
+                - backbone_features: [B, seq_len, backbone_embedding_dim]
+                - backbone_attention_mask: [B, seq_len]
+            action_input: Input containing:
+                - state: [B, state_history_length, max_state_dim]
+                - embodiment_id: [B] (embodiment IDs)
+
+        Returns:
+            BatchFeature containing:
+                - backbone_features: [B, seq_len, backbone_embedding_dim]
+                - state_features: [B, 1, input_embedding_dim]
+        """
+        backbone_output = self.process_backbone_output(backbone_output)
+
+        vl_embeds = backbone_output.backbone_features
+        embodiment_id = action_input.embodiment_id
+
+        state = action_input.state
+        current_T = state.shape[1]
+        assert current_T == self.config.state_history_length, "current_T != state_history_length"
+        # [B, state_history_length, max_state_dim] -> [B, 1, state_history_length * max_state_dim]
+        state = state.view(state.shape[0], 1, -1)
+
+        state_features = self.state_encoder(state, embodiment_id)
+
+        return BatchFeature(data={"backbone_features": vl_embeds, "state_features": state_features})
+
+    @torch.no_grad()
+    def get_action_with_features(
+        self,
+        backbone_features: torch.Tensor,
+        state_features: torch.Tensor,
+        embodiment_id: torch.Tensor,
+        backbone_output: BatchFeature,
+        action_input: BatchFeature,
+        options: dict[str, Any] | None = None,
+    ) -> BatchFeature:
+        """
+        Generate actions using the flow matching diffusion process.
+
+        Args:
+            backbone_features: [B, seq_len, backbone_embedding_dim]
+            state_features: [B, state_horizon, input_embedding_dim]
+            embodiment_id: [B] (embodiment IDs)
+            backbone_output: Output from the backbone model
+        """
+        vl_embeds = backbone_features
+
+        batch_size = vl_embeds.shape[0]
+        device = vl_embeds.device
+        _seed_env = os.environ.get("GR00T_NOISE_SEED")
+        if _seed_env is not None:
+            _gen = torch.Generator(device=device).manual_seed(int(_seed_env))
+            actions = torch.randn(
+                size=(batch_size, self.config.action_horizon, self.action_dim),
+                dtype=vl_embeds.dtype,
+                device=device,
+                generator=_gen,
+            )
+        else:
+            actions = torch.randn(
+                size=(batch_size, self.config.action_horizon, self.action_dim),
+                dtype=vl_embeds.dtype,
+                device=device,
+            )
+
+        dt = 1.0 / self.num_inference_timesteps
+        vel_strength = torch.ones_like(actions)
+
+        if "action" in action_input:
+            # If action in input when doing get action, it means we want to use RTC.
+            # action_horizon is the action horizon of the input action.
+            # rtc_overlap_steps is the number of steps to overlap with the previous action chunks.
+            # rtc_frozen_steps is the policy inference latency expressed as frozen action steps.
+            # rtc_ramp_rate is the rate of the ramp of denoising the actions.
+            assert options is not None, "options is not None"
+            assert "action_horizon" in options, "action_horizon is not in options"
+            assert "rtc_overlap_steps" in options, "rtc_overlap_steps is not in options"
+            assert "rtc_frozen_steps" in options, "rtc_frozen_steps is not in options"
+            assert "rtc_ramp_rate" in options, "rtc_ramp_rate is not in options"
+
+            action_horizon_before_padding = options["action_horizon"]
+
+            # Use previous action instead of pure noise to do inpainting
+            actions[:, : options["rtc_overlap_steps"], :] = action_input["action"][
+                :,
+                action_horizon_before_padding - options["rtc_overlap_steps"] : action_horizon_before_padding,
+                :,
+            ]
+            vel_strength[:, : options["rtc_frozen_steps"], :] = 0.0
+            # NOTE: use an exponential ramp strength to set the remaining unfrozen rtc_steps
+            intermediate_steps = options["rtc_overlap_steps"] - options["rtc_frozen_steps"]
+            # Create exponential ramp from 0 to 1 over intermediate steps
+            t = torch.linspace(0.0, 1.0, intermediate_steps + 2, device=device)
+            ramp = 1 - torch.exp(-options["rtc_ramp_rate"] * t)
+            ramp = ramp / ramp[-1].clamp_min(1e-8)  # normalize to [0,1]
+            ramp = ramp[1:-1]  # we will only take the middle part of the ramp, ignore the 0.0 and 1.0
+            # Apply ramp to the intermediate steps [batch, intermediate_steps, action_dim]
+            vel_strength[
+                :,
+                options["rtc_frozen_steps"] : options["rtc_overlap_steps"],
+                :,
+            ] = ramp[None, :, None].to(device)
+
+        # Run denoising steps.
+        for t in range(self.num_inference_timesteps):
+            t_discretized = t * self.num_timestep_buckets // self.num_inference_timesteps
+
+            # Embed noised action trajectory.
+            timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device)
+            action_features = self.action_encoder(actions, timesteps_tensor, embodiment_id)
+            # Add position embedding.
+            if self.config.add_pos_embed:
+                pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
+                pos_embs = self.position_embedding(pos_ids).unsqueeze(0)
+                action_features = action_features + pos_embs
+
+            # Join vision, language, state and action embedding along sequence dimension.
+            sa_embs = torch.cat((state_features, action_features), dim=1)
+
+            # Run model forward.
+            if self.config.use_alternate_vl_dit:
+                model_output = self.model(
+                    hidden_states=sa_embs,
+                    encoder_hidden_states=vl_embeds,
+                    timestep=timesteps_tensor,
+                    image_mask=backbone_output.image_mask,
+                    backbone_attention_mask=backbone_output.backbone_attention_mask,
+                )
+            else:
+                model_output = self.model(
+                    hidden_states=sa_embs,
+                    encoder_hidden_states=vl_embeds,
+                    timestep=timesteps_tensor,
+                )
+            pred = self.action_decoder(model_output, embodiment_id)
+
+            pred_velocity = pred[:, -self.action_horizon :]
+
+            # Update actions using euler integration.
+            actions = actions + dt * pred_velocity * vel_strength
+
+        return BatchFeature(
+            data={
+                "action_pred": actions,
+                "backbone_features": vl_embeds,
+                "state_features": state_features,
+            }
+        )
+
+    @torch.no_grad()
+    def get_action(
+        self,
+        backbone_output: BatchFeature,
+        action_input: BatchFeature,
+        options: dict[str, Any] | None = None,
+    ) -> BatchFeature:
+        """
+        Generate actions using the flow matching diffusion process.
+
+        Args:
+            backbone_output: Output from the backbone model containing:
+                - backbone_features: [B, seq_len, backbone_embedding_dim]
+                - backbone_attention_mask: [B, seq_len]
+            action_input: Input containing:
+                - state: [B, state_dim]
+                - embodiment_id: [B] (embodiment IDs)
+
+        Returns:
+            BatchFeature containing:
+                - action_pred: [B, action_horizon, action_dim] predicted actions
+        """
+        features = self._encode_features(backbone_output, action_input)
+        return self.get_action_with_features(
+            backbone_features=features.backbone_features,
+            state_features=features.state_features,
+            embodiment_id=action_input.embodiment_id,
+            backbone_output=backbone_output,
+            action_input=action_input,
+            options=options,
+        )
+
+    @property
+    def device(self):
+        return next(iter(self.parameters())).device
+
+    @property
+    def dtype(self):
+        return next(iter(self.parameters())).dtype
+
+    def prepare_input(self, batch: dict) -> BatchFeature:
+        """Prepare input batch for the action head."""
+        return BatchFeature(data=batch)
+
+
+def _load_qwen3_vl_cls() -> type:
+    """Import Qwen3-VL lazily so the missing-dep failure is clear and deferred."""
+    try:
+        from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration
+    except ImportError as exc:
+        raise ImportError(
+            "GR00T-N1.7 requires transformers>=4.57.1 for transformers.models.qwen3_vl "
+            "(the repo's declared floor is 4.56.0). Please upgrade transformers."
+        ) from exc
+    return Qwen3VLForConditionalGeneration
+
+
+class _Qwen3VLBackbone(nn.Module):
+    """GR00T adapter around the shared Qwen3-VL implementation."""
+
+    def __init__(
+        self,
+        model_name: str,
+        select_layer: int,
+        backbone_embedding_dim: int,
+        load_bf16: bool,
+        transformers_loading_kwargs: dict[str, Any] | None = None,
+    ):
+        super().__init__()
+        backbone_config = AutoConfig.from_pretrained(
+            model_name, **(transformers_loading_kwargs or {"trust_remote_code": True})
+        )
+        backbone_config.text_config._attn_implementation = "sdpa"
+        backbone_config.vision_config._attn_implementation = "sdpa"
+        self.model = _load_qwen3_vl_cls()(backbone_config).eval()
+        if load_bf16:
+            self.model.to(dtype=torch.bfloat16)
+
+        target_layers = select_layer if select_layer >= 0 else len(self.model.model.language_model.layers)
+        while len(self.model.model.language_model.layers) > target_layers:
+            self.model.model.language_model.layers.pop(-1)
+
+        # Drop the unused lm_head: feature-extractor only, and .to(dtype) above unties it (~0.58 GiB).
+        self.model.lm_head = None
+
+    def set_frozen_modules_to_eval_mode(self) -> None:
+        self.eval()
+
+    def prepare_input(self, batch: dict) -> BatchFeature:
+        return BatchFeature(data=batch)
+
+    def forward(self, vl_input: BatchFeature) -> BatchFeature:
+        self.set_frozen_modules_to_eval_mode()
+        keys_to_use = [
+            "input_ids",
+            "attention_mask",
+            "pixel_values",
+            "image_grid_thw",
+            "mm_token_type_ids",
+        ]
+        vl_input = {key: vl_input[key] for key in keys_to_use if key in vl_input}
+        # GR00T was trained against the **pre-final-RMSNorm** output of the last
+        # decoder layer (this was `outer.hidden_states[-1]` under HF 4.57.x semantics,
+        # where the @check_model_inputs tie-logic skipped overwriting because
+        # Qwen3VLCausalLMOutputWithPast has no `last_hidden_state` field).
+        # HF >=5.x changed the output_hidden_states tuple plumbing so that
+        # `hidden_states[-1]` is post-norm, which is a different tensor and breaks
+        # the trained action head. Capture the pre-norm tensor via a hook on the
+        # final RMSNorm regardless of transformers version.
+        _captured: list[torch.Tensor] = []
+
+        def _pre_norm_hook(_module, args, _out):
+            _captured.append(args[0])
+
+        norm = self.model.model.language_model.norm
+        handle = norm.register_forward_hook(_pre_norm_hook)
+        try:
+            self.model.model(**vl_input, return_dict=True)
+        finally:
+            handle.remove()
+        if not _captured:
+            raise RuntimeError("Failed to capture pre-norm hidden states from Qwen3-VL backbone")
+        backbone_features = _captured[-1]
+        image_mask = vl_input["input_ids"] == self.model.config.image_token_id
+        attention_mask = vl_input["attention_mask"] == 1
+        return BatchFeature(
+            data={
+                "backbone_features": backbone_features,
+                "backbone_attention_mask": attention_mask,
+                "image_mask": image_mask,
+            }
+        )
+
+
+def get_backbone_cls(config: Gr00tN1d7Config):
+    if "nvidia/Cosmos-Reason2" in config.model_name:
+        return _Qwen3VLBackbone
+    raise ValueError(f"Unsupported model name: {config.model_name}")
+
+
+class Gr00tN1d7(PreTrainedModel):
+    """Gr00tN1d7: VLA model with Cosmos-Reason2-2B (Qwen3-VL) backbone."""
+
+    config_class = Gr00tN1d7Config
+    supports_gradient_checkpointing = True
+    _tp_plan = {}
+
+    @property
+    def all_tied_weights_keys(self) -> dict[str, Any]:
+        return {}
+
+    def __init__(
+        self,
+        config: Gr00tN1d7Config,
+        transformers_loading_kwargs: dict | None = None,
+    ):
+        super().__init__(config)
+        if transformers_loading_kwargs is None:
+            transformers_loading_kwargs = {"trust_remote_code": True}
+        self.config = config
+
+        backbone_cls = get_backbone_cls(config)
+        self.backbone = backbone_cls(
+            model_name=config.model_name,
+            select_layer=config.select_layer,
+            backbone_embedding_dim=config.backbone_embedding_dim,
+            load_bf16=config.load_bf16,
+            transformers_loading_kwargs=transformers_loading_kwargs,
+        )
+
+        # Initialize action head
+        self.action_head = Gr00tN1d7ActionHead(config)
+
+        self.collator = Gr00tN1d7DataCollator(
+            model_name=config.model_name,
+            model_type=config.backbone_model_type,
+            transformers_loading_kwargs=transformers_loading_kwargs,
+        )
+
+    def prepare_input(self, inputs: dict) -> tuple[BatchFeature, BatchFeature]:
+        """Prepare inputs for backbone and action head."""
+
+        # NOTE -- currently the eval code doesn't use collator, so we need to add it here
+        if "vlm_content" in inputs:
+            # Fix for n_envs > 1: Process all environments' VLM content, not just the first
+            vlm_content_list = inputs["vlm_content"]
+            # Ensure vlm_content_list is always a list for consistent processing
+            if not isinstance(vlm_content_list, list):
+                vlm_content_list = [vlm_content_list]
+
+            # Process all VLM contents through the collator
+            prep = self.collator([{"vlm_content": vlm} for vlm in vlm_content_list])["inputs"]
+            inputs.pop("vlm_content")
+            inputs.update(prep)
+
+        backbone_inputs = self.backbone.prepare_input(inputs)
+        action_inputs = self.action_head.prepare_input(inputs)
+
+        backbone_inputs = backbone_inputs.to(device=self.device, dtype=self.dtype)
+        action_inputs = action_inputs.to(device=self.device, dtype=self.dtype)
+
+        return backbone_inputs, action_inputs
+
+    def forward(self, inputs: dict) -> BatchFeature:
+        """
+        Forward pass through the complete model.
+
+        Args:
+            inputs: Dictionary containing:
+                - Action inputs (state, action, embodiment_id, etc.)
+
+        Returns:
+            BatchFeature containing loss and other outputs
+        """
+        # Prepare inputs for backbone and action head
+        backbone_inputs, action_inputs = self.prepare_input(inputs)
+        backbone_outputs = self.backbone(backbone_inputs)
+        action_outputs = self.action_head(backbone_outputs, action_inputs)
+
+        return action_outputs
+
+    def get_action(self, inputs: dict, options: dict[str, Any] | None = None) -> BatchFeature:
+        """
+        Generate actions using the complete model.
+        """
+        # Prepare inputs for backbone and action head
+        backbone_inputs, action_inputs = self.prepare_input(inputs)
+
+        # Forward through backbone
+        backbone_outputs = self.backbone(backbone_inputs)
+        action_outputs = self.action_head.get_action(backbone_outputs, action_inputs, options)
+
+        return action_outputs
+
+    @property
+    def device(self):
+        return next(iter(self.parameters())).device
+
+    @property
+    def dtype(self):
+        return next(iter(self.parameters())).dtype
+
+
+# Register the model with HuggingFace
+AutoConfig.register("Gr00tN1d7", Gr00tN1d7Config)
+AutoModel.register(Gr00tN1d7Config, Gr00tN1d7)
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py
new file mode 100644
index 00000000000..208f01a7cb5
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py
new file mode 100644
index 00000000000..54fbe770d0d
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py
@@ -0,0 +1,344 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+from diffusers import ConfigMixin, ModelMixin
+from diffusers.configuration_utils import register_to_config
+from diffusers.models.attention import Attention, FeedForward
+from diffusers.models.embeddings import SinusoidalPositionalEmbedding, TimestepEmbedding, Timesteps
+from torch import nn
+
+
+class AdaLayerNorm(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        norm_elementwise_affine: bool = False,
+        norm_eps: float = 1e-5,
+        chunk_dim: int = 0,
+    ):
+        super().__init__()
+        self.chunk_dim = chunk_dim
+        output_dim = embedding_dim * 2
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, output_dim)
+        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        temb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        temb = self.linear(self.silu(temb))
+        scale, shift = temb.chunk(2, dim=1)
+        x = self.norm(x) * (1 + scale[:, None]) + shift[:, None]
+        return x
+
+
+class BasicTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: int | None = None,
+        activation_fn: str = "geglu",
+        attention_bias: bool = False,
+        upcast_attention: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_type: str = "layer_norm",
+        norm_eps: float = 1e-5,
+        final_dropout: bool = False,
+        attention_type: str = "default",
+        positional_embeddings: str | None = None,
+        num_positional_embeddings: int | None = None,
+        ff_inner_dim: int | None = None,
+        ff_bias: bool = True,
+        attention_out_bias: bool = True,
+    ):
+        super().__init__()
+        self.norm_type = norm_type
+
+        if positional_embeddings and (num_positional_embeddings is None):
+            raise ValueError(
+                "If `positional_embedding` type is defined, `num_position_embeddings` must also be defined."
+            )
+
+        if positional_embeddings == "sinusoidal":
+            self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)
+        else:
+            self.pos_embed = None
+
+        if norm_type == "ada_norm":
+            self.norm1 = AdaLayerNorm(dim)
+        else:
+            self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
+
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim,
+            upcast_attention=upcast_attention,
+            out_bias=attention_out_bias,
+        )
+
+        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+            inner_dim=ff_inner_dim,
+            bias=ff_bias,
+        )
+        self.final_dropout = nn.Dropout(dropout) if final_dropout else None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        temb: torch.LongTensor | None = None,
+    ) -> torch.Tensor:
+        if self.norm_type == "ada_norm":
+            norm_hidden_states = self.norm1(hidden_states, temb)
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+
+        if self.pos_embed is not None:
+            norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=(encoder_attention_mask if encoder_hidden_states is not None else attention_mask),
+        )
+        if self.final_dropout:
+            attn_output = self.final_dropout(attn_output)
+
+        hidden_states = attn_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        norm_hidden_states = self.norm3(hidden_states)
+        ff_output = self.ff(norm_hidden_states)
+
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+        return hidden_states
+
+
+class TimestepEncoder(nn.Module):
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        dtype = next(self.parameters()).dtype
+        return self.timestep_embedder(self.time_proj(timesteps).to(dtype))
+
+
+class DiT(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 8,
+        attention_head_dim: int = 64,
+        output_dim: int = 26,
+        num_layers: int = 12,
+        dropout: float = 0.1,
+        attention_bias: bool = True,
+        activation_fn: str = "gelu-approximate",
+        num_embeds_ada_norm: int | None = 1000,
+        upcast_attention: bool = False,
+        norm_type: str = "ada_norm",
+        norm_elementwise_affine: bool = False,
+        norm_eps: float = 1e-5,
+        max_num_positional_embeddings: int = 512,
+        compute_dtype=torch.float32,
+        final_dropout: bool = True,
+        positional_embeddings: str | None = "sinusoidal",
+        interleave_self_attention=False,
+        cross_attention_dim: int | None = None,
+    ):
+        super().__init__()
+
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.gradient_checkpointing = False
+        self.timestep_encoder = TimestepEncoder(embedding_dim=self.inner_dim)
+
+        all_blocks = []
+        for idx in range(num_layers):
+            use_self_attn = idx % 2 == 1 and interleave_self_attention
+            curr_cross_attention_dim = cross_attention_dim if not use_self_attn else None
+            all_blocks.append(
+                BasicTransformerBlock(
+                    self.inner_dim,
+                    num_attention_heads,
+                    attention_head_dim,
+                    dropout=dropout,
+                    activation_fn=activation_fn,
+                    attention_bias=attention_bias,
+                    upcast_attention=upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                    norm_eps=norm_eps,
+                    positional_embeddings=positional_embeddings,
+                    num_positional_embeddings=max_num_positional_embeddings,
+                    final_dropout=final_dropout,
+                    cross_attention_dim=curr_cross_attention_dim,
+                )
+            )
+        self.transformer_blocks = nn.ModuleList(all_blocks)
+        self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
+        self.proj_out_2 = nn.Linear(self.inner_dim, output_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.LongTensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        return_all_hidden_states: bool = False,
+    ):
+        temb = self.timestep_encoder(timestep)
+        hidden_states = hidden_states.contiguous()
+        encoder_hidden_states = encoder_hidden_states.contiguous()
+        all_hidden_states = [hidden_states]
+
+        for idx, block in enumerate(self.transformer_blocks):
+            if idx % 2 == 1 and self.config.interleave_self_attention:
+                hidden_states = block(hidden_states, temb=temb)
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    temb=temb,
+                )
+            all_hidden_states.append(hidden_states)
+
+        shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=1)
+        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+        if return_all_hidden_states:
+            return self.proj_out_2(hidden_states), all_hidden_states
+        return self.proj_out_2(hidden_states)
+
+
+class AlternateVLDiT(DiT):
+    def __init__(self, *args, attend_text_every_n_blocks: int = 2, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.attend_text_every_n_blocks = attend_text_every_n_blocks
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.LongTensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        return_all_hidden_states: bool = False,
+        image_mask: torch.Tensor | None = None,
+        backbone_attention_mask: torch.Tensor | None = None,
+    ):
+        assert image_mask is not None, "Image mask is required"
+        assert self.config.interleave_self_attention, "Interleave self attention must be enabled"
+
+        temb = self.timestep_encoder(timestep)
+        hidden_states = hidden_states.contiguous()
+        encoder_hidden_states = encoder_hidden_states.contiguous()
+
+        image_attention_mask = image_mask & backbone_attention_mask
+        non_image_attention_mask = (~image_mask) & backbone_attention_mask
+
+        all_hidden_states = [hidden_states]
+        for idx, block in enumerate(self.transformer_blocks):
+            if idx % 2 == 1:
+                hidden_states = block(hidden_states, temb=temb)
+            else:
+                curr_encoder_attention_mask = (
+                    non_image_attention_mask
+                    if idx % (2 * self.attend_text_every_n_blocks) == 0
+                    else image_attention_mask
+                )
+                hidden_states = block(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=curr_encoder_attention_mask,
+                    temb=temb,
+                )
+            all_hidden_states.append(hidden_states)
+
+        shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=1)
+        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+        if return_all_hidden_states:
+            return self.proj_out_2(hidden_states), all_hidden_states
+        return self.proj_out_2(hidden_states)
+
+
+class SelfAttentionTransformer(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 8,
+        attention_head_dim: int = 64,
+        output_dim: int = 26,
+        num_layers: int = 12,
+        dropout: float = 0.1,
+        attention_bias: bool = True,
+        activation_fn: str = "gelu-approximate",
+        num_embeds_ada_norm: int | None = 1000,
+        upcast_attention: bool = False,
+        max_num_positional_embeddings: int = 512,
+        compute_dtype=torch.float32,
+        final_dropout: bool = True,
+        positional_embeddings: str | None = "sinusoidal",
+        interleave_self_attention=False,
+    ):
+        super().__init__()
+
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.gradient_checkpointing = False
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    self.inner_dim,
+                    num_attention_heads,
+                    attention_head_dim,
+                    dropout=dropout,
+                    activation_fn=activation_fn,
+                    attention_bias=attention_bias,
+                    upcast_attention=upcast_attention,
+                    positional_embeddings=positional_embeddings,
+                    num_positional_embeddings=max_num_positional_embeddings,
+                    final_dropout=final_dropout,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        return_all_hidden_states: bool = False,
+    ):
+        hidden_states = hidden_states.contiguous()
+        all_hidden_states = [hidden_states]
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states)
+            all_hidden_states.append(hidden_states)
+        if return_all_hidden_states:
+            return hidden_states, all_hidden_states
+        return hidden_states
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py
new file mode 100644
index 00000000000..81fcb4726f6
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py
@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm_omni.diffusion.utils.flow_matching import SinusoidalPositionalEncoding, swish
+
+
+class CategorySpecificLinear(nn.Module):
+    """Linear layer with category-specific weights and biases for multi-embodiment support."""
+
+    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.num_categories = num_categories
+        # For each category, we have separate weights and biases.
+        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
+        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))
+
+    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        """Apply per-category linear projection.
+
+        Args:
+            x: ``[B, T, input_dim]`` input tensor.
+            cat_ids: ``[B]`` category / embodiment IDs.
+
+        Returns:
+            ``[B, T, hidden_dim]`` output tensor.
+        """
+        selected_W = self.W[cat_ids]
+        selected_b = self.b[cat_ids]
+        return torch.bmm(x, selected_W) + selected_b.unsqueeze(1)
+
+    def expand_action_dimension(
+        self,
+        old_action_dim: int,
+        new_action_dim: int,
+        expand_input: bool = False,
+        expand_output: bool = False,
+    ) -> None:
+        """Safely expand action dimension with explicit targeting.
+
+        Args:
+            old_action_dim: Original action dimension.
+            new_action_dim: New (larger) action dimension.
+            expand_input: Whether to expand input dimension (dim=1).
+            expand_output: Whether to expand output dimension (dim=2).
+        """
+        if new_action_dim <= old_action_dim:
+            raise ValueError(f"New action dim {new_action_dim} must be larger than old action dim {old_action_dim}")
+
+        # Expand input dimension (dim=1) only if explicitly requested AND dimensions match
+        if expand_input and self.W.shape[1] == old_action_dim:
+            repeat_times = new_action_dim // old_action_dim
+            remainder = new_action_dim % old_action_dim
+
+            new_W_parts = [self.W] * repeat_times
+            if remainder > 0:
+                new_W_parts.append(self.W[:, :remainder, :])
+
+            new_W = torch.cat(new_W_parts, dim=1)
+            self.W = nn.Parameter(new_W)
+
+        # Expand output dimension (dim=2) only if explicitly requested AND dimensions match
+        if expand_output and self.W.shape[2] == old_action_dim:
+            repeat_times = new_action_dim // old_action_dim
+            remainder = new_action_dim % old_action_dim
+
+            new_W_parts = [self.W] * repeat_times
+            if remainder > 0:
+                new_W_parts.append(self.W[:, :, :remainder])
+
+            new_W = torch.cat(new_W_parts, dim=2)
+            self.W = nn.Parameter(new_W)
+
+            # Expand bias for output dimension
+            if self.b.shape[1] == old_action_dim:
+                new_b_parts = [self.b] * repeat_times
+                if remainder > 0:
+                    new_b_parts.append(self.b[:, :remainder])
+
+                new_b = torch.cat(new_b_parts, dim=1)
+                self.b = nn.Parameter(new_b)
+
+
+class CategorySpecificMLP(nn.Module):
+    """Two-layer MLP with category-specific weights for multi-embodiment support."""
+
+    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int, output_dim: int) -> None:
+        super().__init__()
+        self.num_categories = num_categories
+        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
+        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)
+
+    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        """Apply two-layer per-category MLP.
+
+        Args:
+            x: ``[B, T, input_dim]`` input tensor.
+            cat_ids: ``[B]`` category / embodiment IDs.
+
+        Returns:
+            ``[B, T, output_dim]`` output tensor.
+        """
+        hidden = F.relu(self.layer1(x, cat_ids))
+        return self.layer2(hidden, cat_ids)
+
+    def expand_action_dimension(self, old_action_dim: int, new_action_dim: int) -> None:
+        """Expand action dimension by copying weights from existing dimensions.
+
+        Args:
+            old_action_dim: Original action dimension.
+            new_action_dim: New (larger) action dimension.
+        """
+        # self.layer1 does not take action_dim as input, so no expansion needed
+        self.layer2.expand_action_dimension(old_action_dim, new_action_dim, expand_input=False, expand_output=True)
+
+
+class MultiEmbodimentActionEncoder(nn.Module):
+    """Action encoder with multi-embodiment support and sinusoidal positional encoding."""
+
+    def __init__(self, action_dim: int, hidden_size: int, num_embodiments: int) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_embodiments = num_embodiments
+
+        # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w}
+        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)  # (d -> w)
+        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)  # (2w -> w)
+        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)  # (w -> w)
+        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)
+
+    def forward(self, actions: torch.Tensor, timesteps: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        """Encode actions with embodiment-conditioned MLP and time embedding.
+
+        Args:
+            actions: ``[B, T, action_dim]`` action tensor.
+            timesteps: ``[B]`` timesteps (one scalar per batch item).
+            cat_ids: ``[B]`` category / embodiment IDs.
+
+        Returns:
+            ``[B, T, hidden_size]`` encoded action features.
+        """
+        B, T, _ = actions.shape
+
+        # 1) Expand each batch's single scalar time 'tau' across all T steps
+        #    so that shape => (B, T)
+        #    e.g. if timesteps is (B,), replicate across T
+        if timesteps.dim() == 1 and timesteps.shape[0] == B:
+            # shape (B,) => (B,T)
+            timesteps = timesteps.unsqueeze(1).expand(-1, T)
+        else:
+            raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.")
+
+        # 2) Standard action MLP step for shape => (B, T, w)
+        a_emb = self.W1(actions, cat_ids)
+
+        # 3) Get the sinusoidal encoding (B, T, w)
+        tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype)
+
+        # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish
+        x = torch.cat([a_emb, tau_emb], dim=-1)
+        x = swish(self.W2(x, cat_ids))
+
+        # 5) Finally W3 => (B, T, w)
+        x = self.W3(x, cat_ids)
+        return x
+
+    def expand_action_dimension(self, old_action_dim: int, new_action_dim: int) -> None:
+        """Expand action dimension by copying weights from existing dimensions.
+
+        Args:
+            old_action_dim: Original action dimension.
+            new_action_dim: New (larger) action dimension.
+        """
+        # Only W1 takes action_dim as input, so only expand its input dimension
+        self.W1.expand_action_dimension(old_action_dim, new_action_dim, expand_input=True, expand_output=False)
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
new file mode 100644
index 00000000000..ba128f4505b
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
@@ -0,0 +1,555 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import os
+import re
+import warnings
+from copy import deepcopy
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torchvision.transforms.v2 as transforms
+from PIL import Image
+from transformers import AutoProcessor, ProcessorMixin
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils import cached_file
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig
+from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag
+from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor
+from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable
+
+
+class LetterBoxTransform:
+    """Pad image to square dimensions by adding black bars to the smaller side."""
+
+    def __call__(self, img: torch.Tensor) -> torch.Tensor:
+        import math
+
+        *leading_dims, c, h, w = img.shape
+        if h == w:
+            return img
+        max_dim = max(h, w)
+        pad_h, pad_w = max_dim - h, max_dim - w
+        pad_top, pad_left = pad_h // 2, pad_w // 2
+        pad_bottom, pad_right = pad_h - pad_top, pad_w - pad_left
+        if leading_dims:
+            batch_size = math.prod(leading_dims)
+            img_r = img.reshape(batch_size, c, h, w)
+            padded = transforms.functional.pad(img_r, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0)
+            return padded.reshape(leading_dims + [c, max_dim, max_dim])
+        return transforms.functional.pad(img, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0)
+
+
+def _build_eval_image_transform(
+    image_target_size: list[int],
+    image_crop_size: list[int],
+) -> transforms.Compose:
+    """Deterministic eval/inference image transform (letterbox → resize → centercrop → resize)."""
+    return transforms.Compose(
+        [
+            transforms.ToImage(),
+            LetterBoxTransform(),
+            transforms.Resize(size=image_target_size),
+            transforms.CenterCrop(size=image_crop_size),
+            transforms.Resize(size=image_target_size),
+        ]
+    )
+
+
+logger = init_logger(__name__)
+
+# Suppress protobuf deprecation warnings
+warnings.filterwarnings("ignore", category=DeprecationWarning, module="google.protobuf")
+
+EMBODIMENT_TAG_TO_PROJECTOR_INDEX = {
+    # Pretrain embodiment ids
+    "oxe_droid_relative_eef_relative_joint": 24,
+    "xdof_relative_eef_relative_joint": 27,
+    "xdof_relative_eef_relative_joint_subtask": 27,
+    "real_g1_relative_eef_relative_joints": 25,
+    "real_r1_pro_sharpa_relative_eef": 26,
+    "real_r1_pro_sharpa_relative_eef_human": 26,
+    "real_r1_pro_sharpa_relative_eef_maxinsights": 26,
+    "real_r1_pro_sharpa_relative_eef_mecka": 26,
+    # Posttrain embodiment ids
+    "unitree_g1_full_body_with_waist_height_nav_cmd": 25,
+    "unitree_g1_sonic": 11,
+    "simpler_env_google": 0,
+    "simpler_env_widowx": 1,
+    "libero_sim": 2,
+    "new_embodiment": 10,
+}
+
+QWEN3_VL_2B_PROCESSOR = "Qwen/Qwen3-VL-2B-Instruct"
+
+
+def build_processor(model_name: str, transformers_loading_kwargs: dict) -> ProcessorMixin:
+    try:
+        from transformers import Qwen3VLProcessor
+    except ImportError as exc:
+        raise ImportError(
+            "GR00T-N1.7 requires transformers>=4.57.1 for Qwen3VLProcessor "
+            "(the repo's declared floor is 4.56.0). Please upgrade transformers."
+        ) from exc
+    if model_name == "nvidia/Cosmos-Reason2-2B":
+        # Cosmos-Reason2-2B lacks a Qwen3VLProcessor; fall back to upstream artifacts.
+        logger.warning_once(
+            "Substituting processor from %s because %s does not ship one. "
+            "If you fine-tune Cosmos-Reason2-2B's tokenizer/image processor, "
+            "load the processor explicitly instead of relying on this fallback.",
+            QWEN3_VL_2B_PROCESSOR,
+            model_name,
+        )
+        model_name = QWEN3_VL_2B_PROCESSOR
+    return Qwen3VLProcessor.from_pretrained(model_name, **transformers_loading_kwargs)
+
+
+class Gr00tN1d7DataCollator:
+    def __init__(
+        self,
+        model_name: str,
+        model_type: str = "qwen",
+        transformers_loading_kwargs: dict | None = None,
+    ):
+        if transformers_loading_kwargs is None:
+            transformers_loading_kwargs = {}
+        self.processor = build_processor(model_name, transformers_loading_kwargs)
+        self.processor.tokenizer.padding_side = "left"
+        self.model_type = model_type
+        self.model_name = model_name
+
+    def __call__(self, features: list[dict[str, Any]]) -> BatchFeature:
+        batch = {}
+        keys = list(set().union(*(elem.keys() for elem in features)))
+
+        for key in keys:
+            values = [elem[key] for elem in features if key in elem]
+            if key == "vlm_content":
+                text_list = []
+                image_inputs = []
+                for v in values:
+                    curr_text_list = [v["text"]]
+
+                    text_list += curr_text_list
+                    curr_image_inputs = v["images"]
+                    image_inputs += curr_image_inputs
+
+                vlm_inputs = self.processor(
+                    text=text_list,
+                    images=image_inputs,
+                    return_tensors="pt",
+                    padding=True,
+                )
+                for k, v in vlm_inputs.items():
+                    batch[k] = v
+            elif key in (
+                "pixel_values",
+                "image_grid_thw",
+                "attention_mask",
+                "input_ids",
+            ):
+                raise Exception("Not implemented")
+            else:
+                batch[key] = torch.from_numpy(np.stack(values))
+        return BatchFeature(data={"inputs": batch})
+
+    def __str__(self):
+        return f"Gr00tN1d7DataCollator(model_name={self.model_name}, model_type={self.model_type})"
+
+
+class Gr00tN1d7Processor(ProcessorMixin):
+    data_collator_class = Gr00tN1d7DataCollator
+
+    def __init__(
+        self,
+        modality_configs: dict[str, dict[str, ModalityConfig]],
+        statistics: (dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None) = None,
+        use_percentiles: bool = False,
+        clip_outliers: bool = True,
+        image_crop_size: list[int] = None,
+        image_target_size: list[int] = None,
+        shortest_image_edge: int = 256,
+        crop_fraction: float = 0.95,
+        random_rotation_angle: int | None = None,
+        color_jitter_params: dict[str, float] | None = None,
+        formalize_language: bool = True,
+        model_name: str = "nvidia/Cosmos-Reason2-2B",
+        model_type: str = "qwen",
+        max_state_dim: int = 29,
+        max_action_dim: int = 29,
+        max_action_horizon: int = 50,
+        apply_sincos_state_encoding: bool = False,
+        use_relative_action: bool = False,
+        embodiment_id_mapping: dict[str, int] | None = None,
+        transformers_loading_kwargs: dict | None = None,
+        exclude_state: bool = False,
+        # Normalization
+        use_mean_std: bool = False,
+        **kwargs,  # absorb deprecated training-only keys from saved processor_config.json
+    ):
+        if transformers_loading_kwargs is None:
+            transformers_loading_kwargs = {"trust_remote_code": True}
+        if kwargs:
+            logger.debug("Gr00tN1d7Processor: ignoring unknown keys: %s", list(kwargs))
+        self.modality_configs = parse_modality_configs(modality_configs)
+
+        # Initialize StateActionProcessor for state/action normalization
+        self.state_action_processor = StateActionProcessor(
+            modality_configs=modality_configs,
+            statistics=statistics,
+            use_percentiles=use_percentiles,
+            clip_outliers=clip_outliers,
+            apply_sincos_state_encoding=apply_sincos_state_encoding,
+            use_relative_action=use_relative_action,
+        )
+
+        self.use_percentiles = use_percentiles
+        self.use_mean_std = use_mean_std
+        self.clip_outliers = clip_outliers
+        self.apply_sincos_state_encoding = apply_sincos_state_encoding
+        self.use_relative_action = use_relative_action
+
+        self.exclude_state = exclude_state
+
+        self.formalize_language = formalize_language
+        self.model_name = model_name
+        self.model_type = model_type
+
+        self.max_state_dim = max_state_dim
+        self.max_action_dim = max_action_dim
+        self.max_action_horizon = max_action_horizon
+
+        self.image_crop_size = image_crop_size
+        self.image_target_size = image_target_size
+        self.random_rotation_angle = random_rotation_angle
+        self.color_jitter_params = color_jitter_params
+        self.processor = build_processor(model_name, transformers_loading_kwargs)
+        self.processor.tokenizer.padding_side = "left"
+        self.embodiment_id_mapping = embodiment_id_mapping or EMBODIMENT_TAG_TO_PROJECTOR_INDEX
+        for k, v in EMBODIMENT_TAG_TO_PROJECTOR_INDEX.items():
+            if k not in self.embodiment_id_mapping:
+                self.embodiment_id_mapping[k] = v
+        self.shortest_image_edge = shortest_image_edge
+        self.crop_fraction = crop_fraction
+
+        self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {}
+
+        # Eval/inference image transform
+        self.eval_image_transform = _build_eval_image_transform(
+            image_target_size,
+            image_crop_size,
+        )
+        self._collator = self.data_collator_class(
+            model_name=model_name,
+            model_type=model_type,
+            transformers_loading_kwargs=transformers_loading_kwargs,
+        )
+
+    @property
+    def collator(self):
+        return self._collator
+
+    def set_statistics(
+        self,
+        statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]],
+        override: bool = False,
+    ) -> None:
+        """Set dataset statistics for normalization."""
+        for key in statistics:
+            if key not in self.statistics or override:
+                if override:
+                    logger.info("Overriding statistics for %s", key)
+                self.statistics[key] = deepcopy(statistics[key])
+            else:
+                logger.debug("Embodiment tag %s already in statistics, skipping update", key)
+
+        self.state_action_processor.set_statistics(statistics, override=override)
+
+        # Compute action dimensions for convenience
+        self.action_dim = {}
+        for embodiment_tag in self.state_action_processor.statistics:
+            self.action_dim[embodiment_tag] = self.state_action_processor.get_action_dim(embodiment_tag)
+
+    def decode_action(
+        self,
+        action: np.ndarray,
+        embodiment_tag: EmbodimentTag,
+        state: dict[str, np.ndarray] | None = None,
+    ):
+        """Undo action normalization and convert relative actions to absolute."""
+        # Split concatenated action into joint groups
+        out_dict = {}
+        start_idx = 0
+        joint_groups = self.modality_configs[embodiment_tag.value]["action"].modality_keys
+        action_horizon = len(self.modality_configs[embodiment_tag.value]["action"].delta_indices)
+        for key in joint_groups:
+            joint_dim = self.state_action_processor.norm_params[embodiment_tag.value]["action"][key]["dim"].item()
+            out_dict[key] = action[..., :action_horizon, start_idx : start_idx + joint_dim]
+            start_idx += joint_dim
+
+        # Use StateActionProcessor to unnormalize and convert to absolute
+        return self.state_action_processor.unapply_action(out_dict, embodiment_tag.value, state=state)
+
+    def _apply_vlm_processing(self, images: np.ndarray, language: str) -> BatchFeature:
+        pil_images = [Image.fromarray(np.transpose(v, (1, 2, 0))) for v in images]
+        conversation = [
+            {
+                "role": "user",
+                "content": [
+                    *[{"type": "image", "image": img} for img in pil_images],
+                    {"type": "text", "text": language},
+                ],
+            }
+        ]
+
+        text = self.processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False)
+        return {
+            "vlm_content": {
+                "text": text,
+                "images": pil_images,
+                "conversation": conversation,
+            }
+        }
+
+    def __call__(
+        self,
+        messages: list[dict[str, Any]],
+    ):
+        assert len(messages) == 1
+        content = messages[0]["content"]
+        embodiment_tag = content.embodiment
+        action_data = content.actions
+        state_data = content.states
+
+        norm_state_dict, normalized_actions = self.state_action_processor.apply(
+            state=state_data,
+            action=action_data,
+            embodiment_tag=embodiment_tag.value,
+        )
+
+        if normalized_actions:
+            action_keys = self.modality_configs[embodiment_tag.value]["action"].modality_keys
+            normalized_actions = torch.cat(
+                [torch.from_numpy(normalized_actions[key]) for key in action_keys],
+                dim=-1,
+            )  # (t, d)
+            action_dim = normalized_actions.shape[1]
+            normalized_actions = torch.cat(
+                [
+                    normalized_actions,
+                    torch.zeros(
+                        normalized_actions.shape[0],
+                        self.max_action_dim - normalized_actions.shape[1],
+                    ),
+                ],
+                dim=-1,
+            )  # (t, max_action_dim)
+            action_horizon = normalized_actions.shape[0]
+            assert action_horizon <= self.max_action_horizon, (
+                f"Action sequence length {action_horizon} exceeds max_action_horizon"
+                f" {self.max_action_horizon}. Increase model config action_horizon to"
+                f" >= {action_horizon}."
+            )
+            normalized_actions = torch.cat(
+                [
+                    normalized_actions,
+                    torch.zeros(
+                        self.max_action_horizon - normalized_actions.shape[0],
+                        self.max_action_dim,
+                    ),
+                ],
+                dim=0,
+            )  # (max_action_horizon, max_action_dim)
+            action_mask = torch.ones_like(normalized_actions)
+            action_mask[action_horizon:] = 0
+            action_mask[:, action_dim:] = 0
+        else:
+            normalized_actions = None
+            action_mask = None
+
+        state_keys = self.modality_configs[embodiment_tag.value]["state"].modality_keys
+        exclude_state = self.exclude_state or getattr(
+            self.modality_configs[embodiment_tag.value]["state"], "exclude_state", False
+        )
+        if exclude_state:
+            normalized_states = torch.cat(
+                [torch.from_numpy(np.zeros_like(state_data[key])) for key in state_keys], dim=-1
+            )
+        else:
+            normalized_states = torch.cat([torch.from_numpy(norm_state_dict[key]) for key in state_keys], dim=-1)
+        normalized_states = torch.cat(
+            [
+                normalized_states,
+                torch.zeros(
+                    normalized_states.shape[0],
+                    self.max_state_dim - normalized_states.shape[1],
+                ),
+            ],
+            dim=-1,
+        )
+
+        image_transform = self.eval_image_transform
+        image_keys = self.modality_configs[embodiment_tag.value]["video"].modality_keys
+
+        if self.formalize_language:
+            language = content.text.lower()
+            language = re.sub(r"[^\w\s]", "", language)
+        else:
+            language = content.text
+
+        vlm_inputs = self._get_vlm_inputs(
+            image_keys=image_keys,
+            images=content.images,
+            masks=content.masks,
+            image_transform=image_transform,
+            language=language,
+        )
+
+        transformed_inputs = {
+            "state": normalized_states.to(torch.get_default_dtype()),
+        }
+        if normalized_actions is not None:
+            transformed_inputs["action"] = normalized_actions.to(torch.get_default_dtype())
+        # Add VLM inputs
+        transformed_inputs.update(vlm_inputs)
+        if action_mask is not None:
+            transformed_inputs["action_mask"] = action_mask
+        transformed_inputs["embodiment_id"] = self.embodiment_id_mapping[embodiment_tag.value]
+        return transformed_inputs
+
+    def _get_vlm_inputs(
+        self,
+        image_keys: list[str],
+        images: list[Image.Image],
+        masks: dict[str, list[np.ndarray]] | None,
+        image_transform: transforms.Compose,
+        language: str,
+    ):
+        temporal_stacked_images = {}
+
+        if masks is not None:
+            raise ValueError("Mask-based transforms are not supported at inference.")
+        for view in image_keys:
+            assert view in images, f"{view} not in {images}"
+            temporal_stacked_images[view] = torch.stack([image_transform(img) for img in images[view]])  # (T, C, H, W)
+
+        for k, v in temporal_stacked_images.items():
+            assert isinstance(k, str), f"{k} is not a string"
+            assert isinstance(v, torch.Tensor), f"{v} is not a torch tensor"
+            assert v.ndim == 4, f"{v} is not a 4D tensor"
+            assert v.dtype == torch.uint8, f"{v} is not a uint8 tensor"
+            assert v.shape[1] == 3, f"{v} is not a 3 channel tensor"
+
+        stacked_images = (
+            torch.stack([temporal_stacked_images[view] for view in image_keys], dim=1).flatten(0, 1).numpy()
+        )  # (T*V, C, H, W), processor expects numpy array
+
+        vlm_inputs = self._apply_vlm_processing(stacked_images, language)
+        return vlm_inputs
+
+    def save_pretrained(self, save_directory: str | Path) -> list[Path]:
+        save_directory = Path(save_directory)
+        save_directory.mkdir(parents=True, exist_ok=True)
+        main_config_file = save_directory / "processor_config.json"
+        statistics_file = save_directory / "statistics.json"
+        embodiment_id_file = save_directory / "embodiment_id.json"
+
+        config = {
+            "processor_class": self.__class__.__name__,
+            "processor_kwargs": {
+                "modality_configs": to_json_serializable(self.modality_configs),
+                "image_crop_size": self.image_crop_size,
+                "image_target_size": self.image_target_size,
+                "random_rotation_angle": self.random_rotation_angle,
+                "color_jitter_params": self.color_jitter_params,
+                "shortest_image_edge": self.shortest_image_edge,
+                "crop_fraction": self.crop_fraction,
+                "model_name": self.model_name,
+                "model_type": self.model_type,
+                "formalize_language": self.formalize_language,
+                "max_state_dim": self.max_state_dim,
+                "max_action_dim": self.max_action_dim,
+                "max_action_horizon": self.max_action_horizon,
+                "use_percentiles": self.use_percentiles,
+                "use_mean_std": self.use_mean_std,
+                "clip_outliers": self.clip_outliers,
+                "apply_sincos_state_encoding": self.apply_sincos_state_encoding,
+                "use_relative_action": self.use_relative_action,
+                "exclude_state": self.exclude_state,
+            },
+        }
+        with open(main_config_file, "w") as f:
+            json.dump(config, f, indent=2)
+        with open(statistics_file, "w") as f:
+            json.dump(
+                to_json_serializable(self.state_action_processor.statistics),
+                f,
+                indent=2,
+            )
+        with open(embodiment_id_file, "w") as f:
+            json.dump(self.embodiment_id_mapping, f, indent=2)
+        return [main_config_file, statistics_file, embodiment_id_file]
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs):
+        transformers_loading_kwargs = kwargs.pop("transformers_loading_kwargs", {"trust_remote_code": True})
+        pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
+        config_file = pretrained_model_name_or_path / "processor_config.json"
+        statistics_file = pretrained_model_name_or_path / "statistics.json"
+        embodiment_id_file = pretrained_model_name_or_path / "embodiment_id.json"
+        is_local = os.path.isdir(pretrained_model_name_or_path)
+        if not is_local:
+            config_file = Path(cached_file(pretrained_model_name_or_path, "processor_config.json"))
+            statistics_file = Path(cached_file(pretrained_model_name_or_path, "statistics.json"))
+            embodiment_id_file = Path(cached_file(pretrained_model_name_or_path, "embodiment_id.json"))
+
+        with open(config_file) as f:
+            config = json.load(f)
+        with open(statistics_file) as f:
+            statistics = json.load(f)
+        if embodiment_id_file.exists():
+            with open(embodiment_id_file) as f:
+                embodiment_id_mapping = json.load(f)
+        else:
+            embodiment_id_mapping = None
+        processor_kwargs = config["processor_kwargs"]
+        processor_kwargs["statistics"] = statistics
+        processor_kwargs["embodiment_id_mapping"] = embodiment_id_mapping
+
+        # Backfill missing fields from older checkpoints.
+        processor_kwargs.setdefault("model_name", "nvidia/Cosmos-Reason2-2B")
+        processor_kwargs.setdefault("model_type", "qwen")
+        processor_kwargs.setdefault("clip_outliers", True)
+
+        # Directly override other processor kwargs
+        if kwargs:
+            # Override modality configs while keeping pretrained embodiment configs
+            modality_configs = kwargs.pop("modality_configs", {})
+            for embodiment_tag, modality_config in modality_configs.items():
+                processor_kwargs["modality_configs"][embodiment_tag] = modality_config
+            override_keys = [
+                "random_rotation_angle",
+                "color_jitter_params",
+                "use_relative_action",
+                "exclude_state",
+                "use_mean_std",
+                "model_name",
+                "model_type",
+                "max_action_horizon",
+                "max_state_dim",
+                "max_action_dim",
+            ]
+            for key in override_keys:
+                if key in kwargs:
+                    override = kwargs.pop(key)
+                    if override is not None:
+                        processor_kwargs[key] = override
+        return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs)
+
+
+AutoProcessor.register("Gr00tN1d7", Gr00tN1d7Processor)
diff --git a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py
new file mode 100644
index 00000000000..98565a93faf
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py
@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+import numpy as np
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.gr00t.policy import Gr00tPolicy
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
+
+logger = init_logger(__name__)
+
+
+def _to_float32_action_dict(actions: Mapping[str, Any]) -> dict[str, np.ndarray]:
+    converted = {str(key): np.asarray(value, dtype=np.float32) for key, value in actions.items()}
+    if not converted:
+        raise RuntimeError("GR00T policy returned an empty action dict.")
+    return converted
+
+
+class Gr00tN1d7Pipeline(nn.Module):
+    """GR00T N1.7 policy pipeline backed by vLLM-Omni's local GR00T port.
+
+    vLLM-Omni owns the serving integration: OpenPI observations arrive through
+    `sampling_params.extra_args["robot_obs"]`, this pipeline runs GR00T policy
+    inference, and actions are returned through `DiffusionOutput.output["actions"]`.
+    """
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
+        super().__init__()
+        model_config = od_config.model_config
+        self.model_path = od_config.model
+        self.embodiment_tag = str(model_config.get("embodiment_tag") or "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT")
+        self.strict = bool(model_config.get("strict", True))
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        logger.info("Loading GR00T N1.7 policy from %s with embodiment_tag=%s", self.model_path, self.embodiment_tag)
+        self.policy = Gr00tPolicy(
+            model_path=self.model_path,
+            embodiment_tag=self.embodiment_tag,
+            device=self.device,
+            strict=self.strict,
+        )
+        self._validate_policy_server_config(model_config.get("policy_server_config"))
+
+    def _validate_policy_server_config(self, psc: Mapping[str, Any] | None) -> None:
+        """Fail fast if the deploy handshake drifts from the loaded checkpoint.
+
+        ``policy_server_config`` is sent verbatim to the OpenPI client, so its
+        model/embodiment-specific values must match what the loaded policy actually
+        produces; otherwise the client is handed the wrong action contract.
+        """
+        if not isinstance(psc, Mapping):
+            return
+        action_config = self.policy.modality_configs["action"]
+        expected_horizon = len(action_config.delta_indices)
+        expected_keys = set(action_config.modality_keys)
+
+        if "action_horizon" in psc and psc["action_horizon"] != expected_horizon:
+            raise ValueError(
+                f"policy_server_config.action_horizon={psc['action_horizon']} != loaded model's "
+                f"action horizon {expected_horizon}."
+            )
+        if "action_keys" in psc and set(psc["action_keys"]) != expected_keys:
+            raise ValueError(
+                f"policy_server_config.action_keys={list(psc['action_keys'])} != loaded model's "
+                f"action keys {sorted(expected_keys)}."
+            )
+        psc_embodiment = psc.get("embodiment_tag")
+        if psc_embodiment is not None and psc_embodiment != self.embodiment_tag:
+            raise ValueError(
+                f"policy_server_config.embodiment_tag={psc_embodiment!r} != "
+                f"model_config.embodiment_tag={self.embodiment_tag!r}."
+            )
+
+    def reset(self) -> dict[str, Any]:
+        return self.policy.reset() or {}
+
+    @property
+    def weights_sources(self) -> tuple[Any, ...]:
+        return ()
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        consumed = list(weights)
+        if consumed:
+            raise RuntimeError(
+                f"Gr00tN1d7Pipeline.load_weights received {len(consumed)} weight tensors; "
+                "weights_sources=() should prevent this. GR00T weights are loaded directly by Gr00tPolicy."
+            )
+        return set()
+
+    def _dummy_actions(self) -> dict[str, np.ndarray]:
+        embodiment_value = self.policy.embodiment_tag.value
+        action_config = self.policy.modality_configs["action"]
+        horizon = len(action_config.delta_indices)
+        norm_params = self.policy.processor.state_action_processor.norm_params[embodiment_value]["action"]
+        actions = {}
+        for key in action_config.modality_keys:
+            dim = norm_params[key]["dim"]
+            dim = int(dim.item() if hasattr(dim, "item") else dim)
+            actions[key] = np.zeros((1, horizon, dim), dtype=np.float32)
+        return actions
+
+    @torch.inference_mode()
+    def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
+        del kwargs
+        extra_args = req.sampling_params.extra_args or {}
+        robot_obs = extra_args.get("robot_obs")
+        if robot_obs is None:
+            if req.request_id == DUMMY_DIFFUSION_REQUEST_ID:
+                return DiffusionOutput(output={"actions": self._dummy_actions()})
+            return DiffusionOutput(error="Gr00tN1d7Pipeline.forward expects sampling_params.extra_args['robot_obs'].")
+        if not isinstance(robot_obs, Mapping):
+            return DiffusionOutput(error=f"robot_obs must be a dict, got {type(robot_obs).__name__}.")
+
+        if extra_args.get("reset"):
+            self.reset()
+
+        policy_obs = _normalize_observation(robot_obs, language_key=self.policy.language_key)
+        result = self.policy.get_action(policy_obs)
+        actions = result[0] if isinstance(result, tuple) else result
+        if not isinstance(actions, Mapping):
+            return DiffusionOutput(error=f"GR00T policy returned {type(actions).__name__}; expected dict actions.")
+        # Return actions via output.output (like the DreamZero OpenPI policy) so the engine's
+        # empty-output guard passes.
+        return DiffusionOutput(output={"actions": _to_float32_action_dict(actions)})
+
+
+def _normalize_observation(robot_obs: Mapping[str, Any], *, language_key: str) -> dict[str, Any]:
+    obs: dict[str, Any] = {}
+    if "video" in robot_obs:
+        obs["video"] = robot_obs["video"]
+    elif "images" in robot_obs:
+        obs["video"] = robot_obs["images"]
+    if "state" in robot_obs:
+        obs["state"] = robot_obs["state"]
+    if "language" in robot_obs:
+        obs["language"] = robot_obs["language"]
+    else:
+        prompt = robot_obs.get("prompt")
+        if prompt is not None:
+            obs["language"] = {language_key: [[str(prompt)]]}
+    return obs
diff --git a/vllm_omni/diffusion/models/gr00t/policy.py b/vllm_omni/diffusion/models/gr00t/policy.py
new file mode 100644
index 00000000000..c16b344632f
--- /dev/null
+++ b/vllm_omni/diffusion/models/gr00t/policy.py
@@ -0,0 +1,434 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from transformers import AutoModel, AutoProcessor
+
+from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import FINETUNE_ONLY_TAGS, POSTTRAIN_TAGS, EmbodimentTag
+from vllm_omni.diffusion.models.gr00t.dataio.types import MessageType, VLAStepData
+from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7
+from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7Processor
+
+
+def _rec_to_dtype(value: Any, dtype: torch.dtype) -> Any:
+    """Recursively convert floating-point tensors in nested collator output."""
+    if isinstance(value, torch.Tensor):
+        return value.to(dtype=dtype) if torch.is_floating_point(value) else value
+    if isinstance(value, dict) or hasattr(value, "items"):
+        return {key: _rec_to_dtype(item, dtype) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_rec_to_dtype(item, dtype) for item in value]
+    return value
+
+
+class Gr00tPolicy:
+    """Core policy class for Gr00t model inference.
+
+    This policy handles the end-to-end inference pipeline:
+    1. Validates input observations
+    2. Processes observations with pretrained VLA processor
+    3. Runs model inference
+    4. Decodes and returns actions
+
+    The policy expects observations with specific modalities (video, state, language)
+    and returns actions in the format defined by the model's modality configuration.
+    """
+
+    def __init__(
+        self,
+        embodiment_tag: EmbodimentTag | str,
+        model_path: str,
+        *,
+        device: int | str,
+        strict: bool = True,
+    ):
+        """Initialize the Gr00t Policy.
+
+        Args:
+            embodiment_tag: The embodiment tag defining the robot/environment type.
+                Accepts an EmbodimentTag enum or a string (resolved case-insensitively).
+            model_path: Path to the pretrained model checkpoint directory
+            device: Device to run the model on (e.g., 'cuda:0', 0, 'cpu')
+            strict: Whether to enforce strict input validation (default: True)
+        """
+        self.strict = strict
+        if isinstance(embodiment_tag, str):
+            embodiment_tag = EmbodimentTag.resolve(embodiment_tag)
+        model_dir = Path(model_path)
+
+        # Load directly into bfloat16 to avoid a transient fp32 copy at load time.
+        self.model: Gr00tN1d7 = AutoModel.from_pretrained(model_dir, torch_dtype=torch.bfloat16)
+        self.model.eval()
+        self.model.to(device=device)
+
+        # Training saves processor files under a "processor/" subdirectory, but
+        # AutoProcessor expects them at the model root.  Fall back to the
+        # subdirectory when the root lacks a processor_config.json.
+        processor_dir = (
+            model_dir / "processor"
+            if (model_dir / "processor").is_dir() and not (model_dir / "processor_config.json").exists()
+            else model_dir
+        )
+        self.processor: Gr00tN1d7Processor = AutoProcessor.from_pretrained(processor_dir)
+
+        # Store embodiment-specific configurations
+        self.embodiment_tag = embodiment_tag
+        all_modality_configs = self.processor.modality_configs
+        if self.embodiment_tag.value not in all_modality_configs:
+            # Map raw checkpoint tag values to user-friendly enum names where possible.
+            supported_lines = []
+            for tag_value in sorted(all_modality_configs.keys()):
+                enum_name = EmbodimentTag.reverse_lookup(tag_value)
+                if enum_name != tag_value:
+                    supported_lines.append(f"  {enum_name:30s} (--embodiment-tag {enum_name})")
+                else:
+                    supported_lines.append(f"  {tag_value:30s} (internal, no public enum)")
+            supported_str = "\n".join(supported_lines)
+
+            hint = ""
+            if self.embodiment_tag in POSTTRAIN_TAGS:
+                hint = (
+                    f"\n\nHint: '{self.embodiment_tag.name}' is a posttrain tag that requires "
+                    f"a finetuned checkpoint, not the base model. "
+                    f"See the example READMEs for how to finetune and download checkpoints."
+                )
+            elif self.embodiment_tag in FINETUNE_ONLY_TAGS:
+                hint = (
+                    f"\n\nHint: '{self.embodiment_tag.name}' is for finetuning custom robots. "
+                    f"Use it with launch_finetune.py, not with the base model directly."
+                )
+
+            raise ValueError(
+                f"Embodiment tag '{self.embodiment_tag.name}' "
+                f"(value='{self.embodiment_tag.value}') is not supported "
+                f"by this checkpoint.\n\n"
+                f"Supported tags in this checkpoint:\n{supported_str}"
+                f"{hint}"
+            )
+        self.modality_configs = {
+            k: v for k, v in all_modality_configs[self.embodiment_tag.value].items() if k != "rl_info"
+        }
+        self.collate_fn = self.processor.collator
+
+        # Extract and validate language configuration
+        # Embodiments may define multiple language keys (e.g. paraphrases); use only the first at inference.
+        language_keys = self.modality_configs["language"].modality_keys
+        language_delta_indices = self.modality_configs["language"].delta_indices
+        assert len(language_keys) >= 1, "At least one language key is required"
+        assert len(language_delta_indices) == 1, "Only one language delta index is supported"
+        self.language_key = language_keys[0]
+
+    def _unbatch_observation(self, value: dict[str, Any]) -> list[dict[str, Any]]:
+        """Unbatch a batched observation into a list of single observations.
+
+        Args:
+            value: Batched observation with shape (B, ...) for each modality
+
+        Returns:
+            List of B observations, each with the batch dimension removed
+        """
+        unbatched_obs = []
+        # Infer batch size from the first video key
+        batch_size = value["video"][list(value["video"].keys())[0]].shape[0]
+
+        # Split each modality along the batch dimension
+        for i in range(batch_size):
+            unbatched_value = {
+                "video": {k: v[i] for k, v in value["video"].items()},
+                "state": {k: v[i] for k, v in value["state"].items()},
+                "language": {k: v[i] for k, v in value["language"].items()},
+            }
+            unbatched_obs.append(unbatched_value)
+        return unbatched_obs
+
+    def _to_vla_step_data(self, observation: dict[str, Any]) -> VLAStepData:
+        """Convert a single observation into a VLAStepData object for processing.
+
+        Args:
+            observation: Single observation dict with video, state, and language
+
+        Returns:
+            VLAStepData object ready for processor input
+        """
+        return VLAStepData(
+            images=observation["video"],
+            states=observation["state"],
+            actions={},  # No ground truth actions during inference
+            text=observation["language"][self.language_key][0],
+            embodiment=self.embodiment_tag,
+        )
+
+    def check_observation(self, observation: dict[str, Any]) -> None:
+        """Validate that the observation has the correct structure and types.
+
+        This method ensures that all required modalities are present and that their
+        data types, shapes, and dimensions match the model's expectations.
+
+        Expected observation structure:
+            - video: dict[str, np.ndarray[np.uint8, (B, T, H, W, C)]]
+                - B: batch size
+                - T: temporal horizon (number of frames)
+                - H, W: image height and width
+                - C: number of channels (must be 3 for RGB)
+            - state: dict[str, np.ndarray[np.float32, (B, T, D)]]
+                - B: batch size
+                - T: temporal horizon (number of state observations)
+                - D: state dimension
+            - language: dict[str, list[list[str]]]
+                - Shape: (B, T) where each element is a string
+                - T: temporal horizon (typically 1 for language)
+
+        Args:
+            observation: Dictionary containing video, state, and language modalities
+
+        Raises:
+            AssertionError: If any validation check fails
+        """
+        # Check that observation contains all required top-level modality keys
+        for modality in ["video", "state", "language"]:
+            assert modality in observation, f"Observation must contain a '{modality}' key"
+            assert isinstance(observation[modality], dict), (
+                f"Observation '{modality}' must be a dictionary. "
+                f"Got {type(observation[modality])}: {observation[modality]}"
+            )
+
+        # Track batch size across modalities to ensure consistency
+        bs = -1
+
+        for video_key in self.modality_configs["video"].modality_keys:
+            assert video_key in observation["video"], f"Video key '{video_key}' must be in observation"
+
+            # Set or verify batch size consistency across all video keys
+            if bs == -1:
+                bs = len(observation["video"][video_key])
+            else:
+                assert len(observation["video"][video_key]) == bs, (
+                    f"Video key '{video_key}' must have batch size {bs}. Got {len(observation['video'][video_key])}"
+                )
+
+            batched_video = observation["video"][video_key]
+
+            # Verify data type is numpy array
+            assert isinstance(batched_video, np.ndarray), (
+                f"Video key '{video_key}' must be a numpy array. Got {type(batched_video)}"
+            )
+
+            # Verify dtype is uint8 (standard for image data, range 0-255)
+            assert batched_video.dtype == np.uint8, (
+                f"Video key '{video_key}' must be a numpy array of type np.uint8. Got {batched_video.dtype}"
+            )
+
+            # Verify shape has 5 dimensions: (B, T, H, W, C)
+            assert batched_video.ndim == 5, (
+                f"Video key '{video_key}' must be a numpy array of shape (B, T, H, W, C), got {batched_video.shape}"
+            )
+
+            # Verify temporal dimension matches the expected horizon from config
+            assert batched_video.shape[1] == len(self.modality_configs["video"].delta_indices), (
+                f"Video key '{video_key}'s horizon must be "
+                f"{len(self.modality_configs['video'].delta_indices)}. Got {batched_video.shape[1]}"
+            )
+
+            # Verify channel dimension is 3 (RGB images)
+            assert batched_video.shape[-1] == 3, (
+                f"Video key '{video_key}'s channel 'C' must be 3. Got {batched_video.shape[-1]}"
+            )
+
+        for state_key in self.modality_configs["state"].modality_keys:
+            # Check that the expected state key exists in the observation
+            # Must happen before indexing; see video validation above.
+            assert state_key in observation["state"], f"State key '{state_key}' must be in observation"
+
+            # Set or verify batch size consistency across all state keys
+            if bs == -1:
+                bs = len(observation["state"][state_key])
+            else:
+                assert len(observation["state"][state_key]) == bs, (
+                    f"State key '{state_key}' must have batch size {bs}. Got {len(observation['state'][state_key])}"
+                )
+
+            batched_state = observation["state"][state_key]
+
+            # Verify data type is numpy array
+            assert isinstance(batched_state, np.ndarray), (
+                f"State key '{state_key}' must be a numpy array. Got {type(batched_state)}"
+            )
+
+            # Verify dtype is float32 (standard for continuous state values)
+            assert batched_state.dtype == np.float32, (
+                f"State key '{state_key}' must be a numpy array of type np.float32. Got {batched_state.dtype}"
+            )
+
+            # Verify shape has 3 dimensions: (B, T, D)
+            assert batched_state.ndim == 3, (
+                f"State key '{state_key}' must be a numpy array of shape (B, T, D), got {batched_state.shape}"
+            )
+
+            # Verify temporal dimension matches the expected horizon from config
+            assert batched_state.shape[1] == len(self.modality_configs["state"].delta_indices), (
+                f"State key '{state_key}'s horizon must be "
+                f"{len(self.modality_configs['state'].delta_indices)}. Got {batched_state.shape[1]}"
+            )
+
+        for language_key in self.modality_configs["language"].modality_keys:
+            # Check that the expected language key exists in the observation
+            # Must happen before indexing; see video validation above.
+            assert language_key in observation["language"], f"Language key '{language_key}' must be in observation"
+
+            # Set or verify batch size consistency (language uses len instead of .shape)
+            if bs == -1:
+                bs = len(observation["language"][language_key])
+            else:
+                assert len(observation["language"][language_key]) == bs, (
+                    f"Language key '{language_key}' must have batch size {bs}. "
+                    f"Got {len(observation['language'][language_key])}"
+                )
+
+            batched_language: list[list[str]] = observation["language"][language_key]
+
+            # Verify outer structure is a list (batch dimension)
+            assert isinstance(batched_language, list), (
+                f"Language key '{language_key}' must be a list. Got {type(batched_language)}"
+            )
+
+            # Validate each batch item
+            for batch_item in batched_language:
+                # Verify temporal dimension matches expected horizon
+                assert len(batch_item) == len(self.modality_configs["language"].delta_indices), (
+                    f"Language key '{language_key}'s horizon must be "
+                    f"{len(self.modality_configs['language'].delta_indices)}. Got {len(batched_language)}"
+                )
+
+                # Verify inner structure is also a list (temporal dimension)
+                assert isinstance(batch_item, list), f"Language batch item must be a list. Got {type(batch_item)}"
+
+                # Current implementation expects exactly one language instruction per timestep
+                assert len(batch_item) == 1, f"Language batch item must have exactly one item. Got {len(batch_item)}"
+
+                # Verify the instruction itself is a string
+                assert isinstance(batch_item[0], str), (
+                    f"Language batch item must be a string. Got {type(batch_item[0])}"
+                )
+
+    def _get_action(
+        self, observation: dict[str, Any], options: dict[str, Any] | None = None
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Internal method to compute actions from observations.
+
+        Pipeline:
+        1. Unbatch observations into individual samples
+        2. Convert each to VLAStepData and process
+        3. Collate into model input batch
+        4. Run model inference
+        5. Decode and unnormalize actions
+
+        Args:
+            observation: Batched observation dictionary
+            options: Optional parameters (currently unused)
+
+        Returns:
+            Tuple of (actions_dict, info_dict)
+        """
+        # Step 1: Split batched observation into individual observations
+        unbatched_observations = self._unbatch_observation(observation)
+        processed_inputs = []
+
+        # Step 2: Process each observation through the VLA processor
+        states = []
+        for obs in unbatched_observations:
+            vla_step_data = self._to_vla_step_data(obs)
+            states.append(vla_step_data.states)  # dict[str, np.ndarray[np.float32, (T, D)]]
+            messages = [{"type": MessageType.EPISODE_STEP.value, "content": vla_step_data}]
+            processed_inputs.append(self.processor(messages))
+
+        # Step 3: Collate processed inputs into a single batch for model
+        collated_inputs = self.collate_fn(processed_inputs)
+        collated_inputs = _rec_to_dtype(collated_inputs, dtype=torch.bfloat16)
+
+        # Step 4: Run model inference to predict actions
+        with torch.inference_mode():
+            model_pred = self.model.get_action(**collated_inputs)
+        normalized_action = model_pred["action_pred"].float()
+
+        # Step 5: Decode actions from normalized space back to physical units
+        batched_states = {}
+        for k in self.modality_configs["state"].modality_keys:
+            batched_states[k] = np.stack([s[k] for s in states], axis=0)  # (B, T, D)
+        unnormalized_action = self.processor.decode_action(
+            normalized_action.cpu().numpy(), self.embodiment_tag, batched_states
+        )
+
+        # Cast all actions to float32 for consistency
+        casted_action = {key: value.astype(np.float32) for key, value in unnormalized_action.items()}
+        return casted_action, {}
+
+    def check_action(self, action: dict[str, Any]) -> None:
+        """Validate that the action has the correct structure and types.
+
+        This method ensures that all required action keys are present and that their
+        data types, shapes, and dimensions match the model's action space.
+
+        Expected action structure:
+            - action: dict[str, np.ndarray[np.float32, (B, T, D)]]
+                - B: batch size
+                - T: action horizon (number of future action steps)
+                - D: action dimension (e.g., joint positions, velocities, gripper state)
+
+        Args:
+            action: Dictionary containing action arrays for each action key
+
+        Raises:
+            AssertionError: If any validation check fails
+        """
+        # Validate each action key defined in the modality config
+        for action_key in self.modality_configs["action"].modality_keys:
+            # Check that the expected action key exists
+            assert action_key in action, f"Action key '{action_key}' must be in action"
+
+            action_arr = action[action_key]
+
+            # Verify data type is numpy array
+            assert isinstance(action_arr, np.ndarray), (
+                f"Action key '{action_key}' must be a numpy array. Got {type(action_arr)}"
+            )
+
+            # Verify dtype is float32 (standard for continuous actions)
+            assert action_arr.dtype == np.float32, (
+                f"Action key '{action_key}' must be a numpy array of type np.float32. Got {action_arr.dtype}"
+            )
+
+            # Verify shape has 3 dimensions: (B, T, D)
+            assert action_arr.ndim == 3, (
+                f"Action key '{action_key}' must be a numpy array of shape (B, T, D), got {action_arr.shape}"
+            )
+
+            # Verify action horizon matches the expected temporal dimension from config
+            assert action_arr.shape[1] == len(self.modality_configs["action"].delta_indices), (
+                f"Action key '{action_key}'s horizon must be "
+                f"{len(self.modality_configs['action'].delta_indices)}. Got {action_arr.shape[1]}"
+            )
+
+    def get_action(
+        self, observation: dict[str, Any], options: dict[str, Any] | None = None
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        if self.strict:
+            self.check_observation(observation)
+        action, info = self._get_action(observation, options)
+        if self.strict:
+            self.check_action(action)
+        return action, info
+
+    def reset(self, options: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Reset the policy to its initial state.
+
+        Args:
+            options: Dictionary containing the options for the reset
+
+        Returns:
+            Dictionary containing the info after resetting the policy
+        """
+        return {}
diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py
index ed0a50fee38..730c10f1c78 100644
--- a/vllm_omni/diffusion/models/helios/helios_transformer.py
+++ b/vllm_omni/diffusion/models/helios/helios_transformer.py
@@ -11,6 +11,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import FP32LayerNorm
@@ -25,6 +26,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
@@ -602,6 +604,13 @@ class HeliosTransformer3DModel(nn.Module):
     guidance cross-attention, and chunked video generation support.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "blocks": ForwardPattern.Pattern_2,
+        },
+        has_separate_cfg=True,
+    )
+
     _repeated_blocks = ["HeliosTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["blocks"]
     packed_modules_mapping = {
diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py
deleted file mode 100644
index abc50eb8a22..00000000000
--- a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py
+++ /dev/null
@@ -1,1059 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import inspect
-import json
-import math
-import os
-from collections.abc import Callable, Iterable
-from typing import Any
-
-import numpy as np
-import torch
-from diffusers.image_processor import VaeImageProcessor
-from diffusers.models import AutoencoderKL
-from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import deprecate, logging
-from diffusers.utils.torch_utils import randn_tensor
-from torch import nn
-from transformers import (
-    CLIPTextModelWithProjection,
-    CLIPTokenizer,
-    LlamaForCausalLM,
-    PreTrainedTokenizerFast,
-    T5EncoderModel,
-    T5Tokenizer,
-)
-from vllm.model_executor.models.utils import AutoWeightsLoader
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
-from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
-from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel
-from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
-from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
-from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
-
-logger = logging.get_logger(__name__)
-
-
-def get_hidream_image_post_process_func(
-    od_config: OmniDiffusionConfig,
-):
-    model_name = od_config.model
-    if os.path.exists(model_name):
-        model_path = model_name
-    else:
-        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
-    vae_config_path = os.path.join(model_path, "vae/config.json")
-    if not os.path.exists(vae_config_path):
-        raise FileNotFoundError(
-            f"VAE config not found at {vae_config_path}. "
-            "Please ensure the model path contains a valid VAE configuration."
-        )
-    with open(vae_config_path) as f:
-        vae_config = json.load(f)
-        vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
-
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
-
-    def post_process_func(
-        images: torch.Tensor,
-    ):
-        return image_processor.postprocess(images)
-
-    return post_process_func
-
-
-# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
-def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
-):
-    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-    b = base_shift - m * base_seq_len
-    mu = image_seq_len * m + b
-    return mu
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: int | None = None,
-    device: str | torch.device | None = None,
-    timesteps: list[int] | None = None,
-    sigmas: list[float] | None = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`list[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`list[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        timesteps (`torch.Tensor`): The timestep schedule from the scheduler.
-        num_inference_steps (`int`): The number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class HiDreamI1ImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, ProgressBarMixin):
-    def __init__(
-        self,
-        *,
-        od_config: OmniDiffusionConfig,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.od_config = od_config
-        dtype = getattr(od_config, "dtype", torch.bfloat16)
-        llama_path = self.od_config.extras["auxiliary_text_encoder"]
-        if llama_path is None:
-            logger.warning(
-                f"auxiliary_text_encoder is not provided. "
-                f"Attempting to load default LLaMA model: {self.od_config.extras['default_llama_model_id']}"
-            )
-            llama_path = self.od_config.extras["default_llama_model_id"]
-        logger.info(f"Loading LLAMA model from: {llama_path}")
-
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model,
-                subfolder="transformer",
-                revision=None,
-                prefix="transformer.",
-                fall_back_to_pt=True,
-            ),
-        ]
-
-        self.device = get_local_device()
-        model = od_config.model
-        # Check if model is a local path
-        local_files_only = os.path.exists(model)
-
-        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
-        # race; prefetch the in-repo component set before any from_pretrained
-        # (``text_encoder_4`` lives in a separate Llama repo and is unaffected).
-        hidream_subfolders = [
-            "scheduler",
-            "vae",
-            "text_encoder",
-            "tokenizer",
-            "text_encoder_2",
-            "tokenizer_2",
-            "text_encoder_3",
-            "tokenizer_3",
-        ]
-        prefetch_subfolders(model, hidream_subfolders, local_files_only=local_files_only)
-
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            model, subfolder="scheduler", local_files_only=local_files_only
-        )
-        self.vae = from_pretrained_with_prefetch(
-            AutoencoderKL.from_pretrained,
-            model,
-            subfolder="vae",
-            prefetch_list=hidream_subfolders,
-            local_files_only=local_files_only,
-        ).to(self.device)
-        self.text_encoder = from_pretrained_with_prefetch(
-            CLIPTextModelWithProjection.from_pretrained,
-            model,
-            subfolder="text_encoder",
-            prefetch_list=hidream_subfolders,
-            local_files_only=local_files_only,
-        )
-        self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.text_encoder_2 = from_pretrained_with_prefetch(
-            CLIPTextModelWithProjection.from_pretrained,
-            model,
-            subfolder="text_encoder_2",
-            prefetch_list=hidream_subfolders,
-            local_files_only=local_files_only,
-        )
-        self.tokenizer_2 = CLIPTokenizer.from_pretrained(
-            model, subfolder="tokenizer_2", local_files_only=local_files_only
-        )
-        self.text_encoder_3 = from_pretrained_with_prefetch(
-            T5EncoderModel.from_pretrained,
-            model,
-            subfolder="text_encoder_3",
-            prefetch_list=hidream_subfolders,
-            local_files_only=local_files_only,
-        )
-        self.tokenizer_3 = T5Tokenizer.from_pretrained(
-            model, subfolder="tokenizer_3", local_files_only=local_files_only
-        )
-        self.text_encoder_4 = LlamaForCausalLM.from_pretrained(llama_path, output_hidden_states=True, dtype=dtype).to(
-            self.device
-        )
-        self.tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(llama_path, use_fast=False)
-        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HiDreamImageTransformer2DModel)
-        self.transformer = HiDreamImageTransformer2DModel(
-            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
-        )
-
-        self.stage = None
-
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
-        # HiDreamImage latents are turned into 2x2 patches and packed.
-        # This means the latent width and height has to be divisible
-        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
-        self.default_sample_size = 128
-        if getattr(self, "tokenizer_4", None) is not None:
-            self.tokenizer_4.pad_token = self.tokenizer_4.eos_token
-
-        self.setup_diffusion_pipeline_profiler(
-            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
-        )
-
-    def check_inputs(
-        self,
-        prompt,
-        prompt_2,
-        prompt_3,
-        prompt_4,
-        negative_prompt=None,
-        negative_prompt_2=None,
-        negative_prompt_3=None,
-        negative_prompt_4=None,
-        prompt_embeds_t5=None,
-        prompt_embeds_llama3=None,
-        negative_prompt_embeds_t5=None,
-        negative_prompt_embeds_llama3=None,
-        pooled_prompt_embeds=None,
-        negative_pooled_prompt_embeds=None,
-        callback_on_step_end_tensor_inputs=None,
-    ):
-        if prompt is not None and pooled_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. "
-                "Please make sure to only forward one of the two."
-            )
-        elif prompt_2 is not None and pooled_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}."
-                "Please make sure to only forward one of the two."
-            )
-        elif prompt_3 is not None and prompt_embeds_t5 is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}."
-                "Please make sure to only forward one of the two."
-            )
-        elif prompt_4 is not None and prompt_embeds_llama3 is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}."
-                "Please make sure to only forward one of the two."
-            )
-        elif prompt is None and pooled_prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `pooled_prompt_embeds`."
-                "Cannot leave both `prompt` and `pooled_prompt_embeds` undefined."
-            )
-        elif prompt is None and prompt_embeds_t5 is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds_t5`."
-                "Cannot leave both `prompt` and `prompt_embeds_t5` undefined."
-            )
-        elif prompt is None and prompt_embeds_llama3 is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds_llama3`."
-                "Cannot leave both `prompt` and `prompt_embeds_llama3` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
-            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
-        elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)):
-            raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}")
-        elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)):
-            raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}")
-
-        if negative_prompt is not None and negative_pooled_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:"
-                f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two."
-            )
-        elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:"
-                f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two."
-            )
-        elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:"
-                f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two."
-            )
-        elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:"
-                f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two."
-            )
-
-        if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None:
-            if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape:
-                raise ValueError(
-                    "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds`"
-                    " must have the same shape when passed directly, but"
-                    f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`"
-                    f" {negative_pooled_prompt_embeds.shape}."
-                )
-        if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None:
-            if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape:
-                raise ValueError(
-                    "`prompt_embeds_t5` and `negative_prompt_embeds_t5`"
-                    " must have the same shape when passed directly, but"
-                    f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`"
-                    f" {negative_prompt_embeds_t5.shape}."
-                )
-        if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None:
-            if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape:
-                raise ValueError(
-                    "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3`"
-                    " must have the same shape when passed directly, but"
-                    f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`"
-                    f" {negative_prompt_embeds_llama3.shape}."
-                )
-
-    def _get_t5_prompt_embeds(
-        self,
-        prompt: str | list[str] = None,
-        max_sequence_length: int = 128,
-        dtype: torch.dtype | None = None,
-    ):
-        dtype = dtype or self.text_encoder_3.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        text_inputs = self.tokenizer_3(
-            prompt,
-            padding="max_length",
-            max_length=min(max_sequence_length, self.tokenizer_3.model_max_length),
-            truncation=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids
-        attention_mask = text_inputs.attention_mask
-        untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
-
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer_3.batch_decode(
-                untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1]
-            )
-            logger.warning(
-                "The following part of your input was truncated because `max_sequence_length` is set to "
-                f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}"
-            )
-
-        prompt_embeds = self.text_encoder_3(
-            text_input_ids.to(self.device), attention_mask=attention_mask.to(self.device)
-        )[0]
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
-        return prompt_embeds
-
-    def _get_clip_prompt_embeds(
-        self,
-        tokenizer,
-        text_encoder,
-        prompt: str | list[str],
-        max_sequence_length: int = 128,
-        dtype: torch.dtype | None = None,
-    ):
-        dtype = dtype or text_encoder.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        model_max_length = self.tokenizer.model_max_length
-        text_inputs = tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=min(max_sequence_length, model_max_length),
-            truncation=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids
-        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = tokenizer.batch_decode(untruncated_ids[:, model_max_length - 1 : -1])
-            logger.warning(
-                "The following part of your input was truncated because CLIP can only handle sequences up to"
-                f" {model_max_length} tokens: {removed_text}"
-            )
-        prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True)
-
-        # Use pooled output of CLIPTextModel
-        prompt_embeds = prompt_embeds[0]
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
-        return prompt_embeds
-
-    def _get_llama3_prompt_embeds(
-        self,
-        prompt: str | list[str] = None,
-        max_sequence_length: int = 128,
-        dtype: torch.dtype | None = None,
-    ):
-        dtype = dtype or self.text_encoder_4.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        text_inputs = self.tokenizer_4(
-            prompt,
-            padding="max_length",
-            max_length=min(max_sequence_length, self.tokenizer_4.model_max_length),
-            truncation=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids
-        attention_mask = text_inputs.attention_mask
-        untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids
-
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer_4.batch_decode(
-                untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1]
-            )
-            logger.warning(
-                "The following part of your input was truncated because `max_sequence_length` is set to "
-                f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}"
-            )
-
-        outputs = self.text_encoder_4(
-            text_input_ids.to(self.device),
-            attention_mask=attention_mask.to(self.device),
-            output_hidden_states=True,
-        )
-
-        prompt_embeds = outputs.hidden_states[1:]
-        prompt_embeds = torch.stack(prompt_embeds, dim=0)
-        return prompt_embeds
-
-    def encode_prompt(
-        self,
-        prompt: str | list[str] | None = None,
-        prompt_2: str | list[str] | None = None,
-        prompt_3: str | list[str] | None = None,
-        prompt_4: str | list[str] | None = None,
-        dtype: torch.dtype | None = None,
-        num_images_per_prompt: int = 1,
-        do_classifier_free_guidance: bool = True,
-        negative_prompt: str | list[str] | None = None,
-        negative_prompt_2: str | list[str] | None = None,
-        negative_prompt_3: str | list[str] | None = None,
-        negative_prompt_4: str | list[str] | None = None,
-        prompt_embeds_t5: list[torch.FloatTensor] | None = None,
-        prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
-        negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None,
-        negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
-        pooled_prompt_embeds: torch.FloatTensor | None = None,
-        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
-        max_sequence_length: int = 128,
-        lora_scale: float | None = None,
-    ):
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        if prompt is not None:
-            batch_size = len(prompt)
-        else:
-            batch_size = pooled_prompt_embeds.shape[0]
-
-        if pooled_prompt_embeds is None:
-            pooled_prompt_embeds_1 = self._get_clip_prompt_embeds(
-                self.tokenizer, self.text_encoder, prompt, max_sequence_length, dtype
-            )
-
-        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
-
-            if len(negative_prompt) > 1 and len(negative_prompt) != batch_size:
-                raise ValueError(f"negative_prompt must be of length 1 or {batch_size}")
-
-            negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds(
-                self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, dtype
-            )
-
-            if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1:
-                negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1)
-
-        if pooled_prompt_embeds is None:
-            prompt_2 = prompt_2 or prompt
-            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
-
-            if len(prompt_2) > 1 and len(prompt_2) != batch_size:
-                raise ValueError(f"prompt_2 must be of length 1 or {batch_size}")
-
-            pooled_prompt_embeds_2 = self._get_clip_prompt_embeds(
-                self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, dtype
-            )
-
-            if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1:
-                pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1)
-
-        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
-            negative_prompt_2 = negative_prompt_2 or negative_prompt
-            negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
-
-            if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size:
-                raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}")
-
-            negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds(
-                self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, dtype
-            )
-
-            if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1:
-                negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1)
-
-        if pooled_prompt_embeds is None:
-            pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1)
-
-        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
-            negative_pooled_prompt_embeds = torch.cat(
-                [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1
-            )
-
-        if prompt_embeds_t5 is None:
-            prompt_3 = prompt_3 or prompt
-            prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
-
-            if len(prompt_3) > 1 and len(prompt_3) != batch_size:
-                raise ValueError(f"prompt_3 must be of length 1 or {batch_size}")
-
-            prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, dtype)
-
-            if prompt_embeds_t5.shape[0] == 1 and batch_size > 1:
-                prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1)
-
-        if do_classifier_free_guidance and negative_prompt_embeds_t5 is None:
-            negative_prompt_3 = negative_prompt_3 or negative_prompt
-            negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
-
-            if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size:
-                raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}")
-
-            negative_prompt_embeds_t5 = self._get_t5_prompt_embeds(negative_prompt_3, max_sequence_length, dtype)
-
-            if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1:
-                negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1)
-
-        if prompt_embeds_llama3 is None:
-            prompt_4 = prompt_4 or prompt
-            prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4
-
-            if len(prompt_4) > 1 and len(prompt_4) != batch_size:
-                raise ValueError(f"prompt_4 must be of length 1 or {batch_size}")
-
-            prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, dtype)
-
-            if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1:
-                prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
-
-        if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None:
-            negative_prompt_4 = negative_prompt_4 or negative_prompt
-            negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4
-
-            if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size:
-                raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}")
-
-            negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds(
-                negative_prompt_4, max_sequence_length, dtype
-            )
-
-            if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1:
-                negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
-
-        # duplicate pooled_prompt_embeds for each generation per prompt
-        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt)
-        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
-
-        # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt
-        bs_embed, seq_len, _ = prompt_embeds_t5.shape
-        if bs_embed == 1 and batch_size > 1:
-            prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1)
-        elif bs_embed > 1 and bs_embed != batch_size:
-            raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}")
-        prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-        # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt
-        _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape
-        if bs_embed == 1 and batch_size > 1:
-            prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
-        elif bs_embed > 1 and bs_embed != batch_size:
-            raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}")
-        prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1)
-        prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim)
-
-        if do_classifier_free_guidance:
-            # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt
-            bs_embed, seq_len = negative_pooled_prompt_embeds.shape
-            if bs_embed == 1 and batch_size > 1:
-                negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1)
-            elif bs_embed > 1 and bs_embed != batch_size:
-                raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}")
-            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt)
-            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
-
-            # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt
-            bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape
-            if bs_embed == 1 and batch_size > 1:
-                negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1)
-            elif bs_embed > 1 and bs_embed != batch_size:
-                raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}")
-            negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1)
-            negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-            # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt
-            _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape
-            if bs_embed == 1 and batch_size > 1:
-                negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
-            elif bs_embed > 1 and bs_embed != batch_size:
-                raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}")
-            negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1)
-            negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view(
-                -1, batch_size * num_images_per_prompt, seq_len, dim
-            )
-
-        return (
-            prompt_embeds_t5,
-            negative_prompt_embeds_t5,
-            prompt_embeds_llama3,
-            negative_prompt_embeds_llama3,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        )
-
-    def enable_vae_slicing(self):
-        r"""
-        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
-        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        depr_message = (
-            f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this"
-            " method will be removed in a future version. Please use `pipe.vae.enable_slicing()`."
-        )
-        deprecate(
-            "enable_vae_slicing",
-            "0.40.0",
-            depr_message,
-        )
-        self.vae.enable_slicing()
-
-    def disable_vae_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        depr_message = (
-            f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this"
-            " method will be removed in a future version. Please use `pipe.vae.disable_slicing()`."
-        )
-        deprecate(
-            "disable_vae_slicing",
-            "0.40.0",
-            depr_message,
-        )
-        self.vae.disable_slicing()
-
-    def enable_vae_tiling(self):
-        r"""
-        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
-        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
-        processing larger images.
-        """
-        depr_message = (
-            f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}`"
-            " is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`."
-        )
-        deprecate(
-            "enable_vae_tiling",
-            "0.40.0",
-            depr_message,
-        )
-        self.vae.enable_tiling()
-
-    def disable_vae_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        depr_message = (
-            f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is"
-            " deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`."
-        )
-        deprecate(
-            "disable_vae_tiling",
-            "0.40.0",
-            depr_message,
-        )
-        self.vae.disable_tiling()
-
-    def prepare_latents(
-        self,
-        batch_size,
-        num_channels_latents,
-        height,
-        width,
-        dtype,
-        generator,
-        latents=None,
-    ):
-        # VAE applies 8x compression on images but we must also account for packing which requires
-        # latent height and width to be divisible by 2.
-        height = 2 * (int(height) // (self.vae_scale_factor * 2))
-        width = 2 * (int(width) // (self.vae_scale_factor * 2))
-
-        shape = (batch_size, num_channels_latents, height, width)
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype)
-        else:
-            if latents.shape != shape:
-                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
-            latents = latents.to(self.device)
-        return latents
-
-    def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len):
-        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
-        mu = calculate_shift(
-            image_seq_len,
-            self.scheduler.config.get("base_image_seq_len", 256),
-            self.scheduler.config.get("max_image_seq_len", 4096),
-            self.scheduler.config.get("base_shift", 0.5),
-            self.scheduler.config.get("max_shift", 1.15),
-        )
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            self.device,
-            sigmas=sigmas,
-            mu=mu,
-        )
-        return timesteps, num_inference_steps
-
-    def _extract_prompts(self, prompts):
-        """Extract prompt and negative_prompt from OmniPromptType list."""
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
-        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
-            negative_prompt = None
-        elif prompts:
-            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
-        else:
-            negative_prompt = None
-        return prompt, negative_prompt
-
-    def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool):
-        if get_classifier_free_guidance_world_size() == 1:
-            return True
-
-        if true_cfg_scale <= 1:
-            logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.")
-            return False
-
-        if not has_neg_prompt:
-            logger.warning(
-                "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings."
-            )
-            return False
-        return True
-
-    def diffuse(
-        self,
-        prompt_embeds_t5: torch.Tensor,
-        prompt_embeds_llama3: torch.Tensor,
-        pooled_prompt_embeds: torch.Tensor,
-        latents: torch.Tensor,
-        timesteps: torch.Tensor,
-        do_true_cfg: bool,
-    ) -> torch.Tensor:
-        with self.progress_bar(total=len(timesteps)) as pbar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0])
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    timesteps=timestep,
-                    encoder_hidden_states_t5=prompt_embeds_t5,
-                    encoder_hidden_states_llama3=prompt_embeds_llama3,
-                    pooled_embeds=pooled_prompt_embeds,
-                    return_dict=False,
-                )[0]
-                noise_pred = -noise_pred
-
-                # TODO: Modify CFG guidance
-                # perform guidance
-                if self.do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
-
-                pbar.update()
-
-        return latents
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1
-
-    @property
-    def attention_kwargs(self):
-        return self._attention_kwargs
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    def forward(
-        self,
-        req: OmniDiffusionRequest,
-        prompt: str | list[str] = None,
-        prompt_2: str | list[str] | None = None,
-        prompt_3: str | list[str] | None = None,
-        prompt_4: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 5.0,
-        negative_prompt: str | list[str] | None = None,
-        negative_prompt_2: str | list[str] | None = None,
-        negative_prompt_3: str | list[str] | None = None,
-        negative_prompt_4: str | list[str] | None = None,
-        num_images_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds_t5: torch.FloatTensor | None = None,
-        prompt_embeds_llama3: torch.FloatTensor | None = None,
-        negative_prompt_embeds_t5: torch.FloatTensor | None = None,
-        negative_prompt_embeds_llama3: torch.FloatTensor | None = None,
-        pooled_prompt_embeds: torch.FloatTensor | None = None,
-        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 128,
-        **kwargs,
-    ):
-        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
-        prompt = extracted_prompt or prompt
-
-        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
-        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
-
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or guidance_scale
-        if req.sampling_params.guidance_scale_provided:
-            guidance_scale = req.sampling_params.guidance_scale
-        num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
-        )
-
-        prompt_embeds = kwargs.get("prompt_embeds", None)
-        negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None)
-
-        if prompt_embeds is not None:
-            deprecation_message = (
-                "The `prompt_embeds` argument is deprecated."
-                " Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead."
-            )
-            deprecate("prompt_embeds", "0.35.0", deprecation_message)
-            prompt_embeds_t5 = prompt_embeds[0]
-            prompt_embeds_llama3 = prompt_embeds[1]
-
-        if negative_prompt_embeds is not None:
-            deprecation_message = (
-                "The `negative_prompt_embeds` argument is deprecated."
-                "Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead."
-            )
-            deprecate("negative_prompt_embeds", "0.35.0", deprecation_message)
-            negative_prompt_embeds_t5 = negative_prompt_embeds[0]
-            negative_prompt_embeds_llama3 = negative_prompt_embeds[1]
-
-        division = self.vae_scale_factor * 2
-        S_max = (self.default_sample_size * self.vae_scale_factor) ** 2
-        scale = S_max / (width * height)
-        scale = math.sqrt(scale)
-        width, height = int(width * scale // division * division), int(height * scale // division * division)
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            prompt_2,
-            prompt_3,
-            prompt_4,
-            negative_prompt=negative_prompt,
-            negative_prompt_2=negative_prompt_2,
-            negative_prompt_3=negative_prompt_3,
-            negative_prompt_4=negative_prompt_4,
-            prompt_embeds_t5=prompt_embeds_t5,
-            prompt_embeds_llama3=prompt_embeds_llama3,
-            negative_prompt_embeds_t5=negative_prompt_embeds_t5,
-            negative_prompt_embeds_llama3=negative_prompt_embeds_llama3,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._interrupt = False
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        elif pooled_prompt_embeds is not None:
-            batch_size = pooled_prompt_embeds.shape[0]
-
-        # TODO: CFG guidance configuration
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
-        )
-        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
-        self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
-
-        # 3. Encode prompt
-        lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None
-        (
-            prompt_embeds_t5,
-            negative_prompt_embeds_t5,
-            prompt_embeds_llama3,
-            negative_prompt_embeds_llama3,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            prompt_2=prompt_2,
-            prompt_3=prompt_3,
-            prompt_4=prompt_4,
-            negative_prompt=negative_prompt,
-            negative_prompt_2=negative_prompt_2,
-            negative_prompt_3=negative_prompt_3,
-            negative_prompt_4=negative_prompt_4,
-            do_classifier_free_guidance=self.do_classifier_free_guidance,
-            prompt_embeds_t5=prompt_embeds_t5,
-            prompt_embeds_llama3=prompt_embeds_llama3,
-            negative_prompt_embeds_t5=negative_prompt_embeds_t5,
-            negative_prompt_embeds_llama3=negative_prompt_embeds_llama3,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            num_images_per_prompt=num_images_per_prompt,
-            max_sequence_length=max_sequence_length,
-            lora_scale=lora_scale,
-        )
-
-        if self.do_classifier_free_guidance:
-            prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0)
-            prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1)
-            pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
-
-        # 4. Prepare latent variables
-        num_channels_latents = self.transformer.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            pooled_prompt_embeds.dtype,
-            generator,
-            latents,
-        )
-
-        # 5. Prepare timesteps
-        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
-        self._num_timesteps = len(timesteps)
-
-        # 6. Denoising loop
-        latents = self.diffuse(
-            prompt_embeds_t5,
-            prompt_embeds_llama3,
-            pooled_prompt_embeds,
-            latents,
-            timesteps,
-            do_true_cfg,
-        )
-
-        if output_type == "latent":
-            image = latents
-
-        else:
-            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
-
-            image = self.vae.decode(latents, return_dict=False)[0]
-
-        return DiffusionOutput(
-            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
-        )
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
index 5ce3098ad40..e4c9e0b2a1e 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
@@ -11,6 +11,7 @@
 import numpy as np
 import regex as re
 import torch
+from cache_dit import ForwardPattern
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
@@ -59,6 +60,7 @@
     AttentionMetadata,
 )
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
     get_classifier_free_guidance_rank,
@@ -2018,6 +2020,13 @@ def forward(self, x):
 
 
 class HunyuanImage3Model(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "layers": ForwardPattern.Pattern_4,
+        },
+        check_forward_pattern=False,
+    )
+
     _sp_plan = {
         # Split custom_pos_emb tuple elements (cos, sin) at model forward input
         "pre_processor": {
diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
index 9af9e848875..9224973dda8 100644
--- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
+++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
@@ -6,6 +6,7 @@
 
 import torch
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     CombinedTimestepTextProjEmbeddings,
     TimestepEmbedding,
@@ -22,6 +23,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size
@@ -591,6 +593,13 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
     tensor-parallel layers for the 54 main transformer blocks.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_0,
+        },
+        has_separate_cfg=True,
+    )
+
     _repeated_blocks = ["HunyuanVideo15TransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     packed_modules_mapping = {
diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
index 8f0ff446afd..744c1806a5a 100644
--- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
+++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
@@ -6,6 +6,7 @@
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
@@ -18,6 +19,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -581,6 +583,14 @@ class LongCatImageTransformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+            "single_transformer_blocks": ForwardPattern.Pattern_1,
+        },
+        has_separate_cfg=True,
+    )
+
     _repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
index dcf848390d4..6072c04560a 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
@@ -23,6 +23,7 @@
 import torch.distributed
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings, PixArtAlphaTextProjection
 from diffusers.utils import (
     BaseOutput,
@@ -44,6 +45,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
@@ -1449,6 +1451,14 @@ class LTX2VideoTransformer3DModel(nn.Module):
             The normalization layer to use.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_0,
+        },
+        has_separate_cfg=True,
+        check_forward_pattern=False,
+    )
+
     _supports_gradient_checkpointing = True
     _skip_layerwise_casting_patterns = ["norm"]
     _repeated_blocks = ["LTX2VideoTransformerBlock"]
diff --git a/vllm_omni/diffusion/models/sd3/sd3_transformer.py b/vllm_omni/diffusion/models/sd3/sd3_transformer.py
index 89f06157758..d8e03705112 100644
--- a/vllm_omni/diffusion/models/sd3/sd3_transformer.py
+++ b/vllm_omni/diffusion/models/sd3/sd3_transformer.py
@@ -5,6 +5,7 @@
 import torch.nn.functional as F
 
 # TODO replace this with vLLM implementation
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, SD35AdaLayerNormZeroX
@@ -19,6 +20,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 
 logger = init_logger(__name__)
@@ -386,6 +388,12 @@ class SD3Transformer2DModel(nn.Module):
     The Transformer model introduced in [Stable Diffusion 3](https://huggingface.co/papers/2403.03206).
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+        }
+    )
+
     _repeated_blocks = ["SD3TransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
 
diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
index 44446ae8381..c1ba258b28c 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
@@ -534,6 +534,8 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
             self.llm_cfg,
             prefix="language_model",
         )
+        # We define this for Cache DiT compatibility
+        self.transformer = self.language_model.model
 
         # Vision model (understanding branch)
         self.vision_model = NEOVisionModel(self.vis_cfg)
diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
index f9ff8aaf470..53e1a47c239 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
@@ -16,6 +16,7 @@
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from transformers.cache_utils import DynamicCache
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -30,6 +31,7 @@
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, SensenovaCachedAdapter
 
 
 @dataclass
@@ -606,6 +608,14 @@ def forward(
 
 
 class SenseNovaU1Model(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "layers": ForwardPattern.Pattern_3,
+        },
+        has_separate_cfg=True,
+        cached_adapter_cls=SensenovaCachedAdapter,
+    )
+
     def __init__(self, config, quant_config=None, prefix: str = ""):
         super().__init__()
         self.config = config
diff --git a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py b/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py
index 83a669d92aa..c7bc9b2d4f6 100644
--- a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py
+++ b/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py
@@ -9,12 +9,14 @@
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.layers.fourier import GaussianFourierProjection
@@ -464,6 +466,12 @@ class StableAudioDiTModel(nn.Module):
     - Output: [B, out_channels, L]
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_3,
+        },
+    )
+
     _repeated_blocks = ["StableAudioDiTBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     _hsdp_shard_conditions = [is_transformer_block_module]
diff --git a/vllm_omni/diffusion/quantization/__init__.py b/vllm_omni/diffusion/quantization/__init__.py
new file mode 100644
index 00000000000..208f01a7cb5
--- /dev/null
+++ b/vllm_omni/diffusion/quantization/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/diffusion/quantization/hsdp_fp8.py b/vllm_omni/diffusion/quantization/hsdp_fp8.py
new file mode 100644
index 00000000000..c2ddd68ae9c
--- /dev/null
+++ b/vllm_omni/diffusion/quantization/hsdp_fp8.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HSDP/FSDP2 compatibility for online FP8 quantization.
+
+vllm's ``Fp8OnlineLinearMethod.process_weights_after_loading`` ends with
+``layer.weight = qweight.t()`` so that the Cutlass FP8 GEMM kernel sees its
+B operand as column-major ``[K, N]`` (the TN layout required by Hopper FP8
+``wgmma`` instructions). The resulting tensor is a non-contiguous transpose
+view of a ``(out_features, in_features)`` row-major storage.
+
+FSDP2 ``fully_shard`` rejects non-contiguous parameters because dim-0 sharding
+on a column-major view cannot be a contiguous memcpy. The two requirements are
+fundamentally at odds with the same physical buffer.
+
+This module reconciles them by separating the views: keep the parameter as
+the underlying ``(out, in)`` row-major contiguous storage (FSDP-friendly),
+and inject the equivalent ``.t()`` at the GEMM call site so the Cutlass
+kernel still receives a column-major B (zero-copy stride flip).
+
+The transpose is injected inside ``ScaledMMLinearKernel._get_layer_params``,
+which is the single place where ``apply_weights`` reads the weight tensor.
+The rest of ``apply_weights`` -- including its ``output_shape = w.shape[1]``
+computation and the eventual ``apply_scaled_mm(B=w, ...)`` call -- then
+operates on a tensor whose shape ``(K, N)`` and stride ``(1, K)`` match
+what the upstream FP8 kernel expects. No other upstream code is overridden,
+which keeps us robust against future changes inside ``apply_weights``.
+"""
+
+from __future__ import annotations
+
+import types
+
+from torch import nn
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
+
+logger = init_logger(__name__)
+
+
+def _build_transposed_get_layer_params(original_bound_method):
+    """Return a ``_get_layer_params`` replacement that transposes ``w``.
+
+    ``original_bound_method`` is the kernel instance's existing
+    ``_get_layer_params`` (already bound to ``self``); we close over it and
+    apply ``.t()`` to the returned weight, leaving scales untouched.
+    """
+
+    def _get_layer_params(self, layer):
+        w, w_s, x_s, x_s_ub = original_bound_method(layer)
+        return w.t(), w_s, x_s, x_s_ub
+
+    return _get_layer_params
+
+
+def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int:
+    """Make online-FP8 linear layers in ``model`` FSDP2-compatible.
+
+    For every layer whose ``quant_method`` is an :class:`Fp8OnlineLinearMethod`
+    and whose weight is currently a non-contiguous transpose view, this
+    function:
+
+    1. Replaces ``layer.weight`` with the underlying ``(out, in)`` row-major
+       contiguous storage so FSDP2 ``fully_shard`` accepts it.
+    2. Patches the per-layer GEMM kernel's bound ``_get_layer_params`` method
+       to return a ``.t()`` view of the weight, so ``apply_weights`` and the
+       downstream ``apply_scaled_mm`` continue to see a column-major
+       ``(in, out)`` B with zero copies.
+
+    Layers whose weight is already contiguous (e.g. Marlin FP8, offline-
+    quantized checkpoints) or that use a different quant method are skipped.
+
+    Returns:
+        Number of layers rewritten.
+    """
+    n_patched = 0
+    patched_kernel_ids: set[int] = set()
+    for module in model.modules():
+        qm = getattr(module, "quant_method", None)
+        if not isinstance(qm, Fp8OnlineLinearMethod):
+            continue
+
+        weight = getattr(module, "weight", None)
+        if weight is None or weight.is_contiguous():
+            continue
+
+        # ``weight`` here is ``qweight.t()`` (a non-contiguous (in, out) view
+        # of an (out, in) row-major qweight storage). ``weight.t()`` recovers
+        # that storage; ``.contiguous()`` is a zero-copy alias when the source
+        # is already row-major contiguous, but defensively materializes if
+        # upstream ever stacks another view.
+        contig = weight.data.t().contiguous()
+        new_param = nn.Parameter(contig, requires_grad=False)
+        module.weight = new_param
+
+        kernel = qm.fp8_linear
+        # Each linear layer is expected to have its own kernel instance, but
+        # guard against shared instances to avoid stacking multiple ``.t()``
+        # patches on the same object (which would compose into identity).
+        if id(kernel) not in patched_kernel_ids:
+            original_get = kernel._get_layer_params
+            kernel._get_layer_params = types.MethodType(_build_transposed_get_layer_params(original_get), kernel)
+            patched_kernel_ids.add(id(kernel))
+
+        n_patched += 1
+
+    if n_patched:
+        logger.info(
+            "Rewrote %d FP8 linear layer(s) into FSDP2-compatible storage; "
+            "transpose to column-major B is now applied at GEMM time.",
+            n_patched,
+        )
+    return n_patched
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index 996e74b32c9..5ab84702452 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -156,6 +156,11 @@
         "pipeline_internvla_a1",
         "InternVLAA1Pipeline",
     ),
+    "Gr00tN1d7Pipeline": (
+        "gr00t",
+        "pipeline_gr00t",
+        "Gr00tN1d7Pipeline",
+    ),
     "LongCatImageEditPipeline": (
         "longcat_image",
         "pipeline_longcat_image_edit",
diff --git a/vllm_omni/diffusion/utils/flow_matching.py b/vllm_omni/diffusion/utils/flow_matching.py
new file mode 100644
index 00000000000..e81d405ec23
--- /dev/null
+++ b/vllm_omni/diffusion/utils/flow_matching.py
@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import torch
+from torch import nn
+
+
+def swish(x: torch.Tensor) -> torch.Tensor:
+    return x * torch.sigmoid(x)
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    def __init__(self, embedding_dim: int) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        timesteps = timesteps.float()
+        B, T = timesteps.shape
+        device = timesteps.device
+
+        half_dim = self.embedding_dim // 2
+        exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * (math.log(10000.0) / half_dim)
+        freqs = timesteps.unsqueeze(-1) * exponent.exp()  # (B, T, half_dim)
+        enc = torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1)  # (B, T, embedding_dim)
+        return enc
diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
index 278f1df40f5..288973637b7 100644
--- a/vllm_omni/diffusion/worker/diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py
@@ -168,8 +168,17 @@ def get_memory_context():
         # Apply torch.compile if not in eager mode
         if not self.od_config.enforce_eager:
             if current_omni_platform.supports_torch_inductor():
-                self._compile_transformer("transformer")
-                self._compile_transformer("transformer_2")
+                if hasattr(self.pipeline, "setup_compile"):
+                    try:
+                        self.pipeline.setup_compile()
+                    except Exception as exc:
+                        logger.warning(
+                            "Model runner: setup_compile() failed (%s); running without compile.",
+                            exc,
+                        )
+                else:
+                    self._compile_transformer("transformer")
+                    self._compile_transformer("transformer_2")
             else:
                 logger.warning(
                     "Model runner: Platform %s does not support torch inductor, skipping torch.compile.",
@@ -306,8 +315,15 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
                 # num_inference_steps at all (i.e., just resets state and clears
                 # stale residuals).
                 num_inference_steps = req.sampling_params.num_inference_steps
-                if self.od_config.cache_backend == "tea_cache" and num_inference_steps is None:
-                    num_inference_steps = 0
+                if num_inference_steps is None and self.od_config.cache_backend in (
+                    "tea_cache",
+                    "step_cache",
+                ):
+                    # TeaCache refresh ignores the value; step_cache refresh is a
+                    # no-op (per-chunk state resets in the denoise loop). DreamZero often
+                    # leaves sampling_params.num_inference_steps unset and uses the
+                    # pipeline default instead.
+                    num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0
 
                 if num_inference_steps is not None:
                     self.cache_backend.refresh(self.pipeline, num_inference_steps)
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index f2ad13d5076..56a9793967d 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -161,6 +161,12 @@ def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDi
         return vllm_config
 
 
+def _get_cumem_allocator_class() -> type:
+    from vllm.device_allocator.cumem import CuMemAllocator
+
+    return CuMemAllocator
+
+
 def _resolve_ir_op_priority(od_config: OmniDiffusionConfig, vllm_config: VllmConfig) -> Any:
     ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config)
     ir_op_priority_func = get_diffusion_ir_op_priority_func(od_config)
@@ -460,8 +466,7 @@ def sleep(self, level: int = 1) -> bool:
         Args:
             level: Sleep level. Level 1 offloads weights, level 2 also saves buffers.
         """
-        from vllm.device_allocator.cumem import CuMemAllocator
-
+        CuMemAllocator = _get_cumem_allocator_class()
         allocator = CuMemAllocator.get_instance()
 
         usage_before = allocator.get_current_usage()
@@ -513,8 +518,7 @@ def wake_up(self, tags: list[str] | None = None) -> bool:
             tags: List of memory pool tags to re-activate (e.g., ["weights"]
                   to match Level 1 sleep). If None, all pools are re-activated.
         """
-        from vllm.device_allocator.cumem import CuMemAllocator
-
+        CuMemAllocator = _get_cumem_allocator_class()
         allocator = CuMemAllocator.get_instance()
         allocator.wake_up(tags)
         current_omni_platform.synchronize()
@@ -636,8 +640,7 @@ def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
         if is_sleep_enabled:
             current_omni_platform.synchronize()
             gc.collect()
-            from vllm.device_allocator.cumem import CuMemAllocator
-
+            CuMemAllocator = _get_cumem_allocator_class()
             allocator = CuMemAllocator.get_instance()
             if tag == "weights":
                 assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index 721e58c7b7c..ac6140291f1 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -520,7 +520,7 @@ def _ensure_stage_replica_mm_uuids(
             return
 
         mm_data = prompt.get("multi_modal_data")
-        if not isinstance(mm_data, Mapping) or not mm_data:
+        if not isinstance(mm_data, dict) or not mm_data:
             return
 
         from vllm.multimodal.hasher import MultiModalHasher
@@ -842,6 +842,14 @@ def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | Non
                 "mag_max_skip_steps": 5,
                 "mag_retention_ratio": 0.1,
             }
+        if cache_backend in ("step_cache"):
+            return {
+                "step_cache_dit_enabled": True,
+                "velocity_sim_thresholds": [0.95, 0.93],
+                "velocity_skip_countdowns": [4, 2],
+                "step_cache_dit_min_history": 2,
+                "step_cache_dit_max_history": 2,
+            }
         return None
 
     @staticmethod
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 905ff4eded8..984706ae007 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -512,7 +512,7 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu
             "--cache-backend",
             type=str,
             default="none",
-            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache'",
+            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'",
         )
         omni_config_group.add_argument(
             "--cache-config",
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 627a1070524..35b9dc20fac 100755
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -137,7 +137,7 @@
     build_stage_sampling_params_list,
     get_default_sampling_params_list,
 )
-from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
+from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER, FileStorageHandle
 from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url, decode_input_reference
@@ -474,6 +474,9 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         await omni_init_app_state(engine_client, app.state, args)
 
+        # Start background processes
+        await STORAGE_MANAGER.start()
+
         # Conditionally register profiler endpoints based on stage YAML configs
         stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
         if _should_enable_profiler_endpoints(stage_configs):
@@ -2690,12 +2693,11 @@ def _video_error_from_exception(exc: Exception) -> VideoError:
     )
 
 
-def _cleanup_video(video_id: str, output_path: str | None):
+async def _cleanup_video(video_id: str):
     try:
-        if output_path is not None:
-            os.remove(output_path)
-    except OSError:
-        logger.warning("Failed to cleanup partial video file '%s' for id=%s", output_path, video_id)
+        await STORAGE_MANAGER.delete(video_id)
+    except Exception:
+        logger.warning("Failed to cleanup partial video file '%s'", video_id)
 
 
 async def _run_video_generation_job(
@@ -2714,7 +2716,6 @@ async def _run_video_generation_job(
 
     await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS})
     started_at = time.perf_counter()
-    output_path = None
     try:
         video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes(
             request,
@@ -2724,27 +2725,27 @@ async def _run_video_generation_job(
             reference_audio=reference_audio,
         )
 
-        file_name = f"{video_id}.{job.file_extension}"
-        output_path = await STORAGE_MANAGER.save(video_bytes, file_name)
-        logger.info("Video request %s persisted %s output file.", video_id, output_path)
+        save_context = await STORAGE_MANAGER.save(video_bytes, video_id)
+        logger.info("Video request %s persisted %s output file.", video_id, save_context.key)
+
+        updated_fields = {
+            "status": VideoGenerationStatus.COMPLETED,
+            "progress": 100,
+            "file_name": f"{video_id}.{job.file_extension}",
+            "completed_at": save_context.created_at,
+            "inference_time_s": time.perf_counter() - started_at,
+            "stage_durations": stage_durations,
+            "peak_memory_mb": peak_memory_mb,
+            "action": action,
+        }
+        if save_context.expires_at is not None:
+            updated_fields["expires_at"] = save_context.expires_at
 
-        await VIDEO_STORE.update_fields(
-            video_id,
-            {
-                "status": VideoGenerationStatus.COMPLETED,
-                "progress": 100,
-                "file_name": file_name,
-                "completed_at": int(time.time()),
-                "inference_time_s": time.perf_counter() - started_at,
-                "stage_durations": stage_durations,
-                "peak_memory_mb": peak_memory_mb,
-                "action": action,
-            },
-        )
+        await VIDEO_STORE.update_fields(video_id, updated_fields)
     except (EngineGenerateError, EngineDeadError) as exc:
         logger.exception("Video generation failed (engine error) for id=%s", video_id)
 
-        _cleanup_video(video_id, output_path)
+        await _cleanup_video(video_id)
         await VIDEO_STORE.update_fields(
             video_id,
             {
@@ -2764,7 +2765,7 @@ async def _run_video_generation_job(
     except Exception as exc:
         logger.exception("Video generation failed for id=%s", video_id)
 
-        _cleanup_video(video_id, output_path)
+        await _cleanup_video(video_id)
         await VIDEO_STORE.update_fields(
             video_id,
             {
@@ -2775,7 +2776,7 @@ async def _run_video_generation_job(
             },
         )
     except asyncio.CancelledError:
-        _cleanup_video(video_id, output_path)
+        await _cleanup_video(video_id)
         await VIDEO_STORE.pop(video_id)
         raise
     finally:
@@ -3163,7 +3164,7 @@ async def delete_video(video_id: str) -> VideoDeleteResponse:
     elif job.status is VideoGenerationStatus.FAILED:
         if job.file_name is not None:
             try:
-                await STORAGE_MANAGER.delete(job.file_name)
+                await STORAGE_MANAGER.delete(video_id)
             except Exception:
                 logger.warning("Failed to delete stored artifact for failed video job %s", video_id, exc_info=True)
 
@@ -3173,13 +3174,13 @@ async def delete_video(video_id: str) -> VideoDeleteResponse:
     if job.file_name is None:
         raise HTTPException(status_code=409, detail="Video output not yet available. Please try again later.")
 
-    await STORAGE_MANAGER.delete(job.file_name)
+    await STORAGE_MANAGER.delete(video_id)
     await VIDEO_STORE.pop(video_id)
     return VideoDeleteResponse(id=job.id, deleted=True)
 
 
 @router.get("/v1/videos/{video_id}/content")
-async def download_video(video_id: str) -> FileResponse:
+async def download_video(video_id: str) -> Response:
     """Download the generated file for a completed video job.
 
     Args:
@@ -3202,11 +3203,19 @@ async def download_video(video_id: str) -> FileResponse:
     if not job.file_name:
         raise HTTPException(status_code=404, detail="Generation is still in-progress")
 
-    full_path = STORAGE_MANAGER.get_full_file_path(job.file_name)
-    if not os.path.exists(full_path):
+    file_handle = await STORAGE_MANAGER.open(video_id)
+    if file_handle is None:
         raise HTTPException(status_code=404, detail="Generated video file not found on disk")
 
-    return FileResponse(path=full_path, media_type=job.media_type, filename=job.file_name)
+    file_name = job.file_name or f"{video_id}.{job.file_extension}"
+    if isinstance(file_handle, FileStorageHandle):
+        response = FileResponse(path=file_handle.path, media_type=job.media_type, filename=file_name)
+    else:
+        raise HTTPException(
+            status_code=500, detail=f"Server generated an unsupported file storage handle for file id {video_id}"
+        )
+
+    return response
 
 
 @profiler_router.post("/start_profile")
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index 1f7a6b79dfc..f867bfadb45 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/vllm_omni/entrypoints/openai/serving_speech.py
@@ -672,7 +672,6 @@ def _detect_tts_model_type(self) -> str | None:
         if model_stage in _OMNIVOICE_TTS_MODEL_STAGES:
             return "omnivoice"
         if model_stage in _COVO_AUDIO_MODEL_STAGES:
-            model_arch = getattr(self._tts_stage.engine_args, "model_arch", None)
             if model_arch and "CovoAudio" in model_arch:
                 return "covo_audio"
         if model_stage in _VOXCPM2_TTS_MODEL_STAGES:
diff --git a/vllm_omni/entrypoints/openai/storage.py b/vllm_omni/entrypoints/openai/storage.py
index 52aef8948f1..a36688956e1 100644
--- a/vllm_omni/entrypoints/openai/storage.py
+++ b/vllm_omni/entrypoints/openai/storage.py
@@ -1,32 +1,78 @@
 import asyncio
+import contextlib
 import os
+import stat
+import time
 from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
 from tempfile import NamedTemporaryFile
+from typing import Generic, Literal, TypeVar
 
+from vllm.logger import init_logger
 
-class StorageBaseManager(ABC):
+from vllm_omni.config.server_settings import SERVER_SETTINGS_CONFIG, FileBackend
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class SaveContext:
+    key: str
+    created_at: int
+    expires_at: int | None = None
+
+
+@dataclass(frozen=True)
+class BaseStorageHandle:
+    kind: str
+
+
+@dataclass(frozen=True)
+class FileStorageHandle(BaseStorageHandle):
+    path: str
+    kind: Literal["path"] = field(default="path", init=False)
+
+
+K = TypeVar("K", bound=BaseStorageHandle, covariant=True)
+
+
+class StorageBaseManager(Generic[K], ABC):
     @abstractmethod
-    async def save(self, *args, **kwargs) -> str:
+    async def save(self, *args, **kwargs) -> SaveContext:
         pass
 
     @abstractmethod
     async def delete(self, *args, **kwargs) -> bool:
         pass
 
+    async def start(self, *args, **kwargs) -> None:
+        pass
+
+    async def stop(self, *args, **kwargs) -> None:
+        pass
+
+    @abstractmethod
+    async def open(self, storage_key: str) -> K | None:
+        pass
+
 
-class LocalStorageManager(StorageBaseManager):
-    def __init__(self, storage_path: str | None = None):
-        if storage_path is None:
-            storage_path = os.getenv("VLLM_OMNI_STORAGE_PATH", "/tmp/storage")
-        self.storage_path = storage_path
+class LocalStorageManager(StorageBaseManager[FileStorageHandle]):
+    def __init__(self, storage_path: str, max_concurrency: int = 4):
+        self.storage_path = os.path.realpath(storage_path)
         os.makedirs(self.storage_path, exist_ok=True)
 
-        max_concurrency = int(os.getenv("VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "4"))
         self._io_semaphore = asyncio.Semaphore(max(1, max_concurrency))
 
-    def _save_sync(self, data: bytes, file_name: str) -> str:
+    async def open(self, storage_key: str) -> FileStorageHandle | None:
+        local_file = self.get_full_file_path(storage_key)
+        if not os.path.exists(local_file):
+            return None
+        return FileStorageHandle(path=local_file)
+
+    def _save_sync(self, data: bytes, file_name: str) -> SaveContext:
         filename = self.get_full_file_path(file_name)
         tmp_name: str | None = None
+
         try:
             with NamedTemporaryFile("wb", dir=self.storage_path, delete=False) as f:
                 tmp_name = f.name
@@ -34,7 +80,8 @@ def _save_sync(self, data: bytes, file_name: str) -> str:
                 f.flush()
                 os.fsync(f.fileno())
             os.replace(tmp_name, filename)
-            return filename
+            response = SaveContext(key=file_name, created_at=int(time.time()))
+            return response
         except Exception:
             if tmp_name is not None:
                 try:
@@ -43,7 +90,7 @@ def _save_sync(self, data: bytes, file_name: str) -> str:
                     pass
             raise
 
-    async def save(self, data: bytes, file_name: str) -> str:
+    async def save(self, data: bytes, file_name: str) -> SaveContext:
         async with self._io_semaphore:
             return await asyncio.to_thread(self._save_sync, data, file_name)
 
@@ -62,9 +109,102 @@ def exists(self, file_name: str) -> bool:
         return os.path.exists(self.get_full_file_path(file_name))
 
     def get_full_file_path(self, file_name: str) -> str:
-        return os.path.join(self.storage_path, file_name)
+        full = os.path.realpath(os.path.join(self.storage_path, file_name))
+
+        if os.path.commonpath([self.storage_path, full]) != self.storage_path:
+            raise ValueError(f"Illegal storage key: {file_name!r}")
+
+        return full
+
+
+class LocalStorageTTLManager(LocalStorageManager):
+    def __init__(self, ttl_seconds: int, sweep_interval_seconds: int, *args, **kwargs):
+        if ttl_seconds <= 0:
+            raise ValueError("`ttl_seconds` must be greater than or equal to 1.")
+        if sweep_interval_seconds <= 0:
+            raise ValueError("`sweep_interval_seconds` must be greater than or equal to 1.")
+
+        self._ttl_seconds = ttl_seconds
+        self._sweep_interval_seconds = sweep_interval_seconds
+        self._sweeper_task: asyncio.Task[None] | None = None
+
+        super().__init__(*args, **kwargs)
+
+    async def save(self, data: bytes, file_name: str) -> SaveContext:
+        result = await super().save(data, file_name)
+        result.expires_at = result.created_at + self._ttl_seconds
+        return result
+
+    async def _sweep_once(self, cutoff: float) -> int:
+        expired: list[str] = []
+        for entry in os.scandir(self.storage_path):
+            try:
+                if not entry.is_file(follow_symlinks=False):
+                    continue
+                if entry.stat(follow_symlinks=False).st_mtime < cutoff:
+                    expired.append(entry.path)
+            except FileNotFoundError:
+                logger.debug("TTL sweep skipped %s; file removed during scan", entry.path)
+            except OSError:
+                logger.warning("TTL sweep failed to inspect %s", entry.path, exc_info=True)
+
+        deleted = 0
+        for path in expired:
+            try:
+                async with self._io_semaphore:
+                    st = await asyncio.to_thread(os.stat, path, follow_symlinks=False)
+                    if not stat.S_ISREG(st.st_mode):
+                        continue
+                    if st.st_mtime >= cutoff:
+                        continue
+                    await asyncio.to_thread(os.remove, path)
+                    deleted += 1
+            except FileNotFoundError:
+                logger.debug("TTL sweep skipped %s; file already removed", path)
+            except OSError:
+                logger.warning("TTL sweep failed to delete expired file %s", path, exc_info=True)
+
+        return deleted
+
+    async def _sweep_loop(self) -> None:
+        while True:
+            try:
+                cutoff = time.time() - self._ttl_seconds
+                await self._sweep_once(cutoff)
+            except Exception:
+                logger.exception("TTL sweep failed for storage path %s", self.storage_path)
+            await asyncio.sleep(self._sweep_interval_seconds)
+
+    async def start(self) -> None:
+        if self._sweeper_task is None or self._sweeper_task.done():
+            self._sweeper_task = asyncio.create_task(self._sweep_loop())
+
+    async def stop(self) -> None:
+        if self._sweeper_task is None:
+            return
+        self._sweeper_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._sweeper_task
+        self._sweeper_task = None
+
+
+def get_storage_manager(storage_config: FileBackend) -> StorageBaseManager[FileStorageHandle]:
+    if isinstance(storage_config, FileBackend):
+        if storage_config.file_ttl is not None and storage_config.ttl_sweep_interval is not None:
+            manager = LocalStorageTTLManager(
+                storage_path=storage_config.path,
+                max_concurrency=storage_config.file_concurrency,
+                ttl_seconds=storage_config.file_ttl,
+                sweep_interval_seconds=storage_config.ttl_sweep_interval,
+            )
+        else:
+            manager = LocalStorageManager(
+                storage_path=storage_config.path, max_concurrency=storage_config.file_concurrency
+            )
+    else:
+        raise ValueError("No supported storage managers")
+
+    return manager
 
 
-# Should implement TTL management on local file storage
-# and recovery on restart in the *STORE's
-STORAGE_MANAGER = LocalStorageManager()
+STORAGE_MANAGER = get_storage_manager(SERVER_SETTINGS_CONFIG.storage)
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py
index b6d3d3e9f85..ad753a91ed2 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/base.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py
@@ -194,6 +194,3 @@ def extra_body_params(cls) -> frozenset[str]:
     "SpeechServingContext",
     "TTSModelAdapter",
 ]
-
-# Type alias kept for readability in adapter signatures.
-NormalizeFn = Callable[["OpenAICreateSpeechRequest"], None]
diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py
index 69f722ac824..b6526458363 100644
--- a/vllm_omni/entrypoints/openpi/connection.py
+++ b/vllm_omni/entrypoints/openpi/connection.py
@@ -8,7 +8,13 @@
     Infer    -> client sends msgpack(obs), server sends msgpack(ndarray)
     Reset    -> client sends msgpack({endpoint:reset}), server sends msgpack(status)
 
-NumPy values use the msgpack-numpy marker mapping:
+NumPy values use msgpack-numpy marker mappings. Outbound payloads follow the
+openpi-client format so official OpenPI clients can decode action responses:
+    ndarray -> {__ndarray__: true, dtype, shape, data}
+    scalar  -> {__npgeneric__: true, dtype, data}
+
+Inbound payloads accept both the openpi-client markers above and the legacy
+vLLM-native markers:
     ndarray -> {nd: true, type, kind, shape, data}
     scalar  -> {nd: false, type, kind, data}
 """
@@ -41,18 +47,16 @@ def _pack_numpy(obj: Any) -> Any:
         if not obj.flags.c_contiguous:
             obj = np.ascontiguousarray(obj)
         return {
-            b"nd": True,
+            b"__ndarray__": True,
             b"data": obj.tobytes(),
-            b"type": obj.dtype.str,
-            b"kind": obj.dtype.kind,
+            b"dtype": obj.dtype.str,
             b"shape": obj.shape,
         }
     if isinstance(obj, np.generic):
         return {
-            b"nd": False,
-            b"data": obj.tobytes(),
-            b"type": obj.dtype.str,
-            b"kind": obj.dtype.kind,
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
         }
     raise TypeError(f"Unsupported type: {type(obj)!r}")
 
@@ -67,27 +71,67 @@ def _decode_marker_text(value: Any) -> str:
     return str(value)
 
 
+def _is_truthy_marker(value: Any) -> bool:
+    return value is True or value == 1
+
+
+def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any:
+    nd = _mapping_get(obj, "nd", _MISSING)
+    dtype = _mapping_get(obj, "type", _MISSING)
+    kind = _mapping_get(obj, "kind", _MISSING)
+    data = _mapping_get(obj, "data", _MISSING)
+    if nd is _MISSING or dtype is _MISSING or kind is _MISSING or data is _MISSING:
+        return _MISSING
+
+    dtype_obj = np.dtype(_decode_marker_text(dtype))
+    kind_text = _decode_marker_text(kind)
+    if dtype_obj.kind != kind_text:
+        raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}")
+    if dtype_obj.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {dtype_obj}")
+
+    array = np.frombuffer(data, dtype=dtype_obj).copy()
+    if nd:
+        shape = _mapping_get(obj, "shape", _MISSING)
+        if shape is _MISSING:
+            raise ValueError("NumPy ndarray marker is missing shape")
+        return array.reshape(tuple(shape))
+    return array[0]
+
+
+def _decode_openpi_numpy_marker(obj: dict[Any, Any]) -> Any:
+    if _is_truthy_marker(_mapping_get(obj, "__ndarray__", False)):
+        data = _mapping_get(obj, "data", _MISSING)
+        dtype = _mapping_get(obj, "dtype", _MISSING)
+        shape = _mapping_get(obj, "shape", _MISSING)
+        if data is _MISSING or dtype is _MISSING or shape is _MISSING:
+            raise ValueError("OpenPI ndarray marker is missing required fields")
+
+        dtype_obj = np.dtype(_decode_marker_text(dtype))
+        if dtype_obj.kind in ("V", "O", "c"):
+            raise ValueError(f"Unsupported dtype: {dtype_obj}")
+        array = np.frombuffer(data, dtype=dtype_obj).copy()
+        return array.reshape(tuple(shape))
+
+    if _is_truthy_marker(_mapping_get(obj, "__npgeneric__", False)):
+        data = _mapping_get(obj, "data", _MISSING)
+        dtype = _mapping_get(obj, "dtype", _MISSING)
+        if data is _MISSING or dtype is _MISSING:
+            raise ValueError("OpenPI scalar marker is missing required fields")
+        dtype_obj = np.dtype(_decode_marker_text(dtype))
+        if dtype_obj.kind in ("V", "O", "c"):
+            raise ValueError(f"Unsupported dtype: {dtype_obj}")
+        return dtype_obj.type(data)
+
+    return _MISSING
+
+
 def _unpack_numpy(obj: Any) -> Any:
     if isinstance(obj, dict):
-        nd = _mapping_get(obj, "nd", _MISSING)
-        dtype = _mapping_get(obj, "type", _MISSING)
-        kind = _mapping_get(obj, "kind", _MISSING)
-        data = _mapping_get(obj, "data", _MISSING)
-        if nd is not _MISSING and dtype is not _MISSING and kind is not _MISSING and data is not _MISSING:
-            dtype_obj = np.dtype(_decode_marker_text(dtype))
-            kind_text = _decode_marker_text(kind)
-            if dtype_obj.kind != kind_text:
-                raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}")
-            if dtype_obj.kind in ("V", "O", "c"):
-                raise ValueError(f"Unsupported dtype: {dtype_obj}")
-
-            array = np.frombuffer(data, dtype=dtype_obj).copy()
-            if nd:
-                shape = _mapping_get(obj, "shape", _MISSING)
-                if shape is _MISSING:
-                    raise ValueError("NumPy ndarray marker is missing shape")
-                return array.reshape(tuple(shape))
-            return array[0]
+        for decoder in (_decode_vllm_numpy_marker, _decode_openpi_numpy_marker):
+            decoded = decoder(obj)
+            if decoded is not _MISSING:
+                return decoded
         return {key: _unpack_numpy(value) for key, value in obj.items()}
     if isinstance(obj, list):
         return [_unpack_numpy(value) for value in obj]
diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
index 1a387cdcc39..d32ea62f284 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
@@ -479,14 +479,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             # Keep additional information synchronized for async_chunk updates.
             self.enable_update_additional_information = True
 
-            # Expose streaming parameters
-            self.token_overlap_len = self.code2wav.token_overlap_len
-            self.mel_overlap_len = self.code2wav.mel_overlap_len
-            self.mel_window = self.code2wav.mel_window
-            self.mel_cache_len = self.code2wav.mel_cache_len
-            self.source_cache_len = self.code2wav.source_cache_len
-            self.speech_window = self.code2wav.speech_window
-            self._stream_audio_cache_by_req: dict[str, torch.Tensor] = {}
             self._stream_audio_cache_lock = Lock()
             self._stream_vocoder_cache_by_req: dict[str, dict[str, torch.Tensor]] = {}
         else:
@@ -511,27 +503,6 @@ def _create_llm_vllm_config(self, parent_config: VllmConfig) -> VllmConfig:
         # Use parent's cache config - critical for PagedAttention to work correctly
         return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"])
 
-    @staticmethod
-    def _cross_fade_audio(audio: torch.Tensor, prev_tail: torch.Tensor) -> torch.Tensor:
-        """Blend previous chunk tail into current chunk head using a Hamming window.
-
-        This mirrors upstream CosyVoice's `fade_in_out(...)` semantics:
-        update the current head in-place using a 2*overlap window, then
-        concatenate the unchanged remainder.
-        """
-        if audio.numel() == 0 or prev_tail.numel() == 0:
-            return audio
-        overlap = min(int(audio.numel()), int(prev_tail.numel()))
-        if overlap <= 0:
-            return audio
-        window = torch.hamming_window(2 * overlap, periodic=False, dtype=audio.dtype, device=audio.device)
-        fade_in = window[:overlap]
-        fade_out = window[overlap:]
-        blended = audio[:overlap] * fade_in + prev_tail[-overlap:].to(device=audio.device, dtype=audio.dtype) * fade_out
-        if overlap == int(audio.numel()):
-            return blended
-        return torch.cat([blended, audio[overlap:]], dim=0)
-
     def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_finished: bool) -> torch.Tensor:
         """Pass-through stitching for async_chunk.
 
@@ -539,11 +510,9 @@ def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_f
         Applying an additional waveform-domain fade/cache step introduces either
         duplicated overlap (if no tail trim) or duration shrink (if tail trim).
         """
-        if req_id is not None and stream_finished and hasattr(self, "_stream_audio_cache_by_req"):
+        if req_id is not None and stream_finished and hasattr(self, "_stream_vocoder_cache_by_req"):
             with self._stream_audio_cache_lock:
-                self._stream_audio_cache_by_req.pop(req_id, None)
-                if hasattr(self, "_stream_vocoder_cache_by_req"):
-                    self._stream_vocoder_cache_by_req.pop(req_id, None)
+                self._stream_vocoder_cache_by_req.pop(req_id, None)
         return audio
 
     @staticmethod
diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
index cb3228c13a7..49e8e515b25 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
@@ -11,8 +11,6 @@
 
 from __future__ import annotations
 
-from contextlib import nullcontext
-
 import numpy as np
 import torch
 import torch.nn as nn
@@ -175,19 +173,18 @@ def _forward_mel(
         prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32)
         prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32)
 
-        with nullcontext():
-            feat, _ = self.flow_model.inference(
-                token=token,
-                token_len=token_len,
-                prompt_token=prompt_token,
-                prompt_token_len=prompt_token_len,
-                prompt_feat=prompt_feat,
-                prompt_feat_len=prompt_feat_len,
-                embedding=embedding,
-                streaming=streaming,
-                finalize=finalize,
-                n_timesteps=n_timesteps,
-            )
+        feat, _ = self.flow_model.inference(
+            token=token,
+            token_len=token_len,
+            prompt_token=prompt_token,
+            prompt_token_len=prompt_token_len,
+            prompt_feat=prompt_feat,
+            prompt_feat_len=prompt_feat_len,
+            embedding=embedding,
+            streaming=streaming,
+            finalize=finalize,
+            n_timesteps=n_timesteps,
+        )
 
         trim_mel = max(0, int(token_offset_tokens)) * int(self.token_mel_ratio)
         if trim_mel > 0:
@@ -195,28 +192,6 @@ def _forward_mel(
 
         return feat
 
-    @staticmethod
-    def _fade_speech(
-        speech: torch.Tensor,
-        prev_speech: torch.Tensor,
-    ) -> torch.Tensor:
-        """Blend previous speech tail into current speech head."""
-        if speech.numel() == 0 or prev_speech.numel() == 0:
-            return speech
-        overlap = min(int(speech.shape[-1]), int(prev_speech.shape[-1]))
-        if overlap <= 0:
-            return speech
-        window = torch.hamming_window(2 * overlap, periodic=False, dtype=speech.dtype, device=speech.device)
-        fade_in = window[:overlap].view(1, -1)
-        fade_out = window[overlap:].view(1, -1)
-        blended_head = (
-            speech[:, :overlap] * fade_in
-            + prev_speech[:, -overlap:].to(device=speech.device, dtype=speech.dtype) * fade_out
-        )
-        if overlap == int(speech.shape[-1]):
-            return blended_head
-        return torch.cat([blended_head, speech[:, overlap:]], dim=-1)
-
     @torch.inference_mode()
     def forward_streaming(
         self,
@@ -238,17 +213,16 @@ def forward_streaming(
         suffix. That preserves causal look-right handling without double
         trimming or duplicated overlap at chunk boundaries.
         """
-        with nullcontext():
-            feat = self._forward_mel(
-                token=token,
-                prompt_token=prompt_token,
-                prompt_feat=prompt_feat,
-                embedding=embedding,
-                n_timesteps=n_timesteps,
-                token_offset_tokens=token_offset_tokens,
-                streaming=True,
-                finalize=finalize,
-            )
+        feat = self._forward_mel(
+            token=token,
+            prompt_token=prompt_token,
+            prompt_feat=prompt_feat,
+            embedding=embedding,
+            n_timesteps=n_timesteps,
+            token_offset_tokens=token_offset_tokens,
+            streaming=True,
+            finalize=finalize,
+        )
         hift_weight = self.hift.m_source.l_linear.weight
         chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
 
@@ -268,8 +242,7 @@ def forward_streaming(
         if tts_mel.shape[-1] == 0:
             tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype)
         else:
-            with nullcontext():
-                tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
+            tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
 
         tts_speech = tts_speech.reshape(tts_speech.shape[0], -1)
         speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1])))
diff --git a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
index ed35c93bd13..32db8ef5b3b 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
@@ -4,8 +4,10 @@
 
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
 Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
-  * ``sync_process_input_func`` runs when ``deploy.async_chunk=false``:
-    stage 1 builds full-sequence flow input via ``text2flow``.
+  * ``sync_process_input_func`` (``text2flow_token_only``) runs when
+    ``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
+    from the emitted speech tokens. ``text2flow`` is the
+    ``custom_process_input_func`` variant.
   * ``async_chunk_process_next_stage_input_func`` runs when
     ``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
     through the shared-memory connector.
diff --git a/vllm_omni/model_executor/models/cosyvoice3/utils.py b/vllm_omni/model_executor/models/cosyvoice3/utils.py
index 5827e6bbbc2..d2c3c686574 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/utils.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/utils.py
@@ -13,8 +13,6 @@
 
 logger = logging.getLogger(__name__)
 
-IGNORE_ID = -1
-
 
 def dynamic_range_compression_torch(x, c=1, clip_val=1e-5):
     return torch.log(torch.clamp(x, min=clip_val) * c)
@@ -182,26 +180,6 @@ def log_mel_spectrogram(
     return log_spec
 
 
-def extract_speech_token(prompt_wav, speech_tokenizer_session, device):
-    speech = load_wav(prompt_wav, 16000)
-    assert speech.shape[1] / 16000 <= 30, "do not support extract speech token for audio longer than 30s"
-    feat = log_mel_spectrogram(speech, n_mels=128)
-    speech_token = (
-        speech_tokenizer_session.run(
-            None,
-            {
-                speech_tokenizer_session.get_inputs()[0].name: feat.detach().cpu().numpy(),
-                speech_tokenizer_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32),
-            },
-        )[0]
-        .flatten()
-        .tolist()
-    )
-    speech_token = torch.tensor([speech_token], dtype=torch.int32).to(device)
-    speech_token_len = torch.tensor([speech_token.shape[1]], dtype=torch.int32).to(device)
-    return speech_token, speech_token_len
-
-
 def extract_spk_embedding(prompt_wav, campplus_session, device):
     speech = load_wav(prompt_wav, 16000)
     feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000)
diff --git a/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py b/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py
index 8153c2f385e..b9b06857d07 100644
--- a/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py
+++ b/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py
@@ -67,6 +67,5 @@ def __init__(self, **kwargs):
         self.inference = {
             "s_steps": 10,
             "cfg_alpha": 1.0,
-            "dynamic_the": False,
             "rescale_logits": False,
         }
diff --git a/vllm_omni/model_executor/models/covo_audio/token2wav.py b/vllm_omni/model_executor/models/covo_audio/token2wav.py
index 4a6ce4b7e3a..8a96e1f398b 100644
--- a/vllm_omni/model_executor/models/covo_audio/token2wav.py
+++ b/vllm_omni/model_executor/models/covo_audio/token2wav.py
@@ -598,7 +598,7 @@ def inference(
         latent_cond = torch.zeros(b, tgt_lens, self.target_dim).to(device)
         if prefix_target is not None:
             latent_cond[:, : prefix_target.size(1), :] = prefix_target
-        sample, trajectory = self.sample(
+        sample, _ = self.sample(
             tokens=token, audio=latent_cond, steps=s_steps, alpha=cfg_alpha, g_cond=ge, rescale_logits=rescale_logits
         )
         return sample
@@ -612,9 +612,8 @@ def solver(t, z):
                 output = torch.cat([tokens, audio, z], dim=-1)
                 return self.vectorfield_forward(inputs=output, times=t.unsqueeze(0), self_attn_mask=None, g_cond=g_cond)
             tokens_empty = torch.zeros(*audio.shape[:2], self.model_dim, device=tokens.device, dtype=tokens.dtype)
-            audio_empty = audio
             tokens_t = torch.cat([tokens_empty, tokens], dim=0)
-            audio_t = torch.cat([audio_empty, audio], dim=0)
+            audio_t = torch.cat([audio, audio], dim=0)
             audio_noizy_t = torch.cat([z, z], dim=0)
             t_t = torch.stack([t, t], dim=0)
             c = g_cond
diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
index c6e9064f133..9627dc7ce13 100644
--- a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
+++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
@@ -348,18 +348,16 @@ def forward(
 
             if req_ids.numel() < 1:
                 continue
-            flat = req_ids
-            n = flat.numel()
-            if n == 0 or n % q != 0:
-                if n > 0:
-                    logger.warning(
-                        "DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.",
-                        n,
-                        q,
-                    )
+            n = req_ids.numel()
+            if n % q != 0:
+                logger.warning(
+                    "DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.",
+                    n,
+                    q,
+                )
                 continue
             frames = n // q
-            codes_qf = flat.reshape(q, frames)
+            codes_qf = req_ids.reshape(q, frames)
             parsed_ctx_frames[i] = ctx_frames
             parsed_total_frames[i] = frames
             valid_codes_qf.append(codes_qf)
diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
index 46ceee49835..eb056f16a10 100644
--- a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
+++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
@@ -191,7 +191,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class _FastARDecoderLayer(nn.Module):
-    """Transformer decoder layer for Fast AR (SDPA, no KV cache)."""
+    """Transformer decoder layer for Fast AR.
+
+    The batched ``forward`` runs SDPA over the full sequence and is cache-free,
+    while ``forward_one`` decodes a single step using a per-call KV cache.
+    """
 
     def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None:
         super().__init__()
@@ -519,7 +523,7 @@ def forward(
         embed_buf[:bsz, 0, :] = projected
 
         # Position 1: embedding of semantic code.
-        code_embed = self.fast_embeddings(semantic_code.clamp(min=0))
+        code_embed = self.fast_embeddings(semantic_code)
         embed_buf[:bsz, 1, :] = code_embed
 
         use_sampling = do_sample and temperature > 0
@@ -542,8 +546,7 @@ def forward(
             logits = self.fast_output(self.fast_norm(hidden_out))
 
             # Residual codebooks (step >= 1) only have 1024 entries.
-            if step >= 1:
-                logits = logits[:, :residual_codebook_size]
+            logits = logits[:, :residual_codebook_size]
 
             if use_sampling:
                 scaled = logits * inv_temperature
diff --git a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py
index 8b8d8559ead..3918a6f9462 100644
--- a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py
+++ b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py
@@ -132,13 +132,3 @@ def estimate_fish_voice_clone_prompt_len_from_normalized(
         [0] * semantic_len,
     )
     return len(prompt_ids)
-
-
-def estimate_fish_voice_clone_prompt_len(tokenizer: Any, text: str, ref_text: str, semantic_len: int) -> int:
-    normalized_text, normalized_ref_text = normalize_fish_voice_clone_texts(text, ref_text)
-    return estimate_fish_voice_clone_prompt_len_from_normalized(
-        tokenizer,
-        normalized_text,
-        normalized_ref_text,
-        semantic_len,
-    )
diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py
index 5cda6ec45f5..666998e9dda 100644
--- a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py
+++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py
@@ -472,10 +472,6 @@ def __init__(
         self.num_heads = heads
         self.spkr_emb_adaln = spkr_emb_adaln
 
-        # Frame rates for block-causal mask token→mel conversion
-        self._input_frame_rate: float = 25.0
-        self._mel_framerate: float = 50.0
-
         # Speaker embedding dimension for AdaLN
         if spkr_emb_adaln:
             additional_dim = spkr_dim + 256 if use_wavlm_emb else spkr_dim
diff --git a/vllm_omni/model_executor/models/gr00t/__init__.py b/vllm_omni/model_executor/models/gr00t/__init__.py
new file mode 100644
index 00000000000..208f01a7cb5
--- /dev/null
+++ b/vllm_omni/model_executor/models/gr00t/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/model_executor/models/gr00t/pipeline.py b/vllm_omni/model_executor/models/gr00t/pipeline.py
new file mode 100644
index 00000000000..38a76f29769
--- /dev/null
+++ b/vllm_omni/model_executor/models/gr00t/pipeline.py
@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GR00T N1.7 single-stage policy topology."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+GR00T_N1D7_PIPELINE = PipelineConfig(
+    model_type="Gr00tN1d7",
+    model_arch="Gr00tN1d7Pipeline",
+    hf_architectures=("Gr00tN1d7",),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="actions",
+            model_arch="Gr00tN1d7Pipeline",
+        ),
+    ),
+)
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py
index dc5eb12b475..b6eef8dc4a5 100644
--- a/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py
@@ -155,15 +155,3 @@ def __init__(
     @property
     def num_real_codes(self) -> int:
         return self.audio_stream_bos_id  # i.e. 1024 for the boson-ai checkpoint
-
-    # The full vocab size of each per-codebook output head.
-    @property
-    def codebook_output_size(self) -> int:
-        return self.codebook_size  # i.e. 1026 for the boson-ai checkpoint
-
-    # Canonical MusicGen-style delay pattern used by
-    # `HiggsAudioV2DelayPatternLogitsProcessor`: codebook k starts emitting
-    # real codes only after k frames. Kept as a function so subclasses can
-    # override if upstream ever publishes a different convention.
-    def default_delay_pattern(self) -> list[int]:
-        return list(range(self.num_codebooks))
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
index df636a4bbe2..be08163e4d2 100644
--- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
@@ -821,12 +821,12 @@ def forward(
         except Exception:
             self._last_step_query_start_loc = None
 
-        # Voice-clone ref-audio substitution: at the prompt-side audio
-        # placeholder span we overwrite the text-embed lookups with the
-        # multi-codebook embedding of the encoded reference clip carried in
-        # ``model_intermediate_buffer``. No-op for plain-text requests and
-        # for decode steps (which fall through to the audio-feedback path).
         if input_ids is not None and inputs_embeds is None:
+            # Voice-clone ref-audio substitution: at the prompt-side audio
+            # placeholder span we overwrite the text-embed lookups with the
+            # multi-codebook embedding of the encoded reference clip carried in
+            # ``model_intermediate_buffer``. No-op for plain-text requests and
+            # for decode steps (which fall through to the audio-feedback path).
             info_dicts = kwargs.get("model_intermediate_buffer")
             if info_dicts is None:
                 info_dicts = kwargs.get("runtime_additional_information")
@@ -836,15 +836,14 @@ def forward(
                 info_dicts,
             )
 
-        # Audio-feedback substitution: vLLM's AR runner calls forward()
-        # with ``inputs_embeds=None`` and lets the model embed its own
-        # input_ids inline. ``embed_input_ids`` is only invoked via the
-        # supports_mm_inputs/prompt_embeds branches, neither of which fires
-        # for the talker. So we apply the substitution directly here: at
-        # decode positions where input_ids == audio_token_id, swap the text
-        # embedding with embed_audio_codes(last audio frame stashed in
-        # ``self._audio_state[req_id]`` by :meth:`sample`).
-        if input_ids is not None and inputs_embeds is None:
+            # Audio-feedback substitution: vLLM's AR runner calls forward()
+            # with ``inputs_embeds=None`` and lets the model embed its own
+            # input_ids inline. ``embed_input_ids`` is only invoked via the
+            # supports_mm_inputs/prompt_embeds branches, neither of which fires
+            # for the talker. So we apply the substitution directly here: at
+            # decode positions where input_ids == audio_token_id, swap the text
+            # embedding with embed_audio_codes(last audio frame stashed in
+            # ``self._audio_state[req_id]`` by :meth:`sample`).
             hidden_states = self._maybe_apply_audio_feedback(hidden_states, input_ids)
 
         audio_mask = self.audio_token_mask(input_ids) if input_ids is not None else None
@@ -1506,7 +1505,7 @@ def _sample_audio_codes_upstream(
             sampled[~valid] = torch.argmax(logits[~valid].float(), dim=-1)
         return sampled
 
-    def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> int:
+    def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> None:
         """Mask non-audio tokens at audio-mode positions, in-place on ``logits``.
 
         Heuristic: per-request, find the last token seen so far (last of
@@ -1524,7 +1523,7 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -
         the stock sampler.
         """
         if logits is None or logits.ndim != 2:
-            return 0
+            return
         audio_bos = int(self.config.audio_bos_token_id)
         audio_id = int(self.config.audio_token_id)
         # (PyTorch align, corrected): upstream NEVER lets the model
@@ -1542,7 +1541,6 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -
         prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None)
         output_ids = getattr(sampling_metadata, "output_token_ids", None)
         num_rows = int(logits.shape[0])
-        biased = 0
         # Fallback "previous token" source: read each request's LAST input
         # position from ``self._last_step_input_ids`` using ``query_start_loc``
         # from the forward context. The previous version of this code took
@@ -1619,7 +1617,6 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -
                 mask[audio_eos_vocab] = row[audio_eos_vocab]
                 logits[i].copy_(mask)
                 slot_state["should_terminate"] = False
-                biased += 1
                 continue
             # Always force ``audio_token_id`` — upstream mirror.
             allowed: set[int] = {audio_id, *allowed_extra}
@@ -1629,9 +1626,7 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -
                 if 0 <= tok < row.shape[-1]:
                     mask[tok] = row[tok]
             logits[i].copy_(mask)
-            biased += 1
         self._last_first_audio_after_bos = first_after_bos
-        return biased
 
     has_postprocess: bool = True
 
diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py b/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py
deleted file mode 100644
index 2a21aafcd39..00000000000
--- a/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py
+++ /dev/null
@@ -1,111 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Bounded LRU cache for Higgs Audio v3 reference-audio codec encoding.
-
-Voice-clone requests that re-send the same ref_audio are otherwise re-encoded
-through the DAC encoder on every call. The cache keys by a stable content
-hash (size + sample rate + head/mid/tail sentinel bytes) so the same audio
-served via different URLs or repeated base64 payloads also hits warm.
-Tensors are cloned on get/put so callers and the cache cannot alias storage.
-"""
-
-from __future__ import annotations
-
-import hashlib
-import threading
-from collections import OrderedDict
-from collections.abc import Callable
-
-import numpy as np
-import torch
-
-__all__ = ["cached_encode_reference_audio", "cache_stats", "clear_cache"]
-
-_DEFAULT_MAX_ENTRIES = 128
-_SENTINEL_BYTES = 4096
-
-_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
-_cache_lock = threading.Lock()
-_hits = 0
-_misses = 0
-
-
-def _content_hash(wav: np.ndarray, sr: int) -> str:
-    """Stable content hash. Head/mid/tail sentinel keeps cost O(1)."""
-    raw = wav.tobytes()
-    n = len(raw)
-    h = hashlib.sha256()
-    h.update(n.to_bytes(8, "little"))
-    h.update(int(sr).to_bytes(4, "little"))
-    if n == 0:
-        return h.hexdigest()
-    if n <= _SENTINEL_BYTES * 3:
-        h.update(raw)
-        return h.hexdigest()
-    h.update(raw[:_SENTINEL_BYTES])
-    mid = (n - _SENTINEL_BYTES) // 2
-    h.update(raw[mid : mid + _SENTINEL_BYTES])
-    h.update(raw[-_SENTINEL_BYTES:])
-    return h.hexdigest()
-
-
-def cached_encode_reference_audio(
-    wav: np.ndarray,
-    sr: int,
-    encode_fn: Callable[[np.ndarray, int], torch.Tensor],
-    max_entries: int = _DEFAULT_MAX_ENTRIES,
-) -> torch.Tensor:
-    """LRU-cached wrapper around ``encode_reference_audio``.
-
-    ``encode_fn`` is the actual codec encoder; passed in so this module
-    stays decoupled from the tokenizer import (avoids circular imports
-    during stage-input-processor discovery).
-
-    Returns a CPU tensor clone — never the cache's interior storage.
-    """
-    global _hits, _misses
-    key = _content_hash(wav, sr)
-    with _cache_lock:
-        cached = _cache.get(key)
-        if cached is not None:
-            _hits += 1
-            _cache.move_to_end(key)
-            return cached.clone()
-
-    # Miss: encode outside the lock so concurrent misses on different
-    # refs don't serialize on this single mutex.
-    codes = encode_fn(wav, sr)
-    codes_cpu = codes.detach().to("cpu").contiguous()
-    with _cache_lock:
-        if key in _cache:
-            # Concurrent miss winner already stored it; refresh LRU position
-            # and return that copy to avoid divergent tensors per request.
-            _cache.move_to_end(key)
-            _misses += 1  # we still paid the encode cost
-            return _cache[key].clone()
-        _cache[key] = codes_cpu
-        if len(_cache) > max_entries:
-            _cache.popitem(last=False)
-        _misses += 1
-    return codes_cpu.clone()
-
-
-def cache_stats() -> dict:
-    """Return (hits, misses, size, hit_rate) — for /metrics-style reporting."""
-    with _cache_lock:
-        total = _hits + _misses
-        return {
-            "hits": _hits,
-            "misses": _misses,
-            "size": len(_cache),
-            "hit_rate": (_hits / total) if total > 0 else 0.0,
-        }
-
-
-def clear_cache() -> None:
-    """Drop all cached entries. Used by tests."""
-    global _hits, _misses
-    with _cache_lock:
-        _cache.clear()
-        _hits = 0
-        _misses = 0
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/projectors.py b/vllm_omni/model_executor/models/ming_flash_omni/projectors.py
index 42e53d1c635..03527ba6327 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/projectors.py
+++ b/vllm_omni/model_executor/models/ming_flash_omni/projectors.py
@@ -108,20 +108,6 @@ def __init__(
             layers.append(nn.Linear(llm_dim, llm_dim))
         self.proj = nn.Sequential(*layers)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Project audio features with temporal downsampling.
-
-        Args:
-            x: [B, T, audio_dim] audio encoder output (channel-last).
-
-        Returns:
-            [B, T', llm_dim] projected features (channel-last),
-            where T' = (T - ds_kernel_size + 2*(ds_kernel_size//2)) // ds_stride + 1.
-        """
-        # Conv1d expects [B, C, T], so transpose input
-        x = x.transpose(-1, -2)  # [B, audio_dim, T]
-        return self.proj(x)
-
     def forward_packed(
         self,
         packed: torch.Tensor,
@@ -155,19 +141,6 @@ def forward_packed(
         packed_proj = mlp(packed_proj)
         return packed_proj, proj_lens
 
-    def compute_output_length(self, input_length: torch.Tensor) -> torch.Tensor:
-        """Compute output sequence length after Conv1d downsampling.
-
-        Args:
-            input_length: Original mel spectrogram lengths.
-
-        Returns:
-            Output lengths after both convolutions.
-        """
-        length = (input_length - 3 + 2 * 1) // 2 + 1
-        length = (length - self.ds_kernel_size + 2 * (self.ds_kernel_size // 2)) // self.ds_stride + 1
-        return length
-
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py b/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py
index 436b92f1428..e0e8e73e669 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py
+++ b/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py
@@ -15,19 +15,6 @@
 
 logger = init_logger(__name__)
 
-# Ported from
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/toolkit.py
-_TOKENIZE_PATTERN = re.compile(r"(?:[a-zA-Z]\.)+|[a-zA-Z]+(?:['\-][a-zA-Z]+)*|\d+(?:\.\d+)?|[\u4e00-\u9fff]|\s+|\S")
-
-
-def tokenize_mixed_text(text: str) -> list[str]:
-    return re.findall(_TOKENIZE_PATTERN, text)
-
-
-def tokenize_mixed_text_iterator(text: str):
-    for match in _TOKENIZE_PATTERN.finditer(text):
-        yield match.group(0)
-
 
 # Ported from
 # https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/text_segment_cut.py
@@ -218,84 +205,6 @@ def cut_text_by_semantic_length(
     return [f.replace(_DOT_PLACEHOLDER, ".") for f in result_fragments]
 
 
-# Streaming sentence boundary detection
-
-_RE_CJK = re.compile(r"[\u4e00-\u9fff]")
-_RE_DIGIT_LAST = re.compile(r"[0-9]")
-
-
-# Left for reference for now
-# Alternative to `cut_text_by_semantic_length`
-def detect_sentence_boundaries(
-    text: str,
-    max_length: int = 50,
-) -> list[str]:
-    """Accumulate tokens and flush at sentence boundaries.
-
-    Ported from the streaming sentence detection loop from the Ming repo
-    TTS branch, but operates on the full text since we have it available upfront.
-    """
-    sentences: list[str] = []
-    streaming_text: list[str] = []
-    count = 0
-
-    for ele in tokenize_mixed_text_iterator(text):
-        if len(ele) == 0:
-            continue
-
-        should_process = False
-        min_tokens = 12 if count == 0 else 8
-
-        if ele[-1] in "！？。，!?":
-            if len(streaming_text) >= min_tokens:
-                should_process = True
-            streaming_text.append(ele)
-
-        elif ele[-1] == ".":
-            if (
-                len(streaming_text) >= min_tokens
-                and streaming_text
-                and not _RE_DIGIT_LAST.search(streaming_text[-1][-1])
-            ):
-                should_process = True
-            streaming_text.append(ele)
-
-        elif ele[-1] == "\n":
-            if streaming_text:
-                joined = "".join(streaming_text)
-                if _RE_CJK.search(joined):
-                    if _RE_CJK.search(streaming_text[-1][-1]):
-                        ele = "，"
-                        streaming_text.append(ele)
-                else:
-                    if len(ele) > 1 and re.search(r"[a-zA-Z]", ele[-2]):
-                        ele = ele[:-1] + "."
-                    else:
-                        ele = ele[:-1]
-                    streaming_text.append(ele)
-
-            if len(streaming_text) >= min_tokens:
-                should_process = True
-        else:
-            streaming_text.append(ele)
-            continue
-
-        if should_process:
-            joined = "".join(streaming_text)
-            fragments = cut_text_by_semantic_length(joined, max_length)
-            sentences.extend(fragments)
-            streaming_text = []
-            count += 1
-
-    # Flush remaining
-    if streaming_text and re.search(r"[a-zA-Z\u4e00-\u9fff1-9]", "".join(streaming_text)):
-        joined = "".join(streaming_text)
-        fragments = cut_text_by_semantic_length(joined, max_length)
-        sentences.extend(fragments)
-
-    return sentences
-
-
 # number normalization for English. Ported from
 # https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/number_en.py
 
diff --git a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
index b42c4fdca62..8ef8a3a0e36 100644
--- a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
+++ b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
@@ -30,7 +30,6 @@
     KEY_MAX_DECODE_STEPS,
     KEY_MIN_DECODE_STEPS,
     KEY_NEXT_EMBEDS,
-    KEY_PROMPT_LATENT_TAIL,
     KEY_PROMPT_LATENTS,
     KEY_REQUEST_ID,
     KEY_SIGMA,
@@ -184,7 +183,6 @@ def validate(self) -> None:
     "KEY_MAX_DECODE_STEPS",
     "KEY_MIN_DECODE_STEPS",
     "KEY_NEXT_EMBEDS",
-    "KEY_PROMPT_LATENT_TAIL",
     "KEY_PROMPT_LATENTS",
     "KEY_REQUEST_ID",
     "KEY_SIGMA",
diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py
index 2f59e209401..4b3f527da6c 100644
--- a/vllm_omni/model_executor/models/ming_tts/constants.py
+++ b/vllm_omni/model_executor/models/ming_tts/constants.py
@@ -55,7 +55,6 @@
 KEY_LAST_STOP_PROB = "ming_last_stop_prob"
 KEY_NEXT_EMBEDS = "ming_next_embeds"
 KEY_PROMPT_LATENTS = "ming_prompt_latents"
-KEY_PROMPT_LATENT_TAIL = "ming_prompt_latent_tail"
 KEY_SPEAKER_EMBEDDING = "ming_speaker_embedding"
 KEY_REQUEST_ID = "ming_request_id"
 KEY_CHUNK_ID = "ming_chunk_id"
diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts.py b/vllm_omni/model_executor/models/ming_tts/ming_tts.py
index 667aded27d7..4140b403be5 100644
--- a/vllm_omni/model_executor/models/ming_tts/ming_tts.py
+++ b/vllm_omni/model_executor/models/ming_tts/ming_tts.py
@@ -16,7 +16,6 @@
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 
 from .audio_prep import (
-    _coerce_prompt_latents,
     _find_audio_placeholder_positions,
     _find_speaker_placeholder_positions,
     _initial_history,
@@ -32,7 +31,6 @@
     KEY_MAX_DECODE_STEPS,
     KEY_MIN_DECODE_STEPS,
     KEY_NEXT_EMBEDS,
-    KEY_PROMPT_LATENT_TAIL,
     KEY_REQUEST_ID,
     KEY_SIGMA,
     KEY_SPEAKER_EMBEDDING,
@@ -175,7 +173,6 @@ def _prefill_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tenso
             dtype=torch.float32,
         )
         update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous()
-        update[KEY_PROMPT_LATENT_TAIL] = update[KEY_LATENT_HISTORY]
 
         speaker_embedding = info_dict.get(KEY_SPEAKER_EMBEDDING, info_dict.get("speaker_embedding"))
         speaker_embeddings = None
@@ -262,7 +259,6 @@ def _copy_runtime_controls(update: dict[str, Any], info_dict: dict[str, Any]) ->
 __all__ = [
     "MingTTSForConditionalGeneration",
     "_ModelSampleAdapter",
-    "_coerce_prompt_latents",
     "_find_audio_placeholder_positions",
     "_find_speaker_placeholder_positions",
     "_initial_history",
diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
index 92daad967e0..e7e59631f76 100644
--- a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
+++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
@@ -3,7 +3,6 @@
 # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py
 from __future__ import annotations
 
-import warnings
 from collections.abc import Iterable
 from typing import Any
 
@@ -346,10 +345,10 @@ def _decode_one_step(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         params_dict = dict(self.named_parameters(remove_duplicate=False))
         loaded_params = AutoWeightsLoader(self).load_weights(weights, mapper=self.hf_to_vllm_mapper)
-        _warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.", fatal=True)
-        _warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.", fatal=True)
-        _warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.", fatal=True)
-        _warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.", fatal=True)
+        _warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.")
+        _warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.")
+        _warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.")
+        _warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.")
         return loaded_params
 
 
@@ -383,7 +382,6 @@ def _warn_missing_prefix(
     params_dict: dict[str, nn.Parameter],
     loaded_params: set[str],
     prefix: str,
-    fatal: bool = False,
 ) -> None:
     missing = {key for key in params_dict if key.startswith(prefix)} - loaded_params
     if not missing:
@@ -392,9 +390,7 @@ def _warn_missing_prefix(
         f"MingLLMModel: {len(missing)} {module_name} params not loaded "
         f"(prefix={prefix}). First few: {sorted(missing)[:5]}"
     )
-    if fatal:
-        raise RuntimeError(msg)
-    warnings.warn(msg, stacklevel=3)
+    raise RuntimeError(msg)
 
 
 __all__ = [
diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py
index a2cf75b07b1..89c52485a39 100644
--- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py
+++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py
@@ -78,7 +78,7 @@ def num_quantizers(self) -> int:
 
     @property
     def codebook_size(self) -> int:
-        return int(self.quantizer_kwargs.get("codebook_size", 4096))
+        return int(self.quantizer_kwargs.get("codebook_size", 1024))
 
     @property
     def frame_rate(self) -> float:
diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
index e7e100a1893..d32af7c9020 100644
--- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
+++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
@@ -31,8 +31,6 @@
 
 logger = init_logger(__name__)
 
-_AUDIO_PAD_ID = 0  # padding id used before delay slot fires
-
 
 def _maybe_prefix(prefix: str, name: str) -> str:
     return f"{prefix}.{name}" if prefix else name
@@ -153,6 +151,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         # `make_omni_output`, which runs immediately before the sampler).
         self._batch_state: list[dict[str, Any]] | None = None
 
+        # Stacked weight caches built after load_weights(). Avoids a Python
+        # loop over n_vq heads/embeddings at every decode step.
+        # shapes: (n_vq, audio_vocab+1, hidden_size)
+        self._stacked_audio_head_w: torch.Tensor | None = None
+        self._stacked_audio_emb_w: torch.Tensor | None = None
+
     # ------------------------------------------------------------------
     # vLLM-Omni hooks
     # ------------------------------------------------------------------
@@ -394,9 +398,15 @@ def preprocess(
                     chunk = ref_codes[ref_offset:end_off]
                     if chunk.numel() > 0 and chunk.shape[0] == span_len:
                         codes = chunk.to(device=device, dtype=torch.long).clamp_(0, self.audio_vocab_size)
-                        audio_embed = torch.zeros_like(embeds)
-                        for i, emb_layer in enumerate(self.audio_embeddings):
-                            audio_embed = audio_embed + emb_layer(codes[:, i])
+                        if self._stacked_audio_emb_w is not None:
+                            codes_nq = codes.t()  # (n_vq, L)
+                            audio_embed = self._stacked_audio_emb_w[
+                                torch.arange(self.n_vq, device=device)[:, None], codes_nq
+                            ].sum(0)
+                        else:
+                            audio_embed = torch.zeros_like(embeds)
+                            for i, emb_layer in enumerate(self.audio_embeddings):
+                                audio_embed = audio_embed + emb_layer(codes[:, i])
                         embeds = embeds + audio_embed
                         last_ref_row = codes[-1].detach()
 
@@ -436,10 +446,17 @@ def preprocess(
         audio_codes_buf = (info_dict.get("audio_codes", {}) or {}).get("current")
         if isinstance(audio_codes_buf, torch.Tensor) and audio_codes_buf.numel() == self.n_vq:
             codes = audio_codes_buf.to(device=device, dtype=torch.long)
-            audio_embed = torch.zeros_like(text_embed)
-            for i, emb_layer in enumerate(self.audio_embeddings):
-                code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1)
-                audio_embed = audio_embed + emb_layer(code_i)
+            if self._stacked_audio_emb_w is not None:
+                # Single gather: stacked_emb[i, codes[i]] for each i → (n_vq, H) → sum → (1, H)
+                n_vq_idx = torch.arange(self.n_vq, device=device)
+                audio_embed = self._stacked_audio_emb_w[n_vq_idx, codes.clamp(0, self.audio_vocab_size)].sum(
+                    0, keepdim=True
+                )
+            else:
+                audio_embed = torch.zeros_like(text_embed)
+                for i, emb_layer in enumerate(self.audio_embeddings):
+                    code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1)
+                    audio_embed = audio_embed + emb_layer(code_i)
             combined = text_embed + audio_embed
         else:
             combined = text_embed
@@ -501,17 +518,38 @@ def _sample_audio_codes(
         if not bool(sampling_mask.any()):
             return codes
 
-        # Use deploy YAML defaults. The audio sampler isn't routed through
-        # vLLM's sampler so we apply the upstream's audio_top_k locally.
         audio_top_k = 25
         audio_temp = 1.7
 
+        if self._stacked_audio_head_w is not None:
+            # Single batched matmul: (n_vq, V+1, H) @ (H,) → (n_vq, V+1).
+            # Replaces n_vq serial nn.Linear calls.
+            all_logits = self._stacked_audio_head_w @ last_h.reshape(-1)  # (n_vq, V+1)
+            all_logits[:, -1] = float("-inf")  # invalid sentinel
+            all_logits[:, self.audio_pad_code] = float("-inf")
+
+            active_idx = sampling_mask.nonzero(as_tuple=True)[0]  # (n_active,)
+            active_logits = all_logits[active_idx]  # (n_active, V+1)
+            active_logits = active_logits / max(audio_temp, 1e-6)
+            if audio_top_k > 0 and audio_top_k < active_logits.shape[-1]:
+                top_vals, _ = torch.topk(active_logits, audio_top_k, dim=-1)
+                kth = top_vals[:, -1:].expand_as(active_logits)
+                active_logits = torch.where(
+                    active_logits < kth,
+                    torch.full_like(active_logits, float("-inf")),
+                    active_logits,
+                )
+            probs = torch.softmax(active_logits, dim=-1)
+            codes[active_idx] = torch.multinomial(probs, num_samples=1).squeeze(-1).long()
+            return codes
+
+        # Fallback: per-head loop (used before load_weights() completes).
         for i in range(n_vq):
             if not bool(sampling_mask[i]):
                 continue
             head = self.audio_heads[i]
             logits = head(last_h).reshape(-1)  # (V,)
-            logits[..., -1] = float("-inf")  # invalid sentinel
+            logits[..., -1] = float("-inf")
             logits[..., self.audio_pad_code] = float("-inf")
             sampled = self._sample_with_top_k(logits, audio_top_k, audio_temp)
             codes[i] = sampled.long()
@@ -553,7 +591,19 @@ def make_omni_output(
         # rows align with hidden rows, which align with info_dicts in order.
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
 
-        audio_codes_list: list[torch.Tensor] = []
+        # Per-request (start, end) hidden-row spans from the runner. In mixed
+        # prefill+decode steps the per-request token counts differ, so an equal
+        # rows-per-request split would sample codes from the wrong rows; fall
+        # back to that split only when spans are unavailable.
+        spans = kwargs.get("request_token_spans")
+
+        # One accumulated-codes tensor per request, in batch order. The generic
+        # splitter (``to_payload_element``) routes ``per_req_codes[idx]`` to
+        # request ``idx``, so Stage 1 decodes each request's own codes instead
+        # of request 0's. Skipped requests get an empty placeholder so indices
+        # stay aligned (a shorter list would silently fall back to entry 0).
+        per_req_codes: list[torch.Tensor] = []
+        have_codes = False
 
         if hidden.numel() > 0:
             num_rows = hidden.shape[0]
@@ -561,10 +611,16 @@ def make_omni_output(
 
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
+                    per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
-                row_start = i * rows_per_req
-                row_end = min(row_start + rows_per_req, num_rows)
+                if spans is not None and i < len(spans):
+                    row_start, row_end = spans[i]
+                    row_end = min(int(row_end), num_rows)
+                else:
+                    row_start = i * rows_per_req
+                    row_end = min(row_start + rows_per_req, num_rows)
                 if row_start >= row_end:
+                    per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
 
                 # Sample new audio codes from the last hidden state of this request.
@@ -582,21 +638,20 @@ def make_omni_output(
                     "current": new_codes,
                     "accumulated": updated_acc,
                 }
-                audio_codes_list.append(updated_acc)
+                per_req_codes.append(updated_acc)
+                have_codes = True
 
-        if not audio_codes_list:
+        if not have_codes:
             return OmniOutput(
                 text_hidden_states=hidden,
                 multimodal_outputs={},
             )
 
-        # Forward the accumulated audio codes (T, NQ) of the first request to
-        # Stage 1.  Multi-request batching for the codec is not yet supported
-        # by this adapter; the existing code already takes the first stage
-        # output anyway.
+        # Emit one accumulated-codes tensor per request as a list; the splitter
+        # picks ``element[idx]`` per request on the way to Stage 1.
         return OmniOutput(
             text_hidden_states=hidden,
-            multimodal_outputs={"codes": {"audio": audio_codes_list[0]}},
+            multimodal_outputs={"codes": {"audio": per_req_codes}},
         )
 
     # ------------------------------------------------------------------
@@ -661,6 +716,14 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         for n in backbone_loaded:
             loaded.add(f"model.{n}")
 
+        # Build stacked weight caches for vectorised decode ops.
+        self._stacked_audio_head_w = torch.stack(
+            [h.weight.detach() for h in self.audio_heads], dim=0
+        )  # (n_vq, audio_vocab+1, hidden_size)
+        self._stacked_audio_emb_w = torch.stack(
+            [e.weight.detach() for e in self.audio_embeddings], dim=0
+        )  # (n_vq, audio_vocab+1, hidden_size)
+
         return loaded
 
 
@@ -748,6 +811,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.logits_processor = LogitsProcessor(self.text_vocab_size)
         self._batch_state: list[dict[str, Any]] | None = None
 
+        # Stacked weight cache built after load_weights().
+        # shape: (n_vq, audio_vocab_size, hidden_size)
+        self._stacked_audio_emb_w: torch.Tensor | None = None
+
         self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
             ("audio_codes", "current"),
             ("audio_codes", "accumulated"),
@@ -837,6 +904,11 @@ def _build_input_embeds(
         embeds = self.embed_tokens[0](text_ids)
         if audio_codes is None:
             return embeds
+        if self._stacked_audio_emb_w is not None:
+            # Single gather: (n_vq, vocab, H)[arange[:,None], codes.T] → (n_vq, T, H) → sum → (T, H)
+            n_vq_idx = torch.arange(self.n_vq, device=audio_codes.device)
+            codes_t = audio_codes.t().clamp(0, self.audio_vocab_size - 1)  # (n_vq, T)
+            return embeds + self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t].sum(0)
         for i in range(self.n_vq):
             col = audio_codes[:, i].clamp_(0, self.audio_vocab_size - 1)
             embeds = embeds + self.embed_tokens[i + 1](col)
@@ -933,15 +1005,27 @@ def make_omni_output(
 
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
 
-        audio_codes_list: list[torch.Tensor] = []
+        # See delay talker: real per-request row spans from the runner, with the
+        # equal-split as fallback when unavailable.
+        spans = kwargs.get("request_token_spans")
+        # Per-request accumulated codes in batch order, pre-filled with empty
+        # placeholders so every skip path (stopped / eos / bos / pad) keeps
+        # indices aligned for to_payload_element's per-request routing. A
+        # shorter list would silently fall back to request 0's codes.
+        per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts]
+        have_codes = False
         if hidden.numel() > 0 and info_dicts:
             num_rows = hidden.shape[0]
             rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1))
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
                     continue
-                row_start = i * rows_per_req
-                row_end = min(row_start + rows_per_req, num_rows)
+                if spans is not None and i < len(spans):
+                    row_start, row_end = spans[i]
+                    row_end = min(int(row_end), num_rows)
+                else:
+                    row_start = i * rows_per_req
+                    row_end = min(row_start + rows_per_req, num_rows)
                 if row_start >= row_end:
                     continue
 
@@ -972,7 +1056,7 @@ def make_omni_output(
                     history_per_codebook=hist_per_cb,
                 ).squeeze(0)  # (n_vq,)
                 if int(state.get("step", 0)) < 5 or int(state.get("step", 0)) % 50 == 0:
-                    logger.info(
+                    logger.debug(
                         "[MossTTSRealtime make_omni] step=%d ch0=%d cursor=%d/%d",
                         int(state.get("step", 0)),
                         int(new_codes[0].item()),
@@ -1008,23 +1092,25 @@ def make_omni_output(
 
                 info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
                 state["step"] = int(state.get("step", 0)) + 1
-                audio_codes_list.append(updated_acc)
+                per_req_codes[i] = updated_acc
+                have_codes = True
 
-        if not audio_codes_list:
+        if not have_codes:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
 
-        # Forward the first request's accumulated codes (multi-request
-        # batching for the codec is not yet wired up — same as delay).
-        # The realtime variant emits raw codes (no delay pattern), so we
-        # signal that to the chunk processor via a 1-D bool tensor in
-        # ``meta.finished`` (already a tensor field on the schema). The
-        # processor reads its truth value as "skip de-delay" only for
-        # realtime; the delay path doesn't populate ``meta`` at all.
-        device = audio_codes_list[0].device
+        # Emit one accumulated-codes tensor per request as a list; the splitter
+        # routes ``element[idx]`` to request idx, so each Stage-1 request decodes
+        # its own codes instead of request 0's.
+        # The realtime variant emits raw codes (no delay pattern), so we signal
+        # that to the chunk processor via a 1-D bool tensor in ``meta.finished``
+        # (already a tensor field on the schema). The processor reads its truth
+        # value as "skip de-delay" only for realtime; the delay path doesn't
+        # populate ``meta`` at all.
+        device = hidden.device
         return OmniOutput(
             text_hidden_states=hidden,
             multimodal_outputs={
-                "codes": {"audio": audio_codes_list[0]},
+                "codes": {"audio": per_req_codes},
                 "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
             },
         )
@@ -1084,6 +1170,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         for n in backbone_loaded:
             loaded.add(f"model.{n}")
 
+        # Build stacked embedding cache for vectorised decode ops.
+        # embed_tokens[0] is the text embedding; [1..n_vq] are audio codebooks.
+        self._stacked_audio_emb_w = torch.stack(
+            [self.embed_tokens[i + 1].weight.detach() for i in range(self.n_vq)], dim=0
+        )  # (n_vq, audio_vocab_size, hidden_size)
+
         logger.info(
             "[MossTTSRealtime] loaded %d/%d params; skipped=%d (first 5: %s)",
             len(loaded),
diff --git a/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py
deleted file mode 100644
index a5a9f25acd3..00000000000
--- a/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py
+++ /dev/null
@@ -1,60 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Configuration for MOSS-TTS-Nano in vLLM-Omni single-stage pipeline."""
-
-from transformers.configuration_utils import PretrainedConfig
-
-
-class MossTTSNanoConfig(PretrainedConfig):
-    """Config for MOSS-TTS-Nano (OpenMOSS-Team/MOSS-TTS-Nano).
-
-    The model is a Global-Local GPT2-based AR LM (0.1B) paired with the
-    MOSS-Audio-Tokenizer-Nano codec.  Both components run in a single
-    vLLM-Omni generation stage.
-
-    Relevant fields from config.json:
-      gpt2_config  – backbone config (n_layer=12, n_embd=768, n_head=12, ...)
-      n_vq         – number of RVQ codebooks (16)
-      audio_vocab_size – per-codebook vocabulary size (1024)
-      audio_tokenizer_pretrained_name_or_path – HF hub path for codec model
-    """
-
-    model_type = "moss_tts_nano"
-
-    def __init__(self, **kwargs):
-        gpt2_cfg = kwargs.pop("gpt2_config", None) or {}
-        if hasattr(gpt2_cfg, "to_dict"):
-            gpt2_cfg = gpt2_cfg.to_dict()
-
-        super().__init__(**kwargs)
-
-        # --- GPT2 backbone parameters (exposed at top level for vLLM) ---
-        self.hidden_size = gpt2_cfg.get("n_embd", 768)
-        self.num_hidden_layers = gpt2_cfg.get("n_layer", 12)
-        self.num_attention_heads = gpt2_cfg.get("n_head", 12)
-        self.num_key_value_heads = self.num_attention_heads  # no GQA in GPT2
-        self.head_dim = self.hidden_size // self.num_attention_heads
-        self.vocab_size = gpt2_cfg.get("vocab_size", 16384)
-        self.max_position_embeddings = gpt2_cfg.get("n_positions", 32768)
-        self.intermediate_size = gpt2_cfg.get("n_inner", self.hidden_size * 4)
-
-        # --- Audio codec parameters ---
-        self.n_vq: int = getattr(self, "n_vq", 16)
-        self.audio_vocab_size: int = getattr(self, "audio_vocab_size", 1024)
-        self.audio_start_token_id: int = getattr(self, "audio_start_token_id", 6)
-        self.audio_end_token_id: int = getattr(self, "audio_end_token_id", 7)
-        self.audio_user_slot_token_id: int = getattr(self, "audio_user_slot_token_id", 8)
-        self.audio_assistant_slot_token_id: int = getattr(self, "audio_assistant_slot_token_id", 9)
-        self.audio_tokenizer_pretrained_name_or_path: str = getattr(
-            self,
-            "audio_tokenizer_pretrained_name_or_path",
-            "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
-        )
-        self.audio_tokenizer_sample_rate: int = getattr(self, "audio_tokenizer_sample_rate", 48000)
-
-        # vLLM requires speculative_config to be absent or None
-        self.speculative_config = None
-
-    def get_text_config(self, **kwargs):
-        """Return self so vLLM uses our top-level config."""
-        return self
diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py
index f52577f67aa..70596098298 100644
--- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py
+++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py
@@ -172,7 +172,6 @@ class MossTTSNanoForGeneration(nn.Module):
     has_preprocess = False
     has_postprocess = False
     enable_update_additional_information = True
-    inject_omni_request_id_into_runtime_info = True
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         super().__init__()
diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py
index c5026086e23..43c496a1cb0 100644
--- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py
+++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py
@@ -379,13 +379,12 @@ def _forward_generator(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutpu
         target_ids = torch.full((num_codebooks, target_len), mask_id, dtype=torch.long, device=device)
 
         # Conditional: [text] [ref_audio?] [target_mask]
+        cond_audio_start = text_ids.shape[1]
         if ref_audio_tokens is not None:
             ref_tokens = ref_audio_tokens.to(device)  # [8, T_ref]
             cond_ids = torch.cat([text_ids, ref_tokens, target_ids], dim=1)
-            cond_audio_start = text_ids.shape[1]
         else:
             cond_ids = torch.cat([text_ids, target_ids], dim=1)
-            cond_audio_start = text_ids.shape[1]
 
         cond_len = cond_ids.shape[1]
 
diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
index fc5befeb959..7eacfb9bb92 100644
--- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
+++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
@@ -6,8 +6,9 @@
 Generates 8-codebook audio tokens from text via 32-step non-autoregressive
 iterative masked prediction with classifier-free guidance.
 
-Uses vLLM-Omni's DiffusionAttention for optimized full (bidirectional) attention
-via FlashAttention/SageAttention/SDPA backends.
+Uses full bidirectional attention computed directly with PyTorch SDPA
+(torch.nn.functional.scaled_dot_product_attention); no auto-selected
+FlashAttention/SageAttention/DiffusionAttention backend is used.
 """
 
 from __future__ import annotations
@@ -237,7 +238,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge
 
 
 # ---------------------------------------------------------------------------
-# Qwen3-style transformer blocks using DiffusionAttention
+# Qwen3-style transformer blocks using PyTorch SDPA
 # ---------------------------------------------------------------------------
 
 
@@ -256,7 +257,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class OmniVoiceAttention(nn.Module):
-    """Qwen3-style GQA attention using DiffusionAttention backend."""
+    """Qwen3-style GQA attention using PyTorch SDPA (full bidirectional)."""
 
     def __init__(self, config: OmniVoiceConfig):
         super().__init__()
@@ -315,18 +316,17 @@ def forward(
         v = v.permute(0, 2, 1, 3)
 
         # Convert [B, 1, S, S] bool mask to float mask for SDPA
+        # (0.0 where attend, -inf where masked)
         sdpa_mask = None
         if attention_mask is not None:
-            sdpa_mask = attention_mask.to(dtype=q.dtype)
-            sdpa_mask = sdpa_mask.masked_fill(~attention_mask, float("-inf"))
-            sdpa_mask = sdpa_mask.masked_fill(attention_mask, 0.0)
+            sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf"))
 
         out = F.scaled_dot_product_attention(
             q,
             k,
             v,
             attn_mask=sdpa_mask,
-            scale=1.0 / math.sqrt(self.head_dim),
+            scale=self.scale,
         )
 
         # Back to (batch, seq, heads * head_dim)
@@ -351,7 +351,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class OmniVoiceTransformerBlock(nn.Module):
-    """Single Qwen3 transformer block with DiffusionAttention."""
+    """Single Qwen3 transformer block with PyTorch SDPA attention."""
 
     def __init__(self, config: OmniVoiceConfig):
         super().__init__()
@@ -643,10 +643,9 @@ class OmniVoiceGenerator(nn.Module):
     - 32-step iterative unmasking with classifier-free guidance
 
     Optimizations:
-    - DiffusionAttention (FlashAttn/SageAttn/SDPA auto-selected)
-    - TeaCache / Cache-DiT compatible (hook-based, non-intrusive)
+    - Full bidirectional attention via PyTorch SDPA (no auto-selected
+      FlashAttn/SageAttn/DiffusionAttention backend)
     - regionally_compile() compatible for torch.compile on repeated blocks
-    - Sequence parallelism via SP hooks for multi-GPU
     """
 
     # For regionally_compile() support
diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py
index 8bb921a1eae..28ca4f36289 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py
@@ -27,7 +27,7 @@ class Qwen3TTSSpeakerEncoderConfig(PretrainedConfig):
     Args:
         mel_dim (`int`, *optional*, defaults to 128):
             The dimension of the input mel-spectrogram.
-        enc_dim (`int`, *optional*, defaults to 192):
+        enc_dim (`int`, *optional*, defaults to 1024):
             The dimension of the final speaker embedding.
         enc_channels (`list[int]`, *optional*, defaults to `[512, 512, 512, 512, 1536]`):
             A list of output channels for each TDNN/SERes2Net layer in the encoder.
@@ -80,11 +80,11 @@ class Qwen3TTSTalkerCodePredictorConfig(PretrainedConfig):
 
 
     Args:
-        vocab_size (`int`, *optional*, defaults to 151936):
+        vocab_size (`int`, *optional*, defaults to 2048):
             Vocabulary size of the Qwen3TTSTalkerCodePredictor model.
             Defines the number of different tokens that can be represented by the
             `inputs_ids` passed when calling [`Qwen3TTSTalkerCodePredictorModel`]
-        hidden_size (`int`, *optional*, defaults to 4096):
+        hidden_size (`int`, *optional*, defaults to 1024):
             Dimension of the hidden representations.
         intermediate_size (`int`, *optional*, defaults to 22016):
             Dimension of the MLP representations.
@@ -271,11 +271,11 @@ class Qwen3TTSTalkerConfig(PretrainedConfig):
 
 
     Args:
-        vocab_size (`int`, *optional*, defaults to 151936):
+        vocab_size (`int`, *optional*, defaults to 3072):
             Vocabulary size of the Qwen3TTSTalker model.
             Defines the number of different tokens that can be represented by the
             `inputs_ids` passed when calling [`Qwen3TTSTalkerModel`]
-        hidden_size (`int`, *optional*, defaults to 2048):
+        hidden_size (`int`, *optional*, defaults to 1024):
             Dimension of the hidden representations.
         intermediate_size (`int`, *optional*, defaults to 6144):
             Dimension of the MLP representations.
@@ -436,7 +436,6 @@ def __init__(
             self.rope_scaling["rope_type"] = self.rope_scaling["type"]
 
         if code_predictor_config is None:
-            code_predictor_config = {}
             self.code_predictor_config = Qwen3TTSTalkerCodePredictorConfig()
             logger.info("code_predictor_config is None. Initializing code_predictor model with default values")
         elif isinstance(code_predictor_config, Qwen3TTSTalkerCodePredictorConfig):
diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py
index b0294b72583..0d8f032949e 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py
@@ -1320,9 +1320,7 @@ def _get_ref_audio() -> tuple[np.ndarray, int]:
             # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
             spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long)
             spk_embed = codec_embed(spk_tensor)
-            if spk_embed.ndim == 1:
-                spk_embed = spk_embed.view(1, 1, -1)
-            elif spk_embed.ndim == 2:
+            if spk_embed.ndim in (1, 2):
                 spk_embed = spk_embed.view(1, 1, -1)
             speaker_embed = spk_embed
             codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
@@ -1523,9 +1521,7 @@ def _first(x: object, default: object) -> object:
 
                 ref_code_len: int | None = None
                 if isinstance(ref_code, list):
-                    if ref_code and isinstance(ref_code[0], list):
-                        ref_code_len = len(ref_code)
-                    elif ref_code:
+                    if ref_code:
                         ref_code_len = len(ref_code)
                 elif hasattr(ref_code, "shape"):
                     try:
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
index 64eb3d25320..c1bb7c049ad 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
@@ -286,8 +286,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             )
         self._codec_eos_token_id = int(getattr(self.talker_config, "codec_eos_token_id", -1))
 
-        self._eos_logit_bias: float = 0.0
-
         self.have_multimodal_outputs = True
         self.has_preprocess = True
         self.has_postprocess = True
@@ -517,11 +515,6 @@ def compute_logits(
         # Mask out invalid codec ids using the pre-built constant buffer.
         logits = logits.masked_fill(self._codec_disallowed_mask, float("-inf"))
 
-        if self._eos_logit_bias != 0.0:
-            eos_id = self._codec_eos_token_id
-            if 0 <= eos_id < logits.shape[-1]:
-                logits[:, eos_id] = logits[:, eos_id] + self._eos_logit_bias
-
         return logits
 
     # -------------------- Omni multimodal output plumbing --------------------
diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
index 6eb78eeae0d..447823abb11 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
@@ -912,11 +912,6 @@ def enable_cudagraph(
             self._cudagraph_wrapper.compile_shapes,
         )
 
-    def disable_cudagraph(self):
-        self._cudagraph_enabled = False
-        self._cudagraph_wrapper = None
-        logger.info("CUDA Graph disabled for decoder")
-
     def forward(self, codes):
         if codes.shape[1] != self.config.num_quantizers:
             raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py
index 5609abb394c..4014014533d 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py
@@ -33,7 +33,6 @@
 
 import random
 import typing as tp
-from math import ceil
 
 import numpy as np
 import torch
@@ -45,10 +44,6 @@
 logger = init_logger(__name__)
 
 
-def round_up_multiple(num, mult):
-    return ceil(num / mult) * mult
-
-
 def default(val: tp.Any, d: tp.Any) -> tp.Any:
     return val if val is not None else d
 
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2.py b/vllm_omni/model_executor/models/step_audio2/step_audio2.py
index 82e5609ff40..0a94ab9defb 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2.py
+++ b/vllm_omni/model_executor/models/step_audio2/step_audio2.py
@@ -26,17 +26,17 @@ class StepAudio2ForConditionalGeneration(nn.Module, SupportsMultiModal, Supports
     Step-Audio2 Main Controller
 
     Manages two-stage inference pipeline:
-    - Stage 1 (Thinker): Audio understanding and token generation
-    - Stage 2 (Token2Wav): Audio token to waveform synthesis
+    - Stage 0 (Thinker): Audio understanding and token generation
+    - Stage 1 (Token2Wav): Audio token to waveform synthesis
 
     Usage:
-        # Stage 1: Thinker
+        # Stage 0: Thinker
         model = StepAudio2ForConditionalGeneration(
             vllm_config=config,
             model_stage="thinker"
         )
 
-        # Stage 2: Token2Wav
+        # Stage 1: Token2Wav
         model = StepAudio2ForConditionalGeneration(
             vllm_config=config,
             model_stage="token2wav"
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py
index 802cd63bef8..2f2918a1be3 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py
+++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py
@@ -42,39 +42,6 @@ def audio_end(self) -> int:
         return self.audio_start + self.audio_vocab_size - 1
 
 
-@dataclass(frozen=True)
-class StepAudio2EncoderConfig:
-    """
-    Step-Audio2 audio encoder architecture configuration.
-
-    These values define the encoder structure and should match your model weights.
-    """
-
-    # Mel spectrogram parameters
-    n_mels: int = 128
-    """Number of mel frequency bins"""
-
-    # Encoder architecture
-    n_audio_ctx: int = 1500
-    """Audio context length (max sequence length)"""
-
-    n_audio_state: int = 512
-    """Audio encoder hidden state dimension"""
-
-    n_audio_head: int = 8
-    """Number of attention heads in encoder"""
-
-    n_audio_layer: int = 6
-    """Number of encoder layers"""
-
-    # Adapter parameters
-    kernel_size: int = 3
-    """Convolution kernel size for adapter"""
-
-    adapter_stride: int = 2
-    """Stride for adapter convolution (downsampling factor)"""
-
-
 @dataclass(frozen=True)
 class StepAudio2Token2WavConfig:
     """Token2Wav synthesis configuration."""
@@ -141,7 +108,6 @@ class StepAudio2ModelConfig:
 
 # Default configuration instances
 DEFAULT_TOKEN_CONFIG = StepAudio2TokenConfig()
-DEFAULT_ENCODER_CONFIG = StepAudio2EncoderConfig()
 DEFAULT_TOKEN2WAV_CONFIG = StepAudio2Token2WavConfig()
 DEFAULT_STREAM_CONFIG = StepAudio2StreamConfig()
 DEFAULT_MODEL_CONFIG = StepAudio2ModelConfig()
@@ -162,12 +128,10 @@ class StepAudio2ModelConfig:
 
 __all__ = [
     "StepAudio2TokenConfig",
-    "StepAudio2EncoderConfig",
     "StepAudio2Token2WavConfig",
     "StepAudio2StreamConfig",
     "StepAudio2ModelConfig",
     "DEFAULT_TOKEN_CONFIG",
-    "DEFAULT_ENCODER_CONFIG",
     "DEFAULT_TOKEN2WAV_CONFIG",
     "DEFAULT_STREAM_CONFIG",
     "DEFAULT_MODEL_CONFIG",
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py
index 7fd0bf4fbd7..c92f2397813 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py
+++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py
@@ -609,7 +609,7 @@ class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, S
     Architecture:
         AudioEncoder (6 layers, 512 hidden) →
         Adaptor (512 → LLM dim) →
-        Qwen2 LLM (vocab=64012)
+        Qwen2 LLM (vocab=158720)
     """
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -868,31 +868,3 @@ def separate_tokens(
     def has_audio_output(token_ids: list[int], audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start) -> bool:
         """Check if generated tokens contain audio tokens"""
         return any(tid >= audio_start for tid in token_ids)
-
-
-class StepAudio2OutputProcessor:
-    """Helper class to process Step-Audio2 outputs"""
-
-    @staticmethod
-    def process_output(output_ids: torch.Tensor, tokenizer, remove_audio_padding: bool = True) -> dict:
-        """Process model output and separate text/audio tokens"""
-        if isinstance(output_ids, torch.Tensor):
-            output_ids = output_ids.squeeze().tolist()
-
-        # Separate tokens
-        text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(output_ids)
-
-        # Remove audio padding if requested
-        if remove_audio_padding and audio_tokens:
-            audio_tokens = [t for t in audio_tokens if t < DEFAULT_TOKEN_CONFIG.audio_eos]
-
-        # Decode text
-        text = tokenizer.decode(text_tokens, skip_special_tokens=False)
-
-        return {
-            "text": text,
-            "text_tokens": text_tokens,
-            "audio_tokens": audio_tokens,
-            "has_audio": len(audio_tokens) > 0,
-            "all_tokens": output_ids,
-        }
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py
index a59ba3fef19..86109ad9c7d 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py
+++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py
@@ -94,8 +94,6 @@ def __init__(
         self.mel_cache_len = DEFAULT_STREAM_CONFIG.mel_cache_len
         self.source_cache_len = STREAM_SOURCE_CACHE_LEN
         self.speech_window: torch.Tensor | None = None  # created lazily on device
-        self.stream_cache: dict | None = None  # flow model causal cache
-        self.hift_cache_dict: dict = {}  # HiFT vocoder overlap cache
 
     def _ensure_models_loaded(self):
         """Lazy load models on first use"""
@@ -253,125 +251,6 @@ def forward(
         else:
             return wav
 
-    # ------------------------------------------------------------------
-    # Streaming (chunked) inference
-    # ------------------------------------------------------------------
-
-    def setup_stream(self, prompt_wav: str) -> None:
-        """Initialize flow + HiFT caches for chunked streaming inference.
-
-        Must be called once before the first ``stream_chunk()`` call for a
-        given utterance.
-        """
-        if prompt_wav not in self.cache:
-            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
-        (
-            prompt_speech_tokens,
-            _,
-            spk_emb,
-            prompt_mels,
-            _,
-        ) = self.cache[prompt_wav]
-
-        # Lazily create the Hamming window on the correct device
-        if self.speech_window is None or self.speech_window.device != self.device:
-            self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to(
-                device=self.device, dtype=torch.float32
-            )
-
-        # Flow model: setup causal cache with prompt tokens + lookahead
-        pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len
-        self.stream_cache = self.flow.setup_cache(
-            torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]], dim=1),
-            prompt_mels,
-            spk_emb,
-            n_timesteps=self.n_timesteps,
-        )
-
-        # HiFT vocoder: empty overlap buffers
-        self.hift_cache_dict = {
-            "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device),
-            "source": torch.zeros(1, 1, 0, device=self.device),
-            "speech": torch.zeros(1, 0, device=self.device),
-        }
-
-    def stream_chunk(
-        self,
-        audio_tokens: list[int],
-        prompt_wav: str,
-        last_chunk: bool = False,
-    ) -> torch.Tensor:
-        """Process one chunk of audio tokens and return a waveform segment.
-
-        Args:
-            audio_tokens: Token IDs for this chunk (includes lookahead for
-                non-last chunks).
-            prompt_wav: Path to the prompt wav (for cached speaker embeddings).
-            last_chunk: Whether this is the final chunk of the utterance.
-
-        Returns:
-            1-D float tensor of audio samples (24 kHz).
-        """
-        if self.stream_cache is None:
-            raise ValueError("stream_cache not initialised – call setup_stream() first")
-
-        if prompt_wav not in self.cache:
-            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
-        _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
-
-        token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device)
-
-        with torch.amp.autocast(
-            str(self.device.type),
-            dtype=torch.float16 if self.float16 else torch.float32,
-        ):
-            chunk_mel, self.stream_cache = self.flow.inference_chunk(
-                token=token_tensor,
-                spk=spk_emb,
-                cache=self.stream_cache,
-                last_chunk=last_chunk,
-                n_timesteps=self.n_timesteps,
-            )
-
-        # Trim estimator attention cache to avoid unbounded growth
-        keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep
-        est_att = self.stream_cache.get("estimator_att_cache")
-        if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep):
-            self.stream_cache["estimator_att_cache"] = torch.cat(
-                [est_att[:, :, :, :, : prompt_mels.shape[1]], est_att[:, :, :, :, -keep:]],
-                dim=4,
-            )
-
-        # ---- HiFT vocoder with overlap-add ----
-        hift_cache_mel = self.hift_cache_dict["mel"]
-        hift_cache_source = self.hift_cache_dict["source"]
-        hift_cache_speech = self.hift_cache_dict["speech"]
-
-        mel = torch.cat([hift_cache_mel, chunk_mel], dim=2)
-        speech, source = self.hift(mel, hift_cache_source)
-
-        # Cross-fade overlap region for smooth concatenation
-        if hift_cache_speech.shape[-1] > 0:
-            speech = fade_in_out(speech, hift_cache_speech, self.speech_window)
-
-        # Update vocoder cache for the next chunk
-        self.hift_cache_dict = {
-            "mel": mel[..., -self.mel_cache_len :].clone().detach(),
-            "source": source[:, :, -self.source_cache_len :].clone().detach(),
-            "speech": speech[:, -self.source_cache_len :].clone().detach(),
-        }
-
-        # For non-last chunks, trim the tail that will be re-synthesised
-        if not last_chunk:
-            speech = speech[:, : -self.source_cache_len]
-
-        return speech.squeeze(0)  # [samples]
-
-    def reset_stream(self) -> None:
-        """Clear all streaming state after an utterance is complete."""
-        self.stream_cache = None
-        self.hift_cache_dict = {}
-
     # ------------------------------------------------------------------
     # Per-request streaming (does NOT mutate self – uses external state)
     # ------------------------------------------------------------------
diff --git a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py
index b87ec5aafef..13f1d0937d3 100644
--- a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py
+++ b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py
@@ -65,7 +65,6 @@ def __init__(
         prefix: str = "",
     ) -> None:
         super().__init__()
-        self.layer_idx = layer_idx
         self.hidden_size = hidden_size
         self.num_heads = num_attention_heads
         self.head_dim = kv_channels if kv_channels else hidden_size // num_attention_heads
@@ -318,18 +317,12 @@ def precompute_fused_qkv(self) -> None:
                 ).detach()
 
     def compile_selective(self) -> list[str]:
-        """Compile the full model forward as one graph.
-
-        Earlier versions compiled ``layer.mlp`` + ``layer.self_attn.o_proj``
-        (PR #2690) and then the whole ``layer`` (perf/voxcpm2-streaming-vae).
-        Both still paid one Dynamo dispatch per layer per decode step.
-        V3 profiling showed 1,332 per-layer dispatches (~28 layers × ~47
-        decode steps) costing ~726 ms of CPU self-time for a long prompt.
-
-        Compiling ``forward`` at the model level lets Dynamo unroll the
-        28-layer Python loop inside the graph. Graph breaks at
-        PagedAttention produce sub-graphs but Dynamo memoises the whole
-        trace once, so the per-step dispatch drops from 28 to just a few.
+        """Compile the full model ``forward`` as one graph.
+
+        ``torch.compile`` is applied at the model level so Dynamo unrolls the
+        per-layer Python loop inside the graph. Graph breaks at PagedAttention
+        produce sub-graphs, but Dynamo memoises the whole trace once, so the
+        per-decode-step dispatch drops to just a few instead of one per layer.
         """
         if self._compiled_layers:
             return []
diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
index 3fcc683f3aa..9505bfbee38 100644
--- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
+++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
@@ -312,7 +312,6 @@ class _RequestState:
     decode_step_count: int = 0
     request_start_time: float = 0.0
     prefill_completed: bool = False
-    prefill_text: str = ""
     prompt_cache: dict | None = None
     prefill_masks: tuple | None = None
     is_stopping: bool = False
@@ -2510,7 +2509,6 @@ def preprocess(
                 token_ids = token_ids[1:]
 
             state = self._get_or_create_state(req_id)
-            state.prefill_text = ""
             state.decode_pad = None
             state.prefill_completed = False
             state.decode_step_count = 0
@@ -2572,29 +2570,26 @@ def preprocess(
                 )
 
             if state.prompt_cache is None and (ref_audio or (prompt_audio and prompt_text)):
-                if state.prompt_cache is None:
-                    try:
-                        state.prompt_cache = self._build_prompt_cache(
-                            ref_audio=ref_audio,
-                            prompt_audio=prompt_audio,
-                            prompt_text=prompt_text,
+                try:
+                    state.prompt_cache = self._build_prompt_cache(
+                        ref_audio=ref_audio,
+                        prompt_audio=prompt_audio,
+                        prompt_text=prompt_text,
+                    )
+                    if (
+                        voice_name
+                        and state.prompt_cache is not None
+                        and state.prompt_cache.get("mode") == "reference"
+                        and "ref_audio_feat" in state.prompt_cache
+                    ):
+                        _key = self._speaker_cache.make_cache_key(
+                            voice_name, model_type="voxcpm2", created_at=_created_at
                         )
-                        if (
-                            voice_name
-                            and state.prompt_cache is not None
-                            and state.prompt_cache.get("mode") == "reference"
-                            and "ref_audio_feat" in state.prompt_cache
-                        ):
-                            _key = self._speaker_cache.make_cache_key(
-                                voice_name, model_type="voxcpm2", created_at=_created_at
-                            )
-                            self._speaker_cache.put(
-                                _key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()}
-                            )
-                            logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
-                    except Exception as e:
-                        logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e)
-                        state.prompt_cache = None
+                        self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()})
+                        logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
+                except Exception as e:
+                    logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e)
+                    state.prompt_cache = None
 
             inputs = self._build_prefill_inputs(token_ids, dev, req_id)
             tts = self.tts
diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
index d12034c243b..8b6872cfeef 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
+++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
@@ -244,9 +244,6 @@ def embed_multimodal(self, **kwargs):
         # Delegate to generation model for multimodal processing
         return self.model.embed_multimodal(**kwargs)
 
-    def last_index_of(self, list, value):
-        return len(list) - 1 - list[::-1].index(value)
-
     def forward(
         self,
         input_ids: torch.Tensor,
diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
index 2e50e34f32d..ce005936c63 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
+++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
@@ -246,7 +246,6 @@ def __init__(
 
         self.n_local_heads: int = args.n_heads
         self.n_local_kv_heads: int = args.n_kv_heads
-        self.repeats = self.n_local_heads
         self.layer_id = layer_id
 
         self.head_dim = args.head_dim
@@ -276,7 +275,6 @@ def __init__(
             bias=args.use_biases,
         )
 
-        self.softmax_scale: float = self.args.head_dim**-0.5
         self.repeats = self.n_local_heads // self.n_local_kv_heads
 
     def _native_attention(
@@ -420,13 +418,6 @@ def __init__(
         self.acoustic_transformer_args = args
         assert isinstance(self.acoustic_transformer_args, AcousticTransformerArgs)
 
-        # currently assuming always 1 semantic codebook + N acoustic codebooks
-        self.num_non_acoustic_embeddings = 1
-        self.num_acoustic_codebooks = len(self.model_args.get_codebook_sizes()) - self.num_non_acoustic_embeddings
-
-        # flow matching utils
-        self.sigma = args.sigma
-
         # codebook sizes
         acoustic_codebook_sizes = self.model_args.get_codebook_sizes(
             pad_to_multiple=None, include_special_tokens=False
@@ -715,7 +706,6 @@ def __call__(
                 )
 
             text_tokens_for_audio = list[torch.Tensor]()
-            assert audios is not None
             audio_tokens_pt = list[torch.Tensor]()
             for audio_token_array in audio_tokens:
                 assert isinstance(audio_token_array, np.ndarray)
@@ -1008,7 +998,7 @@ def _parse_and_validate_audio_arrays(
 
         if audio_arrays is not None:
             if not isinstance(audio_arrays, (torch.Tensor, list)):
-                raise ValueError(f"Incorrect type of images. Got type: {type(audio_arrays)}")
+                raise ValueError(f"Incorrect type of audio_arrays. Got type: {type(audio_arrays)}")
             if isinstance(audio_arrays, torch.Tensor) and audio_arrays.dim() == 3:
                 audio_arrays = flatten_bn(audio_arrays)
             if isinstance(audio_arrays, torch.Tensor):
@@ -1053,24 +1043,11 @@ def fake_logits_for_audio_tokens(
         fake_logits[~is_eos, self.audio_tok_id] = 1.0
         return fake_logits
 
-    # TODO(chenyo): Remove this
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-    ) -> torch.Tensor | None:
-        text_logits = self.language_model.compute_logits(
-            hidden_states,
-        )
-        assert text_logits is not None
-        return text_logits
-
     def compute_mm_logits(
         self,
         hidden_states: torch.Tensor,
         cfg_alpha: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        audio_codes = None
-        mm_tokens = None
         audio_codes = self.acoustic_transformer(
             llm_hidden=hidden_states,
             cfg_alpha=cfg_alpha,
@@ -1088,7 +1065,6 @@ def compute_mm_logits(
         return fake_eos, mm_tokens
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # fmt: on
         remapping_rules = [
             (r"^acoustic_transformer\.(.*)$", r"\1"),  # noqa: E501
             (r"^audio_tokenizer\.(.*)$", r"\1"),  # noqa: E501
diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py
index fc753a58f37..57a86a5e91c 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py
+++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py
@@ -652,7 +652,7 @@ def __init__(
     def layer_id(self) -> int:
         return self._layer_id
 
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         r = self.attention.forward(self.attention_norm(x))
         if self.post_attention_norm is not None:
             r = self.post_attention_norm(r)
@@ -757,7 +757,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             # transformer
             layer_args = deepcopy(args)
             layer_args.attn_sliding_window_size = cur_window_size
-            assert layer_args.qk_norm, "qk_norm must be True for decoder"
+            assert layer_args.qk_norm, "qk_norm must be True for encoder"
             encoder_transformer = Transformer(
                 args=layer_args,
                 n_layers=n_layers,
diff --git a/vllm_omni/request.py b/vllm_omni/request.py
index f936069be2d..b094f2d39fa 100644
--- a/vllm_omni/request.py
+++ b/vllm_omni/request.py
@@ -30,15 +30,16 @@ class OmniRequest(Request):
 
     def __init__(
         self,
+        *args,
         prompt_embeds: PromptEmbedsPayload | torch.Tensor | None = None,
         # Optional external request ID for tracking
         external_req_id: str | None = None,
         additional_information: AdditionalInformationPayload | None = None,
-        *args,
         **kwargs,
     ):
-        prompt_embeds_tensor = self._maybe_decode_prompt_embeds(prompt_embeds)
-        super().__init__(prompt_embeds=prompt_embeds_tensor, *args, **kwargs)
+        if prompt_embeds is not None:
+            kwargs["prompt_embeds"] = self._maybe_decode_prompt_embeds(prompt_embeds)
+        super().__init__(*args, **kwargs)
         # Preserve serialized prompt embeddings payload (optional)
         self.prompt_embeds_payload: PromptEmbedsPayload | None = (
             prompt_embeds if isinstance(prompt_embeds, PromptEmbedsPayload) else None
diff --git a/vllm_omni/transformers_utils/processors/ming.py b/vllm_omni/transformers_utils/processors/ming.py
index 685909029ad..b32f10d1204 100644
--- a/vllm_omni/transformers_utils/processors/ming.py
+++ b/vllm_omni/transformers_utils/processors/ming.py
@@ -143,7 +143,6 @@ def __call__(
         #      (conv1 has stride=1 and does not change T)
         #   2. AudioProjector Conv1d: kernel=3, stride=2, padding=1
         # Combined: T → ((T-1)//2 + 1 - 1)//2 + 1
-        # See also: AudioProjector.compute_output_length()
         encoder_feats_lengths = ((audio_feats_lengths - 3 + 2 * 1) // 2 + 1 - 3 + 2 * 1) // 2 + 1
         audio_feats = torch.cat(audio_feat_list, dim=0).unsqueeze(0)  # [1, T_total, n_mels]
 
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index 349e1da0411..cdf2286bfe4 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -1313,6 +1313,20 @@ def _build_model_kwargs_extra(self) -> dict:
 
             traceback.print_exc()
 
+        # Per-request (start, end) hidden-row spans so make_omni_output can map
+        # flat hidden rows to the right request in mixed prefill+decode steps,
+        # instead of assuming an equal rows-per-request split (which samples the
+        # wrong rows whenever per-request token counts differ).
+        nstp = self._omni_num_scheduled_tokens_np
+        if nstp is not None and len(nstp) == len(self.input_batch.req_ids):
+            try:
+                model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(nstp)
+            except Exception as e:
+                # Visible on purpose: the fallback is the equal rows-per-request
+                # split, which can re-introduce the cross-request corruption this
+                # plumbing fixes — a silent failure here must not pass unnoticed.
+                logger.warning("[OMNI] Failed to compute request_token_spans: %s", e)
+
         if self._omni_query_start_loc_model_kwarg:
             try:
                 num_reqs = len(self.input_batch.req_ids)

From 16e0b6b00d1d7922afaafa5596e6b5dcec58007d Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 20 Jun 2026 10:55:30 +0000
Subject: [PATCH 27/51] [BDE] Fix import order in runner.py (ruff/isort after
 experimental rename)

The experimental/bde imports sorted before vllm_omni.diffusion.* once the
package moved under experimental/; ruff isort wants diffusion first. Fixes the
pre-commit (ruff-check) CI check on PR #4534.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 vllm_omni/experimental/bde/runner.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/vllm_omni/experimental/bde/runner.py b/vllm_omni/experimental/bde/runner.py
index b68c8d98226..2118a91a98c 100644
--- a/vllm_omni/experimental/bde/runner.py
+++ b/vllm_omni/experimental/bde/runner.py
@@ -20,13 +20,13 @@
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.gather import BDEKVState
-from vllm_omni.experimental.bde.kv_cache.manager import BDEKVCache
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import MAX_DREAMZERO_SESSIONS
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
+from vllm_omni.experimental.bde.kv_cache.gather import BDEKVState
+from vllm_omni.experimental.bde.kv_cache.manager import BDEKVCache
 
 logger = init_logger(__name__)
 

From 371db602043eb5dee0a806aa880777132795f097 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 20 Jun 2026 11:57:12 +0000
Subject: [PATCH 28/51] [BDE] Cache I2V image-token k_img/v_img in the
 cross-attn pool
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Restore BDE's engine-managed cross-attn (§5.2) after the main rebase. Upstream
#4154 added image-token k_img/v_img caching to the I2V cross-attn; BDE's pool only
held text k/v, so it had been falling back to the model-local cache. Now BDE caches
both:

- BDEKVCache: add per-layer _cross_k_img/_cross_v_img pools (2, 257, heads, dim) for
  I2V; write_cross_kv/read_cross_kv handle the image variants (read returns
  k_img/v_img when the image pool is allocated).
- runner: derive cross_attn_img_length (257 for i2v, 0 for t2v) and thread it through
  build_kv_cache.
- pipeline: _kv_get_cross routes back to the pool; _kv_populate_cross projects text
  (text_embedding) and image (img_emb(clip_feature)) K/V, and is moved into the
  image-encode block so clip_feas is available (runs on the first forward and after
  each window-boundary reset).

Validated on H200 (DreamZero-DROID): tests/bde 77 passing; bit-exact parity (KV-off
vs BDE_KV_ENABLE=1, max_abs_diff=0, PSNR=inf) with the cross-attn k_img/v_img pool
active. Image pool ~0.5 GiB.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../models/dreamzero/pipeline_dreamzero.py    | 69 +++++++++++++------
 .../experimental/bde/kv_cache/manager.py      | 36 +++++++++-
 vllm_omni/experimental/bde/runner.py          |  6 ++
 3 files changed, 87 insertions(+), 24 deletions(-)

diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index c9372a8187d..40be4a88e3e 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -134,25 +134,52 @@ def _kv_commit(self):
             s.commit_chunk()
 
     def _kv_get_cross(self, state, is_negative):
-        """Cross-attn cache: always model-local (even under BDE).
+        """Cross-attn cache: pool-backed when BDE active, model-local otherwise.
 
-        Upstream #4154 extended the I2V cross-attn cache to also store image-token
-        ``k_img``/``v_img``; BDE's cross-attn pool only stores text ``k``/``v``, so a
-        pool-backed dict makes the I2V forward ``KeyError`` on ``'k_img'``. BDE
-        manages *self*-attn KV via the paged pool; cross-attn uses the model's own
-        cache (now image-token-caching), which BDE's reset path already clears.
+        The BDE pool caches both text ``k``/``v`` and (for I2V) the image-token
+        ``k_img``/``v_img`` added by #4154, so the returned dict satisfies the I2V
+        forward without recomputation.
         """
+        s = self._bde_kv_state
+        if s is not None:
+            return s.get_cross_kv_caches(is_negative, lambda: state.get_crossattn_caches(is_negative))
         return state.get_crossattn_caches(is_negative)
 
-    def _kv_populate_cross(self, context: torch.Tensor, is_negative: bool) -> None:
-        """No-op: cross-attn is model-local (see :meth:`_kv_get_cross`).
+    def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None:
+        """Eagerly project cross-attn K/V for all layers into the BDE pool.
 
-        Kept as a stable hook. The BDE cross-attn pool is unused after upstream
-        #4154 added ``k_img``/``v_img`` caching, so there is nothing to eagerly
-        populate; the model's own cross-attn cache handles (and now caches) the
-        image-token projections.
+        Caches the session-invariant cross-attn projections once (re-run after a
+        window-boundary reset, when ``_cross_populated`` clears): the text ``k``/``v``
+        from ``text_embedding(context)``, and for I2V the image-token ``k_img``/
+        ``v_img`` from ``img_emb(clip_feature)`` (the 257 image tokens the forward
+        splits off, cached model-side by #4154). Must run after the image is encoded
+        so ``clip_feature`` is available.
         """
-        return
+        s = self._bde_kv_state
+        if s is None or s._cross_populated.get(is_negative, False):
+            return
+        projected = self.transformer.text_embedding(context)
+        img_ctx = None
+        if clip_feature is not None and getattr(self.transformer, "model_type", "t2v") == "i2v":
+            img_ctx = self.transformer.img_emb(clip_feature)
+        for i, block in enumerate(self.transformer.blocks):
+            ca = block.cross_attn
+            n, d = ca.tp_num_heads, ca.head_dim
+            k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
+            v = ca.v(projected).unflatten(2, (n, d))
+            k_img = v_img = None
+            if img_ctx is not None:
+                k_img = ca.norm_k_img(ca.k_img(img_ctx)).unflatten(2, (n, d))
+                v_img = ca.v_img(img_ctx).unflatten(2, (n, d))
+            s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img)
+        s._cross_populated[is_negative] = True
+        _log.info(
+            "BDE CROSS POPULATE [%s]: %d layers, text=%s img=%s",
+            "neg" if is_negative else "pos",
+            len(self.transformer.blocks),
+            tuple(context.shape),
+            None if img_ctx is None else tuple(img_ctx.shape),
+        )
 
     def _kv_reset(self, state, *, clear_video_latents: bool = True):
         """Reset model-local KV and, under BDE, the pooled session window too.
@@ -1094,14 +1121,6 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
                 neg_inputs["attention_mask"].to(device),
             )
 
-        # -- Eager cross-attn population (BDE only) ----------------------------
-        # Project encoder hidden states to K/V once and store in the BDE pool
-        # so every denoising step reads without recomputing.  Re-populated after
-        # each window-boundary reset (_cross_populated clears → guard re-triggers).
-        self._kv_populate_cross(prompt_embeds, is_negative=False)
-        if negative_prompt_embeds is not None:
-            self._kv_populate_cross(negative_prompt_embeds, is_negative=True)
-
         # Extract first/last frame for CLIP + VAE encoding
         if num_frames_raw == 4 or num_frames_raw == 9:
             image = videos[:, :, -1:].transpose(1, 2)
@@ -1118,6 +1137,14 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             state.clip_feas = clip_feas.to(dtype=image.dtype)
             state.ys = ys.to(dtype=image.dtype)
 
+            # Eager cross-attn population (BDE only): cache text + image-token K/V
+            # into the pool now that the image is encoded (clip_feas available).
+            # Runs on the first forward of a session and after each window-boundary
+            # reset (current_start_frame returns to 0); _cross_populated guards re-entry.
+            self._kv_populate_cross(prompt_embeds, state.clip_feas, is_negative=False)
+            if negative_prompt_embeds is not None:
+                self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True)
+
         if state.current_start_frame != 0:
             # Subsequent calls: encode current observation via VAE
             if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
diff --git a/vllm_omni/experimental/bde/kv_cache/manager.py b/vllm_omni/experimental/bde/kv_cache/manager.py
index 58c58de5d9a..1363b7a4397 100644
--- a/vllm_omni/experimental/bde/kv_cache/manager.py
+++ b/vllm_omni/experimental/bde/kv_cache/manager.py
@@ -50,6 +50,7 @@ def __init__(
         max_model_len: int,
         available_bytes: int,
         cross_attn_length: int = 0,
+        cross_attn_img_length: int = 0,
         device: torch.device | None = None,
     ) -> None:
         if not config.enable:
@@ -77,8 +78,15 @@ def __init__(
         # DreamZero I2V).  Both pools draw from the same GPU free memory but
         # the cross-attn pool is sized directly rather than via the block pool
         # budget (no eviction/chunk lifecycle needed).
+        self.cross_attn_img_length = cross_attn_img_length
         self._cross_k: list[torch.Tensor] = []
         self._cross_v: list[torch.Tensor] = []
+        # I2V image-token cross-attn pool. Like the text k/v, the image-token
+        # k_img/v_img are session-invariant (the conditioning image doesn't change),
+        # so they are cached once and read every denoise step — see
+        # WanI2VCrossAttention (#4154 caches these model-side too). Empty for T2V.
+        self._cross_k_img: list[torch.Tensor] = []
+        self._cross_v_img: list[torch.Tensor] = []
         if device is not None and cross_attn_length > 0:
             cross_shape = (2, cross_attn_length, num_kv_heads, head_size)
             bytes_per_element = dtype.itemsize
@@ -102,6 +110,16 @@ def __init__(
                 head_size,
                 cross_bytes / (1024 * 1024),
             )
+            if cross_attn_img_length > 0:
+                img_shape = (2, cross_attn_img_length, num_kv_heads, head_size)
+                for _ in range(num_layers):
+                    self._cross_k_img.append(torch.empty(img_shape, dtype=dtype, device=device))
+                    self._cross_v_img.append(torch.empty(img_shape, dtype=dtype, device=device))
+                _log.info(
+                    "BDE cross-attn IMG pool: %d layers × %d img-tok (I2V)",
+                    num_layers,
+                    cross_attn_img_length,
+                )
 
         self.spec = ChunkWindowSpec(
             block_size=block_size,
@@ -161,28 +179,40 @@ def write_cross_kv(
         is_negative: bool,
         k: torch.Tensor,
         v: torch.Tensor,
+        k_img: torch.Tensor | None = None,
+        v_img: torch.Tensor | None = None,
     ) -> None:
         """Write one layer's cross-attn K/V into the pool.
 
         ``k`` / ``v``: ``(B, text_len, tp_num_heads, head_dim)``; only batch-0
         is copied (B=1 for inference). The ``is_negative`` flag selects the
-        correct CFG branch slot.
+        correct CFG branch slot. ``k_img`` / ``v_img`` (I2V image tokens,
+        ``(B, 257, ...)``) are written when the image pool is allocated.
         """
         branch = 1 if is_negative else 0
         self._cross_k[layer_idx][branch].copy_(k[0])
         self._cross_v[layer_idx][branch].copy_(v[0])
+        if k_img is not None and self._cross_k_img:
+            self._cross_k_img[layer_idx][branch].copy_(k_img[0])
+            self._cross_v_img[layer_idx][branch].copy_(v_img[0])
 
     def read_cross_kv(self, layer_idx: int, is_negative: bool) -> dict:
         """Return a pool-backed cross-attn cache dict for one layer.
 
         The dict matches the ``{"is_init": True, "k": Tensor, "v": Tensor}``
         convention the cross-attention module expects — it reads from the pool
-        slice rather than from the lazy-initialised model-local dict.
+        slice rather than from the lazy-initialised model-local dict. For I2V,
+        ``k_img`` / ``v_img`` are added so the image-token cache (added by #4154)
+        reads from the pool too.
         """
         branch = 1 if is_negative else 0
         k = self._cross_k[layer_idx][branch].unsqueeze(0)  # (1, L, heads, dim)
         v = self._cross_v[layer_idx][branch].unsqueeze(0)
-        return {"is_init": True, "k": k, "v": v}
+        cache = {"is_init": True, "k": k, "v": v}
+        if self._cross_k_img:
+            cache["k_img"] = self._cross_k_img[layer_idx][branch].unsqueeze(0)
+            cache["v_img"] = self._cross_v_img[layer_idx][branch].unsqueeze(0)
+        return cache
 
     # -- request lifecycle ---------------------------------------------------
 
diff --git a/vllm_omni/experimental/bde/runner.py b/vllm_omni/experimental/bde/runner.py
index 2118a91a98c..72c190f604e 100644
--- a/vllm_omni/experimental/bde/runner.py
+++ b/vllm_omni/experimental/bde/runner.py
@@ -84,6 +84,7 @@ def build_kv_cache(
         max_model_len: int,
         available_bytes: int,
         cross_attn_length: int = 0,
+        cross_attn_img_length: int = 0,
     ) -> None:
         """Construct the BDEKVCache (preallocating GPU pools on ``self.device``).
 
@@ -101,6 +102,7 @@ def build_kv_cache(
             max_model_len=max_model_len,
             available_bytes=available_bytes,
             cross_attn_length=cross_attn_length,
+            cross_attn_img_length=cross_attn_img_length,
             device=self.device,
         )
         logger.info(
@@ -159,6 +161,9 @@ def _preallocate_kv_cache(self) -> None:
         # prepended in _forward_blocks are stripped inside the cross-attn
         # forward but are not part of the cached k/v).
         cross_attn_length = int(getattr(t, "text_len", 0))
+        # I2V image-token cross-attn length: the model splits context[:, :257] as the
+        # image tokens (the img_emb(clip_feature) output), caching k_img/v_img. T2V: 0.
+        cross_attn_img_length = 257 if getattr(t, "model_type", "t2v") == "i2v" else 0
 
         # Frame-granular paging: 1 block = 1 frame = frame_seqlen tokens, so the
         # resident window matches max_attention_size exactly (it need not be a whole
@@ -187,6 +192,7 @@ def _preallocate_kv_cache(self) -> None:
             max_model_len=1 << 20,
             available_bytes=free_bytes,
             cross_attn_length=cross_attn_length,
+            cross_attn_img_length=cross_attn_img_length,
         )
 
     def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:

From 4db8e8b637c9e89800abcdae520e458bd0ec2aa9 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 20 Jun 2026 13:05:35 +0000
Subject: [PATCH 29/51] [BDE] Guard diffusion_model_runner_cls override with
 isinstance(str)

The BDE runner override read `getattr(od_config, "diffusion_model_runner_cls",
None) or platform_hook()`. With a Mock od_config (upstream worker tests from #4282,
e.g. test_model_runner_resolved_via_platform / the cuda-profiler tests), the getattr
returns a truthy auto-Mock, shadowing the platform hook so it's never called. Only
honor the override when it's an explicit import-path string, so the platform default
path is unaffected.

Fixes the buildkite CPU "Simple" step (5 tests/diffusion failures); tests/bde +
worker tests green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 vllm_omni/diffusion/worker/diffusion_worker.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index 56a9793967d..0cc855df5c2 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -211,10 +211,14 @@ def __init__(
         self.init_device()
         # Create model runner. An od_config override (e.g. the BDE engine selecting
         # BDEModelRunner) takes precedence; unset falls back to the platform default,
-        # so existing models are unaffected.
+        # so existing models are unaffected. The override must be an explicit import-
+        # path string — guard with isinstance so a non-string (e.g. a Mock od_config
+        # in tests) doesn't shadow the platform hook.
+        runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None)
         model_runner_cls_path = (
-            getattr(self.od_config, "diffusion_model_runner_cls", None)
-            or current_omni_platform.get_diffusion_model_runner_cls()
+            runner_override
+            if isinstance(runner_override, str) and runner_override
+            else current_omni_platform.get_diffusion_model_runner_cls()
         )
         model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path)
         self.model_runner = model_runner_cls(

From 560ba4fdf4bfb91a9e5dda2a86bfb3819ef42cb0 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Wed, 24 Jun 2026 01:24:22 +0000
Subject: [PATCH 30/51] [BDE] Engine owns 100% of KV; consolidate kv_cache into
 4 modules
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two related changes to the BDE KV stack (both verified bit-identical):

Engine-owned KV (removes the last model-local fallbacks):
- Self-attn: build_window_slots returns an empty index for a 0-block window,
  so prefill-start gathers a zero-length window (the concat-identity the
  model expects) straight from the pool. Dropped BDEKVState.get_kv_caches'
  `_committed==0` fallback + its `fallback` param.
- Cross-attn: get_cross_kv_caches now asserts the cross pool is populated
  (raises) instead of falling back to model-local lazy projection.
- pipeline _kv_get/_kv_get_cross drop the model-local lambdas; _kv_create is
  skipped under BDE (pool owns allocation). Model + non-BDE path untouched.
- Verified: KV-off vs KV-on on DreamZero-DROID is bit-identical
  (PSNR=inf, SSIM~1.0, max_abs_diff=0.0, LPIPS=0.0).

Consolidate kv_cache 7 -> 4 files (pure move, public API unchanged):
- config.py (knob) · paged.py (engine-generic slot-mapping + pool I/O +
  chunk-window eviction spec) · manager.py (BDEKVCache + BDERequestAdapter +
  pool builders) · state.py (the model-facing BDEKVState bridge).
- Removed adapter.py / chunk_window.py / gather.py / pool.py / slot_mapping.py;
  __init__ re-exports the same names. Mirrors vLLM's granularity and keeps the
  generic-vs-bridge seam the Cosmos port reuses.

Tests: tests/bde 79 passed; full bde-or-dreamzero suite green on GPU.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_kv_cache.py                    |   2 +-
 tests/bde/test_kv_orchestration.py            |   2 +-
 tests/bde/test_pipeline_bridge.py             | 106 ++++++--
 .../models/dreamzero/pipeline_dreamzero.py    |  15 +-
 .../experimental/bde/kv_cache/__init__.py     |  24 +-
 .../experimental/bde/kv_cache/adapter.py      |  63 -----
 .../experimental/bde/kv_cache/chunk_window.py |  85 ------
 .../experimental/bde/kv_cache/manager.py      | 124 ++++++++-
 vllm_omni/experimental/bde/kv_cache/paged.py  | 257 ++++++++++++++++++
 vllm_omni/experimental/bde/kv_cache/pool.py   |  68 -----
 .../experimental/bde/kv_cache/slot_mapping.py |  71 -----
 .../bde/kv_cache/{gather.py => state.py}      | 135 ++-------
 vllm_omni/experimental/bde/runner.py          |   2 +-
 13 files changed, 505 insertions(+), 449 deletions(-)
 delete mode 100644 vllm_omni/experimental/bde/kv_cache/adapter.py
 delete mode 100644 vllm_omni/experimental/bde/kv_cache/chunk_window.py
 create mode 100644 vllm_omni/experimental/bde/kv_cache/paged.py
 delete mode 100644 vllm_omni/experimental/bde/kv_cache/pool.py
 delete mode 100644 vllm_omni/experimental/bde/kv_cache/slot_mapping.py
 rename vllm_omni/experimental/bde/kv_cache/{gather.py => state.py} (63%)

diff --git a/tests/bde/test_kv_cache.py b/tests/bde/test_kv_cache.py
index 57f6c244908..ad4e1bd1bea 100644
--- a/tests/bde/test_kv_cache.py
+++ b/tests/bde/test_kv_cache.py
@@ -20,7 +20,7 @@
     build_kv_manager,
     compute_num_blocks,
 )
-from vllm_omni.experimental.bde.kv_cache.chunk_window import chunk_window_skipped_tokens
+from vllm_omni.experimental.bde.kv_cache.paged import chunk_window_skipped_tokens
 
 BLOCK = 16
 
diff --git a/tests/bde/test_kv_orchestration.py b/tests/bde/test_kv_orchestration.py
index 266dc12566b..dc5099cfd6f 100644
--- a/tests/bde/test_kv_orchestration.py
+++ b/tests/bde/test_kv_orchestration.py
@@ -73,7 +73,7 @@ def test_state_close_frees_both_branch_blocks():
     evicted, close() must free the blocks both adapters hold, or session churn
     leaks pool ownership (review P1).
     """
-    from vllm_omni.experimental.bde.kv_cache.gather import BDEKVState
+    from vllm_omni.experimental.bde.kv_cache.state import BDEKVState
 
     kv = make_cache(chunk_size=16, window_chunks=2)
     free_total = kv.manager.block_pool.get_num_free_blocks()
diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index ff58962ffe8..203aa6f012a 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -1,17 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 """Tests for the BDEKVState pipeline bridge (Step 5)."""
 
+import pytest
 import torch
 
 from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.gather import BDEKVState
+from vllm_omni.experimental.bde.kv_cache.state import BDEKVState
 
 BLOCK = 16
 N_HEADS = 4
 HEAD_DIM = 64
 
 
-def make_state(num_layers=1, window_chunks=4):
+def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
     cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
     kv = BDEKVCache(
         cfg,
@@ -22,6 +23,7 @@ def make_state(num_layers=1, window_chunks=4):
         block_size=BLOCK,
         max_model_len=4096,
         available_bytes=1 << 24,
+        cross_attn_length=cross_attn_length,
         device=torch.device("cpu"),
     )
     pos = kv.begin_request("r-pos")
@@ -38,18 +40,22 @@ def _window(n_chunks):
 # the resident window back; output mirrors the model-local sliding window.
 
 
-def test_get_returns_fallback_when_empty():
+def test_get_returns_empty_window_when_nothing_committed():
+    """The engine owns the read end-to-end: with nothing committed the adapter has
+    no resident blocks, so gather returns zero-length windows — the concat-identity
+    the model's attention expects on the first step — with no model-local fallback."""
     _, st = make_state()
-    sentinel = ["FALLBACK"]
-    assert st.get_kv_caches(False, lambda: sentinel) is sentinel  # nothing committed
-    assert st.get_kv_caches(True, lambda: sentinel) is sentinel
+    for neg in (False, True):
+        windows = st.get_kv_caches(neg)
+        assert len(windows) == 1
+        assert windows[0].shape == (2, 1, 0, N_HEADS, HEAD_DIM)
 
 
 def test_paged_write_then_gather_roundtrips():
     _, st = make_state(num_layers=1)
     win = _window(2)  # this forward appends 2 frame-chunks
     st.update_kv_cache(0, win, False, seq_len=2 * BLOCK)
-    got = st.get_kv_caches(False, lambda: ["X"])[0]  # (2, 1, W, heads, head_dim)
+    got = st.get_kv_caches(False)[0]  # (2, 1, W, heads, head_dim)
     assert got.shape == (2, 1, 2 * BLOCK, N_HEADS, HEAD_DIM)
     # The gathered window equals the tokens we wrote (bit-exact).
     assert torch.allclose(got[0, 0], win[0, 0], atol=0)
@@ -61,24 +67,24 @@ def test_gather_is_memoized_within_a_forward():
     get_kv_caches must reuse one gather; a write-back (layer 0) invalidates it."""
     _, st = make_state(num_layers=2)
     st.update_kv_cache(0, _window(2), False, seq_len=2 * BLOCK)  # commit one chunk
-    first = st.get_kv_caches(False, lambda: ["X"])
-    again = st.get_kv_caches(False, lambda: ["X"])
+    first = st.get_kv_caches(False)
+    again = st.get_kv_caches(False)
     assert again is first  # cache HIT: identical object, no re-gather
     # A new write-back grows the window -> next get must re-gather a fresh object.
     st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
-    after = st.get_kv_caches(False, lambda: ["X"])
+    after = st.get_kv_caches(False)
     assert after is not first
     # Negative branch is cached independently.
     st.update_kv_cache(0, _window(1), True, seq_len=BLOCK)
-    neg1 = st.get_kv_caches(True, lambda: ["X"])
-    assert st.get_kv_caches(True, lambda: ["X"]) is neg1
+    neg1 = st.get_kv_caches(True)
+    assert st.get_kv_caches(True) is neg1
     assert neg1 is not after
 
 
 def test_reset_clears_gather_cache():
     _, st = make_state(num_layers=1)
     st.update_kv_cache(0, _window(2), False, seq_len=2 * BLOCK)
-    st.get_kv_caches(False, lambda: ["X"])
+    st.get_kv_caches(False)
     assert st._gather_cache[False] is not None
     st.reset()
     assert st._gather_cache == {False: None, True: None}
@@ -88,16 +94,16 @@ def test_eviction_bounds_resident_window():
     kv, st = make_state(num_layers=1, window_chunks=3)
     st.update_kv_cache(0, _window(3), False, seq_len=3 * BLOCK)  # resident 3/3
     st.update_kv_cache(0, _window(5), False, seq_len=2 * BLOCK)  # +2 -> evict to 3
-    got = st.get_kv_caches(False, lambda: ["X"])[0]
+    got = st.get_kv_caches(False)[0]
     assert got.shape[2] <= 3 * BLOCK  # window bounded
 
 
 def test_branches_are_independent():
     _, st = make_state()
     st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
-    # Negative branch never written -> still falls back.
-    assert st.get_kv_caches(True, lambda: ["NEG_EMPTY"]) == ["NEG_EMPTY"]
-    assert st.get_kv_caches(False, lambda: ["X"])[0].shape[2] == BLOCK
+    # Negative branch never written -> empty (zero-length) window.
+    assert st.get_kv_caches(True)[0].shape[2] == 0
+    assert st.get_kv_caches(False)[0].shape[2] == BLOCK
 
 
 def test_reset_clears_session_window():
@@ -111,15 +117,50 @@ def test_reset_clears_session_window():
 
     st.reset()
 
-    # Both branches fall back to empty again, and pool blocks are freed.
+    # Both branches gather empty windows again, and pool blocks are freed.
     assert st._committed == {False: 0, True: 0}
     assert st._pending == {False: [], True: []}
-    assert st.get_kv_caches(False, lambda: ["EMPTY"]) == ["EMPTY"]
-    assert st.get_kv_caches(True, lambda: ["EMPTY"]) == ["EMPTY"]
+    assert st.get_kv_caches(False)[0].shape[2] == 0
+    assert st.get_kv_caches(True)[0].shape[2] == 0
     assert kv.manager.block_pool.get_num_free_blocks() > free_before
     # Adapters are live again under the same ids — a post-reset write works.
     st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
-    assert st.get_kv_caches(False, lambda: ["X"])[0].shape[2] == BLOCK
+    assert st.get_kv_caches(False)[0].shape[2] == BLOCK
+
+
+def test_get_cross_kv_caches_raises_before_populate():
+    """Engine ownership guard: a cross-attn read before _kv_populate_cross must
+    fail loud rather than fall back to the model (which would self-project KV)."""
+    _, st = make_state()  # cross_attn_length == 0, nothing populated
+    for neg in (False, True):
+        with pytest.raises(RuntimeError, match="cross-attn read before"):
+            st.get_cross_kv_caches(neg)
+    # Guards the full AND: _cross_populated alone is not enough without a pool.
+    st._cross_populated[False] = True
+    with pytest.raises(RuntimeError, match="cross-attn read before"):
+        st.get_cross_kv_caches(False)
+
+
+def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
+    """Once the cross pool is filled and marked populated, the read returns the
+    pool-backed dicts (is_init=True) bit-exactly — no model fallback involved."""
+    L = 8
+    kv, st = make_state(num_layers=2, cross_attn_length=L)
+    written = []
+    for i in range(2):
+        k = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        v = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        kv.write_cross_kv(i, False, k, v)
+        written.append((k, v))
+    st._cross_populated[False] = True
+
+    out = st.get_cross_kv_caches(False)
+    assert len(out) == 2
+    for i, (k, v) in enumerate(written):
+        assert out[i]["is_init"] is True
+        assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM)
+        assert torch.equal(out[i]["k"], k)
+        assert torch.equal(out[i]["v"], v)
 
 
 def test_commit_is_noop():
@@ -139,16 +180,23 @@ def test_kv_state_noop():
         assert hasattr(p, attr), f"proxy method {attr} missing"
 
 
-def test_kv_create_always_initializes_model_local():
-    """_kv_create must call state.create_kv_caches (no recursion), and must do so
-    even under BDE so cross-attention (not managed by BDE) is initialized."""
+def test_kv_create_owned_by_engine_under_bde():
+    """_kv_create initializes model-local caches ONLY on the non-BDE path. Under
+    BDE the pool owns all allocation, so state.create_kv_caches must NOT be called."""
     from unittest.mock import MagicMock
 
     from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 
     p = DreamZeroPipeline.__new__(DreamZeroPipeline)
-    for bde_state in (None, object()):  # disabled and enabled both init model-local
-        p._bde_kv_state = bde_state
-        state = MagicMock()
-        p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
-        state.create_kv_caches.assert_called_once_with(1, "float32", "cpu", 24, 4, 64)
+
+    # Non-BDE: model-local caches are created (no recursion).
+    p._bde_kv_state = None
+    state = MagicMock()
+    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+    state.create_kv_caches.assert_called_once_with(1, "float32", "cpu", 24, 4, 64)
+
+    # BDE active: allocation is engine-owned -> no model-local creation.
+    p._bde_kv_state = object()
+    state = MagicMock()
+    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+    state.create_kv_caches.assert_not_called()
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 40be4a88e3e..44f1f980d6d 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -108,15 +108,16 @@ def _kv_get(self, state, is_negative):
         s = self._bde_kv_state
         if s is not None:
             _log.debug("BDE KV get (neg=%s): %d layers", is_negative, s.num_layers)
-            return s.get_kv_caches(is_negative, lambda: state.get_kv_caches(is_negative))
+            return s.get_kv_caches(is_negative)
         return state.get_kv_caches(is_negative)
 
     def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
-        # Always create the model-local caches for the non-BDE fallback path.
-        # Under BDE both self-attn and cross-attn KV are pool-backed — the
-        # state caches created here are unused when _bde_kv_state is set.
-        # Cross-attn is populated eagerly in _kv_populate_cross after text
-        # encoding; self-attn is routed through _kv_get / _kv_update.
+        # Under BDE the pool owns all KV allocation: self-attn is allocated per
+        # chunk in _kv_update and read via _kv_get (empty window at prefill start),
+        # cross-attn is populated eagerly in _kv_populate_cross. No model-local
+        # caches are created. Only the non-BDE path needs them.
+        if self._bde_kv_state is not None:
+            return
         state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
 
     def _kv_update(self, state, layer_idx, updated_kv, is_negative, seq_len=None):
@@ -142,7 +143,7 @@ def _kv_get_cross(self, state, is_negative):
         """
         s = self._bde_kv_state
         if s is not None:
-            return s.get_cross_kv_caches(is_negative, lambda: state.get_crossattn_caches(is_negative))
+            return s.get_cross_kv_caches(is_negative)
         return state.get_crossattn_caches(is_negative)
 
     def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None:
diff --git a/vllm_omni/experimental/bde/kv_cache/__init__.py b/vllm_omni/experimental/bde/kv_cache/__init__.py
index 040652ee758..673c541e824 100644
--- a/vllm_omni/experimental/bde/kv_cache/__init__.py
+++ b/vllm_omni/experimental/bde/kv_cache/__init__.py
@@ -4,21 +4,27 @@
 Thin glue over vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` /
 ``SlidingWindowManager``), used by the Block Diffusion Engine to manage KV for
 AR-diffusion models. See ``BDE_doc/dreamzero_kv_phase1_plan.md``.
+
+Layout: ``config`` (public knob) · ``paged`` (engine-generic paging mechanics +
+chunk-window eviction spec) · ``manager`` (the BDEKVCache orchestrator + its
+request adapter / pool builders) · ``state`` (the model-facing BDEKVState bridge).
 """
 
-from vllm_omni.experimental.bde.kv_cache.adapter import BDERequestAdapter
-from vllm_omni.experimental.bde.kv_cache.chunk_window import ChunkWindowManager, ChunkWindowSpec
 from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.gather import (
-    allocate_kv_pool,
-    pool_gather_window,
-    pool_write_chunk,
+from vllm_omni.experimental.bde.kv_cache.manager import (
+    BDEKVCache,
+    BDERequestAdapter,
+    build_kv_manager,
+    compute_num_blocks,
 )
-from vllm_omni.experimental.bde.kv_cache.manager import BDEKVCache
-from vllm_omni.experimental.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
-from vllm_omni.experimental.bde.kv_cache.slot_mapping import (
+from vllm_omni.experimental.bde.kv_cache.paged import (
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    allocate_kv_pool,
     chunk_slot_mapping,
     compute_slot_mapping,
+    pool_gather_window,
+    pool_write_chunk,
     resident_block_ids,
 )
 
diff --git a/vllm_omni/experimental/bde/kv_cache/adapter.py b/vllm_omni/experimental/bde/kv_cache/adapter.py
deleted file mode 100644
index 2a639e9ce0f..00000000000
--- a/vllm_omni/experimental/bde/kv_cache/adapter.py
+++ /dev/null
@@ -1,63 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Adapt a BDE diffusion request to vLLM's ``KVCacheManager`` request surface."""
-
-from __future__ import annotations
-
-from vllm.v1.request import RequestStatus
-
-
-class BDERequestAdapter:
-    """Duck-types the subset of ``vllm.v1.request.Request`` that the
-    ``KVCacheManager`` reads (``allocate_slots`` / ``get_computed_blocks`` /
-    ``free`` and the coordinator they call into).
-
-    It is intentionally NOT a full ``Request``. The conformance test exercises a
-    real ``KVCacheManager`` against this adapter so the surface cannot silently
-    drift across vLLM versions.
-
-    A BDE request advances one *chunk* at a time: ``allocate_slots`` is called
-    once per chunk and ``num_computed_tokens`` advances only when a chunk is
-    committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk
-    reuse the same slots.
-    """
-
-    def __init__(
-        self,
-        request_id: str,
-        *,
-        chunk_size: int,
-        prefill_prefix_tokens: int = 0,
-    ) -> None:
-        self.request_id = request_id
-        self._chunk_size = chunk_size
-        self._prefill = prefill_prefix_tokens
-        self._completed_chunks = 0
-        # Filled only when cross-request prefix reuse is enabled (Phase 3).
-        self.block_hashes: list = []
-        self.skip_reading_prefix_cache = True
-        self.num_preemptions = 0
-        # vLLM watermark gate reads this; map the request lifecycle onto it.
-        self.status = RequestStatus.WAITING
-
-    @property
-    def num_computed_tokens(self) -> int:
-        """Persistent KV already materialized (committed chunks + prefill)."""
-        return self._prefill + self._completed_chunks * self._chunk_size
-
-    @property
-    def num_tokens(self) -> int:
-        """Total tokens once the in-flight chunk is committed."""
-        return self._prefill + (self._completed_chunks + 1) * self._chunk_size
-
-    @property
-    def num_prompt_tokens(self) -> int:
-        """The prefill prefix length (read by ``cache_blocks`` when caching)."""
-        return self._prefill
-
-    @property
-    def completed_chunks(self) -> int:
-        return self._completed_chunks
-
-    def on_chunk_committed(self) -> None:
-        """Advance by one chunk. Call once per chunk, not per denoise step."""
-        self._completed_chunks += 1
diff --git a/vllm_omni/experimental/bde/kv_cache/chunk_window.py b/vllm_omni/experimental/bde/kv_cache/chunk_window.py
deleted file mode 100644
index 484c4782bda..00000000000
--- a/vllm_omni/experimental/bde/kv_cache/chunk_window.py
+++ /dev/null
@@ -1,85 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Chunk-granularity sliding-window KV spec + manager for the BDE engine.
-
-A world-model keeps only the last ``W`` chunks of KV resident. That is a sliding
-window whose unit is a *chunk*, expressed as a ``SlidingWindowSpec`` subclass with
-``sliding_window = window_chunks * chunk_size`` plus a manager that evicts at
-chunk boundaries. Memory policy, refcounting, and ``null_block`` replacement stay
-in vLLM's ``BlockPool`` — only the token-skip math is overridden here.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
-from vllm.v1.kv_cache_interface import SlidingWindowSpec
-from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
-
-
-def chunk_window_skipped_tokens(
-    num_computed_tokens: int,
-    *,
-    chunk_size: int,
-    sliding_window: int,
-    sink_chunks: int,
-    reset_at_boundary: bool,
-) -> int:
-    """Tokens outside the resident chunk window, snapped to a chunk boundary.
-
-    Pure function so the eviction policy is unit-testable without constructing a
-    manager. Two strategies:
-
-    - ``reset_at_boundary`` (DreamZero): at each chunk boundary everything past
-      the sink is dropped.
-    - otherwise (VGGT-style sliding replace): keep the last ``window`` tokens
-      (plus the sink); the skip count snaps down to a chunk boundary so a chunk
-      is never half-evicted.
-    """
-    sink = sink_chunks * chunk_size
-    if reset_at_boundary:
-        completed = (num_computed_tokens // chunk_size) * chunk_size
-        return max(0, completed - sink)
-    skipped = max(0, num_computed_tokens - sliding_window - sink)
-    return (skipped // chunk_size) * chunk_size
-
-
-class ChunkWindowManager(SlidingWindowManager):
-    """``SlidingWindowManager`` that evicts at chunk boundaries.
-
-    ``self.sliding_window`` is set by the base ``__init__``; the chunk fields are
-    read from ``self.kv_cache_spec`` (a :class:`ChunkWindowSpec`).
-    """
-
-    def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
-        spec = self.kv_cache_spec
-        return chunk_window_skipped_tokens(
-            num_computed_tokens,
-            chunk_size=spec.chunk_size,
-            sliding_window=self.sliding_window,
-            sink_chunks=spec.sink_chunks,
-            reset_at_boundary=spec.reset_at_boundary,
-        )
-
-
-# Register so KVCacheManager resolves ChunkWindowSpec to ChunkWindowManager.
-# Dispatch walks the spec's MRO, so without explicit registration the subclass
-# would silently fall back to the parent SlidingWindowManager (override ignored).
-# uniform_type_base_spec=None => its own KV cache group.
-@register_kv_cache_spec(manager_class=ChunkWindowManager, uniform_type_base_spec=None)
-@dataclass(frozen=True, kw_only=True)
-class ChunkWindowSpec(SlidingWindowSpec):
-    # sliding_window (inherited) MUST equal window_chunks * chunk_size.
-    chunk_size: int
-    window_chunks: int
-    sink_chunks: int = 0
-    reset_at_boundary: bool = False
-
-    def __post_init__(self):
-        super().__post_init__()
-        if self.sliding_window != self.window_chunks * self.chunk_size:
-            raise ValueError(
-                "ChunkWindowSpec.sliding_window must equal "
-                f"window_chunks * chunk_size ({self.window_chunks} * "
-                f"{self.chunk_size}), got {self.sliding_window}"
-            )
diff --git a/vllm_omni/experimental/bde/kv_cache/manager.py b/vllm_omni/experimental/bde/kv_cache/manager.py
index 1363b7a4397..2568d172f8a 100644
--- a/vllm_omni/experimental/bde/kv_cache/manager.py
+++ b/vllm_omni/experimental/bde/kv_cache/manager.py
@@ -12,24 +12,138 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 import torch
 from vllm.logger import init_logger
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheSpec,
+    KVCacheTensor,
+)
+from vllm.v1.request import RequestStatus
 
-from vllm_omni.experimental.bde.kv_cache.adapter import BDERequestAdapter
-from vllm_omni.experimental.bde.kv_cache.chunk_window import ChunkWindowSpec
 from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.gather import (
+from vllm_omni.experimental.bde.kv_cache.paged import (
+    ChunkWindowSpec,
     allocate_kv_pool,
     build_window_slots,
+    chunk_slot_mapping,
     pool_gather_window,
     pool_write_chunk,
+    resident_block_ids,
 )
-from vllm_omni.experimental.bde.kv_cache.pool import build_kv_manager, compute_num_blocks
-from vllm_omni.experimental.bde.kv_cache.slot_mapping import chunk_slot_mapping, resident_block_ids
 
 _log = init_logger(__name__)
 
 
+class BDERequestAdapter:
+    """Duck-types the subset of ``vllm.v1.request.Request`` that the
+    ``KVCacheManager`` reads (``allocate_slots`` / ``get_computed_blocks`` /
+    ``free`` and the coordinator they call into).
+
+    It is intentionally NOT a full ``Request``. The conformance test exercises a
+    real ``KVCacheManager`` against this adapter so the surface cannot silently
+    drift across vLLM versions.
+
+    A BDE request advances one *chunk* at a time: ``allocate_slots`` is called
+    once per chunk and ``num_computed_tokens`` advances only when a chunk is
+    committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk
+    reuse the same slots.
+    """
+
+    def __init__(
+        self,
+        request_id: str,
+        *,
+        chunk_size: int,
+        prefill_prefix_tokens: int = 0,
+    ) -> None:
+        self.request_id = request_id
+        self._chunk_size = chunk_size
+        self._prefill = prefill_prefix_tokens
+        self._completed_chunks = 0
+        # Filled only when cross-request prefix reuse is enabled (Phase 3).
+        self.block_hashes: list = []
+        self.skip_reading_prefix_cache = True
+        self.num_preemptions = 0
+        # vLLM watermark gate reads this; map the request lifecycle onto it.
+        self.status = RequestStatus.WAITING
+
+    @property
+    def num_computed_tokens(self) -> int:
+        """Persistent KV already materialized (committed chunks + prefill)."""
+        return self._prefill + self._completed_chunks * self._chunk_size
+
+    @property
+    def num_tokens(self) -> int:
+        """Total tokens once the in-flight chunk is committed."""
+        return self._prefill + (self._completed_chunks + 1) * self._chunk_size
+
+    @property
+    def num_prompt_tokens(self) -> int:
+        """The prefill prefix length (read by ``cache_blocks`` when caching)."""
+        return self._prefill
+
+    @property
+    def completed_chunks(self) -> int:
+        return self._completed_chunks
+
+    def on_chunk_committed(self) -> None:
+        """Advance by one chunk. Call once per chunk, not per denoise step."""
+        self._completed_chunks += 1
+
+
+def compute_num_blocks(
+    available_bytes: int,
+    gpu_memory_fraction: float,
+    page_size_bytes: int,
+) -> int:
+    """Number of KV blocks that fit in ``fraction`` of the memory budget."""
+    if page_size_bytes <= 0:
+        raise ValueError(f"page_size_bytes must be positive, got {page_size_bytes}")
+    if not 0.0 < gpu_memory_fraction <= 1.0:
+        raise ValueError(f"gpu_memory_fraction must be in (0, 1], got {gpu_memory_fraction}")
+    budget = int(available_bytes * gpu_memory_fraction)
+    return max(0, budget // page_size_bytes)
+
+
+def build_kv_manager(
+    spec: KVCacheSpec,
+    layer_names: Sequence[str],
+    num_blocks: int,
+    max_model_len: int,
+    *,
+    enable_caching: bool = False,
+) -> KVCacheManager:
+    """Build a ``KVCacheManager`` with a single KV cache group for ``spec``.
+
+    Args:
+        spec: The KV cache spec for the group (e.g. a ``ChunkWindowSpec``).
+        layer_names: Attention layers sharing this group's block table.
+        num_blocks: Total physical blocks in the pool.
+        max_model_len: Upper bound on a request's sequence length.
+        enable_caching: Cross-request prefix caching (Phase 3); off in Phase 1.
+    """
+    layer_names = list(layer_names)
+    group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
+    tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)]
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=tensors,
+        kv_cache_groups=[group],
+    )
+    return KVCacheManager(
+        config,
+        max_model_len=max_model_len,
+        scheduler_block_size=spec.block_size,
+        hash_block_size=spec.block_size,
+        enable_caching=enable_caching,
+    )
+
+
 class BDEKVCache:
     """Owns the paged KV pool + per-request lifecycle for a BDE model.
 
diff --git a/vllm_omni/experimental/bde/kv_cache/paged.py b/vllm_omni/experimental/bde/kv_cache/paged.py
new file mode 100644
index 00000000000..97da3c72992
--- /dev/null
+++ b/vllm_omni/experimental/bde/kv_cache/paged.py
@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generic paging mechanics + chunk-window eviction spec for the BDE engine.
+
+Engine-generic, model-agnostic primitives — the layer a second model (e.g. the
+Cosmos port) reuses unchanged. Three concerns live here:
+
+* **Slot mapping** — absolute token positions → physical KV-cache slots, the
+  standard PagedAttention layout ``slot(pos) = block_id(pos) * block_size +
+  (pos % block_size)``. Phase 1 is single-request and gathers the resident
+  window into the model's existing attention rather than calling vLLM's paged
+  kernel, so a thin slot-mapping helper is used instead of vLLM's ``BlockTables``.
+* **Pool I/O** — allocate per-layer flat K/V pools, write a committed chunk into
+  slots, and gather the resident window back as a contiguous tensor. The
+  per-forward gather copy is the deliberate Phase-1 simplification the fused
+  paged backend retires in Phase 2.
+* **Chunk-window eviction** — a ``SlidingWindowSpec`` subclass whose unit is a
+  *chunk* (``sliding_window = window_chunks * chunk_size``) plus a manager that
+  evicts at chunk boundaries. Memory policy / refcounting / ``null_block``
+  replacement stay in vLLM's ``BlockPool``; only the token-skip math is here.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import torch
+from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
+from vllm.v1.kv_cache_interface import SlidingWindowSpec
+from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+
+
+# ── Slot mapping ────────────────────────────────────────────────────────────
+
+
+def compute_slot_mapping(
+    block_ids: Sequence[int],
+    positions: torch.Tensor | Sequence[int],
+    block_size: int,
+) -> torch.Tensor:
+    """Map absolute token positions to physical KV-cache slots.
+
+    Args:
+        block_ids: physical block id per block index (the request's block table).
+        positions: absolute token positions to map (1-D).
+        block_size: tokens per block.
+
+    Returns:
+        ``LongTensor`` of physical slots, one per position.
+    """
+    if block_size <= 0:
+        raise ValueError(f"block_size must be positive, got {block_size}")
+    table = torch.as_tensor(block_ids, dtype=torch.long)
+    pos = torch.as_tensor(positions, dtype=torch.long)
+    block_index = torch.div(pos, block_size, rounding_mode="floor")
+    offset = pos % block_size
+    return table[block_index] * block_size + offset
+
+
+def chunk_slot_mapping(
+    block_ids: Sequence[int],
+    num_computed_tokens: int,
+    chunk_size: int,
+    block_size: int,
+) -> torch.Tensor:
+    """Slot mapping for the in-flight chunk's tokens (the commit write target).
+
+    The chunk occupies absolute positions
+    ``[num_computed_tokens, num_computed_tokens + chunk_size)``.
+    """
+    positions = torch.arange(
+        num_computed_tokens,
+        num_computed_tokens + chunk_size,
+        dtype=torch.long,
+    )
+    return compute_slot_mapping(block_ids, positions, block_size)
+
+
+def resident_block_ids(block_ids: Sequence[int], null_block_id: int) -> list[int]:
+    """Real (non-null) blocks currently resident, in table order.
+
+    These are the blocks the read path gathers the attention window from;
+    out-of-window positions are the shared ``null_block`` and are excluded.
+    """
+    return [int(b) for b in block_ids if int(b) != null_block_id]
+
+
+# ── Pool I/O (allocate / write / gather) ────────────────────────────────────
+
+
+def allocate_kv_pool(
+    num_blocks: int,
+    block_size: int,
+    num_layers: int,
+    num_kv_heads: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+    """Allocate per-layer paged K and V pools on *device*.
+
+    Each pool is a contiguous ``(num_blocks * block_size, num_kv_heads, head_dim)``
+    tensor — a flat address space indexed by ``slot = block_id * block_size + offset``.
+    The slot mapping (above) tells writes and gathers where each token lives.
+    """
+    k_pools: list[torch.Tensor] = []
+    v_pools: list[torch.Tensor] = []
+    for _ in range(num_layers):
+        k_pools.append(torch.empty(num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device))
+        v_pools.append(torch.empty(num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device))
+    return k_pools, v_pools
+
+
+def pool_write_chunk(
+    k_pool: torch.Tensor,
+    v_pool: torch.Tensor,
+    new_k: torch.Tensor,
+    new_v: torch.Tensor,
+    slot_mapping: torch.Tensor,
+) -> None:
+    """Write the committed chunk's K/V into pool slots (in place).
+
+    ``new_k`` / ``new_v`` shape: ``(batch, chunk_size, num_kv_heads, head_dim)``.
+    ``slot_mapping`` shape: ``(chunk_size,)`` — one physical slot per token position.
+    The batch dim is written to the same slots (multi-batch writes the same
+    positions; the caller is responsible for per-sequence bookkeeping).
+
+    This is the *commit* write — called once per chunk, after the last denoise step.
+    """
+    k_pool[slot_mapping] = new_k[0]  # batch index 0
+    v_pool[slot_mapping] = new_v[0]
+
+
+def build_window_slots(window_block_ids: list[int], block_size: int, device: torch.device) -> torch.Tensor:
+    """Flat slot index covering all resident blocks, in table order.
+
+    Vectorized (no Python loop): ``slot = block_id * block_size + offset`` for every
+    ``offset in [0, block_size)``. The window's blocks are identical across all layers
+    within a forward, so this is built once per forward and shared.
+    """
+    if not window_block_ids:
+        # No resident blocks (e.g. prefill start, before the first chunk is
+        # committed): an empty slot index gathers a zero-length window, which is
+        # the concat-identity the model's attention expects as its initial KV.
+        return torch.empty(0, dtype=torch.long, device=device)
+    starts = torch.tensor(window_block_ids, dtype=torch.long, device=device) * block_size
+    offsets = torch.arange(block_size, device=device)
+    return (starts[:, None] + offsets[None, :]).reshape(-1)
+
+
+def pool_gather_window(
+    k_pool: torch.Tensor,
+    v_pool: torch.Tensor,
+    window_block_ids: list[int],
+    block_size: int,
+    max_attention_size: int,
+    *,
+    slots: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Materialize the resident window K/V as a contiguous tensor.
+
+    Gathers all slots belonging to ``window_block_ids`` (in table order), trims to
+    the last ``max_attention_size`` tokens (matching the existing ``cat`` + slice),
+    and returns a ``(2, batch, window_tokens, n_heads, head_dim)`` tensor — exactly
+    the format DreamZero's attention expects as ``kv_cache``. Pass a precomputed
+    ``slots`` (from :func:`build_window_slots`) to share the index across layers.
+    """
+    if slots is None:
+        slots = build_window_slots(window_block_ids, block_size, k_pool.device)
+    # Gather then trim to the attention window.
+    wk = k_pool[slots]  # (window_slots, n_heads, head_dim)
+    wv = v_pool[slots]
+    wk = wk[-max_attention_size:]  # (window, n_heads, head_dim)
+    wv = wv[-max_attention_size:]
+    # Add the batch dim then stack as (2, batch, window, n_heads, head_dim).
+    wk = wk.unsqueeze(0)  # (1, window, n_heads, head_dim)
+    wv = wv.unsqueeze(0)
+    window = torch.stack([wk, wv], dim=0)  # (2, 1, window, n_heads, head_dim)
+    return window
+
+
+def _drop_batch(t: torch.Tensor) -> torch.Tensor:
+    """``(1, L, n_heads, head_dim)`` -> ``(L, n_heads, head_dim)``; pass 3-D through."""
+    if t.dim() == 4 and t.shape[0] == 1:
+        return t[0]
+    return t
+
+
+# ── Chunk-window eviction (spec + manager) ──────────────────────────────────
+
+
+def chunk_window_skipped_tokens(
+    num_computed_tokens: int,
+    *,
+    chunk_size: int,
+    sliding_window: int,
+    sink_chunks: int,
+    reset_at_boundary: bool,
+) -> int:
+    """Tokens outside the resident chunk window, snapped to a chunk boundary.
+
+    Pure function so the eviction policy is unit-testable without constructing a
+    manager. Two strategies:
+
+    - ``reset_at_boundary`` (DreamZero): at each chunk boundary everything past
+      the sink is dropped.
+    - otherwise (VGGT-style sliding replace): keep the last ``window`` tokens
+      (plus the sink); the skip count snaps down to a chunk boundary so a chunk
+      is never half-evicted.
+    """
+    sink = sink_chunks * chunk_size
+    if reset_at_boundary:
+        completed = (num_computed_tokens // chunk_size) * chunk_size
+        return max(0, completed - sink)
+    skipped = max(0, num_computed_tokens - sliding_window - sink)
+    return (skipped // chunk_size) * chunk_size
+
+
+class ChunkWindowManager(SlidingWindowManager):
+    """``SlidingWindowManager`` that evicts at chunk boundaries.
+
+    ``self.sliding_window`` is set by the base ``__init__``; the chunk fields are
+    read from ``self.kv_cache_spec`` (a :class:`ChunkWindowSpec`).
+    """
+
+    def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
+        spec = self.kv_cache_spec
+        return chunk_window_skipped_tokens(
+            num_computed_tokens,
+            chunk_size=spec.chunk_size,
+            sliding_window=self.sliding_window,
+            sink_chunks=spec.sink_chunks,
+            reset_at_boundary=spec.reset_at_boundary,
+        )
+
+
+# Register so KVCacheManager resolves ChunkWindowSpec to ChunkWindowManager.
+# Dispatch walks the spec's MRO, so without explicit registration the subclass
+# would silently fall back to the parent SlidingWindowManager (override ignored).
+# uniform_type_base_spec=None => its own KV cache group.
+@register_kv_cache_spec(manager_class=ChunkWindowManager, uniform_type_base_spec=None)
+@dataclass(frozen=True, kw_only=True)
+class ChunkWindowSpec(SlidingWindowSpec):
+    # sliding_window (inherited) MUST equal window_chunks * chunk_size.
+    chunk_size: int
+    window_chunks: int
+    sink_chunks: int = 0
+    reset_at_boundary: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.sliding_window != self.window_chunks * self.chunk_size:
+            raise ValueError(
+                "ChunkWindowSpec.sliding_window must equal "
+                f"window_chunks * chunk_size ({self.window_chunks} * "
+                f"{self.chunk_size}), got {self.sliding_window}"
+            )
diff --git a/vllm_omni/experimental/bde/kv_cache/pool.py b/vllm_omni/experimental/bde/kv_cache/pool.py
deleted file mode 100644
index 6619b8ecf0e..00000000000
--- a/vllm_omni/experimental/bde/kv_cache/pool.py
+++ /dev/null
@@ -1,68 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Construct a vLLM ``KVCacheManager`` for the BDE engine.
-
-Phase 1 uses static pool sizing (no profiling run): the worker passes the free
-device memory and a budget fraction. Construction is metadata-only — the backing
-KV tensors are allocated by the worker; here we only build the block bookkeeping
-so allocation / eviction / free flow through vLLM's ``KVCacheManager``.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Sequence
-
-from vllm.v1.core.kv_cache_manager import KVCacheManager
-from vllm.v1.kv_cache_interface import (
-    KVCacheConfig,
-    KVCacheGroupSpec,
-    KVCacheSpec,
-    KVCacheTensor,
-)
-
-
-def compute_num_blocks(
-    available_bytes: int,
-    gpu_memory_fraction: float,
-    page_size_bytes: int,
-) -> int:
-    """Number of KV blocks that fit in ``fraction`` of the memory budget."""
-    if page_size_bytes <= 0:
-        raise ValueError(f"page_size_bytes must be positive, got {page_size_bytes}")
-    if not 0.0 < gpu_memory_fraction <= 1.0:
-        raise ValueError(f"gpu_memory_fraction must be in (0, 1], got {gpu_memory_fraction}")
-    budget = int(available_bytes * gpu_memory_fraction)
-    return max(0, budget // page_size_bytes)
-
-
-def build_kv_manager(
-    spec: KVCacheSpec,
-    layer_names: Sequence[str],
-    num_blocks: int,
-    max_model_len: int,
-    *,
-    enable_caching: bool = False,
-) -> KVCacheManager:
-    """Build a ``KVCacheManager`` with a single KV cache group for ``spec``.
-
-    Args:
-        spec: The KV cache spec for the group (e.g. a ``ChunkWindowSpec``).
-        layer_names: Attention layers sharing this group's block table.
-        num_blocks: Total physical blocks in the pool.
-        max_model_len: Upper bound on a request's sequence length.
-        enable_caching: Cross-request prefix caching (Phase 3); off in Phase 1.
-    """
-    layer_names = list(layer_names)
-    group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
-    tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)]
-    config = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=tensors,
-        kv_cache_groups=[group],
-    )
-    return KVCacheManager(
-        config,
-        max_model_len=max_model_len,
-        scheduler_block_size=spec.block_size,
-        hash_block_size=spec.block_size,
-        enable_caching=enable_caching,
-    )
diff --git a/vllm_omni/experimental/bde/kv_cache/slot_mapping.py b/vllm_omni/experimental/bde/kv_cache/slot_mapping.py
deleted file mode 100644
index 6ca40138664..00000000000
--- a/vllm_omni/experimental/bde/kv_cache/slot_mapping.py
+++ /dev/null
@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Slot mapping + block-table access for the BDE engine.
-
-Phase 1 is single-request and owns its KV layout — it gathers the resident window
-into DreamZero's existing attention rather than calling vLLM's paged kernel — so a
-thin slot-mapping helper is used instead of vLLM's batch/Triton ``BlockTables``.
-The layout is the standard PagedAttention one:
-
-    slot(pos) = block_id(pos) * block_size + (pos % block_size)
-
-Write (commit) and read (gather) use the same formula, so the layout is
-self-consistent regardless of vLLM's internal kernel layout.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Sequence
-
-import torch
-
-
-def compute_slot_mapping(
-    block_ids: Sequence[int],
-    positions: torch.Tensor | Sequence[int],
-    block_size: int,
-) -> torch.Tensor:
-    """Map absolute token positions to physical KV-cache slots.
-
-    Args:
-        block_ids: physical block id per block index (the request's block table).
-        positions: absolute token positions to map (1-D).
-        block_size: tokens per block.
-
-    Returns:
-        ``LongTensor`` of physical slots, one per position.
-    """
-    if block_size <= 0:
-        raise ValueError(f"block_size must be positive, got {block_size}")
-    table = torch.as_tensor(block_ids, dtype=torch.long)
-    pos = torch.as_tensor(positions, dtype=torch.long)
-    block_index = torch.div(pos, block_size, rounding_mode="floor")
-    offset = pos % block_size
-    return table[block_index] * block_size + offset
-
-
-def chunk_slot_mapping(
-    block_ids: Sequence[int],
-    num_computed_tokens: int,
-    chunk_size: int,
-    block_size: int,
-) -> torch.Tensor:
-    """Slot mapping for the in-flight chunk's tokens (the commit write target).
-
-    The chunk occupies absolute positions
-    ``[num_computed_tokens, num_computed_tokens + chunk_size)``.
-    """
-    positions = torch.arange(
-        num_computed_tokens,
-        num_computed_tokens + chunk_size,
-        dtype=torch.long,
-    )
-    return compute_slot_mapping(block_ids, positions, block_size)
-
-
-def resident_block_ids(block_ids: Sequence[int], null_block_id: int) -> list[int]:
-    """Real (non-null) blocks currently resident, in table order.
-
-    These are the blocks the read path gathers the attention window from;
-    out-of-window positions are the shared ``null_block`` and are excluded.
-    """
-    return [int(b) for b in block_ids if int(b) != null_block_id]
diff --git a/vllm_omni/experimental/bde/kv_cache/gather.py b/vllm_omni/experimental/bde/kv_cache/state.py
similarity index 63%
rename from vllm_omni/experimental/bde/kv_cache/gather.py
rename to vllm_omni/experimental/bde/kv_cache/state.py
index 09a5e163314..2ad37517bd3 100644
--- a/vllm_omni/experimental/bde/kv_cache/gather.py
+++ b/vllm_omni/experimental/bde/kv_cache/state.py
@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Pool write + window gather for the BDE engine (Phase 1).
+"""BDEKVState — per-request bridge between a model rollout and the BDE KV pool.
+
+This is the one model-facing piece of the KV stack (everything in ``paged.py`` /
+``manager.py`` is engine-generic). A second model (e.g. the Cosmos port) adds a
+sibling here that reuses the same pool primitives.
 
 Phase 1 keeps the existing attention kernel unchanged — it only changes *where*
 the past-window KV comes from (the pool) and *where* the committed chunk's KV is
 stored (the pool). The model's ``self_attn`` still receives a contiguous
 ``(batch, window, n_heads, head_dim)`` K/V pair; the gather materializes that
-tensor from the resident pool blocks. The per-forward gather copy is the
-deliberate Phase-1 simplification the fused paged backend retires in Phase 2.
+tensor from the resident pool blocks.
 """
 
 from __future__ import annotations
@@ -16,103 +19,9 @@
 import torch
 from vllm.logger import init_logger
 
-_log = init_logger(__name__)
-
-
-def allocate_kv_pool(
-    num_blocks: int,
-    block_size: int,
-    num_layers: int,
-    num_kv_heads: int,
-    head_dim: int,
-    dtype: torch.dtype,
-    device: torch.device,
-) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
-    """Allocate per-layer paged K and V pools on *device*.
-
-    Each pool is a contiguous ``(num_blocks * block_size, num_kv_heads, head_dim)``
-    tensor — a flat address space indexed by ``slot = block_id * block_size + offset``.
-    The slot mapping (``bde.kv_cache.slot_mapping``) tells writes and gathers where
-    each token lives.
-    """
-    k_pools: list[torch.Tensor] = []
-    v_pools: list[torch.Tensor] = []
-    for _ in range(num_layers):
-        k_pools.append(torch.empty(num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device))
-        v_pools.append(torch.empty(num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device))
-    return k_pools, v_pools
-
-
-def pool_write_chunk(
-    k_pool: torch.Tensor,
-    v_pool: torch.Tensor,
-    new_k: torch.Tensor,
-    new_v: torch.Tensor,
-    slot_mapping: torch.Tensor,
-) -> None:
-    """Write the committed chunk's K/V into pool slots (in place).
-
-    ``new_k`` / ``new_v`` shape: ``(batch, chunk_size, num_kv_heads, head_dim)``.
-    ``slot_mapping`` shape: ``(chunk_size,)`` — one physical slot per token position.
-    The batch dim is written to the same slots (multi-batch writes the same
-    positions; the caller is responsible for per-sequence bookkeeping).
-
-    This is the *commit* write — called once per chunk, after the last denoise step.
-    """
-    k_pool[slot_mapping] = new_k[0]  # batch index 0
-    v_pool[slot_mapping] = new_v[0]
-
+from vllm_omni.experimental.bde.kv_cache.paged import _drop_batch, pool_write_chunk
 
-def build_window_slots(window_block_ids: list[int], block_size: int, device: torch.device) -> torch.Tensor:
-    """Flat slot index covering all resident blocks, in table order.
-
-    Vectorized (no Python loop): ``slot = block_id * block_size + offset`` for every
-    ``offset in [0, block_size)``. The window's blocks are identical across all layers
-    within a forward, so this is built once per forward and shared.
-    """
-    if not window_block_ids:
-        raise ValueError("window_block_ids is empty — no blocks are resident")
-    starts = torch.tensor(window_block_ids, dtype=torch.long, device=device) * block_size
-    offsets = torch.arange(block_size, device=device)
-    return (starts[:, None] + offsets[None, :]).reshape(-1)
-
-
-def pool_gather_window(
-    k_pool: torch.Tensor,
-    v_pool: torch.Tensor,
-    window_block_ids: list[int],
-    block_size: int,
-    max_attention_size: int,
-    *,
-    slots: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Materialize the resident window K/V as a contiguous tensor.
-
-    Gathers all slots belonging to ``window_block_ids`` (in table order), trims to
-    the last ``max_attention_size`` tokens (matching the existing ``cat`` + slice),
-    and returns a ``(2, batch, window_tokens, n_heads, head_dim)`` tensor — exactly
-    the format DreamZero's attention expects as ``kv_cache``. Pass a precomputed
-    ``slots`` (from :func:`build_window_slots`) to share the index across layers.
-    """
-    if slots is None:
-        slots = build_window_slots(window_block_ids, block_size, k_pool.device)
-    # Gather then trim to the attention window.
-    wk = k_pool[slots]  # (window_slots, n_heads, head_dim)
-    wv = v_pool[slots]
-    wk = wk[-max_attention_size:]  # (window, n_heads, head_dim)
-    wv = wv[-max_attention_size:]
-    # Add the batch dim then stack as (2, batch, window, n_heads, head_dim).
-    wk = wk.unsqueeze(0)  # (1, window, n_heads, head_dim)
-    wv = wv.unsqueeze(0)
-    window = torch.stack([wk, wv], dim=0)  # (2, 1, window, n_heads, head_dim)
-    return window
-
-
-def _drop_batch(t: torch.Tensor) -> torch.Tensor:
-    """``(1, L, n_heads, head_dim)`` -> ``(L, n_heads, head_dim)``; pass 3-D through."""
-    if t.dim() == 4 and t.shape[0] == 1:
-        return t[0]
-    return t
+_log = init_logger(__name__)
 
 
 class BDEKVState:
@@ -162,12 +71,11 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
     def _adapter(self, is_negative: bool):
         return self.neg if is_negative else self.pos
 
-    def get_kv_caches(self, is_negative: bool, fallback) -> list:
+    def get_kv_caches(self, is_negative: bool) -> list:
+        # Engine owns the read end-to-end: at prefill start the adapter has no
+        # resident blocks, so gather returns zero-length windows (the empty seed
+        # the model concatenates against) — no model-local fallback needed.
         branch = "neg" if is_negative else "pos"
-        if self._committed[is_negative] == 0:  # nothing committed yet (prefill start)
-            out = fallback()
-            _log.info("BDE GET   [%s] source=model-local-empty layers=%d", branch, len(out))
-            return out
         cached = self._gather_cache[is_negative] if self._memo_enabled else None
         if cached is not None:  # window unchanged this forward -> reuse (no re-gather)
             return cached
@@ -185,11 +93,20 @@ def get_kv_caches(self, is_negative: bool, fallback) -> list:
         )
         return windows
 
-    def get_cross_kv_caches(self, is_negative: bool, fallback) -> list[dict]:
-        """Return pool-backed cross-attn cache dicts, or fallback if not populated."""
-        if self._cross_populated[is_negative] and self.kv_cache.cross_attn_length > 0:
-            return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
-        return fallback()
+    def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
+        """Return pool-backed cross-attn cache dicts.
+
+        Under BDE the cross-attn pool is always populated by ``_kv_populate_cross``
+        before the first read, so this must never be reached unpopulated — if it
+        were, the model would lazily project and write cross KV itself, violating
+        engine ownership. Fail loud rather than fall back to the model.
+        """
+        if not (self._cross_populated[is_negative] and self.kv_cache.cross_attn_length > 0):
+            raise RuntimeError(
+                "BDE cross-attn read before _kv_populate_cross (neg=%s, cross_attn_length=%d) "
+                "— the engine must own all cross KV" % (is_negative, self.kv_cache.cross_attn_length)
+            )
+        return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
 
     def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative: bool, seq_len) -> None:
         branch = "neg" if is_negative else "pos"
diff --git a/vllm_omni/experimental/bde/runner.py b/vllm_omni/experimental/bde/runner.py
index 72c190f604e..40a74ba360b 100644
--- a/vllm_omni/experimental/bde/runner.py
+++ b/vllm_omni/experimental/bde/runner.py
@@ -25,7 +25,7 @@
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.gather import BDEKVState
+from vllm_omni.experimental.bde.kv_cache.state import BDEKVState
 from vllm_omni.experimental.bde.kv_cache.manager import BDEKVCache
 
 logger = init_logger(__name__)

From 0191f62f716e05dbec4984d01f11c2fd7f942460 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Wed, 24 Jun 2026 09:17:50 +0000
Subject: [PATCH 31/51] [BDE] Slim PR: drop benchmarks/perf scaffolding (kept
 local), trim tests, address review

Thin the PR per amy-why-3459's review (#4534), no KV/engine logic change:

- Untrack benchmarks/bde/ + the two perf examples (bde_perf_compare.py,
  bde_long_rollout.py) and remove the gpu_mem_stats worker RPC (its only
  caller). Files stay on disk locally via .git/info/exclude. Resolves [1]/[2]/[6].
- Drop _log = logger alias in pipeline_dreamzero.py; use logger. Resolves [5].
- Trim tests: drop test_validation_harness.py (only benchmarks consumer) and
  test_full_chain_gpu.py (2-GPU, covered on CPU); remove a few trivial/redundant
  cases in pipeline_bridge / kv_cache / engine_routing. 79 -> 58 tests, core KV
  correctness coverage preserved.

[3] (experimental/ naming) kept as a deliberate staging namespace; [4] already
fixed in 534ed0b70.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 benchmarks/bde/__init__.py                    |   9 --
 benchmarks/bde/artifacts.py                   | 122 ---------------
 benchmarks/bde/dreamzero_config.py            |  54 -------
 benchmarks/bde/metrics.py                     | 104 -------------
 benchmarks/bde/parity_check.py                | 127 ---------------
 benchmarks/bde/profile.py                     |  94 -----------
 .../dreamzero/bde_long_rollout.py             |  64 --------
 .../dreamzero/bde_perf_compare.py             | 139 -----------------
 tests/bde/test_engine_routing.py              |  14 --
 tests/bde/test_full_chain_gpu.py              | 126 ---------------
 tests/bde/test_kv_cache.py                    |  40 -----
 tests/bde/test_pipeline_bridge.py             |  17 --
 tests/bde/test_validation_harness.py          | 146 ------------------
 .../models/dreamzero/pipeline_dreamzero.py    |  11 +-
 .../models/dreamzero/video_export_worker.py   |  15 --
 15 files changed, 5 insertions(+), 1077 deletions(-)
 delete mode 100644 benchmarks/bde/__init__.py
 delete mode 100644 benchmarks/bde/artifacts.py
 delete mode 100644 benchmarks/bde/dreamzero_config.py
 delete mode 100644 benchmarks/bde/metrics.py
 delete mode 100644 benchmarks/bde/parity_check.py
 delete mode 100644 benchmarks/bde/profile.py
 delete mode 100644 examples/offline_inference/dreamzero/bde_long_rollout.py
 delete mode 100644 examples/offline_inference/dreamzero/bde_perf_compare.py
 delete mode 100644 tests/bde/test_full_chain_gpu.py
 delete mode 100644 tests/bde/test_validation_harness.py

diff --git a/benchmarks/bde/__init__.py b/benchmarks/bde/__init__.py
deleted file mode 100644
index 96af27c90eb..00000000000
--- a/benchmarks/bde/__init__.py
+++ /dev/null
@@ -1,9 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""BDE DreamZero validation harness — accuracy metrics, artifacts, parity check.
-
-Implements the plan in ``BDE_doc/dreamzero_kv_phase1_profiling_accuracy.md``:
-compare the BDE KV path (enabled) against the model-local path (disabled) on
-weighted DreamZero, save the generated videos + comparisons, and gate on
-PSNR/SSIM/LPIPS. The model run itself is delegated to the existing offline
-export script; the reusable pieces here are unit-tested on synthetic data.
-"""
diff --git a/benchmarks/bde/artifacts.py b/benchmarks/bde/artifacts.py
deleted file mode 100644
index ddeb484b6a4..00000000000
--- a/benchmarks/bde/artifacts.py
+++ /dev/null
@@ -1,122 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Save BDE DreamZero validation artifacts: videos, comparisons, run.json."""
-
-from __future__ import annotations
-
-import json
-import math
-from pathlib import Path
-
-import cv2
-import numpy as np
-
-
-def _u8(frames) -> np.ndarray:
-    a = np.asarray(frames)
-    if a.dtype != np.uint8:
-        a = (a * 255.0 if float(a.max()) <= 1.5 else a).clip(0, 255).astype(np.uint8)
-    return a
-
-
-def save_mp4(path, frames, fps: int = 5) -> Path:
-    """Write ``(T, H, W, C)`` RGB frames to mp4 (matches the export script writer)."""
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    frames = _u8(frames)
-    h, w = frames.shape[1:3]
-    writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), float(fps), (w, h))
-    if not writer.isOpened():
-        raise RuntimeError(f"failed to open video writer for {path}")
-    try:
-        for frame in frames:
-            writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-    finally:
-        writer.release()
-    return path
-
-
-def load_mp4(path) -> np.ndarray:
-    cap = cv2.VideoCapture(str(path))
-    frames = []
-    while True:
-        ok, frame = cap.read()
-        if not ok:
-            break
-        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-    cap.release()
-    if not frames:
-        raise RuntimeError(f"no frames decoded from {path}")
-    return np.stack(frames, axis=0)
-
-
-def side_by_side(*videos, gap: int = 4) -> np.ndarray:
-    """Stitch videos horizontally (``real | off | on``) into one clip."""
-    vids = [_u8(v) for v in videos]
-    n_frames = min(v.shape[0] for v in vids)
-    height = max(v.shape[1] for v in vids)
-    cols = []
-    for v in vids:
-        v = v[:n_frames]
-        if v.shape[1] < height:
-            pad = np.zeros((n_frames, height - v.shape[1], v.shape[2], 3), np.uint8)
-            v = np.concatenate([v, pad], axis=1)
-        cols.append(v)
-        cols.append(np.zeros((n_frames, height, gap, 3), np.uint8))
-    return np.concatenate(cols[:-1], axis=2)
-
-
-def diff_heatmap(ref, test) -> np.ndarray:
-    """Colorized per-pixel abs-diff video ``|test - ref|`` (RGB)."""
-    r = _u8(ref).astype(np.int16)
-    t = _u8(test).astype(np.int16)
-    n = min(r.shape[0], t.shape[0])
-    d = np.abs(t[:n] - r[:n]).max(axis=-1).astype(np.uint8)  # (T, H, W)
-    out = np.stack([cv2.applyColorMap(d[i], cv2.COLORMAP_JET) for i in range(n)])  # BGR
-    return out[..., ::-1].copy()  # -> RGB
-
-
-def metrics_plot(path, metrics: dict) -> Path:
-    import matplotlib
-
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
-
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    psnr = [min(p, 100.0) for p in metrics["psnr"]]  # cap inf for plotting
-    fig, ax = plt.subplots(2, 1, figsize=(8, 5), sharex=True)
-    ax[0].plot(psnr)
-    ax[0].axhline(40.0, ls="--", c="r")
-    ax[0].set_ylabel("PSNR (dB, capped 100)")
-    ax[1].plot(metrics["ssim"])
-    ax[1].axhline(0.99, ls="--", c="r")
-    ax[1].set_ylabel("SSIM")
-    ax[1].set_xlabel("frame")
-    fig.tight_layout()
-    fig.savefig(path, dpi=100)
-    plt.close(fig)
-    return path
-
-
-def _safe(value):
-    if isinstance(value, float) and (math.isinf(value) or math.isnan(value)):
-        if math.isnan(value):
-            return "NaN"
-        return "Infinity" if value > 0 else "-Infinity"
-    return value
-
-
-def write_run_json(path, **fields) -> Path:
-    """Persist reproducibility metadata + metrics + gates (inf/nan-safe)."""
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-
-    def clean(x):
-        if isinstance(x, dict):
-            return {k: clean(v) for k, v in x.items()}
-        if isinstance(x, (list, tuple)):
-            return [clean(v) for v in x]
-        return _safe(x)
-
-    path.write_text(json.dumps(clean(fields), indent=2))
-    return path
diff --git a/benchmarks/bde/dreamzero_config.py b/benchmarks/bde/dreamzero_config.py
deleted file mode 100644
index 16721b25ac2..00000000000
--- a/benchmarks/bde/dreamzero_config.py
+++ /dev/null
@@ -1,54 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Derive a parity-faithful ``BDEKVConfig`` from DreamZero geometry.
-
-For KV-on to match the model-local ``[-max_attention_size:]`` window exactly, the
-BDE window must equal ``max_attention_size = local_attn_size * frame_seqlen``. See
-``BDE_doc/dreamzero_kv_phase1_profiling_accuracy.md`` §1.
-"""
-
-from __future__ import annotations
-
-from vllm_omni.experimental.bde.kv_cache import BDEKVConfig
-
-
-def bde_config_for_dreamzero(
-    *,
-    num_frame_per_block: int,
-    frame_seqlen: int,
-    local_attn_size: int,
-    gpu_memory_fraction: float = 0.1,
-) -> BDEKVConfig:
-    """Build the parity-faithful ``BDEKVConfig`` for DreamZero.
-
-    Frame-granular paging (matches ``BDEModelRunner._preallocate_kv_cache``): one
-    pool block = one frame, so the resident window equals ``max_attention_size``
-    exactly without rounding to whole causal blocks.
-
-    ``chunk_size      = frame_seqlen``
-    ``window (tokens) = local_attn_size * frame_seqlen   (== max_attention_size)``
-    ``window_chunks   = local_attn_size``
-
-    ``num_frame_per_block`` no longer governs paging (it sets the reset cadence); it
-    is kept as an argument for call-site compatibility.
-    """
-    if num_frame_per_block <= 0 or frame_seqlen <= 0:
-        raise ValueError("num_frame_per_block and frame_seqlen must be > 0")
-    if local_attn_size <= 0:
-        raise ValueError(
-            f"local_attn_size must be > 0 for a bounded BDE window (got {local_attn_size}); "
-            "DreamZero -1 (full attention) needs an explicit window size for Phase 1"
-        )
-
-    chunk_size = frame_seqlen
-    window_chunks = local_attn_size
-    config = BDEKVConfig(
-        enable=True,
-        chunk_size=chunk_size,
-        window_chunks=window_chunks,
-        sink_chunks=0,
-        reset_at_boundary=False,
-        gpu_memory_fraction=gpu_memory_fraction,
-    )
-    # Invariant: the BDE window in tokens equals DreamZero's max_attention_size.
-    assert config.sliding_window == local_attn_size * frame_seqlen
-    return config
diff --git a/benchmarks/bde/metrics.py b/benchmarks/bde/metrics.py
deleted file mode 100644
index 2fe6a85e95c..00000000000
--- a/benchmarks/bde/metrics.py
+++ /dev/null
@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Accuracy metrics for BDE DreamZero validation (KV-on vs KV-off)."""
-
-from __future__ import annotations
-
-import numpy as np
-import torch
-from torchmetrics.functional.image import (
-    peak_signal_noise_ratio,
-    structural_similarity_index_measure,
-)
-
-
-def _to_nchw(frames) -> torch.Tensor:
-    """``(T, H, W, C)`` uint8/float -> ``(T, C, H, W)`` float in ``[0, 1]``."""
-    t = torch.as_tensor(np.asarray(frames))
-    if t.ndim != 4:
-        raise ValueError(f"expected (T, H, W, C) frames, got shape {tuple(t.shape)}")
-    t = t.float()
-    if float(t.max()) > 1.5:  # uint8 range
-        t = t / 255.0
-    return t.permute(0, 3, 1, 2).contiguous()
-
-
-_LPIPS = "unloaded"
-
-
-def _lpips_model():
-    """Lazily build an LPIPS metric; ``None`` if its weights can't load (offline)."""
-    global _LPIPS
-    if _LPIPS == "unloaded":
-        try:
-            from torchmetrics.image import LearnedPerceptualImagePatchSimilarity
-
-            _LPIPS = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).eval()
-        except Exception:  # weights unavailable / no network
-            _LPIPS = None
-    return _LPIPS
-
-
-def frame_metrics(ref, test) -> dict:
-    """Per-frame PSNR / SSIM (+ LPIPS if available) between two videos.
-
-    ``ref`` is the reference (KV-off / model-local); ``test`` is KV-on.
-    """
-    r, t = _to_nchw(ref), _to_nchw(test)
-    if r.shape != t.shape:
-        raise ValueError(f"shape mismatch ref={tuple(r.shape)} test={tuple(t.shape)}")
-    psnr, ssim = [], []
-    for i in range(r.shape[0]):
-        ri, ti = r[i : i + 1], t[i : i + 1]
-        psnr.append(float(peak_signal_noise_ratio(ti, ri, data_range=1.0)))
-        ssim.append(float(structural_similarity_index_measure(ti, ri, data_range=1.0)))
-    out = {
-        "num_frames": int(r.shape[0]),
-        "psnr": psnr,
-        "ssim": ssim,
-        "psnr_min": min(psnr),
-        "psnr_mean": sum(psnr) / len(psnr),
-        "ssim_min": min(ssim),
-        "ssim_mean": sum(ssim) / len(ssim),
-        "max_abs_diff": float((t - r).abs().max()),
-        "lpips": None,
-    }
-    model = _lpips_model()
-    if model is not None:
-        with torch.no_grad():
-            lp = [float(model(t[i : i + 1].clamp(0, 1), r[i : i + 1].clamp(0, 1))) for i in range(r.shape[0])]
-        out["lpips"] = lp
-        out["lpips_max"] = max(lp)
-        out["lpips_mean"] = sum(lp) / len(lp)
-    return out
-
-
-def tensor_parity(ref, test, *, atol: float = 1e-4, rtol: float = 1e-4) -> dict:
-    """Strict numerical parity for attention-input / latent tensors (Tier A.1)."""
-    r = torch.as_tensor(ref).float()
-    t = torch.as_tensor(test).float()
-    if r.shape != t.shape:
-        raise ValueError(f"shape mismatch ref={tuple(r.shape)} test={tuple(t.shape)}")
-    diff = (t - r).abs()
-    return {
-        "shape": list(r.shape),
-        "max_abs": float(diff.max()),
-        "mean_abs": float(diff.mean()),
-        "allclose": bool(torch.allclose(t, r, atol=atol, rtol=rtol)),
-        "atol": atol,
-        "rtol": rtol,
-    }
-
-
-def check_gates(metrics: dict, *, psnr_min: float = 40.0, ssim_min: float = 0.99, lpips_max: float = 0.01) -> dict:
-    """Apply the §5 acceptance gates to a ``frame_metrics`` result."""
-
-    def gate(value, threshold, ok):
-        return {"pass": ok, "value": value, "threshold": threshold}
-
-    gates = {
-        "psnr": gate(metrics["psnr_min"], psnr_min, metrics["psnr_min"] >= psnr_min),
-        "ssim": gate(metrics["ssim_min"], ssim_min, metrics["ssim_min"] >= ssim_min),
-    }
-    if metrics.get("lpips") is not None:
-        gates["lpips"] = gate(metrics["lpips_max"], lpips_max, metrics["lpips_max"] <= lpips_max)
-    return {"passed": all(g["pass"] for g in gates.values()), "gates": gates}
diff --git a/benchmarks/bde/parity_check.py b/benchmarks/bde/parity_check.py
deleted file mode 100644
index 7ee8280f838..00000000000
--- a/benchmarks/bde/parity_check.py
+++ /dev/null
@@ -1,127 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Tier-A numerical parity check for weighted DreamZero: KV-off vs KV-on.
-
-The model run is delegated to the existing offline export script; this module
-orchestrates the two runs, loads the saved videos, computes metrics, saves
-artifacts (side-by-side, diff heatmap, metrics plot, run.json), and gates on
-PSNR/SSIM/LPIPS.
-
-    python -m benchmarks.bde.parity_check \\
-        --deploy-off vllm_omni/deploy/dreamzero.yaml \\
-        --deploy-on  vllm_omni/deploy/dreamzero_bde.yaml \\
-        --input-video <dir> --session-id parity0 --out artifacts/dreamzero/parity
-
-``compare()`` is the model-free core (load -> metrics -> artifacts -> gates) and
-is unit-tested; ``run_export()`` shells out to the weighted model.
-"""
-
-from __future__ import annotations
-
-import argparse
-import subprocess
-import sys
-from pathlib import Path
-
-from benchmarks.bde import artifacts, metrics
-
-EXPORT_SCRIPT = "examples/offline_inference/dreamzero/export_prediction_video.py"
-
-
-def compare(
-    off_video,
-    on_video,
-    out_dir,
-    *,
-    fps: int = 5,
-    label: str = "parity",
-    real_video=None,
-    run_meta: dict | None = None,
-) -> dict:
-    """Load KV-off / KV-on videos, compute metrics, save artifacts, gate."""
-    out = Path(out_dir)
-    out.mkdir(parents=True, exist_ok=True)
-
-    off = artifacts.load_mp4(off_video)
-    on = artifacts.load_mp4(on_video)
-    frame_metrics = metrics.frame_metrics(off, on)
-    gates = metrics.check_gates(frame_metrics)
-
-    panes = [v for v in (artifacts.load_mp4(real_video) if real_video else None, off, on) if v is not None]
-    artifacts.save_mp4(out / f"{label}_sidebyside.mp4", artifacts.side_by_side(*panes), fps=fps)
-    artifacts.save_mp4(out / f"{label}_diff.mp4", artifacts.diff_heatmap(off, on), fps=fps)
-    artifacts.metrics_plot(out / f"{label}_metrics.png", frame_metrics)
-    artifacts.write_run_json(
-        out / "run.json",
-        label=label,
-        metrics=frame_metrics,
-        gates=gates,
-        meta=run_meta or {},
-    )
-    return {"metrics": frame_metrics, "gates": gates, "out_dir": str(out)}
-
-
-def run_export(model: str, deploy_config, session_id: str, out_stem, *, extra=()) -> Path:
-    """Run the offline DreamZero export (saves the generated video). Needs weights."""
-    out_stem = Path(out_stem)
-    out_stem.parent.mkdir(parents=True, exist_ok=True)
-    cmd = [
-        sys.executable,
-        EXPORT_SCRIPT,
-        "--model",
-        model,
-        "--deploy-config",
-        str(deploy_config),
-        "--session-id",
-        session_id,
-        "--output-stem",
-        str(out_stem),
-        "--save-gif",
-        "--save-input-video",
-        "--save-actions",
-        *extra,
-    ]
-    subprocess.run(cmd, check=True)
-    return out_stem.with_suffix(".mp4")
-
-
-def main(argv=None) -> int:
-    ap = argparse.ArgumentParser(description="DreamZero BDE KV parity check (KV-off vs KV-on)")
-    ap.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
-    ap.add_argument("--deploy-off", required=True, help="deploy config with BDE KV disabled (model-local)")
-    ap.add_argument("--deploy-on", required=True, help="deploy config with BDE KV enabled")
-    ap.add_argument("--session-id", default="parity0")
-    ap.add_argument("--out", default="artifacts/dreamzero/parity")
-    ap.add_argument("--fps", type=int, default=5)
-    ap.add_argument("--export-arg", action="append", default=[], help="extra arg forwarded to export")
-    args = ap.parse_args(argv)
-
-    base = Path(args.out)
-    sid = args.session_id
-    off = run_export(args.model, args.deploy_off, sid, base / "off" / sid / "pred", extra=args.export_arg)
-    on = run_export(args.model, args.deploy_on, sid, base / "on" / sid / "pred", extra=args.export_arg)
-
-    result = compare(
-        off,
-        on,
-        base / "compare",
-        fps=args.fps,
-        label=sid,
-        run_meta={
-            "model": args.model,
-            "deploy_off": args.deploy_off,
-            "deploy_on": args.deploy_on,
-            "session_id": sid,
-        },
-    )
-    m, g = result["metrics"], result["gates"]
-    status = "PASS" if g["passed"] else "FAIL"
-    print(
-        f"[parity] {status}  PSNR_min={m['psnr_min']:.2f}  "
-        f"SSIM_min={m['ssim_min']:.4f}  max_abs_diff={m['max_abs_diff']:.4f}"
-    )
-    print(f"[parity] artifacts: {result['out_dir']}")
-    return 0 if g["passed"] else 1
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
diff --git a/benchmarks/bde/profile.py b/benchmarks/bde/profile.py
deleted file mode 100644
index 50914b01242..00000000000
--- a/benchmarks/bde/profile.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Profiling utilities for BDE DreamZero (plan §3).
-
-``kv_memory_plateau`` is runnable without weights — it drives a ``BDEKVCache`` over
-many chunks and records pool occupancy, demonstrating the bounded-memory win (KV
-plateaus at the window vs the model-local grow-then-slice). Latency breakdown
-(allocate / write / gather / attention) requires the weighted model and is
-collected via ``record_function`` hooks during a real rollout — see the plan.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import torch
-
-from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
-
-
-def kv_memory_plateau(
-    config: BDEKVConfig,
-    *,
-    num_chunks: int,
-    num_layers: int,
-    num_kv_heads: int,
-    head_size: int,
-    block_size: int,
-    dtype: torch.dtype = torch.float16,
-    max_model_len: int = 1 << 20,
-    available_bytes: int = 1 << 32,
-    device: torch.device | None = None,
-) -> dict:
-    """Drive a ``BDEKVCache`` over ``num_chunks`` and record pool occupancy.
-
-    Returns per-chunk free / used / resident block counts and whether usage
-    plateaued once the window filled.
-    """
-    kv = BDEKVCache(
-        config,
-        num_layers=num_layers,
-        num_kv_heads=num_kv_heads,
-        head_size=head_size,
-        dtype=dtype,
-        block_size=block_size,
-        max_model_len=max_model_len,
-        available_bytes=available_bytes,
-        device=device,
-    )
-    total = kv.manager.block_pool.get_num_free_blocks()
-    adapter = kv.begin_request("profile")
-    free, resident = [], []
-    for _ in range(num_chunks):
-        kv.allocate_chunk(adapter)
-        free.append(kv.manager.block_pool.get_num_free_blocks())
-        resident.append(len(kv.window_block_ids(adapter)))
-        kv.commit_chunk(adapter)
-    kv.end_request(adapter)
-
-    used = [total - f for f in free]
-    window = config.window_chunks
-    plateaued = used[-1] == used[window] if (window is not None and num_chunks > window) else None
-    return {
-        "num_chunks": num_chunks,
-        "total_blocks": total,
-        "window_chunks": window,
-        "free": free,
-        "used_blocks": used,
-        "resident_blocks": resident,
-        "peak_used": max(used),
-        "plateaued": plateaued,
-    }
-
-
-def plot_memory(path, profile: dict) -> Path:
-    import matplotlib
-
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
-
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    fig, ax = plt.subplots(figsize=(8, 4))
-    ax.plot(profile["used_blocks"], label="pool blocks used (BDE)")
-    ax.plot(profile["resident_blocks"], label="resident (window) blocks")
-    if profile["window_chunks"] is not None:
-        ax.axhline(profile["window_chunks"] + 1, ls="--", c="g", label="window + 1")
-    ax.set_xlabel("chunk")
-    ax.set_ylabel("blocks")
-    ax.set_title("BDE KV pool occupancy (bounded by window)")
-    ax.legend()
-    fig.tight_layout()
-    fig.savefig(path, dpi=100)
-    plt.close(fig)
-    return path
diff --git a/examples/offline_inference/dreamzero/bde_long_rollout.py b/examples/offline_inference/dreamzero/bde_long_rollout.py
deleted file mode 100644
index 2d81345efc3..00000000000
--- a/examples/offline_inference/dreamzero/bde_long_rollout.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Long DreamZero rollout in one session — to exercise BDE KV eviction.
-
-The stock export submits 2 forwards (prefill + 1 chunk = ~12 frames), which never
-reaches the model's attention window (max_attention_size). This harness submits
-``--num-chunks`` chunk observations in a single session so the cumulative KV grows
-past the window and the model (and BDE) must evict — letting us check that BDE's
-paged eviction stays parity-exact with the model-local path under real eviction.
-
-KV-off vs KV-on is toggled by the BDE_KV_ENABLE env var (same as the export):
-
-    BDE_KV_ENABLE=1 CUDA_VISIBLE_DEVICES=0 HF_HOME=/models \\
-      python examples/offline_inference/dreamzero/bde_long_rollout.py \\
-      --num-chunks 8 --output outputs/bde_parity/long_on.mp4
-"""
-
-import argparse
-import os
-import sys
-from pathlib import Path
-
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import export_prediction_video as E  # noqa: E402
-
-
-def main() -> None:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
-    ap.add_argument("--deploy-config", type=Path, default=Path("vllm_omni/deploy/dreamzero.yaml"))
-    ap.add_argument("--video-dir", type=Path, default=Path("/models/dreamzero-assets"))
-    ap.add_argument("--prompt", default="A robot arm manipulates objects on a table.")
-    ap.add_argument("--session-id", default="long-rollout")
-    ap.add_argument("--num-chunks", type=int, default=8)
-    ap.add_argument("--output", type=Path, default=Path("outputs/bde_parity/long.mp4"))
-    ap.add_argument(
-        "--latents",
-        type=Path,
-        default=None,
-        help="if set, torch.save the raw decoded latents here for exact parity checks",
-    )
-    args = ap.parse_args()
-
-    # base = [prefill_obs, chunk_obs]; extend the chunk into a long session.
-    _, base = E._build_observations(args.video_dir, prompt=args.prompt, session_id=args.session_id)
-    observations = [base[0]] + [base[1]] * args.num_chunks
-    print(f"[long-rollout] submitting {len(observations)} forwards in session {args.session_id!r}")
-
-    omni, outputs = E._run_generation(args.model, args.deploy_config, observations)
-    latents = torch.cat([E._extract_latents(o) for o in outputs], dim=2)
-    if args.latents is not None:
-        args.latents.parent.mkdir(parents=True, exist_ok=True)
-        torch.save(latents.detach().cpu(), args.latents)
-        print(f"SAVED_LATENTS={args.latents}  shape={tuple(latents.shape)}")
-    frames = E._decode_with_worker(omni, latents)
-
-    args.output.parent.mkdir(parents=True, exist_ok=True)
-    E._write_mp4(args.output, frames, fps=5)
-    print(f"SAVED_MP4={args.output}  frames={frames.shape[0]}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/offline_inference/dreamzero/bde_perf_compare.py b/examples/offline_inference/dreamzero/bde_perf_compare.py
deleted file mode 100644
index ab4725eb086..00000000000
--- a/examples/offline_inference/dreamzero/bde_perf_compare.py
+++ /dev/null
@@ -1,139 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Timed DreamZero rollout for BDE-vs-main comparison.
-
-Runs one session of (prefill + N chunk) forwards, timing each `omni.generate`
-(the per-forward E2E), saving the decoded latents for a precision check, and
-dumping a timing JSON. Works on both the BDE branch (with BDE_KV_ENABLE=1) and a
-clean main checkout (baseline) — it only uses helpers from export_prediction_video,
-which exist on both.
-
-    BDE_KV_ENABLE=1 CUDA_VISIBLE_DEVICES=0 HF_HOME=/models \\
-      python examples/offline_inference/dreamzero/bde_perf_compare.py \\
-      --num-chunks 12 --tag bde \\
-      --latents outputs/bde_parity/perf_bde.pt --timing outputs/bde_parity/perf_bde.json
-"""
-
-import argparse
-import json
-import os
-import sys
-import time
-from pathlib import Path
-
-import torch
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import export_prediction_video as E  # noqa: E402
-
-
-def main() -> None:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
-    ap.add_argument("--deploy-config", type=Path, default=Path("vllm_omni/deploy/dreamzero.yaml"))
-    ap.add_argument("--video-dir", type=Path, default=Path("/models/dreamzero-assets"))
-    ap.add_argument("--prompt", default="A robot arm manipulates objects on a table.")
-    ap.add_argument("--session-id", default="perf-cmp")
-    ap.add_argument("--num-chunks", type=int, default=12)
-    ap.add_argument("--tag", default="run")
-    ap.add_argument("--latents", type=Path, default=None)
-    ap.add_argument("--video", type=Path, default=None)
-    ap.add_argument("--fps", type=int, default=5)
-    ap.add_argument("--timing", type=Path, default=None)
-    args = ap.parse_args()
-
-    _, observations = E._build_observations(
-        args.video_dir,
-        prompt=args.prompt,
-        session_id=args.session_id,
-        num_chunks=args.num_chunks,
-        repeat_chunk_observations=True,
-    )
-    n = len(observations)
-    print(
-        f"[perf:{args.tag}] {n} forwards (1 prefill + {args.num_chunks} chunk) "
-        f"BDE_KV_ENABLE={os.environ.get('BDE_KV_ENABLE', '0')}"
-    )
-
-    t0 = time.perf_counter()
-    omni = E.Omni(
-        model=args.model,
-        deploy_config=str(args.deploy_config),
-        enforce_eager=True,
-        worker_extension_cls=E.WORKER_EXTENSION,
-    )
-    load_s = time.perf_counter() - t0
-    print(f"[perf:{args.tag}] engine load: {load_s:.2f}s")
-
-    outputs = []
-    per_forward = []
-    for index, obs in enumerate(observations):
-        sp = E.OmniDiffusionSamplingParams(
-            extra_args={"reset": index == 0, "session_id": obs["session_id"], "robot_obs": obs}
-        )
-        s = time.perf_counter()
-        result = omni.generate(obs["prompt"], sampling_params_list=[sp])  # blocking -> E2E
-        dt = time.perf_counter() - s
-        if not result:
-            raise RuntimeError(f"No output for forward {index}")
-        outputs.append(result[0])
-        per_forward.append(dt)
-        print(f"[perf:{args.tag}] forward {index:2d} ({'prefill' if index == 0 else 'chunk'}): {dt * 1000:8.1f} ms")
-
-    prefill = per_forward[0]
-    steady = per_forward[2:] if len(per_forward) > 2 else per_forward[1:]  # drop prefill + 1 warmup
-    mean_steady = sum(steady) / len(steady)
-
-    # Peak GPU memory over the whole rollout, queried from the worker process where
-    # the model + KV pools live. Monotonic high-water mark from process start.
-    peak_reserved_gib = None
-    peak_allocated_gib = None
-    try:
-        engine = getattr(omni.engine.stage_clients[0], "_engine", None)
-        if engine is not None:
-            stats = engine.executor.collective_rpc("gpu_mem_stats", unique_reply_rank=0, exec_all_ranks=True)
-            stats = stats[0] if isinstance(stats, (list, tuple)) else stats
-            peak_reserved_gib = stats["peak_reserved_gib"]
-            peak_allocated_gib = stats["peak_allocated_gib"]
-    except Exception as e:  # noqa: BLE001 — profiling extra must never fail the run
-        print(f"[perf:{args.tag}] WARN: could not read GPU mem stats: {e}")
-
-    summary = {
-        "tag": args.tag,
-        "bde_kv_enable": os.environ.get("BDE_KV_ENABLE", "0"),
-        "bde_kv_no_memo": os.environ.get("BDE_KV_NO_MEMO", "0"),
-        "num_forwards": n,
-        "load_s": load_s,
-        "prefill_s": prefill,
-        "steady_mean_s": mean_steady,
-        "steady_min_s": min(steady),
-        "steady_max_s": max(steady),
-        "total_gen_s": sum(per_forward),
-        "peak_reserved_gib": peak_reserved_gib,
-        "peak_allocated_gib": peak_allocated_gib,
-        "per_forward_s": per_forward,
-    }
-    mem_str = f"  peak_reserved={peak_reserved_gib:.2f}GiB" if peak_reserved_gib is not None else ""
-    print(
-        f"[perf:{args.tag}] prefill={prefill * 1000:.1f}ms  "
-        f"steady_mean={mean_steady * 1000:.1f}ms (n={len(steady)})  total={sum(per_forward):.2f}s{mem_str}"
-    )
-
-    # Decode once (after timing) so latents + the final video are both saved.
-    latents = torch.cat([E._extract_latents(o) for o in outputs], dim=2)
-    if args.latents is not None:
-        args.latents.parent.mkdir(parents=True, exist_ok=True)
-        torch.save(latents.detach().cpu(), args.latents)
-        print(f"[perf:{args.tag}] SAVED_LATENTS={args.latents} shape={tuple(latents.shape)}")
-    if args.video is not None:
-        frames = E._decode_with_worker(omni, latents)
-        args.video.parent.mkdir(parents=True, exist_ok=True)
-        E._write_mp4(args.video, frames, fps=args.fps)
-        print(f"[perf:{args.tag}] SAVED_MP4={args.video} frames={frames.shape[0]}")
-    if args.timing is not None:
-        args.timing.parent.mkdir(parents=True, exist_ok=True)
-        args.timing.write_text(json.dumps(summary, indent=2))
-        print(f"[perf:{args.tag}] SAVED_TIMING={args.timing}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/bde/test_engine_routing.py b/tests/bde/test_engine_routing.py
index 5dea9d4f184..5db07684c69 100644
--- a/tests/bde/test_engine_routing.py
+++ b/tests/bde/test_engine_routing.py
@@ -35,15 +35,6 @@ def test_default_resolves_to_diffusion_engine():
     assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
 
 
-def test_missing_field_resolves_to_diffusion_engine():
-    # Backward compatibility: configs predating the engine_backend field.
-    assert DiffusionEngine.resolve_engine_class(SimpleNamespace()) is DiffusionEngine
-
-
-def test_none_resolves_to_diffusion_engine():
-    assert DiffusionEngine.resolve_engine_class(_cfg(None)) is DiffusionEngine
-
-
 def test_bde_key_resolves_to_bde_engine():
     assert DiffusionEngine.resolve_engine_class(_cfg("bde")) is BDEEngine
 
@@ -153,11 +144,6 @@ def test_apply_runner_default_respects_explicit_choice():
     assert cfg.diffusion_model_runner_cls == "my.custom.Runner"
 
 
-def test_config_runner_cls_field_defaults_none():
-    field = {f.name: f for f in fields(OmniDiffusionConfig)}["diffusion_model_runner_cls"]
-    assert field.default is None
-
-
 def test_worker_runner_selection_prefers_override():
     # Mirrors the worker hook: an od_config override wins over the platform default.
     over = SimpleNamespace(diffusion_model_runner_cls=BDE_MODEL_RUNNER_CLS)
diff --git a/tests/bde/test_full_chain_gpu.py b/tests/bde/test_full_chain_gpu.py
deleted file mode 100644
index 11e88c67ae5..00000000000
--- a/tests/bde/test_full_chain_gpu.py
+++ /dev/null
@@ -1,126 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Full-chain BDE wiring + pool test on real GPUs (GPU 2,3).
-
-Spawns a rank-0 / rank-1 ``DiffusionWorker`` pair (world_size=2, NCCL rendezvous)
-with the BDE runner override, asserts both workers build ``BDEModelRunner``, then
-constructs a ``BDEKVCache`` with GPU-resident pools and runs a write→gather
-roundtrip on each rank. No model weights needed.
-
-Run explicitly (needs >= 4 GPUs):
-    CUDA_VISIBLE_DEVICES=2,3 python -m pytest tests/bde/test_full_chain_gpu.py -v -s
-"""
-
-import multiprocessing as mp
-
-import pytest
-import torch
-
-requires_gpus_2_3 = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.accelerator.device_count() < 2,
-    reason="requires >= 2 GPUs",
-)
-
-RUNNER_OVERRIDE = "vllm_omni.experimental.bde.runner.BDEModelRunner"
-BLOCK = 16
-N_HEADS = 4
-HEAD_DIM = 64
-
-
-def _worker_entry(rank: int, od_config, q) -> None:
-    try:
-        from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
-
-        worker = DiffusionWorker(local_rank=rank, rank=rank, od_config=od_config, skip_load_model=True)
-        device = str(worker.device)
-        runner_name = type(worker.model_runner).__name__
-
-        # --- pool write + gather on the worker's GPU ------------------------
-        from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
-
-        cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
-        kv = BDEKVCache(
-            cfg,
-            num_layers=1,
-            num_kv_heads=N_HEADS,
-            head_size=HEAD_DIM,
-            dtype=torch.float32,
-            block_size=BLOCK,
-            max_model_len=512,
-            available_bytes=1 << 28,
-            device=worker.device,
-        )
-
-        # Verify pools are on the worker's GPU.
-        assert kv._k_pools[0].device == worker.device
-        assert kv._v_pools[0].device == worker.device
-        assert kv._k_pools[0].shape == (kv.num_blocks * BLOCK, N_HEADS, HEAD_DIM)
-
-        # Write a chunk → gather window → verify roundtrip.
-        adapter = kv.begin_request("r")
-        kv.allocate_chunk(adapter)
-        write_slots = kv.chunk_write_slots(adapter)
-
-        new_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, device=worker.device)
-        new_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, device=worker.device)
-        kv.write_chunk_kv(0, new_k, new_v, adapter)
-        kv.commit_chunk(adapter)
-
-        # Read back from the pool at the write slots.
-        pool_read_k = kv._k_pools[0][write_slots]
-        pool_read_v = kv._v_pools[0][write_slots]
-        assert torch.allclose(pool_read_k, new_k[0], atol=1e-5), f"gathered K ≠ written (rank {rank})"
-        assert torch.allclose(pool_read_v, new_v[0], atol=1e-5), f"gathered V ≠ written (rank {rank})"
-
-        # Allocate a second chunk then gather the window.
-        kv.allocate_chunk(adapter)
-        window = kv.gather_window(0, adapter)
-        assert window.device == worker.device
-        assert window.shape == (2, 1, kv.spec.sliding_window, N_HEADS, HEAD_DIM)
-
-        q.put((rank, device, runner_name, "pool_roundtrip_ok", None))
-    except Exception:  # pragma: no cover - reported to the parent
-        import traceback
-
-        q.put((rank, None, None, None, traceback.format_exc()))
-
-
-@requires_gpus_2_3
-def test_bde_full_chain_two_gpu_with_pools():
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-    od_config = OmniDiffusionConfig.from_kwargs(
-        model="dummy",
-        model_class_name="DummyPipeline",
-        num_gpus=2,
-        parallel_config={"cfg_parallel_size": 2},
-        diffusion_load_format="dummy",
-        diffusion_model_runner_cls=RUNNER_OVERRIDE,
-    )
-    assert od_config.diffusion_model_runner_cls == RUNNER_OVERRIDE
-
-    ctx = mp.get_context("spawn")
-    q = ctx.Queue()
-    procs = [ctx.Process(target=_worker_entry, args=(rank, od_config, q)) for rank in (0, 1)]
-    for p in procs:
-        p.start()
-
-    results = {}
-    try:
-        for _ in range(2):
-            rank, device, runner, pool_status, err = q.get(timeout=300)
-            results[rank] = (device, runner, pool_status, err)
-    finally:
-        for p in procs:
-            p.join(timeout=60)
-            if p.is_alive():
-                p.terminate()
-
-    for rank in (0, 1):
-        assert rank in results, f"rank {rank} produced no result (hung?)"
-        device, runner, pool_status, err = results[rank]
-        assert err is None, f"rank {rank} failed:\n{err}"
-        assert runner == "BDEModelRunner", f"rank {rank} built {runner}, expected BDEModelRunner"
-        assert pool_status == "pool_roundtrip_ok"
-
-    assert results[0][0] == "cuda:0"
-    assert results[1][0] == "cuda:1"
diff --git a/tests/bde/test_kv_cache.py b/tests/bde/test_kv_cache.py
index ad4e1bd1bea..56b25e80857 100644
--- a/tests/bde/test_kv_cache.py
+++ b/tests/bde/test_kv_cache.py
@@ -106,10 +106,6 @@ def test_kv_config_sliding_window_property():
     assert BDEKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
 
 
-def test_kv_config_disabled_by_default():
-    assert BDEKVConfig().enable is False
-
-
 # --- BDERequestAdapter ------------------------------------------------------
 
 
@@ -141,13 +137,6 @@ def test_compute_num_blocks():
     assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
 
 
-def test_compute_num_blocks_validates():
-    with pytest.raises(ValueError):
-        compute_num_blocks(1 << 20, 0.5, 0)
-    with pytest.raises(ValueError):
-        compute_num_blocks(1 << 20, 1.5, 16 << 10)
-
-
 def test_build_manager_allocate_free_roundtrip():
     """End-to-end: a BDERequestAdapter drives a real KVCacheManager.
 
@@ -165,32 +154,3 @@ def test_build_manager_allocate_free_roundtrip():
 
     mgr.free(adapter)
     assert mgr.block_pool.get_num_free_blocks() == free_before
-
-
-def test_window_eviction_plateaus_pool_memory():
-    """Multi-chunk rollout: out-of-window blocks are recycled, so pool memory
-    plateaus at ~window size instead of growing with rollout length.
-
-    This is the core promise of the chunk window — drive a real KVCacheManager
-    over many chunks and assert the free-block count stops decreasing once the
-    window is full, and the resident (non-null) blocks stay bounded.
-    """
-    spec = make_spec(chunk_size=BLOCK, window_chunks=2)  # window = 2 chunks
-    mgr = build_kv_manager(spec, ["l0"], num_blocks=64, max_model_len=4096)
-    null_id = mgr.block_pool.null_block.block_id
-    adapter = BDERequestAdapter("req", chunk_size=BLOCK)
-
-    free_after = []
-    distinct_real = []
-    for _ in range(10):
-        mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
-        ids = mgr.get_block_ids(adapter.request_id)[0]
-        distinct_real.append(len({b for b in ids if b != null_id}))
-        free_after.append(mgr.block_pool.get_num_free_blocks())
-        adapter.on_chunk_committed()
-
-    # Free-block count stops decreasing once the window is full (warmup = window).
-    assert free_after[-1] == free_after[spec.window_chunks]
-    assert free_after[-1] == free_after[-5]  # flat across the tail
-    # Resident non-null blocks bounded by window + the in-flight chunk.
-    assert max(distinct_real) <= spec.window_chunks + 1
diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index 203aa6f012a..29b6206b7d5 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -163,23 +163,6 @@ def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
         assert torch.equal(out[i]["v"], v)
 
 
-def test_commit_is_noop():
-    _, st = make_state()
-    st.commit_chunk()  # paged mode keeps resident windows; commit just logs
-
-
-def test_kv_state_noop():
-    """When _bde_kv_state is None, proxy methods fall through — verified by
-    the existing tests passing unchanged (the pipeline's default __init__
-    sets _bde_kv_state = None)."""
-    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
-
-    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
-    assert p._bde_kv_state is None
-    for attr in ("_kv_get", "_kv_create", "_kv_update", "_kv_commit"):
-        assert hasattr(p, attr), f"proxy method {attr} missing"
-
-
 def test_kv_create_owned_by_engine_under_bde():
     """_kv_create initializes model-local caches ONLY on the non-BDE path. Under
     BDE the pool owns all allocation, so state.create_kv_caches must NOT be called."""
diff --git a/tests/bde/test_validation_harness.py b/tests/bde/test_validation_harness.py
deleted file mode 100644
index 2e778f915f2..00000000000
--- a/tests/bde/test_validation_harness.py
+++ /dev/null
@@ -1,146 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Tests for the BDE DreamZero validation harness (metrics / artifacts / config).
-
-Exercises everything except the weighted model run (synthetic videos/tensors).
-"""
-
-import numpy as np
-import pytest
-import torch
-
-from benchmarks.bde import artifacts, metrics
-from benchmarks.bde.dreamzero_config import bde_config_for_dreamzero
-from benchmarks.bde.parity_check import compare
-
-
-def _vid(t=6, h=32, w=32, seed=0):
-    rng = np.random.default_rng(seed)
-    return rng.integers(0, 256, (t, h, w, 3)).astype(np.uint8)
-
-
-# --- metrics ----------------------------------------------------------------
-
-
-def test_frame_metrics_identical():
-    v = _vid()
-    m = metrics.frame_metrics(v, v)
-    assert m["max_abs_diff"] == 0.0
-    assert m["ssim_min"] == pytest.approx(1.0)
-    assert m["psnr_min"] == float("inf")
-
-
-def test_frame_metrics_detects_difference():
-    v = _vid()
-    w = v.copy()
-    w[:, :16] = 0
-    m = metrics.frame_metrics(v, w)
-    assert m["max_abs_diff"] > 0
-    assert m["ssim_min"] < 1.0
-    assert m["psnr_min"] < float("inf")
-
-
-def test_tensor_parity():
-    a = torch.zeros(2, 3, 4)
-    assert metrics.tensor_parity(a, a.clone())["allclose"]
-    b = a.clone()
-    b[0, 0, 0] = 1.0
-    r = metrics.tensor_parity(a, b)
-    assert not r["allclose"] and r["max_abs"] == 1.0
-
-
-def test_check_gates():
-    assert metrics.check_gates({"psnr_min": 50.0, "ssim_min": 0.999})["passed"]
-    assert not metrics.check_gates({"psnr_min": 30.0, "ssim_min": 0.999})["passed"]
-    assert not metrics.check_gates({"psnr_min": 50.0, "ssim_min": 0.5})["passed"]
-
-
-# --- config derivation ------------------------------------------------------
-
-
-def test_config_derivation_matches_window():
-    # Frame-granular: one pool block = one frame, so chunk_size == frame_seqlen and
-    # window_chunks == local_attn_size (matches BDEModelRunner._preallocate_kv_cache,
-    # which overwrites chunk_size to frame_seqlen but preserves window_chunks).
-    cfg = bde_config_for_dreamzero(num_frame_per_block=3, frame_seqlen=220, local_attn_size=21)
-    assert cfg.chunk_size == 220
-    assert cfg.window_chunks == 21
-    assert cfg.sliding_window == 21 * 220  # == max_attention_size
-    assert cfg.enable and not cfg.reset_at_boundary
-
-
-def test_config_window_is_frame_granular():
-    # Regression: frame granularity makes the window exact for any
-    # num_frame_per_block (no divisibility constraint), and window_chunks tracks
-    # local_attn_size rather than local_attn_size // num_frame_per_block — the old
-    # causal-block math the runtime would silently shrink to a too-small window.
-    cfg = bde_config_for_dreamzero(num_frame_per_block=4, frame_seqlen=10, local_attn_size=21)
-    assert cfg.chunk_size == 10
-    assert cfg.window_chunks == 21
-    assert cfg.sliding_window == 21 * 10
-
-
-def test_config_rejects_full_attention():
-    with pytest.raises(ValueError):
-        bde_config_for_dreamzero(num_frame_per_block=3, frame_seqlen=10, local_attn_size=-1)
-
-
-# --- artifacts --------------------------------------------------------------
-
-
-def test_save_load_mp4_roundtrip(tmp_path):
-    v = _vid()
-    path = artifacts.save_mp4(tmp_path / "a.mp4", v, fps=5)
-    back = artifacts.load_mp4(path)
-    assert back.shape == v.shape  # mp4v is lossy; shape must match
-
-
-def test_side_by_side_and_diff_shapes():
-    a, b = _vid(seed=1), _vid(seed=2)
-    sbs = artifacts.side_by_side(a, b)
-    assert sbs.shape[0] == a.shape[0]
-    assert sbs.shape[2] >= a.shape[2] * 2  # two panes + gap
-    d = artifacts.diff_heatmap(a, b)
-    assert d.shape[:3] == a.shape[:3] and d.shape[3] == 3
-
-
-def test_write_run_json_handles_inf(tmp_path):
-    import json
-
-    p = artifacts.write_run_json(tmp_path / "run.json", metrics={"psnr_min": float("inf")})
-    data = json.loads(p.read_text())
-    assert data["metrics"]["psnr_min"] == "Infinity"
-
-
-# --- end-to-end compare pipeline (no model) ---------------------------------
-
-
-def test_compare_pipeline_produces_artifacts(tmp_path):
-    v = _vid()
-    off = artifacts.save_mp4(tmp_path / "off.mp4", v, fps=5)
-    on = artifacts.save_mp4(tmp_path / "on.mp4", v, fps=5)
-    res = compare(off, on, tmp_path / "cmp", label="t")
-    out = tmp_path / "cmp"
-    assert (out / "run.json").exists()
-    assert (out / "t_sidebyside.mp4").exists()
-    assert (out / "t_diff.mp4").exists()
-    assert (out / "t_metrics.png").exists()
-    # Same source through the same codec -> near-perfect parity, gates pass.
-    assert res["metrics"]["ssim_min"] > 0.9
-    assert res["gates"]["passed"]
-
-
-# --- profiling (memory plateau, no model) -----------------------------------
-
-
-def test_kv_memory_plateau(tmp_path):
-    from benchmarks.bde.dreamzero_config import bde_config_for_dreamzero
-    from benchmarks.bde.profile import kv_memory_plateau, plot_memory
-
-    cfg = bde_config_for_dreamzero(num_frame_per_block=1, frame_seqlen=16, local_attn_size=2)
-    prof = kv_memory_plateau(cfg, num_chunks=10, num_layers=2, num_kv_heads=4, head_size=64, block_size=16)
-    # Usage plateaus once the window fills; resident bounded by window + 1.
-    assert prof["plateaued"] is True
-    assert max(prof["resident_blocks"]) <= cfg.window_chunks + 1
-    assert prof["peak_used"] == prof["used_blocks"][cfg.window_chunks]
-    out = plot_memory(tmp_path / "mem.png", prof)
-    assert out.exists()
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 44f1f980d6d..362e0e5b371 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -56,7 +56,6 @@
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
-_log = logger
 MAX_DREAMZERO_SESSIONS = 64
 
 
@@ -107,7 +106,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     def _kv_get(self, state, is_negative):
         s = self._bde_kv_state
         if s is not None:
-            _log.debug("BDE KV get (neg=%s): %d layers", is_negative, s.num_layers)
+            logger.debug("BDE KV get (neg=%s): %d layers", is_negative, s.num_layers)
             return s.get_kv_caches(is_negative)
         return state.get_kv_caches(is_negative)
 
@@ -123,7 +122,7 @@ def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, he
     def _kv_update(self, state, layer_idx, updated_kv, is_negative, seq_len=None):
         s = self._bde_kv_state
         if s is not None:
-            _log.debug("BDE KV write: layer %d neg=%s shape=%s", layer_idx, is_negative, tuple(updated_kv.shape))
+            logger.debug("BDE KV write: layer %d neg=%s shape=%s", layer_idx, is_negative, tuple(updated_kv.shape))
             s.update_kv_cache(layer_idx, updated_kv, is_negative, seq_len)
         else:
             state.update_kv_cache(layer_idx, updated_kv, is_negative=is_negative)
@@ -131,7 +130,7 @@ def _kv_update(self, state, layer_idx, updated_kv, is_negative, seq_len=None):
     def _kv_commit(self):
         s = self._bde_kv_state
         if s is not None:
-            _log.debug("BDE KV commit: pos=%d chunks neg=%d chunks", s.pos.completed_chunks, s.neg.completed_chunks)
+            logger.debug("BDE KV commit: pos=%d chunks neg=%d chunks", s.pos.completed_chunks, s.neg.completed_chunks)
             s.commit_chunk()
 
     def _kv_get_cross(self, state, is_negative):
@@ -174,7 +173,7 @@ def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: b
                 v_img = ca.v_img(img_ctx).unflatten(2, (n, d))
             s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img)
         s._cross_populated[is_negative] = True
-        _log.info(
+        logger.info(
             "BDE CROSS POPULATE [%s]: %d layers, text=%s img=%s",
             "neg" if is_negative else "pos",
             len(self.transformer.blocks),
@@ -448,7 +447,7 @@ def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor
             state = kwargs.get("dreamzero_state", self.state)
             is_neg = kwargs.get("is_negative", False)
             if self._bde_kv_state is not None:
-                _log.info(
+                logger.info(
                     "BDE pipeline predict_noise -> write-back: is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
                     is_neg,
                     kwargs.get("seq_len"),
diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
index 530a3fbec89..d039e53c816 100644
--- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
+++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
@@ -9,21 +9,6 @@
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
-    def gpu_mem_stats(self) -> dict:
-        """Peak GPU memory (GiB) for this worker's device, for profiling.
-
-        ``max_memory_reserved`` is the caching-allocator high-water mark (the best
-        in-process proxy for "peak VRAM"); ``max_memory_allocated`` is the live
-        tensor high-water mark. Both are monotonic from process start, so the
-        end-of-run values capture the whole rollout's peak.
-        """
-        dev = torch.accelerator.current_device_index()
-        return {
-            "device": int(dev),
-            "peak_reserved_gib": torch.accelerator.max_memory_reserved(dev) / (1024**3),
-            "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3),
-        }
-
     @staticmethod
     def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
         decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()

From 2cb7e6565d53282e3dd761f13714c4f13d6b44a4 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Thu, 25 Jun 2026 07:23:03 +0000
Subject: [PATCH 32/51] [BDE] Fix cross-attn KV budget double-count + assert
 frame-aligned seq_len (review)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Addresses zwhzzz0821's review on #4534:

- manager.py: the cross-attn pools (text + I2V image) are torch.empty-allocated
  directly from GPU memory, but their bytes were not subtracted before sizing
  the self-attn paged pool — so total allocation could exceed gpu_memory_fraction
  of free memory and OOM. Now deduct cross_total_bytes from available_bytes
  before compute_num_blocks, so cross + fraction*(free - cross) <= free.
- state.py: update_kv_cache assumed seq_len is a multiple of chunk_size. In
  DreamZero it always is (whole frames, chunk_size == frame_seqlen), but a
  non-frame-aligned model would silently drop the leading remainder. Added an
  explicit assert to fail loud.

Tests: +2 regression tests; tests/bde 60 passed. No behavior change for
DreamZero (assert never fires; cross deduction marginally shrinks the pool),
so the bit-identical KV-off/KV-on parity is unaffected.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/bde/test_kv_cache.py                    | 25 +++++++++++++++++++
 tests/bde/test_pipeline_bridge.py             |  8 ++++++
 .../experimental/bde/kv_cache/manager.py      | 25 +++++++++++--------
 vllm_omni/experimental/bde/kv_cache/state.py  |  7 ++++++
 4 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/tests/bde/test_kv_cache.py b/tests/bde/test_kv_cache.py
index 56b25e80857..845e291897c 100644
--- a/tests/bde/test_kv_cache.py
+++ b/tests/bde/test_kv_cache.py
@@ -13,6 +13,7 @@
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.experimental.bde.kv_cache import (
+    BDEKVCache,
     BDEKVConfig,
     BDERequestAdapter,
     ChunkWindowManager,
@@ -154,3 +155,27 @@ def test_build_manager_allocate_free_roundtrip():
 
     mgr.free(adapter)
     assert mgr.block_pool.get_num_free_blocks() == free_before
+
+
+def test_cross_attn_pool_deducted_from_self_attn_budget():
+    """The cross-attn pool is allocated directly; its bytes are subtracted from
+    the self-attn paged-pool budget so the two together stay within the free
+    memory budget (review: zwhzzz0821)."""
+    L = 512
+    avail = 1 << 30  # 1 GiB
+    kv = BDEKVCache(
+        BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
+        num_layers=2,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=avail,
+        cross_attn_length=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)
+    assert kv.num_blocks == expected
+    assert expected > 2 * (2 + 1) + 2  # above the min-blocks floor, so the deduction is what's tested
diff --git a/tests/bde/test_pipeline_bridge.py b/tests/bde/test_pipeline_bridge.py
index 29b6206b7d5..ba56aa6f6a5 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/bde/test_pipeline_bridge.py
@@ -183,3 +183,11 @@ def test_kv_create_owned_by_engine_under_bde():
     state = MagicMock()
     p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
     state.create_kv_caches.assert_not_called()
+
+
+def test_update_kv_cache_requires_frame_aligned_seqlen():
+    """Engine guards the frame-alignment invariant (review: zwhzzz0821) — a
+    non-multiple seq_len would otherwise silently drop the leading remainder."""
+    _, st = make_state()
+    with pytest.raises(AssertionError, match="frame-aligned"):
+        st.update_kv_cache(0, _window(1), False, seq_len=BLOCK + 1)
diff --git a/vllm_omni/experimental/bde/kv_cache/manager.py b/vllm_omni/experimental/bde/kv_cache/manager.py
index 2568d172f8a..50cb5e62510 100644
--- a/vllm_omni/experimental/bde/kv_cache/manager.py
+++ b/vllm_omni/experimental/bde/kv_cache/manager.py
@@ -201,18 +201,18 @@ def __init__(
         # WanI2VCrossAttention (#4154 caches these model-side too). Empty for T2V.
         self._cross_k_img: list[torch.Tensor] = []
         self._cross_v_img: list[torch.Tensor] = []
+
+        # Bytes the cross-attn pools consume directly (K+V, pos+neg, all layers).
+        # Deducted from the self-attn paged-pool budget below so the two
+        # allocations together stay within the GPU free-memory budget.
+        def _cross_pool_bytes(length: int) -> int:
+            return 2 * 2 * length * num_kv_heads * head_size * dtype.itemsize * num_layers
+
+        cross_total_bytes = 0
         if device is not None and cross_attn_length > 0:
             cross_shape = (2, cross_attn_length, num_kv_heads, head_size)
-            bytes_per_element = dtype.itemsize
-            cross_bytes = (
-                2  # K + V
-                * 2  # pos + neg
-                * cross_attn_length
-                * num_kv_heads
-                * head_size
-                * bytes_per_element
-                * num_layers
-            )
+            cross_bytes = _cross_pool_bytes(cross_attn_length)
+            cross_total_bytes += cross_bytes
             for _ in range(num_layers):
                 self._cross_k.append(torch.empty(cross_shape, dtype=dtype, device=device))
                 self._cross_v.append(torch.empty(cross_shape, dtype=dtype, device=device))
@@ -226,6 +226,7 @@ def __init__(
             )
             if cross_attn_img_length > 0:
                 img_shape = (2, cross_attn_img_length, num_kv_heads, head_size)
+                cross_total_bytes += _cross_pool_bytes(cross_attn_img_length)
                 for _ in range(num_layers):
                     self._cross_k_img.append(torch.empty(img_shape, dtype=dtype, device=device))
                     self._cross_v_img.append(torch.empty(img_shape, dtype=dtype, device=device))
@@ -248,8 +249,10 @@ def __init__(
         )
         # Each pool block spans all layers' K/V, so size against the per-layer
         # page size times the layer count.
+        # Size the self-attn pool against the memory left after the cross-attn
+        # pools (allocated above), so cross + self-attn stays within the budget.
         num_blocks = compute_num_blocks(
-            available_bytes,
+            max(0, available_bytes - cross_total_bytes),
             config.gpu_memory_fraction,
             self.spec.page_size_bytes * num_layers,
         )
diff --git a/vllm_omni/experimental/bde/kv_cache/state.py b/vllm_omni/experimental/bde/kv_cache/state.py
index 2ad37517bd3..7aaac1bf1c9 100644
--- a/vllm_omni/experimental/bde/kv_cache/state.py
+++ b/vllm_omni/experimental/bde/kv_cache/state.py
@@ -117,6 +117,13 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
             # (seq_len), not the cumulative-minus-committed delta. Allocate those
             # frame-blocks once per forward (shared across layers).
             new_count = int(seq_len)
+            # DreamZero appends whole frames each forward and chunk_size ==
+            # frame_seqlen, so seq_len is always a multiple of cs. Guard the
+            # invariant: a non-frame-aligned model would otherwise silently drop
+            # the leading new_count % cs tokens and desync num_computed.
+            assert new_count % cs == 0, (
+                f"BDE expects frame-aligned seq_len (multiple of chunk_size={cs}), got {new_count}"
+            )
             n_chunks = new_count // cs
             slots = []
             for _ in range(n_chunks):

From a58d3e66d6944b84647b9c249b6e1a717967c3aa Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Thu, 25 Jun 2026 07:32:17 +0000
Subject: [PATCH 33/51] [BDE] Fix .gitignore negation path for the
 experimental/bde/kv_cache source dir

The blanket `kv_cache/` ignore had a stale un-ignore (`!vllm_omni/bde/kv_cache/`)
from before the experimental/ move, so the real source dir
`vllm_omni/experimental/bde/kv_cache/` stayed ignored and new files there
required `git add -f`. Point the negation at the current path.

Runtime kv_cache/ dirs elsewhere remain ignored; the 5 tracked source files
are unaffected.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .gitignore | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 7a963cd68b6..1fcb1ea0a8a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -208,7 +208,7 @@ cache/
 .cache/
 diffusion_cache/
 kv_cache/
-!vllm_omni/bde/kv_cache/
+!vllm_omni/experimental/bde/kv_cache/
 
 # Logs
 logs/

From a64befa46ba212513d29fc12f08111b28dd0dbe6 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Fri, 26 Jun 2026 07:16:44 +0000
Subject: [PATCH 34/51] [AR-Diffusion] Rename BDE -> AR-Diffusion Engine;
 DreamZero engine-only, selected via command args

Rename the engine subsystem and drop DreamZero's dual path:

- Rename BDE -> AR-Diffusion Engine: package experimental/bde -> experimental/
  ar_diffusion; classes BDE* -> ARDiffusion* (ARDiffusionEngine / ModelRunner /
  KVCache / KVState / KVConfig / RequestAdapter); engine_backend "bde" ->
  "ar_diffusion"; tests/bde -> tests/ar_diffusion; docstrings "AR-Diffusion Engine".
- DreamZero is engine-only: removed the model-local KV fallback. The pipeline
  _kv_* hooks route only to the engine state; DreamZeroState's KV methods/attrs
  are deleted (non-KV session state kept).
- Engine/KV selection via command/deploy args, not env: dropped the BDE_KV_ENABLE
  gate (_bde_routing_enabled removed). DreamZero auto-routes to ar_diffusion
  unconditionally (overridable via the engine_backend arg); KV params come from a
  new ar_diffusion_kv_config field. The runner always enables KV.

Tests: tests/ar_diffusion 59 passed; E2E DreamZero generates on ar_diffusion with
no env vars (AR-Diffusion KV cache enabled, 17-frame clip).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .gitignore                                    |   2 +-
 tests/{bde => ar_diffusion}/__init__.py       |   0
 .../test_engine_routing.py                    |  82 ++++++-------
 tests/{bde => ar_diffusion}/test_gather.py    |  12 +-
 tests/{bde => ar_diffusion}/test_kv_cache.py  |  34 +++---
 .../test_kv_orchestration.py                  |  20 ++--
 .../test_pipeline_bridge.py                   |  30 ++---
 .../test_slot_mapping.py                      |   8 +-
 vllm_omni/diffusion/data.py                   |  68 +++++------
 vllm_omni/diffusion/diffusion_engine.py       |  10 +-
 .../models/dreamzero/pipeline_dreamzero.py    |  80 ++++---------
 .../models/dreamzero/state_dreamzero.py       |  59 ---------
 .../diffusion/worker/diffusion_worker.py      |   4 +-
 vllm_omni/experimental/__init__.py            |   4 +-
 .../{bde => ar_diffusion}/__init__.py         |   8 +-
 .../{bde => ar_diffusion}/engine.py           |  32 ++---
 .../kv_cache/__init__.py                      |  24 ++--
 .../{bde => ar_diffusion}/kv_cache/config.py  |  10 +-
 .../{bde => ar_diffusion}/kv_cache/manager.py |  76 ++++++------
 .../{bde => ar_diffusion}/kv_cache/paged.py   |   2 +-
 .../{bde => ar_diffusion}/kv_cache/state.py   |  38 +++---
 .../{bde => ar_diffusion}/runner.py           | 112 ++++++++----------
 22 files changed, 295 insertions(+), 420 deletions(-)
 rename tests/{bde => ar_diffusion}/__init__.py (100%)
 rename tests/{bde => ar_diffusion}/test_engine_routing.py (58%)
 rename tests/{bde => ar_diffusion}/test_gather.py (93%)
 rename tests/{bde => ar_diffusion}/test_kv_cache.py (82%)
 rename tests/{bde => ar_diffusion}/test_kv_orchestration.py (75%)
 rename tests/{bde => ar_diffusion}/test_pipeline_bridge.py (87%)
 rename tests/{bde => ar_diffusion}/test_slot_mapping.py (91%)
 rename vllm_omni/experimental/{bde => ar_diffusion}/__init__.py (55%)
 rename vllm_omni/experimental/{bde => ar_diffusion}/engine.py (58%)
 rename vllm_omni/experimental/{bde => ar_diffusion}/kv_cache/__init__.py (58%)
 rename vllm_omni/experimental/{bde => ar_diffusion}/kv_cache/config.py (75%)
 rename vllm_omni/experimental/{bde => ar_diffusion}/kv_cache/manager.py (85%)
 rename vllm_omni/experimental/{bde => ar_diffusion}/kv_cache/paged.py (99%)
 rename vllm_omni/experimental/{bde => ar_diffusion}/kv_cache/state.py (84%)
 rename vllm_omni/experimental/{bde => ar_diffusion}/runner.py (65%)

diff --git a/.gitignore b/.gitignore
index 1fcb1ea0a8a..f6e9f49f8d7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -208,7 +208,7 @@ cache/
 .cache/
 diffusion_cache/
 kv_cache/
-!vllm_omni/experimental/bde/kv_cache/
+!vllm_omni/experimental/ar_diffusion/kv_cache/
 
 # Logs
 logs/
diff --git a/tests/bde/__init__.py b/tests/ar_diffusion/__init__.py
similarity index 100%
rename from tests/bde/__init__.py
rename to tests/ar_diffusion/__init__.py
diff --git a/tests/bde/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py
similarity index 58%
rename from tests/bde/test_engine_routing.py
rename to tests/ar_diffusion/test_engine_routing.py
index 5db07684c69..0c380ab730f 100644
--- a/tests/bde/test_engine_routing.py
+++ b/tests/ar_diffusion/test_engine_routing.py
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Routing tests for BDE engine selection (Phase 1, PR-1).
+"""Routing tests for AR-Diffusion engine selection (Phase 1, PR-1).
 
 Covers ``DiffusionEngine.resolve_engine_class`` — the ``engine_backend``
-config-field dispatcher that routes a request to ``BDEEngine`` vs the base
+config-field dispatcher that routes a request to ``ARDiffusionEngine`` vs the base
 ``DiffusionEngine``. Resolution is tested without constructing an engine
 (construction runs a dummy forward that needs a real model).
 """
@@ -15,15 +15,14 @@
 
 from vllm_omni.diffusion.data import (
     OmniDiffusionConfig,
-    _bde_routing_enabled,
     default_engine_backend_for_model,
     resolve_default_engine_backend,
 )
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
-from vllm_omni.experimental.bde.engine import (
-    BDE_MODEL_RUNNER_CLS,
-    BDEEngine,
-    apply_bde_runner_default,
+from vllm_omni.experimental.ar_diffusion.engine import (
+    AR_DIFFUSION_MODEL_RUNNER_CLS,
+    ARDiffusionEngine,
+    apply_ar_diffusion_runner_default,
 )
 
 
@@ -36,16 +35,16 @@ def test_default_resolves_to_diffusion_engine():
 
 
 def test_bde_key_resolves_to_bde_engine():
-    assert DiffusionEngine.resolve_engine_class(_cfg("bde")) is BDEEngine
+    assert DiffusionEngine.resolve_engine_class(_cfg("ar_diffusion")) is ARDiffusionEngine
 
 
 def test_subclass_type_is_returned():
-    assert DiffusionEngine.resolve_engine_class(_cfg(BDEEngine)) is BDEEngine
+    assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine
 
 
 def test_qualname_string_resolves():
-    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.bde.engine.BDEEngine"))
-    assert cls is BDEEngine
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
+    assert cls is ARDiffusionEngine
 
 
 def test_non_engine_type_raises():
@@ -64,7 +63,7 @@ def test_bad_qualname_raises():
 
 
 def test_bde_engine_is_diffusion_engine_subclass():
-    assert issubclass(BDEEngine, DiffusionEngine)
+    assert issubclass(ARDiffusionEngine, DiffusionEngine)
 
 
 def test_config_engine_backend_field_defaults_to_default():
@@ -72,11 +71,11 @@ def test_config_engine_backend_field_defaults_to_default():
     assert field.default == "default"
 
 
-# --- DreamZero defaults to the BDE engine -----------------------------------
+# --- DreamZero defaults to the AR-Diffusion engine -----------------------------------
 
 
 def test_dreamzero_pipeline_defaults_to_bde():
-    assert default_engine_backend_for_model("DreamZeroPipeline") == "bde"
+    assert default_engine_backend_for_model("DreamZeroPipeline") == "ar_diffusion"
 
 
 def test_unknown_model_defaults_to_none():
@@ -89,64 +88,57 @@ def test_none_model_defaults_to_none():
 
 def test_dreamzero_default_routes_to_bde_engine():
     # End-to-end of the routing: DreamZero's per-model default backend resolves
-    # to BDEEngine through the same dispatcher used by make_engine.
+    # to ARDiffusionEngine through the same dispatcher used by make_engine.
     backend = default_engine_backend_for_model("DreamZeroPipeline")
-    assert DiffusionEngine.resolve_engine_class(_cfg(backend)) is BDEEngine
+    assert DiffusionEngine.resolve_engine_class(_cfg(backend)) is ARDiffusionEngine
 
 
-# --- BDE routing is gated on BDE_KV_ENABLE (PR #4534 blast-radius review) ----
+# --- AR-Diffusion routing: unconditional per-model default; explicit arg overrides ---
 
 
-def test_routing_enabled_reads_env(monkeypatch):
-    monkeypatch.setenv("BDE_KV_ENABLE", "1")
-    assert _bde_routing_enabled() is True
-    monkeypatch.delenv("BDE_KV_ENABLE", raising=False)
-    assert _bde_routing_enabled() is False
+def test_dreamzero_default_applied_unconditionally():
+    # DreamZero + unset ("default") backend -> auto-route to ar_diffusion (no env gate).
+    assert resolve_default_engine_backend("DreamZeroPipeline", "default") == "ar_diffusion"
 
 
-def test_default_backend_applied_only_when_enabled(monkeypatch):
-    # DreamZero + unset ("default") backend: auto-route applies only with BDE on.
-    monkeypatch.setenv("BDE_KV_ENABLE", "1")
-    assert resolve_default_engine_backend("DreamZeroPipeline", "default") == "bde"
-    monkeypatch.delenv("BDE_KV_ENABLE", raising=False)
-    assert resolve_default_engine_backend("DreamZeroPipeline", "default") is None
+def test_explicit_backend_not_overridden():
+    # An explicit engine_backend command arg bypasses the per-model default.
+    assert resolve_default_engine_backend("DreamZeroPipeline", "ar_diffusion") is None
+    assert (
+        resolve_default_engine_backend(
+            "DreamZeroPipeline", "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+        )
+        is None
+    )
 
 
-def test_explicit_backend_not_overridden_even_when_enabled(monkeypatch):
-    # An explicit engine_backend bypasses the auto-route gate entirely.
-    monkeypatch.setenv("BDE_KV_ENABLE", "1")
-    assert resolve_default_engine_backend("DreamZeroPipeline", "bde") is None
-    assert resolve_default_engine_backend("DreamZeroPipeline", "vllm_omni.experimental.bde.engine.BDEEngine") is None
-
-
-def test_non_bde_model_never_routed(monkeypatch):
-    monkeypatch.setenv("BDE_KV_ENABLE", "1")
+def test_non_ar_diffusion_model_not_routed():
     assert resolve_default_engine_backend("WanS2VPipeline", "default") is None
 
 
-# --- BDE worker/runner wiring -----------------------------------------------
+# --- AR-Diffusion worker/runner wiring -----------------------------------------------
 
 
 def test_bde_model_runner_cls_resolves():
-    cls = resolve_obj_by_qualname(BDE_MODEL_RUNNER_CLS)
-    assert cls.__name__ == "BDEModelRunner"
+    cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS)
+    assert cls.__name__ == "ARDiffusionModelRunner"
 
 
 def test_apply_runner_default_sets_when_unset():
     cfg = SimpleNamespace(diffusion_model_runner_cls=None)
-    apply_bde_runner_default(cfg)
-    assert cfg.diffusion_model_runner_cls == BDE_MODEL_RUNNER_CLS
+    apply_ar_diffusion_runner_default(cfg)
+    assert cfg.diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
 
 
 def test_apply_runner_default_respects_explicit_choice():
     cfg = SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner")
-    apply_bde_runner_default(cfg)
+    apply_ar_diffusion_runner_default(cfg)
     assert cfg.diffusion_model_runner_cls == "my.custom.Runner"
 
 
 def test_worker_runner_selection_prefers_override():
     # Mirrors the worker hook: an od_config override wins over the platform default.
-    over = SimpleNamespace(diffusion_model_runner_cls=BDE_MODEL_RUNNER_CLS)
-    assert (getattr(over, "diffusion_model_runner_cls", None) or "PLATFORM") == BDE_MODEL_RUNNER_CLS
+    over = SimpleNamespace(diffusion_model_runner_cls=AR_DIFFUSION_MODEL_RUNNER_CLS)
+    assert (getattr(over, "diffusion_model_runner_cls", None) or "PLATFORM") == AR_DIFFUSION_MODEL_RUNNER_CLS
     unset = SimpleNamespace(diffusion_model_runner_cls=None)
     assert (getattr(unset, "diffusion_model_runner_cls", None) or "PLATFORM") == "PLATFORM"
diff --git a/tests/bde/test_gather.py b/tests/ar_diffusion/test_gather.py
similarity index 93%
rename from tests/bde/test_gather.py
rename to tests/ar_diffusion/test_gather.py
index e61e66bd156..41ebc33b9a4 100644
--- a/tests/bde/test_gather.py
+++ b/tests/ar_diffusion/test_gather.py
@@ -3,9 +3,9 @@
 
 import torch
 
-from vllm_omni.experimental.bde.kv_cache import (
-    BDEKVCache,
-    BDEKVConfig,
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
     allocate_kv_pool,
     pool_gather_window,
     pool_write_chunk,
@@ -76,9 +76,9 @@ def test_pool_write_gather_window_trim():
 
 
 def test_bde_kv_cache_write_gather_integration():
-    """BDEKVCache.write_chunk_kv → gather_window roundtrip through a real manager."""
-    cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
-    kv = BDEKVCache(
+    """ARDiffusionKVCache.write_chunk_kv → gather_window roundtrip through a real manager."""
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
+    kv = ARDiffusionKVCache(
         cfg,
         num_layers=1,
         num_kv_heads=N_HEADS,
diff --git a/tests/bde/test_kv_cache.py b/tests/ar_diffusion/test_kv_cache.py
similarity index 82%
rename from tests/bde/test_kv_cache.py
rename to tests/ar_diffusion/test_kv_cache.py
index 845e291897c..c03f45147fd 100644
--- a/tests/bde/test_kv_cache.py
+++ b/tests/ar_diffusion/test_kv_cache.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Unit tests for the BDE KV cache helpers (Phase 1, PR-2).
+"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
 
 Covers the request adapter, the chunk-window spec/manager (registration + the
 eviction policy), and the pool builder — exercised against the installed vLLM
@@ -12,16 +12,16 @@
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.request import RequestStatus
 
-from vllm_omni.experimental.bde.kv_cache import (
-    BDEKVCache,
-    BDEKVConfig,
-    BDERequestAdapter,
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionRequestAdapter,
     ChunkWindowManager,
     ChunkWindowSpec,
     build_kv_manager,
     compute_num_blocks,
 )
-from vllm_omni.experimental.bde.kv_cache.paged import chunk_window_skipped_tokens
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
 
 BLOCK = 16
 
@@ -99,19 +99,19 @@ def test_reset_at_boundary_drops_completed_past_sink():
     assert skip == 32
 
 
-# --- BDEKVConfig ------------------------------------------------------------
+# --- ARDiffusionKVConfig ------------------------------------------------------------
 
 
 def test_kv_config_sliding_window_property():
-    assert BDEKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
-    assert BDEKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
 
 
-# --- BDERequestAdapter ------------------------------------------------------
+# --- ARDiffusionRequestAdapter ------------------------------------------------------
 
 
 def test_adapter_advances_per_chunk_not_per_step():
-    a = BDERequestAdapter("r0", chunk_size=16)
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16)
     assert a.num_computed_tokens == 0
     assert a.num_tokens == 16  # in-flight chunk
     a.on_chunk_committed()
@@ -120,14 +120,14 @@ def test_adapter_advances_per_chunk_not_per_step():
 
 
 def test_adapter_accounts_for_prefill_prefix():
-    a = BDERequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
     assert a.num_computed_tokens == 4
     assert a.num_prompt_tokens == 4
     assert a.num_tokens == 20
 
 
 def test_adapter_status_is_vllm_enum():
-    assert isinstance(BDERequestAdapter("r0", chunk_size=16).status, RequestStatus)
+    assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus)
 
 
 # --- pool / manager ---------------------------------------------------------
@@ -139,7 +139,7 @@ def test_compute_num_blocks():
 
 
 def test_build_manager_allocate_free_roundtrip():
-    """End-to-end: a BDERequestAdapter drives a real KVCacheManager.
+    """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.
 
     This is the adapter conformance check — if the adapter were missing an
     attribute the manager reads, allocate_slots/free would raise here.
@@ -148,7 +148,7 @@ def test_build_manager_allocate_free_roundtrip():
     mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
     free_before = mgr.block_pool.get_num_free_blocks()
 
-    adapter = BDERequestAdapter("req-0", chunk_size=BLOCK)
+    adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK)
     blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
     assert blocks is not None
     assert mgr.block_pool.get_num_free_blocks() < free_before
@@ -163,8 +163,8 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
     memory budget (review: zwhzzz0821)."""
     L = 512
     avail = 1 << 30  # 1 GiB
-    kv = BDEKVCache(
-        BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
+    kv = ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
         num_layers=2,
         num_kv_heads=4,
         head_size=64,
diff --git a/tests/bde/test_kv_orchestration.py b/tests/ar_diffusion/test_kv_orchestration.py
similarity index 75%
rename from tests/bde/test_kv_orchestration.py
rename to tests/ar_diffusion/test_kv_orchestration.py
index dc5099cfd6f..4ba5cb49bf1 100644
--- a/tests/bde/test_kv_orchestration.py
+++ b/tests/ar_diffusion/test_kv_orchestration.py
@@ -1,28 +1,28 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Tests for BDEKVCache — the engine-level KV orchestration body (Phase 1)."""
+"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1)."""
 
 import pytest
 import torch
 
-from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
 
 DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16)
 
 
 def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
-    cfg = BDEKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
-    return BDEKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+    return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
 
 
 def test_requires_enabled_config():
     with pytest.raises(ValueError):
-        BDEKVCache(BDEKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
+        ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
 
 
 def test_requires_bounded_window():
-    cfg = BDEKVConfig(enable=True, chunk_size=16, window_chunks=None)
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None)
     with pytest.raises(ValueError):
-        BDEKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
+        ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
 
 
 def test_full_request_lifecycle_and_eviction():
@@ -67,20 +67,20 @@ def test_num_computed_advances_per_chunk():
 
 
 def test_state_close_frees_both_branch_blocks():
-    """BDEKVState.close() returns both CFG branches' pool blocks to the pool.
+    """ARDiffusionKVState.close() returns both CFG branches' pool blocks to the pool.
 
     This is the primitive the runner's LRU eviction relies on: when a session is
     evicted, close() must free the blocks both adapters hold, or session churn
     leaks pool ownership (review P1).
     """
-    from vllm_omni.experimental.bde.kv_cache.state import BDEKVState
+    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
     kv = make_cache(chunk_size=16, window_chunks=2)
     free_total = kv.manager.block_pool.get_num_free_blocks()
 
     pos = kv.begin_request("bde__s")
     neg = kv.begin_request("bde__s__neg")
-    state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
+    state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers)
     for _ in range(3):
         for adapter in (pos, neg):
             kv.allocate_chunk(adapter)
diff --git a/tests/bde/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py
similarity index 87%
rename from tests/bde/test_pipeline_bridge.py
rename to tests/ar_diffusion/test_pipeline_bridge.py
index ba56aa6f6a5..fb38b4f8304 100644
--- a/tests/bde/test_pipeline_bridge.py
+++ b/tests/ar_diffusion/test_pipeline_bridge.py
@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Tests for the BDEKVState pipeline bridge (Step 5)."""
+"""Tests for the ARDiffusionKVState pipeline bridge (Step 5)."""
 
 import pytest
 import torch
 
-from vllm_omni.experimental.bde.kv_cache import BDEKVCache, BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.state import BDEKVState
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
 BLOCK = 16
 N_HEADS = 4
@@ -13,8 +13,8 @@
 
 
 def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
-    cfg = BDEKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
-    kv = BDEKVCache(
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
         cfg,
         num_layers=num_layers,
         num_kv_heads=N_HEADS,
@@ -28,7 +28,7 @@ def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
     )
     pos = kv.begin_request("r-pos")
     neg = kv.begin_request("r-neg")
-    return kv, BDEKVState(kv, pos, neg, num_layers=num_layers)
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
 
 
 def _window(n_chunks):
@@ -36,7 +36,7 @@ def _window(n_chunks):
     return torch.randn(2, 1, n_chunks * BLOCK, N_HEADS, HEAD_DIM)
 
 
-# Approach 2: BDEKVState writes the new tokens into paged pool blocks and gathers
+# Approach 2: ARDiffusionKVState writes the new tokens into paged pool blocks and gathers
 # the resident window back; output mirrors the model-local sliding window.
 
 
@@ -163,23 +163,15 @@ def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
         assert torch.equal(out[i]["v"], v)
 
 
-def test_kv_create_owned_by_engine_under_bde():
-    """_kv_create initializes model-local caches ONLY on the non-BDE path. Under
-    BDE the pool owns all allocation, so state.create_kv_caches must NOT be called."""
+def test_kv_create_is_noop_engine_owns_allocation():
+    """The AR-Diffusion engine owns all KV allocation, so _kv_create never touches
+    the model-local state (DreamZero's model-local KV path was removed)."""
     from unittest.mock import MagicMock
 
     from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 
     p = DreamZeroPipeline.__new__(DreamZeroPipeline)
-
-    # Non-BDE: model-local caches are created (no recursion).
-    p._bde_kv_state = None
-    state = MagicMock()
-    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
-    state.create_kv_caches.assert_called_once_with(1, "float32", "cpu", 24, 4, 64)
-
-    # BDE active: allocation is engine-owned -> no model-local creation.
-    p._bde_kv_state = object()
+    p._ar_diffusion_kv_state = object()
     state = MagicMock()
     p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
     state.create_kv_caches.assert_not_called()
diff --git a/tests/bde/test_slot_mapping.py b/tests/ar_diffusion/test_slot_mapping.py
similarity index 91%
rename from tests/bde/test_slot_mapping.py
rename to tests/ar_diffusion/test_slot_mapping.py
index 31464b869c1..061df4d9c64 100644
--- a/tests/bde/test_slot_mapping.py
+++ b/tests/ar_diffusion/test_slot_mapping.py
@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Slot-mapping / block-table tests for the BDE engine (Phase 1, Step 3)."""
+"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3)."""
 
 import pytest
 import torch
 
-from vllm_omni.experimental.bde.kv_cache import (
-    BDERequestAdapter,
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionRequestAdapter,
     ChunkWindowSpec,
     build_kv_manager,
     chunk_slot_mapping,
@@ -61,7 +61,7 @@ def test_blocktable_build_single_request():
     spec = make_spec()
     mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
     null_id = mgr.block_pool.null_block.block_id
-    adapter = BDERequestAdapter("req", chunk_size=BLOCK)
+    adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK)
 
     mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
     block_ids = mgr.get_block_ids(adapter.request_id)[0]
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index ec0f45fb9ae..5c2f3285df4 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -493,50 +493,33 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
     return None
 
 
-# Model classes that default to the BDE (AR-diffusion) engine.
-_BDE_ENGINE_MODEL_CLASSES = frozenset({"DreamZeroPipeline"})
+# Model classes that default to the AR-Diffusion (AR-diffusion) engine.
+_AR_DIFFUSION_ENGINE_MODEL_CLASSES = frozenset({"DreamZeroPipeline"})
 
 
 def default_engine_backend_for_model(model_class_name: str | None) -> str | None:
     """Engine backend a model opts into by default.
 
-    Returns ``"bde"`` for AR-diffusion world models (DreamZero today) so they run
-    on the Block Diffusion Engine with engine-level KV cache management; returns
-    ``None`` to keep the base ``DiffusionEngine``. Extend
-    ``_BDE_ENGINE_MODEL_CLASSES`` as more AR-DiT models migrate onto BDE.
-
-    NOTE: this is the *pure* "what would this model pick" lookup. Whether it is
-    actually applied is gated by :func:`_bde_routing_enabled` at the call sites, so
-    a deploy that hasn't opted into BDE stays on the base engine.
+    Returns ``"ar_diffusion"`` for AR-diffusion world models (DreamZero today) so they run on
+    the AR-Diffusion engine with engine-level KV cache management; ``None`` keeps the
+    base ``DiffusionEngine``. Extend ``_AR_DIFFUSION_ENGINE_MODEL_CLASSES`` as more AR-DiT
+    models migrate. Applied when ``engine_backend`` is unset (see
+    :func:`resolve_default_engine_backend`); an explicit command/deploy arg overrides.
     """
-    if model_class_name in _BDE_ENGINE_MODEL_CLASSES:
-        return "bde"
+    if model_class_name in _AR_DIFFUSION_ENGINE_MODEL_CLASSES:
+        return "ar_diffusion"
     return None
 
 
-def _bde_routing_enabled() -> bool:
-    """Whether to auto-route a BDE-eligible model onto the BDE engine/runner.
-
-    Gated on the same switch as BDE KV (``BDE_KV_ENABLE``, read by
-    ``resolve_bde_kv_config``): a deploy that hasn't opted into BDE KV stays on the
-    base ``DiffusionEngine``/runner, keeping the (still-unvalidated for the
-    production CFG-parallel + CUDA-graph + compile topology) BDE runner off the
-    default path. An explicit ``engine_backend`` selection bypasses this gate
-    entirely (it only governs the ``"default"`` -> per-model auto-route). (PR #4534)
-    """
-    return os.environ.get("BDE_KV_ENABLE") == "1"
-
-
 def resolve_default_engine_backend(model_class_name: str | None, current_backend: str | None) -> str | None:
     """Engine backend to apply for an auto-routed model, or ``None`` to leave as-is.
 
-    Applies the per-model default (:func:`default_engine_backend_for_model`) only
-    when ``current_backend`` is the unset ``"default"`` sentinel *and* BDE routing
-    is opted in (:func:`_bde_routing_enabled`). Returns ``None`` otherwise, so an
-    explicit backend choice and BDE-off deploys both stay untouched. Shared by the
-    two routing sites (``__post_init__`` auto-detect and ``from_kwargs``).
+    Applies the per-model default (:func:`default_engine_backend_for_model`) when
+    ``current_backend`` is the unset ``"default"`` sentinel; an explicit
+    ``engine_backend`` command/deploy arg is left untouched. Shared by the two
+    routing sites (``enrich_config`` and ``from_kwargs``).
     """
-    if current_backend != "default" or not _bde_routing_enabled():
+    if current_backend != "default":
         return None
     return default_engine_backend_for_model(model_class_name)
 
@@ -589,20 +572,26 @@ class OmniDiffusionConfig:
     # Engine backend selection. Resolved by ``DiffusionEngine.make_engine``
     # (mirrors ``distributed_executor_backend`` / ``DiffusionExecutor.get_class``):
     #   "default" -> DiffusionEngine
-    #   "bde"     -> BDEEngine (AR-diffusion engine with engine-level KV cache
+    #   "ar_diffusion"     -> ARDiffusionEngine (AR-diffusion engine with engine-level KV cache
     #                management); also accepts a DiffusionEngine subclass or an
     #                import path string.
     engine_backend: str = "default"
 
     # Optional override for the diffusion model runner class (import path).
-    # The worker uses it instead of the platform default when set; the BDE engine
-    # sets it to BDEModelRunner. Unset -> platform default (existing behavior).
+    # The worker uses it instead of the platform default when set; the AR-Diffusion engine
+    # sets it to ARDiffusionModelRunner. Unset -> platform default (existing behavior).
     diffusion_model_runner_cls: str | None = None
+
+    # AR-Diffusion engine KV-cache parameters (command/deploy arg, e.g. in the
+    # deploy yaml). Dict of window_chunks / gpu_memory_fraction / sink_chunks /
+    # reset_at_boundary; None uses engine defaults. Selecting the AR-Diffusion
+    # engine enables KV (no env gate).
+    ar_diffusion_kv_config: dict | None = None
     # DEFERRED (#4425 / RFC #4021 structured config): when rebasing onto the
     # structured VllmOmniConfig, mirror engine_backend + diffusion_model_runner_cls
     # into its DiffusionConfig (engine/runtime group) and parity test. enrich_config()
     # there only copies fields present on DiffusionConfig, so without the mirror
-    # DreamZero's BDE routing silently drops on the structured-config path. (PR #4534)
+    # DreamZero's AR-Diffusion routing silently drops on the structured-config path. (PR #4534)
 
     # HuggingFace specific parameters
     trust_remote_code: bool = False
@@ -1131,9 +1120,8 @@ def enrich_config(self) -> None:
                         self.set_tf_model_config(TransformerConfig())
                         self.update_multimodal_support()
                         # DreamZero is an AR-diffusion world model: route it to the
-                        # BDE engine (engine-level KV cache management) when BDE is
-                        # opted in (BDE_KV_ENABLE) and the caller did not explicitly
-                        # select an engine backend; otherwise stay on base.
+                        # AR-Diffusion engine (engine-level KV cache management) unless
+                        # the caller explicitly selected an engine backend.
                         backend = resolve_default_engine_backend(self.model_class_name, self.engine_backend)
                         if backend is not None:
                             self.engine_backend = backend
@@ -1210,8 +1198,8 @@ def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
         instance = cls(**filtered_kwargs)
         # Route per-model engines after model_class_name is finalized. This covers
         # the explicit-model_class_name path (e.g. deploy configs) that the
-        # __post_init__ auto-detect branch does not reach. Gated on BDE_KV_ENABLE
-        # so non-BDE deploys stay on the base DiffusionEngine (PR #4534 review).
+        # __post_init__ auto-detect branch does not reach. An explicit engine_backend
+        # command/deploy arg overrides the per-model default.
         backend = resolve_default_engine_backend(instance.model_class_name, instance.engine_backend)
         if backend is not None:
             instance.engine_backend = backend
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index c4eeac40a64..83716718247 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -417,7 +417,7 @@ def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
         """Resolve the engine class selected by ``config.engine_backend``.
 
         Mirrors ``DiffusionExecutor.get_class``: accepts the string keys
-        ``"default"`` / ``"bde"``, a ``DiffusionEngine`` subclass, or an import
+        ``"default"`` / ``"ar_diffusion"``, a ``DiffusionEngine`` subclass, or an import
         path string. Kept separate from :meth:`make_engine` so the selection is
         testable without constructing an engine (which runs a dummy forward).
 
@@ -435,12 +435,12 @@ def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
             return backend
         if backend == "default":
             return DiffusionEngine
-        if backend == "bde":
-            # Lazy import to avoid a circular dependency (bde.engine imports this
+        if backend == "ar_diffusion":
+            # Lazy import to avoid a circular dependency (ar_diffusion.engine imports this
             # module), mirroring how DiffusionExecutor.get_class imports backends.
-            from vllm_omni.experimental.bde.engine import BDEEngine
+            from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
 
-            return BDEEngine
+            return ARDiffusionEngine
         if isinstance(backend, str):
             try:
                 engine_class = resolve_obj_by_qualname(backend)
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 362e0e5b371..e43f465f774 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -95,58 +95,34 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     Multi-output: predict_noise() returns (video_pred, action_pred).
     CFG: video gets standard CFG, action takes positive branch only.
 
-    BDE KV integration: when ``self._bde_kv_state`` is set by a BDEModelRunner
-    before ``forward()``, the pipeline routes its KV access (get / create / update
-    / commit) through the pool-backed ``BDEKVState`` instead of the model-local
-    ``DreamZeroState``. No BDE imports -- purely duck-typed, inactive by default.
+    KV is managed by the AR-Diffusion engine: ``self._ar_diffusion_kv_state`` is set by the
+    runner before ``forward()`` and the pipeline routes all KV access (get / update /
+    commit / reset) through the pool-backed state. Purely duck-typed (no engine import).
     """
 
-    _bde_kv_state = None  # set by BDEModelRunner when KV management is enabled
+    _ar_diffusion_kv_state = None  # set by the runner before each forward
 
     def _kv_get(self, state, is_negative):
-        s = self._bde_kv_state
-        if s is not None:
-            logger.debug("BDE KV get (neg=%s): %d layers", is_negative, s.num_layers)
-            return s.get_kv_caches(is_negative)
-        return state.get_kv_caches(is_negative)
+        return self._ar_diffusion_kv_state.get_kv_caches(is_negative)
 
     def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
-        # Under BDE the pool owns all KV allocation: self-attn is allocated per
-        # chunk in _kv_update and read via _kv_get (empty window at prefill start),
-        # cross-attn is populated eagerly in _kv_populate_cross. No model-local
-        # caches are created. Only the non-BDE path needs them.
-        if self._bde_kv_state is not None:
-            return
-        state.create_kv_caches(batch_size, dtype, device, num_layers, num_heads, head_dim)
+        # The engine owns all KV allocation: self-attn is allocated per chunk in
+        # _kv_update and read via _kv_get (empty window at prefill start); cross-attn
+        # is populated eagerly in _kv_populate_cross. Nothing to create model-side.
+        return
 
     def _kv_update(self, state, layer_idx, updated_kv, is_negative, seq_len=None):
-        s = self._bde_kv_state
-        if s is not None:
-            logger.debug("BDE KV write: layer %d neg=%s shape=%s", layer_idx, is_negative, tuple(updated_kv.shape))
-            s.update_kv_cache(layer_idx, updated_kv, is_negative, seq_len)
-        else:
-            state.update_kv_cache(layer_idx, updated_kv, is_negative=is_negative)
+        self._ar_diffusion_kv_state.update_kv_cache(layer_idx, updated_kv, is_negative, seq_len)
 
     def _kv_commit(self):
-        s = self._bde_kv_state
-        if s is not None:
-            logger.debug("BDE KV commit: pos=%d chunks neg=%d chunks", s.pos.completed_chunks, s.neg.completed_chunks)
-            s.commit_chunk()
+        self._ar_diffusion_kv_state.commit_chunk()
 
     def _kv_get_cross(self, state, is_negative):
-        """Cross-attn cache: pool-backed when BDE active, model-local otherwise.
-
-        The BDE pool caches both text ``k``/``v`` and (for I2V) the image-token
-        ``k_img``/``v_img`` added by #4154, so the returned dict satisfies the I2V
-        forward without recomputation.
-        """
-        s = self._bde_kv_state
-        if s is not None:
-            return s.get_cross_kv_caches(is_negative)
-        return state.get_crossattn_caches(is_negative)
+        """Cross-attn cache from the engine pool (text k/v + I2V image k_img/v_img)."""
+        return self._ar_diffusion_kv_state.get_cross_kv_caches(is_negative)
 
     def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None:
-        """Eagerly project cross-attn K/V for all layers into the BDE pool.
+        """Eagerly project cross-attn K/V for all layers into the AR-Diffusion pool.
 
         Caches the session-invariant cross-attn projections once (re-run after a
         window-boundary reset, when ``_cross_populated`` clears): the text ``k``/``v``
@@ -155,8 +131,8 @@ def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: b
         splits off, cached model-side by #4154). Must run after the image is encoded
         so ``clip_feature`` is available.
         """
-        s = self._bde_kv_state
-        if s is None or s._cross_populated.get(is_negative, False):
+        s = self._ar_diffusion_kv_state
+        if s._cross_populated.get(is_negative, False):
             return
         projected = self.transformer.text_embedding(context)
         img_ctx = None
@@ -174,7 +150,7 @@ def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: b
             s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img)
         s._cross_populated[is_negative] = True
         logger.info(
-            "BDE CROSS POPULATE [%s]: %d layers, text=%s img=%s",
+            "AR-Diffusion CROSS POPULATE [%s]: %d layers, text=%s img=%s",
             "neg" if is_negative else "pos",
             len(self.transformer.blocks),
             tuple(context.shape),
@@ -182,18 +158,14 @@ def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: b
         )
 
     def _kv_reset(self, state, *, clear_video_latents: bool = True):
-        """Reset model-local KV and, under BDE, the pooled session window too.
+        """Reset the engine's pooled session window plus the model's non-KV state.
 
-        DreamZero resets its sliding window at the attention-window boundary; the
-        BDE pool must drop the same window so the next forward starts fresh and
-        stays parity-exact with the model-local path. ``clear_video_latents=False``
-        mirrors ``state.reset_inference_state()`` (window-boundary "inference" reset),
-        which keeps the accumulated video latents for export.
+        DreamZero resets at the attention-window boundary; the engine pool drops the
+        same window so the next forward starts fresh. ``clear_video_latents=False``
+        keeps the accumulated video latents for export.
         """
         state.reset(clear_video_latents=clear_video_latents)
-        s = self._bde_kv_state
-        if s is not None:
-            s.reset()
+        self._ar_diffusion_kv_state.reset()
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.
@@ -446,9 +418,9 @@ def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor
         if kwargs.get("update_kv_cache", False) and updated_kv_caches:
             state = kwargs.get("dreamzero_state", self.state)
             is_neg = kwargs.get("is_negative", False)
-            if self._bde_kv_state is not None:
+            if self._ar_diffusion_kv_state is not None:
                 logger.info(
-                    "BDE pipeline predict_noise -> write-back: is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
+                    "AR-Diffusion pipeline predict_noise -> write-back: is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
                     is_neg,
                     kwargs.get("seq_len"),
                     kwargs.get("current_start_frame"),
@@ -1087,7 +1059,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
         else:
             # Auto-reset based on model state (before accumulation). Both a "session"
             # reset and an "inference" (window-boundary) reset clear the model-local
-            # KV window, so route both through _kv_reset to drop the BDE pool window
+            # KV window, so route both through _kv_reset to drop the AR-Diffusion pool window
             # too and stay parity-exact; "inference" keeps the accumulated video
             # latents (reset_inference_state == reset(clear_video_latents=False)).
             reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
@@ -1137,7 +1109,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             state.clip_feas = clip_feas.to(dtype=image.dtype)
             state.ys = ys.to(dtype=image.dtype)
 
-            # Eager cross-attn population (BDE only): cache text + image-token K/V
+            # Eager cross-attn population (AR-Diffusion only): cache text + image-token K/V
             # into the pool now that the image is encoded (clip_feas available).
             # Runs on the first forward of a session and after each window-boundary
             # reset (current_start_frame returns to 0); _cross_populated guards re-entry.
diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
index 5815d93d4b1..2ddb9d0f31f 100644
--- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
@@ -115,11 +115,6 @@ def reset(self, *, clear_video_latents: bool = True) -> None:
         # Concatenated along T before decode (matches upstream video_across_time).
         self.video_latents_across_time = saved_video_latents
 
-        # KV cache once robot-policy diffusion supports that integration.
-        self.kv_cache = None
-        self.kv_cache_neg = None
-        self.crossattn_cache = None
-        self.crossattn_cache_neg = None
         self.current_start_frame = 0
 
         self.clip_feas = None
@@ -166,57 +161,3 @@ def reset_reason(
     def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
         """Determine if state should be reset before this forward()."""
         return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None
-
-    # ------------------------------------------------------------------
-    # KV cache management
-    # ------------------------------------------------------------------
-
-    def create_kv_caches(
-        self,
-        batch_size: int,
-        dtype: torch.dtype,
-        device: torch.device,
-        num_layers: int,
-        num_heads: int,
-        head_dim: int,
-    ) -> None:
-        """Initialize empty KV caches and cross-attention caches."""
-        self.kv_cache = [
-            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
-        ]
-        self.kv_cache_neg = [
-            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
-        ]
-
-        self.crossattn_cache = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-        self.crossattn_cache_neg = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-
-    def update_kv_cache(
-        self,
-        layer_index: int,
-        updated_kv: torch.Tensor,
-        is_negative: bool = False,
-    ) -> None:
-        """Update a single layer's KV cache after prefill."""
-        cache = self.kv_cache_neg if is_negative else self.kv_cache
-        if cache is None:
-            raise RuntimeError("KV caches not initialized, call create_kv_caches first.")
-        cache[layer_index] = updated_kv.clone()
-
-    def get_kv_caches(self, is_negative: bool = False) -> list[torch.Tensor]:
-        """Get KV caches for the specified branch."""
-        cache = self.kv_cache_neg if is_negative else self.kv_cache
-        if cache is None:
-            raise RuntimeError("KV caches not initialized.")
-        return cache
-
-    def get_crossattn_caches(self, is_negative: bool = False) -> list[dict[str, bool | torch.Tensor | None]]:
-        """Get cross-attention caches for the specified branch."""
-        cache = self.crossattn_cache_neg if is_negative else self.crossattn_cache
-        if cache is None:
-            raise RuntimeError("Cross-attn caches not initialized.")
-        return cache
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index 0cc855df5c2..2bd1c9b50c6 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -209,8 +209,8 @@ def __init__(
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
-        # Create model runner. An od_config override (e.g. the BDE engine selecting
-        # BDEModelRunner) takes precedence; unset falls back to the platform default,
+        # Create model runner. An od_config override (e.g. the AR-Diffusion engine selecting
+        # ARDiffusionModelRunner) takes precedence; unset falls back to the platform default,
         # so existing models are unaffected. The override must be an explicit import-
         # path string — guard with isinstance so a non-string (e.g. a Mock od_config
         # in tests) doesn't shadow the platform hook.
diff --git a/vllm_omni/experimental/__init__.py b/vllm_omni/experimental/__init__.py
index 1be37e247b2..c7d59e2e262 100644
--- a/vllm_omni/experimental/__init__.py
+++ b/vllm_omni/experimental/__init__.py
@@ -3,6 +3,6 @@
 
 Modules under ``vllm_omni.experimental`` are in active development and not yet
 covered by the stability guarantees of the rest of the package — their APIs,
-config surfaces, and defaults may change without notice. The Block Diffusion
-Engine (:mod:`vllm_omni.experimental.bde`) lives here.
+config surfaces, and defaults may change without notice. The AR-Diffusion
+Engine (:mod:`vllm_omni.experimental.ar_diffusion`) lives here.
 """
diff --git a/vllm_omni/experimental/bde/__init__.py b/vllm_omni/experimental/ar_diffusion/__init__.py
similarity index 55%
rename from vllm_omni/experimental/bde/__init__.py
rename to vllm_omni/experimental/ar_diffusion/__init__.py
index 03c88964779..b312660ec72 100644
--- a/vllm_omni/experimental/bde/__init__.py
+++ b/vllm_omni/experimental/ar_diffusion/__init__.py
@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Block Diffusion Engine (BDE).
+"""AR-Diffusion Engine (AR-Diffusion).
 
 The AR-Diffusion engine: a ``DiffusionEngine`` subclass that adds engine-level KV
 cache management for autoregressive / chunked "world-model" diffusion models
 (DreamZero and the AR-DiT family). Selected via
-``OmniDiffusionConfig.engine_backend = "bde"``.
+``OmniDiffusionConfig.engine_backend = "ar_diffusion"``.
 """
 
-from vllm_omni.experimental.bde.engine import BDEEngine
+from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
 
-__all__ = ["BDEEngine"]
+__all__ = ["ARDiffusionEngine"]
diff --git a/vllm_omni/experimental/bde/engine.py b/vllm_omni/experimental/ar_diffusion/engine.py
similarity index 58%
rename from vllm_omni/experimental/bde/engine.py
rename to vllm_omni/experimental/ar_diffusion/engine.py
index e8914ee1a60..f8d09967ef3 100644
--- a/vllm_omni/experimental/bde/engine.py
+++ b/vllm_omni/experimental/ar_diffusion/engine.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""The Block Diffusion Engine (BDE)."""
+"""The AR-Diffusion Engine (AR-Diffusion)."""
 
 from __future__ import annotations
 
@@ -11,41 +11,41 @@
 
 logger = init_logger(__name__)
 
-#: Import path of the runner the BDE engine routes its workers to.
-BDE_MODEL_RUNNER_CLS = "vllm_omni.experimental.bde.runner.BDEModelRunner"
+#: Import path of the runner the AR-Diffusion engine routes its workers to.
+AR_DIFFUSION_MODEL_RUNNER_CLS = "vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner"
 
 
-def apply_bde_runner_default(od_config: OmniDiffusionConfig) -> None:
-    """Route the BDE engine's workers to ``BDEModelRunner``.
+def apply_ar_diffusion_runner_default(od_config: OmniDiffusionConfig) -> None:
+    """Route the AR-Diffusion engine's workers to ``ARDiffusionModelRunner``.
 
     Sets ``od_config.diffusion_model_runner_cls`` in place, unless the caller
     already chose a runner. The worker (in its own process) reads this off the
-    propagated ``od_config`` and builds the BDE runner instead of the platform
-    default, keeping the swap scoped to BDE-routed models.
+    propagated ``od_config`` and builds the AR-Diffusion runner instead of the platform
+    default, keeping the swap scoped to AR-Diffusion-routed models.
     """
     if getattr(od_config, "diffusion_model_runner_cls", None) is None:
-        od_config.diffusion_model_runner_cls = BDE_MODEL_RUNNER_CLS
+        od_config.diffusion_model_runner_cls = AR_DIFFUSION_MODEL_RUNNER_CLS
 
 
-class BDEEngine(DiffusionEngine):
+class ARDiffusionEngine(DiffusionEngine):
     """AR-Diffusion engine with engine-level KV cache management.
 
-    BDE serves autoregressive / chunked blockwise-causal diffusion models
+    AR-Diffusion serves autoregressive / chunked blockwise-causal diffusion models
     (world models, AR-DiT) that materialize persistent attention KV. It reuses
     vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` / ``BlockTables``)
     as a library, driven from the engine rather than hand-rolled inside each
     model. See ``BDE_doc/diffusion_kv_cache_management.md`` for the design and
     ``BDE_doc/dreamzero_kv_phase1_plan.md`` for the rollout.
 
-    It is selected per model via ``OmniDiffusionConfig.engine_backend = "bde"``
+    It is selected per model via ``OmniDiffusionConfig.engine_backend = "ar_diffusion"``
     (resolved by :meth:`DiffusionEngine.make_engine`), so models that do not opt
     in keep using the base ``DiffusionEngine`` unchanged.
 
     Architecture note: in the multiproc setup the KV cache lives in the worker /
     runner process (GPU side), co-located with the model and KV tensors — so the
-    actual KV *body* is :class:`~vllm_omni.experimental.bde.kv_cache.manager.BDEKVCache`, owned
-    by :class:`~vllm_omni.experimental.bde.runner.BDEModelRunner`. ``BDEEngine`` itself is the
-    thin selection / injection seam; it wires the BDE executor → worker → runner
+    actual KV *body* is :class:`~vllm_omni.experimental.ar_diffusion.kv_cache.manager.ARDiffusionKVCache`, owned
+    by :class:`~vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner`. ``ARDiffusionEngine`` itself is the
+    thin selection / injection seam; it wires the AR-Diffusion executor → worker → runner
     so DreamZero's rollout runs against the runner-owned KV cache.
     """
 
@@ -54,7 +54,7 @@ def __init__(
         od_config: OmniDiffusionConfig,
         scheduler: SchedulerInterface | None = None,
     ) -> None:
-        # Route this engine's workers to BDEModelRunner before the base __init__
+        # Route this engine's workers to ARDiffusionModelRunner before the base __init__
         # builds the executor and spawns workers (od_config is propagated to them).
-        apply_bde_runner_default(od_config)
+        apply_ar_diffusion_runner_default(od_config)
         super().__init__(od_config, scheduler=scheduler)
diff --git a/vllm_omni/experimental/bde/kv_cache/__init__.py b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
similarity index 58%
rename from vllm_omni/experimental/bde/kv_cache/__init__.py
rename to vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
index 673c541e824..31311a706ce 100644
--- a/vllm_omni/experimental/bde/kv_cache/__init__.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
@@ -1,23 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
-"""BDE engine-level KV cache helpers.
+"""AR-Diffusion engine-level KV cache helpers.
 
 Thin glue over vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` /
-``SlidingWindowManager``), used by the Block Diffusion Engine to manage KV for
+``SlidingWindowManager``), used by the AR-Diffusion Engine to manage KV for
 AR-diffusion models. See ``BDE_doc/dreamzero_kv_phase1_plan.md``.
 
 Layout: ``config`` (public knob) · ``paged`` (engine-generic paging mechanics +
-chunk-window eviction spec) · ``manager`` (the BDEKVCache orchestrator + its
-request adapter / pool builders) · ``state`` (the model-facing BDEKVState bridge).
+chunk-window eviction spec) · ``manager`` (the ARDiffusionKVCache orchestrator + its
+request adapter / pool builders) · ``state`` (the model-facing ARDiffusionKVState bridge).
 """
 
-from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.manager import (
-    BDEKVCache,
-    BDERequestAdapter,
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.manager import (
+    ARDiffusionKVCache,
+    ARDiffusionRequestAdapter,
     build_kv_manager,
     compute_num_blocks,
 )
-from vllm_omni.experimental.bde.kv_cache.paged import (
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
     ChunkWindowManager,
     ChunkWindowSpec,
     allocate_kv_pool,
@@ -29,9 +29,9 @@
 )
 
 __all__ = [
-    "BDEKVCache",
-    "BDEKVConfig",
-    "BDERequestAdapter",
+    "ARDiffusionKVCache",
+    "ARDiffusionKVConfig",
+    "ARDiffusionRequestAdapter",
     "ChunkWindowManager",
     "ChunkWindowSpec",
     "allocate_kv_pool",
diff --git a/vllm_omni/experimental/bde/kv_cache/config.py b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
similarity index 75%
rename from vllm_omni/experimental/bde/kv_cache/config.py
rename to vllm_omni/experimental/ar_diffusion/kv_cache/config.py
index 84cb05780c1..d3bb5c00ea2 100644
--- a/vllm_omni/experimental/bde/kv_cache/config.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Configuration for BDE engine-level KV cache management."""
+"""Configuration for AR-Diffusion engine-level KV cache management."""
 
 from __future__ import annotations
 
@@ -7,10 +7,10 @@
 
 
 @dataclass
-class BDEKVConfig:
-    """Settings for the BDE paged KV cache.
+class ARDiffusionKVConfig:
+    """Settings for the AR-Diffusion paged KV cache.
 
-    Disabled by default: when ``enable`` is False the BDE engine behaves exactly
+    Disabled by default: when ``enable`` is False the AR-Diffusion engine behaves exactly
     like the base ``DiffusionEngine`` (no pool, no paged KV).
     """
 
@@ -24,7 +24,7 @@ class BDEKVConfig:
     sink_chunks: int = 0
     # DreamZero-style window reset vs. VGGT-style sliding replace.
     reset_at_boundary: bool = False
-    # Fraction of free device memory budgeted for the BDE KV pool.
+    # Fraction of free device memory budgeted for the AR-Diffusion KV pool.
     gpu_memory_fraction: float = 0.1
 
     @property
diff --git a/vllm_omni/experimental/bde/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
similarity index 85%
rename from vllm_omni/experimental/bde/kv_cache/manager.py
rename to vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
index 50cb5e62510..3c88b053053 100644
--- a/vllm_omni/experimental/bde/kv_cache/manager.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-"""BDEKVCache — the engine-level KV cache orchestrator for one BDE model.
+"""ARDiffusionKVCache — the engine-level KV cache orchestrator for one AR-Diffusion model.
 
-This is the *body* of BDE's KV management: it owns a vLLM ``KVCacheManager`` (a
+This is the *body* of AR-Diffusion's KV management: it owns a vLLM ``KVCacheManager`` (a
 single chunk-window group) and the per-request adapter lifecycle, and exposes the
 per-chunk operations a rollout needs — allocate, slot mapping, commit, window
 lookup, free. It lives in the model runner (worker / GPU side), co-located with
 the model and the KV tensors; the DreamZero pipeline calls these methods during a
-rollout. The main-process ``BDEEngine`` only selects the engine and is otherwise
+rollout. The main-process ``ARDiffusionEngine`` only selects the engine and is otherwise
 thin.
 """
 
@@ -25,8 +25,8 @@
 )
 from vllm.v1.request import RequestStatus
 
-from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.paged import (
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
     ChunkWindowSpec,
     allocate_kv_pool,
     build_window_slots,
@@ -39,7 +39,7 @@
 _log = init_logger(__name__)
 
 
-class BDERequestAdapter:
+class ARDiffusionRequestAdapter:
     """Duck-types the subset of ``vllm.v1.request.Request`` that the
     ``KVCacheManager`` reads (``allocate_slots`` / ``get_computed_blocks`` /
     ``free`` and the coordinator they call into).
@@ -48,7 +48,7 @@ class BDERequestAdapter:
     real ``KVCacheManager`` against this adapter so the surface cannot silently
     drift across vLLM versions.
 
-    A BDE request advances one *chunk* at a time: ``allocate_slots`` is called
+    A AR-Diffusion request advances one *chunk* at a time: ``allocate_slots`` is called
     once per chunk and ``num_computed_tokens`` advances only when a chunk is
     committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk
     reuse the same slots.
@@ -144,8 +144,8 @@ def build_kv_manager(
     )
 
 
-class BDEKVCache:
-    """Owns the paged KV pool + per-request lifecycle for a BDE model.
+class ARDiffusionKVCache:
+    """Owns the paged KV pool + per-request lifecycle for a AR-Diffusion model.
 
     Build once per loaded model (dimensions known); then per request:
     ``begin_request`` → per chunk (``allocate_chunk`` → ``chunk_write_slots`` →
@@ -154,7 +154,7 @@ class BDEKVCache:
 
     def __init__(
         self,
-        config: BDEKVConfig,
+        config: ARDiffusionKVConfig,
         *,
         num_layers: int,
         num_kv_heads: int,
@@ -168,11 +168,11 @@ def __init__(
         device: torch.device | None = None,
     ) -> None:
         if not config.enable:
-            raise ValueError("BDEKVCache built with a disabled BDEKVConfig")
+            raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig")
         if config.window_chunks is None:
             raise ValueError("Phase 1 requires a bounded window (window_chunks)")
         if config.chunk_size <= 0:
-            raise ValueError("BDEKVConfig.chunk_size must be set (> 0)")
+            raise ValueError("ARDiffusionKVConfig.chunk_size must be set (> 0)")
 
         self.config = config
         self.block_size = block_size
@@ -182,7 +182,7 @@ def __init__(
         self.dtype = dtype
         self.cross_attn_length = cross_attn_length
         self.device = device or torch.device("cpu")
-        self._adapters: dict[str, BDERequestAdapter] = {}
+        self._adapters: dict[str, ARDiffusionRequestAdapter] = {}
 
         # -- cross-attention pool (static, one-time-fill) -----------------------
         # Cross-attn KV is computed once from the text encoder and never changes.
@@ -217,7 +217,7 @@ def _cross_pool_bytes(length: int) -> int:
                 self._cross_k.append(torch.empty(cross_shape, dtype=dtype, device=device))
                 self._cross_v.append(torch.empty(cross_shape, dtype=dtype, device=device))
             _log.info(
-                "BDE cross-attn pool: %d layers × (%d tok × %d heads × %d) = %.1f MiB",
+                "AR-Diffusion cross-attn pool: %d layers × (%d tok × %d heads × %d) = %.1f MiB",
                 num_layers,
                 cross_attn_length,
                 num_kv_heads,
@@ -231,7 +231,7 @@ def _cross_pool_bytes(length: int) -> int:
                     self._cross_k_img.append(torch.empty(img_shape, dtype=dtype, device=device))
                     self._cross_v_img.append(torch.empty(img_shape, dtype=dtype, device=device))
                 _log.info(
-                    "BDE cross-attn IMG pool: %d layers × %d img-tok (I2V)",
+                    "AR-Diffusion cross-attn IMG pool: %d layers × %d img-tok (I2V)",
                     num_layers,
                     cross_attn_img_length,
                 )
@@ -266,7 +266,7 @@ def _cross_pool_bytes(length: int) -> int:
         min_blocks = 2 * (config.window_chunks + 1) + 2
         if num_blocks < min_blocks:
             _log.warning(
-                "BDE KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
+                "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
                 "floor (2 CFG branches x (window_chunks=%d + 1) + 2 headroom)",
                 num_blocks,
                 min_blocks,
@@ -333,19 +333,19 @@ def read_cross_kv(self, layer_idx: int, is_negative: bool) -> dict:
 
     # -- request lifecycle ---------------------------------------------------
 
-    def begin_request(self, request_id: str, *, prefill_prefix_tokens: int = 0) -> BDERequestAdapter:
-        adapter = BDERequestAdapter(
+    def begin_request(self, request_id: str, *, prefill_prefix_tokens: int = 0) -> ARDiffusionRequestAdapter:
+        adapter = ARDiffusionRequestAdapter(
             request_id,
             chunk_size=self.spec.chunk_size,
             prefill_prefix_tokens=prefill_prefix_tokens,
         )
         self._adapters[request_id] = adapter
-        _log.debug("BDE begin_request: req=%s prefill=%d", request_id, prefill_prefix_tokens)
+        _log.debug("AR-Diffusion begin_request: req=%s prefill=%d", request_id, prefill_prefix_tokens)
         return adapter
 
-    def end_request(self, adapter: BDERequestAdapter) -> None:
+    def end_request(self, adapter: ARDiffusionRequestAdapter) -> None:
         _log.debug(
-            "BDE end_request: req=%s chunks=%d free=%d",
+            "AR-Diffusion end_request: req=%s chunks=%d free=%d",
             adapter.request_id,
             adapter.completed_chunks,
             self.manager.block_pool.get_num_free_blocks(),
@@ -355,18 +355,18 @@ def end_request(self, adapter: BDERequestAdapter) -> None:
 
     # -- per-chunk operations ------------------------------------------------
 
-    def allocate_chunk(self, adapter: BDERequestAdapter) -> list[int]:
+    def allocate_chunk(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
         """Allocate a chunk's blocks (evicting out-of-window blocks first).
 
         Returns the request's full block table (incl. null_block placeholders).
         """
         blocks = self.manager.allocate_slots(adapter, num_new_tokens=self.spec.chunk_size)
         if blocks is None:
-            raise RuntimeError("BDE KV pool exhausted while allocating a chunk")
+            raise RuntimeError("AR-Diffusion KV pool exhausted while allocating a chunk")
         table = self.block_table(adapter)
         resident = resident_block_ids(table, self.null_block_id)
         _log.debug(
-            "BDE allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d",
+            "AR-Diffusion allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d",
             adapter.request_id,
             adapter.completed_chunks,
             len(table),
@@ -375,10 +375,10 @@ def allocate_chunk(self, adapter: BDERequestAdapter) -> list[int]:
         )
         return table
 
-    def block_table(self, adapter: BDERequestAdapter) -> list[int]:
+    def block_table(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
         return list(self.manager.get_block_ids(adapter.request_id)[0])
 
-    def chunk_write_slots(self, adapter: BDERequestAdapter) -> torch.Tensor:
+    def chunk_write_slots(self, adapter: ARDiffusionRequestAdapter) -> torch.Tensor:
         """Slot mapping for the in-flight chunk — the K/V write target."""
         return chunk_slot_mapping(
             self.block_table(adapter),
@@ -387,23 +387,23 @@ def chunk_write_slots(self, adapter: BDERequestAdapter) -> torch.Tensor:
             self.block_size,
         )
 
-    def window_block_ids(self, adapter: BDERequestAdapter) -> list[int]:
+    def window_block_ids(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
         """Resident (non-null) blocks the read path gathers the window from."""
         return resident_block_ids(self.block_table(adapter), self.null_block_id)
 
-    def commit_chunk(self, adapter: BDERequestAdapter) -> None:
+    def commit_chunk(self, adapter: ARDiffusionRequestAdapter) -> None:
         """Advance the adapter by one chunk after its K/V is written.
 
         This is the standalone per-chunk advance used by the unit lifecycle and the
         ``tests/bde`` suite. In the DreamZero production path the advance happens
-        elsewhere: :meth:`BDEKVState.update_kv_cache` (gather.py) calls
+        elsewhere: :meth:`ARDiffusionKVState.update_kv_cache` (gather.py) calls
         ``adapter.on_chunk_committed()`` directly as it allocates each frame-chunk,
-        and the pipeline's ``_kv_commit`` -> :meth:`BDEKVState.commit_chunk` is only
+        and the pipeline's ``_kv_commit`` -> :meth:`ARDiffusionKVState.commit_chunk` is only
         a no-op log. Call once per chunk, not per denoise step.
         """
-        _log.debug("BDE commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
+        _log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
         adapter.on_chunk_committed()
-        _log.debug("BDE commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
+        _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
 
     # -- pool-backed K/V access (Step 4 — gather / write) --------------------
 
@@ -412,12 +412,12 @@ def write_chunk_kv(
         layer_index: int,
         new_k: torch.Tensor,
         new_v: torch.Tensor,
-        adapter: BDERequestAdapter,
+        adapter: ARDiffusionRequestAdapter,
     ) -> None:
         """Write one layer's committed-chunk K/V into the pool."""
         slots = self.chunk_write_slots(adapter)
         _log.debug(
-            "BDE write: req=%s layer=%d chunk=%d shapes=%s dev=%s",
+            "AR-Diffusion write: req=%s layer=%d chunk=%d shapes=%s dev=%s",
             adapter.request_id,
             layer_index,
             adapter.completed_chunks,
@@ -432,7 +432,7 @@ def write_chunk_kv(
             slots,
         )
 
-    def gather_window(self, layer_index: int, adapter: BDERequestAdapter) -> torch.Tensor:
+    def gather_window(self, layer_index: int, adapter: ARDiffusionRequestAdapter) -> torch.Tensor:
         """Gather the resident-window K/V for one layer.
 
         Returns a ``(2, 1, window, n_heads, head_dim)`` tensor — the format
@@ -447,7 +447,7 @@ def gather_window(self, layer_index: int, adapter: BDERequestAdapter) -> torch.T
             self.spec.sliding_window,
         )
         _log.debug(
-            "BDE gather: req=%s layer=%d blocks=%s window=%s dev=%s",
+            "AR-Diffusion gather: req=%s layer=%d blocks=%s window=%s dev=%s",
             adapter.request_id,
             layer_index,
             window_ids,
@@ -456,7 +456,7 @@ def gather_window(self, layer_index: int, adapter: BDERequestAdapter) -> torch.T
         )
         return window
 
-    def gather_window_all_layers(self, adapter: BDERequestAdapter) -> list[torch.Tensor]:
+    def gather_window_all_layers(self, adapter: ARDiffusionRequestAdapter) -> list[torch.Tensor]:
         """Gather the resident-window K/V for every layer in one shot.
 
         The window's block table is identical across layers within a forward, so
@@ -479,7 +479,7 @@ def gather_window_all_layers(self, adapter: BDERequestAdapter) -> list[torch.Ten
             for i in range(self.num_layers)
         ]
         _log.debug(
-            "BDE gather-all: req=%s layers=%d blocks=%d window=%s",
+            "AR-Diffusion gather-all: req=%s layers=%d blocks=%d window=%s",
             adapter.request_id,
             self.num_layers,
             len(window_ids),
diff --git a/vllm_omni/experimental/bde/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py
similarity index 99%
rename from vllm_omni/experimental/bde/kv_cache/paged.py
rename to vllm_omni/experimental/ar_diffusion/kv_cache/paged.py
index 97da3c72992..3ef4e6288bb 100644
--- a/vllm_omni/experimental/bde/kv_cache/paged.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Generic paging mechanics + chunk-window eviction spec for the BDE engine.
+"""Generic paging mechanics + chunk-window eviction spec for the AR-Diffusion engine.
 
 Engine-generic, model-agnostic primitives — the layer a second model (e.g. the
 Cosmos port) reuses unchanged. Three concerns live here:
diff --git a/vllm_omni/experimental/bde/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
similarity index 84%
rename from vllm_omni/experimental/bde/kv_cache/state.py
rename to vllm_omni/experimental/ar_diffusion/kv_cache/state.py
index 7aaac1bf1c9..19b48094190 100644
--- a/vllm_omni/experimental/bde/kv_cache/state.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""BDEKVState — per-request bridge between a model rollout and the BDE KV pool.
+"""ARDiffusionKVState — per-request bridge between a model rollout and the AR-Diffusion KV pool.
 
 This is the one model-facing piece of the KV stack (everything in ``paged.py`` /
 ``manager.py`` is engine-generic). A second model (e.g. the Cosmos port) adds a
@@ -19,21 +19,21 @@
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.experimental.bde.kv_cache.paged import _drop_batch, pool_write_chunk
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import _drop_batch, pool_write_chunk
 
 _log = init_logger(__name__)
 
 
-class BDEKVState:
-    """Per-request bridge between a DreamZero rollout and the BDE KV pool.
+class ARDiffusionKVState:
+    """Per-request bridge between a DreamZero rollout and the AR-Diffusion KV pool.
 
     Replaces the model-local ``DreamZeroState`` KV methods::
       state.get_kv_caches(neg)    →  bde_state.get_kv_caches(neg)
       state.update_kv_cache(i,kv) →  bde_state.update_kv_cache(i,kv,neg)
       state.create_kv_caches(...) →  no-op (pool is pre-allocated)
 
-    One ``BDEKVState`` per request holds the ``BDEKVCache`` and both CFG adapters
-    (positive / negative). The runner sets ``pipeline._bde_kv_state`` before
+    One ``ARDiffusionKVState`` per request holds the ``ARDiffusionKVCache`` and both CFG adapters
+    (positive / negative). The runner sets ``pipeline._ar_diffusion_kv_state`` before
     ``pipeline.forward(req)``; the pipeline delegates every KV call to it when
     the attribute is present and not None.
 
@@ -59,10 +59,10 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
         # re-materializing the same tensors ~16x. Memoize the gathered window per
         # branch; invalidate on write-back (window grows) and on reset.
         self._gather_cache: dict[bool, list | None] = {False: None, True: None}
-        # Profiling escape hatch: BDE_KV_NO_MEMO=1 disables the per-forward gather
+        # Profiling escape hatch: AR_DIFFUSION_KV_NO_MEMO=1 disables the per-forward gather
         # memoization so every denoise step re-gathers the window (the pre-memo
         # behavior), for A/B measurement of the memoization win. Default: enabled.
-        self._memo_enabled = os.environ.get("BDE_KV_NO_MEMO") != "1"
+        self._memo_enabled = os.environ.get("AR_DIFFUSION_KV_NO_MEMO") != "1"
         # Cross-attn KV is populated once after text encoding (eager), then
         # read every denoising step.  Reset clears these flags so the next
         # forward repopulates from the new text encoding.
@@ -84,7 +84,7 @@ def get_kv_caches(self, is_negative: bool) -> list:
         if self._memo_enabled:
             self._gather_cache[is_negative] = windows
         _log.info(
-            "BDE GET   [%s] source=paged-gather layers=%d window0=%s resident_blocks=%d/%d",
+            "AR-Diffusion GET   [%s] source=paged-gather layers=%d window0=%s resident_blocks=%d/%d",
             branch,
             self.num_layers,
             tuple(windows[0].shape),
@@ -96,14 +96,14 @@ def get_kv_caches(self, is_negative: bool) -> list:
     def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
         """Return pool-backed cross-attn cache dicts.
 
-        Under BDE the cross-attn pool is always populated by ``_kv_populate_cross``
+        Under AR-Diffusion the cross-attn pool is always populated by ``_kv_populate_cross``
         before the first read, so this must never be reached unpopulated — if it
         were, the model would lazily project and write cross KV itself, violating
         engine ownership. Fail loud rather than fall back to the model.
         """
         if not (self._cross_populated[is_negative] and self.kv_cache.cross_attn_length > 0):
             raise RuntimeError(
-                "BDE cross-attn read before _kv_populate_cross (neg=%s, cross_attn_length=%d) "
+                "AR-Diffusion cross-attn read before _kv_populate_cross (neg=%s, cross_attn_length=%d) "
                 "— the engine must own all cross KV" % (is_negative, self.kv_cache.cross_attn_length)
             )
         return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
@@ -122,7 +122,7 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
             # invariant: a non-frame-aligned model would otherwise silently drop
             # the leading new_count % cs tokens and desync num_computed.
             assert new_count % cs == 0, (
-                f"BDE expects frame-aligned seq_len (multiple of chunk_size={cs}), got {new_count}"
+                f"AR-Diffusion expects frame-aligned seq_len (multiple of chunk_size={cs}), got {new_count}"
             )
             n_chunks = new_count // cs
             slots = []
@@ -136,7 +136,7 @@ def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative:
             # on the next forward's first read of this branch.
             self._gather_cache[is_negative] = None
             _log.info(
-                "BDE WRITE [%s] new_tokens=%d -> %d frame-chunks  resident=%d/%d  storing %d layers",
+                "AR-Diffusion WRITE [%s] new_tokens=%d -> %d frame-chunks  resident=%d/%d  storing %d layers",
                 branch,
                 new_count,
                 n_chunks,
@@ -165,13 +165,13 @@ def commit_chunk(self) -> None:
         frame-chunk. This method exists so the pipeline's ``_kv_commit`` bridge has
         a symmetric call; the resident window is retained across forwards.
         """
-        _log.info("BDE COMMIT (paged; resident window retained across forwards)")
+        _log.info("AR-Diffusion COMMIT (paged; resident window retained across forwards)")
 
     def close(self) -> None:
         """Final teardown — free both branches' resident pool blocks.
 
         Unlike :meth:`reset`, this does **not** restart the adapters: the caller is
-        dropping this ``BDEKVState`` for good (session eviction / shutdown), so the
+        dropping this ``ARDiffusionKVState`` for good (session eviction / shutdown), so the
         blocks return to the ``BlockPool`` and the dead adapters are discarded with
         the state. Frees pool ownership that would otherwise leak when a session is
         evicted from the runner's session map.
@@ -183,9 +183,9 @@ def reset(self) -> None:
         """Drop this session's KV — mirrors the model-local ``state.reset()``.
 
         DreamZero resets at the attention-window boundary (``should_reset``); the
-        model starts a fresh sliding window, so BDE must free the resident pool
-        blocks and start a fresh adapter for each branch. The ``BDEKVState``
-        object (and thus ``pipeline._bde_kv_state``) is preserved across the reset
+        model starts a fresh sliding window, so AR-Diffusion must free the resident pool
+        blocks and start a fresh adapter for each branch. The ``ARDiffusionKVState``
+        object (and thus ``pipeline._ar_diffusion_kv_state``) is preserved across the reset
         so the runner's session mapping stays valid — only the pool-backed state
         is recycled.
         """
@@ -197,4 +197,4 @@ def reset(self) -> None:
         self._pending = {False: [], True: []}
         self._gather_cache = {False: None, True: None}
         self._cross_populated = {False: False, True: False}
-        _log.info("BDE RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)
+        _log.info("AR-Diffusion RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)
diff --git a/vllm_omni/experimental/bde/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
similarity index 65%
rename from vllm_omni/experimental/bde/runner.py
rename to vllm_omni/experimental/ar_diffusion/runner.py
index 40a74ba360b..3003e06713f 100644
--- a/vllm_omni/experimental/bde/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -1,20 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
-"""BDEModelRunner — the diffusion model runner for the BDE engine.
+"""ARDiffusionModelRunner — the diffusion model runner for the AR-Diffusion engine.
 
 Subclasses ``DiffusionModelRunner`` and owns the engine-level KV cache
-(:class:`BDEKVCache`). It brackets a request's rollout with the KV lifecycle
-(``begin_request`` / ``end_request``) and exposes the live ``BDEKVCache`` so the
+(:class:`ARDiffusionKVCache`). It brackets a request's rollout with the KV lifecycle
+(``begin_request`` / ``end_request``) and exposes the live ``ARDiffusionKVCache`` so the
 model pipeline (DreamZero) can drive the per-chunk allocate / slot-mapping /
 gather / commit operations during ``pipeline.forward``.
 
-When KV management is disabled (the Phase-1 default), the runner is
-behavior-preserving — it simply defers to the base ``DiffusionModelRunner``.
+Selecting the AR-Diffusion engine enables KV management (no env gate); the KV
+parameters come from ``od_config.ar_diffusion_kv_config`` (deploy/command arg).
 """
 
 from __future__ import annotations
 
 import dataclasses
-import os
 from collections import OrderedDict
 
 import torch
@@ -24,54 +23,45 @@
 from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import MAX_DREAMZERO_SESSIONS
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
-from vllm_omni.experimental.bde.kv_cache.config import BDEKVConfig
-from vllm_omni.experimental.bde.kv_cache.state import BDEKVState
-from vllm_omni.experimental.bde.kv_cache.manager import BDEKVCache
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
 
 logger = init_logger(__name__)
 
 
-def resolve_bde_kv_config(od_config: OmniDiffusionConfig) -> BDEKVConfig:
-    """Resolve the BDE KV config (disabled unless explicitly enabled).
+def resolve_ar_diffusion_kv_config(od_config: OmniDiffusionConfig) -> ARDiffusionKVConfig:
+    """Resolve the AR-Diffusion KV config from the deploy/command arg.
 
-    Precedence: an ``OmniDiffusionConfig`` field (``bde_kv_config`` /
-    ``kv_block_config``), then the ``BDE_KV_ENABLE=1`` env switch (with optional
-    ``BDE_KV_WINDOW_CHUNKS``), else disabled. ``chunk_size`` / ``window_chunks``
-    are finalized from the model geometry at load (see ``_preallocate_kv_cache``).
+    Reads ``od_config.ar_diffusion_kv_config`` (a ``ARDiffusionKVConfig``, a dict of overrides such
+    as ``window_chunks`` / ``gpu_memory_fraction``, or None). Selecting the
+    AR-Diffusion engine implies KV management, so the result is always enabled;
+    ``chunk_size`` / ``window_chunks`` are finalized from the model geometry at load
+    (see ``_preallocate_kv_cache``).
     """
-    raw = getattr(od_config, "bde_kv_config", None)
-    if raw is None:
-        raw = getattr(od_config, "kv_block_config", None)
-    if isinstance(raw, BDEKVConfig):
-        return raw
+    raw = getattr(od_config, "ar_diffusion_kv_config", None)
+    if isinstance(raw, ARDiffusionKVConfig):
+        return dataclasses.replace(raw, enable=True)
     if isinstance(raw, dict):
-        return BDEKVConfig(**raw)
-    if os.environ.get("BDE_KV_ENABLE") == "1":
-        window = os.environ.get("BDE_KV_WINDOW_CHUNKS")
-        gpu_frac = os.environ.get("BDE_KV_GPU_FRACTION")
-        return BDEKVConfig(
-            enable=True,
-            window_chunks=int(window) if window else None,
-            gpu_memory_fraction=float(gpu_frac) if gpu_frac else 0.1,
-        )
-    return BDEKVConfig()
+        return ARDiffusionKVConfig(**{**raw, "enable": True})
+    return ARDiffusionKVConfig(enable=True)
 
 
-class BDEModelRunner(DiffusionModelRunner):
+class ARDiffusionModelRunner(DiffusionModelRunner):
     def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None:
         super().__init__(vllm_config, od_config, device)
-        self.bde_kv_config = resolve_bde_kv_config(od_config)
+        self.ar_diffusion_kv_config = resolve_ar_diffusion_kv_config(od_config)
         # Built after the model is loaded (dimensions known); stays None while KV
         # management is disabled.
-        self.kv_cache: BDEKVCache | None = None
+        self.kv_cache: ARDiffusionKVCache | None = None
         # DreamZero KV is session-scoped (persists across a session's forwards),
-        # so BDE KV state is keyed by session_id and reused, not created per request.
+        # so AR-Diffusion KV state is keyed by session_id and reused, not created per request.
         # Bounded by the same LRU cap as DreamZeroPipeline._states: an evicted
-        # session's BDEKVState owns pool blocks (two CFG adapters), so without a
+        # session's ARDiffusionKVState owns pool blocks (two CFG adapters), so without a
         # bound, session-id churn in a long-running server would leak pool ownership
         # until the KV pool is exhausted. Evicting frees those blocks (state.close()).
-        self._bde_states: OrderedDict[str, BDEKVState] = OrderedDict()
-        self._max_bde_states = MAX_DREAMZERO_SESSIONS
+        self._ar_diffusion_states: OrderedDict[str, ARDiffusionKVState] = OrderedDict()
+        self._max_ar_diffusion_states = MAX_DREAMZERO_SESSIONS
 
     def build_kv_cache(
         self,
@@ -86,14 +76,14 @@ def build_kv_cache(
         cross_attn_length: int = 0,
         cross_attn_img_length: int = 0,
     ) -> None:
-        """Construct the BDEKVCache (preallocating GPU pools on ``self.device``).
+        """Construct the ARDiffusionKVCache (preallocating GPU pools on ``self.device``).
 
         No-op when KV management is disabled.
         """
-        if not self.bde_kv_config.enable:
+        if not self.ar_diffusion_kv_config.enable:
             return
-        self.kv_cache = BDEKVCache(
-            self.bde_kv_config,
+        self.kv_cache = ARDiffusionKVCache(
+            self.ar_diffusion_kv_config,
             num_layers=num_layers,
             num_kv_heads=num_kv_heads,
             head_size=head_size,
@@ -106,11 +96,11 @@ def build_kv_cache(
             device=self.device,
         )
         logger.info(
-            "BDE KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s, "
+            "AR-Diffusion KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s, "
             "layers=%d kv_heads=%d head_dim=%d block_size=%d cross_attn_len=%d device=%s",
             self.kv_cache.num_blocks,
-            self.bde_kv_config.chunk_size,
-            self.bde_kv_config.window_chunks,
+            self.ar_diffusion_kv_config.chunk_size,
+            self.ar_diffusion_kv_config.window_chunks,
             num_layers,
             num_kv_heads,
             head_size,
@@ -123,7 +113,7 @@ def build_kv_cache(
 
     def load_model(self, *args, **kwargs):
         super().load_model(*args, **kwargs)
-        if self.bde_kv_config.enable and self.pipeline is not None:
+        if self.ar_diffusion_kv_config.enable and self.pipeline is not None:
             self._preallocate_kv_cache()
 
     def _infer_frame_seqlen(self) -> int:
@@ -149,7 +139,7 @@ def _preallocate_kv_cache(self) -> None:
         inferred = self._infer_frame_seqlen()
         if frame_seqlen != inferred:
             logger.warning(
-                "BDE frame_seqlen mismatch: transformer=%d deploy-config=%d; using transformer",
+                "AR-Diffusion frame_seqlen mismatch: transformer=%d deploy-config=%d; using transformer",
                 frame_seqlen,
                 inferred,
             )
@@ -169,12 +159,12 @@ def _preallocate_kv_cache(self) -> None:
         # resident window matches max_attention_size exactly (it need not be a whole
         # number of num_frame_per_block causal blocks).
         chunk_size = frame_seqlen
-        window_chunks = self.bde_kv_config.window_chunks or (max_attention_size // frame_seqlen)
+        window_chunks = self.ar_diffusion_kv_config.window_chunks or (max_attention_size // frame_seqlen)
 
-        self.bde_kv_config = dataclasses.replace(self.bde_kv_config, chunk_size=chunk_size, window_chunks=window_chunks)
+        self.ar_diffusion_kv_config = dataclasses.replace(self.ar_diffusion_kv_config, chunk_size=chunk_size, window_chunks=window_chunks)
         free_bytes = torch.cuda.mem_get_info(self.device)[0]
         logger.info(
-            "BDE preallocating (paged): frame_seqlen=%d num_frame_per_block=%d "
+            "AR-Diffusion preallocating (paged): frame_seqlen=%d num_frame_per_block=%d "
             "local_attn_size=%d -> chunk_size=%d window_chunks=%d (window=%d tokens)",
             frame_seqlen,
             num_frame_per_block,
@@ -202,35 +192,35 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
 
         kv = self.kv_cache
         # DreamZero KV is session-scoped (the model state persists across a
-        # session's forwards), so the BDE KV state is keyed by session_id and
+        # session's forwards), so the AR-Diffusion KV state is keyed by session_id and
         # reused — matching how pipeline.forward resolves the model-local state.
         extra_args = req.sampling_params.extra_args or {}
         session_id = str(extra_args.get("session_id") or "default")
-        state = self._bde_states.get(session_id)
+        state = self._ar_diffusion_states.get(session_id)
         if state is None:
             pos = kv.begin_request(f"bde__{session_id}")
             neg = kv.begin_request(f"bde__{session_id}__neg")
-            state = BDEKVState(kv, pos, neg, num_layers=kv.num_layers)
-            self._bde_states[session_id] = state
+            state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers)
+            self._ar_diffusion_states[session_id] = state
             # Evict the least-recently-used session(s) past the cap, freeing their
-            # pool blocks — mirrors DreamZeroPipeline._states so the BDE pool can't
+            # pool blocks — mirrors DreamZeroPipeline._states so the AR-Diffusion pool can't
             # outlive the model-local state map it shadows.
-            while len(self._bde_states) > self._max_bde_states:
-                old_id, old_state = self._bde_states.popitem(last=False)
+            while len(self._ar_diffusion_states) > self._max_ar_diffusion_states:
+                old_id, old_state = self._ar_diffusion_states.popitem(last=False)
                 old_state.close()
-                logger.debug("BDE evicted session=%s (LRU); freed pool blocks", old_id)
+                logger.debug("AR-Diffusion evicted session=%s (LRU); freed pool blocks", old_id)
         # Track recency on every access (hit or miss) so the LRU order is correct.
-        self._bde_states.move_to_end(session_id)
+        self._ar_diffusion_states.move_to_end(session_id)
         logger.debug(
-            "BDE execute_model: req=%s session=%s chunk_size=%d window_chunks=%s num_blocks=%d",
+            "AR-Diffusion execute_model: req=%s session=%s chunk_size=%d window_chunks=%s num_blocks=%d",
             req.request_id,
             session_id,
             kv.spec.chunk_size,
             kv.config.window_chunks,
             kv.num_blocks,
         )
-        self.pipeline._bde_kv_state = state
+        self.pipeline._ar_diffusion_kv_state = state
         try:
             return super().execute_model(req)
         finally:
-            self.pipeline._bde_kv_state = None
+            self.pipeline._ar_diffusion_kv_state = None

From 3dae6d1197996021d1a577fde767127a4e6f0cb9 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Fri, 26 Jun 2026 09:16:24 +0000
Subject: [PATCH 35/51] [AR-Diffusion] Minimize public-base footprint:
 config-driven engine selection
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Push all AR-Diffusion/DreamZero specifics out of the public diffusion base into
the deploy config + experimental package, so the public diff is a tiny generic seam:

- diffusion_engine.py: resolve_engine_class/make_engine are now a generic factory
  ("default" / DiffusionEngine subclass / import-path string). Removed the hardcoded
  "ar_diffusion" branch and its experimental import — no coupling to experimental/.
- data.py: removed default_engine_backend_for_model, resolve_default_engine_backend,
  _AR_DIFFUSION_ENGINE_MODEL_CLASSES, the two engine_backend-rewrite sites in
  enrich_config + from_kwargs, the ar_diffusion_kv_config field, and the structured-
  config mirror note. Public engine-routing addition is now just the generic
  engine_backend / diffusion_model_runner_cls fields. from_kwargs returns to its
  clean cls(**filtered_kwargs) tail (no #4425 conflict).
- dreamzero.yaml: select the engine via engine_backend=<ARDiffusionEngine qualname>.
- runner: KV overrides read from model_config["ar_diffusion_kv_config"] (defaults work).
- tests: generic dispatch tests (qualname -> ARDiffusionEngine); dropped per-model
  default tests. 51 passed; E2E DreamZero generates via the yaml with no env.

Trade-off: a bare load without the deploy yaml no longer auto-selects the engine;
DreamZero is selected via dreamzero.yaml (its single supported path).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/ar_diffusion/test_engine_routing.py     | 78 ++++---------------
 vllm_omni/deploy/dreamzero.yaml               |  2 +
 vllm_omni/diffusion/data.py                   | 69 ++--------------
 vllm_omni/diffusion/diffusion_engine.py       | 14 +---
 vllm_omni/experimental/ar_diffusion/runner.py | 18 +++--
 5 files changed, 36 insertions(+), 145 deletions(-)

diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py
index 0c380ab730f..5254eb1a3fe 100644
--- a/tests/ar_diffusion/test_engine_routing.py
+++ b/tests/ar_diffusion/test_engine_routing.py
@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Routing tests for AR-Diffusion engine selection (Phase 1, PR-1).
+"""Tests for the generic engine-backend dispatcher + AR-Diffusion runner wiring.
 
-Covers ``DiffusionEngine.resolve_engine_class`` — the ``engine_backend``
-config-field dispatcher that routes a request to ``ARDiffusionEngine`` vs the base
-``DiffusionEngine``. Resolution is tested without constructing an engine
-(construction runs a dummy forward that needs a real model).
+``DiffusionEngine.resolve_engine_class`` is a generic dispatcher (``"default"`` / a
+``DiffusionEngine`` subclass / an import-path string). DreamZero selects the
+AR-Diffusion engine via its deploy config's ``engine_backend`` qualname — no
+DreamZero/ar_diffusion knowledge lives in the public base, so the routing check
+here is simply that the ``ARDiffusionEngine`` qualname resolves correctly.
 """
 
 from dataclasses import fields
@@ -13,11 +14,7 @@
 import pytest
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
-from vllm_omni.diffusion.data import (
-    OmniDiffusionConfig,
-    default_engine_backend_for_model,
-    resolve_default_engine_backend,
-)
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.experimental.ar_diffusion.engine import (
     AR_DIFFUSION_MODEL_RUNNER_CLS,
@@ -34,16 +31,15 @@ def test_default_resolves_to_diffusion_engine():
     assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
 
 
-def test_bde_key_resolves_to_bde_engine():
-    assert DiffusionEngine.resolve_engine_class(_cfg("ar_diffusion")) is ARDiffusionEngine
-
-
 def test_subclass_type_is_returned():
     assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine
 
 
-def test_qualname_string_resolves():
-    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
+def test_ar_diffusion_qualname_resolves():
+    # How DreamZero's deploy config selects the engine: a full import-path string.
+    cls = DiffusionEngine.resolve_engine_class(
+        _cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine")
+    )
     assert cls is ARDiffusionEngine
 
 
@@ -62,7 +58,7 @@ def test_bad_qualname_raises():
         DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine"))
 
 
-def test_bde_engine_is_diffusion_engine_subclass():
+def test_ar_diffusion_engine_is_diffusion_engine_subclass():
     assert issubclass(ARDiffusionEngine, DiffusionEngine)
 
 
@@ -71,55 +67,10 @@ def test_config_engine_backend_field_defaults_to_default():
     assert field.default == "default"
 
 
-# --- DreamZero defaults to the AR-Diffusion engine -----------------------------------
-
-
-def test_dreamzero_pipeline_defaults_to_bde():
-    assert default_engine_backend_for_model("DreamZeroPipeline") == "ar_diffusion"
-
-
-def test_unknown_model_defaults_to_none():
-    assert default_engine_backend_for_model("WanS2VPipeline") is None
-
-
-def test_none_model_defaults_to_none():
-    assert default_engine_backend_for_model(None) is None
-
-
-def test_dreamzero_default_routes_to_bde_engine():
-    # End-to-end of the routing: DreamZero's per-model default backend resolves
-    # to ARDiffusionEngine through the same dispatcher used by make_engine.
-    backend = default_engine_backend_for_model("DreamZeroPipeline")
-    assert DiffusionEngine.resolve_engine_class(_cfg(backend)) is ARDiffusionEngine
-
-
-# --- AR-Diffusion routing: unconditional per-model default; explicit arg overrides ---
-
-
-def test_dreamzero_default_applied_unconditionally():
-    # DreamZero + unset ("default") backend -> auto-route to ar_diffusion (no env gate).
-    assert resolve_default_engine_backend("DreamZeroPipeline", "default") == "ar_diffusion"
-
-
-def test_explicit_backend_not_overridden():
-    # An explicit engine_backend command arg bypasses the per-model default.
-    assert resolve_default_engine_backend("DreamZeroPipeline", "ar_diffusion") is None
-    assert (
-        resolve_default_engine_backend(
-            "DreamZeroPipeline", "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
-        )
-        is None
-    )
-
-
-def test_non_ar_diffusion_model_not_routed():
-    assert resolve_default_engine_backend("WanS2VPipeline", "default") is None
-
-
 # --- AR-Diffusion worker/runner wiring -----------------------------------------------
 
 
-def test_bde_model_runner_cls_resolves():
+def test_ar_diffusion_model_runner_cls_resolves():
     cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS)
     assert cls.__name__ == "ARDiffusionModelRunner"
 
@@ -137,7 +88,6 @@ def test_apply_runner_default_respects_explicit_choice():
 
 
 def test_worker_runner_selection_prefers_override():
-    # Mirrors the worker hook: an od_config override wins over the platform default.
     over = SimpleNamespace(diffusion_model_runner_cls=AR_DIFFUSION_MODEL_RUNNER_CLS)
     assert (getattr(over, "diffusion_model_runner_cls", None) or "PLATFORM") == AR_DIFFUSION_MODEL_RUNNER_CLS
     unset = SimpleNamespace(diffusion_model_runner_cls=None)
diff --git a/vllm_omni/deploy/dreamzero.yaml b/vllm_omni/deploy/dreamzero.yaml
index 079101ba0e0..62b3a5eed9b 100644
--- a/vllm_omni/deploy/dreamzero.yaml
+++ b/vllm_omni/deploy/dreamzero.yaml
@@ -14,6 +14,8 @@ stages:
     max_num_seqs: 1
     enforce_eager: false
     model_class_name: DreamZeroPipeline
+    # DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache).
+    engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine
     model_config:
       default_robot_embodiment: roboarena
       policy_server_config:
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 5c2f3285df4..d08b4aa73ff 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -493,37 +493,6 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
     return None
 
 
-# Model classes that default to the AR-Diffusion (AR-diffusion) engine.
-_AR_DIFFUSION_ENGINE_MODEL_CLASSES = frozenset({"DreamZeroPipeline"})
-
-
-def default_engine_backend_for_model(model_class_name: str | None) -> str | None:
-    """Engine backend a model opts into by default.
-
-    Returns ``"ar_diffusion"`` for AR-diffusion world models (DreamZero today) so they run on
-    the AR-Diffusion engine with engine-level KV cache management; ``None`` keeps the
-    base ``DiffusionEngine``. Extend ``_AR_DIFFUSION_ENGINE_MODEL_CLASSES`` as more AR-DiT
-    models migrate. Applied when ``engine_backend`` is unset (see
-    :func:`resolve_default_engine_backend`); an explicit command/deploy arg overrides.
-    """
-    if model_class_name in _AR_DIFFUSION_ENGINE_MODEL_CLASSES:
-        return "ar_diffusion"
-    return None
-
-
-def resolve_default_engine_backend(model_class_name: str | None, current_backend: str | None) -> str | None:
-    """Engine backend to apply for an auto-routed model, or ``None`` to leave as-is.
-
-    Applies the per-model default (:func:`default_engine_backend_for_model`) when
-    ``current_backend`` is the unset ``"default"`` sentinel; an explicit
-    ``engine_backend`` command/deploy arg is left untouched. Shared by the two
-    routing sites (``enrich_config`` and ``from_kwargs``).
-    """
-    if current_backend != "default":
-        return None
-    return default_engine_backend_for_model(model_class_name)
-
-
 @dataclass
 class OmniDiffusionConfig:
     # Model and path configuration (for convenience)
@@ -569,30 +538,15 @@ class OmniDiffusionConfig:
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None
 
-    # Engine backend selection. Resolved by ``DiffusionEngine.make_engine``
-    # (mirrors ``distributed_executor_backend`` / ``DiffusionExecutor.get_class``):
-    #   "default" -> DiffusionEngine
-    #   "ar_diffusion"     -> ARDiffusionEngine (AR-diffusion engine with engine-level KV cache
-    #                management); also accepts a DiffusionEngine subclass or an
-    #                import path string.
+    # Engine backend selection, resolved generically by ``DiffusionEngine.make_engine``
+    # (mirrors ``DiffusionExecutor.get_class``): "default" -> DiffusionEngine, a
+    # DiffusionEngine subclass, or an import-path string (set e.g. by a deploy config).
     engine_backend: str = "default"
 
-    # Optional override for the diffusion model runner class (import path).
-    # The worker uses it instead of the platform default when set; the AR-Diffusion engine
-    # sets it to ARDiffusionModelRunner. Unset -> platform default (existing behavior).
+    # Optional override for the diffusion model runner class (import path). The worker
+    # uses it instead of the platform default when set; unset -> platform default.
     diffusion_model_runner_cls: str | None = None
 
-    # AR-Diffusion engine KV-cache parameters (command/deploy arg, e.g. in the
-    # deploy yaml). Dict of window_chunks / gpu_memory_fraction / sink_chunks /
-    # reset_at_boundary; None uses engine defaults. Selecting the AR-Diffusion
-    # engine enables KV (no env gate).
-    ar_diffusion_kv_config: dict | None = None
-    # DEFERRED (#4425 / RFC #4021 structured config): when rebasing onto the
-    # structured VllmOmniConfig, mirror engine_backend + diffusion_model_runner_cls
-    # into its DiffusionConfig (engine/runtime group) and parity test. enrich_config()
-    # there only copies fields present on DiffusionConfig, so without the mirror
-    # DreamZero's AR-Diffusion routing silently drops on the structured-config path. (PR #4534)
-
     # HuggingFace specific parameters
     trust_remote_code: bool = False
     revision: str | None = None
@@ -1119,12 +1073,6 @@ def enrich_config(self) -> None:
                         self.model_class_name = "DreamZeroPipeline"
                         self.set_tf_model_config(TransformerConfig())
                         self.update_multimodal_support()
-                        # DreamZero is an AR-diffusion world model: route it to the
-                        # AR-Diffusion engine (engine-level KV cache management) unless
-                        # the caller explicitly selected an engine backend.
-                        backend = resolve_default_engine_backend(self.model_class_name, self.engine_backend)
-                        if backend is not None:
-                            self.engine_backend = backend
                     else:
                         raise
                 elif architectures and len(architectures) == 1:
@@ -1196,13 +1144,6 @@ def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
         filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
 
         instance = cls(**filtered_kwargs)
-        # Route per-model engines after model_class_name is finalized. This covers
-        # the explicit-model_class_name path (e.g. deploy configs) that the
-        # __post_init__ auto-detect branch does not reach. An explicit engine_backend
-        # command/deploy arg overrides the per-model default.
-        backend = resolve_default_engine_backend(instance.model_class_name, instance.engine_backend)
-        if backend is not None:
-            instance.engine_backend = backend
         return instance
 
 
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 83716718247..f80fb624ad8 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -416,10 +416,10 @@ def _handle_finished_requests(
     def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
         """Resolve the engine class selected by ``config.engine_backend``.
 
-        Mirrors ``DiffusionExecutor.get_class``: accepts the string keys
-        ``"default"`` / ``"ar_diffusion"``, a ``DiffusionEngine`` subclass, or an import
-        path string. Kept separate from :meth:`make_engine` so the selection is
-        testable without constructing an engine (which runs a dummy forward).
+        Mirrors ``DiffusionExecutor.get_class``: accepts ``"default"``, a
+        ``DiffusionEngine`` subclass, or an import-path string (e.g. a deploy
+        config's ``engine_backend``). Kept separate from :meth:`make_engine` so the
+        selection is testable without constructing an engine (which runs a dummy forward).
 
         Args:
             config: The configuration for the diffusion engine.
@@ -435,12 +435,6 @@ def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
             return backend
         if backend == "default":
             return DiffusionEngine
-        if backend == "ar_diffusion":
-            # Lazy import to avoid a circular dependency (ar_diffusion.engine imports this
-            # module), mirroring how DiffusionExecutor.get_class imports backends.
-            from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
-
-            return ARDiffusionEngine
         if isinstance(backend, str):
             try:
                 engine_class = resolve_obj_by_qualname(backend)
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index 3003e06713f..6320c4cdb7d 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -7,8 +7,8 @@
 model pipeline (DreamZero) can drive the per-chunk allocate / slot-mapping /
 gather / commit operations during ``pipeline.forward``.
 
-Selecting the AR-Diffusion engine enables KV management (no env gate); the KV
-parameters come from ``od_config.ar_diffusion_kv_config`` (deploy/command arg).
+Selecting the AR-Diffusion engine enables KV management (no env gate); optional KV
+parameters come from the deploy yaml's ``model_config["ar_diffusion_kv_config"]``.
 """
 
 from __future__ import annotations
@@ -33,13 +33,17 @@
 def resolve_ar_diffusion_kv_config(od_config: OmniDiffusionConfig) -> ARDiffusionKVConfig:
     """Resolve the AR-Diffusion KV config from the deploy/command arg.
 
-    Reads ``od_config.ar_diffusion_kv_config`` (a ``ARDiffusionKVConfig``, a dict of overrides such
-    as ``window_chunks`` / ``gpu_memory_fraction``, or None). Selecting the
-    AR-Diffusion engine implies KV management, so the result is always enabled;
-    ``chunk_size`` / ``window_chunks`` are finalized from the model geometry at load
-    (see ``_preallocate_kv_cache``).
+    Selecting the AR-Diffusion engine implies KV management, so the result is always
+    enabled. Optional overrides (``window_chunks`` / ``gpu_memory_fraction`` /
+    ``sink_chunks`` / ``reset_at_boundary``) may be supplied via the deploy yaml as
+    ``model_config["ar_diffusion_kv_config"]``. ``chunk_size`` / ``window_chunks`` are
+    finalized from the model geometry at load (see ``_preallocate_kv_cache``).
     """
     raw = getattr(od_config, "ar_diffusion_kv_config", None)
+    if raw is None:
+        model_config = getattr(od_config, "model_config", None)
+        if isinstance(model_config, dict):
+            raw = model_config.get("ar_diffusion_kv_config")
     if isinstance(raw, ARDiffusionKVConfig):
         return dataclasses.replace(raw, enable=True)
     if isinstance(raw, dict):

From cb8467b3acc6bb2006838c09319de2587193faf7 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 29 Jun 2026 02:30:19 +0000
Subject: [PATCH 36/51] [AR-Diffusion] Restore profiling/benchmark scaffolding
 on a side branch

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../models/dreamzero/video_export_worker.py       | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
index d039e53c816..530a3fbec89 100644
--- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
+++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
@@ -9,6 +9,21 @@
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
+    def gpu_mem_stats(self) -> dict:
+        """Peak GPU memory (GiB) for this worker's device, for profiling.
+
+        ``max_memory_reserved`` is the caching-allocator high-water mark (the best
+        in-process proxy for "peak VRAM"); ``max_memory_allocated`` is the live
+        tensor high-water mark. Both are monotonic from process start, so the
+        end-of-run values capture the whole rollout's peak.
+        """
+        dev = torch.accelerator.current_device_index()
+        return {
+            "device": int(dev),
+            "peak_reserved_gib": torch.accelerator.max_memory_reserved(dev) / (1024**3),
+            "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3),
+        }
+
     @staticmethod
     def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
         decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()

From 61d4f3c343d4476688e9abc29f878c204ecc1633 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 29 Jun 2026 08:16:30 +0000
Subject: [PATCH 37/51] [torch-profiler-analyze] Introduce new skill for
 analyzing PyTorch profiler artifacts

- Added `torch-profiler-analyze` skill to read and interpret PyTorch profiler outputs, including timeline and flamegraph generation.
- Enhanced `SKILL.md` documentation with usage instructions and command examples.
- Implemented `trace_reader.py` for summarizing trace events with multi-rank aggregation and named-region attribution.
- Created `flamegraph_reader.py` for generating flamegraphs from folded stack files.
- Added artifact format references and NVTX region documentation for better user guidance.

This commit enhances the profiling capabilities for vLLM Omni models, providing a comprehensive toolset for performance analysis.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../models/dreamzero/video_export_worker.py      | 14 ++++++++++++++
 vllm_omni/experimental/ar_diffusion/runner.py    | 16 +++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
index 530a3fbec89..ba980a33e08 100644
--- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
+++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
@@ -24,6 +24,20 @@ def gpu_mem_stats(self) -> dict:
             "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3),
         }
 
+    def ar_diffusion_perf_stats(self, reset: bool = True) -> dict:
+        """Per-request server-side forward E2E times (s) recorded by the AR-Diffusion runner.
+
+        The worker-side compute time per ``execute_model`` — the analog of the
+        upstream server's per-request E2E — which excludes the engine<->worker IPC
+        that the client's ``omni.generate`` wall time includes. Returns ``[]`` when
+        the runner is the base (KV-disabled) runner that does not record timings.
+        """
+        runner = self.model_runner
+        times = list(getattr(runner, "_perf_e2e_times", []) or [])
+        if reset and hasattr(runner, "_perf_e2e_times"):
+            runner._perf_e2e_times = []
+        return {"server_e2e_s": times}
+
     @staticmethod
     def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
         decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index 6320c4cdb7d..66d74bae7af 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import dataclasses
+import time
 from collections import OrderedDict
 
 import torch
@@ -66,6 +67,12 @@ def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None:
         # until the KV pool is exhausted. Evicting frees those blocks (state.close()).
         self._ar_diffusion_states: OrderedDict[str, ARDiffusionKVState] = OrderedDict()
         self._max_ar_diffusion_states = MAX_DREAMZERO_SESSIONS
+        # Per-request server-side forward E2E times (seconds), recorded in
+        # execute_model. This is the worker-side compute time — the analog of the
+        # upstream server's per-request E2E — and excludes the engine<->worker IPC
+        # that the client's omni.generate wall time includes. Read (and cleared)
+        # by the ar_diffusion_perf_stats worker RPC for the perf-compare summary.
+        self._perf_e2e_times: list[float] = []
 
     def build_kv_cache(
         self,
@@ -224,7 +231,14 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
             kv.num_blocks,
         )
         self.pipeline._ar_diffusion_kv_state = state
+        # Time the worker-side forward (true per-request compute E2E). The CUDA
+        # sync below makes the stop reflect completed GPU work, not just dispatch.
+        _e2e_t0 = time.perf_counter()
         try:
-            return super().execute_model(req)
+            out = super().execute_model(req)
         finally:
             self.pipeline._ar_diffusion_kv_state = None
+        if self.device is not None and torch.cuda.is_available():
+            torch.cuda.synchronize(self.device)
+        self._perf_e2e_times.append(time.perf_counter() - _e2e_t0)
+        return out

From 0c18e08d25a7226e226c1c85c504127ee2f26d89 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Mon, 29 Jun 2026 09:41:57 +0000
Subject: [PATCH 38/51] [AR-Diffusion] Move CUDA-graph capture off the serving
 hot path via a startup warm-up
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The DiT is compiled dynamic=False and the paged KV window grows 1->3->5->7->9
frames over the first chunks, so each new shape recompiles + recaptures — a
~40s tax on the first user requests. Add ARDiffusionModelRunner._warmup_ar_rollout():
at load (gated by not enforce_eager + ARDiffusionKVConfig.warmup_cudagraph), drive a
synthetic DreamZero rollout through one window-fill on a throwaway session, then free
it (KV pool restored exactly; cleaned from both _ar_diffusion_states and pipeline._states).
Result: serving is fast from chunk 0 (first forward ~31x faster), warm-up moves to
startup (+~21s one-time), numerically inert (warm-up on/off diff <= same-config baseline).

Also add --seed to ar_diffusion_perf_compare.py for deterministic numerics checks.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../ar_diffusion/kv_cache/config.py           |   6 +
 vllm_omni/experimental/ar_diffusion/runner.py | 119 ++++++++++++++++++
 2 files changed, 125 insertions(+)

diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/config.py b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
index d3bb5c00ea2..4b3c737751d 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
@@ -26,6 +26,12 @@ class ARDiffusionKVConfig:
     reset_at_boundary: bool = False
     # Fraction of free device memory budgeted for the AR-Diffusion KV pool.
     gpu_memory_fraction: float = 0.1
+    # When CUDA graph / torch.compile is on (not enforce_eager), pre-capture the
+    # DiT graphs for every window-fill shape at load time via a synthetic rollout,
+    # so the serving run is fast from the first chunk. No effect when eager.
+    warmup_cudagraph: bool = True
+    # Also capture the post-window-boundary (reset-cycle) forward during warm-up.
+    warmup_capture_reset: bool = False
 
     @property
     def sliding_window(self) -> int | None:
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index 66d74bae7af..09a259a0a0d 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -14,9 +14,11 @@
 from __future__ import annotations
 
 import dataclasses
+import math
 import time
 from collections import OrderedDict
 
+import numpy as np
 import torch
 from vllm.logger import init_logger
 
@@ -126,6 +128,10 @@ def load_model(self, *args, **kwargs):
         super().load_model(*args, **kwargs)
         if self.ar_diffusion_kv_config.enable and self.pipeline is not None:
             self._preallocate_kv_cache()
+            # Pre-capture the CUDA graphs for every window-fill shape at load time
+            # (only when compiling/cuda-graph is on) so serving is fast from chunk 0.
+            if not self.od_config.enforce_eager and self.ar_diffusion_kv_config.warmup_cudagraph:
+                self._warmup_ar_rollout()
 
     def _infer_frame_seqlen(self) -> int:
         """frame_seqlen = (H//8)*(W//8)//4 from the configured image_resolution."""
@@ -242,3 +248,116 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
             torch.cuda.synchronize(self.device)
         self._perf_e2e_times.append(time.perf_counter() - _e2e_t0)
         return out
+
+    # -- cuda-graph warm-up ----------------------------------------------------
+
+    #: Dedicated session id for the warm-up rollout; cleaned from both state maps.
+    _WARMUP_SID = "__ardiffusion_warmup__"
+
+    def _synth_robot_obs(self, h: int, w: int, n_frames: int) -> dict:
+        """Zero-filled roboarena observation matching the client's obs layout.
+
+        Mirrors ``examples/.../client_schedule.make_obs_from_video`` (2 ext + wrist
+        cameras, 7/6/1 state vectors) so it survives ``_transform_robot_obs`` exactly
+        like a real request. One frame -> ``(H,W,3)``; chunk -> ``(n_frames,H,W,3)``.
+        """
+        img = (
+            np.zeros((h, w, 3), dtype=np.uint8)
+            if n_frames == 1
+            else np.zeros((n_frames, h, w, 3), dtype=np.uint8)
+        )
+        return {
+            "observation/exterior_image_0_left": img,
+            "observation/exterior_image_1_left": img,
+            "observation/wrist_image_left": img,
+            "observation/joint_position": np.zeros(7, dtype=np.float32),
+            "observation/cartesian_position": np.zeros(6, dtype=np.float32),
+            "observation/gripper_position": np.zeros(1, dtype=np.float32),
+            "prompt": "warmup",
+            "session_id": self._WARMUP_SID,
+        }
+
+    def _warmup_ar_rollout(self) -> None:
+        """Drive a synthetic rollout through one window-fill so every DiT graph is
+        captured at load (off the serving hot path). Frees the synthetic session and
+        leaves the KV pool exactly as found. Never raises — falls back to lazy capture.
+        """
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        kv = self.kv_cache
+        if kv is None:
+            return
+        sid = self._WARMUP_SID
+        t = self.pipeline.transformer
+        nfpb = int(t.num_frame_per_block)
+        window_chunks = int(kv.config.window_chunks or 0)
+        if window_chunks <= 0:
+            return
+        # Distinct DiT window shapes are walked by one fill: prefill + ceil((W-1)/nfpb)
+        # chunk forwards. +1 optionally to also capture the reset-cycle forward.
+        n_forwards = 1 + math.ceil(max(0, window_chunks - 1) / nfpb)
+        if self.ar_diffusion_kv_config.warmup_capture_reset:
+            n_forwards += 1
+        res = self._image_resolution()
+        h, w = int(res[0]), int(res[1])
+
+        free_before = kv.manager.block_pool.get_num_free_blocks()
+        # The step cache is left active: the DiT input shape is per-chunk (the resident
+        # window length), and the first denoise step of every chunk always computes, so
+        # every window shape is captured regardless of how many later steps are skipped.
+        logger.info(
+            "AR-Diffusion cudagraph warm-up: up to %d synthetic forwards (window_chunks=%d, nfpb=%d)",
+            n_forwards, window_chunks, nfpb,
+        )
+        try:
+            for i in range(n_forwards):
+                n_frames = 1 if i == 0 else 4  # client convention: 1-frame prefill, 4-frame chunks
+                obs = self._synth_robot_obs(h, w, n_frames)
+                sp = OmniDiffusionSamplingParams(
+                    extra_args={"reset": i == 0, "session_id": sid, "robot_obs": obs}
+                )
+                req = OmniDiffusionRequest(
+                    prompts=["warmup"], sampling_params=sp, request_id=f"ardiffusion-warmup-{i}"
+                )
+                self.execute_model(req)
+                # Stop once the resident window is full — the remaining shapes are
+                # already captured (the window caps/resets through the same set).
+                ar_state = self._ar_diffusion_states.get(sid)
+                if (
+                    not self.ar_diffusion_kv_config.warmup_capture_reset
+                    and ar_state is not None
+                    and len(kv.window_block_ids(ar_state.pos)) >= window_chunks
+                ):
+                    break
+        except Exception as e:  # noqa: BLE001 — warm-up must never break model load
+            logger.warning("AR-Diffusion cudagraph warm-up failed (%s); using lazy capture.", e)
+        finally:
+            # Free the synthetic session from BOTH maps (KV pool + pipeline state).
+            ar_state = self._ar_diffusion_states.pop(sid, None)
+            if ar_state is not None:
+                try:
+                    ar_state.close()  # return KV blocks to the BlockPool
+                except Exception as e:  # noqa: BLE001
+                    logger.warning("AR-Diffusion warm-up: state.close() failed (%s)", e)
+            try:
+                self.pipeline._states.pop(sid, None)
+            except Exception:  # noqa: BLE001
+                pass
+            # Drop any per-forward timings the warm-up recorded so serving stats start clean.
+            self._perf_e2e_times.clear()
+            free_after = kv.manager.block_pool.get_num_free_blocks()
+            if free_after != free_before:
+                logger.warning(
+                    "AR-Diffusion warm-up: KV pool not restored (free %d -> %d) — possible leak",
+                    free_before, free_after,
+                )
+            else:
+                logger.info(
+                    "AR-Diffusion cudagraph warm-up complete; KV pool restored (%d free blocks)",
+                    free_after,
+                )
+
+    def _image_resolution(self) -> list[int]:
+        mc = getattr(self.od_config, "model_config", None) or {}
+        psc = (mc.get("policy_server_config") if isinstance(mc, dict) else None) or {}
+        return psc.get("image_resolution", [180, 320])

From 2abab0fbb38187c1c763733d617d03c99dd82ea2 Mon Sep 17 00:00:00 2001
From: amy-why-3459 <wuhaiyan17@huawei.com>
Date: Fri, 26 Jun 2026 23:00:28 +0800
Subject: [PATCH 39/51] Fix vae encoder's torch.compile

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
(cherry picked from commit 65d0bc5fe3ceb4df31d03491da714aac33445376)
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../models/dreamzero/pipeline_dreamzero.py    | 169 ++++++++++++++++--
 .../models/dreamzero/state_dreamzero.py       |   8 +
 2 files changed, 160 insertions(+), 17 deletions(-)

diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index e43f465f774..87b666d0212 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -442,6 +442,8 @@ def setup_compile(self) -> None:
         Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False``
         on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and
         compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop).
+        Incremental VAE encode (``_vae_encode_encoder_chunk``) stays eager because
+        Wan ``feat_cache`` mutation is incompatible with CUDAGraph capture.
         DiT blocks use per-block ``fullgraph=True``.
         """
         if not torch.cuda.is_available():
@@ -476,9 +478,9 @@ def setup_compile(self) -> None:
             logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
 
         try:
-            self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
+            self.vae._encode = torch.compile(self.vae._encode, **compile_ro)
         except Exception as exc:
-            logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
+            logger.warning("DreamZero: vae._encode compile failed (%s); skipping.", exc)
 
         compiled_blocks = 0
         for block in self.transformer.blocks:
@@ -517,10 +519,24 @@ def warmup_compile(self) -> None:
 
             try:
                 image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
-                self._encode_image(image, self.num_frames, 180, 320)
+                self._encode_image(image, self.num_frames, 180, 320, state=self.state)
             except Exception as exc:
                 logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
 
+            try:
+                self.state.reset_vae_encoder_stream()
+                dummy_video = torch.zeros(1, 3, 1, 180, 320, dtype=torch.bfloat16, device=device)
+                self._vae_stream_seed(self.state, dummy_video[:, :, :1])
+                for _ in range(4):
+                    self._vae_stream_append_frame(self.state, dummy_video[:, :, :1])
+                self._vae_stream_get_observation_latents(
+                    self.state,
+                    self.num_frame_per_block,
+                    dtype=torch.bfloat16,
+                )
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (vae encode stream) skipped: %s", exc)
+
             try:
                 latent_h, latent_w = 180 // 8, 320 // 8
                 dummy_latent = torch.zeros(
@@ -624,6 +640,8 @@ def _encode_image(
         num_frames: int,
         height: int,
         width: int,
+        *,
+        state: DreamZeroState | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Encode first frame via CLIP + VAE.
         Returns: (clip_feas, ys, image_latent)
@@ -659,11 +677,140 @@ def _encode_image(
             new_image = y[:, :, 0:1]
             y = torch.concat([msk, y], dim=1)
 
+            if state is not None:
+                # Seed the AR streaming encoder after the compiled full encode above;
+                # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph.
+                self._vae_stream_seed(state, image_input[:, :, :1])
+
         return clip_context, y, new_image
 
+    def _vae_patchify(self, videos: torch.Tensor) -> torch.Tensor:
+        if self.vae.config.patch_size is not None:
+            from diffusers.models.autoencoders.autoencoder_kl_wan import patchify
+
+            return patchify(videos, patch_size=self.vae.config.patch_size)
+        return videos
+
+    @staticmethod
+    def _vae_clone_feat_map(feat_map: list[torch.Tensor | None]) -> list[torch.Tensor | None]:
+        return [entry.clone() if isinstance(entry, torch.Tensor) else entry for entry in feat_map]
+
+    def _vae_init_enc_feat_map(self) -> list[torch.Tensor | None]:
+        self.vae.clear_cache()
+        return self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_encode_encoder_chunk(
+        self,
+        chunk: torch.Tensor,
+        feat_map: list[torch.Tensor | None],
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """Run one Wan encoder chunk while mutating ``feat_map`` in place.
+
+        Must stay eager (not ``torch.compile``): Wan causal ``feat_cache`` updates
+        conflict with ``reduce-overhead`` CUDAGraph capture.
+        """
+        self.vae._enc_feat_map = feat_map
+        self.vae._enc_conv_idx = [0]
+        out = self.vae.encoder(
+            chunk,
+            feat_cache=self.vae._enc_feat_map,
+            feat_idx=self.vae._enc_conv_idx,
+        )
+        return out, self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_quantize_encoder_out(self, encoder_out: torch.Tensor) -> torch.Tensor:
+        enc = self.vae.quant_conv(encoder_out)
+        mu, _ = enc.chunk(2, dim=1)
+        mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
+        inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype)
+        return (mu - mean) * inv_std
+
+    def _vae_stream_seed(self, state: DreamZeroState, first_frame: torch.Tensor) -> None:
+        """Seed incremental VAE encode with the first observation frame."""
+        state.reset_vae_encoder_stream()
+        feat_map = self._vae_init_enc_feat_map()
+        chunk = self._vae_patchify(first_frame.to(dtype=self.vae.dtype))
+        self._cudagraph_mark_step_begin()
+        encoder_out, feat_map = self._vae_encode_encoder_chunk(chunk[:, :, :1], feat_map)
+        state.vae_enc_feat_map = feat_map
+        state.vae_encoder_out = encoder_out
+        state.vae_pending_body_frames = None
+        state.vae_stream_initialized = True
+
+    def _vae_stream_append_frame(self, state: DreamZeroState, new_frame: torch.Tensor) -> None:
+        """Append one pixel frame and encode a 4-frame body chunk when ready."""
+        if not state.vae_stream_initialized or state.vae_enc_feat_map is None or state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream is not initialized.")
+
+        frame = new_frame.to(dtype=self.vae.dtype)
+        if state.vae_pending_body_frames is None:
+            state.vae_pending_body_frames = frame
+        else:
+            state.vae_pending_body_frames = torch.cat([state.vae_pending_body_frames, frame], dim=2)
+
+        while state.vae_pending_body_frames is not None and state.vae_pending_body_frames.shape[2] >= 4:
+            body = state.vae_pending_body_frames[:, :, :4]
+            if state.vae_pending_body_frames.shape[2] > 4:
+                state.vae_pending_body_frames = state.vae_pending_body_frames[:, :, 4:]
+            else:
+                state.vae_pending_body_frames = None
+            chunk = self._vae_patchify(body)
+            encoder_chunk, feat_map = self._vae_encode_encoder_chunk(chunk, state.vae_enc_feat_map)
+            state.vae_encoder_out = torch.cat([state.vae_encoder_out, encoder_chunk], dim=2)
+            state.vae_enc_feat_map = feat_map
+
+    def _vae_stream_get_observation_latents(
+        self,
+        state: DreamZeroState,
+        num_latent_frames: int,
+        *,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream has no accumulated encoder output.")
+        latents = self._vae_quantize_encoder_out(state.vae_encoder_out).to(dtype=dtype)
+        if latents.shape[2] >= num_latent_frames:
+            return latents[:, :, -num_latent_frames:]
+        pad_count = num_latent_frames - latents.shape[2]
+        pad = latents[:, :, -1:].expand(-1, -1, pad_count, -1, -1)
+        return torch.cat([pad, latents], dim=2)
+
+    def _preprocess_vae_observation_window(self, videos: torch.Tensor) -> torch.Tensor:
+        _, _, num_frames_raw, _, _ = videos.shape
+        if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
+            return videos
+        if num_frames_raw // 4 != self.num_frame_per_block:
+            repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
+            videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
+            first_frame = videos[:, :, 0:1]
+            return torch.cat([first_frame, videos], dim=2)
+        first_frame = videos[:, :, 0:1]
+        return torch.cat([first_frame, videos], dim=2)
+
+    def _encode_observation_latents(
+        self,
+        state: DreamZeroState,
+        videos: torch.Tensor,
+        *,
+        latent_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Encode current robot observation into normalized VAE latents."""
+        if state.vae_stream_initialized:
+            self._cudagraph_mark_step_begin()
+            self._vae_stream_append_frame(state, videos[:, :, -1:])
+            return self._vae_stream_get_observation_latents(
+                state,
+                self.num_frame_per_block,
+                dtype=latent_dtype,
+            )
+
+        videos = self._preprocess_vae_observation_window(videos)
+        return self._encode_vae_latents(videos).to(dtype=latent_dtype)
+
     def _encode_vae_latents(self, videos: torch.Tensor) -> torch.Tensor:
         """Encode videos into normalized VAE latents."""
         input_dtype = videos.dtype
+        self._cudagraph_mark_step_begin()
         hidden = self.vae._encode(videos.to(dtype=self.vae.dtype))
         mu, _ = hidden.chunk(2, dim=1)
         mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
@@ -1105,6 +1252,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
                 self.num_frames,
                 height,
                 width,
+                state=state,
             )
             state.clip_feas = clip_feas.to(dtype=image.dtype)
             state.ys = ys.to(dtype=image.dtype)
@@ -1118,22 +1266,9 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
                 self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True)
 
         if state.current_start_frame != 0:
-            # Subsequent calls: encode current observation via VAE
-            if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
-                pass
-            elif num_frames_raw // 4 != self.num_frame_per_block:
-                repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
-                videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
-            else:
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
-
             latent_dtype = videos.dtype
             with torch.no_grad():
-                image = self._encode_vae_latents(videos)
-            image = image.to(dtype=latent_dtype)
+                image = self._encode_observation_latents(state, videos, latent_dtype=latent_dtype)
 
         batch_size = image.shape[0]
         generator = torch.Generator(device=device).manual_seed(self.seed)
diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
index 2ddb9d0f31f..a6299efab6a 100644
--- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
@@ -119,9 +119,17 @@ def reset(self, *, clear_video_latents: bool = True) -> None:
 
         self.clip_feas = None
         self.ys = None
+        self.reset_vae_encoder_stream()
         if clear_video_latents:
             self.language = None
 
+    def reset_vae_encoder_stream(self) -> None:
+        """Clear incremental Wan VAE encoder state used across AR steps."""
+        self.vae_stream_initialized = False
+        self.vae_enc_feat_map: list[torch.Tensor | None] | None = None
+        self.vae_encoder_out: torch.Tensor | None = None
+        self.vae_pending_body_frames: torch.Tensor | None = None
+
     def reset_inference_state(self) -> None:
         """Reset KV/frame state after local attention rolls without dropping video latents."""
         self.reset(clear_video_latents=False)

From 7dfb21e16a5772ca93c41b4857e4b9344aa631e5 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Thu, 2 Jul 2026 08:21:46 +0000
Subject: [PATCH 40/51] [AR-Diffusion][DreamZero] Cache session text embeds;
 keep VAE stream across window resets

Follow-ups to the PR #4750 cherry-pick, from a DZ_PHASE_TIMING breakdown
(text encode ~20 ms/chunk, VAE obs encode 0.9 ms, KV prefill ~100 ms,
denoise ~430-480 ms):

- Cache prompt embeds per session in DreamZeroState (cleared on "session"
  reset with state.language) and the constant negative-prompt embeds on
  the pipeline. UMT5 encode is deterministic -> numerically identical,
  saves ~20 ms on every chunk.
- Keep the incremental VAE encoder stream alive across "inference"
  (attention-window) resets: the Wan feat_cache history is independent of
  the DiT KV window. On a window restart with a live stream, append the
  restart observation instead of reseeding from scratch.
- Add env-gated DZ_PHASE_TIMING=1 per-forward stage timing (syncs CUDA at
  each mark; off for timed runs).

8-scene steady mean: 689 -> 632 ms (2.66x vs NVIDIA upstream FA2,
4.18x vs vllm-omni v0.22.0 aligned-config baseline).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 .../models/dreamzero/pipeline_dreamzero.py    | 80 +++++++++++++++----
 .../models/dreamzero/state_dreamzero.py       |  6 +-
 2 files changed, 69 insertions(+), 17 deletions(-)

diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 87b666d0212..b02ddfb29e7 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -299,6 +299,9 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         self.max_action_dim: int = ah_config["max_action_dim"]
 
         self.negative_prompt: str = model_config.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT)
+        # The negative prompt is a model constant: encode it once, lazily, and
+        # reuse across every forward/session (UMT5 encode is deterministic).
+        self._negative_prompt_embeds_cache: torch.Tensor | None = None
 
         # Embodiment name -> numeric ID mapping (model knowledge)
         self.embodiment_name_to_id: dict[str, int] = model_config.get(
@@ -678,9 +681,16 @@ def _encode_image(
             y = torch.concat([msk, y], dim=1)
 
             if state is not None:
-                # Seed the AR streaming encoder after the compiled full encode above;
-                # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph.
-                self._vae_stream_seed(state, image_input[:, :, :1])
+                if not state.vae_stream_initialized:
+                    # Seed the AR streaming encoder after the compiled full encode above;
+                    # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph.
+                    self._vae_stream_seed(state, image_input[:, :, :1])
+                else:
+                    # Window ("inference") restart with a live stream: keep the real
+                    # frame history in the Wan feat_cache and append the restart
+                    # observation instead of reseeding from scratch.
+                    self._cudagraph_mark_step_begin()
+                    self._vae_stream_append_frame(state, image_input[:, :, :1])
 
         return clip_context, y, new_image
 
@@ -1223,22 +1233,45 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
         videos = self._preprocess_video(videos)  # -> [B,C,T,H,W] bf16
         _, _, num_frames_raw, height, width = videos.shape
 
-        prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        # Optional phase timing (DZ_PHASE_TIMING=1): logs per-forward stage costs
+        # (text encode / obs VAE encode / KV prefill / denoise) at INFO. Each mark
+        # synchronizes CUDA, so leave it off for timed benchmark runs.
+        _pt = None
+        if os.environ.get("DZ_PHASE_TIMING"):
+            import time as _time
+
+            torch.cuda.synchronize()
+            _pt = {"time": _time, "t0": _time.perf_counter(), "marks": []}
+
+        def _pt_mark(name: str) -> None:
+            if _pt is not None:
+                torch.cuda.synchronize()
+                _pt["marks"].append((name, _pt["time"].perf_counter()))
+
+        # Prompt embeds are constant within a session (a prompt change triggers a
+        # "session" reset above, which clears this cache alongside state.language).
+        if state.prompt_embeds is None:
+            state.prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        prompt_embeds = state.prompt_embeds
         # Negative prompt for CFG uncond branch (model constant)
         negative_prompt_embeds = None
         if self.cfg_scale > 1.0:
-            neg_inputs = self.tokenizer(
-                self.negative_prompt,
-                max_length=512,
-                padding="max_length",
-                truncation=True,
-                return_tensors="pt",
-                add_special_tokens=True,
-            )
-            negative_prompt_embeds = self._encode_text(
-                neg_inputs["input_ids"].to(device),
-                neg_inputs["attention_mask"].to(device),
-            )
+            if self._negative_prompt_embeds_cache is None:
+                neg_inputs = self.tokenizer(
+                    self.negative_prompt,
+                    max_length=512,
+                    padding="max_length",
+                    truncation=True,
+                    return_tensors="pt",
+                    add_special_tokens=True,
+                )
+                self._negative_prompt_embeds_cache = self._encode_text(
+                    neg_inputs["input_ids"].to(device),
+                    neg_inputs["attention_mask"].to(device),
+                )
+            negative_prompt_embeds = self._negative_prompt_embeds_cache
+
+        _pt_mark("text_encode")
 
         # Extract first/last frame for CLIP + VAE encoding
         if num_frames_raw == 4 or num_frames_raw == 9:
@@ -1269,6 +1302,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             latent_dtype = videos.dtype
             with torch.no_grad():
                 image = self._encode_observation_latents(state, videos, latent_dtype=latent_dtype)
+        _pt_mark("obs_vae_encode")
 
         batch_size = image.shape[0]
         generator = torch.Generator(device=device).manual_seed(self.seed)
@@ -1309,6 +1343,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             do_true_cfg,
             state,
         )
+        _pt_mark("prefill_kv")
 
         sample_scheduler = copy.deepcopy(self.scheduler)
         sample_scheduler_action = copy.deepcopy(self.scheduler)
@@ -1350,6 +1385,19 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
             state_features=state_features,
             embodiment_id=embodiment_id,
         )
+        if _pt is not None:
+            _pt_mark("diffuse")
+            prev_t = _pt["t0"]
+            parts = []
+            for name, t in _pt["marks"]:
+                parts.append(f"{name}={1000 * (t - prev_t):.1f}ms")
+                prev_t = t
+            logger.info(
+                "DZ_PHASE_TIMING csf=%s total=%.1fms %s",
+                state.current_start_frame,
+                1000 * (prev_t - _pt["t0"]),
+                " ".join(parts),
+            )
 
         if state.current_start_frame == 1:
             video_out = torch.cat([image, video_out], dim=1)
diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
index a6299efab6a..4c9209808a4 100644
--- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
@@ -119,9 +119,13 @@ def reset(self, *, clear_video_latents: bool = True) -> None:
 
         self.clip_feas = None
         self.ys = None
-        self.reset_vae_encoder_stream()
         if clear_video_latents:
+            # Session reset: drop the prompt-embed cache and VAE encoder stream.
             self.language = None
+            self.prompt_embeds = None
+            self.reset_vae_encoder_stream()
+        # Window ("inference") resets keep both: the prompt is unchanged, and the
+        # Wan feat_cache history is independent of the DiT attention window.
 
     def reset_vae_encoder_stream(self) -> None:
         """Clear incremental Wan VAE encoder state used across AR steps."""

From 2e713c20999f7a38c3de9ebce8fd8e18ca9a0bd9 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Thu, 2 Jul 2026 11:56:03 +0000
Subject: [PATCH 41/51] [AR-Diffusion][DreamZero] Keep text cross-attn KV
 across window restarts

Torch-profiler flame graphs show the window-restart forward spends
~192-215 ms in _kv_populate_cross re-projecting 40 layers x 2 CFG
branches of cross-attn K/V. The text half depends only on the session
prompt (unchanged across window rolls) and lives in static pool buffers
untouched by the block-pool reset -- re-projecting it is pure waste.

- Split _cross_populated into text/img halves; window ("inference")
  resets keep the text flags (ARDiffusionKVState.reset(keep_cross_text)),
  session resets clear both.
- _kv_populate_cross projects only the missing halves; write_cross_kv
  accepts k/v=None to write the image half alone.
- Numerically identical: the kept text K/V are the same values the
  restart would have re-projected (same prompt embeds, same weights).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/ar_diffusion/test_pipeline_bridge.py    |  6 +--
 .../models/dreamzero/pipeline_dreamzero.py    | 48 ++++++++++++-------
 .../ar_diffusion/kv_cache/manager.py          | 17 ++++---
 .../ar_diffusion/kv_cache/state.py            | 41 ++++++++++++----
 4 files changed, 76 insertions(+), 36 deletions(-)

diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py
index fb38b4f8304..9e49e6c30ba 100644
--- a/tests/ar_diffusion/test_pipeline_bridge.py
+++ b/tests/ar_diffusion/test_pipeline_bridge.py
@@ -135,8 +135,8 @@ def test_get_cross_kv_caches_raises_before_populate():
     for neg in (False, True):
         with pytest.raises(RuntimeError, match="cross-attn read before"):
             st.get_cross_kv_caches(neg)
-    # Guards the full AND: _cross_populated alone is not enough without a pool.
-    st._cross_populated[False] = True
+    # Guards the full AND: the populated flag alone is not enough without a pool.
+    st._cross_text_populated[False] = True
     with pytest.raises(RuntimeError, match="cross-attn read before"):
         st.get_cross_kv_caches(False)
 
@@ -152,7 +152,7 @@ def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
         v = torch.randn(1, L, N_HEADS, HEAD_DIM)
         kv.write_cross_kv(i, False, k, v)
         written.append((k, v))
-    st._cross_populated[False] = True
+    st._cross_text_populated[False] = True
 
     out = st.get_cross_kv_caches(False)
     assert len(out) == 2
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index b02ddfb29e7..3d840dacfbe 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -124,36 +124,46 @@ def _kv_get_cross(self, state, is_negative):
     def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None:
         """Eagerly project cross-attn K/V for all layers into the AR-Diffusion pool.
 
-        Caches the session-invariant cross-attn projections once (re-run after a
-        window-boundary reset, when ``_cross_populated`` clears): the text ``k``/``v``
-        from ``text_embedding(context)``, and for I2V the image-token ``k_img``/
-        ``v_img`` from ``img_emb(clip_feature)`` (the 257 image tokens the forward
-        splits off, cached model-side by #4154). Must run after the image is encoded
-        so ``clip_feature`` is available.
+        Caches the session-invariant cross-attn projections once, per half: the text
+        ``k``/``v`` from ``text_embedding(context)`` survive window-boundary resets
+        (prompt unchanged within a session — only session resets clear them), while
+        the I2V image-token ``k_img``/``v_img`` from ``img_emb(clip_feature)`` (the
+        257 image tokens the forward splits off, cached model-side by #4154) are
+        re-projected on every window restart from the fresh CLIP features. Must run
+        after the image is encoded so ``clip_feature`` is available.
         """
         s = self._ar_diffusion_kv_state
-        if s._cross_populated.get(is_negative, False):
+        need_text = not s._cross_text_populated.get(is_negative, False)
+        need_img = (
+            clip_feature is not None
+            and getattr(self.transformer, "model_type", "t2v") == "i2v"
+            and not s._cross_img_populated.get(is_negative, False)
+        )
+        if not need_text and not need_img:
             return
-        projected = self.transformer.text_embedding(context)
-        img_ctx = None
-        if clip_feature is not None and getattr(self.transformer, "model_type", "t2v") == "i2v":
-            img_ctx = self.transformer.img_emb(clip_feature)
+        projected = self.transformer.text_embedding(context) if need_text else None
+        img_ctx = self.transformer.img_emb(clip_feature) if need_img else None
         for i, block in enumerate(self.transformer.blocks):
             ca = block.cross_attn
             n, d = ca.tp_num_heads, ca.head_dim
-            k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
-            v = ca.v(projected).unflatten(2, (n, d))
+            k = v = None
+            if projected is not None:
+                k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
+                v = ca.v(projected).unflatten(2, (n, d))
             k_img = v_img = None
             if img_ctx is not None:
                 k_img = ca.norm_k_img(ca.k_img(img_ctx)).unflatten(2, (n, d))
                 v_img = ca.v_img(img_ctx).unflatten(2, (n, d))
             s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img)
-        s._cross_populated[is_negative] = True
+        if need_text:
+            s._cross_text_populated[is_negative] = True
+        if need_img:
+            s._cross_img_populated[is_negative] = True
         logger.info(
             "AR-Diffusion CROSS POPULATE [%s]: %d layers, text=%s img=%s",
             "neg" if is_negative else "pos",
             len(self.transformer.blocks),
-            tuple(context.shape),
+            "kept" if projected is None else tuple(context.shape),
             None if img_ctx is None else tuple(img_ctx.shape),
         )
 
@@ -163,9 +173,13 @@ def _kv_reset(self, state, *, clear_video_latents: bool = True):
         DreamZero resets at the attention-window boundary; the engine pool drops the
         same window so the next forward starts fresh. ``clear_video_latents=False``
         keeps the accumulated video latents for export.
+
+        ``clear_video_latents=False`` also marks a window ("inference") reset: the
+        prompt is unchanged, so the pool keeps the text cross-attn K/V and only the
+        image half repopulates on the restart forward.
         """
         state.reset(clear_video_latents=clear_video_latents)
-        self._ar_diffusion_kv_state.reset()
+        self._ar_diffusion_kv_state.reset(keep_cross_text=not clear_video_latents)
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.
@@ -1293,7 +1307,7 @@ def _pt_mark(name: str) -> None:
             # Eager cross-attn population (AR-Diffusion only): cache text + image-token K/V
             # into the pool now that the image is encoded (clip_feas available).
             # Runs on the first forward of a session and after each window-boundary
-            # reset (current_start_frame returns to 0); _cross_populated guards re-entry.
+            # reset (current_start_frame returns to 0); the populate guards (text/img halves) gate re-entry.
             self._kv_populate_cross(prompt_embeds, state.clip_feas, is_negative=False)
             if negative_prompt_embeds is not None:
                 self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True)
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
index 3c88b053053..bbf0af90f58 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
@@ -294,21 +294,24 @@ def write_cross_kv(
         self,
         layer_idx: int,
         is_negative: bool,
-        k: torch.Tensor,
-        v: torch.Tensor,
+        k: torch.Tensor | None,
+        v: torch.Tensor | None,
         k_img: torch.Tensor | None = None,
         v_img: torch.Tensor | None = None,
     ) -> None:
         """Write one layer's cross-attn K/V into the pool.
 
         ``k`` / ``v``: ``(B, text_len, tp_num_heads, head_dim)``; only batch-0
-        is copied (B=1 for inference). The ``is_negative`` flag selects the
-        correct CFG branch slot. ``k_img`` / ``v_img`` (I2V image tokens,
-        ``(B, 257, ...)``) are written when the image pool is allocated.
+        is copied (B=1 for inference). ``None`` skips the text half (window
+        restarts keep the still-valid text K/V and rewrite only the image half).
+        The ``is_negative`` flag selects the correct CFG branch slot. ``k_img`` /
+        ``v_img`` (I2V image tokens, ``(B, 257, ...)``) are written when the
+        image pool is allocated.
         """
         branch = 1 if is_negative else 0
-        self._cross_k[layer_idx][branch].copy_(k[0])
-        self._cross_v[layer_idx][branch].copy_(v[0])
+        if k is not None:
+            self._cross_k[layer_idx][branch].copy_(k[0])
+            self._cross_v[layer_idx][branch].copy_(v[0])
         if k_img is not None and self._cross_k_img:
             self._cross_k_img[layer_idx][branch].copy_(k_img[0])
             self._cross_v_img[layer_idx][branch].copy_(v_img[0])
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
index 19b48094190..161faedb152 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
@@ -64,9 +64,12 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
         # behavior), for A/B measurement of the memoization win. Default: enabled.
         self._memo_enabled = os.environ.get("AR_DIFFUSION_KV_NO_MEMO") != "1"
         # Cross-attn KV is populated once after text encoding (eager), then
-        # read every denoising step.  Reset clears these flags so the next
-        # forward repopulates from the new text encoding.
-        self._cross_populated: dict[bool, bool] = {False: False, True: False}
+        # read every denoising step. Tracked per half: the text K/V depend only
+        # on the session prompt (survive window-boundary resets), while the I2V
+        # image-token K/V come from the current observation's CLIP features and
+        # must be re-projected on every window restart.
+        self._cross_text_populated: dict[bool, bool] = {False: False, True: False}
+        self._cross_img_populated: dict[bool, bool] = {False: False, True: False}
 
     def _adapter(self, is_negative: bool):
         return self.neg if is_negative else self.pos
@@ -101,10 +104,17 @@ def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
         were, the model would lazily project and write cross KV itself, violating
         engine ownership. Fail loud rather than fall back to the model.
         """
-        if not (self._cross_populated[is_negative] and self.kv_cache.cross_attn_length > 0):
+        img_ok = not self.kv_cache._cross_k_img or self._cross_img_populated[is_negative]
+        if not (self._cross_text_populated[is_negative] and img_ok and self.kv_cache.cross_attn_length > 0):
             raise RuntimeError(
-                "AR-Diffusion cross-attn read before _kv_populate_cross (neg=%s, cross_attn_length=%d) "
-                "— the engine must own all cross KV" % (is_negative, self.kv_cache.cross_attn_length)
+                "AR-Diffusion cross-attn read before _kv_populate_cross (neg=%s, cross_attn_length=%d, "
+                "text=%s img=%s) — the engine must own all cross KV"
+                % (
+                    is_negative,
+                    self.kv_cache.cross_attn_length,
+                    self._cross_text_populated[is_negative],
+                    self._cross_img_populated[is_negative],
+                )
             )
         return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
 
@@ -179,7 +189,7 @@ def close(self) -> None:
         self.kv_cache.end_request(self.pos)
         self.kv_cache.end_request(self.neg)
 
-    def reset(self) -> None:
+    def reset(self, *, keep_cross_text: bool = False) -> None:
         """Drop this session's KV — mirrors the model-local ``state.reset()``.
 
         DreamZero resets at the attention-window boundary (``should_reset``); the
@@ -188,6 +198,12 @@ def reset(self) -> None:
         object (and thus ``pipeline._ar_diffusion_kv_state``) is preserved across the reset
         so the runner's session mapping stays valid — only the pool-backed state
         is recycled.
+
+        Args:
+            keep_cross_text: On window ("inference") resets the session prompt is
+                unchanged, so the text cross-attn K/V in the static pool buffers are
+                still valid — keep them and only force the image half to repopulate.
+                Session resets (prompt change / explicit reset) pass ``False``.
         """
         pos_id, neg_id = self.pos.request_id, self.neg.request_id
         self.close()  # free resident blocks for both branches
@@ -196,5 +212,12 @@ def reset(self) -> None:
         self._committed = {False: 0, True: 0}
         self._pending = {False: [], True: []}
         self._gather_cache = {False: None, True: None}
-        self._cross_populated = {False: False, True: False}
-        _log.info("AR-Diffusion RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)
+        if not keep_cross_text:
+            self._cross_text_populated = {False: False, True: False}
+        self._cross_img_populated = {False: False, True: False}
+        _log.info(
+            "AR-Diffusion RESET [%s/%s] session KV cleared (window boundary; cross text %s)",
+            pos_id,
+            neg_id,
+            "kept" if keep_cross_text else "cleared",
+        )

From e9d628f2e2de103e0111117d4257a178d7323c77 Mon Sep 17 00:00:00 2001
From: Acerak01-fy <wfy2003324@163.com>
Date: Thu, 2 Jul 2026 10:35:56 +0000
Subject: [PATCH 42/51] Add AR-Diffusion paged self-attention backend
 Signed-off-by: Acerak01-fy <wfy2003324@163.com> Signed-off-by: Taichang Zhou
 <tzhouam@connect.ust.hk>

---
 tests/ar_diffusion/test_gather.py             | 110 -------
 tests/ar_diffusion/test_kv_cache.py           |  34 +-
 tests/ar_diffusion/test_paged_attention.py    | 279 ++++++++++++++++
 tests/ar_diffusion/test_pipeline_bridge.py    | 146 ++++----
 .../models/dreamzero/causal_wan_model.py      | 123 +++++--
 .../models/dreamzero/pipeline_dreamzero.py    |  56 ++--
 .../ar_diffusion/kv_cache/__init__.py         |  14 +-
 .../ar_diffusion/kv_cache/manager.py          | 141 ++++----
 .../ar_diffusion/kv_cache/paged.py            |  98 ++----
 .../ar_diffusion/kv_cache/paged_attention.py  | 311 ++++++++++++++++++
 .../ar_diffusion/kv_cache/state.py            | 217 +++++-------
 11 files changed, 1008 insertions(+), 521 deletions(-)
 delete mode 100644 tests/ar_diffusion/test_gather.py
 create mode 100644 tests/ar_diffusion/test_paged_attention.py
 create mode 100644 vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py

diff --git a/tests/ar_diffusion/test_gather.py b/tests/ar_diffusion/test_gather.py
deleted file mode 100644
index 41ebc33b9a4..00000000000
--- a/tests/ar_diffusion/test_gather.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Tests for pool write + window gather (Step 4)."""
-
-import torch
-
-from vllm_omni.experimental.ar_diffusion.kv_cache import (
-    ARDiffusionKVCache,
-    ARDiffusionKVConfig,
-    allocate_kv_pool,
-    pool_gather_window,
-    pool_write_chunk,
-)
-
-BLOCK = 16
-N_HEADS = 4
-HEAD_DIM = 64
-
-
-def test_pool_write_gather_roundtrip():
-    """Write a single chunk to the pool then gather — K/V comes back identical."""
-    num_blocks = 4
-    kp, vp = allocate_kv_pool(num_blocks, BLOCK, 1, N_HEADS, HEAD_DIM, torch.float32, "cpu")
-    new_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
-    new_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
-    slot_mapping = torch.arange(num_blocks * BLOCK)  # identity mapping for the test
-    slot_mapping = slot_mapping[:BLOCK]
-
-    pool_write_chunk(kp[0], vp[0], new_k, new_v, slot_mapping)
-    window = pool_gather_window(kp[0], vp[0], window_block_ids=[0], block_size=BLOCK, max_attention_size=BLOCK)
-
-    assert window.shape == (2, 1, BLOCK, N_HEADS, HEAD_DIM)
-    assert torch.allclose(window[0, 0], new_k[0])
-    assert torch.allclose(window[1, 0], new_v[0])
-
-
-def test_pool_write_gather_multiple_blocks():
-    """Write 2 blocks (distinct positions) and gather both from the window."""
-    num_blocks = 8
-    kp, vp = allocate_kv_pool(num_blocks, BLOCK, 1, N_HEADS, HEAD_DIM, torch.float32, "cpu")
-
-    # Write block 0
-    k0 = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
-    v0 = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
-    s0 = torch.arange(BLOCK)  # block 0 starts at slot 0
-    pool_write_chunk(kp[0], vp[0], k0, v0, s0)
-
-    # Write block 2 (holes OK, block id doesn't need to be contiguous in slot space)
-    k2 = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
-    v2 = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
-    s2 = 2 * BLOCK + torch.arange(BLOCK)  # block 2
-    pool_write_chunk(kp[0], vp[0], k2, v2, s2)
-
-    # Gather blocks 0 and 2 — they appear in table order.
-    window = pool_gather_window(kp[0], vp[0], window_block_ids=[0, 2], block_size=BLOCK, max_attention_size=2 * BLOCK)
-    assert window.shape == (2, 1, 2 * BLOCK, N_HEADS, HEAD_DIM)
-    assert torch.allclose(window[0, 0, :BLOCK], k0[0])
-    assert torch.allclose(window[1, 0, :BLOCK], v0[0])
-    assert torch.allclose(window[0, 0, BLOCK:], k2[0])
-    assert torch.allclose(window[1, 0, BLOCK:], v2[0])
-
-
-def test_pool_write_gather_window_trim():
-    """Gather with max_attention_size smaller than the resident window trims to tail."""
-    num_blocks = 4
-    kp, vp = allocate_kv_pool(num_blocks, BLOCK, 1, N_HEADS, HEAD_DIM, torch.float32, "cpu")
-    k = torch.randn(1, 2 * BLOCK, N_HEADS, HEAD_DIM)
-    v = torch.randn(1, 2 * BLOCK, N_HEADS, HEAD_DIM)
-    slots = torch.cat([torch.arange(0 * BLOCK, 1 * BLOCK), torch.arange(1 * BLOCK, 2 * BLOCK)])
-    pool_write_chunk(kp[0], vp[0], k[:, :BLOCK], v[:, :BLOCK], slots[:BLOCK])  # block 0
-    pool_write_chunk(kp[0], vp[0], k[:, BLOCK:], v[:, BLOCK:], slots[BLOCK:])  # block 1
-    # max_attention = BLOCK — only last block.
-    window = pool_gather_window(kp[0], vp[0], window_block_ids=[0, 1], block_size=BLOCK, max_attention_size=BLOCK)
-    assert window.shape == (2, 1, BLOCK, N_HEADS, HEAD_DIM)
-    assert torch.allclose(window[0, 0], k[0, BLOCK:])
-    assert torch.allclose(window[1, 0], v[0, BLOCK:])
-
-
-def test_bde_kv_cache_write_gather_integration():
-    """ARDiffusionKVCache.write_chunk_kv → gather_window roundtrip through a real manager."""
-    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2)
-    kv = ARDiffusionKVCache(
-        cfg,
-        num_layers=1,
-        num_kv_heads=N_HEADS,
-        head_size=HEAD_DIM,
-        dtype=torch.float32,
-        block_size=BLOCK,
-        max_model_len=512,
-        available_bytes=1 << 24,
-        device=torch.device("cpu"),
-    )
-    adapter = kv.begin_request("r")
-    kv.allocate_chunk(adapter)
-    # Use the orchestrator's own slot mapping (which resolves the actual physical
-    # block id from the request's block table).
-    write_slots = kv.chunk_write_slots(adapter)
-
-    new_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
-    new_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM)
-    kv.write_chunk_kv(0, new_k, new_v, adapter)
-    kv.commit_chunk(adapter)
-
-    # Read back the chunk's K/V from the pool at exactly the slots it was written to.
-    assert torch.allclose(kv._k_pools[0][write_slots], new_k[0])
-    assert torch.allclose(kv._v_pools[0][write_slots], new_v[0])
-
-    # Allocate a second chunk so there's a past-window + current chunk.
-    kv.allocate_chunk(adapter)
-    window = kv.gather_window(0, adapter)
-    assert window.shape == (2, 1, kv.spec.sliding_window, N_HEADS, HEAD_DIM)
diff --git a/tests/ar_diffusion/test_kv_cache.py b/tests/ar_diffusion/test_kv_cache.py
index c03f45147fd..7698081014a 100644
--- a/tests/ar_diffusion/test_kv_cache.py
+++ b/tests/ar_diffusion/test_kv_cache.py
@@ -9,15 +9,21 @@
 import pytest
 import torch
 from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
-from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.request import RequestStatus
 
+try:
+    from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+except ModuleNotFoundError:
+    KVCacheSpecRegistry = None
+    from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
 from vllm_omni.experimental.ar_diffusion.kv_cache import (
     ARDiffusionKVCache,
     ARDiffusionKVConfig,
     ARDiffusionRequestAdapter,
     ChunkWindowManager,
     ChunkWindowSpec,
+    allocate_kv_pool_with_views,
     build_kv_manager,
     compute_num_blocks,
 )
@@ -46,7 +52,11 @@ def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boun
 def test_spec_registration_resolves_to_chunk_window_manager():
     # Without explicit registration the MRO walk would fall back to the parent
     # SlidingWindowManager; assert the subclass manager wins.
-    assert KVCacheSpecRegistry.get_manager_class(make_spec()) is ChunkWindowManager
+    spec = make_spec()
+    if KVCacheSpecRegistry is not None:
+        assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+    else:
+        assert spec_manager_map[type(spec)] is ChunkWindowManager
 
 
 def test_spec_kind_is_sliding_window():
@@ -138,6 +148,26 @@ def test_compute_num_blocks():
     assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
 
 
+def test_paged_pool_layout_exposes_flat_slot_views():
+    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
+        num_blocks=4,
+        block_size=BLOCK,
+        num_layers=1,
+        num_kv_heads=4,
+        head_dim=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
+    assert k_pools[0].shape == (4 * BLOCK, 4, 64)
+    assert v_pools[0].shape == (4 * BLOCK, 4, 64)
+
+    k_pools[0][BLOCK + 3].fill_(7)
+    v_pools[0][2 * BLOCK + 5].fill_(11)
+    assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
+    assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
+
+
 def test_build_manager_allocate_free_roundtrip():
     """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.
 
diff --git a/tests/ar_diffusion/test_paged_attention.py b/tests/ar_diffusion/test_paged_attention.py
new file mode 100644
index 00000000000..975732877de
--- /dev/null
+++ b/tests/ar_diffusion/test_paged_attention.py
@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for AR-Diffusion paged self-attention contexts."""
+
+from __future__ import annotations
+
+import subprocess
+from importlib.util import find_spec
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+    ar_diffusion_paged_attention,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=dtype,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 26,
+        device=device,
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5)
+    probs = torch.softmax(scores, dim=-1).to(value.dtype)
+    return torch.einsum("bhqk,bkhd->bqhd", probs, value)
+
+
+def _cuda_flash_attn_usable() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        spec = find_spec("vllm.vllm_flash_attn")
+        if spec is None or spec.origin is None:
+            return True
+        fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so"
+        linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5)
+    except Exception:
+        return True
+    if "libcudart.so.13" not in linked:
+        return True
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            text=True,
+            timeout=5,
+        )
+        driver_major = int(out.splitlines()[0].split(".")[0])
+    except Exception:
+        return True
+    return driver_major >= 580
+
+
+def _commit_video_span(
+    kv: ARDiffusionKVCache,
+    st: ARDiffusionKVState,
+    *,
+    is_negative: bool,
+    n_chunks: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
+    st.commit_paged_context(is_negative)
+    return k, v
+
+
+def test_paged_context_allocates_lazily_and_commits_after_forward():
+    _, st = make_state()
+
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+    assert st.pos.completed_chunks == 0
+    assert ctx.current_video_slot_mapping is None
+
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 1
+
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 1
+    assert st._committed[False] == BLOCK
+
+
+def test_scratch_video_and_action_blocks_do_not_commit():
+    kv, st = make_state()
+
+    ctx = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    ctx.ensure_action_slots(3, torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 2)
+    assert ctx.action_scratch_block_ids == kv.scratch_block_ids(False, 2, 1)
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+
+
+def test_pipeline_kv_get_paged_path_has_no_gather_backend():
+    kv, st = make_state()
+    assert not hasattr(kv, "gather_window_all_layers")
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._ar_diffusion_kv_state = st
+    contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False)
+
+    assert len(contexts) == 1
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+
+
+@pytest.mark.parametrize("history_chunks", [0, 1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k_parts: list[torch.Tensor] = []
+    history_v_parts: list[torch.Tensor] = []
+    if history_chunks:
+        k, v = _commit_video_span(
+            kv,
+            st,
+            is_negative=False,
+            n_chunks=history_chunks,
+            dtype=dtype,
+            device=device,
+        )
+        history_k_parts.append(k)
+        history_v_parts.append(v)
+
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len,
+        query_len=query.shape[1],
+        device=device,
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    if history_k_parts:
+        history_k = torch.cat(history_k_parts, dim=1)
+        history_v = torch.cat(history_v_parts, dim=1)
+    else:
+        history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+    before = st.pos.completed_chunks
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == before + (1 if commit_current else 0)
+
+
+@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
+@pytest.mark.parametrize("history_chunks", [1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
+    pytest.importorskip("vllm.vllm_flash_attn")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k, history_v = _commit_video_span(
+        kv,
+        st,
+        is_negative=False,
+        n_chunks=history_chunks,
+        dtype=dtype,
+        device=device,
+    )
+
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len,
+        query_len=query.shape[1],
+        device=device,
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py
index 9e49e6c30ba..1ad90d4ba4e 100644
--- a/tests/ar_diffusion/test_pipeline_bridge.py
+++ b/tests/ar_diffusion/test_pipeline_bridge.py
@@ -1,10 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Tests for the ARDiffusionKVState pipeline bridge (Step 5)."""
+"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
 
 import pytest
 import torch
 
-from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+)
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
 BLOCK = 16
@@ -31,106 +35,92 @@ def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
     return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
 
 
-def _window(n_chunks):
-    """A model-style window tensor (2, 1, n_chunks*BLOCK, heads, head_dim)."""
-    return torch.randn(2, 1, n_chunks * BLOCK, N_HEADS, HEAD_DIM)
+def _prepare_and_commit(st: ARDiffusionKVState, is_negative: bool, n_chunks: int) -> None:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    st.commit_paged_context(is_negative)
 
 
-# Approach 2: ARDiffusionKVState writes the new tokens into paged pool blocks and gathers
-# the resident window back; output mirrors the model-local sliding window.
+def test_get_returns_paged_layer_contexts_when_nothing_committed():
+    _, st = make_state(num_layers=2)
+    for neg in (False, True):
+        contexts = st.get_kv_caches(neg, seq_len=BLOCK, commit_current=False)
+        assert len(contexts) == 2
+        assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts)
+        assert contexts[0].history_block_ids == []
 
 
-def test_get_returns_empty_window_when_nothing_committed():
-    """The engine owns the read end-to-end: with nothing committed the adapter has
-    no resident blocks, so gather returns zero-length windows — the concat-identity
-    the model's attention expects on the first step — with no model-local fallback."""
-    _, st = make_state()
-    for neg in (False, True):
-        windows = st.get_kv_caches(neg)
-        assert len(windows) == 1
-        assert windows[0].shape == (2, 1, 0, N_HEADS, HEAD_DIM)
+def test_managed_paged_context_commits_only_after_forward():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
 
+    assert st.pos.completed_chunks == 0
+    assert kv.window_block_ids(st.pos) == []
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 2
 
-def test_paged_write_then_gather_roundtrips():
-    _, st = make_state(num_layers=1)
-    win = _window(2)  # this forward appends 2 frame-chunks
-    st.update_kv_cache(0, win, False, seq_len=2 * BLOCK)
-    got = st.get_kv_caches(False)[0]  # (2, 1, W, heads, head_dim)
-    assert got.shape == (2, 1, 2 * BLOCK, N_HEADS, HEAD_DIM)
-    # The gathered window equals the tokens we wrote (bit-exact).
-    assert torch.allclose(got[0, 0], win[0, 0], atol=0)
-    assert torch.allclose(got[1, 0], win[1, 0], atol=0)
+    st.commit_paged_context(False)
 
+    assert st.pos.completed_chunks == 2
+    assert st._committed[False] == 2 * BLOCK
+    assert len(kv.window_block_ids(st.pos)) == 2
 
-def test_gather_is_memoized_within_a_forward():
-    """The resident window is frozen across a forward's denoise steps, so repeated
-    get_kv_caches must reuse one gather; a write-back (layer 0) invalidates it."""
-    _, st = make_state(num_layers=2)
-    st.update_kv_cache(0, _window(2), False, seq_len=2 * BLOCK)  # commit one chunk
-    first = st.get_kv_caches(False)
-    again = st.get_kv_caches(False)
-    assert again is first  # cache HIT: identical object, no re-gather
-    # A new write-back grows the window -> next get must re-gather a fresh object.
-    st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
-    after = st.get_kv_caches(False)
-    assert after is not first
-    # Negative branch is cached independently.
-    st.update_kv_cache(0, _window(1), True, seq_len=BLOCK)
-    neg1 = st.get_kv_caches(True)
-    assert st.get_kv_caches(True) is neg1
-    assert neg1 is not after
-
-
-def test_reset_clears_gather_cache():
-    _, st = make_state(num_layers=1)
-    st.update_kv_cache(0, _window(2), False, seq_len=2 * BLOCK)
-    st.get_kv_caches(False)
-    assert st._gather_cache[False] is not None
-    st.reset()
-    assert st._gather_cache == {False: None, True: None}
+
+def test_scratch_paged_context_does_not_commit():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
+    ctx = contexts[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 1)
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+    assert kv.window_block_ids(st.pos) == []
 
 
 def test_eviction_bounds_resident_window():
     kv, st = make_state(num_layers=1, window_chunks=3)
-    st.update_kv_cache(0, _window(3), False, seq_len=3 * BLOCK)  # resident 3/3
-    st.update_kv_cache(0, _window(5), False, seq_len=2 * BLOCK)  # +2 -> evict to 3
-    got = st.get_kv_caches(False)[0]
-    assert got.shape[2] <= 3 * BLOCK  # window bounded
+    _prepare_and_commit(st, False, 3)
+    assert len(kv.window_block_ids(st.pos)) == 3
+
+    _prepare_and_commit(st, False, 2)
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    visible_blocks, video_len = ctx.video_block_table(torch.device("cpu"))
+    assert len(visible_blocks) == 3
+    assert video_len == 3 * BLOCK
 
 
 def test_branches_are_independent():
-    _, st = make_state()
-    st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
-    # Negative branch never written -> empty (zero-length) window.
-    assert st.get_kv_caches(True)[0].shape[2] == 0
-    assert st.get_kv_caches(False)[0].shape[2] == BLOCK
+    kv, st = make_state()
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+    assert kv.window_block_ids(st.neg) == []
 
 
 def test_reset_clears_session_window():
-    """reset() drops the resident window so the next forward starts fresh —
-    mirrors DreamZero's window-boundary state.reset()."""
     kv, st = make_state(num_layers=1, window_chunks=4)
-    st.update_kv_cache(0, _window(2), False, seq_len=2 * BLOCK)
-    st.update_kv_cache(0, _window(1), True, seq_len=BLOCK)
+    _prepare_and_commit(st, False, 2)
+    _prepare_and_commit(st, True, 1)
     assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK
     free_before = kv.manager.block_pool.get_num_free_blocks()
 
     st.reset()
 
-    # Both branches gather empty windows again, and pool blocks are freed.
     assert st._committed == {False: 0, True: 0}
-    assert st._pending == {False: [], True: []}
-    assert st.get_kv_caches(False)[0].shape[2] == 0
-    assert st.get_kv_caches(True)[0].shape[2] == 0
+    assert st._paged_pending == {False: None, True: None}
+    assert kv.window_block_ids(st.pos) == []
+    assert kv.window_block_ids(st.neg) == []
     assert kv.manager.block_pool.get_num_free_blocks() > free_before
-    # Adapters are live again under the same ids — a post-reset write works.
-    st.update_kv_cache(0, _window(1), False, seq_len=BLOCK)
-    assert st.get_kv_caches(False)[0].shape[2] == BLOCK
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
 
 
 def test_get_cross_kv_caches_raises_before_populate():
-    """Engine ownership guard: a cross-attn read before _kv_populate_cross must
-    fail loud rather than fall back to the model (which would self-project KV)."""
+    """Engine ownership guard: a cross-attn read before _kv_populate_cross must fail loud."""
     _, st = make_state()  # cross_attn_length == 0, nothing populated
     for neg in (False, True):
         with pytest.raises(RuntimeError, match="cross-attn read before"):
@@ -142,8 +132,6 @@ def test_get_cross_kv_caches_raises_before_populate():
 
 
 def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
-    """Once the cross pool is filled and marked populated, the read returns the
-    pool-backed dicts (is_init=True) bit-exactly — no model fallback involved."""
     L = 8
     kv, st = make_state(num_layers=2, cross_attn_length=L)
     written = []
@@ -164,8 +152,6 @@ def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
 
 
 def test_kv_create_is_noop_engine_owns_allocation():
-    """The AR-Diffusion engine owns all KV allocation, so _kv_create never touches
-    the model-local state (DreamZero's model-local KV path was removed)."""
     from unittest.mock import MagicMock
 
     from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
@@ -177,9 +163,7 @@ def test_kv_create_is_noop_engine_owns_allocation():
     state.create_kv_caches.assert_not_called()
 
 
-def test_update_kv_cache_requires_frame_aligned_seqlen():
-    """Engine guards the frame-alignment invariant (review: zwhzzz0821) — a
-    non-multiple seq_len would otherwise silently drop the leading remainder."""
+def test_get_kv_cache_requires_frame_aligned_seqlen():
     _, st = make_state()
     with pytest.raises(AssertionError, match="frame-aligned"):
-        st.update_kv_cache(0, _window(1), False, seq_len=BLOCK + 1)
+        st.get_kv_caches(False, seq_len=BLOCK + 1, commit_current=True)
diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
index f058ba2bd3a..1e399935444 100644
--- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
+++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
@@ -510,7 +510,7 @@ def forward(
         freqs_action: torch.Tensor,
         freqs_state: torch.Tensor,
         action_register_length: int | None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -567,29 +567,108 @@ def forward(
             action_v = v[:, -action_register_length:]
             v = v[:, :-action_register_length]
 
-        updated_k = kv_cache[0]
-        updated_v = kv_cache[1]
-        new_k = torch.cat([updated_k, roped_key], dim=1)
-        new_v = torch.cat([updated_v, v], dim=1)
-        new_k = new_k[:, -self.max_attention_size :]
-        new_v = new_v[:, -self.max_attention_size :]
-
-        if action_register_length is not None:
-            q_cat = torch.cat([roped_query, roped_action_query], dim=1)
-            k_cat = torch.cat([new_k, roped_action_key], dim=1)
-            v_cat = torch.cat([new_v, action_v], dim=1)
+        if getattr(kv_cache, "is_ar_diffusion_paged_context", False):
+            x = self._forward_ar_diffusion_paged(
+                roped_query=roped_query,
+                roped_key=roped_key,
+                value=v,
+                roped_action_query=roped_action_query,
+                roped_action_key=roped_action_key,
+                action_value=action_v,
+                action_register_length=action_register_length,
+                kv_cache=kv_cache,
+            )
         else:
-            q_cat = roped_query
-            k_cat = new_k
-            v_cat = new_v
+            updated_k = kv_cache[0]
+            updated_v = kv_cache[1]
+            new_k = torch.cat([updated_k, roped_key], dim=1)
+            new_v = torch.cat([updated_v, v], dim=1)
+            new_k = new_k[:, -self.max_attention_size :]
+            new_v = new_v[:, -self.max_attention_size :]
+
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_cat = torch.cat([new_k, roped_action_key], dim=1)
+                v_cat = torch.cat([new_v, action_v], dim=1)
+            else:
+                q_cat = roped_query
+                k_cat = new_k
+                v_cat = new_v
 
-        x = self.attn(q_cat, k_cat, v_cat)
-        updated_kv_cache = torch.stack([new_k, new_v], dim=0)
+            x = self.attn(q_cat, k_cat, v_cat)
+            updated_kv_cache = torch.stack([new_k, new_v], dim=0)
 
         x = x.flatten(2)
         x = self.o(x)
         return x, updated_kv_cache
 
+    def _forward_ar_diffusion_paged(
+        self,
+        *,
+        roped_query: torch.Tensor,
+        roped_key: torch.Tensor,
+        value: torch.Tensor,
+        roped_action_query: torch.Tensor | None,
+        roped_action_key: torch.Tensor | None,
+        action_value: torch.Tensor | None,
+        action_register_length: int | None,
+        kv_cache: Any,
+    ) -> torch.Tensor:
+        from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ar_diffusion_paged_attention
+
+        if roped_key.shape[0] != 1:
+            raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1")
+        if roped_key.shape[1] != kv_cache.seq_len:
+            raise RuntimeError(
+                "AR-Diffusion paged context seq_len="
+                f"{kv_cache.seq_len} but current video KV has {roped_key.shape[1]} tokens"
+            )
+
+        ctx = kv_cache.forward_ctx
+        layer_idx = kv_cache.layer_idx
+        device = roped_key.device
+        ctx.ensure_video_slots(device)
+
+        k_pool = ctx.kv_cache._k_pools[layer_idx]
+        v_pool = ctx.kv_cache._v_pools[layer_idx]
+        video_slots = ctx.current_video_slot_mapping
+        if video_slots is None:
+            raise RuntimeError("AR-Diffusion paged context did not prepare current video slots")
+        k_pool[video_slots] = roped_key[0].to(k_pool.dtype)
+        v_pool[video_slots] = value[0].to(v_pool.dtype)
+
+        action_len = int(action_register_length or 0)
+        if action_len > 0:
+            if roped_action_key is None or action_value is None or roped_action_query is None:
+                raise RuntimeError("AR-Diffusion paged action path missing action query/key/value tensors")
+            ctx.ensure_action_slots(action_len, device)
+            action_slots = ctx.action_slot_mapping
+            if action_slots is None:
+                raise RuntimeError("AR-Diffusion paged context did not prepare action slots")
+            k_pool[action_slots] = roped_action_key[0].to(k_pool.dtype)
+            v_pool[action_slots] = action_value[0].to(v_pool.dtype)
+            q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+        else:
+            q_cat = roped_query
+
+        block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+            action_len=action_len,
+            query_len=q_cat.shape[1],
+            device=device,
+        )
+        return ar_diffusion_paged_attention(
+            q_cat,
+            ctx.kv_cache.key_cache(layer_idx),
+            ctx.kv_cache.value_cache(layer_idx),
+            block_table=block_table,
+            query_start_loc=query_start_loc,
+            seq_lens=seq_lens,
+            max_query_len=max_query_len,
+            max_seq_len=max_seq_len,
+            softmax_scale=self.head_dim**-0.5,
+            causal=False,
+        )
+
 
 # ── Attention Block ─────────────────────────────────────────────────
 
@@ -646,7 +725,7 @@ def forward(
         freqs_state: torch.Tensor,
         context: torch.Tensor,
         action_register_length: int | None = None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         crossattn_cache: dict | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
@@ -925,10 +1004,10 @@ def _forward_blocks(
         action: torch.Tensor | None,
         timestep_action: torch.Tensor | None,
         state: torch.Tensor | None,
-        kv_cache: list[torch.Tensor],
+        kv_cache: list[Any],
         current_start_frame: int,
         crossattn_cache: list[dict] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor | None]]:
         x = x.flatten(start_dim=2).transpose(1, 2)
         B = x.shape[0]
         F_t = timestep.shape[1]
@@ -966,7 +1045,7 @@ def _forward_blocks(
             clip_embedding = self.img_emb(clip_feature)
             context = torch.cat([clip_embedding, context], dim=1)
 
-        updated_kv_caches: list[torch.Tensor] = []
+        updated_kv_caches: list[torch.Tensor | None] = []
         for block_index, block in enumerate(self.blocks):
             x, updated_kv_cache = block(
                 x=x,
@@ -1000,7 +1079,7 @@ def _forward_inference(
         timestep: torch.Tensor,
         context: torch.Tensor,
         seq_len: int,
-        kv_cache: list[torch.Tensor],
+        kv_cache: list[Any],
         current_start_frame: int,
         crossattn_cache: list[dict] | None = None,
         y: torch.Tensor | None = None,
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 3d840dacfbe..aeb2ff8edce 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -102,20 +102,21 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     _ar_diffusion_kv_state = None  # set by the runner before each forward
 
-    def _kv_get(self, state, is_negative):
-        return self._ar_diffusion_kv_state.get_kv_caches(is_negative)
+    def _kv_get(self, state, is_negative, seq_len=None, update_kv_cache=False):
+        return self._ar_diffusion_kv_state.get_kv_caches(
+            is_negative,
+            seq_len=seq_len,
+            commit_current=update_kv_cache,
+        )
 
     def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
-        # The engine owns all KV allocation: self-attn is allocated per chunk in
-        # _kv_update and read via _kv_get (empty window at prefill start); cross-attn
-        # is populated eagerly in _kv_populate_cross. Nothing to create model-side.
+        # The engine owns all KV allocation: self-attn is allocated lazily from
+        # paged contexts, and cross-attn is populated eagerly in _kv_populate_cross.
+        # Nothing is created model-side.
         return
 
-    def _kv_update(self, state, layer_idx, updated_kv, is_negative, seq_len=None):
-        self._ar_diffusion_kv_state.update_kv_cache(layer_idx, updated_kv, is_negative, seq_len)
-
-    def _kv_commit(self):
-        self._ar_diffusion_kv_state.commit_chunk()
+    def _kv_commit(self, is_negative: bool):
+        self._ar_diffusion_kv_state.commit_paged_context(is_negative)
 
     def _kv_get_cross(self, state, is_negative):
         """Cross-attn cache from the engine pool (text k/v + I2V image k_img/v_img)."""
@@ -432,20 +433,17 @@ def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor
             state=kwargs.get("state_features"),
             embodiment_id=kwargs.get("embodiment_id"),
         )
-        if kwargs.get("update_kv_cache", False) and updated_kv_caches:
-            state = kwargs.get("dreamzero_state", self.state)
+        if kwargs.get("update_kv_cache", False):
             is_neg = kwargs.get("is_negative", False)
-            if self._ar_diffusion_kv_state is not None:
-                logger.info(
-                    "AR-Diffusion pipeline predict_noise -> write-back: is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
-                    is_neg,
-                    kwargs.get("seq_len"),
-                    kwargs.get("current_start_frame"),
-                    len(updated_kv_caches),
-                )
-            for i, kv in enumerate(updated_kv_caches):
-                self._kv_update(state, i, kv, is_neg, kwargs.get("seq_len"))
-            self._kv_commit()
+            logger.info(
+                "AR-Diffusion pipeline predict_noise -> commit paged context: "
+                "is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
+                is_neg,
+                kwargs.get("seq_len"),
+                kwargs.get("current_start_frame"),
+                len(updated_kv_caches),
+            )
+            self._kv_commit(is_neg)
 
         return video_pred, action_pred
 
@@ -921,7 +919,7 @@ def _prefill_kv_cache(
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=self._kv_get(state, False),
+                kv_cache=self._kv_get(state, False, seq_len=frame_seqlen, update_kv_cache=True),
                 crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
@@ -929,7 +927,7 @@ def _prefill_kv_cache(
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=self._kv_get(state, True),
+                    kv_cache=self._kv_get(state, True, seq_len=frame_seqlen, update_kv_cache=True),
                     crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
@@ -971,7 +969,7 @@ def _prefill_kv_cache(
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=self._kv_get(state, False),
+                kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=True),
                 crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
@@ -979,7 +977,7 @@ def _prefill_kv_cache(
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=self._kv_get(state, True),
+                    kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=True),
                     crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
@@ -1073,7 +1071,7 @@ def diffuse(
                     hidden_states=noisy_input.transpose(1, 2),
                     timestep_video=timestep,
                     encoder_hidden_states=prompt_embeds,
-                    kv_cache=self._kv_get(state, False),
+                    kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=False),
                     crossattn_cache=self._kv_get_cross(state, False),
                     y=y,
                     clip_feature=state.clip_feas,
@@ -1088,7 +1086,7 @@ def diffuse(
                         hidden_states=noisy_input.transpose(1, 2),
                         timestep_video=timestep,
                         encoder_hidden_states=negative_prompt_embeds,
-                        kv_cache=self._kv_get(state, True),
+                        kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=False),
                         crossattn_cache=self._kv_get_cross(state, True),
                         y=y,
                         clip_feature=state.clip_feas,
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
index 31311a706ce..41d3dddde09 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
@@ -20,26 +20,32 @@
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
     ChunkWindowManager,
     ChunkWindowSpec,
-    allocate_kv_pool,
+    allocate_kv_pool_with_views,
     chunk_slot_mapping,
     compute_slot_mapping,
-    pool_gather_window,
     pool_write_chunk,
     resident_block_ids,
 )
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+    ARDiffusionPagedForwardContext,
+    ARDiffusionPagedLayerContext,
+    ar_diffusion_paged_attention,
+)
 
 __all__ = [
     "ARDiffusionKVCache",
     "ARDiffusionKVConfig",
+    "ARDiffusionPagedForwardContext",
+    "ARDiffusionPagedLayerContext",
     "ARDiffusionRequestAdapter",
     "ChunkWindowManager",
     "ChunkWindowSpec",
-    "allocate_kv_pool",
+    "allocate_kv_pool_with_views",
+    "ar_diffusion_paged_attention",
     "build_kv_manager",
     "chunk_slot_mapping",
     "compute_num_blocks",
     "compute_slot_mapping",
-    "pool_gather_window",
     "pool_write_chunk",
     "resident_block_ids",
 ]
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
index bbf0af90f58..28854328790 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
@@ -12,6 +12,8 @@
 
 from __future__ import annotations
 
+import inspect
+import os
 from collections.abc import Sequence
 
 import torch
@@ -28,10 +30,8 @@
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
     ChunkWindowSpec,
-    allocate_kv_pool,
-    build_window_slots,
+    allocate_kv_pool_with_views,
     chunk_slot_mapping,
-    pool_gather_window,
     pool_write_chunk,
     resident_block_ids,
 )
@@ -135,13 +135,13 @@ def build_kv_manager(
         kv_cache_tensors=tensors,
         kv_cache_groups=[group],
     )
-    return KVCacheManager(
-        config,
-        max_model_len=max_model_len,
-        scheduler_block_size=spec.block_size,
-        hash_block_size=spec.block_size,
-        enable_caching=enable_caching,
-    )
+    kwargs = dict(max_model_len=max_model_len, hash_block_size=spec.block_size, enable_caching=enable_caching)
+    params = inspect.signature(KVCacheManager).parameters
+    if "scheduler_block_size" in params:
+        kwargs["scheduler_block_size"] = spec.block_size
+    if "max_num_batched_tokens" in params:
+        kwargs["max_num_batched_tokens"] = max_model_len
+    return KVCacheManager(config, **kwargs)
 
 
 class ARDiffusionKVCache:
@@ -273,17 +273,34 @@ def _cross_pool_bytes(length: int) -> int:
                 config.window_chunks,
             )
             num_blocks = min_blocks
-        layer_names = [f"bde.layer.{i}" for i in range(num_layers)]
+        layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
         self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
+        self.managed_num_blocks = num_blocks
         self.num_blocks = num_blocks
+        # Scratch blocks are outside KVCacheManager ownership. They hold current
+        # denoise-step video KV when update_kv_cache=False and action/state KV in
+        # both modes; they are reused every forward and never committed.
+        scratch_per_branch = int(os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "8"))
+        if scratch_per_branch <= 0:
+            raise ValueError("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH must be positive")
+        self.scratch_blocks_per_branch = scratch_per_branch
+        self.scratch_num_blocks = 2 * scratch_per_branch
+        self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks
         self.null_block_id = self.manager.block_pool.null_block.block_id
 
         # Allocate the per-layer paged K/V pools on the given device.
+        self._kv_pools: list[torch.Tensor] = []
         self._k_pools: list[torch.Tensor] = []
         self._v_pools: list[torch.Tensor] = []
         if device is not None:
-            self._k_pools, self._v_pools = allocate_kv_pool(
-                num_blocks, block_size, num_layers, num_kv_heads, head_size, dtype, device
+            self._kv_pools, self._k_pools, self._v_pools = allocate_kv_pool_with_views(
+                self.num_blocks_total,
+                block_size,
+                num_layers,
+                num_kv_heads,
+                head_size,
+                dtype,
+                device,
             )
 
     # -- cross-attention pool access -------------------------------------------
@@ -378,6 +395,15 @@ def allocate_chunk(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
         )
         return table
 
+    def allocate_token_slots(self, adapter: ARDiffusionRequestAdapter, num_tokens: int) -> list[int]:
+        """Allocate managed blocks for an in-flight video span without committing it."""
+        if num_tokens <= 0:
+            raise ValueError(f"num_tokens must be positive, got {num_tokens}")
+        blocks = self.manager.allocate_slots(adapter, num_new_tokens=num_tokens)
+        if blocks is None:
+            raise RuntimeError("AR-Diffusion KV pool exhausted while allocating paged attention slots")
+        return self.block_table(adapter)
+
     def block_table(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
         return list(self.manager.get_block_ids(adapter.request_id)[0])
 
@@ -390,25 +416,43 @@ def chunk_write_slots(self, adapter: ARDiffusionRequestAdapter) -> torch.Tensor:
             self.block_size,
         )
 
+    def scratch_block_ids(self, is_negative: bool, start: int, count: int) -> list[int]:
+        """Return branch-local scratch block ids outside manager ownership."""
+        if count < 0 or start < 0:
+            raise ValueError(f"scratch start/count must be non-negative, got start={start}, count={count}")
+        if start + count > self.scratch_blocks_per_branch:
+            raise RuntimeError(
+                "AR-Diffusion paged attention scratch blocks exhausted: "
+                f"need [{start}, {start + count}) of {self.scratch_blocks_per_branch}. "
+                "Increase AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH."
+            )
+        branch_offset = self.scratch_blocks_per_branch if is_negative else 0
+        base = self.managed_num_blocks + branch_offset + start
+        return list(range(base, base + count))
+
+    def key_cache(self, layer_idx: int) -> torch.Tensor:
+        return self._kv_pools[layer_idx][0]
+
+    def value_cache(self, layer_idx: int) -> torch.Tensor:
+        return self._kv_pools[layer_idx][1]
+
     def window_block_ids(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
-        """Resident (non-null) blocks the read path gathers the window from."""
+        """Resident (non-null) managed blocks visible to paged attention."""
         return resident_block_ids(self.block_table(adapter), self.null_block_id)
 
     def commit_chunk(self, adapter: ARDiffusionRequestAdapter) -> None:
         """Advance the adapter by one chunk after its K/V is written.
 
-        This is the standalone per-chunk advance used by the unit lifecycle and the
-        ``tests/bde`` suite. In the DreamZero production path the advance happens
-        elsewhere: :meth:`ARDiffusionKVState.update_kv_cache` (gather.py) calls
-        ``adapter.on_chunk_committed()`` directly as it allocates each frame-chunk,
-        and the pipeline's ``_kv_commit`` -> :meth:`ARDiffusionKVState.commit_chunk` is only
-        a no-op log. Call once per chunk, not per denoise step.
+        This standalone primitive is used by low-level manager tests. In the
+        DreamZero paged-attention path, :meth:`ARDiffusionKVState.commit_paged_context`
+        advances the adapter only after the forward succeeds. Call once per
+        committed chunk, not per denoise step.
         """
         _log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
         adapter.on_chunk_committed()
         _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
 
-    # -- pool-backed K/V access (Step 4 — gather / write) --------------------
+    # -- pool-backed K/V access --------------------------------------------
 
     def write_chunk_kv(
         self,
@@ -434,58 +478,3 @@ def write_chunk_kv(
             new_v,
             slots,
         )
-
-    def gather_window(self, layer_index: int, adapter: ARDiffusionRequestAdapter) -> torch.Tensor:
-        """Gather the resident-window K/V for one layer.
-
-        Returns a ``(2, 1, window, n_heads, head_dim)`` tensor — the format
-        DreamZero's existing attention expects as its ``kv_cache`` argument.
-        """
-        window_ids = self.window_block_ids(adapter)
-        window = pool_gather_window(
-            self._k_pools[layer_index],
-            self._v_pools[layer_index],
-            window_ids,
-            self.block_size,
-            self.spec.sliding_window,
-        )
-        _log.debug(
-            "AR-Diffusion gather: req=%s layer=%d blocks=%s window=%s dev=%s",
-            adapter.request_id,
-            layer_index,
-            window_ids,
-            tuple(window.shape),
-            window.device,
-        )
-        return window
-
-    def gather_window_all_layers(self, adapter: ARDiffusionRequestAdapter) -> list[torch.Tensor]:
-        """Gather the resident-window K/V for every layer in one shot.
-
-        The window's block table is identical across layers within a forward, so
-        the block-id lookup and the flat slot index are computed **once** and shared
-        across all per-layer gathers — removing the per-layer Python work that the
-        single-layer ``gather_window`` repeats. Returns one
-        ``(2, 1, window, n_heads, head_dim)`` tensor per layer.
-        """
-        window_ids = self.window_block_ids(adapter)
-        slots = build_window_slots(window_ids, self.block_size, self._k_pools[0].device)
-        windows = [
-            pool_gather_window(
-                self._k_pools[i],
-                self._v_pools[i],
-                window_ids,
-                self.block_size,
-                self.spec.sliding_window,
-                slots=slots,
-            )
-            for i in range(self.num_layers)
-        ]
-        _log.debug(
-            "AR-Diffusion gather-all: req=%s layers=%d blocks=%d window=%s",
-            adapter.request_id,
-            self.num_layers,
-            len(window_ids),
-            tuple(windows[0].shape),
-        )
-        return windows
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py
index 3ef4e6288bb..cb84467a3e4 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py
@@ -6,13 +6,9 @@
 
 * **Slot mapping** — absolute token positions → physical KV-cache slots, the
   standard PagedAttention layout ``slot(pos) = block_id(pos) * block_size +
-  (pos % block_size)``. Phase 1 is single-request and gathers the resident
-  window into the model's existing attention rather than calling vLLM's paged
-  kernel, so a thin slot-mapping helper is used instead of vLLM's ``BlockTables``.
-* **Pool I/O** — allocate per-layer flat K/V pools, write a committed chunk into
-  slots, and gather the resident window back as a contiguous tensor. The
-  per-forward gather copy is the deliberate Phase-1 simplification the fused
-  paged backend retires in Phase 2.
+  (pos % block_size)``.
+* **Pool I/O** — allocate per-layer FlashAttention-compatible paged K/V pools
+  plus flat compatibility views for direct slot writes.
 * **Chunk-window eviction** — a ``SlidingWindowSpec`` subclass whose unit is a
   *chunk* (``sliding_window = window_chunks * chunk_size``) plus a manager that
   evicts at chunk boundaries. Memory policy / refcounting / ``null_block``
@@ -27,7 +23,21 @@
 import torch
 from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
 from vllm.v1.kv_cache_interface import SlidingWindowSpec
-from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+
+try:
+    from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+except ModuleNotFoundError:
+
+    def register_kv_cache_spec(*, manager_class, uniform_type_base_spec=None):
+        """Compatibility shim for vLLM versions without kv_cache_spec_registry."""
+
+        def decorator(spec_class):
+            from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
+            spec_manager_map[spec_class] = manager_class
+            return spec_class
+
+        return decorator
 
 
 # ── Slot mapping ────────────────────────────────────────────────────────────
@@ -85,10 +95,10 @@ def resident_block_ids(block_ids: Sequence[int], null_block_id: int) -> list[int
     return [int(b) for b in block_ids if int(b) != null_block_id]
 
 
-# ── Pool I/O (allocate / write / gather) ────────────────────────────────────
+# ── Pool I/O (allocate / write) ─────────────────────────────────────────────
 
 
-def allocate_kv_pool(
+def allocate_kv_pool_with_views(
     num_blocks: int,
     block_size: int,
     num_layers: int,
@@ -96,19 +106,23 @@ def allocate_kv_pool(
     head_dim: int,
     dtype: torch.dtype,
     device: torch.device,
-) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
-    """Allocate per-layer paged K and V pools on *device*.
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    """Allocate vLLM-style paged KV pools plus flat slot-write views.
 
-    Each pool is a contiguous ``(num_blocks * block_size, num_kv_heads, head_dim)``
-    tensor — a flat address space indexed by ``slot = block_id * block_size + offset``.
-    The slot mapping (above) tells writes and gathers where each token lives.
+    The owning tensor follows FlashAttention's block-table cache layout:
+    ``(2, num_blocks, block_size, num_kv_heads, head_dim)``, where dim-0 is
+    ``[K, V]``.  The flat K/V views keep ``slot = block_id * block_size +
+    offset`` writes simple without changing the kernel-facing cache layout.
     """
+    kv_pools: list[torch.Tensor] = []
     k_pools: list[torch.Tensor] = []
     v_pools: list[torch.Tensor] = []
     for _ in range(num_layers):
-        k_pools.append(torch.empty(num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device))
-        v_pools.append(torch.empty(num_blocks * block_size, num_kv_heads, head_dim, dtype=dtype, device=device))
-    return k_pools, v_pools
+        kv = torch.empty(2, num_blocks, block_size, num_kv_heads, head_dim, dtype=dtype, device=device)
+        kv_pools.append(kv)
+        k_pools.append(kv[0].reshape(num_blocks * block_size, num_kv_heads, head_dim))
+        v_pools.append(kv[1].reshape(num_blocks * block_size, num_kv_heads, head_dim))
+    return kv_pools, k_pools, v_pools
 
 
 def pool_write_chunk(
@@ -131,54 +145,6 @@ def pool_write_chunk(
     v_pool[slot_mapping] = new_v[0]
 
 
-def build_window_slots(window_block_ids: list[int], block_size: int, device: torch.device) -> torch.Tensor:
-    """Flat slot index covering all resident blocks, in table order.
-
-    Vectorized (no Python loop): ``slot = block_id * block_size + offset`` for every
-    ``offset in [0, block_size)``. The window's blocks are identical across all layers
-    within a forward, so this is built once per forward and shared.
-    """
-    if not window_block_ids:
-        # No resident blocks (e.g. prefill start, before the first chunk is
-        # committed): an empty slot index gathers a zero-length window, which is
-        # the concat-identity the model's attention expects as its initial KV.
-        return torch.empty(0, dtype=torch.long, device=device)
-    starts = torch.tensor(window_block_ids, dtype=torch.long, device=device) * block_size
-    offsets = torch.arange(block_size, device=device)
-    return (starts[:, None] + offsets[None, :]).reshape(-1)
-
-
-def pool_gather_window(
-    k_pool: torch.Tensor,
-    v_pool: torch.Tensor,
-    window_block_ids: list[int],
-    block_size: int,
-    max_attention_size: int,
-    *,
-    slots: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Materialize the resident window K/V as a contiguous tensor.
-
-    Gathers all slots belonging to ``window_block_ids`` (in table order), trims to
-    the last ``max_attention_size`` tokens (matching the existing ``cat`` + slice),
-    and returns a ``(2, batch, window_tokens, n_heads, head_dim)`` tensor — exactly
-    the format DreamZero's attention expects as ``kv_cache``. Pass a precomputed
-    ``slots`` (from :func:`build_window_slots`) to share the index across layers.
-    """
-    if slots is None:
-        slots = build_window_slots(window_block_ids, block_size, k_pool.device)
-    # Gather then trim to the attention window.
-    wk = k_pool[slots]  # (window_slots, n_heads, head_dim)
-    wv = v_pool[slots]
-    wk = wk[-max_attention_size:]  # (window, n_heads, head_dim)
-    wv = wv[-max_attention_size:]
-    # Add the batch dim then stack as (2, batch, window, n_heads, head_dim).
-    wk = wk.unsqueeze(0)  # (1, window, n_heads, head_dim)
-    wv = wv.unsqueeze(0)
-    window = torch.stack([wk, wv], dim=0)  # (2, 1, window, n_heads, head_dim)
-    return window
-
-
 def _drop_batch(t: torch.Tensor) -> torch.Tensor:
     """``(1, L, n_heads, head_dim)`` -> ``(L, n_heads, head_dim)``; pass 3-D through."""
     if t.dim() == 4 and t.shape[0] == 1:
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
new file mode 100644
index 00000000000..ead80c4a1c9
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
@@ -0,0 +1,311 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Paged self-attention helpers for AR-Diffusion DreamZero KV reuse."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, ClassVar
+
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
+
+
+def _ceil_div(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
+def _slot_mapping_for_blocks(
+    block_ids: list[int],
+    token_count: int,
+    block_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    if token_count == 0:
+        return torch.empty(0, dtype=torch.long, device=device)
+    positions = torch.arange(token_count, dtype=torch.long)
+    return compute_slot_mapping(block_ids, positions, block_size).to(device=device)
+
+
+@dataclass
+class ARDiffusionPagedForwardContext:
+    """Mutable branch-level state shared by all layer contexts in one forward."""
+
+    kv_cache: Any
+    adapter: Any
+    is_negative: bool
+    history_block_ids: list[int]
+    seq_len: int
+    commit_current: bool
+    max_video_tokens: int
+    current_video_block_ids: list[int] = field(default_factory=list)
+    current_video_slot_mapping: torch.Tensor | None = None
+    action_scratch_block_ids: list[int] = field(default_factory=list)
+    action_slot_mapping: torch.Tensor | None = None
+    query_len: int = 0
+    kv_len: int = 0
+    _allocated_video: bool = False
+    _committed: bool = False
+    _action_len: int = 0
+
+    @property
+    def block_size(self) -> int:
+        return int(self.kv_cache.block_size)
+
+    @property
+    def num_current_video_blocks(self) -> int:
+        if self.seq_len % self.block_size != 0:
+            raise AssertionError(
+                "AR-Diffusion paged attention expects frame-aligned seq_len "
+                f"(multiple of block_size={self.block_size}), got {self.seq_len}"
+            )
+        return self.seq_len // self.block_size
+
+    def ensure_video_slots(self, device: torch.device) -> None:
+        """Allocate/write targets for the current video tokens, once per branch."""
+        if self._allocated_video:
+            return
+
+        n_blocks = self.num_current_video_blocks
+        if self.commit_current:
+            start = int(self.adapter.num_computed_tokens)
+            self.kv_cache.allocate_token_slots(self.adapter, self.seq_len)
+            table = self.kv_cache.block_table(self.adapter)
+            start_block = start // self.block_size
+            self.current_video_block_ids = [int(b) for b in table[start_block : start_block + n_blocks]]
+            positions = torch.arange(start, start + self.seq_len, dtype=torch.long)
+            self.current_video_slot_mapping = compute_slot_mapping(table, positions, self.block_size).to(device=device)
+        else:
+            self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.is_negative, 0, n_blocks)
+            self.current_video_slot_mapping = _slot_mapping_for_blocks(
+                self.current_video_block_ids,
+                self.seq_len,
+                self.block_size,
+                device,
+            )
+        self._allocated_video = True
+
+    def ensure_action_slots(self, action_len: int, device: torch.device) -> None:
+        """Reserve scratch slots for action/state K/V, if present."""
+        if action_len <= 0:
+            self.action_scratch_block_ids = []
+            self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
+            self._action_len = 0
+            return
+
+        self.ensure_video_slots(device)
+        if self.action_slot_mapping is not None and self._action_len == action_len:
+            return
+
+        action_blocks = _ceil_div(action_len, self.block_size)
+        scratch_offset = 0 if self.commit_current else len(self.current_video_block_ids)
+        self.action_scratch_block_ids = self.kv_cache.scratch_block_ids(
+            self.is_negative,
+            scratch_offset,
+            action_blocks,
+        )
+        self.action_slot_mapping = _slot_mapping_for_blocks(
+            self.action_scratch_block_ids,
+            action_len,
+            self.block_size,
+            device,
+        )
+        self._action_len = action_len
+
+    def video_block_table(self, device: torch.device) -> tuple[list[int], int]:
+        self.ensure_video_slots(device)
+        if self.max_video_tokens % self.block_size != 0:
+            raise AssertionError(
+                "AR-Diffusion paged attention requires max_video_tokens to be block-aligned, "
+                f"got max_video_tokens={self.max_video_tokens}, block_size={self.block_size}"
+            )
+        all_video_blocks = self.history_block_ids + self.current_video_block_ids
+        max_video_blocks = self.max_video_tokens // self.block_size
+        visible_video_blocks = all_video_blocks[-max_video_blocks:] if max_video_blocks else []
+        video_len = min(len(all_video_blocks) * self.block_size, self.max_video_tokens)
+        return visible_video_blocks, video_len
+
+    def build_block_table(
+        self,
+        *,
+        action_len: int,
+        query_len: int,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+        """Build FlashAttention block-table metadata for one self-attn call."""
+        video_blocks, video_len = self.video_block_table(device)
+        self.ensure_action_slots(action_len, device)
+        action_blocks = self.action_scratch_block_ids if action_len > 0 else []
+        block_ids = video_blocks + action_blocks
+        if not block_ids:
+            raise RuntimeError("AR-Diffusion paged attention needs at least current video KV blocks")
+
+        self.query_len = int(query_len)
+        self.kv_len = int(video_len + action_len)
+        block_table = torch.tensor([block_ids], dtype=torch.int32, device=device)
+        query_start_loc = torch.tensor([0, self.query_len], dtype=torch.int32, device=device)
+        seq_lens = torch.tensor([self.kv_len], dtype=torch.int32, device=device)
+        return block_table, query_start_loc, seq_lens, self.query_len, self.kv_len
+
+    def mark_committed(self) -> None:
+        self._committed = True
+
+
+@dataclass
+class ARDiffusionPagedLayerContext:
+    """Layer-specific handle passed through DreamZero's existing ``kv_cache`` slot."""
+
+    is_ar_diffusion_paged_context: ClassVar[bool] = True
+    layer_idx: int
+    forward_ctx: ARDiffusionPagedForwardContext
+
+    @property
+    def kv_cache(self):
+        return self.forward_ctx.kv_cache
+
+    @property
+    def adapter(self):
+        return self.forward_ctx.adapter
+
+    @property
+    def is_negative(self) -> bool:
+        return self.forward_ctx.is_negative
+
+    @property
+    def history_block_ids(self) -> list[int]:
+        return self.forward_ctx.history_block_ids
+
+    @property
+    def current_video_block_ids(self) -> list[int]:
+        return self.forward_ctx.current_video_block_ids
+
+    @property
+    def current_video_slot_mapping(self) -> torch.Tensor | None:
+        return self.forward_ctx.current_video_slot_mapping
+
+    @property
+    def action_scratch_block_ids(self) -> list[int]:
+        return self.forward_ctx.action_scratch_block_ids
+
+    @property
+    def action_slot_mapping(self) -> torch.Tensor | None:
+        return self.forward_ctx.action_slot_mapping
+
+    @property
+    def seq_len(self) -> int:
+        return self.forward_ctx.seq_len
+
+    @property
+    def query_len(self) -> int:
+        return self.forward_ctx.query_len
+
+    @property
+    def kv_len(self) -> int:
+        return self.forward_ctx.kv_len
+
+    @property
+    def commit_current(self) -> bool:
+        return self.forward_ctx.commit_current
+
+
+def is_ar_diffusion_paged_context(value: object) -> bool:
+    return isinstance(value, ARDiffusionPagedLayerContext)
+
+
+def _reference_paged_attention(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    softmax_scale: float,
+    *,
+    causal: bool,
+) -> torch.Tensor:
+    if causal:
+        raise NotImplementedError("AR-Diffusion paged self-attention uses causal=False")
+    outs: list[torch.Tensor] = []
+    block_size = key_cache.shape[1]
+    for i in range(seq_lens.shape[0]):
+        q_start = int(query_start_loc[i].item())
+        q_end = int(query_start_loc[i + 1].item())
+        kv_len = int(seq_lens[i].item())
+        q = query[q_start:q_end]
+        positions = torch.arange(kv_len, device=query.device)
+        logical_blocks = torch.div(positions, block_size, rounding_mode="floor")
+        offsets = positions % block_size
+        physical_blocks = block_table[i, logical_blocks].long()
+        k = key_cache[physical_blocks, offsets]
+        v = value_cache[physical_blocks, offsets]
+        scores = torch.einsum("qhd,khd->hqk", q.float(), k.float()) * float(softmax_scale)
+        probs = torch.softmax(scores, dim=-1).to(v.dtype)
+        outs.append(torch.einsum("hqk,khd->qhd", probs, v))
+    return torch.cat(outs, dim=0)
+
+
+def ar_diffusion_paged_attention(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    *,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    max_query_len: int,
+    max_seq_len: int,
+    softmax_scale: float,
+    causal: bool = False,
+) -> torch.Tensor:
+    """Run non-causal paged attention over a vLLM block table.
+
+    ``query`` may be ``(B, L, H, D)`` or already flattened as ``(T, H, D)``.
+    ``key_cache`` / ``value_cache`` are ``(num_blocks, block_size, H, D)``.
+    """
+    batched = query.dim() == 4
+    if batched:
+        batch, q_len = query.shape[:2]
+        query_flat = query.reshape(batch * q_len, *query.shape[2:])
+    else:
+        query_flat = query
+
+    if not query_flat.is_cuda:
+        out = _reference_paged_attention(
+            query_flat,
+            key_cache,
+            value_cache,
+            block_table,
+            query_start_loc,
+            seq_lens,
+            softmax_scale,
+            causal=causal,
+        )
+    else:
+        from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+        try:
+            from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
+
+            fa_version = get_flash_attn_version(requires_alibi=False, head_size=query_flat.shape[-1])
+        except Exception:
+            fa_version = 2
+
+        out = torch.empty_like(query_flat)
+        flash_attn_varlen_func(
+            q=query_flat,
+            k=key_cache,
+            v=value_cache,
+            out=out,
+            cu_seqlens_q=query_start_loc,
+            max_seqlen_q=int(max_query_len),
+            seqused_k=seq_lens,
+            max_seqlen_k=int(max_seq_len),
+            softmax_scale=float(softmax_scale),
+            causal=causal,
+            block_table=block_table,
+            fa_version=fa_version,
+        )
+
+    if batched:
+        return out.reshape(query.shape)
+    return out
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
index 161faedb152..d3f1216755b 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
@@ -1,25 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
-"""ARDiffusionKVState — per-request bridge between a model rollout and the AR-Diffusion KV pool.
-
-This is the one model-facing piece of the KV stack (everything in ``paged.py`` /
-``manager.py`` is engine-generic). A second model (e.g. the Cosmos port) adds a
-sibling here that reuses the same pool primitives.
-
-Phase 1 keeps the existing attention kernel unchanged — it only changes *where*
-the past-window KV comes from (the pool) and *where* the committed chunk's KV is
-stored (the pool). The model's ``self_attn`` still receives a contiguous
-``(batch, window, n_heads, head_dim)`` K/V pair; the gather materializes that
-tensor from the resident pool blocks.
-"""
+"""DreamZero-facing AR-Diffusion KV bridge backed by paged self-attention."""
 
 from __future__ import annotations
 
-import os
-
-import torch
 from vllm.logger import init_logger
 
-from vllm_omni.experimental.ar_diffusion.kv_cache.paged import _drop_batch, pool_write_chunk
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+    ARDiffusionPagedForwardContext,
+    ARDiffusionPagedLayerContext,
+)
 
 _log = init_logger(__name__)
 
@@ -27,18 +16,10 @@
 class ARDiffusionKVState:
     """Per-request bridge between a DreamZero rollout and the AR-Diffusion KV pool.
 
-    Replaces the model-local ``DreamZeroState`` KV methods::
-      state.get_kv_caches(neg)    →  bde_state.get_kv_caches(neg)
-      state.update_kv_cache(i,kv) →  bde_state.update_kv_cache(i,kv,neg)
-      state.create_kv_caches(...) →  no-op (pool is pre-allocated)
-
-    One ``ARDiffusionKVState`` per request holds the ``ARDiffusionKVCache`` and both CFG adapters
-    (positive / negative). The runner sets ``pipeline._ar_diffusion_kv_state`` before
-    ``pipeline.forward(req)``; the pipeline delegates every KV call to it when
-    the attribute is present and not None.
-
-    Call ``commit_chunk()`` after all layers have been written for a chunk
-    (once per chunk, not per denoise step).
+    The self-attention path is paged-only: ``get_kv_caches`` returns lightweight
+    per-layer contexts, and each attention layer writes current K/V into the
+    pool before calling the FlashAttention block-table kernel.  No resident
+    window is gathered into a contiguous tensor.
     """
 
     def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
@@ -46,23 +27,8 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
         self.pos = pos_adapter
         self.neg = neg_adapter
         self.num_layers = num_layers
-        # --- Approach 2: true chunk-paged write / gather -----------------------
-        # Each get gathers the resident window from the paged pool blocks; each
-        # update writes only the genuinely-new tokens (seq_len growth) into newly
-        # allocated frame-blocks, evicting out-of-window blocks via the
-        # ChunkWindowManager. ``_committed`` is the absolute token count written to
-        # the pool per branch; ``_pending`` holds this forward's new-chunk slot maps.
         self._committed: dict[bool, int] = {False: 0, True: 0}
-        self._pending: dict[bool, list] = {False: [], True: []}
-        # The resident window is frozen across a forward's denoise steps (KV is
-        # only written back after the last step), so the per-step gather is
-        # re-materializing the same tensors ~16x. Memoize the gathered window per
-        # branch; invalidate on write-back (window grows) and on reset.
-        self._gather_cache: dict[bool, list | None] = {False: None, True: None}
-        # Profiling escape hatch: AR_DIFFUSION_KV_NO_MEMO=1 disables the per-forward gather
-        # memoization so every denoise step re-gathers the window (the pre-memo
-        # behavior), for A/B measurement of the memoization win. Default: enabled.
-        self._memo_enabled = os.environ.get("AR_DIFFUSION_KV_NO_MEMO") != "1"
+        self._paged_pending: dict[bool, ARDiffusionPagedForwardContext | None] = {False: None, True: None}
         # Cross-attn KV is populated once after text encoding (eager), then
         # read every denoising step. Tracked per half: the text K/V depend only
         # on the session prompt (survive window-boundary resets), while the I2V
@@ -74,27 +40,82 @@ def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
     def _adapter(self, is_negative: bool):
         return self.neg if is_negative else self.pos
 
-    def get_kv_caches(self, is_negative: bool) -> list:
-        # Engine owns the read end-to-end: at prefill start the adapter has no
-        # resident blocks, so gather returns zero-length windows (the empty seed
-        # the model concatenates against) — no model-local fallback needed.
-        branch = "neg" if is_negative else "pos"
-        cached = self._gather_cache[is_negative] if self._memo_enabled else None
-        if cached is not None:  # window unchanged this forward -> reuse (no re-gather)
-            return cached
+    def get_kv_caches(
+        self,
+        is_negative: bool,
+        seq_len: int | None = None,
+        commit_current: bool = False,
+    ) -> list[ARDiffusionPagedLayerContext]:
+        if seq_len is None:
+            raise ValueError("AR-Diffusion paged self-attention requires seq_len in get_kv_caches()")
+        return self.prepare_paged_context(is_negative, seq_len, commit_current)
+
+    def prepare_paged_context(
+        self,
+        is_negative: bool,
+        seq_len: int,
+        commit_current: bool,
+    ) -> list[ARDiffusionPagedLayerContext]:
+        """Return per-layer paged attention contexts for one branch forward.
+
+        Allocation is lazy: constructing both CFG branches' kwargs must not
+        allocate manager blocks for a branch that a CFG-parallel rank will not run.
+        The first self-attn layer that consumes the context allocates the current
+        video slots, then all layers share that branch-level state.
+        """
+        cs = self.kv_cache.spec.chunk_size
+        if int(seq_len) % cs != 0:
+            raise AssertionError(
+                f"AR-Diffusion expects frame-aligned seq_len (multiple of chunk_size={cs}), got {seq_len}"
+            )
+
+        pending = self._paged_pending.get(is_negative)
+        if pending is not None and pending.commit_current and pending._allocated_video and not pending._committed:
+            raise RuntimeError("AR-Diffusion paged context replaced before its managed current chunk was committed")
+
         adapter = self._adapter(is_negative)
-        windows = self.kv_cache.gather_window_all_layers(adapter)
-        if self._memo_enabled:
-            self._gather_cache[is_negative] = windows
+        branch = "neg" if is_negative else "pos"
+        forward_ctx = ARDiffusionPagedForwardContext(
+            kv_cache=self.kv_cache,
+            adapter=adapter,
+            is_negative=is_negative,
+            history_block_ids=self.kv_cache.window_block_ids(adapter),
+            seq_len=int(seq_len),
+            commit_current=bool(commit_current),
+            max_video_tokens=int(self.kv_cache.spec.sliding_window),
+        )
+        self._paged_pending[is_negative] = forward_ctx
         _log.info(
-            "AR-Diffusion GET   [%s] source=paged-gather layers=%d window0=%s resident_blocks=%d/%d",
+            "AR-Diffusion GET   [%s] source=paged-attn layers=%d history_blocks=%d seq_len=%d commit_current=%s",
             branch,
             self.num_layers,
-            tuple(windows[0].shape),
-            len(self.kv_cache.window_block_ids(adapter)),
-            self.kv_cache.spec.window_chunks,
+            len(forward_ctx.history_block_ids),
+            int(seq_len),
+            bool(commit_current),
         )
-        return windows
+        return [ARDiffusionPagedLayerContext(layer_idx=i, forward_ctx=forward_ctx) for i in range(self.num_layers)]
+
+    def commit_paged_context(self, is_negative: bool) -> None:
+        """Commit the managed current video blocks after a successful forward."""
+        ctx = self._paged_pending.get(is_negative)
+        if ctx is None:
+            return
+        branch = "neg" if is_negative else "pos"
+        if ctx.commit_current and ctx._allocated_video:
+            n_chunks = ctx.seq_len // self.kv_cache.spec.chunk_size
+            for _ in range(n_chunks):
+                ctx.adapter.on_chunk_committed()
+            self._committed[is_negative] += ctx.seq_len
+            _log.info(
+                "AR-Diffusion COMMIT [%s] paged-attn new_tokens=%d chunks=%d resident=%d/%d",
+                branch,
+                ctx.seq_len,
+                n_chunks,
+                len(self.kv_cache.window_block_ids(ctx.adapter)),
+                self.kv_cache.spec.window_chunks,
+            )
+        ctx.mark_committed()
+        self._paged_pending[is_negative] = None
 
     def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
         """Return pool-backed cross-attn cache dicts.
@@ -118,76 +139,11 @@ def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
             )
         return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
 
-    def update_kv_cache(self, layer_idx: int, updated_kv: torch.Tensor, is_negative: bool, seq_len) -> None:
-        branch = "neg" if is_negative else "pos"
-        adapter = self._adapter(is_negative)
-        cs = self.kv_cache.spec.chunk_size
-        if layer_idx == 0:
-            # The K appended this forward is the current observation's length
-            # (seq_len), not the cumulative-minus-committed delta. Allocate those
-            # frame-blocks once per forward (shared across layers).
-            new_count = int(seq_len)
-            # DreamZero appends whole frames each forward and chunk_size ==
-            # frame_seqlen, so seq_len is always a multiple of cs. Guard the
-            # invariant: a non-frame-aligned model would otherwise silently drop
-            # the leading new_count % cs tokens and desync num_computed.
-            assert new_count % cs == 0, (
-                f"AR-Diffusion expects frame-aligned seq_len (multiple of chunk_size={cs}), got {new_count}"
-            )
-            n_chunks = new_count // cs
-            slots = []
-            for _ in range(n_chunks):
-                self.kv_cache.allocate_chunk(adapter)  # allocate + evict old
-                slots.append(self.kv_cache.chunk_write_slots(adapter))
-                adapter.on_chunk_committed()
-            self._pending[is_negative] = slots
-            self._committed[is_negative] += new_count
-            # The window just grew; the memoized gather is stale -> force re-gather
-            # on the next forward's first read of this branch.
-            self._gather_cache[is_negative] = None
-            _log.info(
-                "AR-Diffusion WRITE [%s] new_tokens=%d -> %d frame-chunks  resident=%d/%d  storing %d layers",
-                branch,
-                new_count,
-                n_chunks,
-                len(self.kv_cache.window_block_ids(adapter)),
-                self.kv_cache.spec.window_chunks,
-                self.num_layers,
-            )
-        slots = self._pending[is_negative]
-        if not slots:
-            return
-        n = len(slots) * cs
-        k_all = _drop_batch(updated_kv[0])[-n:]  # (n, n_heads, head_dim) — the new tokens
-        v_all = _drop_batch(updated_kv[1])[-n:]
-        kpool = self.kv_cache._k_pools[layer_idx]
-        vpool = self.kv_cache._v_pools[layer_idx]
-        for c, sm in enumerate(slots):
-            k = k_all[c * cs : (c + 1) * cs].unsqueeze(0).to(kpool.dtype)
-            v = v_all[c * cs : (c + 1) * cs].unsqueeze(0).to(vpool.dtype)
-            pool_write_chunk(kpool, vpool, k, v, sm)
-
-    def commit_chunk(self) -> None:
-        """No-op for DreamZero (logs only).
-
-        The per-chunk advance already happened in :meth:`update_kv_cache`, which
-        calls ``adapter.on_chunk_committed()`` at layer 0 while allocating each
-        frame-chunk. This method exists so the pipeline's ``_kv_commit`` bridge has
-        a symmetric call; the resident window is retained across forwards.
-        """
-        _log.info("AR-Diffusion COMMIT (paged; resident window retained across forwards)")
-
     def close(self) -> None:
-        """Final teardown — free both branches' resident pool blocks.
-
-        Unlike :meth:`reset`, this does **not** restart the adapters: the caller is
-        dropping this ``ARDiffusionKVState`` for good (session eviction / shutdown), so the
-        blocks return to the ``BlockPool`` and the dead adapters are discarded with
-        the state. Frees pool ownership that would otherwise leak when a session is
-        evicted from the runner's session map.
-        """
+        """Final teardown: free both branches' resident pool blocks."""
         self.kv_cache.end_request(self.pos)
         self.kv_cache.end_request(self.neg)
+        self._paged_pending = {False: None, True: None}
 
     def reset(self, *, keep_cross_text: bool = False) -> None:
         """Drop this session's KV — mirrors the model-local ``state.reset()``.
@@ -206,12 +162,11 @@ def reset(self, *, keep_cross_text: bool = False) -> None:
                 Session resets (prompt change / explicit reset) pass ``False``.
         """
         pos_id, neg_id = self.pos.request_id, self.neg.request_id
-        self.close()  # free resident blocks for both branches
-        self.pos = self.kv_cache.begin_request(pos_id)  # fresh, empty window
+        self.close()
+        self.pos = self.kv_cache.begin_request(pos_id)
         self.neg = self.kv_cache.begin_request(neg_id)
         self._committed = {False: 0, True: 0}
-        self._pending = {False: [], True: []}
-        self._gather_cache = {False: None, True: None}
+        self._paged_pending = {False: None, True: None}
         if not keep_cross_text:
             self._cross_text_populated = {False: False, True: False}
         self._cross_img_populated = {False: False, True: False}

From 6776c35b5da6f6f55a119183acee4d0ae7ad2c16 Mon Sep 17 00:00:00 2001
From: Acerak01-fy <wfy2003324@163.com>
Date: Fri, 3 Jul 2026 03:01:32 +0000
Subject: [PATCH 43/51] Simplify paged attention slot mapping helpers
 Signed-off-by: Acerak01-fy <wfy2003324@163.com> Signed-off-by: Taichang Zhou
 <tzhouam@connect.ust.hk>

---
 .../ar_diffusion/kv_cache/paged_attention.py  | 34 +++++--------------
 1 file changed, 9 insertions(+), 25 deletions(-)

diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
index ead80c4a1c9..01bfedf9508 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
@@ -11,22 +11,6 @@
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
 
 
-def _ceil_div(a: int, b: int) -> int:
-    return (a + b - 1) // b
-
-
-def _slot_mapping_for_blocks(
-    block_ids: list[int],
-    token_count: int,
-    block_size: int,
-    device: torch.device,
-) -> torch.Tensor:
-    if token_count == 0:
-        return torch.empty(0, dtype=torch.long, device=device)
-    positions = torch.arange(token_count, dtype=torch.long)
-    return compute_slot_mapping(block_ids, positions, block_size).to(device=device)
-
-
 @dataclass
 class ARDiffusionPagedForwardContext:
     """Mutable branch-level state shared by all layer contexts in one forward."""
@@ -77,12 +61,12 @@ def ensure_video_slots(self, device: torch.device) -> None:
             self.current_video_slot_mapping = compute_slot_mapping(table, positions, self.block_size).to(device=device)
         else:
             self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.is_negative, 0, n_blocks)
-            self.current_video_slot_mapping = _slot_mapping_for_blocks(
+            positions = torch.arange(self.seq_len, dtype=torch.long)
+            self.current_video_slot_mapping = compute_slot_mapping(
                 self.current_video_block_ids,
-                self.seq_len,
+                positions,
                 self.block_size,
-                device,
-            )
+            ).to(device=device)
         self._allocated_video = True
 
     def ensure_action_slots(self, action_len: int, device: torch.device) -> None:
@@ -97,19 +81,19 @@ def ensure_action_slots(self, action_len: int, device: torch.device) -> None:
         if self.action_slot_mapping is not None and self._action_len == action_len:
             return
 
-        action_blocks = _ceil_div(action_len, self.block_size)
+        action_blocks = (action_len + self.block_size - 1) // self.block_size
         scratch_offset = 0 if self.commit_current else len(self.current_video_block_ids)
         self.action_scratch_block_ids = self.kv_cache.scratch_block_ids(
             self.is_negative,
             scratch_offset,
             action_blocks,
         )
-        self.action_slot_mapping = _slot_mapping_for_blocks(
+        positions = torch.arange(action_len, dtype=torch.long)
+        self.action_slot_mapping = compute_slot_mapping(
             self.action_scratch_block_ids,
-            action_len,
+            positions,
             self.block_size,
-            device,
-        )
+        ).to(device=device)
         self._action_len = action_len
 
     def video_block_table(self, device: torch.device) -> tuple[list[int], int]:

From e8b5b21d16f53b62bd73d6056f1410d2ecb373ec Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Fri, 3 Jul 2026 23:24:31 +0800
Subject: [PATCH 44/51] [AR-Diffusion] Make paged self-attention
 torch.compile-native (fused custom op)

PR #4863's paged path regressed the compiled stack +6.2% (CFG2 613->651
ms/chunk): the eager attention region graph-broke every DiT block
(~320x/forward), rebuilt identical block-table metadata per layer x step x
branch, and starved the GPU (idle 23.7%->33.5% with equal busy time; the
0.1-0.5 ms inter-kernel gap class grew 569->3,246).

- Register the fused write+attend as a custom op
  (vllm_omni::ar_diffusion_paged_write_attn): K/V slot writes + the
  FlashAttention block-table kernel execute inside one opaque graph node;
  pools come from a module-level registry (vLLM unified-attention pattern),
  never as graph inputs (the flat/paged pool views alias one storage).
  fa_version resolution is memoized (get_device_capability is not
  dynamo-traceable). DiT blocks return to fullgraph=True.
- Hoist allocation + metadata out of the compiled region:
  ForwardContext.prepare() runs once per branch forward in the uncompiled
  _forward_blocks (CFG-parallel lazy-allocation preserved) and builds the
  block table ONCE for all 40 layers, tail-padded to a fixed width so only
  tensor values change across window growth (no per-window-size recompiles;
  keeps future CUDA-graph capture possible).
- Pass per-layer state into the compiled block as a NamedTuple of tensors
  (ARDiffusionPagedLayerInputs). layer_idx is a 0-dim tensor, NOT an int:
  all blocks share one code object, and an int becomes a per-layer dynamo
  value guard -- 40 cache variants that blow the recompile limit.
- Demote per-step/per-commit hot-path INFO logs to debug.

Verified on H200: tests/ar_diffusion 68 passed (+8 GPU paged-vs-dense parity
combos force-run past the driver gate; op output bit-exact vs the direct
kernel call); zero graph breaks in block.forward, recompiles bounded to the
pre-existing current_start_frame/action-shape set covered by warm-up.
Steady per-chunk: TP1 1175.5 -> 1095.2 ms (gather baseline 1104.3);
CFG2 651.1 -> 611.1 ms (gather 613.1). Peak reserved 76.3 GiB (paged memory
win kept: -13% vs gather TP1). Gap census: 0.1-0.5 ms gaps 3,246/422 ms ->
421/58 ms, below the gather baseline.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/ar_diffusion/test_paged_attention.py    | 153 +++++++++--
 tests/ar_diffusion/test_pipeline_bridge.py    |  13 +
 .../models/dreamzero/causal_wan_model.py      | 133 ++++-----
 .../models/dreamzero/pipeline_dreamzero.py    |   4 +-
 .../ar_diffusion/kv_cache/__init__.py         |   6 +
 .../ar_diffusion/kv_cache/paged_attention.py  | 260 +++++++++++++++++-
 .../ar_diffusion/kv_cache/state.py            |   4 +-
 7 files changed, 463 insertions(+), 110 deletions(-)

diff --git a/tests/ar_diffusion/test_paged_attention.py b/tests/ar_diffusion/test_paged_attention.py
index 975732877de..37507064b9d 100644
--- a/tests/ar_diffusion/test_paged_attention.py
+++ b/tests/ar_diffusion/test_paged_attention.py
@@ -15,7 +15,9 @@
     ARDiffusionKVCache,
     ARDiffusionKVConfig,
     ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
     ar_diffusion_paged_attention,
+    paged_write_attn,
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
@@ -235,39 +237,47 @@ def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len,
         device=device,
     )
 
-    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
-    ctx.ensure_video_slots(device)
+    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
     current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
     current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
-    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
-
     action_k = action_v = None
     if action_len:
-        ctx.ensure_action_slots(action_len, device)
         action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
         action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
-        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
 
     query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
-        action_len=action_len,
-        query_len=query.shape[1],
-        device=device,
-    )
-    paged = ar_diffusion_paged_attention(
+
+    # The production path: once-per-forward host prep, then the fused
+    # write+attend custom op consuming the NamedTuple payload.
+    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
+    inputs = layer_ctx.to_layer_inputs()
+    assert isinstance(inputs, ARDiffusionPagedLayerInputs)
+    paged = paged_write_attn(
+        inputs,
+        query[0],
+        current_k[0],
+        current_v[0],
+        action_k[0] if action_len else None,
+        action_v[0] if action_len else None,
+        HEAD_DIM**-0.5,
+    ).unsqueeze(0)
+
+    # Direct python-fn call on the same (already written) pools must be
+    # bit-exact: identical kernel, identical inputs.
+    direct = ar_diffusion_paged_attention(
         query,
         kv.key_cache(0),
         kv.value_cache(0),
-        block_table=block_table,
-        query_start_loc=query_start_loc,
-        seq_lens=seq_lens,
-        max_query_len=max_query_len,
-        max_seq_len=max_seq_len,
+        block_table=ctx.block_table,
+        query_start_loc=ctx.query_start_loc,
+        seq_lens=ctx.seq_lens,
+        max_query_len=ctx.max_query_len,
+        max_seq_len=ctx.max_seq_len,
         softmax_scale=HEAD_DIM**-0.5,
         causal=False,
     )
+    assert torch.equal(paged, direct)
 
     new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
     new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
@@ -277,3 +287,108 @@ def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len,
     ref = _dense_attention(query, new_k, new_v)
 
     torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
+def test_block_table_padded_to_fixed_width():
+    """Shapes must be constant across window growth: only values change."""
+    device = torch.device("cpu")
+    kv, st = make_state(window_chunks=2)
+
+    ctx1 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    ctx2 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    # 1-block vs 2-block visible history: same table width, same max_seq_len.
+    assert ctx1.block_table.shape == ctx2.block_table.shape
+    assert ctx1.max_seq_len == ctx2.max_seq_len
+    expected_width = kv.spec.sliding_window // kv.block_size + 1
+    assert ctx1.block_table.shape == (1, expected_width)
+    # Real lengths live in seq_lens, not the padded table.
+    assert int(ctx1.seq_lens[0]) == BLOCK
+    assert int(ctx2.seq_lens[0]) == 2 * BLOCK
+
+
+def test_prepare_is_idempotent_and_layers_share_metadata():
+    device = torch.device("cpu")
+    _, st = make_state(num_layers=2)
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
+    fctx = contexts[0].forward_ctx
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    table = fctx.block_table
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    assert fctx.block_table is table  # memoized, not rebuilt
+
+    i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs()
+    # 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer
+    # value guards on the shared block code object.
+    assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0
+    assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1
+    # All layers share the same metadata tensor objects.
+    assert i0.block_table is i1.block_table
+    assert i0.seq_lens is i1.seq_lens
+    assert i0.video_slots is i1.video_slots
+
+
+def test_layer_inputs_before_prepare_raises():
+    _, st = make_state()
+    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0]
+    with pytest.raises(RuntimeError, match="before prepare"):
+        layer_ctx.to_layer_inputs()
+
+
+def test_custom_op_registration_idempotent():
+    import importlib
+    import sys
+
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+    mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"
+    saved = sys.modules.pop(mod)
+    try:
+        importlib.import_module(mod)  # re-registration must not raise
+    finally:
+        sys.modules[mod] = saved
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+
+
+def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
+    """The op must trace as one opaque node: fullgraph OK, and changed tensor
+    VALUES (new slots / block ids) must not trigger recompilation."""
+    import torch._dynamo
+
+    device = torch.device("cpu")
+    kv, st = make_state(num_layers=2, window_chunks=2)
+
+    def run_one_forward(commit):
+        contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit)
+        fctx = contexts[0].forward_ctx
+        fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+        q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        # Both layers through ONE compiled fn: layer_idx is a tensor, so a
+        # different layer must NOT recompile (all 40 DiT blocks share the
+        # block-forward code object in production).
+        for layer_ctx in contexts:
+            out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
+        st.commit_paged_context(False)
+        return out
+
+    torch._dynamo.reset()
+    from torch._dynamo.testing import CompileCounter
+
+    counter = CompileCounter()
+
+    def fn(inputs, q, k, v):
+        return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
+
+    compiled = torch.compile(fn, backend=counter, fullgraph=True)
+
+    run_one_forward(commit=True)   # history grows between calls ->
+    run_one_forward(commit=True)   # block-table VALUES change, shapes don't
+    run_one_forward(commit=False)
+
+    assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py
index 1ad90d4ba4e..4d34e293bb6 100644
--- a/tests/ar_diffusion/test_pipeline_bridge.py
+++ b/tests/ar_diffusion/test_pipeline_bridge.py
@@ -167,3 +167,16 @@ def test_get_kv_cache_requires_frame_aligned_seqlen():
     _, st = make_state()
     with pytest.raises(AssertionError, match="frame-aligned"):
         st.get_kv_caches(False, seq_len=BLOCK + 1, commit_current=True)
+
+
+def test_prepare_one_branch_allocates_nothing_for_the_other():
+    """CFG-parallel laziness: a rank prepares only the branch it runs."""
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    pos_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    neg_ctx = st.get_kv_caches(True, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+
+    pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK)
+
+    assert pos_ctx._allocated_video
+    assert not neg_ctx._allocated_video
+    assert kv.window_block_ids(st.neg) == []
diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
index 1e399935444..c0daf0c95f5 100644
--- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
+++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
@@ -37,6 +37,20 @@
     MultiEmbodimentActionEncoder,
 )
 
+# AR-Diffusion paged self-attention (in-tree experimental engine). Import at
+# module level so the isinstance check + custom-op call trace cleanly inside
+# the fullgraph-compiled DiT block (an import inside the traced region would
+# graph-break). The model still works without the engine: the payload type is
+# only ever constructed by the AR-Diffusion runner.
+try:
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+        ARDiffusionPagedLayerInputs,
+        paged_write_attn,
+    )
+except ImportError:  # pragma: no cover - experimental package always ships in-tree
+    ARDiffusionPagedLayerInputs = None
+    paged_write_attn = None
+
 # ── RoPE utilities ──────────────────────────────────────────────────
 
 
@@ -567,17 +581,25 @@ def forward(
             action_v = v[:, -action_register_length:]
             v = v[:, :-action_register_length]
 
-        if getattr(kv_cache, "is_ar_diffusion_paged_context", False):
-            x = self._forward_ar_diffusion_paged(
-                roped_query=roped_query,
-                roped_key=roped_key,
-                value=v,
-                roped_action_query=roped_action_query,
-                roped_action_key=roped_action_key,
-                action_value=action_v,
-                action_register_length=action_register_length,
-                kv_cache=kv_cache,
-            )
+        if ARDiffusionPagedLayerInputs is not None and isinstance(kv_cache, ARDiffusionPagedLayerInputs):
+            # Fused write+attend custom op: one opaque node in the compiled
+            # graph (slot writes + FlashAttention block-table kernel inside).
+            # Metadata tensors were prepared once per forward in _forward_blocks.
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_act, v_act = roped_action_key[0], action_v[0]
+            else:
+                q_cat = roped_query
+                k_act = v_act = None
+            x = paged_write_attn(
+                kv_cache,
+                q_cat[0],
+                roped_key[0],
+                v[0],
+                k_act,
+                v_act,
+                self.head_dim**-0.5,
+            ).unsqueeze(0)
         else:
             updated_k = kv_cache[0]
             updated_v = kv_cache[1]
@@ -602,73 +624,6 @@ def forward(
         x = self.o(x)
         return x, updated_kv_cache
 
-    def _forward_ar_diffusion_paged(
-        self,
-        *,
-        roped_query: torch.Tensor,
-        roped_key: torch.Tensor,
-        value: torch.Tensor,
-        roped_action_query: torch.Tensor | None,
-        roped_action_key: torch.Tensor | None,
-        action_value: torch.Tensor | None,
-        action_register_length: int | None,
-        kv_cache: Any,
-    ) -> torch.Tensor:
-        from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ar_diffusion_paged_attention
-
-        if roped_key.shape[0] != 1:
-            raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1")
-        if roped_key.shape[1] != kv_cache.seq_len:
-            raise RuntimeError(
-                "AR-Diffusion paged context seq_len="
-                f"{kv_cache.seq_len} but current video KV has {roped_key.shape[1]} tokens"
-            )
-
-        ctx = kv_cache.forward_ctx
-        layer_idx = kv_cache.layer_idx
-        device = roped_key.device
-        ctx.ensure_video_slots(device)
-
-        k_pool = ctx.kv_cache._k_pools[layer_idx]
-        v_pool = ctx.kv_cache._v_pools[layer_idx]
-        video_slots = ctx.current_video_slot_mapping
-        if video_slots is None:
-            raise RuntimeError("AR-Diffusion paged context did not prepare current video slots")
-        k_pool[video_slots] = roped_key[0].to(k_pool.dtype)
-        v_pool[video_slots] = value[0].to(v_pool.dtype)
-
-        action_len = int(action_register_length or 0)
-        if action_len > 0:
-            if roped_action_key is None or action_value is None or roped_action_query is None:
-                raise RuntimeError("AR-Diffusion paged action path missing action query/key/value tensors")
-            ctx.ensure_action_slots(action_len, device)
-            action_slots = ctx.action_slot_mapping
-            if action_slots is None:
-                raise RuntimeError("AR-Diffusion paged context did not prepare action slots")
-            k_pool[action_slots] = roped_action_key[0].to(k_pool.dtype)
-            v_pool[action_slots] = action_value[0].to(v_pool.dtype)
-            q_cat = torch.cat([roped_query, roped_action_query], dim=1)
-        else:
-            q_cat = roped_query
-
-        block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
-            action_len=action_len,
-            query_len=q_cat.shape[1],
-            device=device,
-        )
-        return ar_diffusion_paged_attention(
-            q_cat,
-            ctx.kv_cache.key_cache(layer_idx),
-            ctx.kv_cache.value_cache(layer_idx),
-            block_table=block_table,
-            query_start_loc=query_start_loc,
-            seq_lens=seq_lens,
-            max_query_len=max_query_len,
-            max_seq_len=max_seq_len,
-            softmax_scale=self.head_dim**-0.5,
-            causal=False,
-        )
-
 
 # ── Attention Block ─────────────────────────────────────────────────
 
@@ -1045,6 +1000,28 @@ def _forward_blocks(
             clip_embedding = self.img_emb(clip_feature)
             context = torch.cat([clip_embedding, context], dim=1)
 
+        # AR-Diffusion paged path: host-side prep ONCE per branch forward, outside
+        # the compiled blocks — allocate current video/action slots, build the
+        # padded block-table metadata all 40 layers share, and hand the compiled
+        # region plain tensors (ARDiffusionPagedLayerInputs) instead of a Python
+        # context object. Lazy-allocation contract preserved: only the branch this
+        # CFG-parallel rank executes reaches its _forward_blocks.
+        if kv_cache and getattr(kv_cache[0], "is_ar_diffusion_paged_context", False):
+            if B != 1:
+                raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1")
+            fctx = kv_cache[0].forward_ctx
+            if seq_len != fctx.seq_len:
+                raise RuntimeError(
+                    f"AR-Diffusion paged context seq_len={fctx.seq_len} "
+                    f"but current video KV has {seq_len} tokens"
+                )
+            fctx.prepare(
+                device=x.device,
+                action_len=int(action_register_length or 0),
+                query_len=int(x.shape[1]),
+            )
+            kv_cache = [c.to_layer_inputs() for c in kv_cache]
+
         updated_kv_caches: list[torch.Tensor | None] = []
         for block_index, block in enumerate(self.blocks):
             x, updated_kv_cache = block(
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index aeb2ff8edce..c9d3f7e2b48 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -435,7 +435,7 @@ def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor
         )
         if kwargs.get("update_kv_cache", False):
             is_neg = kwargs.get("is_negative", False)
-            logger.info(
+            logger.debug(
                 "AR-Diffusion pipeline predict_noise -> commit paged context: "
                 "is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
                 is_neg,
@@ -473,6 +473,8 @@ def setup_compile(self) -> None:
 
         compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
         # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
+        # The AR-Diffusion paged self-attention is a registered custom op, so the
+        # block stays fullgraph even on that path.
         dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
 
         logger.info(
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
index 41d3dddde09..ab29f6093ab 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
@@ -29,7 +29,10 @@
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
     ARDiffusionPagedForwardContext,
     ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
     ar_diffusion_paged_attention,
+    paged_write_attn,
+    set_current_paged_kv_cache,
 )
 
 __all__ = [
@@ -37,6 +40,7 @@
     "ARDiffusionKVConfig",
     "ARDiffusionPagedForwardContext",
     "ARDiffusionPagedLayerContext",
+    "ARDiffusionPagedLayerInputs",
     "ARDiffusionRequestAdapter",
     "ChunkWindowManager",
     "ChunkWindowSpec",
@@ -46,6 +50,8 @@
     "chunk_slot_mapping",
     "compute_num_blocks",
     "compute_slot_mapping",
+    "paged_write_attn",
     "pool_write_chunk",
     "resident_block_ids",
+    "set_current_paged_kv_cache",
 ]
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
index 01bfedf9508..1c757124aea 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
@@ -4,12 +4,61 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, ClassVar
+from typing import Any, ClassVar, NamedTuple
 
 import torch
 
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
 
+# Set by ARDiffusionPagedForwardContext.prepare() before each branch forward and
+# read by the fused write+attend custom op below. The pools are process-lifetime
+# allocations owned by one ARDiffusionKVCache per worker process (CFG-parallel /
+# TP ranks are separate processes), mirroring vLLM's forward-context pattern for
+# keeping multi-GiB cache tensors out of the compiled graph's inputs.
+_CURRENT_PAGED_KV_CACHE: Any = None
+
+
+def set_current_paged_kv_cache(kv_cache: Any) -> None:
+    global _CURRENT_PAGED_KV_CACHE
+    _CURRENT_PAGED_KV_CACHE = kv_cache
+
+
+_LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {}
+
+
+def _layer_idx_tensor(layer_idx: int) -> torch.Tensor:
+    t = _LAYER_IDX_TENSORS.get(layer_idx)
+    if t is None:
+        t = torch.tensor(layer_idx, dtype=torch.int64)
+        _LAYER_IDX_TENSORS[layer_idx] = t
+    return t
+
+
+class ARDiffusionPagedLayerInputs(NamedTuple):
+    """Compiled-region payload for one layer's paged self-attention.
+
+    A NamedTuple of plain tensors + ints so ``torch.compile`` treats every field
+    as a pytree graph input (no object-attribute guards, no recompiles when only
+    tensor *values* change). All layers of one branch forward share the same
+    metadata tensor objects, built once by ``prepare()``.
+
+    ``layer_idx`` is a 0-dim CPU tensor, NOT a python int: all 40 DiT blocks
+    share one compiled code object, and an int here becomes a per-layer dynamo
+    value guard (``layer_idx == k``) — 40 cache variants that blow the
+    recompile limit. A tensor input guards on shape/dtype only, so one graph
+    serves every layer.
+    """
+
+    layer_idx: torch.Tensor
+    seq_len: int
+    video_slots: torch.Tensor
+    action_slots: torch.Tensor
+    block_table: torch.Tensor
+    query_start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    max_query_len: int
+    max_seq_len: int
+
 
 @dataclass
 class ARDiffusionPagedForwardContext:
@@ -31,6 +80,13 @@ class ARDiffusionPagedForwardContext:
     _allocated_video: bool = False
     _committed: bool = False
     _action_len: int = 0
+    # Set once by prepare(); shared by all layers of the branch forward.
+    block_table: torch.Tensor | None = None
+    query_start_loc: torch.Tensor | None = None
+    seq_lens: torch.Tensor | None = None
+    max_query_len: int = 0
+    max_seq_len: int = 0
+    _prepared: bool = False
 
     @property
     def block_size(self) -> int:
@@ -116,7 +172,14 @@ def build_block_table(
         query_len: int,
         device: torch.device,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
-        """Build FlashAttention block-table metadata for one self-attn call."""
+        """Build FlashAttention block-table metadata for one self-attn call.
+
+        The block table is tail-padded to a fixed width and ``max_seq_len`` is a
+        constant upper bound, so across window growth only tensor *values* change
+        — tensor shapes and the int consts stay stable for ``torch.compile``
+        (and, later, CUDA-graph capture). The kernel only dereferences the first
+        ``ceil(seq_lens/block_size)`` entries, so padding is never read.
+        """
         video_blocks, video_len = self.video_block_table(device)
         self.ensure_action_slots(action_len, device)
         action_blocks = self.action_scratch_block_ids if action_len > 0 else []
@@ -124,12 +187,57 @@ def build_block_table(
         if not block_ids:
             raise RuntimeError("AR-Diffusion paged attention needs at least current video KV blocks")
 
+        # Fixed capacity: full visible video window + one action-capacity block.
+        action_capacity_blocks = max(1, (action_len + self.block_size - 1) // self.block_size)
+        width = max(self.max_video_tokens // self.block_size + action_capacity_blocks, len(block_ids))
+        padded = block_ids + [0] * (width - len(block_ids))
+
         self.query_len = int(query_len)
         self.kv_len = int(video_len + action_len)
-        block_table = torch.tensor([block_ids], dtype=torch.int32, device=device)
+        max_seq_len = int(self.max_video_tokens + action_capacity_blocks * self.block_size)
+        block_table = torch.tensor([padded], dtype=torch.int32, device=device)
         query_start_loc = torch.tensor([0, self.query_len], dtype=torch.int32, device=device)
         seq_lens = torch.tensor([self.kv_len], dtype=torch.int32, device=device)
-        return block_table, query_start_loc, seq_lens, self.query_len, self.kv_len
+        return block_table, query_start_loc, seq_lens, self.query_len, max_seq_len
+
+    def prepare(self, device: torch.device, action_len: int, query_len: int) -> None:
+        """Host-side, once-per-branch-forward setup (called OUTSIDE torch.compile).
+
+        Allocates the current video/action slots (still lazy: only the branch a
+        CFG-parallel rank actually runs reaches its ``_forward_blocks``), builds
+        the padded block-table metadata ONCE for all layers, and publishes the
+        pool registry for the fused custom op. The compiled per-layer code then
+        only consumes prebuilt tensors via ``ARDiffusionPagedLayerInputs``.
+        """
+        if getattr(self, "_prepared", False):
+            return
+        self.ensure_video_slots(device)
+        (
+            self.block_table,
+            self.query_start_loc,
+            self.seq_lens,
+            self.max_query_len,
+            self.max_seq_len,
+        ) = self.build_block_table(action_len=action_len, query_len=query_len, device=device)
+        if self.action_slot_mapping is None:
+            self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
+        set_current_paged_kv_cache(self.kv_cache)
+        self._prepared = True
+
+    def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs:
+        if not getattr(self, "_prepared", False):
+            raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()")
+        return ARDiffusionPagedLayerInputs(
+            layer_idx=_layer_idx_tensor(layer_idx),
+            seq_len=int(self.seq_len),
+            video_slots=self.current_video_slot_mapping,
+            action_slots=self.action_slot_mapping,
+            block_table=self.block_table,
+            query_start_loc=self.query_start_loc,
+            seq_lens=self.seq_lens,
+            max_query_len=int(self.max_query_len),
+            max_seq_len=int(self.max_seq_len),
+        )
 
     def mark_committed(self) -> None:
         self._committed = True
@@ -191,6 +299,10 @@ def kv_len(self) -> int:
     def commit_current(self) -> bool:
         return self.forward_ctx.commit_current
 
+    def to_layer_inputs(self) -> ARDiffusionPagedLayerInputs:
+        """Compiled-region payload; requires ``forward_ctx.prepare()`` first."""
+        return self.forward_ctx.layer_inputs(self.layer_idx)
+
 
 def is_ar_diffusion_paged_context(value: object) -> bool:
     return isinstance(value, ARDiffusionPagedLayerContext)
@@ -228,6 +340,24 @@ def _reference_paged_attention(
     return torch.cat(outs, dim=0)
 
 
+_FA_VERSION_BY_HEAD_SIZE: dict[int, int] = {}
+
+
+def _resolve_fa_version(head_size: int) -> int:
+    # get_flash_attn_version -> current_platform.get_device_capability() is not
+    # dynamo-traceable, and the answer is fixed per head size for the process.
+    version = _FA_VERSION_BY_HEAD_SIZE.get(head_size)
+    if version is None:
+        try:
+            from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
+
+            version = int(get_flash_attn_version(requires_alibi=False, head_size=head_size) or 2)
+        except Exception:
+            version = 2
+        _FA_VERSION_BY_HEAD_SIZE[head_size] = version
+    return version
+
+
 def ar_diffusion_paged_attention(
     query: torch.Tensor,
     key_cache: torch.Tensor,
@@ -267,12 +397,7 @@ def ar_diffusion_paged_attention(
     else:
         from vllm.vllm_flash_attn import flash_attn_varlen_func
 
-        try:
-            from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
-
-            fa_version = get_flash_attn_version(requires_alibi=False, head_size=query_flat.shape[-1])
-        except Exception:
-            fa_version = 2
+        fa_version = _resolve_fa_version(query_flat.shape[-1])
 
         out = torch.empty_like(query_flat)
         flash_attn_varlen_func(
@@ -293,3 +418,118 @@ def ar_diffusion_paged_attention(
     if batched:
         return out.reshape(query.shape)
     return out
+
+
+# ── Fused write+attend custom op (torch.compile-safe) ──────────────────────
+#
+# One opaque op per layer keeps the compiled DiT block fullgraph: dynamo treats
+# it as a single graph node (no eager island, no graph breaks), and the K/V slot
+# writes happen inside the op so write→read ordering with the block-table kernel
+# is internal — no hidden-mutation ordering hazards against a separate reader op
+# and no multi-GiB pool tensors as graph inputs (the flat and paged pool views
+# alias the same storage, which AOTAutograd handles poorly as inputs). Pattern
+# follows sage_attn3.py in-repo and vLLM's own unified attention ops.
+def _paged_write_attn_impl(
+    query: torch.Tensor,
+    k_curr: torch.Tensor,
+    v_curr: torch.Tensor,
+    k_act: torch.Tensor | None,
+    v_act: torch.Tensor | None,
+    layer_idx: torch.Tensor,
+    video_slots: torch.Tensor,
+    action_slots: torch.Tensor,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    max_query_len: int,
+    max_seq_len: int,
+    softmax_scale: float,
+) -> torch.Tensor:
+    kv = _CURRENT_PAGED_KV_CACHE
+    if kv is None:
+        raise RuntimeError("ar_diffusion_paged_write_attn called before prepare() set the KV pool registry")
+    layer_idx = int(layer_idx)
+    k_pool = kv._k_pools[layer_idx]
+    v_pool = kv._v_pools[layer_idx]
+    k_pool[video_slots] = k_curr.to(k_pool.dtype)
+    v_pool[video_slots] = v_curr.to(v_pool.dtype)
+    if k_act is not None and v_act is not None and k_act.shape[0] > 0:
+        k_pool[action_slots] = k_act.to(k_pool.dtype)
+        v_pool[action_slots] = v_act.to(v_pool.dtype)
+    return ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(layer_idx),
+        kv.value_cache(layer_idx),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=softmax_scale,
+        causal=False,
+    )
+
+
+# hasattr guard keeps registration idempotent across test re-imports that pop
+# the module from sys.modules (same as sage_attn3.py).
+if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
+
+    @torch.library.custom_op("vllm_omni::ar_diffusion_paged_write_attn", mutates_args=())
+    def _paged_write_attn_op(
+        query: torch.Tensor,
+        k_curr: torch.Tensor,
+        v_curr: torch.Tensor,
+        k_act: torch.Tensor | None,
+        v_act: torch.Tensor | None,
+        layer_idx: torch.Tensor,
+        video_slots: torch.Tensor,
+        action_slots: torch.Tensor,
+        block_table: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        max_query_len: int,
+        max_seq_len: int,
+        softmax_scale: float,
+    ) -> torch.Tensor:
+        return _paged_write_attn_impl(
+            query,
+            k_curr,
+            v_curr,
+            k_act,
+            v_act,
+            layer_idx,
+            video_slots,
+            action_slots,
+            block_table,
+            query_start_loc,
+            seq_lens,
+            max_query_len,
+            max_seq_len,
+            softmax_scale,
+        )
+
+    @_paged_write_attn_op.register_fake
+    def _(query, k_curr, v_curr, k_act, v_act, layer_idx, video_slots, action_slots,
+          block_table, query_start_loc, seq_lens, max_query_len, max_seq_len, softmax_scale):
+        return torch.empty_like(query)
+
+
+def paged_write_attn(inputs: ARDiffusionPagedLayerInputs, query, k_curr, v_curr, k_act, v_act,
+                     softmax_scale: float) -> torch.Tensor:
+    """Model-facing entry: routes through the custom op (traceable in fullgraph)."""
+    return torch.ops.vllm_omni.ar_diffusion_paged_write_attn(
+        query,
+        k_curr,
+        v_curr,
+        k_act,
+        v_act,
+        inputs.layer_idx,
+        inputs.video_slots,
+        inputs.action_slots,
+        inputs.block_table,
+        inputs.query_start_loc,
+        inputs.seq_lens,
+        inputs.max_query_len,
+        inputs.max_seq_len,
+        softmax_scale,
+    )
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
index d3f1216755b..a92bbe5b48d 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
@@ -85,7 +85,7 @@ def prepare_paged_context(
             max_video_tokens=int(self.kv_cache.spec.sliding_window),
         )
         self._paged_pending[is_negative] = forward_ctx
-        _log.info(
+        _log.debug(
             "AR-Diffusion GET   [%s] source=paged-attn layers=%d history_blocks=%d seq_len=%d commit_current=%s",
             branch,
             self.num_layers,
@@ -106,7 +106,7 @@ def commit_paged_context(self, is_negative: bool) -> None:
             for _ in range(n_chunks):
                 ctx.adapter.on_chunk_committed()
             self._committed[is_negative] += ctx.seq_len
-            _log.info(
+            _log.debug(
                 "AR-Diffusion COMMIT [%s] paged-attn new_tokens=%d chunks=%d resident=%d/%d",
                 branch,
                 ctx.seq_len,

From 469b85d33d22582e1eff35e999872ed727f96907 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 4 Jul 2026 01:28:48 +0800
Subject: [PATCH 45/51] [AR-Diffusion] Branch-aware KV pool sizing (-14 GiB/GPU
 under CFG-parallel)

The pool floor (2*(window_chunks+1)+2 = 22 blocks) and scratch reservation
(2 x 8 = 16) were hard-coded for two CFG branches sharing one pool, but under
CFG-parallel x2 each rank executes exactly one branch (rank0 pos, rank1 neg)
and the other branch's lazy contexts never allocate on that rank -- ~13 GiB
of the 25.6 GiB preallocation was idle capacity per rank. Measured scratch
high-water is 3 blocks/branch (2 in-flight video + 1 action) vs 8 reserved.

- ARDiffusionKVCache gains local_branches (1 under CFG-parallel, 2 for
  single-process runs) and num_frame_per_block; floor becomes
  local_branches * (window_chunks + num_frame_per_block) + 2 (the old
  per-branch "+1" under-counted the 2-frame in-flight chunk and leaned on
  headroom). Scratch: local_branches * per-branch, default per-branch 8 -> 4;
  a single-branch rank maps either CFG side to scratch slot 0.
- Runner derives local_branches from get_classifier_free_guidance_world_size()
  and threads the geometry through build_kv_cache.
- 4 new sizing tests (branch-aware floor, slot mapping, exhaustion, validation).

Verified on H200: tests/ar_diffusion 72 passed + 8 GPU parity combos; CFG2
scene1 8-chunk: 76.3 -> 62.2 GiB reserved/GPU (pool step 26.7 -> 12.6 GiB),
steady 614.5 ms (ref 611); TP1: 76.3 -> 72.7 GiB, steady 1099.5 ms (ref 1095);
window restarts cross cleanly, pool restored after warm-up. Latents PSNR
39.5 dB vs pre-change -- within the stack's own run-to-run envelope (an
identical-code rerun is also not bit-exact: fresh-process compile/cuBLAS/
FA-splitkv variation amplifies ulps over the autoregressive rollout).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/ar_diffusion/test_kv_cache.py           | 51 +++++++++++++++++++
 .../ar_diffusion/kv_cache/manager.py          | 33 +++++++++---
 vllm_omni/experimental/ar_diffusion/runner.py | 19 +++++++
 3 files changed, 95 insertions(+), 8 deletions(-)

diff --git a/tests/ar_diffusion/test_kv_cache.py b/tests/ar_diffusion/test_kv_cache.py
index 7698081014a..cbb9656e54e 100644
--- a/tests/ar_diffusion/test_kv_cache.py
+++ b/tests/ar_diffusion/test_kv_cache.py
@@ -209,3 +209,54 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
     expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
     assert kv.num_blocks == expected
     assert expected > 2 * (2 + 1) + 2  # above the min-blocks floor, so the deduction is what's tested
+
+
+def _make_kv(*, local_branches, num_frame_per_block=2, window_chunks=9):
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+        num_layers=1,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 16,  # tiny -> the floor binds
+        device=torch.device("cpu"),
+        local_branches=local_branches,
+        num_frame_per_block=num_frame_per_block,
+    )
+
+
+def test_pool_floor_is_branch_aware():
+    """CFG-parallel rank (one local branch) sizes for one window + in-flight chunk;
+    a single-process run (both branches) sizes for two. Scratch scales the same way."""
+    one = _make_kv(local_branches=1)
+    two = _make_kv(local_branches=2)
+
+    assert one.managed_num_blocks == 1 * (9 + 2) + 2  # 13
+    assert two.managed_num_blocks == 2 * (9 + 2) + 2  # 24
+    assert one.scratch_num_blocks == one.scratch_blocks_per_branch
+    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_branch
+    assert one.num_blocks_total == 13 + one.scratch_blocks_per_branch
+
+
+def test_scratch_maps_to_slot_zero_with_one_local_branch():
+    """A CFG-parallel rank runs exactly one branch: whichever CFG side it is,
+    its scratch lands in the rank's single slot (no dead second slot)."""
+    one = _make_kv(local_branches=1)
+    assert one.scratch_block_ids(True, 0, 2) == one.scratch_block_ids(False, 0, 2)
+
+    two = _make_kv(local_branches=2)
+    assert two.scratch_block_ids(True, 0, 2) != two.scratch_block_ids(False, 0, 2)
+
+
+def test_scratch_exhaustion_still_raises():
+    one = _make_kv(local_branches=1)
+    cap = one.scratch_blocks_per_branch
+    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
+        one.scratch_block_ids(False, 0, cap + 1)
+
+
+def test_invalid_local_branches_rejected():
+    with pytest.raises(ValueError, match="local_branches"):
+        _make_kv(local_branches=3)
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
index 28854328790..4b3b5d788d4 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
@@ -166,6 +166,8 @@ def __init__(
         cross_attn_length: int = 0,
         cross_attn_img_length: int = 0,
         device: torch.device | None = None,
+        local_branches: int = 2,
+        num_frame_per_block: int = 1,
     ) -> None:
         if not config.enable:
             raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig")
@@ -173,8 +175,17 @@ def __init__(
             raise ValueError("Phase 1 requires a bounded window (window_chunks)")
         if config.chunk_size <= 0:
             raise ValueError("ARDiffusionKVConfig.chunk_size must be set (> 0)")
+        if local_branches not in (1, 2):
+            raise ValueError(f"local_branches must be 1 or 2, got {local_branches}")
 
         self.config = config
+        # How many CFG branches allocate from THIS rank's pool. Under
+        # CFG-parallel x2 each rank executes exactly one branch (rank0 pos,
+        # rank1 neg) and the other branch's lazy contexts never allocate, so
+        # sizing for both would leave ~half the pool as idle capacity. A
+        # single-process run (TP1 / offline) executes both branches -> 2.
+        self.local_branches = local_branches
+        self.num_frame_per_block = max(1, int(num_frame_per_block))
         self.block_size = block_size
         self.num_layers = num_layers
         self.num_kv_heads = num_kv_heads
@@ -256,21 +267,23 @@ def _cross_pool_bytes(length: int) -> int:
             config.gpu_memory_fraction,
             self.spec.page_size_bytes * num_layers,
         )
-        # Floor: one forward needs the resident window plus the in-flight chunk for
-        # BOTH CFG branches (pos/neg are independent adapters allocating from the
-        # same pool), with a little eviction-transient headroom. The memory-fraction
+        # Floor: one forward needs the resident window plus the in-flight chunk
+        # (num_frame_per_block frame-blocks) for every branch THIS rank runs,
+        # with a little eviction-transient headroom. The memory-fraction
         # heuristic can under-size this once block_size grows — e.g. frame-granular
         # paging at the true frame_seqlen makes each block larger and the pool
         # fewer-blocks — so guarantee the minimum the rollout cannot run without,
         # otherwise allocate_chunk hits an exhausted pool mid-forward.
-        min_blocks = 2 * (config.window_chunks + 1) + 2
+        min_blocks = self.local_branches * (config.window_chunks + self.num_frame_per_block) + 2
         if num_blocks < min_blocks:
             _log.warning(
                 "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
-                "floor (2 CFG branches x (window_chunks=%d + 1) + 2 headroom)",
+                "floor (%d local CFG branch(es) x (window_chunks=%d + num_frame_per_block=%d) + 2 headroom)",
                 num_blocks,
                 min_blocks,
+                self.local_branches,
                 config.window_chunks,
+                self.num_frame_per_block,
             )
             num_blocks = min_blocks
         layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
@@ -280,11 +293,13 @@ def _cross_pool_bytes(length: int) -> int:
         # Scratch blocks are outside KVCacheManager ownership. They hold current
         # denoise-step video KV when update_kv_cache=False and action/state KV in
         # both modes; they are reused every forward and never committed.
-        scratch_per_branch = int(os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "8"))
+        # Default 4: measured high-water is num_frame_per_block video blocks (2)
+        # + 1 action block per branch per forward, +1 margin.
+        scratch_per_branch = int(os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "4"))
         if scratch_per_branch <= 0:
             raise ValueError("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH must be positive")
         self.scratch_blocks_per_branch = scratch_per_branch
-        self.scratch_num_blocks = 2 * scratch_per_branch
+        self.scratch_num_blocks = self.local_branches * scratch_per_branch
         self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks
         self.null_block_id = self.manager.block_pool.null_block.block_id
 
@@ -426,7 +441,9 @@ def scratch_block_ids(self, is_negative: bool, start: int, count: int) -> list[i
                 f"need [{start}, {start + count}) of {self.scratch_blocks_per_branch}. "
                 "Increase AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH."
             )
-        branch_offset = self.scratch_blocks_per_branch if is_negative else 0
+        # With one local branch (CFG-parallel x2) the rank's only branch maps to
+        # slot 0 whichever CFG side it is; the other branch never allocates here.
+        branch_offset = self.scratch_blocks_per_branch if (is_negative and self.local_branches == 2) else 0
         base = self.managed_num_blocks + branch_offset + start
         return list(range(base, base + count))
 
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index 09a259a0a0d..fb6aa8eee17 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -88,6 +88,8 @@ def build_kv_cache(
         available_bytes: int,
         cross_attn_length: int = 0,
         cross_attn_img_length: int = 0,
+        local_branches: int = 2,
+        num_frame_per_block: int = 1,
     ) -> None:
         """Construct the ARDiffusionKVCache (preallocating GPU pools on ``self.device``).
 
@@ -106,6 +108,8 @@ def build_kv_cache(
             available_bytes=available_bytes,
             cross_attn_length=cross_attn_length,
             cross_attn_img_length=cross_attn_img_length,
+            local_branches=local_branches,
+            num_frame_per_block=num_frame_per_block,
             device=self.device,
         )
         logger.info(
@@ -180,6 +184,19 @@ def _preallocate_kv_cache(self) -> None:
 
         self.ar_diffusion_kv_config = dataclasses.replace(self.ar_diffusion_kv_config, chunk_size=chunk_size, window_chunks=window_chunks)
         free_bytes = torch.cuda.mem_get_info(self.device)[0]
+        # Under CFG-parallel each rank executes exactly ONE branch (rank0 pos,
+        # rank1 neg; the other branch's lazy contexts never allocate on this
+        # rank), so its pool only needs one branch's capacity. A single-process
+        # run (cfg world 1) executes both branches from one pool.
+        try:
+            from vllm_omni.diffusion.distributed.parallel_state import (
+                get_classifier_free_guidance_world_size,
+            )
+
+            cfg_world = int(get_classifier_free_guidance_world_size())
+        except Exception:
+            cfg_world = 1
+        local_branches = 1 if cfg_world >= 2 else 2
         logger.info(
             "AR-Diffusion preallocating (paged): frame_seqlen=%d num_frame_per_block=%d "
             "local_attn_size=%d -> chunk_size=%d window_chunks=%d (window=%d tokens)",
@@ -200,6 +217,8 @@ def _preallocate_kv_cache(self) -> None:
             available_bytes=free_bytes,
             cross_attn_length=cross_attn_length,
             cross_attn_img_length=cross_attn_img_length,
+            local_branches=local_branches,
+            num_frame_per_block=num_frame_per_block,
         )
 
     def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:

From ef866cd8cf7c8170bddf53ec0510579212a20a92 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 4 Jul 2026 12:48:51 +0800
Subject: [PATCH 46/51] [AR-Diffusion] Fail fast on unsupported configs (PR
 #4534 review round 2)

Address hsliuustc0106's inline review:
- DreamZeroPipeline raises at init when engine_backend is not the
  AR-Diffusion engine (previously a stale/programmatic 'default' config
  only crashed mid-forward on the first KV access).
- ARDiffusionModelRunner rejects max_num_seqs > 1 at KV-pool init: the
  paged self-attention path is single-sequence (writes batch index 0).
- Document the pos/neg CFG-branch adapter semantics on ARDiffusionKVState.

+2 unit tests (74 passing).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/ar_diffusion/test_engine_routing.py     | 29 +++++++++++++++++++
 .../models/dreamzero/pipeline_dreamzero.py    | 12 ++++++++
 .../ar_diffusion/kv_cache/state.py            |  9 ++++++
 vllm_omni/experimental/ar_diffusion/runner.py | 11 +++++++
 4 files changed, 61 insertions(+)

diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py
index 5254eb1a3fe..a43fdc65de1 100644
--- a/tests/ar_diffusion/test_engine_routing.py
+++ b/tests/ar_diffusion/test_engine_routing.py
@@ -92,3 +92,32 @@ def test_worker_runner_selection_prefers_override():
     assert (getattr(over, "diffusion_model_runner_cls", None) or "PLATFORM") == AR_DIFFUSION_MODEL_RUNNER_CLS
     unset = SimpleNamespace(diffusion_model_runner_cls=None)
     assert (getattr(unset, "diffusion_model_runner_cls", None) or "PLATFORM") == "PLATFORM"
+
+
+def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine():
+    """Review (hsliuustc0106): a stale config with engine_backend='default'
+    must fail at init, not mid-forward on the first KV access."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace(engine_backend="default"))
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace())
+
+
+def test_runner_rejects_multi_sequence_config():
+    """Review (hsliuustc0106): the paged path writes batch index 0 only, so
+    max_num_seqs > 1 must be rejected at init."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
+    with pytest.raises(ValueError, match="max_num_seqs=1"):
+        ARDiffusionModelRunner._preallocate_kv_cache(fake)
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index c9d3f7e2b48..3cd97332944 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -204,6 +204,18 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """
         super().__init__()
 
+        # DreamZero is engine-only: every KV access in forward() routes through
+        # the AR-Diffusion engine's pool-backed state. Fail fast here — a stale
+        # or programmatic config that leaves engine_backend="default" would
+        # otherwise only crash mid-forward on the first KV access.
+        engine_backend = str(getattr(od_config, "engine_backend", "") or "")
+        if "ar_diffusion" not in engine_backend.lower().replace("-", "_"):
+            raise ValueError(
+                "DreamZeroPipeline requires the AR-Diffusion engine; set "
+                "engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine "
+                f"in the deploy config (got engine_backend={engine_backend!r})."
+            )
+
         model_path = od_config.model
         model_config = od_config.model_config
         local_files_only = os.path.exists(model_path)
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
index a92bbe5b48d..9e9c0d877dc 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
@@ -23,6 +23,15 @@ class ARDiffusionKVState:
     """
 
     def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
+        """One state per DreamZero session, holding BOTH CFG branches' adapters.
+
+        ``pos``/``neg`` are the classifier-free-guidance branches: ``pos`` is
+        the conditional pass (real prompt), ``neg`` the unconditional pass
+        (negative prompt). Each owns an independent request adapter — and
+        therefore independent pool blocks — because the two branches produce
+        different K/V for the same frames. Every ``is_negative`` flag in this
+        module selects between them (``False`` -> ``pos``, ``True`` -> ``neg``).
+        """
         self.kv_cache = kv_cache
         self.pos = pos_adapter
         self.neg = neg_adapter
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index fb6aa8eee17..5bdbfa99b9c 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -147,6 +147,17 @@ def _infer_frame_seqlen(self) -> int:
 
     def _preallocate_kv_cache(self) -> None:
         """Build the KV pool once, from the loaded DreamZero transformer geometry."""
+        # The paged self-attention path carries one rollout per request: slot
+        # mappings/adapters are single-sequence and K/V writes address batch
+        # index 0 only. Reject a multi-sequence config at init instead of
+        # failing (or silently dropping sequences) mid-forward.
+        max_num_seqs = int(getattr(self.od_config, "max_num_seqs", 1) or 1)
+        if max_num_seqs > 1:
+            raise ValueError(
+                "AR-Diffusion paged KV supports max_num_seqs=1 (single-sequence "
+                f"rollouts); got max_num_seqs={max_num_seqs}. Per-sequence "
+                "adapters/slot maps are not implemented."
+            )
         t = self.pipeline.transformer
         num_layers = int(t.num_layers)
         num_kv_heads = int(getattr(t.blocks[0].self_attn, "tp_num_heads", t.num_heads))

From be4727e1e1ce8292b23ded5a26919af4c7e5d38a Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 4 Jul 2026 12:56:28 +0800
Subject: [PATCH 47/51] [AR-Diffusion] Tear down the session on a failed
 forward (review follow-up)

A forward that dies partway may have written some layers' K/V into
allocated-but-uncommitted blocks and advanced model-local state; previously
the session stayed poisoned until an explicit reset (the next request hit the
pending-commit guard). execute_model's exception path now frees the session's
pool blocks (ARDiffusionKVState.close()) and drops the model-local state, so
the next request with the same session id starts clean. Completes the
transactional-commit story from hsliuustc0106's state.py review comment.

+1 unit test (75 passing).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/ar_diffusion/test_pipeline_bridge.py    | 40 +++++++++++++++++++
 vllm_omni/experimental/ar_diffusion/runner.py | 19 +++++++++
 2 files changed, 59 insertions(+)

diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py
index 4d34e293bb6..83c710efa8e 100644
--- a/tests/ar_diffusion/test_pipeline_bridge.py
+++ b/tests/ar_diffusion/test_pipeline_bridge.py
@@ -180,3 +180,43 @@ def test_prepare_one_branch_allocates_nothing_for_the_other():
     assert pos_ctx._allocated_video
     assert not neg_ctx._allocated_video
     assert kv.window_block_ids(st.neg) == []
+
+
+def test_failed_forward_tears_down_session(monkeypatch):
+    """Review follow-up (hsliuustc0106): a forward that dies partway must not
+    leave allocated-but-uncommitted KV behind — the session is torn down (pool
+    blocks freed, model-local state dropped) and the error propagates."""
+    from types import SimpleNamespace
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)  # session owns pool blocks
+    free_before_failure = kv.manager.block_pool.get_num_free_blocks()
+
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.kv_cache = kv
+    runner._ar_diffusion_states = __import__("collections").OrderedDict({"s1": st})
+    runner._max_ar_diffusion_states = 4
+    runner.pipeline = SimpleNamespace(_states={"s1": object()}, _ar_diffusion_kv_state=None)
+    runner.device = None
+    runner._perf_e2e_times = []
+
+    def boom(self, req):
+        raise RuntimeError("layer 17 exploded")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
+    req = SimpleNamespace(
+        request_id="r0",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "s1"}),
+    )
+
+    with pytest.raises(RuntimeError, match="layer 17 exploded"):
+        ARDiffusionModelRunner.execute_model(runner, req)
+
+    assert "s1" not in runner._ar_diffusion_states
+    assert "s1" not in runner.pipeline._states
+    assert runner.pipeline._ar_diffusion_kv_state is None
+    # close() freed the session's resident blocks.
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before_failure
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index 5bdbfa99b9c..4b1427bd236 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -272,6 +272,25 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
         _e2e_t0 = time.perf_counter()
         try:
             out = super().execute_model(req)
+        except Exception:
+            # Transactional containment: a forward that died partway may have
+            # written some layers' K/V into allocated-but-uncommitted blocks
+            # and advanced model-local state. Tear the whole session down
+            # (engine pool blocks freed via close(); model-local state dropped)
+            # so the next request with this session id starts clean instead of
+            # tripping the pending-commit guard or reading half-written KV.
+            broken = self._ar_diffusion_states.pop(session_id, None)
+            if broken is not None:
+                broken.close()
+            states = getattr(self.pipeline, "_states", None)
+            if states is not None:
+                states.pop(session_id, None)
+            logger.warning(
+                "AR-Diffusion execute_model failed for session=%s; session state torn down "
+                "(pool blocks freed) — the next request starts a fresh session",
+                session_id,
+            )
+            raise
         finally:
             self.pipeline._ar_diffusion_kv_state = None
         if self.device is not None and torch.cuda.is_available():

From b07269e98a8aff63cf6f749d663afba2ce3fce47 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 4 Jul 2026 13:11:57 +0800
Subject: [PATCH 48/51] [AR-Diffusion] Consolidate engine->runner routing;
 widen engine_backend contract

Address hsliuustc0106's review round 3:

- Routing decision now lives in ONE place: the engine class declares its
  runner via DiffusionEngine.default_diffusion_model_runner_cls (ARDiffusionEngine
  sets ARDiffusionModelRunner). The worker resolves a single precedence chain
  (explicit od_config.diffusion_model_runner_cls override > engine-declared
  runner > platform default) through resolve_engine_class. Engines no longer
  mutate od_config at init; apply_ar_diffusion_runner_default is gone.
- engine_backend annotation widened to `str | type` to match what
  resolve_engine_class actually accepts; docstring pins the config-file
  contract to strings (structured-config mirrors expose the string form).

Tests updated: engine-class declaration + the full selection chain (74 passing).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/ar_diffusion/test_engine_routing.py     | 48 ++++++++++++-------
 vllm_omni/diffusion/data.py                   | 19 +++++---
 vllm_omni/diffusion/diffusion_engine.py       |  8 ++++
 .../diffusion/worker/diffusion_worker.py      | 35 ++++++++++----
 vllm_omni/experimental/ar_diffusion/engine.py | 27 ++---------
 5 files changed, 79 insertions(+), 58 deletions(-)

diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py
index a43fdc65de1..fa798670a84 100644
--- a/tests/ar_diffusion/test_engine_routing.py
+++ b/tests/ar_diffusion/test_engine_routing.py
@@ -19,7 +19,6 @@
 from vllm_omni.experimental.ar_diffusion.engine import (
     AR_DIFFUSION_MODEL_RUNNER_CLS,
     ARDiffusionEngine,
-    apply_ar_diffusion_runner_default,
 )
 
 
@@ -75,23 +74,36 @@ def test_ar_diffusion_model_runner_cls_resolves():
     assert cls.__name__ == "ARDiffusionModelRunner"
 
 
-def test_apply_runner_default_sets_when_unset():
-    cfg = SimpleNamespace(diffusion_model_runner_cls=None)
-    apply_ar_diffusion_runner_default(cfg)
-    assert cfg.diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
-
-
-def test_apply_runner_default_respects_explicit_choice():
-    cfg = SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner")
-    apply_ar_diffusion_runner_default(cfg)
-    assert cfg.diffusion_model_runner_cls == "my.custom.Runner"
-
-
-def test_worker_runner_selection_prefers_override():
-    over = SimpleNamespace(diffusion_model_runner_cls=AR_DIFFUSION_MODEL_RUNNER_CLS)
-    assert (getattr(over, "diffusion_model_runner_cls", None) or "PLATFORM") == AR_DIFFUSION_MODEL_RUNNER_CLS
-    unset = SimpleNamespace(diffusion_model_runner_cls=None)
-    assert (getattr(unset, "diffusion_model_runner_cls", None) or "PLATFORM") == "PLATFORM"
+def test_engine_class_declares_its_runner():
+    """Routing lives on the engine class (review: hsliuustc0106) — engines
+    never mutate od_config; the worker resolves the runner off the class."""
+    assert ARDiffusionEngine.default_diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
+    assert DiffusionEngine.default_diffusion_model_runner_cls is None
+
+
+def _select_runner(od_config, platform_default="PLATFORM"):
+    """Mirror of the worker's runner-selection chain (override > engine > platform)."""
+    runner_override = getattr(od_config, "diffusion_model_runner_cls", None)
+    if isinstance(runner_override, str) and runner_override:
+        return runner_override
+    engine_cls = DiffusionEngine.resolve_engine_class(od_config)
+    engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+    if isinstance(engine_runner, str) and engine_runner:
+        return engine_runner
+    return platform_default
+
+
+def test_runner_selection_chain():
+    ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+    # engine-declared runner wins over platform default
+    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar)) == AR_DIFFUSION_MODEL_RUNNER_CLS
+    # explicit override wins over the engine's declaration
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar))
+        == "my.custom.Runner"
+    )
+    # default engine -> platform default
+    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend="default")) == "PLATFORM"
 
 
 def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine():
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index d08b4aa73ff..54b850fb921 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -538,13 +538,18 @@ class OmniDiffusionConfig:
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None
 
-    # Engine backend selection, resolved generically by ``DiffusionEngine.make_engine``
-    # (mirrors ``DiffusionExecutor.get_class``): "default" -> DiffusionEngine, a
-    # DiffusionEngine subclass, or an import-path string (set e.g. by a deploy config).
-    engine_backend: str = "default"
-
-    # Optional override for the diffusion model runner class (import path). The worker
-    # uses it instead of the platform default when set; unset -> platform default.
+    # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class``
+    # (mirrors ``DiffusionExecutor.get_class``). Config files use a string:
+    # "default" -> DiffusionEngine, or an import path (set e.g. by a deploy
+    # config). Programmatic callers may pass a DiffusionEngine subclass
+    # directly — hence ``str | type`` (structured-config mirrors should expose
+    # the string form only).
+    engine_backend: str | type = "default"
+
+    # Optional override for the diffusion model runner class (import path).
+    # Precedence in the worker: this override > the runner declared by the
+    # selected engine class (``default_diffusion_model_runner_cls``) > the
+    # platform default. Never mutated by engines.
     diffusion_model_runner_cls: str | None = None
 
     # HuggingFace specific parameters
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index f80fb624ad8..8ce24bd4f1d 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -102,6 +102,14 @@ class _RpcTask:
 class DiffusionEngine:
     """The diffusion engine for vLLM-Omni diffusion models."""
 
+    #: Import path of the model runner this engine's workers should build, or
+    #: ``None`` for the platform default. Subclasses declare their runner here
+    #: (e.g. the AR-Diffusion engine), the worker resolves it through
+    #: :meth:`resolve_engine_class` — so engine->runner routing lives on the
+    #: engine class itself, and ``od_config.diffusion_model_runner_cls``
+    #: remains a pure explicit user override (never mutated by engines).
+    default_diffusion_model_runner_cls: str | None = None
+
     def __init__(
         self,
         od_config: OmniDiffusionConfig,
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index 2bd1c9b50c6..a9f65bd31cf 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -209,17 +209,32 @@ def __init__(
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
-        # Create model runner. An od_config override (e.g. the AR-Diffusion engine selecting
-        # ARDiffusionModelRunner) takes precedence; unset falls back to the platform default,
-        # so existing models are unaffected. The override must be an explicit import-
-        # path string — guard with isinstance so a non-string (e.g. a Mock od_config
-        # in tests) doesn't shadow the platform hook.
+        # Create model runner — one decision chain, in precedence order:
+        #   1. explicit od_config.diffusion_model_runner_cls (user override),
+        #   2. the runner declared by the engine class that engine_backend
+        #      selects (e.g. ARDiffusionEngine -> ARDiffusionModelRunner),
+        #   3. the platform default.
+        # Routing policy therefore lives on the engine class / config surface;
+        # engines never mutate od_config. Overrides must be import-path
+        # strings — guard with isinstance so a non-string (e.g. a Mock
+        # od_config in tests) doesn't shadow the platform hook.
         runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None)
-        model_runner_cls_path = (
-            runner_override
-            if isinstance(runner_override, str) and runner_override
-            else current_omni_platform.get_diffusion_model_runner_cls()
-        )
+        engine_runner = None
+        if not (isinstance(runner_override, str) and runner_override):
+            try:
+                from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+                engine_cls = DiffusionEngine.resolve_engine_class(self.od_config)
+                engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+            except Exception:
+                logger.warning("Worker %s: engine_backend resolution failed; using platform runner", self.rank)
+                engine_runner = None
+        if isinstance(runner_override, str) and runner_override:
+            model_runner_cls_path = runner_override
+        elif isinstance(engine_runner, str) and engine_runner:
+            model_runner_cls_path = engine_runner
+        else:
+            model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls()
         model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path)
         self.model_runner = model_runner_cls(
             vllm_config=self.vllm_config,
diff --git a/vllm_omni/experimental/ar_diffusion/engine.py b/vllm_omni/experimental/ar_diffusion/engine.py
index f8d09967ef3..6fb6ee7c940 100644
--- a/vllm_omni/experimental/ar_diffusion/engine.py
+++ b/vllm_omni/experimental/ar_diffusion/engine.py
@@ -5,9 +5,7 @@
 
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
-from vllm_omni.diffusion.sched import SchedulerInterface
 
 logger = init_logger(__name__)
 
@@ -15,18 +13,6 @@
 AR_DIFFUSION_MODEL_RUNNER_CLS = "vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner"
 
 
-def apply_ar_diffusion_runner_default(od_config: OmniDiffusionConfig) -> None:
-    """Route the AR-Diffusion engine's workers to ``ARDiffusionModelRunner``.
-
-    Sets ``od_config.diffusion_model_runner_cls`` in place, unless the caller
-    already chose a runner. The worker (in its own process) reads this off the
-    propagated ``od_config`` and builds the AR-Diffusion runner instead of the platform
-    default, keeping the swap scoped to AR-Diffusion-routed models.
-    """
-    if getattr(od_config, "diffusion_model_runner_cls", None) is None:
-        od_config.diffusion_model_runner_cls = AR_DIFFUSION_MODEL_RUNNER_CLS
-
-
 class ARDiffusionEngine(DiffusionEngine):
     """AR-Diffusion engine with engine-level KV cache management.
 
@@ -49,12 +35,7 @@ class ARDiffusionEngine(DiffusionEngine):
     so DreamZero's rollout runs against the runner-owned KV cache.
     """
 
-    def __init__(
-        self,
-        od_config: OmniDiffusionConfig,
-        scheduler: SchedulerInterface | None = None,
-    ) -> None:
-        # Route this engine's workers to ARDiffusionModelRunner before the base __init__
-        # builds the executor and spawns workers (od_config is propagated to them).
-        apply_ar_diffusion_runner_default(od_config)
-        super().__init__(od_config, scheduler=scheduler)
+    #: Workers of this engine build the AR-Diffusion runner (unless the config
+    #: sets an explicit ``diffusion_model_runner_cls`` override). Declared on
+    #: the class so routing needs no od_config mutation at engine init.
+    default_diffusion_model_runner_cls = AR_DIFFUSION_MODEL_RUNNER_CLS

From 3740e9159239f09ec0c5dd6786e4f46a6d97fef0 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 4 Jul 2026 14:35:52 +0800
Subject: [PATCH 49/51] [AR-Diffusion] Fix pre-commit: TID251
 accelerator.synchronize, UP031, E501, ruff format

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/ar_diffusion/test_engine_routing.py     |  9 +++---
 tests/ar_diffusion/test_paged_attention.py    |  4 +--
 .../models/dreamzero/causal_wan_model.py      |  3 +-
 .../models/dreamzero/pipeline_dreamzero.py    |  4 +--
 .../ar_diffusion/kv_cache/paged_attention.py  | 23 ++++++++++++---
 .../ar_diffusion/kv_cache/state.py            | 12 +++-----
 vllm_omni/experimental/ar_diffusion/runner.py | 29 +++++++++----------
 7 files changed, 46 insertions(+), 38 deletions(-)

diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py
index fa798670a84..fc59e54ac64 100644
--- a/tests/ar_diffusion/test_engine_routing.py
+++ b/tests/ar_diffusion/test_engine_routing.py
@@ -36,9 +36,7 @@ def test_subclass_type_is_returned():
 
 def test_ar_diffusion_qualname_resolves():
     # How DreamZero's deploy config selects the engine: a full import-path string.
-    cls = DiffusionEngine.resolve_engine_class(
-        _cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine")
-    )
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
     assert cls is ARDiffusionEngine
 
 
@@ -96,7 +94,10 @@ def _select_runner(od_config, platform_default="PLATFORM"):
 def test_runner_selection_chain():
     ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
     # engine-declared runner wins over platform default
-    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar)) == AR_DIFFUSION_MODEL_RUNNER_CLS
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar))
+        == AR_DIFFUSION_MODEL_RUNNER_CLS
+    )
     # explicit override wins over the engine's declaration
     assert (
         _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar))
diff --git a/tests/ar_diffusion/test_paged_attention.py b/tests/ar_diffusion/test_paged_attention.py
index 37507064b9d..1836656f940 100644
--- a/tests/ar_diffusion/test_paged_attention.py
+++ b/tests/ar_diffusion/test_paged_attention.py
@@ -387,8 +387,8 @@ def fn(inputs, q, k, v):
 
     compiled = torch.compile(fn, backend=counter, fullgraph=True)
 
-    run_one_forward(commit=True)   # history grows between calls ->
-    run_one_forward(commit=True)   # block-table VALUES change, shapes don't
+    run_one_forward(commit=True)  # history grows between calls ->
+    run_one_forward(commit=True)  # block-table VALUES change, shapes don't
     run_one_forward(commit=False)
 
     assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
index c0daf0c95f5..ca92538918f 100644
--- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
+++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
@@ -1012,8 +1012,7 @@ def _forward_blocks(
             fctx = kv_cache[0].forward_ctx
             if seq_len != fctx.seq_len:
                 raise RuntimeError(
-                    f"AR-Diffusion paged context seq_len={fctx.seq_len} "
-                    f"but current video KV has {seq_len} tokens"
+                    f"AR-Diffusion paged context seq_len={fctx.seq_len} but current video KV has {seq_len} tokens"
                 )
             fctx.prepare(
                 device=x.device,
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index f61ebdbd3d7..107638f8179 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -1267,12 +1267,12 @@ def forward(self, req: DiffusionRequestBatch, **kwargs) -> DiffusionOutput:
         if os.environ.get("DZ_PHASE_TIMING"):
             import time as _time
 
-            torch.cuda.synchronize()
+            torch.accelerator.synchronize()
             _pt = {"time": _time, "t0": _time.perf_counter(), "marks": []}
 
         def _pt_mark(name: str) -> None:
             if _pt is not None:
-                torch.cuda.synchronize()
+                torch.accelerator.synchronize()
                 _pt["marks"].append((name, _pt["time"].perf_counter()))
 
         # Prompt embeds are constant within a session (a prompt change triggers a
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
index 1c757124aea..2038e496d31 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
@@ -509,13 +509,28 @@ def _paged_write_attn_op(
         )
 
     @_paged_write_attn_op.register_fake
-    def _(query, k_curr, v_curr, k_act, v_act, layer_idx, video_slots, action_slots,
-          block_table, query_start_loc, seq_lens, max_query_len, max_seq_len, softmax_scale):
+    def _(
+        query,
+        k_curr,
+        v_curr,
+        k_act,
+        v_act,
+        layer_idx,
+        video_slots,
+        action_slots,
+        block_table,
+        query_start_loc,
+        seq_lens,
+        max_query_len,
+        max_seq_len,
+        softmax_scale,
+    ):
         return torch.empty_like(query)
 
 
-def paged_write_attn(inputs: ARDiffusionPagedLayerInputs, query, k_curr, v_curr, k_act, v_act,
-                     softmax_scale: float) -> torch.Tensor:
+def paged_write_attn(
+    inputs: ARDiffusionPagedLayerInputs, query, k_curr, v_curr, k_act, v_act, softmax_scale: float
+) -> torch.Tensor:
     """Model-facing entry: routes through the custom op (traceable in fullgraph)."""
     return torch.ops.vllm_omni.ar_diffusion_paged_write_attn(
         query,
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
index 9e9c0d877dc..510bbe043a0 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
@@ -137,14 +137,10 @@ def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
         img_ok = not self.kv_cache._cross_k_img or self._cross_img_populated[is_negative]
         if not (self._cross_text_populated[is_negative] and img_ok and self.kv_cache.cross_attn_length > 0):
             raise RuntimeError(
-                "AR-Diffusion cross-attn read before _kv_populate_cross (neg=%s, cross_attn_length=%d, "
-                "text=%s img=%s) — the engine must own all cross KV"
-                % (
-                    is_negative,
-                    self.kv_cache.cross_attn_length,
-                    self._cross_text_populated[is_negative],
-                    self._cross_img_populated[is_negative],
-                )
+                f"AR-Diffusion cross-attn read before _kv_populate_cross (neg={is_negative}, "
+                f"cross_attn_length={self.kv_cache.cross_attn_length}, "
+                f"text={self._cross_text_populated[is_negative]} "
+                f"img={self._cross_img_populated[is_negative]}) — the engine must own all cross KV"
             )
         return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
 
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index 4b1427bd236..c8b8e93e1d8 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -27,8 +27,8 @@
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
-from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
 logger = init_logger(__name__)
 
@@ -193,7 +193,9 @@ def _preallocate_kv_cache(self) -> None:
         chunk_size = frame_seqlen
         window_chunks = self.ar_diffusion_kv_config.window_chunks or (max_attention_size // frame_seqlen)
 
-        self.ar_diffusion_kv_config = dataclasses.replace(self.ar_diffusion_kv_config, chunk_size=chunk_size, window_chunks=window_chunks)
+        self.ar_diffusion_kv_config = dataclasses.replace(
+            self.ar_diffusion_kv_config, chunk_size=chunk_size, window_chunks=window_chunks
+        )
         free_bytes = torch.cuda.mem_get_info(self.device)[0]
         # Under CFG-parallel each rank executes exactly ONE branch (rank0 pos,
         # rank1 neg; the other branch's lazy contexts never allocate on this
@@ -294,7 +296,7 @@ def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
         finally:
             self.pipeline._ar_diffusion_kv_state = None
         if self.device is not None and torch.cuda.is_available():
-            torch.cuda.synchronize(self.device)
+            torch.accelerator.synchronize(self.device)
         self._perf_e2e_times.append(time.perf_counter() - _e2e_t0)
         return out
 
@@ -310,11 +312,7 @@ def _synth_robot_obs(self, h: int, w: int, n_frames: int) -> dict:
         cameras, 7/6/1 state vectors) so it survives ``_transform_robot_obs`` exactly
         like a real request. One frame -> ``(H,W,3)``; chunk -> ``(n_frames,H,W,3)``.
         """
-        img = (
-            np.zeros((h, w, 3), dtype=np.uint8)
-            if n_frames == 1
-            else np.zeros((n_frames, h, w, 3), dtype=np.uint8)
-        )
+        img = np.zeros((h, w, 3), dtype=np.uint8) if n_frames == 1 else np.zeros((n_frames, h, w, 3), dtype=np.uint8)
         return {
             "observation/exterior_image_0_left": img,
             "observation/exterior_image_1_left": img,
@@ -356,18 +354,16 @@ def _warmup_ar_rollout(self) -> None:
         # every window shape is captured regardless of how many later steps are skipped.
         logger.info(
             "AR-Diffusion cudagraph warm-up: up to %d synthetic forwards (window_chunks=%d, nfpb=%d)",
-            n_forwards, window_chunks, nfpb,
+            n_forwards,
+            window_chunks,
+            nfpb,
         )
         try:
             for i in range(n_forwards):
                 n_frames = 1 if i == 0 else 4  # client convention: 1-frame prefill, 4-frame chunks
                 obs = self._synth_robot_obs(h, w, n_frames)
-                sp = OmniDiffusionSamplingParams(
-                    extra_args={"reset": i == 0, "session_id": sid, "robot_obs": obs}
-                )
-                req = OmniDiffusionRequest(
-                    prompts=["warmup"], sampling_params=sp, request_id=f"ardiffusion-warmup-{i}"
-                )
+                sp = OmniDiffusionSamplingParams(extra_args={"reset": i == 0, "session_id": sid, "robot_obs": obs})
+                req = OmniDiffusionRequest(prompts=["warmup"], sampling_params=sp, request_id=f"ardiffusion-warmup-{i}")
                 self.execute_model(req)
                 # Stop once the resident window is full — the remaining shapes are
                 # already captured (the window caps/resets through the same set).
@@ -398,7 +394,8 @@ def _warmup_ar_rollout(self) -> None:
             if free_after != free_before:
                 logger.warning(
                     "AR-Diffusion warm-up: KV pool not restored (free %d -> %d) — possible leak",
-                    free_before, free_after,
+                    free_before,
+                    free_after,
                 )
             else:
                 logger.info(

From 51d70304dd0f4bbf98dbf6b45e299895e49f60ba Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 4 Jul 2026 15:12:05 +0800
Subject: [PATCH 50/51] Credit co-authors of the AR-Diffusion Phase-1 work

This PR is the joint work of the AR-Diffusion / BDE group; crediting all
contributors so the squash merge carries every co-author.

Co-authored-by: amy-why-3459 <221522807+amy-why-3459@users.noreply.github.com>
Co-authored-by: Yangshen Deng <47519522+TKONIY@users.noreply.github.com>
Co-authored-by: Wang Fuyin <wfy2003324@163.com>
Co-authored-by: Honghan Zhu <nenbaying@163.com>
Co-authored-by: Match_yc <57976772+matchyc@users.noreply.github.com>
Co-authored-by: Hyoseop Song <lvsy@proton.me>
Co-authored-by: Wang Yiwen <121547057+yiwen101@users.noreply.github.com>
Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: Wu Jiazhen <83007646+asukaqaq-s@users.noreply.github.com>
Co-authored-by: Yukim1 <121286183+zwhzzz0821@users.noreply.github.com>
Co-authored-by: wtz2333 <wtz2333@mail.ustc.edu.cn>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>

From de9da444b530887ce95a79e73ee4aa7385092d62 Mon Sep 17 00:00:00 2001
From: Taichang Zhou <tzhouam@connect.ust.hk>
Date: Sat, 4 Jul 2026 16:02:05 +0800
Subject: [PATCH 51/51] [AR-Diffusion] Update DreamZeroState test for
 engine-owned KV

test_dreamzero_state_cache_access_requires_initialization exercised the
model-local get_kv_caches/update_kv_cache/get_crossattn_caches accessors that
the paged self-attention backend removed (all KV is engine-owned now). Replace
it with the inverse contract: those accessors must be GONE from DreamZeroState
so nothing can silently bypass the engine pool. Fixes the Simple unit-test CI
failure; the remaining two failures in that bucket (serve_cli mock, voxcpm2
engine_extras) reproduce on origin/main and are unrelated to this PR.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
---
 tests/dreamzero/test_pipeline_state.py | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/tests/dreamzero/test_pipeline_state.py b/tests/dreamzero/test_pipeline_state.py
index a3a5aaec03c..f00caf53ca7 100644
--- a/tests/dreamzero/test_pipeline_state.py
+++ b/tests/dreamzero/test_pipeline_state.py
@@ -4,7 +4,6 @@
 from collections import OrderedDict
 
 import pytest
-import torch
 
 from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
@@ -46,14 +45,11 @@ def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None:
     assert "session-b" not in pipeline._states
 
 
-def test_dreamzero_state_cache_access_requires_initialization() -> None:
+def test_dreamzero_state_owns_no_kv_caches() -> None:
+    """KV (self- and cross-attention) is engine-owned since the AR-Diffusion
+    paged backend: the model-local cache accessors must be gone so nothing can
+    silently bypass the engine pool."""
     state = DreamZeroState()
 
-    with pytest.raises(RuntimeError, match="KV caches not initialized"):
-        state.get_kv_caches()
-
-    with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"):
-        state.get_crossattn_caches()
-
-    with pytest.raises(RuntimeError, match="create_kv_caches first"):
-        state.update_kv_cache(0, torch.empty(0))
+    for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"):
+        assert not hasattr(state, removed)