[Model] Add HunyuanVideo-1.5 T2V and I2V support (#1516)

Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.17.0"
+          VLLM_VERSION: "v0.18.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: true

.buildkite/test-amd.yaml

@@ -1,7 +1,7 @@
 steps:
 
 - label: "Diffusion Model Test"
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -11,7 +11,7 @@ steps:
     - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
 
 - label: "Diffusion Images API LoRA E2E"
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]

.buildkite/test-merge.yml

@@ -17,7 +17,7 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
-    timeout_in_minutes: 20
+    timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
@@ -35,7 +35,7 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Images API LoRA E2E"
-    timeout_in_minutes: 20
+    timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
     commands:
       - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py

.buildkite/test-ready.yml

@@ -36,7 +36,7 @@ steps:
   - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
     commands:
-      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+      - timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
     agents:
       queue: "gpu_1_queue"
     plugins:

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.17.0
+ARG VLLM_BASE_TAG=v0.18.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
 #    we should swap over to use stable release ASAP.
 #    We should avoid relying on custom commits.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
+ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 

docker/Dockerfile.xpu

@@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.17.0
+ARG VLLM_VERSION=v0.18.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/getting_started/installation/gpu.md

@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, 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, and 0.18.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"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -34,13 +34,13 @@ uv pip install vllm-omni
 # --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.17.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
 
 ```bash
-uv pip install vllm==0.17.0 --torch-backend=auto
+uv pip install vllm==0.18.0 --torch-backend=auto
 ```
 
-The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
+The 0.18.0 release of vLLM ships CUDA 12.9-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
@@ -56,11 +56,11 @@ 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.17.0
+git checkout v0.18.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
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.16.0 \
+    vllm/vllm-omni:v0.18.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -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.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
+uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --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.17.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
 
 ```bash
-uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
+uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
 ```
 
-The release 0.17.0 of vLLM requires ROCm 7.0 environment.
+The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. 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.17.0
+git checkout v0.18.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.16.0 \
+  vllm/vllm-omni-rocm:v0.18.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.16.0
+  vllm/vllm-omni-rocm:v0.18.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/installation/npu/npu.inc.md

@@ -68,18 +68,18 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
-# Pin vLLM version to 0.17.0
+# Pin vLLM version to 0.18.0
 cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
-git checkout v0.17.0
+git checkout v0.18.0
 VLLM_TARGET_DEVICE=empty pip install -v -e .
 
 # Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
 cd /vllm-workspace/vllm-ascend
 git pull origin main
 git fetch origin --tags
-git checkout v0.17.0
+git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
 pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.17.0 --torch-backend=auto
+uv pip install vllm==0.18.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
+uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/models/supported_models.md

@@ -52,6 +52,8 @@ th {
 |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
 |`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
 |`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
+| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` |
+| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` |
 |`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` |
 
 

docs/user_guide/diffusion/quantization/fp8.md

@@ -63,6 +63,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
 | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
 | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
+| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
 
 ## Combining with Other Features
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -2,15 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2.
+Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5.
 
 Supports:
 - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
 - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
 - LTX2 image-to-video pipeline
+- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning
 
 Usage:
-    # I2V-A14B (MoE)
+    # Wan I2V-A14B (MoE)
     python image_to_video.py --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
         --image input.jpg --prompt "A cat playing with yarn"
 
@@ -24,6 +25,11 @@ Usage:
         --image input.jpg --prompt "A cinematic dolly shot of a boat" \
         --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
         --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
+
+    # HunyuanVideo-1.5 I2V (480p)
+    python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
+        --image input.jpg --prompt "A cat playing with yarn" \
+        --flow-shift 5.0 --guidance-scale 6.0
 """
 
 import argparse
@@ -43,11 +49,11 @@ from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.")
+    parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-        help="Diffusers Wan2.2 I2V model ID or local path.",
+        help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).",
     )
     parser.add_argument(
         "--model-class-name",

examples/offline_inference/text_to_video/text_to_video.md

@@ -1,10 +1,19 @@
 # Text-To-Video
 
-The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts. This script can also be used
-for `Lightricks/LTX-2` to generate video+audio.
+A unified script for text-to-video generation. Supports multiple models with model-aware defaults.
+
+## Supported Models
+
+| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
+|---|---|---|---|---|---|
+| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
+| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
+| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 
 ## Local CLI Usage
 
+### Wan2.2 (default)
+
 ```bash
 python text_to_video.py \
   --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
@@ -36,22 +45,100 @@ python text_to_video.py \
   --output ltx2_out.mp4
 ```
 
-Key arguments:
+### HunyuanVideo-1.5 (480p)
 
+```bash
+python text_to_video.py \
+  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+  --prompt "A cat walks through a sunlit garden, flowers swaying gently in the breeze." \
+  --height 480 \
+  --width 832 \
+  --num-frames 121 \
+  --guidance-scale 6.0 \
+  --flow-shift 5.0 \
+  --num-inference-steps 50 \
+  --fps 24 \
+  --output hunyuan_video_15_output.mp4
+```
+
+### HunyuanVideo-1.5 (720p)
+
+```bash
+python text_to_video.py \
+  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v \
+  --prompt "A serene lakeside sunrise with mist over the water." \
+  --height 720 \
+  --width 1280 \
+  --num-frames 121 \
+  --guidance-scale 6.0 \
+  --flow-shift 9.0 \
+  --num-inference-steps 50 \
+  --fps 24 \
+  --output hunyuan_720p.mp4
+```
+
+### HunyuanVideo-1.5 with FP8 Quantization
+
+```bash
+python text_to_video.py \
+  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+  --prompt "A dog running across a field of golden wheat." \
+  --quantization fp8 \
+  --height 480 --width 832 --num-frames 121 \
+  --guidance-scale 6.0 --flow-shift 5.0 \
+  --output hunyuan_fp8.mp4
+```
+
+Quick test (smaller resolution, fewer frames):
+
+```bash
+python text_to_video.py \
+  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+  --prompt "A serene lakeside sunrise with mist over the water." \
+  --height 320 --width 576 --num-frames 17 --num-inference-steps 30 \
+  --flow-shift 5.0 \
+  --output quick_test.mp4
+```
+
+## Key Arguments
+
+### Common
+
+- `--model`: Diffusers model ID or local path.
 - `--prompt`: text description (string).
-- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
-- `--num-frames`: Number of frames (Wan default is 81).
-- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
-- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
-- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
-- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
+- `--height/--width`: output resolution. Default depends on model.
+- `--num-frames`: number of frames. Default depends on model.
+- `--guidance-scale`: CFG scale. Default depends on model.
+- `--num-inference-steps`: sampling steps. Default depends on model.
+- `--fps`: frames per second for the saved MP4.
 - `--output`: path to save the generated video.
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
 - `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
+- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
+- `--flow-shift`: scheduler flow_shift parameter.
 
-> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
+### Wan2.2-specific
+
+- `--negative-prompt`: artifacts to suppress.
+- `--guidance-scale-high`: separate CFG scale for high-noise stage.
+- `--boundary-ratio`: boundary split for low/high DiT (default 0.875).
+- `--flow-shift`: scheduler flow_shift (5.0 for 720p, 12.0 for 480p).
+- `--cache-backend`: `cache_dit` for acceleration.
+
+### HunyuanVideo-1.5 Optimal Configs
+
+| Variant | flow_shift | guidance_scale | steps |
+|---------|-----------|----------------|-------|
+| 480p T2V | 5.0 | 6.0 | 50 |
+| 720p T2V | 9.0 | 6.0 | 50 |
+| 480p I2V | 5.0 | 6.0 | 50 |
+| 720p I2V | 7.0 | 6.0 | 50 |
+| CFG-distilled | (same) | 1.0 | 50 |
+
+> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`, `--enable-cpu-offload`, or `--quantization fp8`.

examples/offline_inference/text_to_video/text_to_video.py

@@ -15,56 +15,92 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
+_MODEL_PRESETS = {
+    "wan": {
+        "height": 720,
+        "width": 1280,
+        "num_frames": 81,
+        "num_inference_steps": 40,
+        "guidance_scale": 4.0,
+        "fps": 24,
+        "output": "wan22_output.mp4",
+    },
+    "hunyuan": {
+        "height": 480,
+        "width": 832,
+        "num_frames": 121,
+        "num_inference_steps": 50,
+        "guidance_scale": 6.0,
+        "fps": 24,
+        "output": "hunyuan_video_15_output.mp4",
+    },
+}
+
+
+def _detect_preset(model: str) -> dict:
+    model_lower = model.lower()
+    if "hunyuan" in model_lower:
+        return _MODEL_PRESETS["hunyuan"]
+    return _MODEL_PRESETS["wan"]
+
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video with Wan2.2 T2V.")
+    parser = argparse.ArgumentParser(
+        description="Generate a video from a text prompt. "
+        "Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models."
+    )
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-        help="Diffusers Wan2.2 model ID or local path.",
+        help="Diffusers model ID or local path. "
+        "Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
+        "hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
     )
     parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
-    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt (Wan2.2 only).")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
-    parser.add_argument("--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise.")
-    parser.add_argument("--height", type=int, default=720, help="Video height.")
-    parser.add_argument("--width", type=int, default=1280, help="Video width.")
-    parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
+    parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
+    parser.add_argument(
+        "--guidance-scale-high", type=float, default=None, help="Separate CFG for high-noise stage (Wan2.2 only)."
+    )
+    parser.add_argument("--height", type=int, default=None, help="Video height. Default: model-specific.")
+    parser.add_argument("--width", type=int, default=None, help="Video width. Default: model-specific.")
+    parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
     parser.add_argument(
         "--frame-rate",
         type=float,
         default=None,
         help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
     )
-    parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
+    parser.add_argument(
+        "--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
+    )
     parser.add_argument(
         "--boundary-ratio",
         type=float,
-        default=0.875,
-        help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).",
+        default=None,
+        help="(Wan2.2) Boundary split ratio for low/high DiT. Default 0.875.",
     )
     parser.add_argument(
-        "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
+        "--flow-shift",
+        type=float,
+        default=None,
+        help="Scheduler flow_shift. Wan2.2: 5.0(720p)/12.0(480p). HunyuanVideo-1.5: 5.0(480p)/9.0(720p).",
     )
     parser.add_argument(
         "--cache-backend",
         type=str,
         default=None,
         choices=["cache_dit"],
-        help=(
-            "Cache backend to use for acceleration. "
-            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer). "
-            "Default: None (no cache acceleration)."
-        ),
+        help="Cache backend for acceleration (Wan2.2). Default: None.",
     )
     parser.add_argument(
         "--enable-cache-dit-summary",
         action="store_true",
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
-    parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
-    parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
+    parser.add_argument("--output", type=str, default=None, help="Output path (mp4). Default: model-specific.")
+    parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
     parser.add_argument(
         "--vae-use-slicing",
         action="store_true",
@@ -137,32 +173,41 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "gguf"],
+        help="Quantization method for the transformer (fp8 for online FP8 quantization).",
+    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-    frame_rate = args.frame_rate if args.frame_rate is not None else float(args.fps)
 
-    # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
+    preset = _detect_preset(args.model)
+    for key, default_val in preset.items():
+        if getattr(args, key.replace("-", "_"), None) is None:
+            setattr(args, key.replace("-", "_"), default_val)
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    # Cache-dit config (Wan2.2 only)
     cache_config = None
     if args.cache_backend == "cache_dit":
         cache_config = {
-            # DBCache parameters [cache-dit only]
-            "Fn_compute_blocks": 1,  # Optimized for single-transformer models
-            "Bn_compute_blocks": 0,  # Number of backward compute blocks
-            "max_warmup_steps": 4,  # Maximum warmup steps (works for few-step models)
+            "Fn_compute_blocks": 1,
+            "Bn_compute_blocks": 0,
+            "max_warmup_steps": 4,
             "max_cached_steps": 20,
-            "residual_diff_threshold": 0.24,  # Higher threshold for more aggressive caching
-            "max_continuous_cached_steps": 3,  # Limit to prevent precision degradation
-            # TaylorSeer parameters [cache-dit only]
-            "enable_taylorseer": False,  # Disabled by default (not suitable for few-step models)
-            "taylorseer_order": 1,  # TaylorSeer polynomial order
-            # SCM (Step Computation Masking) parameters [cache-dit only]
-            "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
-            "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
+            "residual_diff_threshold": 0.24,
+            "max_continuous_cached_steps": 3,
+            "enable_taylorseer": False,
+            "taylorseer_order": 1,
+            "scm_steps_mask_policy": None,
+            "scm_steps_policy": "dynamic",
         }
+
     # Configure parallel settings
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
@@ -176,14 +221,11 @@ def main():
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
-    omni = Omni(
+    omni_kwargs = dict(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
-        boundary_ratio=args.boundary_ratio,
-        flow_shift=args.flow_shift,
-        enable_cache_dit_summary=args.enable_cache_dit_summary,
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
@@ -191,6 +233,18 @@ def main():
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
+    if args.boundary_ratio is not None:
+        omni_kwargs["boundary_ratio"] = args.boundary_ratio
+    if args.flow_shift is not None:
+        omni_kwargs["flow_shift"] = args.flow_shift
+    if args.quantization is not None:
+        omni_kwargs["quantization"] = args.quantization
+    if args.cache_backend is not None:
+        omni_kwargs["cache_backend"] = args.cache_backend
+        omni_kwargs["cache_config"] = cache_config
+        omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
+
+    omni = Omni(**omni_kwargs)
 
     if profiler_enabled:
         print("[Profiler] Starting profiling...")
@@ -210,22 +264,25 @@ def main():
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 
+    prompt_dict = {"prompt": args.prompt}
+    if args.negative_prompt:
+        prompt_dict["negative_prompt"] = args.negative_prompt
+
+    sampling_kwargs = dict(
+        height=args.height,
+        width=args.width,
+        generator=generator,
+        guidance_scale=args.guidance_scale,
+        num_inference_steps=args.num_inference_steps,
+        num_frames=args.num_frames,
+    )
+    if args.guidance_scale_high is not None:
+        sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
+
     generation_start = time.perf_counter()
     frames = omni.generate(
-        {
-            "prompt": args.prompt,
-            "negative_prompt": args.negative_prompt,
-        },
-        OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            generator=generator,
-            guidance_scale=args.guidance_scale,
-            guidance_scale_2=args.guidance_scale_high,
-            num_inference_steps=args.num_inference_steps,
-            num_frames=args.num_frames,
-            frame_rate=frame_rate,
-        ),
+        prompt_dict,
+        OmniDiffusionSamplingParams(**sampling_kwargs),
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start

examples/online_serving/image_to_video/run_curl_hunyuan_video_15.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+# HunyuanVideo-1.5 image-to-video curl example using the async video job API.
+
+set -euo pipefail
+
+INPUT_IMAGE="${INPUT_IMAGE:-test_input.jpg}"
+BASE_URL="${BASE_URL:-http://localhost:8099}"
+OUTPUT_PATH="${OUTPUT_PATH:-hunyuan_video_15_i2v.mp4}"
+POLL_INTERVAL="${POLL_INTERVAL:-2}"
+
+if [ ! -f "$INPUT_IMAGE" ]; then
+    echo "Input image not found: $INPUT_IMAGE"
+    echo "Provide an image via INPUT_IMAGE env var."
+    exit 1
+fi
+
+create_response=$(
+  curl -sS -X POST "${BASE_URL}/v1/videos" \
+    -H "Accept: application/json" \
+    -F "prompt=The camera follows the puppy as it runs forward on the grass, its four legs alternating steps, its tail held high and wagging side to side." \
+    -F "input_reference=@${INPUT_IMAGE}" \
+    -F "size=832x480" \
+    -F "num_frames=33" \
+    -F "fps=24" \
+    -F "num_inference_steps=30" \
+    -F "guidance_scale=6.0" \
+    -F "flow_shift=5.0" \
+    -F "seed=42"
+)
+
+video_id="$(echo "${create_response}" | jq -r '.id')"
+if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
+  echo "Failed to create video job:"
+  echo "${create_response}" | jq .
+  exit 1
+fi
+
+echo "Created video job ${video_id}"
+echo "${create_response}" | jq .
+
+while true; do
+  status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
+  status="$(echo "${status_response}" | jq -r '.status')"
+
+  case "${status}" in
+    queued|in_progress)
+      echo "Video job ${video_id} status: ${status}"
+      sleep "${POLL_INTERVAL}"
+      ;;
+    completed)
+      echo "${status_response}" | jq .
+      break
+      ;;
+    failed)
+      echo "Video generation failed:"
+      echo "${status_response}" | jq .
+      exit 1
+      ;;
+    *)
+      echo "Unexpected status response:"
+      echo "${status_response}" | jq .
+      exit 1
+      ;;
+  esac
+done
+
+curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
+echo "Saved video to ${OUTPUT_PATH}"

examples/online_serving/image_to_video/run_server_hunyuan_video_15.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+# HunyuanVideo-1.5 image-to-video online serving startup script
+#
+# 480p: ~35 GB VRAM (BF16), fits 1x A100 80GB
+# 720p: needs FP8 + VAE tiling, ~35 GB VRAM
+
+MODEL="${MODEL:-hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v}"
+PORT="${PORT:-8099}"
+FLOW_SHIFT="${FLOW_SHIFT:-5.0}"
+QUANTIZATION="${QUANTIZATION:-}"
+CACHE_BACKEND="${CACHE_BACKEND:-none}"
+
+echo "Starting HunyuanVideo-1.5 I2V server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+echo "Flow shift: $FLOW_SHIFT"
+echo "Quantization: ${QUANTIZATION:-none}"
+echo "Cache backend: $CACHE_BACKEND"
+
+EXTRA_FLAGS=""
+if [ -n "$QUANTIZATION" ]; then
+    EXTRA_FLAGS="$EXTRA_FLAGS --quantization $QUANTIZATION"
+fi
+if [ "$CACHE_BACKEND" != "none" ]; then
+    EXTRA_FLAGS="$EXTRA_FLAGS --cache-backend $CACHE_BACKEND"
+fi
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT" \
+    --flow-shift "$FLOW_SHIFT" \
+    $EXTRA_FLAGS

examples/online_serving/text_to_video/run_curl_hunyuan_video_15.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+# HunyuanVideo-1.5 text-to-video curl example using the async video job API.
+
+set -euo pipefail
+
+BASE_URL="${BASE_URL:-http://localhost:8098}"
+OUTPUT_PATH="${OUTPUT_PATH:-hunyuan_video_15_t2v.mp4}"
+POLL_INTERVAL="${POLL_INTERVAL:-2}"
+
+create_response=$(
+  curl -sS -X POST "${BASE_URL}/v1/videos" \
+    -H "Accept: application/json" \
+    -F "prompt=A little girl wearing a straw hat runs through a summer meadow full of wildflowers. A wide shot is used, with the camera panning right to follow her." \
+    -F "size=832x480" \
+    -F "num_frames=33" \
+    -F "fps=24" \
+    -F "num_inference_steps=30" \
+    -F "guidance_scale=6.0" \
+    -F "flow_shift=5.0" \
+    -F "seed=42"
+)
+
+video_id="$(echo "${create_response}" | jq -r '.id')"
+if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
+  echo "Failed to create video job:"
+  echo "${create_response}" | jq .
+  exit 1
+fi
+
+echo "Created video job ${video_id}"
+echo "${create_response}" | jq .
+
+while true; do
+  status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
+  status="$(echo "${status_response}" | jq -r '.status')"
+
+  case "${status}" in
+    queued|in_progress)
+      echo "Video job ${video_id} status: ${status}"
+      sleep "${POLL_INTERVAL}"
+      ;;
+    completed)
+      echo "${status_response}" | jq .
+      break
+      ;;
+    failed)
+      echo "Video generation failed:"
+      echo "${status_response}" | jq .
+      exit 1
+      ;;
+    *)
+      echo "Unexpected status response:"
+      echo "${status_response}" | jq .
+      exit 1
+      ;;
+  esac
+done
+
+curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
+echo "Saved video to ${OUTPUT_PATH}"

examples/online_serving/text_to_video/run_server_hunyuan_video_15.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+# HunyuanVideo-1.5 text-to-video online serving startup script
+#
+# 480p: ~35 GB VRAM (BF16), fits 1x A100 80GB
+# 720p: needs FP8 + VAE tiling, ~35 GB VRAM
+
+MODEL="${MODEL:-hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v}"
+PORT="${PORT:-8098}"
+FLOW_SHIFT="${FLOW_SHIFT:-5.0}"
+QUANTIZATION="${QUANTIZATION:-}"
+CACHE_BACKEND="${CACHE_BACKEND:-none}"
+
+echo "Starting HunyuanVideo-1.5 T2V server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+echo "Flow shift: $FLOW_SHIFT"
+echo "Quantization: ${QUANTIZATION:-none}"
+echo "Cache backend: $CACHE_BACKEND"
+
+EXTRA_FLAGS=""
+if [ -n "$QUANTIZATION" ]; then
+    EXTRA_FLAGS="$EXTRA_FLAGS --quantization $QUANTIZATION"
+fi
+if [ "$CACHE_BACKEND" != "none" ]; then
+    EXTRA_FLAGS="$EXTRA_FLAGS --cache-backend $CACHE_BACKEND"
+fi
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT" \
+    --flow-shift "$FLOW_SHIFT" \
+    $EXTRA_FLAGS

tests/conftest.py

@@ -1910,6 +1910,7 @@ class OmniRunner:
     def _cleanup_process(self):
         try:
             keywords = ["enginecore"]
+            matched = []
 
             for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
                 try:
@@ -1922,16 +1923,32 @@ class OmniRunner:
 
                     if is_process:
                         print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
-
-                        try:
-                            proc.terminate()
-                            time.sleep(2)
-                        except Exception:
-                            proc.kill()
-
+                        matched.append(proc)
                 except (psutil.NoSuchProcess, psutil.AccessDenied):
                     pass
 
+            for proc in matched:
+                try:
+                    proc.terminate()
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+            _, still_alive = psutil.wait_procs(matched, timeout=5)
+            for proc in still_alive:
+                try:
+                    proc.kill()
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+            if still_alive:
+                _, stubborn = psutil.wait_procs(still_alive, timeout=3)
+                if stubborn:
+                    print(f"Warning: failed to kill residual vllm pids: {[p.pid for p in stubborn]}")
+                else:
+                    print(f"Force-killed residual vllm pids: {[p.pid for p in still_alive]}")
+            elif matched:
+                print(f"Terminated vllm pids: {[p.pid for p in matched]}")
+
         except Exception as e:
             print(f"Error in psutil vllm cleanup: {e}")
 

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -24,6 +24,7 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 # This test is specific to Z-Image LoRA behavior. Keep it focused on a single
 # model to reduce runtime and avoid extra downloads.
 models = ["Tongyi-MAI/Z-Image-Turbo"]
+DIFFUSION_INIT_TIMEOUT_S = 600
 
 
 @pytest.mark.parametrize("model_name", models)
@@ -76,7 +77,11 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
         )
         return str(adapter_dir)
 
-    m = Omni(model=model_name)
+    m = Omni(
+        model=model_name,
+        stage_init_timeout=DIFFUSION_INIT_TIMEOUT_S,
+        init_timeout=DIFFUSION_INIT_TIMEOUT_S,
+    )
     try:
         # high resolution may cause OOM on L4
         height = 256

tests/e2e/online_serving/test_hunyuan_video_15_expansion.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests of common diffusion feature combinations in online serving mode
+for HunyuanVideo-1.5-T2V (480p).
+
+Coverage (H100, since model cannot fit L4):
+- TeaCache + Layerwise CPU offloading (1 GPU)
+- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs)
+"""
+
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+)
+from tests.utils import hardware_marks
+
+PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion."
+NEGATIVE_PROMPT = "low quality, blurry, distorted"
+
+MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+
+SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+def _get_diffusion_feature_cases(model: str):
+    """Return L4 diffusion feature cases for HunyuanVideo-1.5.
+
+    Designed for 2x H100 environment per issue #1832.
+    """
+    return [
+        # (1 GPU) TeaCache + Layerwise CPU offloading
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "tea_cache",
+                    "--enable-layerwise-offload",
+                ],
+            ),
+            id="single_card_teacache_layerwise",
+            marks=SINGLE_CARD_MARKS,
+        ),
+        # (2 GPUs) CacheDiT + TP=2 + VAE patch parallel=2
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--tensor-parallel-size",
+                    "2",
+                    "--vae-patch-parallel-size",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="parallel_cachedit_tp2_vae2",
+            marks=PARALLEL_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(MODEL),
+    indirect=True,
+)
+def test_hunyuan_video_15_t2v(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    """L4 diffusion feature coverage for HunyuanVideo-1.5-T2V on H100."""
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 480,
+            "width": 640,
+            "num_frames": 5,
+            "num_inference_steps": 2,
+            "guidance_scale": 6.0,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_images_generations_lora.py

@@ -28,6 +28,7 @@ from tests.utils import hardware_test
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 MODEL = "Tongyi-MAI/Z-Image-Turbo"
+DIFFUSION_INIT_TIMEOUT_S = 600
 
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
@@ -37,7 +38,17 @@ SEED = 42
 
 @pytest.fixture(scope="module")
 def omni_server():
-    with OmniServer(MODEL, ["--num-gpus", "1"]) as server:
+    with OmniServer(
+        MODEL,
+        [
+            "--num-gpus",
+            "1",
+            "--stage-init-timeout",
+            str(DIFFUSION_INIT_TIMEOUT_S),
+            "--init-timeout",
+            str(DIFFUSION_INIT_TIMEOUT_S),
+        ],
+    ) as server:
         yield server
 
 

tests/e2e/stage_configs/qwen2_5_omni_ci.yaml

@@ -45,7 +45,7 @@ stage_args:
       max_model_len: 16384
       max_num_batched_tokens: 16384
       max_num_seqs: 1
-      gpu_memory_utilization: 0.9
+      gpu_memory_utilization: 0.4
       skip_mm_profiling: true
       enforce_eager: true
       trust_remote_code: true
@@ -72,7 +72,7 @@ stage_args:
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.9 #increase the gpu memory utilization to enable the test on H800
+      gpu_memory_utilization: 0.5 #increase the gpu memory utilization to enable the test on H800
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false

tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py

@@ -123,6 +123,7 @@ def make_mock_model(hidden: int = 8):
     cfg.video_token_index = VIDEO_TOKEN_ID
     cfg.audio_token_index = AUDIO_TOKEN_ID
     model.config = cfg
+    model._has_oov_mm_tokens = False
 
     def fake_lm_embed(ids: torch.Tensor) -> torch.Tensor:
         # Use .clone() so the tensor is contiguous (expand() creates a strided
@@ -137,13 +138,12 @@ def make_mock_model(hidden: int = 8):
 
     model._embed_text_input_ids = lambda *a, **kw: SupportsMultiModal._embed_text_input_ids(model, *a, **kw)
 
-    def fake_super_embed(ids, mm_embs=None, *, is_multimodal=None, handle_oov_mm_token=False):
+    def fake_super_embed(ids, mm_embs=None, *, is_multimodal=None):
         return SupportsMultiModal.embed_input_ids(
             model,
             ids,
             mm_embs,
             is_multimodal=is_multimodal,
-            handle_oov_mm_token=handle_oov_mm_token,
         )
 
     model.embed_input_ids = lambda *a, **kw: Qwen2_5OmniThinkerForConditionalGeneration.embed_input_ids(model, *a, **kw)

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -364,6 +364,7 @@ class OmniARScheduler(VLLMScheduler):
         if stopped_preempted_reqs:
             # This is a rare case and unlikely to impact performance.
             self.waiting.remove_requests(stopped_preempted_reqs)
+            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
 
         # [Main] Handle failed KV load requests
         if failed_kv_load_req_ids and not self.recompute_kv_load_failures:

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -247,6 +247,15 @@ class OmniGenerationScheduler(VLLMScheduler):
         )
 
         total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
+
+        # Record the request ids scheduled in this step (v0.14.0 behavior).
+        self.prev_step_scheduled_req_ids.clear()
+        self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
+
+        new_block_ids_to_zero = (
+            (self.kv_cache_manager.take_new_block_ids() or None) if self.needs_kv_cache_zeroing else None
+        )
+
         scheduler_output = SchedulerOutput(
             scheduled_new_reqs=new_reqs_data,
             scheduled_cached_reqs=cached_reqs_data,
@@ -258,12 +267,9 @@ class OmniGenerationScheduler(VLLMScheduler):
             finished_req_ids=self.finished_req_ids,
             free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
             preempted_req_ids=set(),
+            new_block_ids_to_zero=new_block_ids_to_zero,
         )
 
-        # Record the request ids scheduled in this step (v0.14.0 behavior).
-        self.prev_step_scheduled_req_ids.clear()
-        self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
-
         # KVTransfer: package metadata
         if self.connector is not None:
             meta = self.connector.build_connector_meta(scheduler_output)
@@ -496,6 +502,7 @@ class OmniGenerationScheduler(VLLMScheduler):
         if stopped_preempted_reqs:
             # This is a rare case and unlikely to impact performance.
             self.waiting.remove_requests(stopped_preempted_reqs)
+            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
 
         # Handle failed KV load requests
         if failed_kv_load_req_ids and not self.recompute_kv_load_failures:

vllm_omni/diffusion/models/hunyuan_video/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HunyuanVideo-1.5 diffusion model components (T2V and I2V)."""
+
+from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import (
+    HunyuanVideo15Transformer3DModel,
+)
+from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import (
+    HunyuanVideo15Pipeline,
+    get_hunyuan_video_15_post_process_func,
+)
+from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v import (
+    HunyuanVideo15I2VPipeline,
+    get_hunyuan_video_15_i2v_post_process_func,
+    get_hunyuan_video_15_i2v_pre_process_func,
+)
+
+__all__ = [
+    "HunyuanVideo15Pipeline",
+    "HunyuanVideo15I2VPipeline",
+    "HunyuanVideo15Transformer3DModel",
+    "get_hunyuan_video_15_post_process_func",
+    "get_hunyuan_video_15_i2v_post_process_func",
+    "get_hunyuan_video_15_i2v_pre_process_func",
+]

vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py

@@ -0,0 +1,791 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.embeddings import (
+    CombinedTimestepTextProjEmbeddings,
+    TimestepEmbedding,
+    Timesteps,
+    get_1d_rotary_pos_embed,
+)
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero
+from torch import nn
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
+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
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
+
+logger = init_logger(__name__)
+
+
+class HunyuanVideo15PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int | tuple[int, int, int] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+    ) -> None:
+        super().__init__()
+        patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.proj(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)  # BCFHW -> BNC
+        return hidden_states
+
+
+class HunyuanVideo15AdaNorm(nn.Module):
+    def __init__(self, in_features: int, out_features: int | None = None) -> None:
+        super().__init__()
+        out_features = out_features or 2 * in_features
+        self.linear = nn.Linear(in_features, out_features)
+        self.nonlinearity = nn.SiLU()
+
+    def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        temb = self.linear(self.nonlinearity(temb))
+        gate_msa, gate_mlp = temb.chunk(2, dim=1)
+        gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
+        return gate_msa, gate_mlp
+
+
+class HunyuanVideo15TimeEmbedding(nn.Module):
+    def __init__(self, embedding_dim: int, use_meanflow: bool = False):
+        super().__init__()
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+        self.use_meanflow = use_meanflow
+        self.time_proj_r = None
+        self.timestep_embedder_r = None
+        if use_meanflow:
+            self.time_proj_r = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+            self.timestep_embedder_r = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        timestep_r: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=timestep.dtype))
+
+        if timestep_r is not None and self.timestep_embedder_r is not None:
+            timesteps_proj_r = self.time_proj_r(timestep_r)
+            timesteps_emb_r = self.timestep_embedder_r(timesteps_proj_r.to(dtype=timestep.dtype))
+            timesteps_emb = timesteps_emb + timesteps_emb_r
+
+        return timesteps_emb
+
+
+class HunyuanVideo15RotaryPosEmbed(nn.Module):
+    def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], theta: float = 256.0) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.patch_size_t = patch_size_t
+        self.rope_dim = rope_dim
+        self.theta = theta
+
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        rope_sizes = [num_frames // self.patch_size_t, height // self.patch_size, width // self.patch_size]
+
+        axes_grids = []
+        for i in range(len(rope_sizes)):
+            grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32)
+            axes_grids.append(grid)
+        grid = torch.meshgrid(*axes_grids, indexing="ij")
+        grid = torch.stack(grid, dim=0)  # [3, T', H', W']
+
+        freqs = []
+        for i in range(3):
+            # use_real=False returns complex [seq, dim/2]; we extract cos/sin manually
+            # to get half-dim tensors compatible with RotaryEmbedding (which doubles internally).
+            freq_cis = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=False)
+            freqs.append((freq_cis.real, freq_cis.imag))
+
+        freqs_cos = torch.cat([f[0] for f in freqs], dim=1).float()
+        freqs_sin = torch.cat([f[1] for f in freqs], dim=1).float()
+        return freqs_cos, freqs_sin
+
+
+class HunyuanVideo15IndividualTokenRefinerBlock(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_width_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
+
+        # Not TP
+        from diffusers.models.attention_processor import Attention as DiffusersAttention
+
+        self.attn = DiffusersAttention(
+            query_dim=hidden_size,
+            cross_attention_dim=None,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            bias=attention_bias,
+        )
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
+
+        from diffusers.models.attention import FeedForward as DiffusersFeedForward
+
+        self.ff = DiffusersFeedForward(
+            hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate
+        )
+
+        self.norm_out = HunyuanVideo15AdaNorm(hidden_size, 2 * hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        norm_hidden_states = self.norm1(hidden_states)
+
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=attention_mask,
+        )
+
+        gate_msa, gate_mlp = self.norm_out(temb)
+        hidden_states = hidden_states + attn_output * gate_msa
+
+        ff_output = self.ff(self.norm2(hidden_states))
+        hidden_states = hidden_states + ff_output * gate_mlp
+
+        return hidden_states
+
+
+class HunyuanVideo15IndividualTokenRefiner(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        num_layers: int,
+        mlp_width_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        self.refiner_blocks = nn.ModuleList(
+            [
+                HunyuanVideo15IndividualTokenRefinerBlock(
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_width_ratio=mlp_width_ratio,
+                    mlp_drop_rate=mlp_drop_rate,
+                    attention_bias=attention_bias,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        self_attn_mask = None
+        if attention_mask is not None:
+            batch_size = attention_mask.shape[0]
+            seq_len = attention_mask.shape[1]
+            attention_mask = attention_mask.to(hidden_states.device).bool()
+            self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
+            self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
+            self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
+
+        for block in self.refiner_blocks:
+            hidden_states = block(hidden_states, temb, self_attn_mask)
+
+        return hidden_states
+
+
+class HunyuanVideo15TokenRefiner(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        num_layers: int,
+        mlp_ratio: float = 4.0,
+        mlp_drop_rate: float = 0.0,
+        attention_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.time_text_embed = CombinedTimestepTextProjEmbeddings(
+            embedding_dim=hidden_size, pooled_projection_dim=in_channels
+        )
+        self.proj_in = nn.Linear(in_channels, hidden_size, bias=True)
+        self.token_refiner = HunyuanVideo15IndividualTokenRefiner(
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            num_layers=num_layers,
+            mlp_width_ratio=mlp_ratio,
+            mlp_drop_rate=mlp_drop_rate,
+            attention_bias=attention_bias,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        attention_mask: torch.LongTensor | None = None,
+    ) -> torch.Tensor:
+        if attention_mask is None:
+            pooled_projections = hidden_states.mean(dim=1)
+        else:
+            original_dtype = hidden_states.dtype
+            mask_float = attention_mask.float().unsqueeze(-1)
+            pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+            pooled_projections = pooled_projections.to(original_dtype)
+
+        temb = self.time_text_embed(timestep, pooled_projections)
+        hidden_states = self.proj_in(hidden_states)
+        hidden_states = self.token_refiner(hidden_states, temb, attention_mask)
+
+        return hidden_states
+
+
+class HunyuanVideo15ByT5TextProjection(nn.Module):
+    def __init__(self, in_features: int, hidden_size: int, out_features: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(in_features)
+        self.linear_1 = nn.Linear(in_features, hidden_size)
+        self.linear_2 = nn.Linear(hidden_size, hidden_size)
+        self.linear_3 = nn.Linear(hidden_size, out_features)
+        self.act_fn = nn.GELU()
+
+    def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(encoder_hidden_states)
+        hidden_states = self.linear_1(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+        hidden_states = self.linear_3(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideo15ImageProjection(nn.Module):
+    def __init__(self, in_channels: int, hidden_size: int):
+        super().__init__()
+        self.norm_in = nn.LayerNorm(in_channels)
+        self.linear_1 = nn.Linear(in_channels, in_channels)
+        self.act_fn = nn.GELU()
+        self.linear_2 = nn.Linear(in_channels, hidden_size)
+        self.norm_out = nn.LayerNorm(hidden_size)
+
+    def forward(self, image_embeds: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm_in(image_embeds)
+        hidden_states = self.linear_1(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        hidden_states = self.norm_out(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideo15Attention(nn.Module):
+    """Dual-stream joint attention with TP optimization.
+
+    Key difference from FluxAttention: RoPE is applied **only** to the video
+    stream Q/K *before* concatenation with the encoder stream, whereas Flux
+    applies RoPE to the concatenated Q/K.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        bias: bool = True,
+        added_kv_proj_dim: int | None = None,
+        added_proj_bias: bool | None = True,
+        out_bias: bool = True,
+        eps: float = 1e-6,
+        out_dim: int | None = None,
+    ):
+        super().__init__()
+
+        self.head_dim = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.added_kv_proj_dim = added_kv_proj_dim
+
+        self.norm_q = RMSNorm(dim_head, eps=eps)
+        self.norm_k = RMSNorm(dim_head, eps=eps)
+
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=query_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.heads,
+            bias=bias,
+        )
+
+        self.to_out = nn.ModuleList(
+            [
+                RowParallelLinear(
+                    self.inner_dim,
+                    self.out_dim,
+                    bias=out_bias,
+                    input_is_parallel=True,
+                    return_bias=False,
+                ),
+                nn.Identity(),  # placeholder for dropout (none used)
+            ]
+        )
+
+        if added_kv_proj_dim is not None:
+            self.norm_added_q = RMSNorm(dim_head, eps=eps)
+            self.norm_added_k = RMSNorm(dim_head, eps=eps)
+
+            self.add_kv_proj = QKVParallelLinear(
+                hidden_size=self.added_kv_proj_dim,
+                head_size=self.head_dim,
+                total_num_heads=self.heads,
+                bias=added_proj_bias,
+            )
+
+            self.to_add_out = RowParallelLinear(
+                self.inner_dim,
+                query_dim,
+                bias=out_bias,
+                input_is_parallel=True,
+                return_bias=False,
+            )
+
+        self.rope = RotaryEmbedding(is_neox_style=False)
+        self.attn = Attention(
+            num_heads=self.to_qkv.num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.to_qkv.num_kv_heads,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        qkv, _ = self.to_qkv(hidden_states)
+        q_size = self.to_qkv.num_heads * self.head_dim
+        kv_size = self.to_qkv.num_kv_heads * self.head_dim
+        query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
+
+        query = query.unflatten(-1, (self.to_qkv.num_heads, -1))
+        key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
+        value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+            cos = cos.to(query.dtype)
+            sin = sin.to(query.dtype)
+            query = self.rope(query, cos, sin)
+            key = self.rope(key, cos, sin)
+
+        if encoder_hidden_states is not None:
+            encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
+            add_q_size = self.add_kv_proj.num_heads * self.head_dim
+            add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
+            encoder_query, encoder_key, encoder_value = encoder_qkv.split(
+                [add_q_size, add_kv_size, add_kv_size], dim=-1
+            )
+
+            encoder_query = encoder_query.unflatten(-1, (self.add_kv_proj.num_heads, -1))
+            encoder_key = encoder_key.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
+            encoder_value = encoder_value.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
+
+            encoder_query = self.norm_added_q(encoder_query)
+            encoder_key = self.norm_added_k(encoder_key)
+
+            query = torch.cat([query, encoder_query], dim=1)
+            key = torch.cat([key, encoder_key], dim=1)
+            value = torch.cat([value, encoder_value], dim=1)
+
+        attn_metadata = None
+        if attention_mask is not None:
+            seq_len = query.shape[1]
+            # Pad mask to full sequence length (video + encoder tokens)
+            # Keep mask 2D (batch_size, seq_len) — each attention backend
+            # handles reshaping internally (flash_attn uses unpadding,
+            # SDPA expands to 4D via _maybe_reshape_attn_mask).
+            attention_mask = F.pad(attention_mask, (seq_len - attention_mask.shape[1], 0), value=True)
+            attention_mask = attention_mask.bool()
+            attn_metadata = AttentionMetadata(attn_mask=attention_mask)
+
+        hidden_states = self.attn(query, key, value, attn_metadata)
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
+                [hidden_states.shape[1] - encoder_hidden_states.shape[1], encoder_hidden_states.shape[1]], dim=1
+            )
+            hidden_states = self.to_out[0](hidden_states)
+            encoder_hidden_states = self.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            hidden_states = self.to_out[0](hidden_states)
+            return hidden_states
+
+
+class HunyuanVideo15TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float = 4.0,
+        qk_norm: str = "rms_norm",
+    ) -> None:
+        super().__init__()
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
+        self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
+
+        self.attn = HunyuanVideo15Attention(
+            query_dim=hidden_size,
+            added_kv_proj_dim=hidden_size,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=hidden_size,
+            bias=True,
+            eps=1e-6,
+        )
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio)
+
+        self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+            encoder_hidden_states, emb=temb
+        )
+
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            attention_mask=attention_mask,
+            image_rotary_emb=freqs_cis,
+        )
+
+        hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1)
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+        ff_output = self.ff(norm_hidden_states)
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+
+        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+
+        return hidden_states, encoder_hidden_states
+
+
+class HunyuanVideo15Transformer3DModel(nn.Module):
+    """
+    HunyuanVideo-1.5 Transformer with TP-optimized dual-stream attention.
+
+    Ported from diffusers ``HunyuanVideo15Transformer3DModel`` with vllm-omni
+    tensor-parallel layers for the 54 main transformer blocks.
+    """
+
+    _repeated_blocks = ["HunyuanVideo15TransformerBlock"]
+    _layerwise_offload_blocks_attr = "transformer_blocks"
+    packed_modules_mapping = {
+        "to_qkv": ["to_q", "to_k", "to_v"],
+        "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],
+    }
+
+    @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,
+        od_config: OmniDiffusionConfig,
+        in_channels: int = 65,
+        out_channels: int = 32,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 128,
+        num_layers: int = 54,
+        num_refiner_layers: int = 2,
+        mlp_ratio: float = 4.0,
+        patch_size: int = 1,
+        patch_size_t: int = 1,
+        qk_norm: str = "rms_norm",
+        text_embed_dim: int = 3584,
+        text_embed_2_dim: int = 1472,
+        image_embed_dim: int = 1152,
+        rope_theta: float = 256.0,
+        rope_axes_dim: tuple[int, ...] = (16, 56, 56),
+        target_size: int = 640,
+        task_type: str = "i2v",
+        use_meanflow: bool = False,
+    ):
+        super().__init__()
+
+        # Allow tf_model_config to override num_layers
+        model_config = od_config.tf_model_config
+        num_layers = getattr(model_config, "num_layers", num_layers)
+
+        self.parallel_config = od_config.parallel_config
+        self.in_channels = in_channels
+        self.out_channels = out_channels or in_channels
+        inner_dim = num_attention_heads * attention_head_dim
+        self.inner_dim = inner_dim
+        self.patch_size = patch_size
+        self.patch_size_t = patch_size_t
+
+        self.x_embedder = HunyuanVideo15PatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
+        self.image_embedder = HunyuanVideo15ImageProjection(image_embed_dim, inner_dim)
+
+        self.context_embedder = HunyuanVideo15TokenRefiner(
+            text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers
+        )
+        self.context_embedder_2 = HunyuanVideo15ByT5TextProjection(text_embed_2_dim, 2048, inner_dim)
+
+        self.time_embed = HunyuanVideo15TimeEmbedding(inner_dim, use_meanflow=use_meanflow)
+        self.cond_type_embed = nn.Embedding(3, inner_dim)
+
+        self.rope = HunyuanVideo15RotaryPosEmbed(patch_size, patch_size_t, list(rope_axes_dim), rope_theta)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                HunyuanVideo15TransformerBlock(
+                    num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * self.out_channels)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        timestep_r: torch.LongTensor | None = None,
+        encoder_hidden_states_2: torch.Tensor | None = None,
+        encoder_attention_mask_2: torch.Tensor | None = None,
+        image_embeds: torch.Tensor | None = None,
+        image_embeds_mask: torch.Tensor | None = None,
+        attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+    ) -> torch.Tensor | Transformer2DModelOutput:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size_t, self.patch_size, self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        image_rotary_emb = self.rope(hidden_states)
+
+        temb = self.time_embed(timestep, timestep_r=timestep_r)
+
+        hidden_states = self.x_embedder(hidden_states)
+
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask)
+
+        encoder_hidden_states_cond_emb = self.cond_type_embed(
+            torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb
+
+        encoder_hidden_states_2 = self.context_embedder_2(encoder_hidden_states_2)
+
+        encoder_hidden_states_2_cond_emb = self.cond_type_embed(
+            torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb
+
+        encoder_hidden_states_3 = self.image_embedder(image_embeds)
+        if image_embeds_mask is not None:
+            encoder_attention_mask_3 = image_embeds_mask
+        else:
+            # Fallback: detect T2V by checking if image_embeds are all zeros
+            is_t2v = torch.all(image_embeds == 0)
+            if is_t2v:
+                encoder_hidden_states_3 = encoder_hidden_states_3 * 0.0
+                encoder_attention_mask_3 = torch.zeros(
+                    (batch_size, encoder_hidden_states_3.shape[1]),
+                    dtype=encoder_attention_mask.dtype,
+                    device=encoder_attention_mask.device,
+                )
+            else:
+                encoder_attention_mask_3 = torch.ones(
+                    (batch_size, encoder_hidden_states_3.shape[1]),
+                    dtype=encoder_attention_mask.dtype,
+                    device=encoder_attention_mask.device,
+                )
+        encoder_hidden_states_3_cond_emb = self.cond_type_embed(
+            2 * torch.ones_like(encoder_hidden_states_3[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb
+
+        # Token reordering: [valid_image, valid_byte5, valid_mllm, padding]
+        encoder_attention_mask = encoder_attention_mask.bool()
+        encoder_attention_mask_2 = encoder_attention_mask_2.bool()
+        encoder_attention_mask_3 = encoder_attention_mask_3.bool()
+        new_encoder_hidden_states = []
+        new_encoder_attention_mask = []
+
+        for text, text_mask, text_2, text_mask_2, image, image_mask in zip(
+            encoder_hidden_states,
+            encoder_attention_mask,
+            encoder_hidden_states_2,
+            encoder_attention_mask_2,
+            encoder_hidden_states_3,
+            encoder_attention_mask_3,
+        ):
+            new_encoder_hidden_states.append(
+                torch.cat(
+                    [
+                        image[image_mask],
+                        text_2[text_mask_2],
+                        text[text_mask],
+                        image[~image_mask],
+                        torch.zeros_like(text_2[~text_mask_2]),
+                        torch.zeros_like(text[~text_mask]),
+                    ],
+                    dim=0,
+                )
+            )
+            new_encoder_attention_mask.append(
+                torch.cat(
+                    [
+                        image_mask[image_mask],
+                        text_mask_2[text_mask_2],
+                        text_mask[text_mask],
+                        image_mask[~image_mask],
+                        text_mask_2[~text_mask_2],
+                        text_mask[~text_mask],
+                    ],
+                    dim=0,
+                )
+            )
+
+        encoder_hidden_states = torch.stack(new_encoder_hidden_states)
+        encoder_attention_mask = torch.stack(new_encoder_attention_mask)
+
+        for block in self.transformer_blocks:
+            hidden_states, encoder_hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                temb,
+                encoder_attention_mask,
+                image_rotary_emb,
+            )
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_num_frames, post_patch_height, post_patch_width, -1, p_t, p_h, p_w
+        )
+        hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
+        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        if not return_dict:
+            return (hidden_states,)
+
+        return Transformer2DModelOutput(sample=hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            # self-attn
+            (".to_qkv", ".to_q", "q"),
+            (".to_qkv", ".to_k", "k"),
+            (".to_qkv", ".to_v", "v"),
+            # cross-attn
+            (".add_kv_proj", ".add_q_proj", "q"),
+            (".add_kv_proj", ".add_k_proj", "k"),
+            (".add_kv_proj", ".add_v_proj", "v"),
+        ]
+
+        params_dict = dict(self.named_parameters())
+
+        # Load buffers for beta and eps (if any normalizations use them)
+        for name, buffer in self.named_buffers():
+            if name.endswith(".beta") or name.endswith(".eps"):
+                params_dict[name] = buffer
+
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            original_name = name
+            lookup_name = name
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in original_name:
+                    continue
+                lookup_name = original_name.replace(weight_name, param_name)
+                if lookup_name not in params_dict:
+                    break
+                param = params_dict[lookup_name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Handle .to_out.0. -> .to_out. remapping
+                if lookup_name not in params_dict and ".to_out.0." in lookup_name:
+                    lookup_name = lookup_name.replace(".to_out.0.", ".to_out.")
+                if lookup_name not in params_dict:
+                    continue
+                param = params_dict[lookup_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(original_name)
+            loaded_params.add(lookup_name)
+        return loaded_params

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py

@@ -0,0 +1,513 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+import torch
+from diffusers import AutoencoderKLHunyuanVideo15
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from torch import nn
+from transformers import AutoConfig, ByT5Tokenizer, Qwen2_5_VLTextModel, Qwen2Tokenizer
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel
+from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.platforms import current_omni_platform
+
+logger = logging.getLogger(__name__)
+
+
+def retrieve_latents(
+    encoder_output: torch.Tensor,
+    generator: torch.Generator | None = None,
+    sample_mode: str = "sample",
+) -> torch.Tensor:
+    """Extract latents from VAE encoder output."""
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def format_text_input(prompt: list[str], system_message: str) -> list[list[dict[str, str]]]:
+    return [
+        [{"role": "system", "content": system_message}, {"role": "user", "content": p if p else " "}] for p in prompt
+    ]
+
+
+def extract_glyph_texts(prompt: str) -> str | None:
+    pattern = r'"(.*?)"|' + r"\u201c(.*?)\u201d"
+    matches = re.findall(pattern, prompt)
+    result = [match[0] or match[1] for match in matches]
+    result = list(dict.fromkeys(result)) if len(result) > 1 else result
+
+    if result:
+        return ". ".join([f'Text "{text}"' for text in result]) + ". "
+    return None
+
+
+def get_hunyuan_video_15_post_process_func(od_config: OmniDiffusionConfig):
+    video_processor = VideoProcessor(vae_scale_factor=16)
+
+    def post_process_func(video: torch.Tensor, output_type: str = "np"):
+        if output_type == "latent":
+            return video
+        return video_processor.postprocess_video(video, output_type=output_type)
+
+    return post_process_func
+
+
+class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+
+        self.device = get_local_device()
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+
+        model = od_config.model
+        local_files_only = os.path.exists(model)
+
+        self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+        self.text_encoder = Qwen2_5_VLTextModel.from_pretrained(
+            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
+        ).to(self.device)
+
+        self.tokenizer_2 = ByT5Tokenizer.from_pretrained(
+            model, subfolder="tokenizer_2", local_files_only=local_files_only
+        )
+        t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only)
+        self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device)
+
+        self.vae = AutoencoderKLHunyuanVideo15.from_pretrained(
+            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        ).to(self.device)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+
+        # Override the scheduler's shift if flow_shift is explicitly provided.
+        # Must set _shift directly since the property has no setter.
+        if od_config.flow_shift is not None:
+            self.scheduler._shift = od_config.flow_shift
+
+        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel)
+        self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs)
+
+        # Check if model uses meanflow (distilled variants)
+        self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False)
+
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="text_encoder_2",
+                revision=None,
+                prefix="text_encoder_2.",
+                fall_back_to_pt=True,
+            ),
+        ]
+
+        self.vae_scale_factor_temporal = (
+            self.vae.temporal_compression_ratio if hasattr(self.vae, "temporal_compression_ratio") else 4
+        )
+        self.vae_scale_factor_spatial = (
+            self.vae.spatial_compression_ratio if hasattr(self.vae, "spatial_compression_ratio") else 16
+        )
+        self.num_channels_latents = self.vae.config.latent_channels if hasattr(self.vae, "config") else 32
+
+        # fmt: off
+        # NOTE: whitespace must match the diffusers reference exactly because
+        # prompt_template_encode_start_idx=108 depends on the tokenised length of
+        # this system message.  The backslash-continuation includes leading spaces.
+        self.system_message = "You are a helpful assistant. Describe the video by detailing the following aspects: \
+        1. The main content and theme of the video. \
+        2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \
+        3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \
+        4. background environment, light, style and atmosphere. \
+        5. camera angles, movements, and transitions used in the video."
+        # fmt: on
+        self.prompt_template_encode_start_idx = 108
+        self.tokenizer_max_length = 1000
+        self.tokenizer_2_max_length = 256
+        self.vision_num_semantic_tokens = 729
+        self.vision_states_dim = 1152
+
+        self._guidance_scale = None
+        self._num_timesteps = None
+        self._current_timestep = None
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    def _get_mllm_prompt_embeds(
+        self,
+        prompt: list[str],
+        device: torch.device,
+        dtype: torch.dtype,
+        num_hidden_layers_to_skip: int = 2,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prompt_formatted = format_text_input(prompt, self.system_message)
+
+        text_inputs = self.tokenizer.apply_chat_template(
+            prompt_formatted,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            padding="max_length",
+            max_length=self.tokenizer_max_length + self.prompt_template_encode_start_idx,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids.to(device=device)
+        prompt_attention_mask = text_inputs.attention_mask.to(device=device)
+
+        prompt_embeds = self.text_encoder(
+            input_ids=text_input_ids,
+            attention_mask=prompt_attention_mask,
+            output_hidden_states=True,
+        ).hidden_states[-(num_hidden_layers_to_skip + 1)]
+
+        crop_start = self.prompt_template_encode_start_idx
+        if crop_start is not None and crop_start > 0:
+            prompt_embeds = prompt_embeds[:, crop_start:]
+            prompt_attention_mask = prompt_attention_mask[:, crop_start:]
+
+        return prompt_embeds.to(dtype=dtype), prompt_attention_mask
+
+    def _get_byte5_prompt_embeds(
+        self,
+        prompt: list[str],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prompt_embeds_list = []
+        prompt_embeds_mask_list = []
+
+        for p in prompt:
+            glyph_text = extract_glyph_texts(p)
+
+            if glyph_text is None:
+                glyph_text_embeds = torch.zeros(
+                    (1, self.tokenizer_2_max_length, self.text_encoder_2.config.d_model),
+                    device=device,
+                    dtype=dtype,
+                )
+                glyph_text_embeds_mask = torch.zeros((1, self.tokenizer_2_max_length), device=device, dtype=torch.int64)
+            else:
+                txt_tokens = self.tokenizer_2(
+                    glyph_text,
+                    padding="max_length",
+                    max_length=self.tokenizer_2_max_length,
+                    truncation=True,
+                    add_special_tokens=True,
+                    return_tensors="pt",
+                ).to(device)
+
+                glyph_text_embeds = self.text_encoder_2(
+                    input_ids=txt_tokens.input_ids,
+                    attention_mask=txt_tokens.attention_mask.float(),
+                )[0]
+                glyph_text_embeds = glyph_text_embeds.to(device=device, dtype=dtype)
+                glyph_text_embeds_mask = txt_tokens.attention_mask.to(device=device)
+
+            prompt_embeds_list.append(glyph_text_embeds)
+            prompt_embeds_mask_list.append(glyph_text_embeds_mask)
+
+        return torch.cat(prompt_embeds_list, dim=0), torch.cat(prompt_embeds_mask_list, dim=0)
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        device: torch.device,
+        dtype: torch.dtype,
+        negative_prompt: str | list[str] | None = None,
+        do_classifier_free_guidance: bool = False,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        prompt_embeds, prompt_embeds_mask = self._get_mllm_prompt_embeds(prompt, device, dtype)
+        prompt_embeds_2, prompt_embeds_mask_2 = self._get_byte5_prompt_embeds(prompt, device, dtype)
+
+        # Cast masks to model dtype to match diffusers behaviour
+        prompt_embeds_mask = prompt_embeds_mask.to(dtype=dtype)
+        prompt_embeds_mask_2 = prompt_embeds_mask_2.to(dtype=dtype)
+
+        negative_prompt_embeds = None
+        negative_prompt_embeds_mask = None
+        negative_prompt_embeds_2 = None
+        negative_prompt_embeds_mask_2 = None
+
+        if do_classifier_free_guidance and negative_prompt is not None:
+            negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_embeds, negative_prompt_embeds_mask = self._get_mllm_prompt_embeds(
+                negative_prompt, device, dtype
+            )
+            negative_prompt_embeds_2, negative_prompt_embeds_mask_2 = self._get_byte5_prompt_embeds(
+                negative_prompt, device, dtype
+            )
+            negative_prompt_embeds_mask = negative_prompt_embeds_mask.to(dtype=dtype)
+            negative_prompt_embeds_mask_2 = negative_prompt_embeds_mask_2.to(dtype=dtype)
+        elif do_classifier_free_guidance:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self._get_mllm_prompt_embeds([""], device, dtype)
+            negative_prompt_embeds_2, negative_prompt_embeds_mask_2 = self._get_byte5_prompt_embeds([""], device, dtype)
+            negative_prompt_embeds_mask = negative_prompt_embeds_mask.to(dtype=dtype)
+            negative_prompt_embeds_mask_2 = negative_prompt_embeds_mask_2.to(dtype=dtype)
+
+        return (
+            prompt_embeds,
+            prompt_embeds_mask,
+            prompt_embeds_2,
+            prompt_embeds_mask_2,
+            negative_prompt_embeds,
+            negative_prompt_embeds_mask,
+            negative_prompt_embeds_2,
+            negative_prompt_embeds_mask_2,
+        )
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        shape = (
+            batch_size,
+            self.num_channels_latents,
+            (num_frames - 1) // self.vae_scale_factor_temporal + 1,
+            int(height) // self.vae_scale_factor_spatial,
+            int(width) // self.vae_scale_factor_spatial,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.")
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        return latents
+
+    def prepare_cond_latents_and_mask(
+        self, latents: torch.Tensor, dtype: torch.dtype, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare zero condition latents and mask for T2V mode."""
+        batch, channels, frames, height, width = latents.shape
+        cond_latents = torch.zeros(batch, channels, frames, height, width, dtype=dtype, device=device)
+        mask = torch.zeros(batch, 1, frames, height, width, dtype=dtype, device=device)
+        return cond_latents, mask
+
+    def predict_noise(self, **kwargs: Any) -> torch.Tensor:
+        return self.transformer(**kwargs)[0]
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 6.0,
+        height: int = 480,
+        width: int = 832,
+        num_frames: int = 121,
+        output_type: str | None = "np",
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        **kwargs,
+    ) -> DiffusionOutput:
+        if len(req.prompts) > 1:
+            raise ValueError("This model only supports a single prompt per request.")
+        if len(req.prompts) == 1:
+            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
+            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
+        else:
+            raise ValueError("Prompt is required for HunyuanVideo-1.5 generation.")
+
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_frames_val = req.sampling_params.num_frames if req.sampling_params.num_frames else num_frames
+        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        self._guidance_scale = guidance_scale
+
+        do_cfg = guidance_scale > 1.0
+
+        device = self.device
+        dtype = self.transformer.transformer_blocks[0].norm1.linear.weight.dtype
+
+        if generator is None:
+            generator = req.sampling_params.generator
+        if generator is None and req.sampling_params.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
+
+        (
+            prompt_embeds,
+            prompt_embeds_mask,
+            prompt_embeds_2,
+            prompt_embeds_mask_2,
+            negative_prompt_embeds,
+            negative_prompt_embeds_mask,
+            negative_prompt_embeds_2,
+            negative_prompt_embeds_mask_2,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            device=device,
+            dtype=dtype,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=do_cfg,
+        )
+
+        batch_size = prompt_embeds.shape[0]
+
+        latents = self.prepare_latents(
+            batch_size=batch_size,
+            height=height,
+            width=width,
+            num_frames=num_frames_val,
+            dtype=dtype,
+            device=device,
+            generator=generator,
+            latents=req.sampling_params.latents,
+        )
+        cond_latents, mask = self.prepare_cond_latents_and_mask(latents, dtype, device)
+
+        # Image embeds (zeros for T2V, no mask — let the transformer detect T2V
+        # via the all-zeros check, matching the diffusers reference behaviour)
+        image_embeds = torch.zeros(
+            batch_size,
+            self.vision_num_semantic_tokens,
+            self.vision_states_dim,
+            dtype=dtype,
+            device=device,
+        )
+
+        # Timesteps — the scheduler handles flow_shift via its config (`shift` param).
+        # We just provide unshifted linear sigmas, matching the diffusers reference.
+        sigmas = np.linspace(1.0, 0.0, num_steps + 1)[:-1]
+        self.scheduler.set_timesteps(sigmas=sigmas, device=device)
+        timesteps = self.scheduler.timesteps
+        self._num_timesteps = len(timesteps)
+
+        for i, t in enumerate(timesteps):
+            self._current_timestep = t
+
+            latent_model_input = torch.cat([latents, cond_latents, mask], dim=1)
+            timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
+
+            timestep_r = None
+            if self.use_meanflow:
+                if i == len(timesteps) - 1:
+                    timestep_r = torch.tensor([0.0], device=device)
+                else:
+                    timestep_r = timesteps[i + 1]
+                timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype)
+
+            positive_kwargs = {
+                "hidden_states": latent_model_input,
+                "timestep": timestep,
+                "timestep_r": timestep_r,
+                "encoder_hidden_states": prompt_embeds,
+                "encoder_attention_mask": prompt_embeds_mask,
+                "encoder_hidden_states_2": prompt_embeds_2,
+                "encoder_attention_mask_2": prompt_embeds_mask_2,
+                "image_embeds": image_embeds,
+                "return_dict": False,
+            }
+
+            negative_kwargs = None
+            if do_cfg and negative_prompt_embeds is not None:
+                negative_kwargs = {
+                    "hidden_states": latent_model_input,
+                    "timestep": timestep,
+                    "timestep_r": timestep_r,
+                    "encoder_hidden_states": negative_prompt_embeds,
+                    "encoder_attention_mask": negative_prompt_embeds_mask,
+                    "encoder_hidden_states_2": negative_prompt_embeds_2,
+                    "encoder_attention_mask_2": negative_prompt_embeds_mask_2,
+                    "image_embeds": image_embeds,
+                    "return_dict": False,
+                }
+
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_cfg and negative_kwargs is not None,
+                true_cfg_scale=guidance_scale,
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+                cfg_normalize=req.sampling_params.cfg_normalize,
+            )
+
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred,
+                t,
+                latents,
+                do_true_cfg=do_cfg and negative_kwargs is not None,
+            )
+
+        self._current_timestep = None
+
+        if current_omni_platform.is_available():
+            current_omni_platform.empty_cache()
+
+        if output_type == "latent":
+            output = latents
+        else:
+            latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor
+            output = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(output=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py

@@ -0,0 +1,594 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Iterable
+from typing import Any, cast
+
+import numpy as np
+import PIL.Image
+import torch
+from diffusers import AutoencoderKLHunyuanVideo15
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from torch import nn
+from transformers import (
+    AutoConfig,
+    ByT5Tokenizer,
+    Qwen2_5_VLTextModel,
+    Qwen2Tokenizer,
+    SiglipImageProcessor,
+    SiglipVisionModel,
+)
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel
+from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import (
+    extract_glyph_texts,
+    format_text_input,
+    get_hunyuan_video_15_post_process_func,
+    retrieve_latents,
+)
+from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.platforms import current_omni_platform
+
+logger = logging.getLogger(__name__)
+
+
+def get_hunyuan_video_15_i2v_post_process_func(od_config: OmniDiffusionConfig):
+    return get_hunyuan_video_15_post_process_func(od_config)
+
+
+def get_hunyuan_video_15_i2v_pre_process_func(od_config: OmniDiffusionConfig):
+    """Pre-process function for I2V: load and resize image."""
+
+    max_area = 480 * 832
+    divisor = 16  # Must be divisible by VAE spatial compression
+
+    def pre_process_func(req: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        if not req.prompts:
+            return req
+
+        prompt_data = req.prompts[0]
+        if isinstance(prompt_data, str):
+            return req
+
+        multi_modal_data = prompt_data.get("multi_modal_data", {})
+        raw_image = multi_modal_data.get("image", None)
+
+        if raw_image is None:
+            return req
+
+        if isinstance(raw_image, list):
+            raw_image = raw_image[0]
+
+        if isinstance(raw_image, str):
+            image = PIL.Image.open(raw_image).convert("RGB")
+        elif isinstance(raw_image, PIL.Image.Image):
+            image = raw_image.convert("RGB")
+        else:
+            return req
+
+        height = req.sampling_params.height
+        width = req.sampling_params.width
+
+        if height is None or width is None:
+            w, h = image.size
+            aspect = w / h
+            target_h = int((max_area / aspect) ** 0.5)
+            target_w = int(target_h * aspect)
+            height = height or ((target_h + divisor - 1) // divisor * divisor)
+            width = width or ((target_w + divisor - 1) // divisor * divisor)
+            req.sampling_params.height = height
+            req.sampling_params.width = width
+
+        return req
+
+    return pre_process_func
+
+
+class HunyuanVideo15I2VPipeline(nn.Module, CFGParallelMixin, SupportImageInput):
+    support_image_input = True
+    color_format = "RGB"
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+
+        self.device = get_local_device()
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+
+        model = od_config.model
+        local_files_only = os.path.exists(model)
+
+        self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+        self.text_encoder = Qwen2_5_VLTextModel.from_pretrained(
+            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
+        ).to(self.device)
+
+        self.tokenizer_2 = ByT5Tokenizer.from_pretrained(
+            model, subfolder="tokenizer_2", local_files_only=local_files_only
+        )
+        t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only)
+        self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device)
+
+        self.image_encoder = SiglipVisionModel.from_pretrained(
+            model, subfolder="image_encoder", torch_dtype=dtype, local_files_only=local_files_only
+        ).to(self.device)
+        self.feature_extractor = SiglipImageProcessor.from_pretrained(
+            model, subfolder="feature_extractor", local_files_only=local_files_only
+        )
+
+        self.vae = AutoencoderKLHunyuanVideo15.from_pretrained(
+            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        ).to(self.device)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+
+        # Override the scheduler's shift if flow_shift is explicitly provided.
+        # Must set _shift directly since the property has no setter.
+        if od_config.flow_shift is not None:
+            self.scheduler._shift = od_config.flow_shift
+
+        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel)
+        self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs)
+
+        self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False)
+
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="text_encoder_2",
+                revision=None,
+                prefix="text_encoder_2.",
+                fall_back_to_pt=True,
+            ),
+        ]
+
+        self.vae_scale_factor_temporal = (
+            self.vae.temporal_compression_ratio if hasattr(self.vae, "temporal_compression_ratio") else 4
+        )
+        self.vae_scale_factor_spatial = (
+            self.vae.spatial_compression_ratio if hasattr(self.vae, "spatial_compression_ratio") else 16
+        )
+        self.num_channels_latents = self.vae.config.latent_channels if hasattr(self.vae, "config") else 32
+
+        # fmt: off
+        # NOTE: whitespace must match the diffusers reference exactly because
+        # prompt_template_encode_start_idx=108 depends on the tokenised length of
+        # this system message.  The backslash-continuation includes leading spaces.
+        self.system_message = "You are a helpful assistant. Describe the video by detailing the following aspects: \
+        1. The main content and theme of the video. \
+        2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \
+        3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \
+        4. background environment, light, style and atmosphere. \
+        5. camera angles, movements, and transitions used in the video."
+        # fmt: on
+        self.prompt_template_encode_start_idx = 108
+        self.tokenizer_max_length = 1000
+        self.tokenizer_2_max_length = 256
+        self.vision_num_semantic_tokens = 729
+        self.vision_states_dim = 1152
+
+        self._guidance_scale = None
+        self._num_timesteps = None
+        self._current_timestep = None
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    def _get_mllm_prompt_embeds(
+        self,
+        prompt: list[str],
+        device: torch.device,
+        dtype: torch.dtype,
+        num_hidden_layers_to_skip: int = 2,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prompt_formatted = format_text_input(prompt, self.system_message)
+
+        text_inputs = self.tokenizer.apply_chat_template(
+            prompt_formatted,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            padding="max_length",
+            max_length=self.tokenizer_max_length + self.prompt_template_encode_start_idx,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids.to(device=device)
+        prompt_attention_mask = text_inputs.attention_mask.to(device=device)
+
+        prompt_embeds = self.text_encoder(
+            input_ids=text_input_ids,
+            attention_mask=prompt_attention_mask,
+            output_hidden_states=True,
+        ).hidden_states[-(num_hidden_layers_to_skip + 1)]
+
+        crop_start = self.prompt_template_encode_start_idx
+        if crop_start is not None and crop_start > 0:
+            prompt_embeds = prompt_embeds[:, crop_start:]
+            prompt_attention_mask = prompt_attention_mask[:, crop_start:]
+
+        return prompt_embeds.to(dtype=dtype), prompt_attention_mask
+
+    def _get_byte5_prompt_embeds(
+        self,
+        prompt: list[str],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prompt_embeds_list = []
+        prompt_embeds_mask_list = []
+
+        for p in prompt:
+            glyph_text = extract_glyph_texts(p)
+
+            if glyph_text is None:
+                glyph_text_embeds = torch.zeros(
+                    (1, self.tokenizer_2_max_length, self.text_encoder_2.config.d_model),
+                    device=device,
+                    dtype=dtype,
+                )
+                glyph_text_embeds_mask = torch.zeros((1, self.tokenizer_2_max_length), device=device, dtype=torch.int64)
+            else:
+                txt_tokens = self.tokenizer_2(
+                    glyph_text,
+                    padding="max_length",
+                    max_length=self.tokenizer_2_max_length,
+                    truncation=True,
+                    add_special_tokens=True,
+                    return_tensors="pt",
+                ).to(device)
+
+                glyph_text_embeds = self.text_encoder_2(
+                    input_ids=txt_tokens.input_ids,
+                    attention_mask=txt_tokens.attention_mask.float(),
+                )[0]
+                glyph_text_embeds = glyph_text_embeds.to(device=device, dtype=dtype)
+                glyph_text_embeds_mask = txt_tokens.attention_mask.to(device=device)
+
+            prompt_embeds_list.append(glyph_text_embeds)
+            prompt_embeds_mask_list.append(glyph_text_embeds_mask)
+
+        return torch.cat(prompt_embeds_list, dim=0), torch.cat(prompt_embeds_mask_list, dim=0)
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        device: torch.device,
+        dtype: torch.dtype,
+        negative_prompt: str | list[str] | None = None,
+        do_classifier_free_guidance: bool = False,
+    ) -> tuple:
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        prompt_embeds, prompt_embeds_mask = self._get_mllm_prompt_embeds(prompt, device, dtype)
+        prompt_embeds_2, prompt_embeds_mask_2 = self._get_byte5_prompt_embeds(prompt, device, dtype)
+
+        # Cast masks to model dtype to match diffusers behaviour
+        prompt_embeds_mask = prompt_embeds_mask.to(dtype=dtype)
+        prompt_embeds_mask_2 = prompt_embeds_mask_2.to(dtype=dtype)
+
+        negative_prompt_embeds = None
+        negative_prompt_embeds_mask = None
+        negative_prompt_embeds_2 = None
+        negative_prompt_embeds_mask_2 = None
+
+        if do_classifier_free_guidance and negative_prompt is not None:
+            negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_embeds, negative_prompt_embeds_mask = self._get_mllm_prompt_embeds(
+                negative_prompt, device, dtype
+            )
+            negative_prompt_embeds_2, negative_prompt_embeds_mask_2 = self._get_byte5_prompt_embeds(
+                negative_prompt, device, dtype
+            )
+            negative_prompt_embeds_mask = negative_prompt_embeds_mask.to(dtype=dtype)
+            negative_prompt_embeds_mask_2 = negative_prompt_embeds_mask_2.to(dtype=dtype)
+        elif do_classifier_free_guidance:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self._get_mllm_prompt_embeds([""], device, dtype)
+            negative_prompt_embeds_2, negative_prompt_embeds_mask_2 = self._get_byte5_prompt_embeds([""], device, dtype)
+            negative_prompt_embeds_mask = negative_prompt_embeds_mask.to(dtype=dtype)
+            negative_prompt_embeds_mask_2 = negative_prompt_embeds_mask_2.to(dtype=dtype)
+
+        return (
+            prompt_embeds,
+            prompt_embeds_mask,
+            prompt_embeds_2,
+            prompt_embeds_mask_2,
+            negative_prompt_embeds,
+            negative_prompt_embeds_mask,
+            negative_prompt_embeds_2,
+            negative_prompt_embeds_mask_2,
+        )
+
+    def _get_image_embeds(
+        self,
+        image: PIL.Image.Image,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Get semantic image embeddings from SigLIP."""
+        image_encoder_dtype = next(self.image_encoder.parameters()).dtype
+        pixel_values = self.feature_extractor(images=image, do_resize=True, return_tensors="pt").pixel_values.to(
+            device=device, dtype=image_encoder_dtype
+        )
+        image_enc_hidden_states = self.image_encoder(pixel_values).last_hidden_state
+        return image_enc_hidden_states
+
+    def _get_image_latents(
+        self,
+        image: PIL.Image.Image,
+        height: int,
+        width: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Encode image to VAE latents for first-frame conditioning."""
+        video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+        image_tensor = video_processor.preprocess(image, height=height, width=width)
+        image_tensor = image_tensor.unsqueeze(2).to(device=device, dtype=self.vae.dtype)  # [B, C, 1, H, W]
+        image_latents = retrieve_latents(self.vae.encode(image_tensor), sample_mode="argmax")
+        image_latents = image_latents * self.vae.config.scaling_factor
+        return image_latents
+
+    def prepare_cond_latents_and_mask(
+        self,
+        latents: torch.Tensor,
+        image: PIL.Image.Image,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare condition latents and mask for I2V: first frame = image, rest = zeros."""
+        batch, channels, frames, lat_height, lat_width = latents.shape
+
+        image_latents = self._get_image_latents(image, height, width, device)
+
+        # Repeat image latent across all frames, then zero out non-first frames
+        latent_condition = image_latents.repeat(batch, 1, frames, 1, 1)
+        latent_condition[:, :, 1:, :, :] = 0
+        latent_condition = latent_condition.to(device=device, dtype=dtype)
+
+        # Mask: 1.0 for first frame (conditioned), 0.0 for rest (to denoise)
+        latent_mask = torch.zeros(batch, 1, frames, lat_height, lat_width, dtype=dtype, device=device)
+        latent_mask[:, :, 0, :, :] = 1.0
+
+        return latent_condition, latent_mask
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        num_channels = self.num_channels_latents
+        shape = (
+            batch_size,
+            num_channels,
+            (num_frames - 1) // self.vae_scale_factor_temporal + 1,
+            int(height) // self.vae_scale_factor_spatial,
+            int(width) // self.vae_scale_factor_spatial,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.")
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        return latents
+
+    def predict_noise(self, **kwargs: Any) -> torch.Tensor:
+        return self.transformer(**kwargs)[0]
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 6.0,
+        height: int = 480,
+        width: int = 832,
+        num_frames: int = 121,
+        output_type: str | None = "np",
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        **kwargs,
+    ) -> DiffusionOutput:
+        if len(req.prompts) > 1:
+            raise ValueError("This model only supports a single prompt per request.")
+        if len(req.prompts) == 1:
+            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
+            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
+        else:
+            raise ValueError("Prompt is required for HunyuanVideo-1.5 I2V generation.")
+
+        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else {}
+        raw_image = multi_modal_data.get("image", None)
+        if isinstance(raw_image, list):
+            if len(raw_image) > 1:
+                logger.warning("Received a list of images. Only the first image will be used.")
+            raw_image = raw_image[0]
+
+        if raw_image is None:
+            raise ValueError(
+                "Image is required for HunyuanVideo-1.5 I2V generation. Pass it via multi_modal_data={'image': <image>}"
+            )
+
+        if isinstance(raw_image, str):
+            image = PIL.Image.open(raw_image).convert("RGB")
+        elif isinstance(raw_image, PIL.Image.Image):
+            image = raw_image.convert("RGB")
+        else:
+            image = cast(PIL.Image.Image, raw_image)
+
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_frames_val = req.sampling_params.num_frames if req.sampling_params.num_frames else num_frames
+        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        self._guidance_scale = guidance_scale
+
+        do_cfg = guidance_scale > 1.0
+
+        device = self.device
+        dtype = self.transformer.transformer_blocks[0].norm1.linear.weight.dtype
+
+        if generator is None:
+            generator = req.sampling_params.generator
+        if generator is None and req.sampling_params.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
+
+        (
+            prompt_embeds,
+            prompt_embeds_mask,
+            prompt_embeds_2,
+            prompt_embeds_mask_2,
+            negative_prompt_embeds,
+            negative_prompt_embeds_mask,
+            negative_prompt_embeds_2,
+            negative_prompt_embeds_mask_2,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            device=device,
+            dtype=dtype,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=do_cfg,
+        )
+
+        batch_size = prompt_embeds.shape[0]
+
+        # Get image embeddings from SigLIP (no explicit mask — the transformer
+        # detects I2V via non-zero image_embeds, matching diffusers behaviour)
+        image_embeds = self._get_image_embeds(image, device).to(dtype=dtype)
+
+        latents = self.prepare_latents(
+            batch_size=batch_size,
+            height=height,
+            width=width,
+            num_frames=num_frames_val,
+            dtype=dtype,
+            device=device,
+            generator=generator,
+            latents=req.sampling_params.latents,
+        )
+
+        cond_latents, mask = self.prepare_cond_latents_and_mask(latents, image, height, width, dtype, device)
+
+        # Timesteps — the scheduler handles flow_shift via its config (`shift` param).
+        sigmas = np.linspace(1.0, 0.0, num_steps + 1)[:-1]
+        self.scheduler.set_timesteps(sigmas=sigmas, device=device)
+        timesteps = self.scheduler.timesteps
+        self._num_timesteps = len(timesteps)
+
+        for i, t in enumerate(timesteps):
+            self._current_timestep = t
+
+            latent_model_input = torch.cat([latents, cond_latents, mask], dim=1)
+            timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
+
+            timestep_r = None
+            if self.use_meanflow:
+                if i == len(timesteps) - 1:
+                    timestep_r = torch.tensor([0.0], device=device)
+                else:
+                    timestep_r = timesteps[i + 1]
+                timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype)
+
+            positive_kwargs = {
+                "hidden_states": latent_model_input,
+                "timestep": timestep,
+                "timestep_r": timestep_r,
+                "encoder_hidden_states": prompt_embeds,
+                "encoder_attention_mask": prompt_embeds_mask,
+                "encoder_hidden_states_2": prompt_embeds_2,
+                "encoder_attention_mask_2": prompt_embeds_mask_2,
+                "image_embeds": image_embeds,
+                "return_dict": False,
+            }
+
+            negative_kwargs = None
+            if do_cfg and negative_prompt_embeds is not None:
+                # For I2V CFG, negative still uses image embeds (only text is unconditional)
+                negative_kwargs = {
+                    "hidden_states": latent_model_input,
+                    "timestep": timestep,
+                    "timestep_r": timestep_r,
+                    "encoder_hidden_states": negative_prompt_embeds,
+                    "encoder_attention_mask": negative_prompt_embeds_mask,
+                    "encoder_hidden_states_2": negative_prompt_embeds_2,
+                    "encoder_attention_mask_2": negative_prompt_embeds_mask_2,
+                    "image_embeds": image_embeds,
+                    "return_dict": False,
+                }
+
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_cfg and negative_kwargs is not None,
+                true_cfg_scale=guidance_scale,
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+                cfg_normalize=req.sampling_params.cfg_normalize,
+            )
+
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred,
+                t,
+                latents,
+                do_true_cfg=do_cfg and negative_kwargs is not None,
+            )
+
+        self._current_timestep = None
+
+        if current_omni_platform.is_available():
+            current_omni_platform.empty_cache()
+
+        if output_type == "latent":
+            output = latents
+        else:
+            latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor
+            output = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(output=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/quantization/base.py

@@ -31,14 +31,15 @@ class DiffusionQuantizationConfig(ABC):
     # The underlying vLLM config instance
     _vllm_config: "QuantizationConfig | None" = None
 
-    def get_name(self) -> str:
+    @classmethod
+    def get_name(cls) -> str:
         """Return the quantization method name (e.g., 'fp8', 'int8').
 
-        By default, delegates to the underlying vLLM config instance.
+        Delegates to the underlying vLLM config class's get_name().
         """
-        if self._vllm_config is not None:
-            return self._vllm_config.get_name()
-        raise NotImplementedError("Subclass must initialize _vllm_config or override get_name().")
+        if cls.quant_config_cls is not None:
+            return cls.quant_config_cls.get_name()
+        raise NotImplementedError("Subclass must set quant_config_cls or override get_name().")
 
     def get_vllm_quant_config(self) -> "QuantizationConfig | None":
         """Return the underlying vLLM QuantizationConfig for linear layers."""

vllm_omni/diffusion/registry.py

@@ -142,6 +142,16 @@ _DIFFUSION_MODELS = {
         "pipeline_dreamid_omni",
         "DreamIDOmniPipeline",
     ),
+    "HunyuanVideo15Pipeline": (
+        "hunyuan_video",
+        "pipeline_hunyuan_video_1_5",
+        "HunyuanVideo15Pipeline",
+    ),
+    "HunyuanVideo15ImageToVideoPipeline": (
+        "hunyuan_video",
+        "pipeline_hunyuan_video_1_5_i2v",
+        "HunyuanVideo15I2VPipeline",
+    ),
 }
 
 
@@ -316,6 +326,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "HeliosPipeline": "get_helios_post_process_func",
     "HeliosPyramidPipeline": "get_helios_post_process_func",
     "Flux2Pipeline": "get_flux2_post_process_func",
+    "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
+    "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func",
 }
 
 _DIFFUSION_PRE_PROCESS_FUNCS = {
@@ -332,6 +344,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "OmniGen2Pipeline": "get_omnigen2_pre_process_func",
     "HeliosPipeline": "get_helios_pre_process_func",
     "HeliosPyramidPipeline": "get_helios_pre_process_func",
+    "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_pre_process_func",
 }
 
 

vllm_omni/distributed/kv_transfer/monkey_patch.py

@@ -8,6 +8,7 @@ so the decode side references the correct KV entry.
 
 from __future__ import annotations
 
+import importlib
 import logging
 import sys
 from dataclasses import dataclass
@@ -30,30 +31,23 @@ class PatchedRecvReqMeta:
 
 def _import_mooncake_module():
     """Import MooncakeConnector module, supporting both vLLM >=0.16 and older."""
-    try:
-        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import mooncake_connector
-
-        return mooncake_connector
-    except ImportError:
-        pass
-    try:
-        from vllm.distributed.kv_transfer.kv_connector.v1 import mooncake_connector
-
-        return mooncake_connector
-    except ImportError:
-        return None
+    for mod_path in (
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector",
+    ):
+        try:
+            return importlib.import_module(mod_path)
+        except (ImportError, ModuleNotFoundError):
+            continue
+    return None
 
 
 def _create_patched_mooncake_connector():
     """Return a subclass of MooncakeConnector with remote_request_id support."""
-    try:
-        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
-            MooncakeConnector as _OriginalMooncakeConnector,
-        )
-    except (ImportError, AttributeError):
-        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector import (
-            MooncakeConnector as _OriginalMooncakeConnector,
-        )
+    _mc_mod = _import_mooncake_module()
+    if _mc_mod is None:
+        raise ImportError("Cannot import MooncakeConnector from upstream vLLM")
+    _OriginalMooncakeConnector = _mc_mod.MooncakeConnector
 
     class PatchedMooncakeConnector(_OriginalMooncakeConnector):
         """Fixes request-ID mismatch in PD disaggregation by injecting

vllm_omni/engine/stage_engine_core_client.py

@@ -38,8 +38,12 @@ class StageEngineCoreClient(AsyncMPClient):
         self,
         vllm_config: Any,
         executor_class: type,
-        metadata: StageMetadata,
+        log_stats: bool = False,
         client_addresses: dict[str, str] | None = None,
+        client_count: int = 1,
+        client_index: int = 0,
+        *,
+        metadata: StageMetadata | None = None,
         engine_manager: Any = None,
         coordinator: Any = None,
     ):
@@ -51,17 +55,18 @@ class StageEngineCoreClient(AsyncMPClient):
         and calls super().__init__().
         """
         # -------- Stage metadata (public fields used at runtime) --------
-        self.stage_id = metadata.stage_id
-        self.stage_type = metadata.stage_type
-        self.engine_output_type = metadata.engine_output_type
-        self.is_comprehension = metadata.is_comprehension
-        self.requires_multimodal_data = metadata.requires_multimodal_data
-        self.engine_input_source = metadata.engine_input_source
-        self.final_output = metadata.final_output
-        self.final_output_type = metadata.final_output_type
-        self.default_sampling_params = metadata.default_sampling_params
-        self.custom_process_input_func = metadata.custom_process_input_func
-        self.model_stage = metadata.model_stage
+        if metadata is not None:
+            self.stage_id = metadata.stage_id
+            self.stage_type = metadata.stage_type
+            self.engine_output_type = metadata.engine_output_type
+            self.is_comprehension = metadata.is_comprehension
+            self.requires_multimodal_data = metadata.requires_multimodal_data
+            self.engine_input_source = metadata.engine_input_source
+            self.final_output = metadata.final_output
+            self.final_output_type = metadata.final_output_type
+            self.default_sampling_params = metadata.default_sampling_params
+            self.custom_process_input_func = metadata.custom_process_input_func
+            self.model_stage = metadata.model_stage
 
         self.engine_outputs: Any = None
 
@@ -73,8 +78,10 @@ class StageEngineCoreClient(AsyncMPClient):
             super().__init__(
                 vllm_config,
                 executor_class,
-                log_stats=False,
+                log_stats=log_stats,
                 client_addresses=client_addresses,
+                client_count=client_count,
+                client_index=client_index,
             )
             if engine_manager is not None:
                 self.resources.engine_manager = engine_manager

vllm_omni/entrypoints/async_omni.py

@@ -29,6 +29,7 @@ from vllm_omni.outputs import OmniRequestOutput
 if TYPE_CHECKING:
     from vllm.inputs.preprocess import InputPreprocessor
     from vllm.tokenizers import TokenizerLike
+    from vllm.v1.engine import PauseMode
 
     from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams
 
@@ -66,7 +67,7 @@ class AsyncOmni(EngineClient, OmniBase):
         ...     print(output)
     """
 
-    def __init__(self, model: str, **kwargs: Any) -> None:
+    def __init__(self, *args: Any, model: str = "", **kwargs: Any) -> None:
         OmniBase.__init__(self, model=model, **kwargs)
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
@@ -130,9 +131,13 @@ class AsyncOmni(EngineClient, OmniBase):
     async def generate(
         self,
         prompt: OmniPromptType,
-        request_id: str,
-        sampling_params_list: Sequence[OmniSamplingParams] | None = None,
+        sampling_params: Any = None,
+        request_id: str = "",
         *,
+        prompt_text: str | None = None,
+        lora_request: Any = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        sampling_params_list: Sequence[OmniSamplingParams] | None = None,
         output_modalities: list[str] | None = None,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         """Generate outputs for the given prompt asynchronously.
@@ -226,6 +231,7 @@ class AsyncOmni(EngineClient, OmniBase):
         trace_headers: dict[str, str] | None = None,
         priority: int = 0,
         tokenization_kwargs: dict[str, Any] | None = None,
+        reasoning_ended: bool | None = None,
     ) -> AsyncGenerator[PoolingRequestOutput, None]:
         """EngineClient.encode() stub.
 
@@ -396,6 +402,7 @@ class AsyncOmni(EngineClient, OmniBase):
     async def pause_generation(
         self,
         *,
+        mode: PauseMode = "abort",
         wait_for_inflight_requests: bool = False,
         clear_cache: bool = True,
     ) -> None:
@@ -467,7 +474,7 @@ class AsyncOmni(EngineClient, OmniBase):
         logger.warning("[AsyncOmni] reset_prefix_cache not yet supported with Orchestrator process")
         return True
 
-    async def sleep(self, level: int = 1) -> None:
+    async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         """Sleep all stages.
 
         Best-effort: unsupported stages will emit a TODO result.
@@ -584,7 +591,7 @@ class AsyncOmni(EngineClient, OmniBase):
 
     # ==================== Shutdown ====================
 
-    def shutdown(self) -> None:
+    def shutdown(self, timeout: float | None = None) -> None:
         """Shutdown the engine."""
         if self.final_output_task is not None:
             self.final_output_task.cancel()

vllm_omni/entrypoints/cli/benchmark/main.py

@@ -9,7 +9,7 @@ from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase
 
 if typing.TYPE_CHECKING:
-    from vllm.utils import FlexibleArgumentParser
+    from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
 class OmniBenchmarkSubcommand(CLISubcommand):

vllm_omni/entrypoints/omni_base.py

@@ -53,6 +53,11 @@ def omni_snapshot_download(model_id: str) -> str:
         )
     except huggingface_hub.errors.RepositoryNotFoundError:
         logger.warning("Repository not found for '%s'.", model_id)
+    except PermissionError:
+        logger.warning(
+            "Permission denied when downloading '%s'. Assuming the model is already cached locally.",
+            model_id,
+        )
 
     return model_id
 

vllm_omni/entrypoints/openai/api_server.py

@@ -35,13 +35,6 @@ from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
 from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
-
-# vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
-# Keep a fallback for older/newer upstream layouts during rebase windows.
-try:
-    from vllm.entrypoints.serve.instrumentator.basic import base
-except ModuleNotFoundError:
-    from vllm.entrypoints.openai.basic.api_router import base
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponse,
@@ -68,10 +61,15 @@ from vllm.entrypoints.openai.speech_to_text.serving import (
 )
 from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.pooling.classify.serving import ServingClassification
-from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding
+from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
 from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
 from vllm.entrypoints.pooling.score.serving import ServingScores
 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 OpenAIServingTokenization
 from vllm.entrypoints.utils import (
     load_aware_call,
@@ -329,6 +327,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
             ssl_certfile=args.ssl_certfile,
             ssl_ca_certs=args.ssl_ca_certs,
             ssl_cert_reqs=args.ssl_cert_reqs,
+            ssl_ciphers=args.ssl_ciphers,
             h11_max_incomplete_event_size=args.h11_max_incomplete_event_size,
             h11_max_header_count=args.h11_max_header_count,
             **uvicorn_kwargs,
@@ -572,7 +571,13 @@ async def omni_init_app_state(
                             else vllm_config.model_config
                         )
                         io_processor_plugin = model_config.io_processor_plugin
-                        engine_client.io_processor = get_io_processor(vllm_config, io_processor_plugin)
+                        renderer = getattr(engine_client, "renderer", None)
+                        if renderer is None:
+                            from vllm.renderers import renderer_from_config
+
+                            renderer = renderer_from_config(vllm_config)
+                            engine_client.renderer = renderer
+                        engine_client.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin)
                         logger.info("Initialized io_processor for AsyncOmni")
                 else:
                     logger.warning("Cannot initialize processors: tokenizer is None. OpenAIServingModels may fail.")
@@ -591,6 +596,22 @@ async def omni_init_app_state(
     )
     await state.openai_serving_models.init_static_loras()
 
+    state.openai_serving_render = OpenAIServingRender(
+        model_config=engine_client.model_config,
+        renderer=engine_client.renderer,
+        io_processor=engine_client.io_processor,
+        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,
+        trust_request_chat_template=args.trust_request_chat_template,
+        enable_auto_tools=args.enable_auto_tool_choice,
+        exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
+        tool_parser=args.tool_call_parser,
+        default_chat_template_kwargs=args.default_chat_template_kwargs,
+        log_error_stack=args.log_error_stack,
+    )
+
     state.openai_serving_responses = (
         OpenAIServingResponses(
             engine_client,
@@ -606,7 +627,6 @@ async def omni_init_app_state(
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_force_include_usage=args.enable_force_include_usage,
             enable_log_outputs=args.enable_log_outputs,
-            log_error_stack=args.log_error_stack,
         )
         if "generate" in supported_tasks
         else None
@@ -616,6 +636,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             args.response_role,
+            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,
@@ -630,24 +651,23 @@ async def omni_init_app_state(
             enable_force_include_usage=args.enable_force_include_usage,
             enable_log_outputs=args.enable_log_outputs,
             enable_log_deltas=args.enable_log_deltas,
-            log_error_stack=args.log_error_stack,
         )
         if "generate" in supported_tasks
         else None
     )
     # Warm up chat template processing to avoid first-request latency
     if state.openai_serving_chat is not None:
-        await state.openai_serving_chat.warmup()
+        state.openai_serving_chat.warmup()
 
     state.openai_serving_completion = (
         OpenAIServingCompletion(
             engine_client,
             state.openai_serving_models,
+            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,
             enable_force_include_usage=args.enable_force_include_usage,
-            log_error_stack=args.log_error_stack,
         )
         if "generate" in supported_tasks
         else None
@@ -661,7 +681,6 @@ async def omni_init_app_state(
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
             trust_request_chat_template=args.trust_request_chat_template,
-            log_error_stack=args.log_error_stack,
         )
         if any(task in POOLING_TASKS for task in supported_tasks)
         else None
@@ -674,7 +693,6 @@ async def omni_init_app_state(
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
             trust_request_chat_template=args.trust_request_chat_template,
-            log_error_stack=args.log_error_stack,
         )
         if "embed" in supported_tasks
         else None
@@ -687,7 +705,6 @@ async def omni_init_app_state(
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
             trust_request_chat_template=args.trust_request_chat_template,
-            log_error_stack=args.log_error_stack,
         )
         if "classify" in supported_tasks
         else None
@@ -700,7 +717,7 @@ async def omni_init_app_state(
             score_template=resolved_chat_template,
             log_error_stack=args.log_error_stack,
         )
-        if ("embed" in supported_tasks or "score" in supported_tasks)
+        if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
         else None
     )
     state.openai_serving_tokenization = OpenAIServingTokenization(
@@ -709,15 +726,14 @@ async def omni_init_app_state(
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
+        default_chat_template_kwargs=args.default_chat_template_kwargs,
         trust_request_chat_template=args.trust_request_chat_template,
-        log_error_stack=args.log_error_stack,
     )
     state.openai_serving_transcription = (
         OpenAIServingTranscription(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            log_error_stack=args.log_error_stack,
             enable_force_include_usage=args.enable_force_include_usage,
         )
         if "transcription" in supported_tasks
@@ -728,7 +744,6 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            log_error_stack=args.log_error_stack,
             enable_force_include_usage=args.enable_force_include_usage,
         )
         if "transcription" in supported_tasks
@@ -739,6 +754,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             args.response_role,
+            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,
@@ -758,7 +774,6 @@ async def omni_init_app_state(
             state.openai_serving_models,
             request_logger=request_logger,
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
-            log_error_stack=args.log_error_stack,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_log_outputs=args.enable_log_outputs,
             force_no_detokenize=args.tokens_only,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -256,7 +256,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 )
             else:
                 should_include_tools = tool_dicts is not None
-                conversation, engine_prompts = self._make_request_with_harmony(request, should_include_tools)
+                conversation, engine_prompts = self.openai_serving_render._make_request_with_harmony(
+                    request, should_include_tools
+                )
 
         except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
             logger.exception("Error in preprocessing prompt inputs")

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1433,7 +1433,6 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         multimodal_embeddings: MultiModalEmbeddings | None = None,
         *,
         is_multimodal: torch.Tensor | None = None,
-        handle_oov_mm_token: bool = False,
     ) -> torch.Tensor:
         """Embed input IDs with optional multimodal embeddings."""
         # Get text embeddings

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -773,7 +773,6 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         multimodal_embeddings: MultiModalEmbeddings | None = None,
         *,
         is_multimodal: torch.Tensor | None = None,
-        handle_oov_mm_token: bool = False,
     ) -> torch.Tensor:
         # This is to satisfy the type checker for each overload
         if multimodal_embeddings is None or is_multimodal is None:
@@ -783,7 +782,6 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
             input_ids,
             multimodal_embeddings=multimodal_embeddings,
             is_multimodal=is_multimodal,
-            handle_oov_mm_token=handle_oov_mm_token,
         )
 
     def base_local_forward(

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py

@@ -109,7 +109,6 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         multimodal_embeddings: MultiModalEmbeddings | None = None,
         *,
         is_multimodal: torch.Tensor | None = None,
-        handle_oov_mm_token: bool = False,
     ) -> torch.Tensor:
         # This is to satisfy the type checker for each overload
         if multimodal_embeddings is None or is_multimodal is None:
@@ -119,7 +118,6 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
             input_ids,
             multimodal_embeddings=multimodal_embeddings,
             is_multimodal=is_multimodal,
-            handle_oov_mm_token=handle_oov_mm_token,
         )
 
     def forward(
@@ -232,7 +230,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
             return []
 
         # The result multimodal_embeddings is tuple of tensors, with each
-        # tensor correspoending to a multimodal data item (image or video).
+        # tensor corresponding to a multimodal data item (image or video).
         multimodal_embeddings: tuple[torch.Tensor, ...] = ()
 
         # NOTE: It is important to iterate over the keys in this dictionary

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -102,6 +102,15 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 tokenizer = self.info.get_tokenizer()
                 audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
                 use_audio_in_video = audio_pad_id not in prompt_ids
+            # for mutilmodality cache
+            if any(item is None for item in mm_kwargs["video"]):
+                video_token_id = self.info.get_hf_config().video_token_id
+                audio_token_id = self.info.get_hf_config().audio_token_id
+                video_audio_item_num = sum(id in (video_token_id, audio_token_id) for id in prompt_ids)
+                audio_updates_num = len(mm_prompt_updates.get("audio", []))
+                video_updates_num = len(mm_prompt_updates.get("video", []))
+                if video_audio_item_num != video_updates_num + audio_updates_num:
+                    use_audio_in_video = True
 
         if is_update_applied:
             mm_placeholders = self._find_mm_placeholders(
@@ -586,11 +595,19 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         multimodal_embeddings: MultiModalEmbeddings | None = None,
         *,
         is_multimodal: torch.Tensor | None = None,
-        handle_oov_mm_token: bool = False,
     ) -> torch.Tensor:
         if multimodal_embeddings is None or is_multimodal is None:
             return super().embed_input_ids(input_ids)
 
+        inputs_embeds = self._embed_text_input_ids(
+            input_ids,
+            self.get_language_model().embed_input_ids,
+            is_multimodal=is_multimodal,
+        )
+
+        if len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
         # Check for audio-in-video: interleaved video and audio tokens
         # in the multimodal region. Only use the interleaved path when
         # needed; otherwise fall back to the default parent implementation.
@@ -608,7 +625,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                 input_ids,
                 self.get_language_model().embed_input_ids,
                 is_multimodal=is_multimodal,
-                handle_oov_mm_token=handle_oov_mm_token,
             )
             return merge_interleaved_embeddings(
                 inputs_embeds,
@@ -625,7 +641,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             input_ids,
             multimodal_embeddings=multimodal_embeddings,
             is_multimodal=is_multimodal,
-            handle_oov_mm_token=handle_oov_mm_token,
         )
 
     def forward(

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -341,7 +341,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
         )
 
         # The result multimodal_embeddings is tuple of tensors, with each
-        # tensor correspoending to a multimodal data item (image or video).
+        # tensor corresponding to a multimodal data item (image or video).
         dummy_multimodal_embeddings: tuple[torch.Tensor, ...] = ()
 
         # NOTE: It is important to iterate over the keys in this dictionary

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -29,6 +29,7 @@ from typing import Any, Literal, cast
 import numpy as np
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from packaging.version import Version
 from transformers import PretrainedConfig
 from transformers import __version__ as TRANSFORMERS_VERSION
@@ -139,44 +140,66 @@ class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer):
         rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw)
 
         if isinstance(grid_thw, torch.Tensor):
-            grid_thw_np = grid_thw.cpu().numpy().astype(np.int32)
+            grid_thw_tensor = grid_thw.to(self.device)
         else:
-            grid_thw_np = np.array(grid_thw, dtype=np.int32)
+            grid_thw_tensor = torch.as_tensor(grid_thw, dtype=torch.int32, device=self.device)
 
-        cu_seqlens = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(axis=0, dtype=np.int32)
-        cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
-
-        sequence_lengths = MMEncoderAttention.maybe_compute_sequence_lengths(self.attn_backend, cu_seqlens)
-        if sequence_lengths is not None:
-            sequence_lengths = torch.from_numpy(sequence_lengths).to(self.device, non_blocking=True)
-        max_seqlen = torch.tensor(
-            MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens),
-            dtype=torch.int32,
-            device=self.device,
-        )
-        cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
-            self.attn_backend,
-            cu_seqlens,
-            self.hidden_size,
-            self.tp_size,
-        )
-        cu_seqlens = torch.from_numpy(cu_seqlens).to(self.device, non_blocking=True)
+        try:
+            cu_seqlens = torch.repeat_interleave(
+                grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2],
+                grid_thw_tensor[:, 0],
+            ).cumsum(
+                dim=0,
+                dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32,
+            )
+            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+        except RuntimeError:
+            logger.warning(
+                "torch.repeat_interleave not executable, switching to vectorized searchsorted implementation."
+            )
+            repeat_counts = grid_thw_tensor[:, 0]
+            values = grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2]
+            repeat_cumsum = repeat_counts.cumsum(0)
+            total_items = repeat_cumsum[-1].item()
+            indices = torch.searchsorted(
+                repeat_cumsum,
+                torch.arange(total_items, device=grid_thw_tensor.device),
+                right=True,
+            )
+            cu_seqlens = values[indices].cumsum(
+                dim=0,
+                dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32,
+            )
+            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
 
         hidden_states = hidden_states.unsqueeze(1)
+        rotary_pos_emb_cos = rotary_pos_emb_cos.to(hidden_states.device)
+        rotary_pos_emb_sin = rotary_pos_emb_sin.to(hidden_states.device)
+        max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
+
+        grid_thw_np = grid_thw_tensor.cpu().numpy().astype(np.int32)
+        cu_seqlens_np = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(
+            axis=0, dtype=np.int32
+        )
+        cu_seqlens_np = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens_np])
+        sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
+            self.attn_backend,
+            cu_seqlens_np,
+            self.device,
+        )
 
         hidden_states_list = []
         deepstack_visual_indexes = self.deepstack_visual_indexes
 
         for layer_num, blk in enumerate(self.blocks):
-            hidden_states = hidden_states + blk.attn(
-                blk.norm1(hidden_states),
+            hidden_states = blk(
+                hidden_states,
                 cu_seqlens=cu_seqlens,
                 rotary_pos_emb_cos=rotary_pos_emb_cos,
                 rotary_pos_emb_sin=rotary_pos_emb_sin,
                 max_seqlen=max_seqlen,
                 sequence_lengths=sequence_lengths,
             )
-            hidden_states = hidden_states + blk.mlp(blk.norm2(hidden_states))
 
             if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes:
                 hidden_states_list.append(hidden_states)
@@ -475,6 +498,15 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
                 tokenizer = self.info.get_tokenizer()
                 audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
                 use_audio_in_video = audio_pad_id not in prompt_ids
+            # for mutilmodality cache
+            if any(item is None for item in mm_kwargs["video"]):
+                video_token_id = self.info.get_hf_config().video_token_id
+                audio_token_id = self.info.get_hf_config().audio_token_id
+                video_audio_item_num = sum(id in (video_token_id, audio_token_id) for id in prompt_ids)
+                audio_updates_num = len(mm_prompt_updates.get("audio", []))
+                video_updates_num = len(mm_prompt_updates.get("video", []))
+                if video_audio_item_num != video_updates_num + audio_updates_num:
+                    use_audio_in_video = True
 
         # normal case with `use_audio_in_video=False`
         if is_update_applied:
@@ -487,19 +519,13 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
                 mm_item_counts,
             )
         else:
-            if use_audio_in_video:
-                # When use_audio_in_video=True, audio is extracted from video and embedded
-                # in the video placeholder tokens. We should:
-                # 1. Filter out audio from prompt updates (audio has no separate placeholder)
-                # 2. Apply remaining updates (video, image, etc.)
-                # 3. Derive audio placeholders from video placeholders
+            if use_audio_in_video and "audio" in mm_prompt_updates:
                 filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
                 prompt_ids, mm_placeholders = self._apply_prompt_updates(
                     prompt_ids,
                     filtered_updates,
                 )
-                # Derive audio placeholders from video placeholders
-                mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_item_counts)
+                mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
             else:
                 prompt_ids, mm_placeholders = self._apply_prompt_updates(
                     prompt_ids,
@@ -618,7 +644,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
         def get_replacement_qwen2_use_audio_in_video(item_idx: int):
             nonlocal audio_in_video_item_idx
-            audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx]
+            audio_num_features = audio_output_lengths[audio_in_video_item_idx]
             video_grid_thw = out_mm_data["video_grid_thw"][item_idx]
 
             audio_in_video_item_idx += 1
@@ -664,27 +690,17 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
     def _derive_audio_from_video_placeholders(
         self,
         placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
-        mm_item_counts: Mapping[str, int],
+        mm_prompt_updates: MultiModalPromptUpdates,
     ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
         """
         Helper to derive audio placeholders from video placeholders when
         use_audio_in_video=True.
-
-        In use_audio_in_video mode, audio is extracted from video and embedded
-        within the video placeholder tokens. This function creates audio placeholder
-        info by extracting the audio token positions from video placeholders.
-
-        Args:
-            placeholders: Current placeholders (should contain "video")
-            mm_item_counts: Counts of multimodal items from mm_items.get_all_counts()
         """
         if "video" not in placeholders:
             return placeholders
 
-        # Validate audio and video counts match
-        # In use_audio_in_video mode, audio count comes from mm_items (extracted from video)
         num_videos = len(placeholders["video"])
-        num_audios = mm_item_counts.get("audio", 0)
+        num_audios = len(mm_prompt_updates.get("audio", []))
         if num_audios != num_videos:
             raise ValueError(
                 f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}"
@@ -942,7 +958,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             return []
 
         # The result multimodal_embeddings is tuple of tensors, with each
-        # tensor correspoending to a multimodal data item (image or video).
+        # tensor corresponding to a multimodal data item (image or video).
         multimodal_embeddings: tuple[torch.Tensor, ...] = ()
 
         # NOTE: It is important to iterate over the keys in this dictionary
@@ -966,13 +982,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         multimodal_embeddings: MultiModalEmbeddings | None = None,
         *,
         is_multimodal: torch.Tensor | None = None,
-        handle_oov_mm_token: bool = False,
     ) -> torch.Tensor:
         inputs_embeds = self._embed_text_input_ids(
             input_ids,
             self.language_model.embed_input_ids,
             is_multimodal=is_multimodal,
-            handle_oov_mm_token=handle_oov_mm_token,
         )
 
         if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
@@ -1064,7 +1078,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             input_ids,
             multimodal_embeddings=multimodal_embeddings,
             is_multimodal=is_multimodal,
-            handle_oov_mm_token=handle_oov_mm_token,
         )
 
     def forward(

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -26,6 +26,7 @@ from vllm.v1.outputs import (
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.structured_output.utils import apply_grammar_bitmask
 from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker import mamba_utils
 from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm_ascend.ascend_forward_context import set_ascend_forward_context
@@ -237,6 +238,23 @@ class NPUARModelRunner(OmniNPUModelRunner):
 
                 pad_attn = cudagraph_mode == CUDAGraphMode.FULL
 
+                # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877,
+                # there should be a corresponding 'postprocess_mamba'. However, it is called inside
+                # '_update_states_after_model_execute', which is not overridden in vLLM-Ascend.
+                # We simply utilize the implementation in vLLM.
+                if self.cache_config.mamba_cache_mode == "align":
+                    mamba_utils.preprocess_mamba(
+                        scheduler_output,
+                        self.kv_cache_config,
+                        self.cache_config,
+                        self.mamba_state_idx,
+                        self.input_batch,
+                        self.requests,
+                        self.compilation_config.static_forward_context,
+                        self.model.get_mamba_state_copy_func(),
+                        self._get_mamba_copy_bufs(),
+                    )
+
                 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
                 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
 
@@ -339,7 +357,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 skip_compiled=has_encoder_input,
             ),
             self.maybe_get_kv_connector_output(
-                scheduler_output, clear_metadata=clear_kv_metadata
+                scheduler_output, defer_finalize=not clear_kv_metadata
             ) as kv_connector_output,
         ):
             hidden_states = self._model_forward(
@@ -606,6 +624,34 @@ class NPUARModelRunner(OmniNPUModelRunner):
             hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output
         )
 
+        # Pre-copy multimodal tensors to CPU once (not per-request) to avoid
+        # redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU.
+        mm_cpu: dict[str, object] = {}
+        if isinstance(multimodal_outputs, dict) and multimodal_outputs:
+            for k, v in multimodal_outputs.items():
+                try:
+                    if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
+                        mm_cpu[k] = v.detach().to("cpu").contiguous()
+                    elif isinstance(v, dict):
+                        sub_dict: dict[str, torch.Tensor] = {}
+                        for sk, sv in v.items():
+                            if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
+                                sub_dict[str(sk)] = sv.detach().to("cpu").contiguous()
+                        if sub_dict:
+                            mm_cpu[k] = sub_dict
+                    elif isinstance(v, list):
+                        if len(v) == 0:
+                            continue
+                        cpu_list = []
+                        for elem in v:
+                            if isinstance(elem, torch.Tensor):
+                                cpu_list.append(elem.detach().to("cpu").contiguous())
+                            else:
+                                cpu_list.append(elem)
+                        mm_cpu[k] = cpu_list
+                except Exception as e:
+                    logger.error(f"Error in merge multimodal outputs: {e}")
+
         pooler_output: list[dict[str, object]] = []
         for rid in req_ids_output_copy:
             idx = req_id_to_index_output_copy[rid]
@@ -614,28 +660,26 @@ class NPUARModelRunner(OmniNPUModelRunner):
             end = start + sched
             hidden_slice = hidden_states_cpu[start:end]
             payload: dict[str, object] = {"hidden": hidden_slice}
-            if isinstance(multimodal_outputs, dict) and multimodal_outputs:
+            if mm_cpu:
                 mm_payload: dict[str, object] = {}
-                for k, v in multimodal_outputs.items():
-                    try:
-                        if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
-                            mm_payload[k] = v.detach().to("cpu")[start:end].contiguous()
-                        elif isinstance(v, dict):
-                            sub_dict: dict[str, torch.Tensor] = {}
-                            for sk, sv in v.items():
-                                if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
-                                    sub_dict[str(sk)] = sv.detach().to("cpu")[start:end].contiguous()
-                            if sub_dict:
-                                mm_payload[k] = sub_dict
-                        elif isinstance(v, list):
-                            element = v[0]
-                            if isinstance(element, torch.Tensor):
-                                element = element.detach().to("cpu").contiguous()
-                            mm_payload[k] = element
-                    except Exception as e:
-                        logger.error(f"Error in merge multimodal outputs: {e}")
-                if mm_payload:
-                    payload.update(mm_payload)
+                for k, v in mm_cpu.items():
+                    if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
+                        mm_payload[k] = v[start:end].contiguous()
+                    elif isinstance(v, dict):
+                        mm_payload[k] = {sk: sv[start:end].contiguous() for sk, sv in v.items()}
+                    elif isinstance(v, list):
+                        element = v[idx] if idx < len(v) else v[0]
+                        # Clone tensors to avoid cross-request aliasing
+                        if isinstance(element, torch.Tensor):
+                            element = element.clone()
+                        mm_payload[k] = element
+                    elif isinstance(v, torch.Tensor):
+                        # List-derived tensor payloads are request-invariant; clone to
+                        # avoid accidental cross-request aliasing on downstream mutation.
+                        mm_payload[k] = v.clone()
+                    else:
+                        mm_payload[k] = v
+                payload.update(mm_payload)
             pooler_output.append(payload)
 
         model_runner_output = OmniModelRunnerOutput(

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -20,6 +20,7 @@ from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
 from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker import mamba_utils
 from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
@@ -195,6 +196,23 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
 
                 pad_attn = cudagraph_mode == CUDAGraphMode.FULL
 
+                # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877,
+                # there should be a corresponding 'postprocess_mamba'. However, it is called inside
+                # '_update_states_after_model_execute', which is not overridden in vLLM-Ascend.
+                # We simply utilize the implementation in vLLM.
+                if self.cache_config.mamba_cache_mode == "align":
+                    mamba_utils.preprocess_mamba(
+                        scheduler_output,
+                        self.kv_cache_config,
+                        self.cache_config,
+                        self.mamba_state_idx,
+                        self.input_batch,
+                        self.requests,
+                        self.compilation_config.static_forward_context,
+                        self.model.get_mamba_state_copy_func(),
+                        self._get_mamba_copy_bufs(),
+                    )
+
                 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
                 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
 
@@ -300,7 +318,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                 skip_compiled=has_encoder_input,
             ),
             self.maybe_get_kv_connector_output(
-                scheduler_output, clear_metadata=clear_kv_metadata
+                scheduler_output, defer_finalize=not clear_kv_metadata
             ) as kv_connector_output,
         ):
             #  -------------------------------------- Omni-new -------------------------------------------------
@@ -503,6 +521,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         remove_lora: bool = True,
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
+        profile_seq_lens: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # only support eager mode and piecewise graph now
         assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
@@ -609,11 +628,14 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len
             # in inference. This will be removed once npu_fused_infer_attention_score
             # outperforms _npu_paged_attention on all cases.
-            seq_lens = (
-                SEQ_LEN_WITH_MAX_PA_WORKSPACE
-                if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config)
-                else max_query_len
-            )  # type: ignore[assignment]
+            if profile_seq_lens is not None:
+                seq_lens = profile_seq_lens
+            else:
+                seq_lens = (
+                    SEQ_LEN_WITH_MAX_PA_WORKSPACE
+                    if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config)
+                    else max_query_len
+                )  # type: ignore[assignment]
             self.seq_lens.np[:num_reqs_padded] = seq_lens
             self.seq_lens.np[num_reqs_padded:] = 0
             self.seq_lens.copy_to_gpu()

vllm_omni/platforms/npu/worker/npu_model_runner.py

@@ -35,9 +35,11 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
         enable_sp(self.vllm_config)
         # TODO move this model specific logic to a separate class
         # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
+        self.has_talker_mtp = False
         talker_mtp = getattr(self.model, "talker_mtp", None)
         if talker_mtp is not None:
             self.talker_mtp = talker_mtp  # type: ignore[assignment]
+            self.has_talker_mtp = True
             cudagraph_mode = self.compilation_config.cudagraph_mode
             assert cudagraph_mode is not None
             # Only wrap talker_mtp in CUDAGraphWrapper for Omni models that
@@ -73,6 +75,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
         remove_lora: bool = True,
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
+        profile_seq_lens: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # only support eager mode and piecewise graph now
         assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
@@ -179,11 +182,14 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
             # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len
             # in inference. This will be removed once npu_fused_infer_attention_score
             # outperforms _npu_paged_attention on all cases.
-            seq_lens = (
-                SEQ_LEN_WITH_MAX_PA_WORKSPACE
-                if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config)
-                else max_query_len
-            )  # type: ignore[assignment]
+            if profile_seq_lens is not None:
+                seq_lens = profile_seq_lens
+            else:
+                seq_lens = (
+                    SEQ_LEN_WITH_MAX_PA_WORKSPACE
+                    if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config)
+                    else max_query_len
+                )  # type: ignore[assignment]
             self.seq_lens.np[:num_reqs_padded] = seq_lens
             self.seq_lens.np[num_reqs_padded:] = 0
             self.seq_lens.copy_to_gpu()
@@ -283,7 +289,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
                 model_instance=self.model,
             ):
                 # ---------------------------------------Omni-new----------------------------------------------
-                if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
+                if getattr(self.model, "talker", None) is not None and self.has_talker_mtp:
                     num_tokens_padded_talker_mtp = num_tokens_padded
                     if num_tokens_padded_talker_mtp == self.max_num_tokens:
                         num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
@@ -362,7 +368,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
 
         # Omni-specific: wrap output if needed
         if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
-            model_output = self.model.make_omni_output(model_output, **model_kwargs_extra)
+            model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra)
 
         # Omni-specific: cache model output for later sample_tokens
         self._omni_last_model_output = model_output
@@ -415,12 +421,14 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
             None, self.vllm_config, aclgraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
         ):
             req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step)
-        # update the inputs_embeds and code_predictor_codes
-        code_predictor_codes_cpu = code_predictor_codes.detach().to("cpu").contiguous()
+        # code_predictor_codes stays on GPU here; _update_intermediate_buffer
+        # keeps it device-resident when the key is in gpu_resident_buffer_keys.
+        # D2H is deferred to sample_tokens where hidden_states.to("cpu") already
+        # syncs the stream, avoiding a per-step cudaStreamSynchronize.
         out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes")
         for idx, req_id in enumerate(decode_req_ids):
             req_index = self.input_batch.req_ids.index(req_id)
             start_offset = int(self.query_start_loc.cpu[req_index])
             inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {out_key: code_predictor_codes_cpu[idx : idx + 1]}
+            update_dict = {out_key: code_predictor_codes[idx : idx + 1]}
             self._merge_additional_information_update(req_id, update_dict)

vllm_omni/worker/gpu_ar_model_runner.py

@@ -123,7 +123,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             request_id_resolver=self._resolve_global_request_id,
         )
 
-        if self.vllm_config.model_config.enable_return_routed_experts:
+        if self.routed_experts_initialized:
             capturer = RoutedExpertsCapturer.get_instance()
             if capturer is not None:
                 capturer.clear_buffer()  # noqa
@@ -141,7 +141,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             # Update persistent batch states.
             self._update_states(scheduler_output)
 
-            if has_ec_transfer() and get_ec_transfer().is_producer:
+            if has_ec_transfer() and not get_ec_transfer().is_consumer:
                 with self.maybe_get_ec_connector_output(
                     scheduler_output,
                     encoder_cache=self.encoder_cache,
@@ -276,9 +276,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         # Run the model.
         # Use persistent buffers for CUDA graphs.
-        # When spec decode is enabled, delay clearing connector metadata
-        # until after draft model runs in sample_tokens.
-        clear_kv_metadata = self.speculative_config is None
+        # When spec decode is enabled, defer connector finalization
+        # (wait_for_save + clear metadata) until after draft model runs.
+        defer_kv_connector_finalize = self.speculative_config is not None
         with (
             set_forward_context(
                 attn_metadata,
@@ -292,7 +292,8 @@ class GPUARModelRunner(OmniGPUModelRunner):
             ),
             record_function_or_nullcontext("gpu_model_runner: forward"),
             self.maybe_get_kv_connector_output(
-                scheduler_output, clear_metadata=clear_kv_metadata
+                scheduler_output,
+                defer_finalize=defer_kv_connector_finalize,
             ) as kv_connector_output,
         ):
             model_output = self._model_forward(
@@ -533,11 +534,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
             # tokens on the CPU, so they are run after bookkeeping.
             propose_draft_token_ids(valid_sampled_token_ids)
 
-        # Clear KV connector metadata after draft model runs (if spec decode).
-        # This was deferred from target model forward to allow draft model
-        # to also save its KV cache.
+        # Finalize KV connector (wait_for_save + clear metadata) after
+        # draft model runs. Deferred from target model forward to allow
+        # draft model to also save its KV cache.
         if self.speculative_config is not None:
-            self.clear_kv_connector_metadata()
+            self.finalize_kv_connector()
 
         with record_function_or_nullcontext("gpu_model_runner: eplb"):
             self.eplb_step()
@@ -617,7 +618,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 payload.update(mm_payload)
             pooler_output.append(payload)
         with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
-            if self.model_config.enable_return_routed_experts:
+            if self.routed_experts_initialized:
                 capturer = RoutedExpertsCapturer.get_instance()
                 if capturer is not None:
                     capturer.save_captured_experts(indices=self.slot_mapping)  # noqa

vllm_omni/worker/gpu_ar_worker.py

@@ -48,17 +48,17 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
                 # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
                 self.local_rank += dp_local_rank * tp_pp_world_size
-                assert self.local_rank < torch.cuda.device_count(), (
+                assert self.local_rank < torch.accelerator.device_count(), (
                     f"DP adjusted local rank {self.local_rank} is out of bounds. "
                 )
-                visible_device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
+                visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0
                 assert self.parallel_config.local_world_size <= visible_device_count, (
                     f"local_world_size ({self.parallel_config.local_world_size}) must "
                     f"be less than or equal to the number of visible devices "
                     f"({visible_device_count})."
                 )
             self.device = torch.device(f"cuda:{self.local_rank}")
-            current_platform.set_device(self.device)
+            torch.accelerator.set_device_index(self.device)
 
             current_platform.check_if_supports_dtype(self.model_config.dtype)
 
@@ -79,7 +79,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
             # Now take memory snapshot after NCCL is initialized
             gc.collect()
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
 
             # take current memory snapshot
             self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)

vllm_omni/worker/gpu_generation_model_runner.py

@@ -87,7 +87,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
-        if self.vllm_config.model_config.enable_return_routed_experts:
+        if self.routed_experts_initialized:
             capturer = RoutedExpertsCapturer.get_instance()
             if capturer is not None:
                 capturer.clear_buffer()  # noqa
@@ -108,7 +108,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             if not scheduler_output.total_num_scheduled_tokens:
                 return EMPTY_MODEL_RUNNER_OUTPUT
 
-            if has_ec_transfer() and get_ec_transfer().is_producer:
+            if has_ec_transfer() and not get_ec_transfer().is_consumer:
                 with self.maybe_get_ec_connector_output(
                     scheduler_output,
                     encoder_cache=self.encoder_cache,
@@ -265,9 +265,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
 
         # Run the model.
         # Use persistent buffers for CUDA graphs.
-        # When spec decode is enabled, delay clearing connector metadata
-        # until after draft model runs in sample_tokens.
-        clear_kv_metadata = self.speculative_config is None
+        # When spec decode is enabled, defer connector finalization
+        # (wait_for_save + clear metadata) until after draft model runs.
+        defer_kv_connector_finalize = self.speculative_config is not None
         with (
             set_forward_context(
                 attn_metadata,
@@ -281,7 +281,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             ),
             record_function_or_nullcontext("Forward"),
             self.maybe_get_kv_connector_output(
-                scheduler_output, clear_metadata=clear_kv_metadata
+                scheduler_output,
+                defer_finalize=defer_kv_connector_finalize,
             ) as kv_connector_output,
         ):
             outputs = self._run_generation_model(
@@ -351,9 +352,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         ) = self.execute_model_state
         self.execute_model_state = None
 
-        # Clear KV connector metadata after draft model runs (if spec decode).
+        # Finalize KV connector (wait_for_save + clear metadata) after
+        # draft model runs. Deferred from target model forward.
         if self.speculative_config is not None:
-            self.clear_kv_connector_metadata()
+            self.finalize_kv_connector()
 
         pooler_output: list[object] = []
         if isinstance(multimodal_outputs, torch.Tensor):
@@ -476,6 +478,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         remove_lora: bool = True,
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
+        profile_seq_lens: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -500,6 +503,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             remove_lora: If False, dummy LoRAs are not destroyed after the run
             num_active_loras: Number of distinct active LoRAs to capture for.
                 LoRA is activated when num_active_loras > 0.
+            profile_seq_lens: If provided, use this value for seq_lens instead
+                of max_query_len. Used to profile attention workspace that
+                scales with context length.
         """
         mm_config = self.vllm_config.model_config.multimodal_config
         if mm_config and mm_config.mm_encoder_only:
@@ -621,11 +627,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             # If force_attention is True, we always capture attention.
             # Otherwise, it only happens for cudagraph_runtime_mode=FULL.
             if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-                if create_mixed_batch:
+                if profile_seq_lens is not None:
+                    seq_lens = profile_seq_lens  # type: ignore[assignment]
+                elif create_mixed_batch:
                     # In the mixed batch mode (used for FI warmup), we use
                     # shorter sequence lengths to run faster.
                     # TODO(luka) better system for describing dummy batches
-                    seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
+                    seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]  # type: ignore[assignment]
                 else:
                     seq_lens = max_query_len  # type: ignore[assignment]
                 self.seq_lens.np[:num_reqs] = seq_lens

vllm_omni/worker/gpu_generation_worker.py

@@ -48,17 +48,17 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
                 # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
                 self.local_rank += dp_local_rank * tp_pp_world_size
-                assert self.local_rank < torch.cuda.device_count(), (
+                assert self.local_rank < torch.accelerator.device_count(), (
                     f"DP adjusted local rank {self.local_rank} is out of bounds. "
                 )
-                visible_device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
+                visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0
                 assert self.parallel_config.local_world_size <= visible_device_count, (
                     f"local_world_size ({self.parallel_config.local_world_size}) must "
                     f"be less than or equal to the number of visible devices "
                     f"({visible_device_count})."
                 )
             self.device = torch.device(f"cuda:{self.local_rank}")
-            current_platform.set_device(self.device)
+            torch.accelerator.set_device_index(self.device)
 
             current_platform.check_if_supports_dtype(self.model_config.dtype)
 
@@ -79,7 +79,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
             # Now take memory snapshot after NCCL is initialized
             gc.collect()
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
 
             # take current memory snapshot
             self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)

vllm_omni/worker/gpu_model_runner.py

@@ -257,6 +257,8 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.requests.pop(req_id, None)
             self.model_intermediate_buffer.pop(req_id, None)
             self.num_prompt_logprobs.pop(req_id, None)
+        if hasattr(self, "late_interaction_runner"):
+            self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids)
         # Remove the finished requests from the persistent batch.
         # NOTE(woosuk): There could be an edge case where finished_req_ids and
         # scheduled_req_ids overlap. This happens when a request is aborted and
@@ -266,6 +268,11 @@ class OmniGPUModelRunner(GPUModelRunner):
         for req_id in scheduler_output.finished_req_ids:
             self.input_batch.remove_request(req_id)
 
+        # Zero GPU memory for freshly allocated cache blocks to prevent
+        # stale NaN/data from corrupting attention or SSM computation.
+        if hasattr(scheduler_output, "new_block_ids_to_zero") and scheduler_output.new_block_ids_to_zero:
+            self._zero_block_ids(scheduler_output.new_block_ids_to_zero)
+
         # Free the cached encoder outputs.
         for mm_hash in scheduler_output.free_encoder_mm_hashes:
             self.encoder_cache.pop(mm_hash, None)
@@ -333,6 +340,8 @@ class OmniGPUModelRunner(GPUModelRunner):
                 lora_request=new_req_data.lora_request,
             )
             self.requests[req_id] = req_state
+            if hasattr(self, "late_interaction_runner"):
+                self.late_interaction_runner.register_request(req_id, pooling_params)
 
             # If prompt embeddings are provided, decode and attach to inter_data
             try:
@@ -418,13 +427,13 @@ class OmniGPUModelRunner(GPUModelRunner):
                 # prev_num_draft_len is used in async scheduling mode with
                 # spec decode. it indicates if need to update num_computed_tokens
                 # of the request. for example:
-                # fist step: num_computed_tokens = 0, spec_tokens = [],
+                # first step: num_computed_tokens = 0, spec_tokens = [],
                 # prev_num_draft_len = 0.
                 # second step: num_computed_tokens = 100(prompt length),
                 # spec_tokens = [a,b], prev_num_draft_len = 0.
                 # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d],
                 # prev_num_draft_len = 2.
-                # num_computed_tokens in first step and second step does't contain
+                # num_computed_tokens in first step and second step doesn't contain
                 # the spec tokens length, but in third step it contains the
                 # spec tokens length. we only need to update num_computed_tokens
                 # when prev_num_draft_len > 0.
@@ -551,6 +560,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         remove_lora: bool = True,
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
+        profile_seq_lens: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -575,6 +585,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             remove_lora: If False, dummy LoRAs are not destroyed after the run
             num_active_loras: Number of distinct active LoRAs to capture for.
                 LoRA is activated when num_active_loras > 0.
+            profile_seq_lens: If provided, use this value for seq_lens instead
+                of max_query_len. Used to profile attention workspace that
+                scales with context length.
         """
         mm_config = self.vllm_config.model_config.multimodal_config
         if mm_config and mm_config.mm_encoder_only:
@@ -695,11 +708,13 @@ class OmniGPUModelRunner(GPUModelRunner):
             # If force_attention is True, we always capture attention.
             # Otherwise, it only happens for cudagraph_runtime_mode=FULL.
             if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-                if create_mixed_batch:
+                if profile_seq_lens is not None:
+                    seq_lens = profile_seq_lens  # type: ignore[assignment]
+                elif create_mixed_batch:
                     # In the mixed batch mode (used for FI warmup), we use
                     # shorter sequence lengths to run faster.
                     # TODO(luka) better system for describing dummy batches
-                    seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
+                    seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]  # type: ignore[assignment]
                 else:
                     seq_lens = max_query_len  # type: ignore[assignment]
                 self.seq_lens.np[:num_reqs] = seq_lens