ci: raise merge-tier entrypoint timeouts (H100 60->90, L4 30->45)

Both entrypoint jobs are killed mid-run by timeout_in_minutes while tests are still passing, not by any failing test. The budgets had no headroom: on main the H100 job timed out in 2 of the last 3 gate builds (2648, 2650) and passed 2651 at 56/60 min; the L4 job passed at 26/30 min only with a warm image cache. On the vLLM 0.25 wheel, model weight loading is slower on the CI runners (like-for-like: +10 min total on H100, 1.7B TTS load 11s -> 132s on L4), which turns the coin flip into a guaranteed timeout (build 2655). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-07-12 02:52:46 +08:00
54 changed files with 138 additions and 2236 deletions
@@ -108,7 +108,10 @@ steps:
    depends_on: upload-merge-pipeline
    steps:
      - label: "Entrypoint Test with H100"
-        timeout_in_minutes: 60
+        # The suite relaunches an OmniServer per test class (~30 model loads,
+        # incl. ten ~27.5 GiB Bagel loads); 60 min passed only ~1 in 3 runs
+        # even on main, and vLLM 0.25 weight loading needs ~10 min more.
+        timeout_in_minutes: 90
        commands:
          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
        agents:
@@ -146,7 +149,9 @@ steps:
                      type: DirectoryOrCreate

      - label: "Entrypoint Test with L4"
-        timeout_in_minutes: 30
+        # 30 min left no headroom for a cold image pull (~11 min for new base
+        # layers) plus per-class engine relaunches; main passed at 26/30.
+        timeout_in_minutes: 45
        commands:
          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
        agents:
@@ -650,8 +650,6 @@ steps:
            tests/e2e/online_serving/test_qwen_image_edit_expansion.py
            tests/e2e/online_serving/test_qwen_image_layered_expansion.py
            tests/e2e/online_serving/test_bagel_expansion.py
-            tests/e2e/online_serving/test_krea2_expansion.py
-            tests/e2e/offline_inference/test_krea2_expansion.py
            -m "full_model and diffusion and H100 and not distributed_cuda"
            --run-level "full_model"
        agents:
@@ -695,7 +693,6 @@ steps:
            pytest -sv
            tests/e2e/online_serving/test_qwen_image_expansion.py
            tests/e2e/online_serving/test_qwen_image_edit_expansion.py
-            tests/e2e/offline_inference/test_krea2_expansion.py
            -m "full_model and diffusion and H100 and distributed_cuda"
            --run-level "full_model"
        agents:
@@ -16,7 +16,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---

 *Latest News* 🔥
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
 - [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
 - [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
 - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
@@ -32,9 +31,9 @@ Easy, fast, and cheap omni-modality model serving for everyone

 [vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:

- **Omni-modality**: Text, image, audio, video, and action data processing
+- **Omni-modality**: Text, image, video, and audio data processing
 - **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs
+- **Heterogeneous outputs**: from traditional text generation to multimodal outputs

 <p align="center">
  <picture>
@@ -58,10 +57,9 @@ vLLM-Omni is flexible and easy to use with:

 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:

- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
+- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
 - **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
 - **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
- **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)

 ## Getting Started

@@ -2,7 +2,7 @@
 # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
 # be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.25.0
+ARG VLLM_BASE_TAG=v0.24.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -28,9 +28,9 @@ Easy, fast, and cheap omni-modality model serving for everyone

 [vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:

- **Omni-modality**: Text, image, audio, video, and action data processing
+- **Omni-modality**: Text, image, video, and audio data processing
 - **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs
+- **Heterogeneous outputs**: from traditional text generation to multimodal outputs

 <p align="center">
  <picture>
@@ -55,10 +55,8 @@ vLLM-Omni is flexible and easy to use with:

 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:

- Omni-modality models (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
- TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- Diffusion models — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
- Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
+- Omni-modality models (e.g. Qwen2.5-Omni, Qwen3-Omni)
+- Multi-modality generation models (e.g. Qwen-Image)

 For more information, checkout the following:

@@ -13,7 +13,7 @@ This document outlines the architectural design for vLLM-Omni.

 The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:

-* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, video, and action trajectories, alongside text.
+* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text.
 * **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
 * **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
 * **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

 ### Pre-built wheels

-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, and 0.25.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, and 0.24.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).

 === "NVIDIA CUDA"

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.25.0 --torch-backend=auto
+uv pip install vllm==0.24.0 --torch-backend=auto
 ```

 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]

 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.24.0 release version of the library

 ```bash
-uv pip install vllm==0.25.0 --torch-backend=auto
+uv pip install vllm==0.24.0 --torch-backend=auto
 ```

-The 0.25.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.24.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.25.0
+git checkout v0.24.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.25.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.25.0/vllm-0.25.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.24.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.24.0/vllm-0.24.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
    --env "HF_TOKEN=$HF_TOKEN" \
    -p 8091:8091 \
    --ipc=host \
-    vllm/vllm-omni:v0.25.0 \
+    vllm/vllm-omni:v0.24.0 \
    vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
+uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
 ```

 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]

 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.24.0 release version of the library

 ```bash
-uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
+uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
 ```

-The pre-built 0.25.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.24.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.25.0
+git checkout v0.24.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.25.0 \
+  vllm/vllm-omni-rocm:v0.24.0 \
  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

@@ -149,7 +149,7 @@ docker run --rm -it \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.25.0
+  vllm/vllm-omni-rocm:v0.24.0
 ```

 # --8<-- [end:pre-built-images]
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate

 # On CUDA
-uv pip install vllm==0.25.0 --torch-backend=auto
+uv pip install vllm==0.24.0 --torch-backend=auto

 # On ROCm
-uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
+uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722

 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
    It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.

-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.25.0` and vLLM Omni `0.25.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.24.0` and vLLM Omni `0.24.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.

 ## Offline Inference

@@ -34,7 +34,6 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -121,7 +121,6 @@ The following tables show which models support each feature:
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
-| **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -145,7 +144,6 @@ The following tables show which models support each feature:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
 > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
 > 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
-> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.

 ### VideoGen

@@ -36,7 +36,6 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
 | `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
 | `HiDream-I1-Full` | 1024 x 1024 | 63.7  | 57.7 |
-| `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | 1024 (Raw) / 2048 (Turbo) | — | ~30 |

 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -197,37 +196,6 @@ python examples/offline_inference/text_to_image/text_to_image.py \
  --output /output.png
 ```

-### Krea 2
-
-Krea 2 is published as diffusers-format checkpoints: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw`.
-The pipeline class (`Krea2Pipeline`) is auto-detected from `model_index.json`. The distilled Turbo checkpoint sets
-`is_distilled=true` and is best run at 2048x2048 with few steps and guidance disabled:
-
-```bash
-# Few-step distilled (Turbo) checkpoint
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model krea/Krea-2-Turbo \
-  --prompt "a fox in the snow" \
-  --num-inference-steps 8 \
-  --guidance-scale 0.0 \
-  --height 2048 --width 2048 \
-  --output krea2.png
-```
-
-For the Raw checkpoint, use more steps with CFG enabled:
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model krea/Krea-2-Raw \
-  --prompt "a fox in the snow" \
-  --num-inference-steps 28 \
-  --guidance-scale 4.5 \
-  --output krea2.png
-```
-
-`guidance-scale` follows the Krea 2 convention (`velocity = cond + guidance_scale * (cond - uncond)`); CFG is enabled
-whenever `guidance-scale > 0`.
-
 ### Batch Requests (Multiple Prompts)

 You can pass multiple prompts in a single `Omni.generate` call. `Omni`
@@ -44,7 +44,6 @@ recipes/
 | [`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) |
 | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
-| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | 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) |
 | [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU |
@@ -1,143 +0,0 @@
-# Krea 2 for text-to-image (base + LoRA)
-
-## Summary
-
- Vendor: Krea
- Model: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw`
- Task: Text-to-image (T2I)
- Mode: Offline engine and online OpenAI-compatible serving
- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe for a known-good starting point to run the Krea 2 single-stream
-MMDiT with vLLM-Omni on a **single 80 GB NVIDIA H100** — both the offline
-`text_to_image.py` entrypoint and the online `/v1/images/generations` route,
-including PEFT LoRA adapters. `Krea2Pipeline` is auto-detected from
-`model_index.json`; guidance is checkpoint-aware (distilled Turbo runs no-CFG,
-Raw runs CFG).
-
-## References
-
- Related example under `examples/`:
-  [`examples/offline_inference/text_to_image/README.md`](../../examples/offline_inference/text_to_image/README.md)
-  (see the "Krea 2" section)
- Feature support:
-  [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md)
- Related PR: [#4730](https://github.com/vllm-project/vllm-omni/pull/4730)
-
-## Hardware Support
-
-This recipe documents a **single-GPU** CUDA layout on H100 80 GB. Add more
-platforms (ROCm / NPU) or multi-GPU layouts as community validation lands.
-
-## GPU
-
-### 1× H100 80GB
-
-#### Environment
-
-Versions from a working editable install (activate `vllm-omni/.venv` or your equivalent):
-
- OS: Linux
- Python: 3.12
- Driver / runtime: NVIDIA driver **595.71.05** (CUDA 13.x)
- torch: **2.11.0+cu130**
- vLLM: **0.24.0**
- vLLM-Omni: editable install from this repo (Git **`960ce8f6`**)
- Transformers: **5.12.1**
-
-Krea 2 weights are ~30 GB and fit comfortably on a single 80 GB H100.
-
-#### Command
-
-Offline — few-step distilled (Turbo) checkpoint (2048×2048, no CFG):
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model krea/Krea-2-Turbo \
-  --prompt "a fox in the snow" \
-  --num-inference-steps 8 \
-  --guidance-scale 0.0 \
-  --height 2048 --width 2048 \
-  --output krea2_turbo.png
-```
-
-Offline — Raw checkpoint (1024×1024, CFG enabled):
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model krea/Krea-2-Raw \
-  --prompt "a fox in the snow" \
-  --num-inference-steps 28 \
-  --guidance-scale 4.5 \
-  --height 1024 --width 1024 \
-  --output krea2_raw.png
-```
-
-LoRA (offline) — pass a PEFT adapter with `--lora-path` / `--lora-scale`.
-`NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA` is a vLLM-Omni-compatible repackaging
-of `krea/Krea-2-LoRA-darkbrush`; its trigger word is `monochrome ink wash style`:
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model krea/Krea-2-Turbo \
-  --prompt "a fox in the snow, monochrome ink wash style" \
-  --num-inference-steps 8 --guidance-scale 0.0 \
-  --lora-path NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA --lora-scale 1.0 \
-  --output krea2_darkbrush.png
-```
-
-Online — serve and query the OpenAI-compatible image route:
-
-```bash
-vllm serve krea/Krea-2-Turbo --omni --enforce-eager --port 8091
-```
-
-```bash
-curl -s http://localhost:8091/v1/images/generations \
-  -H "Content-Type: application/json" \
-  -d '{
-    "prompt": "a fox in the snow, monochrome ink wash style",
-    "size": "1024x1024",
-    "num_inference_steps": 8,
-    "guidance_scale": 0.0,
-    "seed": 42,
-    "lora": {"name": "darkbrush", "path": "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA", "scale": 1.0}
-  }' | jq -r '.data[0].b64_json' | base64 -d > krea2_online.png
-```
-
-Cache-DiT acceleration (offline) — add `--cache-backend cache_dit`:
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model krea/Krea-2-Turbo --prompt "a fox in the snow" \
-  --num-inference-steps 8 --guidance-scale 0.0 \
-  --cache-backend cache_dit --output krea2_cachedit.png
-```
-
-#### Verification
-
-```bash
-# Offline: the script prints "Saved generated image to <path>"; confirm the file exists.
-ls -lh krea2_turbo.png
-```
-
-For the online route, decode the returned `b64_json` (see the `curl` above) and
-confirm a valid 1024×1024 PNG is written.
-
-#### Notes
-
- Key flags: distilled Turbo → few steps (~8) with `guidance_scale=0.0`; Raw →
-  more steps (~28) with `guidance_scale>0` (CFG). `guidance-scale` follows the
-  Krea 2 convention `velocity = cond + guidance_scale * (cond - uncond)`.
- LoRA: adapters must be PEFT format (`adapter_config.json` +
-  `adapter_model.safetensors`); the projections are `ReplicatedLinear`
-  (264 modules). Offline uses `--lora-path` (which sets
-  `OmniDiffusionSamplingParams.lora_request` / `lora_scale`); online uses the
-  top-level `lora` object. Style LoRAs like darkbrush need their trigger word in
-  the prompt.
- Cache-DiT is supported; `has_separate_cfg` is checkpoint-aware (False for Turbo
-  no-CFG, True for Raw CFG).
- Known limitations: TP / SP / CFG-Parallel are not yet wired for Krea 2. HSDP,
-  layerwise CPU offload, and VAE-patch-parallel (decode) are supported.
@@ -9,9 +9,14 @@ 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

+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,
@@ -48,7 +53,10 @@ 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.
    spec = make_spec()
-    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+    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():
@@ -62,8 +62,8 @@ def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> di
        ),
        pytest.param(
            True,
-            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "1"},
-            ("Invalid layers", "2", "10"),
+            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"},
+            ("Invalid layers", "3", "1"),
            400,
            id="invalid_layers",
        ),
@@ -87,8 +87,8 @@ def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
        pytest.param(
            {"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max"
        ),
-        pytest.param({"layers": 1}, ("layers", "value_error", "2", "10"), id="layers_below_min"),
-        pytest.param({"layers": 11}, ("layers", "value_error", "2", "10"), id="layers_above_max"),
+        pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"),
+        pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"),
        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649),
        pytest.param(
            {"seed": 2**32},
@@ -1,155 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""L4 functionality tests for the Krea 2 text-to-image diffusion pipeline (``Krea2Pipeline``).
-
-Runs against the public few-step distilled checkpoint ``krea/Krea-2-Turbo`` by default (override with the
-``KREA2_MODEL`` environment variable, e.g. ``krea/Krea-2-Raw`` for the Raw checkpoint or a local diffusers
-directory). They cover a basic functional smoke plus the layerwise-CPU-offload path that the pipeline declares
-via ``SupportsComponentDiscovery`` / ``_layerwise_offload_blocks_attrs``.
-"""
-
-import os
-
-import numpy as np
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunnerHandler
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.lora.request import LoRARequest
-from vllm_omni.lora.utils import stable_lora_int_id
-
-MODEL = os.environ.get("KREA2_MODEL", "krea/Krea-2-Turbo")
-# vLLM-Omni-compatible PEFT repackaging of krea/Krea-2-LoRA-darkbrush (264 modules, r=alpha=32).
-LORA = os.environ.get("KREA2_LORA", "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA")
-PROMPT = "a fox in the snow, photorealistic"
-# darkbrush is a trigger-word LoRA: the trained monochrome ink-wash style only manifests when
-# its trigger phrase is present, so the LoRA test prompt appends it for a strong, stable signal.
-LORA_TRIGGER = "monochrome ink wash style"
-LORA_PROMPT = f"a fox in the snow, {LORA_TRIGGER}"
-
-pytestmark = [
-    pytest.mark.diffusion,
-    pytest.mark.full_model,
-]
-
-
-def _sampling() -> OmniDiffusionSamplingParams:
-    # Small resolution + few steps to keep the L4 case light. guidance_scale resolution is checkpoint-aware inside
-    # the pipeline (distilled -> no-CFG, Raw -> CFG), so this stays agnostic to which checkpoint KREA2_MODEL points at.
-    return OmniDiffusionSamplingParams(
-        height=512,
-        width=512,
-        num_inference_steps=8,
-        guidance_scale=0.0,
-        seed=42,
-    )
-
-
-@hardware_test(res={"cuda": "H100"})
-@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
-def test_krea2_text_to_image_001(omni_runner_handler: OmniRunnerHandler) -> None:
-    """Basic functional smoke: a single prompt produces a decoded image."""
-    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
-
-
-@hardware_test(res={"cuda": "H100"})
-@pytest.mark.parametrize(
-    "omni_runner",
-    [(MODEL, None, {"enable_layerwise_offload": True})],
-    indirect=True,
-)
-def test_krea2_layerwise_offload(omni_runner_handler: OmniRunnerHandler) -> None:
-    """Exercise layerwise CPU offload on the DiT (SupportsComponentDiscovery + _layerwise_offload_blocks_attrs)."""
-    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
-
-
-@hardware_test(res={"cuda": "H100"})
-@pytest.mark.cache
-@pytest.mark.parametrize("omni_runner", [(MODEL, None, {"cache_backend": "cache_dit"})], indirect=True)
-def test_krea2_cache_dit(omni_runner_handler: OmniRunnerHandler) -> None:
-    """Exercise Cache-DiT on Krea 2 via the custom Krea2Pipeline enabler (``enable_cache_for_krea2``).
-
-    Validates the docs' Cache-DiT support claim for Krea 2. ``has_separate_cfg`` is checkpoint-aware
-    (False for the distilled Turbo no-CFG path, True for the Raw CFG path); the default Turbo checkpoint
-    exercises the no-CFG branch.
-    """
-    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
-
-
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-@pytest.mark.parametrize(
-    "omni_runner",
-    # Standalone HSDP: use_hsdp=True with an explicit hsdp_shard_size shards the DiT weights
-    # across 2 ranks (replicate_size defaults to 1 -> world_size = 1 * 2 = 2). HSDP must keep
-    # tensor_parallel_size/data_parallel_size at 1 (see vllm_omni/diffusion/data.py __post_init__).
-    # Also validated locally at 8 GPUs (hsdp_shard_size=8); CI shard provisions 2.
-    [(MODEL, None, {"use_hsdp": True, "hsdp_shard_size": 2})],
-    indirect=True,
-)
-def test_krea2_hsdp(omni_runner_handler: OmniRunnerHandler) -> None:
-    """Exercise Hybrid Sharded Data Parallel (HSDP) weight sharding on the DiT across 2 GPUs.
-
-    Validates the docs' HSDP-support claim for Krea 2. Scheduled on the multi-GPU (distributed_cuda)
-    shard by ``num_cards=2``; skipped automatically on boxes with fewer than 2 CUDA devices.
-    """
-    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
-
-
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-@pytest.mark.parametrize(
-    "omni_runner",
-    # VAE patch-parallel decode: partition the VAE decode across the TP world. Mirrors the online
-    # "tp_vae_patch" case (tensor_parallel_size=2 + vae_patch_parallel_size=2) from the wan2.2/z-image
-    # expansion suites, so vae_patch_parallel_size must equal the world size established by TP.
-    # Also validated locally at 8 GPUs (tp=8, vae_patch=8); CI shard provisions 2.
-    [(MODEL, None, {"tensor_parallel_size": 2, "vae_patch_parallel_size": 2})],
-    indirect=True,
-)
-def test_krea2_vae_patch_parallel(omni_runner_handler: OmniRunnerHandler) -> None:
-    """Exercise VAE patch-parallel (decode) across 2 GPUs.
-
-    Validates the docs' VAE-patch-parallel (decode) support claim for Krea 2. Scheduled on the
-    multi-GPU (distributed_cuda) shard by ``num_cards=2``; skipped automatically on boxes with
-    fewer than 2 CUDA devices.
-    """
-    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
-
-
-def _generate(handler: OmniRunnerHandler, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> np.ndarray:
-    sp = _sampling()
-    sp.lora_request = lora_request
-    sp.lora_scale = lora_scale
-    resp = handler.send_diffusion_request({"model": MODEL, "prompt": LORA_PROMPT, "sampling_params": sp})
-    return np.asarray(resp.images[0].convert("RGB"), dtype=np.int16)
-
-
-@hardware_test(res={"cuda": "H100"})
-@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
-def test_krea2_lora(omni_runner_handler: OmniRunnerHandler) -> None:
-    """Validate diffusion LoRA on Krea 2: visible effect, scale sensitivity, clean deactivation.
-
-    Uses ``LORA`` (a vLLM-Omni-compatible PEFT repackaging of ``krea/Krea-2-LoRA-darkbrush``,
-    264 modules matched via the ReplicatedLinear projections) passed per-request through
-    ``OmniDiffusionSamplingParams.lora_request`` / ``lora_scale`` — the same fields the
-    diffusion LoRA manager reads to activate/deactivate the adapter.
-    """
-    lora_request = LoRARequest(lora_name="darkbrush", lora_int_id=stable_lora_int_id(LORA), lora_path=LORA)
-
-    baseline = _generate(omni_runner_handler, None)
-    img_1x = _generate(omni_runner_handler, lora_request, lora_scale=1.0)
-    img_2x = _generate(omni_runner_handler, lora_request, lora_scale=2.0)
-    restored = _generate(omni_runner_handler, None)
-
-    diff_1x = np.abs(baseline - img_1x).mean()
-    diff_2x = np.abs(baseline - img_2x).mean()
-    diff_restored = np.abs(baseline - restored).mean()
-
-    # (a) Adapter has a visible effect at both scales.
-    assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
-    assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
-    # (b) Scale changes the output (stronger adapter -> larger perturbation).
-    assert not np.isclose(diff_1x, diff_2x, atol=1.0), f"LoRA scale had no effect: 1x={diff_1x:.2f}, 2x={diff_2x:.2f}"
-    # (c) Passing lora_request=None cleanly deactivates: same seed -> back to baseline (modulo fp drift).
-    assert diff_restored < 5.0, f"LoRA did not deactivate cleanly: diff_restored={diff_restored:.2f}"
@@ -1,152 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""L4 online-serving tests for the Krea 2 text-to-image diffusion pipeline.
-
-Exercises the OpenAI-compatible ``/v1/images/generations`` route (served via
-``vllm serve <model> --omni --enforce-eager``) rather than the offline engine path
-covered by ``tests/e2e/offline_inference/test_krea2_expansion.py``.
-
-Runs against the public few-step distilled checkpoint ``krea/Krea-2-Turbo`` by default
-(override with the ``KREA2_MODEL`` environment variable, e.g. ``krea/Krea-2-Raw`` for the
-Raw checkpoint or a local diffusers directory). Krea 2 shares the Qwen-Image VAE and a
-similar single-stream MMDiT, so this test mirrors the structure/fixtures/markers of
-``test_qwen_image_expansion.py``.
-
-Coverage mirrors the offline suite: a basic functional smoke plus the layerwise-CPU-offload
-path that the pipeline declares via ``SupportsComponentDiscovery`` /
-``_layerwise_offload_blocks_attrs``. Sizing is kept light (512x512, few steps) so the case
-fits single-GPU H100 CI rather than the full 2048x2048 distilled resolution.
-"""
-
-import base64
-import os
-from io import BytesIO
-
-import numpy as np
-import pytest
-import requests
-from PIL import Image
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-MODEL = os.environ.get("KREA2_MODEL", "krea/Krea-2-Turbo")
-# vLLM-Omni-compatible PEFT repackaging of krea/Krea-2-LoRA-darkbrush (264 modules, r=alpha=32).
-LORA = os.environ.get("KREA2_LORA", "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA")
-PROMPT = "a fox in the snow, photorealistic"
-# darkbrush is a trigger-word LoRA; include its trigger phrase so the LoRA case has a strong signal.
-LORA_PROMPT = "a fox in the snow, monochrome ink wash style"
-
-SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
-
-
-def _get_diffusion_feature_cases(model: str):
-    """Return single-card L4 online-serving cases for Krea 2.
-
-    Mirrors the offline suite: a basic functional smoke plus layerwise CPU offload.
-    """
-    return [
-        # Basic functional smoke (single-card).
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--enforce-eager"],
-            ),
-            id="single_card_basic",
-            marks=SINGLE_CARD_FEATURE_MARKS,
-        ),
-        # Layerwise CPU offload (single-card): SupportsComponentDiscovery +
-        # _layerwise_offload_blocks_attrs.
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--enforce-eager", "--enable-layerwise-offload"],
-            ),
-            id="single_card_layerwise_offload",
-            marks=SINGLE_CARD_FEATURE_MARKS,
-        ),
-    ]
-
-
-@pytest.mark.parametrize(
-    "omni_server",
-    _get_diffusion_feature_cases(MODEL),
-    indirect=True,
-)
-def test_krea2(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Serve Krea 2 and validate the OpenAI-compatible image-generation route.
-
-    guidance_scale resolution is checkpoint-aware inside the pipeline (distilled -> no-CFG,
-    Raw -> CFG), so guidance_scale=0.0 stays agnostic to which checkpoint KREA2_MODEL points at.
-    Validation is delegated to assert_diffusion_response, which checks output dimensions and
-    basic correctness.
-    """
-    messages = dummy_messages_from_mix_data(content_text=PROMPT)
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "extra_body": {
-            "height": 512,
-            "width": 512,
-            "num_inference_steps": 8,
-            "guidance_scale": 0.0,
-            "seed": 42,
-        },
-    }
-    openai_client.send_diffusion_request(request_config)
-
-
-def _basic_lora_payload(lora: dict | None = None) -> dict:
-    payload = {
-        "prompt": LORA_PROMPT,
-        "n": 1,
-        "size": "512x512",
-        "num_inference_steps": 8,
-        "guidance_scale": 0.0,
-        "seed": 42,
-    }
-    if lora is not None:
-        payload["lora"] = lora
-    return payload
-
-
-def _post_image(server: OmniServer, payload: dict) -> np.ndarray:
-    url = f"http://{server.host}:{server.port}/v1/images/generations"
-    resp = requests.post(url, json=payload, headers={"Authorization": "Bearer EMPTY"}, timeout=900)
-    resp.raise_for_status()
-    b64 = resp.json()["data"][0]["b64_json"]
-    img = Image.open(BytesIO(base64.b64decode(b64)))
-    img.load()
-    return np.asarray(img.convert("RGB"), dtype=np.int16)
-
-
-@pytest.mark.parametrize(
-    "omni_server",
-    [
-        pytest.param(
-            OmniServerParams(model=MODEL, server_args=["--enforce-eager"]),
-            id="single_card_lora",
-            marks=SINGLE_CARD_FEATURE_MARKS,
-        )
-    ],
-    indirect=True,
-)
-def test_krea2_lora(omni_server: OmniServer):
-    """Serve Krea 2 and validate per-request LoRA over ``/v1/images/generations``.
-
-    Passes ``LORA`` via the top-level ``lora`` object (``{name, path, scale}``) the diffusion
-    image route parses into a ``LoRARequest`` — no server-side LoRA flag needed. Confirms the
-    adapter has a visible effect and that omitting ``lora`` cleanly restores the base output.
-    """
-    baseline = _post_image(omni_server, _basic_lora_payload())
-    with_lora = _post_image(omni_server, _basic_lora_payload({"name": "darkbrush", "path": LORA, "scale": 1.0}))
-    restored = _post_image(omni_server, _basic_lora_payload())
-
-    diff_lora = np.abs(baseline - with_lora).mean()
-    diff_restored = np.abs(baseline - restored).mean()
-
-    assert diff_lora > 0.5, f"LoRA had no visible effect over the serving path: diff={diff_lora}"
-    assert diff_restored < 5.0, f"LoRA did not deactivate cleanly: diff_restored={diff_restored:.2f}"
@@ -874,54 +874,3 @@ def test_build_engine_args_dict_uses_diffusion_attention_config_key():

    assert "attention_config" not in engine_args_dict
    assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN"
-
-
-def test_omni_master_server_allocates_globally_unique_route_ports(monkeypatch):
-    """Regression: two stages must never draw the same ZMQ port.
-
-    ``get_open_ports_list`` only dedups within a single call, so per-route
-    allocation used to let a later stage reuse an earlier stage's port. The
-    second engine to ``bind()`` then died with ``zmq.error.ZMQError: Address
-    already in use`` (flaky multi-stage startup, e.g. Qwen3-Omni thinker/talker/
-    code2wav). ``OmniMasterServer`` now dedups every port it hands out.
-    """
-    import itertools
-
-    from vllm_omni.engine import stage_engine_startup as ses
-
-    # A colliding prefix (repeats + the master port 9000) followed by an endless
-    # fresh stream, so the only way to succeed is to redraw the collisions.
-    supply = itertools.chain([9000, 9000, 9001, 9001, 9000], itertools.count(9002))
-
-    def fake_get_open_ports_list(count):
-        return [next(supply) for _ in range(count)]
-
-    monkeypatch.setattr(ses, "get_open_ports_list", fake_get_open_ports_list)
-
-    server = ses.OmniMasterServer(
-        master_address="127.0.0.1",
-        master_port=9000,  # seed: a route must not reuse the registration port
-        stage_ids=[0, 1, 2],
-    )
-
-    ports = []
-    for sid in (0, 1, 2):
-        alloc = server.get_allocation(sid)
-        for addr in (
-            alloc.handshake_bind_address,
-            alloc.input_bind_address,
-            alloc.output_bind_address,
-        ):
-            ports.append(ses._port_from_zmq_address(addr))
-
-    assert len(ports) == len(set(ports)), f"duplicate route ports allocated: {ports}"
-    assert 9000 not in ports, "route reused the master registration port"
-
-
-def test_port_from_zmq_address_parsing():
-    from vllm_omni.engine.stage_engine_startup import _port_from_zmq_address
-
-    assert _port_from_zmq_address("tcp://127.0.0.1:34277") == 34277
-    assert _port_from_zmq_address(None) is None
-    assert _port_from_zmq_address("ipc:///tmp/sock") is None
-    assert _port_from_zmq_address("tcp://host:not-a-port") is None
@@ -198,28 +198,12 @@ def _assert_realtime_smoke(result: dict) -> None:
    assert result["output_sample_rate"] > 0


-def _assert_realtime_accuracy(
-    result: dict,
-    whisper_model_size: str = "large-v3",
-) -> None:
-    """Assert that whisper transcription of audio output matches model text.
-
-    Args:
-        result: Roundtrip result dict from ``_run_realtime_audio_roundtrip``.
-        whisper_model_size: Whisper model used to transcribe the generated audio
-                   for the accuracy check. Defaults to ``large-v3``: the default
-                   ``small`` model mishears short Chinese TTS clips (observed:
-                   北京→韦京 and a dropped leading sentence, sim=0.443), which
-                   caused spurious sim<0.8 failures under async_chunk codec
-                   variability even though audio generation was correct. large-v3
-                   transcribes these clips reliably, so a failure here now points
-                   at the model, not the ASR grader.
-    """
+def _assert_realtime_accuracy(result: dict) -> None:
    final_text = (result["transcription_text"] or "").strip()
    assert final_text, "Expected non-empty transcription (model text stream)"

    wav_out = _wav_bytes_from_pcm16(result["output_pcm"], result["output_sample_rate"])
-    whisper_text = convert_audio_bytes_to_text(wav_out, model_size=whisper_model_size).strip()
+    whisper_text = convert_audio_bytes_to_text(wav_out).strip()
    assert whisper_text, "Whisper returned empty string for synthesized output audio"

    sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
@@ -87,8 +87,8 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M
    serving_chat.chat_template = None
    serving_chat.chat_template_content_format = "string"
    serving_chat.default_chat_template_kwargs = {}
-    serving_chat.online_renderer = mocker.MagicMock()
-    serving_chat.online_renderer.validate_chat_template.return_value = None
+    serving_chat.openai_serving_render = mocker.MagicMock()
+    serving_chat.openai_serving_render.validate_chat_template.return_value = None
    serving_chat._effective_chat_template_kwargs = mocker.MagicMock(return_value={})
    serving_chat._preprocess_chat = mocker.AsyncMock(
        side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian")
@@ -144,7 +144,7 @@ def _build_serving_chat():
        engine_client=mock_engine,
        models=models,
        response_role="assistant",
-        online_renderer=mock_render,
+        openai_serving_render=mock_render,
        request_logger=None,
        chat_template=None,
        chat_template_content_format="auto",
@@ -2826,7 +2826,7 @@ def iter_omni_server(
        # be ideal to cleanup consistently everywhere.
        original_model = model_prefix + params.model
        model = original_model
-        if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
+        if run_level == "core_model":
            model = resolve_tiny_model_path(model)
        port = params.port
        stage_config_path = stage_config_path_for_run_level(params.stage_config_path, run_level)
@@ -2917,7 +2917,7 @@ def iter_omni_runner(
            extra_omni_kwargs = dict(extra) if extra is not None else {}
        stage_config_path = stage_config_path_for_run_level(stage_config_path, run_level)
        model = model_prefix + model
-        if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
+        if run_level == "core_model":
            model = resolve_tiny_model_path(model)
        with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
            print("OmniRunner started successfully")
@@ -52,10 +52,6 @@ def _make_bare_talker():
    talker._cfm_buffers = None
    talker._last_audio_output_req_ids = []
    talker._batched_fsq_fusion_max_batch = 32
-    # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to
-    # default (False) to avoid AttributeError when test exercises forward
-    # path that checks this flag.
-    talker._enable_unified_decode_graph = False
    return talker


@@ -333,27 +329,23 @@ class TestDecodeBatchContract:

        for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True):
            torch.testing.assert_close(actual_res, expected_res)
-            torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden)
+            torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"])
        assert [state.decode_step_count for state in batch_states] == [1, 1]

    def test_batch_decode_matches_per_request_state_updates(self) -> None:
-        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
-            _DecodeResidualMeta,
-        )
-
        seq_talker = self._make_decode_talker()
        batch_talker = self._make_decode_talker()

        seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
        batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
        metas = [
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
        ]
        batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])

        for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True):
-            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0))
+            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu"))

        batch_talker._finish_decode_batch(
            [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)],
@@ -368,16 +360,12 @@ class TestDecodeBatchContract:
            torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu)

    def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None:
-        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
-            _DecodeResidualMeta,
-        )
-
        talker = self._make_decode_talker()
        talker._enable_batched_cfm = False
        states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
        metas = [
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
        ]
        batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])

@@ -1,102 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from collections import defaultdict
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.model_executor.models.ming_tts.constants import LATENT_DIM, PATCH_SIZE
-from vllm_omni.model_executor.stage_input_processors.ming_tts import (
-    MING_EMIT_PATCH_COUNT_KEY,
-    llm2audio_vae_async_chunk,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_transfer_manager(*, chunk_size: int = 5, initial_chunk_size: int = 2):
-    return SimpleNamespace(
-        code_prompt_token_ids=defaultdict(list),
-        put_req_chunk=defaultdict(int),
-        request_payload={},
-        connector=SimpleNamespace(
-            config={
-                "extra": {
-                    "latent_chunk_size": chunk_size,
-                    "initial_latent_chunk_size": initial_chunk_size,
-                    "latent_left_context": 0,
-                }
-            }
-        ),
-    )
-
-
-def _make_request(req_id: str = "req", *, finished: bool = False):
-    return SimpleNamespace(
-        external_req_id=req_id,
-        is_finished=lambda: finished,
-    )
-
-
-def _make_output(value: float):
-    return {
-        "ming_has_patch": torch.tensor([1], dtype=torch.bool),
-        "ming_latent_patch": torch.full((1, PATCH_SIZE, LATENT_DIM), value),
-    }
-
-
-def _append_patch(tm, req_id: str, idx: int, *, finished: bool = False):
-    return llm2audio_vae_async_chunk(
-        tm,
-        _make_output(float(idx)),
-        _make_request(req_id, finished=finished),
-        is_finished=finished,
-    )
-
-
-def test_ming_async_chunk_emits_small_initial_chunk_then_steady_chunks():
-    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
-
-    assert _append_patch(tm, "req", 0) is None
-
-    first = _append_patch(tm, "req", 1)
-    assert first is not None
-    assert first.latent.shape == (2, PATCH_SIZE, LATENT_DIM)
-    assert first.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 2
-    assert first.meta.finished.item() is False
-
-    for idx in range(2, 6):
-        assert _append_patch(tm, "req", idx) is None
-
-    second = _append_patch(tm, "req", 6)
-    assert second is not None
-    assert second.latent.shape == (5, PATCH_SIZE, LATENT_DIM)
-    assert second.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 5
-
-
-def test_ming_async_chunk_flushes_short_final_chunk():
-    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
-
-    payload = _append_patch(tm, "req", 0, finished=True)
-
-    assert payload is not None
-    assert payload.latent.shape == (1, PATCH_SIZE, LATENT_DIM)
-    assert payload.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 1
-    assert payload.meta.finished.item() is True
-
-
-def test_ming_async_chunk_flushes_leftover_after_initial_chunk():
-    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
-
-    assert _append_patch(tm, "req", 0) is None
-    assert _append_patch(tm, "req", 1) is not None
-    assert _append_patch(tm, "req", 2) is None
-
-    final = _append_patch(tm, "req", 3, finished=True)
-
-    assert final is not None
-    assert final.latent.shape == (2, PATCH_SIZE, LATENT_DIM)
-    assert final.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 2
-    assert final.meta.finished.item() is True
@@ -56,7 +56,6 @@ EXCLUDED_MODELS = [
    "NextStep11Pipeline",
    "FluxPipeline",
    "FluxDMD2Pipeline",
-    "Krea2Pipeline",
    "QwenImageDMD2Pipeline",
    "OmniGen2Pipeline",
    "HeliosPipeline",
@@ -314,6 +314,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):

        outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
        spec_decoding_stats: SpecDecodingStats | None = None
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        if kv_connector_stats and self.connector:
+            kv_stats = self.connector.get_kv_connector_stats()
+            if kv_stats:
+                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)

        failed_kv_load_req_ids = None
        if kv_connector_output and kv_connector_output.invalid_block_ids:
@@ -556,20 +563,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
        if kv_connector_output:
            self._update_from_kv_xfer_finished(kv_connector_output)

-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
        # collect KV cache events from KV cache manager
        events = self.kv_cache_manager.take_events()

@@ -431,6 +431,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):

        outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
        spec_decoding_stats: SpecDecodingStats | None = None
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        # Merge connector-side stats (align with v0.14.0)
+        if kv_connector_stats and self.connector:
+            kv_stats = self.connector.get_kv_connector_stats()
+            if kv_stats:
+                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)

        failed_kv_load_req_ids = None
        if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None):
@@ -638,20 +646,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
        if kv_connector_output:
            self._update_from_kv_xfer_finished(kv_connector_output)

-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
        # Collect and publish KV cache events (align with v0.14.0)
        events = self.kv_cache_manager.take_events()
        if self.connector is not None:
@@ -15,7 +15,6 @@ connectors:
    name: SharedMemoryConnector
    extra:
      latent_chunk_size: 25
-      initial_latent_chunk_size: 4
      latent_left_context: 0

 stages:
@@ -23,7 +23,6 @@ connectors:
    name: SharedMemoryConnector
    extra:
      latent_chunk_size: 25
-      initial_latent_chunk_size: 4
      latent_left_context: 0

 stages:
@@ -975,38 +975,6 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheCo
    return enable_cache_for_dit(pipeline, cache_config, block_adapter)


-def enable_cache_for_krea2(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc:
-    """Enable cache-dit for Krea 2.
-
-    Krea 2 is a single-stream MMDiT: each ``Krea2TransformerBlock`` takes and returns only
-    ``hidden_states`` (text is fused into the token stream), so the blocks follow
-    ``ForwardPattern.Pattern_3``.
-
-    ``has_separate_cfg`` is checkpoint-dependent, which is why this needs a custom enabler
-    rather than a static ``_cache_dit_adapter_config``: the distilled Turbo checkpoint runs
-    no-CFG (a single transformer forward per denoise step), while the Raw checkpoint runs CFG
-    as two separate forwards. cache-dit tells cond/uncond apart purely by transformer-forward
-    parity, so the flag must match the actual per-step forward count — ``True`` only for the
-    CFG (Raw) path. The pipeline exposes this via ``is_distilled`` (read from ``model_index.json``).
-
-    Args:
-        pipeline: The Krea2Pipeline 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 = default_get_pipeline_transformer(pipeline)
-    block_adapter = BlockAdapter(
-        transformer=transformer,
-        blocks=[transformer.transformer_blocks],
-        forward_pattern=[ForwardPattern.Pattern_3],
-        has_separate_cfg=not pipeline.is_distilled,
-        check_forward_pattern=False,
-    )
-    return enable_cache_for_dit(pipeline, cache_config, block_adapter)
-
-
 # Register custom cache-dit enablers after function definitions
 CUSTOM_DIT_ENABLERS.update(
    {
@@ -1016,7 +984,6 @@ CUSTOM_DIT_ENABLERS.update(
        "Wan22VACEPipeline": enable_cache_for_wan22,
        "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
        "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
-        "Krea2Pipeline": enable_cache_for_krea2,
    }
 )

@@ -392,77 +392,16 @@ class DiffusersPipelineLoader:

        return model.eval()

-    @staticmethod
-    def _has_online_quant(model: nn.Module) -> bool:
-        """Whether any layer uses an online-quant method that defers weight
-        materialization onto the ``meta`` device (upstream vLLM
-        ``uses_meta_device=True``, e.g. online FP8)."""
-        for module in model.modules():
-            quant_method = getattr(module, "quant_method", None)
-            if getattr(quant_method, "uses_meta_device", False):
-                return True
-        return False
-
    def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None:
        """Process weights after loading for quantization methods.

        This handles vLLM's quantization methods that need to process weights
        after loading (e.g., FP8 online quantization from BF16/FP16 weights).
        """
-        # Newer upstream vLLM online-quant methods (uses_meta_device=True) create
-        # weights on the ``meta`` device and materialize them just-in-time as each
-        # layer's weights finish loading (via the layerwise online-process loader).
-        # Any "straggler" layers whose weights were not fully materialized during
-        # load (padded / partially-loaded layers) remain on ``meta``. Upstream's
-        # base_loader calls finalize_layerwise_processing() to materialize them;
-        # the diffusion loader must mirror that, otherwise the module.to() below
-        # raises "Cannot copy out of meta tensor; no data!". This whole meta-device
-        # handling is gated on online quant actually being in use, so that the
-        # proven code path for everything else (in particular FSDP/HSDP-sharded
-        # params, whose per-parameter .data cannot be cross-device reassigned) is
-        # left untouched. Import lazily so older vLLM (no meta-device quant) is
-        # unaffected.
-        has_online_quant = self._has_online_quant(model)
-        if has_online_quant:
-            from vllm.model_executor.model_loader.reload.layerwise import (
-                finalize_layerwise_processing,
-            )
-
-            # model_config is only dereferenced by finalize for vLLM Attention /
-            # MLAAttention layers; diffusion DiT models use their own attention and
-            # have none, so passing None is safe here.
-            finalize_layerwise_processing(model, model_config=None)
-
        for _, module in model.named_modules():
            quant_method = getattr(module, "quant_method", None)
-            if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
-                continue
-
-            if has_online_quant:
-                # Online quant may leave straggler params on the ``meta`` device.
-                # Move only real (non-meta) params onto the target device for
-                # processing and restore them afterward, mirroring upstream vLLM's
-                # device_loading_context — a blanket module.to(target_device) would
-                # raise NotImplementedError on meta params. Online quant initializes
-                # on the accelerator, so params are normally already on the target
-                # device and this loop is a no-op move; the point is to skip meta.
-                original_devices: dict[str, torch.device] = {}
-                for name, param in module.named_parameters():
-                    if param.device.type != "meta" and param.device != target_device:
-                        original_devices[name] = param.device
-                        param.data = param.data.to(target_device)
-
-                quant_method.process_weights_after_loading(module)
-
-                # Restore pre-existing params to their original device; leave any
-                # newly created (e.g. quantized) params on the target device.
-                for name, param in module.named_parameters():
-                    if name in original_devices:
-                        param.data = param.data.to(original_devices[name])
-            else:
-                # No meta params possible here. Preserve the original FSDP/HSDP-aware
-                # whole-module move (module.to()), which correctly handles sharded
-                # DTensor params that per-parameter .data reassignment cannot.
+            if quant_method is not None and isinstance(quant_method, QuantizeMethodBase):
+                # Move module to target device for processing if needed
                module_device = next(module.parameters(), None)
                if module_device is not None:
                    module_device = module_device.device
@@ -19,7 +19,6 @@ 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.configs.gr00t_n1d7 import Gr00tN1d7Config
 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
@@ -553,9 +552,4 @@ class Gr00tN1d7Processor(ProcessorMixin):
        return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs)


-# transformers >= 5.5 requires the config *class* (not a model_type string) as the
-# first arg to AutoProcessor.register — it does `key.__module__` internally. Passing
-# the string "Gr00tN1d7" raises AttributeError: 'str' object has no attribute
-# '__module__'. Match the AutoModel.register(Gr00tN1d7Config, ...) call in
-# modeling/gr00t_n1d7.py.
-AutoProcessor.register(Gr00tN1d7Config, Gr00tN1d7Processor)
+AutoProcessor.register("Gr00tN1d7", Gr00tN1d7Processor)
@@ -1,15 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Krea 2 diffusion model components."""
-
-from vllm_omni.diffusion.models.krea2.krea2_transformer import Krea2Transformer2DModel
-from vllm_omni.diffusion.models.krea2.pipeline_krea2 import (
-    Krea2Pipeline,
-    get_krea2_post_process_func,
-)
-
-__all__ = [
-    "Krea2Pipeline",
-    "Krea2Transformer2DModel",
-    "get_krea2_post_process_func",
-]
@@ -1,571 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import math
-from collections.abc import Iterable
-from typing import TYPE_CHECKING
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.models.embeddings import apply_rotary_emb, get_1d_rotary_pos_embed
-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.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.attention.layer import Attention
-
-if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-logger = init_logger(__name__)
-
-# Devices that do not support float64; RoPE frequencies fall back to float32 there.
-_FP64_UNSUPPORTED_DEVICES = frozenset({"mps", "npu", "neuron"})
-
-
-def _rope_freqs_dtype(device: torch.device) -> torch.dtype:
-    return torch.float32 if device.type in _FP64_UNSUPPORTED_DEVICES else torch.float64
-
-
-def _join_prefix(prefix: str, name: str) -> str:
-    return f"{prefix}.{name}" if prefix else name
-
-
-def _linear(
-    in_features: int,
-    out_features: int,
-    bias: bool,
-    quant_config: "QuantizationConfig | None",
-    prefix: str,
-) -> ReplicatedLinear:
-    # ReplicatedLinear (not nn.Linear) so the diffusion LoRA manager can wrap these projections.
-    return ReplicatedLinear(
-        in_features,
-        out_features,
-        bias=bias,
-        quant_config=quant_config,
-        prefix=prefix,
-        return_bias=False,
-    )
-
-
-class Krea2RMSNorm(nn.Module):
-    """RMSNorm with a zero-centered scale: the effective multiplier is ``1 + weight``, matching the Krea 2 checkpoint
-    format. Activations are upcast so the normalization runs in float32."""
-
-    def __init__(self, dim: int, eps: float = 1e-5) -> None:
-        super().__init__()
-        self.dim = dim
-        self.eps = eps
-        self.weight = nn.Parameter(torch.zeros(dim))
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        dtype = hidden_states.dtype
-        hidden_states = F.rms_norm(hidden_states.float(), (self.dim,), weight=self.weight + 1.0, eps=self.eps)
-        return hidden_states.to(dtype)
-
-
-class Krea2Attention(nn.Module):
-    """Self-attention with grouped-query projections, q/k RMSNorm, rotary embeddings and a sigmoid output gate.
-
-    Q/K/V layout is ``[B, seq, heads, head_dim]``; ``attention_mask`` is a 2D boolean key-padding mask
-    ``(batch, key_seq_len)`` that the backend broadcasts.
-    """
-
-    def __init__(
-        self,
-        hidden_size: int,
-        num_heads: int,
-        num_kv_heads: int | None = None,
-        eps: float = 1e-5,
-        role: str = "self",
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        if hidden_size % num_heads != 0:
-            raise ValueError(f"hidden_size={hidden_size} must be divisible by num_heads={num_heads}")
-        self.hidden_size = hidden_size
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
-        self.head_dim = hidden_size // num_heads
-
-        q_dim = self.head_dim * self.num_heads
-        kv_dim = self.head_dim * self.num_kv_heads
-        self.to_q = _linear(hidden_size, q_dim, False, quant_config, _join_prefix(prefix, "to_q"))
-        self.to_k = _linear(hidden_size, kv_dim, False, quant_config, _join_prefix(prefix, "to_k"))
-        self.to_v = _linear(hidden_size, kv_dim, False, quant_config, _join_prefix(prefix, "to_v"))
-        self.to_gate = _linear(hidden_size, hidden_size, False, quant_config, _join_prefix(prefix, "to_gate"))
-        self.norm_q = Krea2RMSNorm(self.head_dim, eps=eps)
-        self.norm_k = Krea2RMSNorm(self.head_dim, eps=eps)
-        self.to_out = nn.ModuleList(
-            [_linear(hidden_size, hidden_size, False, quant_config, _join_prefix(prefix, "to_out.0")), nn.Dropout(0.0)]
-        )
-
-        self.attn = Attention(
-            num_heads=self.num_heads,
-            head_size=self.head_dim,
-            softmax_scale=1.0 / (self.head_dim**0.5),
-            causal=False,
-            num_kv_heads=self.num_kv_heads,
-            role=role,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
-    ) -> torch.Tensor:
-        query = self.to_q(hidden_states).unflatten(-1, (self.num_heads, self.head_dim))
-        key = self.to_k(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim))
-        value = self.to_v(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim))
-        gate = self.to_gate(hidden_states)
-
-        query = self.norm_q(query)
-        key = self.norm_k(key)
-
-        if image_rotary_emb is not None:
-            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
-            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
-        query = query.to(value.dtype)
-        key = key.to(value.dtype)
-
-        attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
-        hidden_states = self.attn(query, key, value, attn_metadata)
-        hidden_states = hidden_states.flatten(2, 3)
-        hidden_states = hidden_states * torch.sigmoid(gate)
-        return self.to_out[0](hidden_states)
-
-
-class Krea2SwiGLU(nn.Module):
-    """SwiGLU feed-forward network."""
-
-    def __init__(
-        self, dim: int, hidden_dim: int, quant_config: "QuantizationConfig | None" = None, prefix: str = ""
-    ) -> None:
-        super().__init__()
-        self.gate = _linear(dim, hidden_dim, False, quant_config, _join_prefix(prefix, "gate"))
-        self.up = _linear(dim, hidden_dim, False, quant_config, _join_prefix(prefix, "up"))
-        self.down = _linear(hidden_dim, dim, False, quant_config, _join_prefix(prefix, "down"))
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return self.down(F.silu(self.gate(hidden_states)) * self.up(hidden_states))
-
-
-class Krea2TextFusionBlock(nn.Module):
-    """Pre-norm transformer block (no rotary embeddings, no time modulation) used by the text fusion stage."""
-
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        num_kv_heads: int,
-        intermediate_size: int,
-        eps: float,
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.norm1 = Krea2RMSNorm(dim, eps=eps)
-        self.norm2 = Krea2RMSNorm(dim, eps=eps)
-        self.attn = Krea2Attention(
-            dim,
-            num_heads,
-            num_kv_heads,
-            eps=eps,
-            role="self",
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "attn"),
-        )
-        self.ff = Krea2SwiGLU(dim, intermediate_size, quant_config=quant_config, prefix=_join_prefix(prefix, "ff"))
-
-    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
-        hidden_states = hidden_states + self.attn(self.norm1(hidden_states), attention_mask=attention_mask)
-        hidden_states = hidden_states + self.ff(self.norm2(hidden_states))
-        return hidden_states
-
-
-class Krea2TextFusion(nn.Module):
-    """Fuses the stack of tapped text-encoder hidden states into a single sequence of text features.
-
-    Two ``layerwise_blocks`` attend across the ``num_text_layers`` axis independently for every token, a linear
-    ``projector`` collapses that axis, and two ``refiner_blocks`` attend across the token sequence.
-    """
-
-    def __init__(
-        self,
-        num_text_layers: int,
-        dim: int,
-        num_heads: int,
-        num_kv_heads: int,
-        intermediate_size: int,
-        num_layerwise_blocks: int,
-        num_refiner_blocks: int,
-        eps: float,
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.layerwise_blocks = nn.ModuleList(
-            [
-                Krea2TextFusionBlock(
-                    dim,
-                    num_heads,
-                    num_kv_heads,
-                    intermediate_size,
-                    eps,
-                    quant_config=quant_config,
-                    prefix=_join_prefix(prefix, f"layerwise_blocks.{i}"),
-                )
-                for i in range(num_layerwise_blocks)
-            ]
-        )
-        self.projector = _linear(num_text_layers, 1, False, quant_config, _join_prefix(prefix, "projector"))
-        self.refiner_blocks = nn.ModuleList(
-            [
-                Krea2TextFusionBlock(
-                    dim,
-                    num_heads,
-                    num_kv_heads,
-                    intermediate_size,
-                    eps,
-                    quant_config=quant_config,
-                    prefix=_join_prefix(prefix, f"refiner_blocks.{i}"),
-                )
-                for i in range(num_refiner_blocks)
-            ]
-        )
-
-    def forward(self, encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
-        batch_size, seq_len, num_text_layers, dim = encoder_hidden_states.shape
-
-        hidden_states = encoder_hidden_states.reshape(batch_size * seq_len, num_text_layers, dim)
-        for block in self.layerwise_blocks:
-            hidden_states = block(hidden_states.contiguous())
-
-        hidden_states = hidden_states.reshape(batch_size, seq_len, num_text_layers, dim).permute(0, 1, 3, 2)
-        hidden_states = self.projector(hidden_states).squeeze(-1)
-
-        for block in self.refiner_blocks:
-            hidden_states = block(hidden_states, attention_mask=attention_mask)
-
-        return hidden_states
-
-
-class Krea2TransformerBlock(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        intermediate_size: int,
-        num_heads: int,
-        num_kv_heads: int,
-        norm_eps: float,
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.scale_shift_table = nn.Parameter(torch.zeros(6, hidden_size))
-        self.norm1 = Krea2RMSNorm(hidden_size, eps=norm_eps)
-        self.norm2 = Krea2RMSNorm(hidden_size, eps=norm_eps)
-        self.attn = Krea2Attention(
-            hidden_size,
-            num_heads,
-            num_kv_heads,
-            eps=norm_eps,
-            role="self",
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "attn"),
-        )
-        self.ff = Krea2SwiGLU(
-            hidden_size, intermediate_size, quant_config=quant_config, prefix=_join_prefix(prefix, "ff")
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        # temb: (B, 1, 6 * hidden_size), shared across all blocks; each block only learns an additive table.
-        modulation = temb.unflatten(-1, (6, -1)) + self.scale_shift_table
-        prescale, preshift, pregate, postscale, postshift, postgate = modulation.unbind(-2)
-
-        attn_out = self.attn(
-            (1.0 + prescale) * self.norm1(hidden_states) + preshift,
-            attention_mask=attention_mask,
-            image_rotary_emb=image_rotary_emb,
-        )
-        hidden_states = hidden_states + pregate * attn_out
-        ff_out = self.ff((1.0 + postscale) * self.norm2(hidden_states) + postshift)
-        hidden_states = hidden_states + postgate * ff_out
-        return hidden_states
-
-
-class Krea2TimestepEmbedding(nn.Module):
-    """Sinusoidal flow-time embedding (cos-first, input scaled by 1000) followed by a two-layer MLP.
-
-    Keeps the sequence dimension at size 1 so the per-block modulations broadcast over tokens.
-    """
-
-    def __init__(
-        self, embed_dim: int, hidden_size: int, quant_config: "QuantizationConfig | None" = None, prefix: str = ""
-    ) -> None:
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.linear_1 = _linear(embed_dim, hidden_size, True, quant_config, _join_prefix(prefix, "linear_1"))
-        self.linear_2 = _linear(hidden_size, hidden_size, True, quant_config, _join_prefix(prefix, "linear_2"))
-
-    def forward(self, timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
-        half = self.embed_dim // 2
-        freqs = torch.exp(-math.log(1e4) * torch.arange(half, dtype=torch.float32, device=timestep.device) / half)
-        args = (timestep.float() * 1e3)[:, None, None] * freqs
-        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(dtype)
-        return self.linear_2(F.gelu(self.linear_1(emb), approximate="tanh"))
-
-
-class Krea2TextProjection(nn.Module):
-    """Projects the fused text features into the transformer width."""
-
-    def __init__(
-        self,
-        text_dim: int,
-        hidden_size: int,
-        eps: float,
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.norm = Krea2RMSNorm(text_dim, eps=eps)
-        self.linear_1 = _linear(text_dim, hidden_size, True, quant_config, _join_prefix(prefix, "linear_1"))
-        self.linear_2 = _linear(hidden_size, hidden_size, True, quant_config, _join_prefix(prefix, "linear_2"))
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.linear_1(self.norm(hidden_states))
-        return self.linear_2(F.gelu(hidden_states, approximate="tanh"))
-
-
-class Krea2FinalLayer(nn.Module):
-    """Final adaptive RMSNorm and output projection."""
-
-    def __init__(
-        self,
-        hidden_size: int,
-        out_channels: int,
-        eps: float,
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.scale_shift_table = nn.Parameter(torch.zeros(2, hidden_size))
-        self.norm = Krea2RMSNorm(hidden_size, eps=eps)
-        self.linear = _linear(hidden_size, out_channels, True, quant_config, _join_prefix(prefix, "linear"))
-
-    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
-        modulation = temb + self.scale_shift_table
-        scale, shift = modulation.chunk(2, dim=1)
-        hidden_states = (1.0 + scale) * self.norm(hidden_states) + shift
-        return self.linear(hidden_states)
-
-
-class Krea2RotaryPosEmbed(nn.Module):
-    """Multi-axis (t, h, w) rotary position embedding, following the Flux/Krea 2 convention."""
-
-    def __init__(self, theta: int, axes_dim: list[int]):
-        super().__init__()
-        self.theta = theta
-        self.axes_dim = axes_dim
-
-    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        n_axes = ids.shape[-1]
-        cos_out = []
-        sin_out = []
-        pos = ids.float()
-        freqs_dtype = _rope_freqs_dtype(ids.device)
-        for i in range(n_axes):
-            cos, sin = get_1d_rotary_pos_embed(
-                self.axes_dim[i],
-                pos[:, i],
-                theta=self.theta,
-                repeat_interleave_real=True,
-                use_real=True,
-                freqs_dtype=freqs_dtype,
-            )
-            cos_out.append(cos)
-            sin_out.append(sin)
-        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
-        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
-        return freqs_cos, freqs_sin
-
-
-class Krea2Transformer2DModel(nn.Module):
-    r"""The single-stream MMDiT flow-matching backbone used by the Krea 2 pipeline (vLLM-Omni port).
-
-    Text conditioning enters as a stack of hidden states tapped from several layers of a multimodal text encoder. A
-    small text-fusion transformer collapses the layer axis and refines the token sequence; the result is concatenated
-    with the patchified image latents into a single ``[text, image]`` sequence processed by the transformer blocks.
-    The timestep conditions every block through one shared modulation vector plus per-block learned tables.
-    """
-
-    _repeated_blocks = ["Krea2TransformerBlock"]
-    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
-
-    @staticmethod
-    def _is_transformer_block(name: str, module) -> bool:
-        return "transformer_blocks" in name and name.split(".")[-1].isdigit()
-
-    _hsdp_shard_conditions = [_is_transformer_block]
-
-    def __init__(
-        self,
-        in_channels: int = 64,
-        num_layers: int = 28,
-        attention_head_dim: int = 128,
-        num_attention_heads: int = 48,
-        num_key_value_heads: int = 12,
-        intermediate_size: int = 16384,
-        timestep_embed_dim: int = 256,
-        text_hidden_dim: int = 2560,
-        num_text_layers: int = 12,
-        text_num_attention_heads: int = 20,
-        text_num_key_value_heads: int = 20,
-        text_intermediate_size: int = 6912,
-        num_layerwise_text_blocks: int = 2,
-        num_refiner_text_blocks: int = 2,
-        axes_dims_rope: tuple[int, int, int] = (32, 48, 48),
-        rope_theta: float = 1000.0,
-        norm_eps: float = 1e-5,
-        od_config: "OmniDiffusionConfig | None" = None,
-        quant_config: "QuantizationConfig | None" = None,
-    ) -> None:
-        super().__init__()
-        # The pipeline reads ``self.dtype`` to cast inputs before the transformer forward.
-        self.dtype = od_config.dtype if od_config is not None else torch.get_default_dtype()
-        self.od_config = od_config
-
-        hidden_size = attention_head_dim * num_attention_heads
-        if sum(axes_dims_rope) != attention_head_dim:
-            raise ValueError(
-                f"sum(axes_dims_rope)={sum(axes_dims_rope)} must equal attention_head_dim={attention_head_dim}"
-            )
-
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.hidden_size = hidden_size
-        self.num_layers = num_layers
-
-        self.img_in = _linear(in_channels, hidden_size, True, quant_config, "img_in")
-        self.time_embed = Krea2TimestepEmbedding(timestep_embed_dim, hidden_size, quant_config, "time_embed")
-        self.time_mod_proj = _linear(hidden_size, 6 * hidden_size, True, quant_config, "time_mod_proj")
-        self.text_fusion = Krea2TextFusion(
-            num_text_layers=num_text_layers,
-            dim=text_hidden_dim,
-            num_heads=text_num_attention_heads,
-            num_kv_heads=text_num_key_value_heads,
-            intermediate_size=text_intermediate_size,
-            num_layerwise_blocks=num_layerwise_text_blocks,
-            num_refiner_blocks=num_refiner_text_blocks,
-            eps=norm_eps,
-            quant_config=quant_config,
-            prefix="text_fusion",
-        )
-        self.txt_in = Krea2TextProjection(text_hidden_dim, hidden_size, norm_eps, quant_config, "txt_in")
-        self.rotary_emb = Krea2RotaryPosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope))
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                Krea2TransformerBlock(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_heads=num_attention_heads,
-                    num_kv_heads=num_key_value_heads,
-                    norm_eps=norm_eps,
-                    quant_config=quant_config,
-                    prefix=f"transformer_blocks.{i}",
-                )
-                for i in range(num_layers)
-            ]
-        )
-
-        self.final_layer = Krea2FinalLayer(hidden_size, in_channels, norm_eps, quant_config, "final_layer")
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.Tensor,
-        position_ids: torch.Tensor,
-        encoder_attention_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        r"""Predict the flow-matching velocity for the image tokens.
-
-        Args:
-            hidden_states: Packed (patchified) noisy image latents ``(batch, image_seq_len, in_channels)``.
-            encoder_hidden_states: Tapped text-encoder hidden states
-                ``(batch, text_seq_len, num_text_layers, text_hidden_dim)``.
-            timestep: Flow-matching time in ``[0, 1]`` of shape ``(batch,)``.
-            position_ids: ``(t, h, w)`` rotary coordinates ``(text_seq_len + image_seq_len, 3)``.
-            encoder_attention_mask: Boolean mask marking valid text tokens ``(batch, text_seq_len)`` or ``None``.
-        """
-        if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
-            raise ValueError(f"`position_ids` must have shape (sequence_length, 3), got {tuple(position_ids.shape)}.")
-
-        batch_size, image_seq_len, _ = hidden_states.shape
-
-        temb = self.time_embed(timestep, dtype=hidden_states.dtype)
-        temb_mod = self.time_mod_proj(F.gelu(temb, approximate="tanh"))
-
-        # 2D key-padding masks: padded text tokens are excluded as keys; their lanes are dropped at the output slice.
-        text_attention_mask = None
-        attention_mask = None
-        if encoder_attention_mask is not None:
-            text_attention_mask = encoder_attention_mask.bool()
-            image_mask = text_attention_mask.new_ones((batch_size, image_seq_len))
-            attention_mask = torch.cat([text_attention_mask, image_mask], dim=1)
-
-        encoder_hidden_states = self.text_fusion(encoder_hidden_states, attention_mask=text_attention_mask)
-        encoder_hidden_states = self.txt_in(encoder_hidden_states)
-
-        text_seq_len = encoder_hidden_states.shape[1]
-        hidden_states = self.img_in(hidden_states)
-        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
-
-        image_rotary_emb = self.rotary_emb(position_ids)
-
-        for block in self.transformer_blocks:
-            hidden_states = block(hidden_states, temb_mod, image_rotary_emb, attention_mask)
-
-        hidden_states = hidden_states[:, text_seq_len:]
-        return self.final_layer(hidden_states, temb)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # Module/parameter names mirror the diffusers ``Krea2Transformer2DModel`` exactly, so weights map 1:1.
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if name not in params_dict:
-                continue
-            param = params_dict[name]
-            weight_loader = getattr(param, "weight_loader", default_weight_loader)
-            weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
@@ -1,551 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-import logging
-import os
-from collections.abc import Iterable
-from typing import ClassVar
-
-import numpy as np
-import torch
-from diffusers.image_processor import VaeImageProcessor
-from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler,
-)
-from diffusers.utils.torch_utils import randn_tensor
-from torch import nn
-from transformers import AutoTokenizer, Qwen3VLModel
-from vllm.model_executor.models.utils import AutoWeightsLoader
-from vllm.transformers_utils.config import get_hf_file_to_dict
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
-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.interface import SupportsComponentDiscovery
-from vllm_omni.diffusion.models.krea2.krea2_transformer import Krea2Transformer2DModel
-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
-
-logger = logging.getLogger(__name__)
-
-# Text-encoder layer taps used when model_index.json does not specify ``text_encoder_select_layers``.
-DEFAULT_TEXT_ENCODER_SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
-
-
-def get_krea2_post_process_func(od_config: OmniDiffusionConfig):
-    model_name = od_config.model
-    # Read only the VAE config (``get_hf_file_to_dict`` resolves a local dir via
-    # ``Path(model)/file`` and a hub repo via a single-file download) instead of
-    # pulling all weights just to compute ``vae_scale_factor``.
-    vae_config = get_hf_file_to_dict("vae/config.json", model_name) or {}
-    vae_scale_factor = 2 ** len(vae_config["temperal_downsample"]) if "temperal_downsample" in vae_config else 8
-
-    patch_size = 2
-    try:
-        model_index = get_hf_file_to_dict("model_index.json", model_name) or {}
-        patch_size = int(model_index.get("patch_size", patch_size))
-    except Exception:  # noqa: BLE001
-        pass
-
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * patch_size)
-
-    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,
-) -> tuple[torch.Tensor, int]:
-    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 Krea2Pipeline(nn.Module, DiffusionPipelineProfilerMixin, ProgressBarMixin, SupportsComponentDiscovery):
-    """The Krea 2 text-to-image pipeline (vLLM-Omni port of ``diffusers.Krea2Pipeline``).
-
-    Components:
-    - ``scheduler``: ``FlowMatchEulerDiscreteScheduler`` (resolution-aware exponential time shift).
-    - ``vae``: the Qwen-Image VAE (f8, 16 latent channels).
-    - ``text_encoder``: a Qwen3-VL model; the pipeline consumes a stack of hidden states tapped from several
-      decoder layers rather than the last hidden state.
-    - ``tokenizer``: the tokenizer paired with the text encoder.
-    - ``transformer``: the Krea 2 single-stream MMDiT that predicts the flow-matching velocity.
-    """
-
-    _dit_modules: ClassVar[list[str]] = ["transformer"]
-    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
-    _vae_modules: ClassVar[list[str]] = ["vae"]
-
-    _PROFILER_TARGETS: ClassVar[list[str]] = ["text_encoder.forward", "diffuse", "vae.decode"]
-
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        super().__init__()
-        self.od_config = od_config
-        self.parallel_config = od_config.parallel_config
-        # Only the transformer goes through the standard weights loader; text encoder and VAE load eagerly below.
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model,
-                subfolder="transformer",
-                revision=od_config.revision,
-                prefix="transformer.",
-                fall_back_to_pt=True,
-            )
-        ]
-
-        self.device = get_local_device()
-        model = od_config.model
-        local_files_only = os.path.isdir(model)
-
-        # Pipeline-level config lives in model_index.json, not the transformer config.
-        model_index = get_hf_file_to_dict("model_index.json", model) or {}
-        select_layers = model_index.get("text_encoder_select_layers") or DEFAULT_TEXT_ENCODER_SELECT_LAYERS
-        self.text_encoder_select_layers = tuple(int(i) for i in select_layers)
-        self.is_distilled = bool(model_index.get("is_distilled", False))
-        self.patch_size = int(model_index.get("patch_size", 2))
-
-        subfolders = ["scheduler", "text_encoder", "vae", "tokenizer"]
-        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
-        prefetch_subfolders(model, subfolders, local_files_only=local_files_only)
-
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            model, subfolder="scheduler", local_files_only=local_files_only
-        )
-
-        self.text_encoder = from_pretrained_with_prefetch(
-            Qwen3VLModel.from_pretrained,
-            model,
-            subfolder="text_encoder",
-            prefetch_list=subfolders,
-            local_files_only=local_files_only,
-            torch_dtype=od_config.dtype,
-        )
-        # Drop the unused Qwen3-VL vision tower before moving to GPU so it never consumes GPU memory.
-        if hasattr(self.text_encoder, "visual"):
-            del self.text_encoder.visual
-        else:
-            logger.warning("Krea2: vision tower not found on text encoder; skipping drop")
-        self.text_encoder = self.text_encoder.to(self.device)
-
-        self.vae = from_pretrained_with_prefetch(
-            DistributedAutoencoderKLQwenImage.from_pretrained,
-            model,
-            subfolder="vae",
-            prefetch_list=subfolders,
-            local_files_only=local_files_only,
-            torch_dtype=od_config.dtype,
-        ).to(self.device)
-
-        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Krea2Transformer2DModel)
-        self.transformer = Krea2Transformer2DModel(
-            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
-        )
-
-        self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-
-        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * self.patch_size)
-        self.default_sample_size = 1024
-
-        # Chat template tokenized as a fixed-length block: prompt padded to a fixed length first, assistant suffix
-        # appended after the padding; the first ``prompt_template_encode_start_idx`` system tokens are later dropped.
-        self.prompt_template_encode_prefix = (
-            "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, "
-            "spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"
-        )
-        self.prompt_template_encode_suffix = "<|im_end|>\n<|im_start|>assistant\n"
-        self.prompt_template_encode_start_idx = 34
-        self.prompt_template_encode_num_suffix_tokens = 5
-
-        self.setup_diffusion_pipeline_profiler(
-            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
-        )
-
-    def get_text_hidden_states(
-        self,
-        prompt: str | list[str],
-        max_sequence_length: int = 512,
-        device: torch.device | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Tokenize ``prompt`` into the fixed-length Krea 2 layout and tap the selected encoder hidden states.
-
-        Returns ``(hidden_states, attention_mask)`` of shapes ``(batch, text_seq_len, num_text_layers,
-        text_hidden_dim)`` and ``(batch, text_seq_len)`` (bool).
-        """
-        device = device or self.device
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        prefix_idx = self.prompt_template_encode_start_idx
-        text = [self.prompt_template_encode_prefix + e for e in prompt]
-        text_tokens = self.tokenizer(
-            text,
-            truncation=True,
-            padding="max_length",
-            max_length=max_sequence_length + prefix_idx - self.prompt_template_encode_num_suffix_tokens,
-            return_tensors="pt",
-        ).to(device)
-        suffix_tokens = self.tokenizer([self.prompt_template_encode_suffix] * len(text), return_tensors="pt").to(device)
-
-        input_ids = torch.cat([text_tokens.input_ids, suffix_tokens.input_ids], dim=1)
-        attention_mask = torch.cat([text_tokens.attention_mask, suffix_tokens.attention_mask], dim=1).bool()
-
-        # Padding sits in the middle (``[prefix | prompt | PAD | suffix]``), so positions must count only real tokens
-        # (padding consumes no position). Broadcast across the 3 mRoPE axes (T/H/W equal for text).
-        position_ids = (attention_mask.long().cumsum(dim=-1) - 1).clamp(min=0)
-        position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
-
-        outputs = self.text_encoder(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            output_hidden_states=True,
-        )
-        hidden_states = torch.stack([outputs.hidden_states[i] for i in self.text_encoder_select_layers], dim=2)
-
-        hidden_states = hidden_states[:, prefix_idx:]
-        attention_mask = attention_mask[:, prefix_idx:]
-        return hidden_states, attention_mask
-
-    def encode_prompt(
-        self,
-        prompt: str | list[str],
-        device: torch.device | None = None,
-        num_images_per_prompt: int = 1,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        max_sequence_length: int = 512,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        device = device or self.device
-
-        if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self.get_text_hidden_states(prompt, max_sequence_length, device)
-
-        batch_size, seq_len, num_text_layers, dim = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, num_text_layers, dim)
-        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt)
-        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
-
-        return prompt_embeds, prompt_embeds_mask
-
-    def _pack_latents(self, latents, batch_size, num_channels_latents, height, width):
-        p = self.patch_size
-        latents = latents.view(batch_size, num_channels_latents, height // p, p, width // p, p)
-        latents = latents.permute(0, 2, 4, 1, 3, 5)
-        latents = latents.reshape(batch_size, (height // p) * (width // p), num_channels_latents * p * p)
-        return latents
-
-    def _unpack_latents(self, latents, height, width):
-        batch_size, _, channels = latents.shape
-        p = self.patch_size
-
-        height = p * (int(height) // (self.vae_scale_factor * p))
-        width = p * (int(width) // (self.vae_scale_factor * p))
-
-        latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p)
-        latents = latents.permute(0, 3, 1, 4, 2, 5)
-        latents = latents.reshape(batch_size, channels // (p * p), 1, height, width)
-        return latents
-
-    @staticmethod
-    def prepare_position_ids(text_seq_len: int, grid_height: int, grid_width: int, device: torch.device):
-        """Build the ``(text_seq_len + grid_height * grid_width, 3)`` rotary coordinates: text tokens sit at the
-        origin, image tokens carry their ``(0, h, w)`` latent-grid coordinates."""
-        text_ids = torch.zeros(text_seq_len, 3, device=device)
-        image_ids = torch.zeros(grid_height, grid_width, 3, device=device)
-        image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None]
-        image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :]
-        image_ids = image_ids.reshape(grid_height * grid_width, 3)
-        return torch.cat([text_ids, image_ids], dim=0)
-
-    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-
-        latent_height = height // self.vae_scale_factor
-        latent_width = width // self.vae_scale_factor
-        shape = (batch_size, num_channels_latents, latent_height, latent_width)
-
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        latents = self._pack_latents(latents, batch_size, num_channels_latents, latent_height, latent_width)
-        return latents
-
-    def _extract_prompts(self, prompts):
-        """Extract prompt and negative_prompt from the request's 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
-
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 0
-
-    def diffuse(
-        self,
-        latents: torch.Tensor,
-        timesteps: torch.Tensor,
-        position_ids: torch.Tensor,
-        prompt_embeds: torch.Tensor,
-        prompt_embeds_mask: torch.Tensor,
-        negative_prompt_embeds: torch.Tensor | None,
-        negative_prompt_embeds_mask: torch.Tensor | None,
-        guidance_scale: float,
-    ) -> torch.Tensor:
-        """Run the flow-matching denoising loop and return the final latents."""
-        self.scheduler.set_begin_index(0)
-        transformer_dtype = self.transformer.dtype
-        with self.progress_bar(total=len(timesteps)) as progress_bar:
-            for _, t in enumerate(timesteps):
-                timestep = (t / self.scheduler.config.num_train_timesteps).expand(latents.shape[0])
-                timestep = timestep.to(transformer_dtype)
-                latent_input = latents.to(transformer_dtype)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_input,
-                    encoder_hidden_states=prompt_embeds.to(transformer_dtype),
-                    timestep=timestep,
-                    position_ids=position_ids,
-                    encoder_attention_mask=prompt_embeds_mask,
-                )
-
-                if self.do_classifier_free_guidance:
-                    neg_noise_pred = self.transformer(
-                        hidden_states=latent_input,
-                        encoder_hidden_states=negative_prompt_embeds.to(transformer_dtype),
-                        timestep=timestep,
-                        position_ids=position_ids,
-                        encoder_attention_mask=negative_prompt_embeds_mask,
-                    )
-                    noise_pred = noise_pred + guidance_scale * (noise_pred - neg_noise_pred)
-
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-                progress_bar.update()
-
-        return latents
-
-    def forward(
-        self,
-        req: OmniDiffusionRequest,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        height: int = 1024,
-        width: int = 1024,
-        num_inference_steps: int = 28,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 4.5,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        max_sequence_length: int = 512,
-    ) -> DiffusionOutput:
-        extracted_prompt, extracted_negative = self._extract_prompts(req.prompts)
-        prompt = extracted_prompt or prompt
-        if extracted_negative is not None:
-            negative_prompt = extracted_negative
-
-        sp = req.sampling_params
-        height = sp.height or height
-        width = sp.width or width
-        num_inference_steps = sp.num_inference_steps or num_inference_steps
-        sigmas = sp.sigmas or sigmas
-        max_sequence_length = sp.max_sequence_length or max_sequence_length
-        generator = sp.generator or generator
-        if sp.guidance_scale_provided:
-            guidance_scale = sp.guidance_scale
-        else:
-            # Distilled checkpoint is trained guidance-free (and the request layer coerces an explicit
-            # guidance_scale=0 to "not provided"), so default it to no CFG.
-            guidance_scale = 0.0 if self.is_distilled else guidance_scale
-        num_images_per_prompt = sp.num_outputs_per_prompt if sp.num_outputs_per_prompt > 0 else num_images_per_prompt
-
-        multiple = self.vae_scale_factor * self.patch_size
-        if height % multiple != 0 or width % multiple != 0:
-            rounded_height = ((height + multiple - 1) // multiple) * multiple
-            rounded_width = ((width + multiple - 1) // multiple) * multiple
-            logger.warning(
-                "`height` and `width` must be multiples of %d; rounding up from %dx%d to %dx%d.",
-                multiple,
-                height,
-                width,
-                rounded_height,
-                rounded_width,
-            )
-            height, width = rounded_height, rounded_width
-
-        device = self.device
-        self._guidance_scale = guidance_scale
-
-        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)
-        else:
-            batch_size = 1
-
-        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
-            prompt=prompt,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            max_sequence_length=max_sequence_length,
-        )
-        negative_prompt_embeds = None
-        negative_prompt_embeds_mask = None
-        if self.do_classifier_free_guidance:
-            if negative_prompt is None:
-                negative_prompt = ""
-            if isinstance(negative_prompt, str):
-                negative_prompt = [negative_prompt] * batch_size
-            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
-                prompt=negative_prompt,
-                device=device,
-                num_images_per_prompt=num_images_per_prompt,
-                max_sequence_length=max_sequence_length,
-            )
-
-        num_channels_latents = self.transformer.in_channels // (self.patch_size**2)
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-        grid_height = height // (self.vae_scale_factor * self.patch_size)
-        grid_width = width // (self.vae_scale_factor * self.patch_size)
-        position_ids = self.prepare_position_ids(prompt_embeds.shape[1], grid_height, grid_width, device)
-
-        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
-        image_seq_len = latents.shape[1]
-        if self.is_distilled:
-            mu = 1.15
-        else:
-            mu = calculate_shift(
-                image_seq_len,
-                self.scheduler.config.get("base_image_seq_len", 256),
-                self.scheduler.config.get("max_image_seq_len", 6400),
-                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, device, sigmas=sigmas, mu=mu
-        )
-
-        latents = self.diffuse(
-            latents=latents,
-            timesteps=timesteps,
-            position_ids=position_ids,
-            prompt_embeds=prompt_embeds,
-            prompt_embeds_mask=prompt_embeds_mask,
-            negative_prompt_embeds=negative_prompt_embeds,
-            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
-            guidance_scale=guidance_scale,
-        )
-
-        if output_type == "latent":
-            image = latents
-        else:
-            latents = self._unpack_latents(latents, height, width)
-            latents = latents.to(self.vae.dtype)
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
-                latents.device, latents.dtype
-            )
-            latents = latents / latents_std + latents_mean
-            image = self.vae.decode(latents, return_dict=False)[0][:, :, 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)
-        loaded_weights = loader.load_weights(weights)
-        # VAE and text encoder load eagerly in __init__; record their params so strict load checks pass.
-        loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
-        if self.text_encoder is not None:
-            loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
-        return loaded_weights
@@ -306,11 +306,6 @@ _DIFFUSION_MODELS = {
        "pipeline_sdxl",
        "StableDiffusionXLPipeline",
    ),
-    "Krea2Pipeline": (
-        "krea2",
-        "pipeline_krea2",
-        "Krea2Pipeline",
-    ),
 }


@@ -538,7 +533,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
    "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func",
    "HiDreamImagePipeline": "get_hidream_image_post_process_func",
    "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func",
-    "Krea2Pipeline": "get_krea2_post_process_func",
 }

 _DIFFUSION_ACTION_POST_PROCESS_FUNCS = {
@@ -1568,18 +1568,16 @@ class AsyncOmniEngine:
        # and can cause CUDA OOM for subsequent engine instances (especially
        # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB).
        #
-        # CuMemAllocator.sleep() is NOT idempotent — calling it on already-
-        # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator
-        # cuMemRelease (double-free of the memory handle).  Use release_pools()
-        # instead, which is the designed cleanup path: it drops MemPool refs
-        # and lets the destructor/free path handle asleep entries correctly
-        # (returns a null handle so the C extension skips unmap/release).
+        # Discard mode (level=2) is correct at shutdown: there is no benefit to
+        # keeping a CPU backup when the engine is being torn down.
        try:
            from vllm.device_allocator.cumem import CuMemAllocator, cumem_available

            if cumem_available:
                allocator = CuMemAllocator.get_instance()
-                allocator.release_pools()
+                # Sleep at level 2 discards all pool memory from the GPU
+                # without creating CPU backups — cheapest and fastest.
+                allocator.sleep()
                logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown")
        except Exception:
            pass
@@ -28,12 +28,6 @@ from vllm_omni.outputs import OmniRequestOutput

 logger = init_logger(__name__)

-# Keys whose values may arrive as scalar tensors (e.g. sample rate) but
-# represent metadata, not concatenatable data.  When these appear in
-# ``MultimodalPayload.tensors`` they must be relocated to ``.metadata``
-# before consolidation so the tensor concatenation path does not attempt
-# a dimensioned cat on a 0-d tensor.
-_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"})

 # ---------------------------------------------------------------------------
 # Module-level helpers
@@ -218,18 +212,6 @@ class OmniRequestState(RequestState):
            modality = OutputModality.TEXT
        strategy = get_accumulation_strategy(modality)

-        # Move known metadata scalar tensors (e.g. sample rate) from
-        # tensors to metadata so they are handled by the REPLACE path
-        # below instead of failing tensor concatenation.
-        for key in _METADATA_TENSOR_KEYS:
-            v = self.mm_accumulated.tensors.pop(key, None)
-            if v is not None:
-                if isinstance(v, list):
-                    # Keep the last value (REPLACE semantics)
-                    self.mm_accumulated.metadata[key] = v[-1]
-                else:
-                    self.mm_accumulated.metadata[key] = v
-
        try:
            for k, v in list(self.mm_accumulated.tensors.items()):
                if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
@@ -138,21 +138,6 @@ class StageReplicaResources:
 # ---------------------------------------------------------------------------


-def _port_from_zmq_address(address: str | None) -> int | None:
-    """Extract the TCP port from a ``tcp://host:port`` ZMQ address, else ``None``.
-
-    Non-TCP transports (``ipc://``, ``inproc://``) and unparsable values return
-    ``None`` so callers can skip seeding them into the port-dedup set.
-    """
-    if not address:
-        return None
-    tail = address.rsplit(":", 1)[-1]
-    try:
-        return int(tail)
-    except ValueError:
-        return None
-
-
 class OmniMasterServer:
    """Registration server for single-stage engine startup."""

@@ -173,25 +158,10 @@ class OmniMasterServer:
        self._stage_configs: dict[StageRoute, Any] = {}
        self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {}
        self._stage_config_events: dict[StageRoute, threading.Event] = {}
-        # Ports already handed out by *this* server. ``get_open_ports_list`` only
-        # guarantees uniqueness *within* a single call; a per-route call cannot
-        # see ports drawn for other routes, so two stages/replicas can draw the
-        # same ephemeral port and the second engine to ``bind()`` it dies with
-        # ``zmq.error.ZMQError: Address already in use``. Multi-stage models
-        # (e.g. Qwen3-Omni: thinker/talker/code2wav) allocate many routes and hit
-        # this intermittently. We dedup every allocated port against this set,
-        # seeded with the registration port (and coordinator ROUTER port) that
-        # are already bound on the same host.
-        self._allocated_ports: set[int] = set()
-        if master_port:
-            self._allocated_ports.add(int(master_port))
        # Coordinator ROUTER address echoed back in every registration reply
        # so OmniCoordClientForStage knows where to connect from inside the
        # engine subprocess.
        self._coordinator_router_address = coordinator_router_address
-        coord_port = _port_from_zmq_address(coordinator_router_address)
-        if coord_port is not None:
-            self._allocated_ports.add(coord_port)
        # Fires only for *newly assigned* (auto-assigned) replicas, not for
        # head-side pre-allocated slots that already have head-side clients.
        self._on_register = on_register
@@ -265,7 +235,7 @@ class OmniMasterServer:

        self._stage_config_events[route] = threading.Event()
        self._stage_coordinator_addresses[route] = StageCoordinatorAddresses()
-        hs_port, inp_port, out_port = self._alloc_unique_ports(3)
+        hs_port, inp_port, out_port = get_open_ports_list(count=3)
        alloc = StageAllocation(
            handshake_bind_address=f"tcp://{self._address}:{hs_port}",
            handshake_connect_address=f"tcp://{self._address}:{hs_port}",
@@ -277,32 +247,6 @@ class OmniMasterServer:
        self._stage_routes[route] = alloc
        return alloc

-    def _alloc_unique_ports(self, count: int) -> list[int]:
-        """Return ``count`` open ports unique across every route this server owns.
-
-        ``get_open_ports_list`` dedups only within its own call, so we redraw any
-        port that collides with one already recorded in ``self._allocated_ports``
-        (registration/coordinator ports plus every previously-allocated route
-        port) and register the winners before returning. Callers already hold the
-        relevant allocation context (``__init__`` is single-threaded;
-        registration holds ``self._alloc_lock``).
-        """
-        picked: list[int] = []
-        # Bounded retry budget so a pathological host (near port exhaustion)
-        # fails loudly instead of spinning forever.
-        for _ in range(64):
-            for port in get_open_ports_list(count=count - len(picked)):
-                if port in self._allocated_ports:
-                    continue
-                self._allocated_ports.add(port)
-                picked.append(port)
-                if len(picked) == count:
-                    return picked
-        raise RuntimeError(
-            f"[OmniMasterServer] Could not allocate {count} unique open ports "
-            f"(host={self._address}); {len(self._allocated_ports)} already in use."
-        )
-
    def _next_free_replica_id(self, stage_id: int) -> int:
        """Return the next replica id to assign for an auto-assign registration.

@@ -32,7 +32,7 @@ from starlette.routing import Route
 from starlette.types import ASGIApp, Receive, Scope, Send
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
-from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
+from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.launcher import serve_http, terminate_if_errored
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
@@ -60,11 +60,12 @@ from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
 from vllm.entrypoints.pooling.scoring.serving import ServingScores
-from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
+from vllm.entrypoints.serve.disagg.serving import ServingTokens

 # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
 # Keep a fallback for older/newer upstream layouts during rebase windows.
 from vllm.entrypoints.serve.instrumentator.basic import base
+from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
 from vllm.entrypoints.serve.utils.api_utils import (
    load_aware_call,
@@ -84,7 +85,6 @@ from vllm.entrypoints.speech_to_text.translation.serving import (
    OpenAIServingTranslation,
 )
 from vllm.logger import init_logger
-from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.tasks import POOLING_TASKS
 from vllm.tool_parsers import ToolParserManager
 from vllm.utils import random_uuid
@@ -451,7 +451,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
    # Add process-specific prefix to stdout and stderr.
    decorate_logs("APIServer", skip_if_decorated=True)

-    listen_address, sock = setup_openai_server(args, reuse_port=False)
+    listen_address, sock = setup_openai_server(args)

    # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models
    await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
@@ -830,12 +830,6 @@ async def omni_init_app_state(
        if tokenizer is None or getattr(tokenizer, "chat_template", None) is None:
            resolved_chat_template = _load_model_chat_template_json(args.model)

-    chat_template_config = ChatTemplateConfig(
-        chat_template=resolved_chat_template,
-        chat_template_content_format=args.chat_template_content_format,
-        trust_request_chat_template=args.trust_request_chat_template,
-    )
-
    if args.tool_server == "demo":
        tool_server: ToolServer | None = DemoToolServer()
        assert isinstance(tool_server, DemoToolServer)
@@ -901,10 +895,11 @@ async def omni_init_app_state(
    await state.openai_serving_models.init_static_loras()

    # NOTE: kept aligned with upstream `init_app_state`:
-    # Use OnlineRenderer (replaced OpenAIServingRender which was removed upstream).
-    state.online_renderer = OnlineRenderer(
+    # Use OpenAIServingRender (replaces the old OnlineRenderer + OnlineDerenderer + ServingRender split).
+    state.openai_serving_render = OpenAIServingRender(
        model_config=engine_client.model_config,
        renderer=engine_client.renderer,
+        model_registry=state.openai_serving_models.registry,
        request_logger=request_logger,
        chat_template=resolved_chat_template,
        chat_template_content_format=args.chat_template_content_format,
@@ -921,7 +916,7 @@ async def omni_init_app_state(
        OpenAIServingResponses(
            engine_client,
            state.openai_serving_models,
-            state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
            request_logger=request_logger,
            chat_template=resolved_chat_template,
            chat_template_content_format=args.chat_template_content_format,
@@ -942,7 +937,7 @@ async def omni_init_app_state(
            engine_client,
            state.openai_serving_models,
            args.response_role,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
            request_logger=request_logger,
            chat_template=resolved_chat_template,
            chat_template_content_format=args.chat_template_content_format,
@@ -969,7 +964,7 @@ async def omni_init_app_state(
        OpenAIServingCompletion(
            engine_client,
            state.openai_serving_models,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
            request_logger=request_logger,
            return_tokens_as_token_ids=args.return_tokens_as_token_ids,
            enable_prompt_tokens_details=args.enable_prompt_tokens_details,
@@ -984,7 +979,9 @@ async def omni_init_app_state(
            state.openai_serving_models,
            supported_tasks=tuple(supported_tasks),
            request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
        )
        if any(task in POOLING_TASKS for task in supported_tasks)
        else None
@@ -994,7 +991,9 @@ async def omni_init_app_state(
            engine_client,
            state.openai_serving_models,
            request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
        )
        if "embed" in supported_tasks
        else None
@@ -1004,7 +1003,9 @@ async def omni_init_app_state(
            engine_client,
            state.openai_serving_models,
            request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
        )
        if "classify" in supported_tasks
        else None
@@ -1013,9 +1014,8 @@ async def omni_init_app_state(
        ServingScores(
            engine_client,
            state.openai_serving_models,
-            supported_tasks=tuple(supported_tasks),
            request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            score_template=resolved_chat_template,
            log_error_stack=args.log_error_stack,
        )
        if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
@@ -1023,7 +1023,7 @@ async def omni_init_app_state(
    )
    state.serving_tokenization = ServingTokenization(
        state.openai_serving_models,
-        state.online_renderer,
+        state.openai_serving_render,
        request_logger=request_logger,
        chat_template=resolved_chat_template,
        chat_template_content_format=args.chat_template_content_format,
@@ -1055,7 +1055,7 @@ async def omni_init_app_state(
            engine_client,
            state.openai_serving_models,
            args.response_role,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
            request_logger=request_logger,
            chat_template=resolved_chat_template,
            chat_template_content_format=args.chat_template_content_format,
@@ -1074,7 +1074,7 @@ async def omni_init_app_state(
        ServingTokens(
            engine_client,
            state.openai_serving_models,
-            state.online_renderer,
+            state.openai_serving_render,
            request_logger=request_logger,
            return_tokens_as_token_ids=args.return_tokens_as_token_ids,
            enable_prompt_tokens_details=args.enable_prompt_tokens_details,
@@ -3,7 +3,7 @@ import asyncio
 import torch
 from fastapi import Request
 from fastapi.responses import Response
-from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.logger import init_logger
 from vllm.utils import random_uuid

@@ -41,7 +41,6 @@ except ImportError:
    soundfile = None


-from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionNamedToolChoiceParam,
@@ -65,6 +64,7 @@ from vllm.entrypoints.openai.engine.protocol import (
    ToolCall,
    UsageInfo,
 )
+from vllm.entrypoints.openai.engine.serving import clamp_prompt_logprobs
 from vllm.entrypoints.openai.parser.harmony_utils import (
    get_streamable_parser_for_assistant,
 )
@@ -150,21 +150,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
    # Harmony flag (always False for vllm-omni models)
    use_harmony: bool = False

-    @property
-    def tool_call_id_type(self) -> str:
-        """Return the tool call ID type, delegating to model config.
-
-        Upstream vLLM removed the stored ``tool_call_id_type`` attribute
-        from ``OpenAIServingChat`` after the ParserManager refactor; the
-        field is now resolved on demand via ``get_tool_call_id_type``.
-        """
-        try:
-            from vllm.entrypoints.chat_utils import get_tool_call_id_type
-
-            return get_tool_call_id_type(self.model_config)
-        except Exception:
-            return "random"
-
    def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest) -> bool:
        """Check if streamed tokens should go through the tool-call parser.

@@ -383,7 +368,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                tool_dicts = [tool.model_dump() for tool in request.tools]

            if not self.use_harmony:
-                error_check_ret = self.online_renderer.validate_chat_template(
+                error_check_ret = self.openai_serving_render.validate_chat_template(
                    request_chat_template=request.chat_template,
                    chat_template_kwargs=request.chat_template_kwargs,
                    trust_request_chat_template=self.trust_request_chat_template,
@@ -411,7 +396,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                )
            else:
                should_include_tools = tool_dicts is not None
-                conversation, engine_prompts = self.online_renderer._make_request_with_harmony(
+                conversation, engine_prompts = self.openai_serving_render._make_request_with_harmony(
                    request, should_include_tools
                )

@@ -21,12 +21,12 @@ import torch
 from fastapi import HTTPException, Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
 from transformers.utils.hub import cached_file
-from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.engine.protocol import (
    ErrorResponse,
    RequestResponseMetadata,
 )
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.inputs import tokens_input
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
@@ -23,7 +23,22 @@ 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
+
+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 ────────────────────────────────────────────────────────────

@@ -22,7 +22,6 @@ from .constants import (
    DEFAULT_SIGMA,
    DEFAULT_TEMPERATURE,
    HISTORY_PATCH_SIZE,
-    INITIAL_LATENT_CHUNK_SIZE,
    KEY_CFG,
    KEY_CHUNK_ID,
    KEY_DECODE_STEP,
@@ -240,7 +239,6 @@ class MingTTSConfig:
    max_decode_steps: int = MAX_DECODE_STEPS

    latent_chunk_size: int = LATENT_CHUNK_SIZE
-    initial_latent_chunk_size: int = INITIAL_LATENT_CHUNK_SIZE
    latent_left_context: int = LATENT_LEFT_CONTEXT

    text_eos_token_id: int = TEXT_EOS_TOKEN_ID
@@ -58,7 +58,6 @@ DEFAULT_TEMPERATURE = 0.0

 # Connector / Stage-2 streaming defaults (runtime tuning)
 LATENT_CHUNK_SIZE = 25
-INITIAL_LATENT_CHUNK_SIZE = 4
 LATENT_LEFT_CONTEXT = 0
 MAX_DECODE_STEPS = 200

@@ -172,9 +172,6 @@ def validate_ming_tts_config(cfg: Any) -> None:

    if cfg.latent_chunk_size <= 0:
        raise ValueError(f"latent_chunk_size must be > 0, got {cfg.latent_chunk_size}.")
-    initial_latent_chunk_size = getattr(cfg, "initial_latent_chunk_size", 0)
-    if initial_latent_chunk_size < 0:
-        raise ValueError(f"initial_latent_chunk_size must be >= 0, got {initial_latent_chunk_size}.")
    if cfg.latent_left_context < 0:
        raise ValueError(f"latent_left_context must be >= 0, got {cfg.latent_left_context}.")
    if cfg.max_decode_steps <= 0:
@@ -12,7 +12,6 @@ from vllm.logger import init_logger
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
-    INITIAL_LATENT_CHUNK_SIZE,
    KEY_CHUNK_ID,
    KEY_REQUEST_ID,
    LATENT_CHUNK_SIZE,
@@ -117,7 +116,7 @@ def _decode_stop_reason(value: Any) -> str | None:
    return MING_STOP_REASON_BY_CODE.get(int(value))


-def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int, int]:
+def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]:
    connector = getattr(transfer_manager, "connector", None)
    raw_cfg = getattr(connector, "config", {}) or {}
    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, Mapping) else {}
@@ -132,19 +131,9 @@ def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int, int]:
            LATENT_LEFT_CONTEXT,
        )
    chunk_size = int(cfg.get("latent_chunk_size", LATENT_CHUNK_SIZE))
-    initial_chunk_size = int(cfg.get("initial_latent_chunk_size", INITIAL_LATENT_CHUNK_SIZE))
    left_context = int(cfg.get("latent_left_context", LATENT_LEFT_CONTEXT))
    if chunk_size <= 0:
        raise ValueError(f"Invalid Ming latent_chunk_size={chunk_size}")
-    if initial_chunk_size < 0:
-        raise ValueError(f"Invalid Ming initial_latent_chunk_size={initial_chunk_size}")
-    if initial_chunk_size > chunk_size:
-        logger.warning(
-            "Ming initial_latent_chunk_size=%d > latent_chunk_size=%d, clamping to latent_chunk_size.",
-            initial_chunk_size,
-            chunk_size,
-        )
-        initial_chunk_size = chunk_size
    # Stage-2 VAE caches past_key_values and stream_state by request_id.
    # Replaying left-context latents would double-feed cached decoder state.
    if left_context != 0:
@@ -153,7 +142,7 @@ def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int, int]:
            "Ming boundary continuity is handled by per-request decoder state cache, not "
            f"latent replay. Got latent_left_context={left_context}."
        )
-    return chunk_size, initial_chunk_size, left_context
+    return chunk_size, left_context


 def _build_chunk_observability(
@@ -207,7 +196,7 @@ def llm2audio_vae_async_chunk(
    if patch is not None:
        transfer_manager.code_prompt_token_ids[request_id].append(patch)

-    chunk_size, initial_chunk_size, _ = _get_async_chunk_config(transfer_manager)
+    chunk_size, _ = _get_async_chunk_config(transfer_manager)

    patches = transfer_manager.code_prompt_token_ids[request_id]
    seen_patch_len = int(state.get("seen_patch_len", 0))
@@ -236,17 +225,11 @@ def llm2audio_vae_async_chunk(
            )
        return None

-    use_first_chunk = 0 < initial_chunk_size < chunk_size and seen_patch_len == 0
-    if finished:
-        emit_count = length
-    elif use_first_chunk:
-        if length < initial_chunk_size:
-            return None
-        emit_count = initial_chunk_size
-    else:
-        if length < chunk_size:
-            return None
-        emit_count = chunk_size
+    chunk_length = length % chunk_size
+    if chunk_length != 0 and not finished:
+        return None
+
+    emit_count = chunk_length if chunk_length != 0 else chunk_size
    emit_patches = list(new_patches[:emit_count])
    state["seen_patch_len"] = seen_patch_len + len(emit_patches)
    latent_patches = torch.stack(emit_patches, dim=0)
@@ -417,70 +417,3 @@ def _patch_fp8_use_quack_fused_bias():


 _patch_fp8_use_quack_fused_bias()
-
-
-# =============================================================================
-# Patch CuMemAllocator._python_free_callback to fix CUDA double-free on shutdown
-# =============================================================================
-# WHY: CuMemAllocator._python_free_callback guards the asleep-entry double-free
-# skip with ``if data.is_asleep and current_platform.is_rocm():`` — only ROCm
-# gets the safe empty-handle return.  On CUDA, the callback falls through and
-# returns the original handle, causing cuMemRelease on already-freed memory
-# (CUDA_ERROR_INVALID_VALUE) during EngineCore subprocess atexit cleanup.
-#
-# This happens because ``sleep()`` calls ``unmap_and_release()`` on ALL
-# platforms, then sets ``is_asleep = True``.  When the atexit handler
-# (``_shutdown_singleton`` -> ``release_pools()`` -> GC) triggers the free
-# callback, the asleep entries must return an empty chunk list so the C
-# extension skips ``cuMemRelease`` — exactly the same logic already used by
-# the ROCm guard.
-#
-# The fix removes the ``current_platform.is_rocm()`` condition so the guard
-# applies to CUDA (and any future platform that implements cumem).
-#
-# FRAGILITY: Relies on ``_python_free_callback`` being a regular method
-# (not a slot or C extension).  If CuMemAllocator is rewritten in C/Cython,
-# this monkey-patch will silently become a no-op.
-def _patch_cumem_free_callback_cuda() -> None:
-    try:
-        from vllm.device_allocator.cumem import CuMemAllocator
-    except ImportError:
-        _PATCH_LOGGER.debug("[cumem-cuda] CuMemAllocator not available; skipping patch")
-        return
-
-    _original_free_callback = CuMemAllocator._python_free_callback
-
-    if getattr(_original_free_callback, "_omni_cumem_cuda_patched", False):
-        return
-
-    # The upstream bug: `_python_free_callback` only skips the double-free
-    # for ROCm (line ~206).  We wrap the method to extend the guard to all
-    # platforms.
-    def _patched_free_callback(self, ptr: int) -> tuple:
-        data = self.pointer_to_data.pop(ptr)
-        if data.cpu_backup_tensor is not None:
-            data.cpu_backup_tensor = None
-        if data.is_asleep:
-            # sleep() already called unmap_and_release() on this allocation.
-            # Return an empty chunk list so the C extension skips
-            # cuMemRelease, avoiding a double-free.  Same logic as the
-            # existing ROCm guard, but applied to all platforms.
-            device, size, d_mem, _ = data.handle
-            result = (device, size, d_mem, [])
-            _PATCH_LOGGER.debug(
-                "[cumem-cuda] Free callback: asleep entry %s -> empty handle",
-                ptr,
-            )
-            return result
-        # Drain pending kernels before the C extension's cuMemUnmap.
-        torch.accelerator.synchronize(data.handle[0])
-        return data.handle
-
-    _patched_free_callback._omni_cumem_cuda_patched = True
-    CuMemAllocator._python_free_callback = _patched_free_callback
-    _PATCH_LOGGER.info(
-        "[cumem-cuda] CuMemAllocator._python_free_callback patched: asleep guard extended to all platforms."
-    )
-
-
-_patch_cumem_free_callback_cuda()
@@ -159,11 +159,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
        async_output_copy_stream = kwargs.pop("async_output_copy_stream")
        vocab_size = kwargs.pop("vocab_size")
        routed_experts = kwargs.pop("routed_experts", None)
-        # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault
-        # for EP all2all fault tolerance (PR #43637). Omni doesn't use this
-        # feature but must consume the kwarg to prevent TypeError from stray
-        # kwargs and initialize the attribute so super().get_output() works.
-        kwargs.pop("check_ep_fault", False)
        if kwargs:
            raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}")

@@ -175,7 +170,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
        self.vocab_size = vocab_size
        self._logprobs_tensors = logprobs_tensors
        self._routed_experts = routed_experts
-        self._has_fault: torch.Tensor | None = None

        default_stream = torch.cuda.current_stream()
        with torch.cuda.stream(async_output_copy_stream):
@@ -226,12 +220,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
            if background_exception is not None:
                raise background_exception
        self._build_model_runner_output_once()
-        # Upstream AsyncGPUModelRunnerOutput.get_output() accesses
-        # self._has_fault for EP all2all fault tolerance (PR #43637).
-        # Ensure the attribute exists even when __init__ was bypassed
-        # (e.g. unit tests using object.__new__).
-        if not hasattr(self, "_has_fault"):
-            self._has_fault = None
        with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"):
            return super().get_output()