Dev/rebase 0.14.0 (#847)

Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: root <root@hk01dgx039.cm.cluster>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: TangPeng <85704592@qq.com>
Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: Sihyeon Jang <sihyeon.jang@navercorp.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: iwzbi <wzbi@zju.edu.cn>
Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Signed-off-by: Yuhan Liu <30294295+liuyuhanalex@users.noreply.github.com>
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: wangyu31577 <wangyu31577@hundsun.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: John Liu BUAA <liukecheng97@gmail.com>
Signed-off-by: Dinesh G <G.Dinesh@ibm.com>
Signed-off-by: gDINESH13 <dinesh13g@gmail.com>
Signed-off-by: root <root@hk01dgx028.cm.cluster>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: JaredforReal <w13431838023@gmail.com>
Co-authored-by: root <root@hk01dgx039.cm.cluster>
Co-authored-by: JustQJ <37905360+JustQJ@users.noreply.github.com>
Co-authored-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Co-authored-by: Sihyeon Jang <uneedsihyeon@gmail.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: catcat <108673086+iwzbi@users.noreply.github.com>
Co-authored-by: Ziming Huang <hzm414167@alibaba-inc.com>
Co-authored-by: Yuhan Liu <30294295+liuyuhanalex@users.noreply.github.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: wangyu31577 <wangyu31577@hundsun.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: John Liu BUAA <liukecheng97@gmail.com>
Co-authored-by: D!NE$H <67671800+gDINESH13@users.noreply.github.com>
Co-authored-by: root <root@hk01dgx028.cm.cluster>

.buildkite/scripts/simple_test.sh

@@ -47,8 +47,9 @@ fi
 VENV_PYTHON="${VENV_DIR}/bin/python"
 [[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
 
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.12.0
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.14.0
 "${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
+
 "${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
 "${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
 "${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py

README.md

@@ -15,7 +15,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
-
+- [2026/01] We released [0.14.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0).
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 

docker/Dockerfile.ci

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

docker/Dockerfile.rocm

@@ -2,7 +2,7 @@ ARG BASE_IMAGE=rocm/vllm-dev:nightly_main_20251205
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app
-ARG VLLM_VERSION=v0.12.0
+ARG VLLM_VERSION=v0.14.0
 ARG PYTORCH_ROCM_ARCH="gfx942;gfx950"
 
 WORKDIR ${COMMON_WORKDIR}

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.12.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 

docs/design/architecture_overview.md

@@ -192,4 +192,4 @@ curl -sS -X POST http://localhost:8091/v1/chat/completions \
 }
 ```
 
-For more usages, please refer to [examples](../examples/README.md).
+For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).

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

@@ -58,11 +58,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.12.0
+git checkout v0.14.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.12.0/vllm-0.12.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.14.0/vllm-0.14.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -93,7 +93,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.12.0rc1 \
+    vllm/vllm-omni:v0.14.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -68,7 +68,7 @@ docker run -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
-  vllm/vllm-omni-rocm:v0.12.0rc1 \
+  vllm/vllm-omni-rocm:v0.14.0 \
   vllm serve --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -86,7 +86,7 @@ docker run -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
-  vllm/vllm-omni-rocm:v0.12.0rc1 \
+  vllm/vllm-omni-rocm:v0.14.0 \
   bash
 ```
 

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

@@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0
 export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -42,7 +42,7 @@ source ~/.bashrc
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.12.0rc1 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -278,9 +278,9 @@ def get_audio_query(question: str = None, audio_path: str | None = None, samplin
 
 
 query_map = {
-    "mixed_modalities": get_mixed_modalities_query,
+    "use_mixed_modalities": get_mixed_modalities_query,
     "use_audio_in_video": get_use_audio_in_video_query,
-    "multi_audios": get_multi_audios_query,
+    "use_multi_audios": get_multi_audios_query,
     "use_image": get_image_query,
     "use_video": get_video_query,
     "use_audio": get_audio_query,
@@ -434,7 +434,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="mixed_modalities",
+        default="use_mixed_modalities",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

@@ -1,2 +1,2 @@
 python end2end.py --output-wav output_audio \
-                  --query-type mixed_modalities
+                  --query-type use_mixed_modalities

examples/offline_inference/qwen3_omni/end2end.py

@@ -12,6 +12,7 @@ from typing import NamedTuple
 import librosa
 import numpy as np
 import soundfile as sf
+import vllm
 from PIL import Image
 from vllm import SamplingParams
 from vllm.assets.audio import AudioAsset
@@ -225,18 +226,76 @@ def get_multi_audios_query() -> QueryResult:
     )
 
 
+# def get_use_audio_in_video_query(video_path: str | None = None) -> QueryResult:
+# question = (
+#     "Describe the content of the video in details, then convert what the "
+#     "baby say into text."
+# )
+# prompt = (
+#     f"<|im_start|>system\n{default_system}<|im_end|>\n"
+#     "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+#     f"{question}<|im_end|>\n"
+#     f"<|im_start|>assistant\n"
+# )
+# if video_path:
+#     if not os.path.exists(video_path):
+#         raise FileNotFoundError(f"Video file not found: {video_path}")
+#     video_frames = video_to_ndarrays(video_path, num_frames=16)
+# else:
+#     video_frames = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays
+# audio = extract_video_audio(video_path, sampling_rate=16000)
+# return QueryResult(
+#     inputs={
+#         "prompt": prompt,
+#         "multi_modal_data": {
+#             "video": video_frames,
+#             "audio": audio,
+#         },
+#         "mm_processor_kwargs": {
+#             "use_audio_in_video": True,
+#         },
+#     },
+#     limit_mm_per_prompt={"audio": 1, "video": 1},
+# )
+def get_use_audio_in_video_query() -> QueryResult:
+    question = "Describe the content of the video in details, then convert what the baby say into text."
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+    asset = VideoAsset(name="baby_reading", num_frames=16)
+    audio = asset.get_audio(sampling_rate=16000)
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "video": asset.np_ndarrays,
+                "audio": audio,
+            },
+            "mm_processor_kwargs": {
+                "use_audio_in_video": True,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1, "video": 1},
+    )
+
+
 query_map = {
     "text": get_text_query,
     "use_audio": get_audio_query,
     "use_image": get_image_query,
     "use_video": get_video_query,
-    "multi_audios": get_multi_audios_query,
-    "mixed_modalities": get_mixed_modalities_query,
+    "use_multi_audios": get_multi_audios_query,
+    "use_mixed_modalities": get_mixed_modalities_query,
+    "use_audio_in_video": get_use_audio_in_video_query,
 }
 
 
 def main(args):
     model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+    print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)
 
     # Get paths from args
     video_path = getattr(args, "video_path", None)
@@ -259,6 +318,10 @@ def main(args):
             num_frames=getattr(args, "num_frames", 16),
             sampling_rate=getattr(args, "sampling_rate", 16000),
         )
+    elif args.query_type == "multi_audios":
+        query_result = query_func()
+    elif args.query_type == "use_audio_in_video":
+        query_result = query_func()
     else:
         query_result = query_func()
 
@@ -270,7 +333,7 @@ def main(args):
     )
 
     thinker_sampling_params = SamplingParams(
-        temperature=0.4,
+        temperature=0.9,
         top_p=0.9,
         top_k=-1,
         max_tokens=1200,
@@ -331,6 +394,8 @@ def main(args):
     total_requests = len(prompts)
     processed_count = 0
 
+    print(f"query type: {args.query_type}")
+
     for stage_outputs in omni_generator:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
@@ -385,7 +450,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="mixed_modalities",
+        default="use_mixed_modalities",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -305,6 +305,26 @@ def get_multi_audios_query(custom_prompt: str | None = None):
     }
 
 
+def get_use_audio_in_video_query(
+    video_path: str | None = None,
+    audio_path: str | None = None,
+    custom_prompt: str | None = None,
+):
+    question = custom_prompt or (
+        "Describe the content of the video in details, then convert what the baby say into text."
+    )
+    video_url = get_video_url_from_path(video_path)
+    audio_url = get_audio_url_from_path(audio_path)
+    return {
+        "role": "user",
+        "content": [
+            {"type": "video_url", "video_url": {"url": video_url}},
+            {"type": "audio_url", "audio_url": {"url": audio_url}},
+            {"type": "text", "text": question},
+        ],
+    }
+
+
 query_map = {
     "text": get_text_query,
     "use_audio": get_audio_query,
@@ -312,6 +332,7 @@ query_map = {
     "use_video": get_video_query,
     "use_mixed_modalities": get_mixed_modalities_query,
     "use_multi_audios": get_multi_audios_query,
+    "use_audio_in_video": get_use_audio_in_video_query,
 }
 
 
@@ -372,6 +393,12 @@ def run_multimodal_generation(args) -> None:
         prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
     elif args.query_type == "text":
         prompt = query_func(custom_prompt=custom_prompt)
+    elif args.query_type == "use_audio_in_video":
+        prompt = query_func(
+            video_path=video_path,
+            audio_path=audio_path,
+            custom_prompt=custom_prompt,
+        )
     else:
         prompt = query_func()
 

mkdocs.yml

@@ -90,6 +90,7 @@ plugins:
       api_root_uri: "api"
       nav_item_prefix: ""  # No prefix in navigation tree (clean names)
       show_full_namespace: false  # Show only module name, not full path
+      on_implicit_namespace_package: skip  # Skip directories without __init__.py (e.g., assets)
       exclude:
         - "re:vllm_omni\\._.*"  # Internal modules
         - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building

pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "cache-dit==1.2.0",
     "tqdm>=4.66.0",
     "torchsde>=0.2.6",  # Required for Stable Audio scheduler
+    "openai-whisper>=20250625",
     # "vllm==0.12.0",  # TODO: fix the entrypoints overwrite problem
 ]
 

tests/conftest.py

@@ -1,5 +1,11 @@
-import base64
 import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+# Set CPU device for CI environments without GPU
+if "VLLM_TARGET_DEVICE" not in os.environ:
+    os.environ["VLLM_TARGET_DEVICE"] = "cpu"
+
+import base64
 import socket
 import subprocess
 import sys
@@ -13,11 +19,33 @@ import torch
 import whisper
 import yaml
 from vllm.logger import init_logger
-from vllm.utils import get_open_port
+from vllm.utils.network_utils import get_open_port
 
 logger = init_logger(__name__)
 
 
+@pytest.fixture(autouse=True)
+def default_vllm_config():
+    """Set a default VllmConfig for all tests.
+
+    This fixture is auto-used for all tests to ensure that any test
+    that directly instantiates vLLM CustomOps (e.g., RMSNorm, LayerNorm)
+    or model components has the required VllmConfig context.
+
+    This fixture is required for vLLM 0.14.0+ where CustomOp initialization
+    requires a VllmConfig context set via set_current_vllm_config().
+    """
+    from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+    # Use CPU device if no GPU is available (e.g., in CI environments)
+    has_gpu = torch.cuda.is_available() and torch.cuda.device_count() > 0
+    device = "cuda" if has_gpu else "cpu"
+    device_config = DeviceConfig(device=device)
+
+    with set_current_vllm_config(VllmConfig(device_config=device_config)):
+        yield
+
+
 @pytest.fixture(autouse=True)
 def clean_gpu_memory_between_tests():
     if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":

tests/diffusion/attention/test_sequence_parallel.py

@@ -170,6 +170,11 @@ def test_sequence_parallel(
     """Test Ulysses attention by comparing with and without SP enabled."""
     sequence_parallel_size = ulysses_degree * ring_degree
 
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < sequence_parallel_size:
+        pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available")
+
     # Create temporary files to share results between processes
     with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
         baseline_output_file = f.name

tests/diffusion/distributed/test_comm.py

@@ -48,6 +48,11 @@ def test_4d_identity(
     use_sync: bool,
 ):
     """Test that two consecutive all-to-all operations return the original input."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     # Ensure num_heads is divisible by world_size
     if num_heads % world_size != 0:
         pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
@@ -177,6 +182,11 @@ def test_5d_identity(
     use_sync: bool,
 ):
     """Test that two consecutive all-to-all operations return the original input."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     # Ensure num_heads is divisible by world_size
     if num_heads % world_size != 0:
         pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
@@ -305,6 +315,11 @@ def test_ring_p2p(
     head_size: int,
 ):
     """Test Ring P2P communication (send_recv)."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     torch.multiprocessing.spawn(
         _test_ring_p2p_worker,
         args=(world_size, dtype, batch_size, num_heads, head_size),

tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml

@@ -78,6 +78,7 @@ stage_args:
       trust_remote_code: true
       enable_prefix_caching: false
       engine_output_type: audio
+      max_num_batched_tokens: 4069
     engine_input_source: [1]
     final_output: true
     final_output_type: audio

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -45,7 +45,7 @@ stage_args:
        worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
-       enforce_eager: false
+       enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
       #  tensor_parallel_size: 2
@@ -84,6 +84,7 @@ stage_args:
       max_num_batched_tokens: 1000000
       hf_config_name: thinker_config
       load_format: dummy
+      async_scheduling: false
     engine_input_source: [1]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
     final_output: true

tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -84,6 +84,7 @@ stage_args:
       max_num_batched_tokens: 1000000
       hf_config_name: thinker_config
       load_format: dummy
+      async_scheduling: false
     engine_input_source: [1]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
     final_output: true

tests/e2e/offline_inference/test_qwen2_5_omni.py

@@ -35,7 +35,7 @@ test_params = [(model, stage_config) for model in models]
 
 @pytest.mark.core_model
 @pytest.mark.parametrize("test_config", test_params)
-@create_new_process_for_each_test()
+@create_new_process_for_each_test("spawn")
 def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
     """Test processing audio, image, and video together, generating audio output."""
     model, stage_config_path = test_config
@@ -95,7 +95,7 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t
 
 @pytest.mark.core_model
 @pytest.mark.parametrize("test_config", test_params)
-@create_new_process_for_each_test()
+@create_new_process_for_each_test("spawn")
 def test_mixed_modalities_to_text_only(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
     """Test processing audio, image, and video together, generating audio output."""
     model, stage_config_path = test_config

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -5,6 +5,8 @@ E2E offline tests for Omni model with video input and audio output.
 """
 
 import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 from pathlib import Path
 
 import pytest
@@ -14,8 +16,6 @@ from vllm_omni.utils import is_rocm
 
 from .conftest import OmniRunner
 
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -30,8 +30,6 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.distributed.parallel_state import device_count
 from vllm_omni.diffusion.envs import get_device_name
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
 # Use random weights model for testing
 models = ["riverclouds/qwen_image_random"]
 

tests/e2e/offline_inference/test_stable_audio_model.py

@@ -1,4 +1,3 @@
-import os
 import sys
 from pathlib import Path
 
@@ -15,8 +14,6 @@ if str(REPO_ROOT) not in sys.path:
 
 from vllm_omni import Omni
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
 # Use random weights model for CI testing (small, no authentication required)
 models = ["linyueqian/stable_audio_random"]
 

tests/e2e/offline_inference/test_t2v_model.py

@@ -11,8 +11,10 @@ if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
 
 from vllm_omni import Omni
+from vllm_omni.outputs import OmniRequestOutput
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 models = ["Wan-AI/Wan2.2-T2V-A14B-Diffusers"]
 
@@ -45,10 +47,10 @@ def test_video_diffusion_model(model_name: str):
         raise ValueError("No request_output found in OmniRequestOutput")
 
     req_out = first_output.request_output[0]
-    if not isinstance(req_out, dict) or "images" not in req_out:
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 
-    frames = req_out["images"][0]
+    frames = req_out.images[0]
 
     assert frames is not None
     assert hasattr(frames, "shape")

tests/e2e/offline_inference/test_zimage_tensor_parallel.py

@@ -23,7 +23,8 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import is_npu, is_rocm
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
 
@@ -129,8 +130,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
     if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
         pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")
 
-    height = 256
-    width = 256
+    height = 512
+    width = 512
     num_inference_steps = 2
     seed = 42
 

tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml

@@ -77,6 +77,7 @@ stage_args:
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
+      async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: audio  # Final output: audio waveform
       gpu_memory_utilization: 0.1

tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -84,6 +84,7 @@ stage_args:
       max_num_batched_tokens: 1000000
       hf_config_name: thinker_config
       load_format: dummy
+      async_scheduling: false
     engine_input_source: [1]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
     final_output: true

tests/e2e/online_serving/test_async_omni.py

@@ -1,4 +1,6 @@
 import asyncio
+import os
+import sys
 from contextlib import ExitStack
 from pathlib import Path
 
@@ -7,6 +9,8 @@ from vllm.inputs import PromptType
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni, ClientRequestState
 
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
 SEED = 42
 
 stage_config = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_thinker_ci.yaml")
@@ -55,13 +59,19 @@ async def generate(
 @pytest.mark.asyncio
 async def test_abort():
     with ExitStack() as after:
-        engine = AsyncOmni(model=model, stage_configs_path=stage_config)
+        # Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS.
+        engine = AsyncOmni(
+            model=model,
+            stage_configs_path=stage_config,
+            shm_threshold_bytes=sys.maxsize,
+        )
         after.callback(engine.shutdown)
 
-        NUM_REQUESTS = 5
-        NUM_EXPECTED_TOKENS = 100
-        NUM_EXPECTED_TOKENS_LONG = 1000
-        REQUEST_IDS_TO_ABORT = [1, 2, 3]
+        # Keep token counts modest to reduce flakiness on slow test hardware.
+        NUM_REQUESTS = 3
+        NUM_EXPECTED_TOKENS = 64
+        NUM_EXPECTED_TOKENS_LONG = 256
+        REQUEST_IDS_TO_ABORT = [1]
 
         prompt = "Hello my name is Robert and "
 
@@ -74,8 +84,10 @@ async def test_abort():
             tasks.append(asyncio.create_task(generate(engine, request_id, prompt, max_tokens)))
 
         # API server cancels requests when they disconnect.
+        # Explicitly abort in the engine to avoid orphaned requests hanging.
         for idx in REQUEST_IDS_TO_ABORT:
             tasks[idx].cancel()
+            await engine.abort(request_ids[idx])
             await asyncio.sleep(0.1)
 
         # Confirm the other requests are okay.
@@ -83,10 +95,10 @@ async def test_abort():
             # Confirm that it was actually canceled.
             if idx in REQUEST_IDS_TO_ABORT:
                 with pytest.raises((asyncio.CancelledError, GeneratorExit)):
-                    await task
+                    await asyncio.wait_for(task, timeout=60)
             else:
                 # Otherwise, make sure the request was not impacted.
-                num_generated_tokens, request_id = await task
+                num_generated_tokens, request_id = await asyncio.wait_for(task, timeout=180)
                 expected_tokens = NUM_EXPECTED_TOKENS
                 assert num_generated_tokens == expected_tokens, (
                     f"{request_id} generated {num_generated_tokens} but expected {expected_tokens}"
@@ -116,11 +128,16 @@ async def test_build_and_log_summary(monkeypatch):
     monkeypatch.setattr("vllm_omni.entrypoints.client_request_state.ClientRequestState", MockCRS)
 
     with ExitStack() as after:
-        engine = AsyncOmni(model=model, stage_configs_path=stage_config)
+        # Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS.
+        engine = AsyncOmni(
+            model=model,
+            stage_configs_path=stage_config,
+            shm_threshold_bytes=sys.maxsize,
+        )
         after.callback(engine.shutdown)
         prompt = "Hello my name is Robert and "
-        NUM_EXPECTED_TOKENS = 100
-        NUM_REQUESTS = 5
+        NUM_EXPECTED_TOKENS = 64
+        NUM_REQUESTS = 3
         request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
 
         # Create concurrent requests.

tests/e2e/online_serving/test_i2i_multi_image_input.py

@@ -17,7 +17,7 @@ import openai
 import pytest
 from PIL import Image
 from vllm.assets.image import ImageAsset
-from vllm.utils import get_open_port
+from vllm.utils.network_utils import get_open_port
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 

tests/e2e/online_serving/test_qwen3_omni.py

@@ -4,10 +4,13 @@
 E2E Online tests for Qwen3-Omni model with video input and audio output.
 """
 
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
 import base64
 import concurrent.futures
 import ctypes
-import os
 import signal
 import socket
 import subprocess
@@ -18,12 +21,10 @@ from pathlib import Path
 import openai
 import pytest
 from vllm.assets.video import VideoAsset
-from vllm.utils import get_open_port
+from vllm.utils.network_utils import get_open_port
 
 from vllm_omni.utils import is_rocm
 
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -6,6 +6,8 @@ E2E Online tests for Qwen3-Omni model.
 
 import concurrent.futures
 import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 import time
 from pathlib import Path
 
@@ -20,8 +22,6 @@ from tests.conftest import (
     modify_stage_config,
 )
 
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 # CI stage config for 2*H100-80G GPUs

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -81,6 +81,7 @@ stage_args:
     distributed_executor_backend: "mp"
     max_num_batched_tokens: 1000000
     hf_config_name: thinker_config
+    async_scheduling: false
   engine_input_source: [1]
   custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
   final_output: true

tests/worker/test_omni_gpu_model_runner.py

@@ -69,7 +69,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
 
     # Provide a minimal implementation that returns the expected 4-tuple.
     def _determine_batch_execution_and_padding(**kwargs):
-        return None, object(), None, None
+        return None, object(), None, None, None
 
     runner._determine_batch_execution_and_padding = _determine_batch_execution_and_padding
 

vllm_omni/assets/video.py

@@ -0,0 +1,16 @@
+import librosa
+import numpy as np
+from vllm.assets.video import VideoAsset
+
+
+def extract_video_audio(path: str = None, sampling_rate: int = 16000) -> np.ndarray:
+    """This function extracts the audio from a video file path and returns the audio as a numpy array.
+    Args:
+        path: The path to the video file.
+    Returns:
+        The audio as a numpy array.
+    """
+    if not path:
+        path = VideoAsset(name="baby_reading").video_path
+    audio_signal, sr = librosa.load(path, sr=sampling_rate)
+    return audio_signal

vllm_omni/config/model.py

@@ -1,20 +1,12 @@
 import warnings
-from importlib.util import find_spec
 from typing import Any
 
 import torch
-import vllm.envs as envs
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
-from vllm.attention.backends.registry import AttentionBackendEnum
 from vllm.config import ModelConfig, config
 from vllm.config.model import (
     _RUNNER_CONVERTS,
-    _RUNNER_TASKS,
-    ConvertOption,
-    ConvertType,
-    RunnerOption,
-    TaskOption,
     _get_and_verify_dtype,
     get_served_model_name,
 )
@@ -29,12 +21,11 @@ from vllm.transformers_utils.config import (
     get_pooling_config,
 )
 from vllm.transformers_utils.gguf_utils import (
+    is_gguf,
     maybe_patch_hf_config_from_gguf,
 )
-from vllm.transformers_utils.utils import (
-    is_gguf,
-    maybe_model_redirect,
-)
+from vllm.transformers_utils.utils import maybe_model_redirect
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 import vllm_omni.model_executor.models as me_models
 
@@ -146,14 +137,6 @@ class OmniModelConfig(ModelConfig):
 
         self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
 
-        if (backend := envs.VLLM_ATTENTION_BACKEND) and backend == "FLASHINFER" and find_spec("flashinfer") is None:
-            raise ValueError(
-                "VLLM_ATTENTION_BACKEND is set to FLASHINFER, but flashinfer "
-                "module was not found. See "
-                "https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile "  # noqa: E501
-                "for instructions on how to install it."
-            )
-
         if self.override_attention_dtype is not None and not current_platform.is_rocm():
             warnings.warn(
                 "override-attention-dtype is set but not using ROCm platform",
@@ -183,111 +166,16 @@ class OmniModelConfig(ModelConfig):
         self.hf_text_config = self.draw_hf_text_config()
         self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None)
         self.encoder_config = self._get_encoder_config()
-        # Try to load image processor config, but allow it to fail for stages that don't need it
-        try:
-            self.hf_image_processor_config = get_hf_image_processor_config(
-                self.model, hf_token=self.hf_token, revision=self.revision
-            )
-        except (OSError, ValueError, IndexError) as e:
-            # Some stages (e.g., code2wav, talker) don't need image processor
-            # Log warning but allow initialization to continue
-            logger.warning(
-                f"Failed to load image processor config for model '{self.model}': {e}. "
-                "This is expected for stages that don't require image processing."
-            )
-            self.hf_image_processor_config = None
+        self.hf_image_processor_config = get_hf_image_processor_config(
+            self.model, hf_token=self.hf_token, revision=self.revision
+        )
+        self.model_arch_config = self.get_model_arch_config()
 
         architectures = self.architectures
         registry = self.registry
         is_generative_model = registry.is_text_generation_model(architectures, self)
         is_pooling_model = registry.is_pooling_model(architectures, self)
 
-        def _task_to_convert(task: TaskOption) -> ConvertType:
-            if task == "embedding" or task == "embed":
-                return "embed"
-            if task == "classify":
-                return "classify"
-            if task == "reward":
-                return "reward"
-            if task == "score":
-                new_task = self._get_default_pooling_task(architectures)
-                return "classify" if new_task == "classify" else "embed"
-
-            return "none"
-
-        if self.task is not None:
-            runner: RunnerOption = "auto"
-            convert: ConvertOption = "auto"
-            msg_prefix = (
-                "The 'task' option has been deprecated and will be removed in v0.13.0 or v1.0, whichever comes first."
-            )
-            msg_hint = "Please remove this option."
-
-            is_generative_task = self.task in _RUNNER_TASKS["generate"]
-            is_pooling_task = self.task in _RUNNER_TASKS["pooling"]
-
-            if is_generative_model and is_pooling_model:
-                if is_generative_task:
-                    runner = "generate"
-                    convert = "auto"
-                    msg_hint = (
-                        "Please replace this option with `--runner "
-                        "generate` to continue using this model "
-                        "as a generative model."
-                    )
-                elif is_pooling_task:
-                    runner = "pooling"
-                    convert = "auto"
-                    msg_hint = (
-                        "Please replace this option with `--runner "
-                        "pooling` to continue using this model "
-                        "as a pooling model."
-                    )
-                else:  # task == "auto"
-                    pass
-            elif is_generative_model or is_pooling_model:
-                if is_generative_task:
-                    runner = "generate"
-                    convert = "auto"
-                    msg_hint = "Please remove this option"
-                elif is_pooling_task:
-                    runner = "pooling"
-                    convert = _task_to_convert(self.task)
-                    msg_hint = (
-                        "Please replace this option with `--convert "
-                        f"{convert}` to continue using this model "
-                        "as a pooling model."
-                    )
-                else:  # task == "auto"
-                    pass
-            else:
-                # Neither generative nor pooling model - try to convert if possible
-                if is_pooling_task:
-                    runner = "pooling"
-                    convert = _task_to_convert(self.task)
-                    msg_hint = (
-                        "Please replace this option with `--runner pooling "
-                        f"--convert {convert}` to continue using this model "
-                        "as a pooling model."
-                    )
-                else:
-                    debug_info = {
-                        "architectures": architectures,
-                        "is_generative_model": is_generative_model,
-                        "is_pooling_model": is_pooling_model,
-                    }
-                    raise AssertionError(
-                        "The model should be a generative or "
-                        "pooling model when task is set to "
-                        f"{self.task!r}. Found: {debug_info}"
-                    )
-
-            self.runner = runner
-            self.convert = convert
-
-            msg = f"{msg_prefix} {msg_hint}"
-            warnings.warn(msg, DeprecationWarning, stacklevel=2)
-
         self.runner_type = self._get_runner_type(architectures, self.runner)
         self.convert_type = self._get_convert_type(architectures, self.runner_type, self.convert)
 
@@ -325,9 +213,12 @@ class OmniModelConfig(ModelConfig):
                     if getattr(self.pooler_config, k) is None:
                         setattr(self.pooler_config, k, v)
 
-            default_pooling_type = self._model_info.default_pooling_type
-            if self.pooler_config.pooling_type is None:
-                self.pooler_config.pooling_type = default_pooling_type
+            default_seq_pooling_type = self._model_info.default_seq_pooling_type
+            if self.pooler_config.seq_pooling_type is None:
+                self.pooler_config.seq_pooling_type = default_seq_pooling_type
+            default_tok_pooling_type = self._model_info.default_tok_pooling_type
+            if self.pooler_config.tok_pooling_type is None:
+                self.pooler_config.tok_pooling_type = default_tok_pooling_type
 
         self.dtype: torch.dtype = _get_and_verify_dtype(
             self.model,
@@ -339,6 +230,11 @@ class OmniModelConfig(ModelConfig):
 
         self.original_max_model_len = self.max_model_len
         self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
+
+        if self.is_encoder_decoder:
+            self.mm_processor_cache_gb = 0
+            logger.info("Encoder-decoder model detected, disabling mm processor cache.")
+
         # Init multimodal config if needed
         if self._model_info.supports_multimodal:
             if mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data:
@@ -382,7 +278,7 @@ class OmniModelConfig(ModelConfig):
 
         # Avoid running try_verify_and_update_config multiple times
         self.config_updated = False
-
+        self._try_verify_and_update_model_config()
         self._verify_quantization()
         self._verify_cuda_graph()
         self._verify_bnb_config()

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 from collections import defaultdict
 from time import time
 
+import numpy as np
+from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
@@ -10,10 +12,13 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import check_stop, remove_all
 from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
+logger = init_logger(__name__)
+
 
 class OmniARScheduler(VLLMScheduler):
     """
@@ -73,6 +78,11 @@ class OmniARScheduler(VLLMScheduler):
         pooler_outputs = model_runner_output.pooler_output
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
+        cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
+
+        perf_stats: PerfStats | None = None
+        if self.perf_metrics and self.perf_metrics.is_enabled():
+            perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
@@ -131,11 +141,14 @@ class OmniARScheduler(VLLMScheduler):
                     spec_decoding_stats,
                     num_draft_tokens=num_draft_tokens,
                     num_accepted_tokens=num_accepted,
+                    num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
+                    request_id=req_id,
                 )
 
             stopped = False
             new_logprobs = None
             new_token_ids = generated_token_ids
+            pooler_output = pooler_outputs[req_index] if pooler_outputs else None
             kv_transfer_params = None
             status_before_stop = request.status
 
@@ -143,14 +156,31 @@ class OmniARScheduler(VLLMScheduler):
             if new_token_ids:
                 new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
 
-            # Stop checking for pooler models.
-            pooler_output = None
-            if pooler_outputs:
-                pooler_output = pooler_outputs[req_index]
+            if pooler_output:
+                # Note: For multimodal outputs, we skip intermediate stop checks.
                 if request.output_token_ids:
-                    stopped = check_stop(request, self.max_model_len, pooler_output)
-
+                    stopped = check_stop(request, self.max_model_len)
+            routed_experts = None
             if stopped:
+                if self.vllm_config.model_config.enable_return_routed_experts:
+                    kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
+                    block_ids = kv_blocks.get_block_ids()[0]
+                    num_tokens = request.num_tokens - 1
+
+                    # compute slot mapping
+                    block_ids_array = np.array(block_ids, dtype=np.int32)
+                    num_blocks = len(block_ids)
+                    block_size = self.block_size
+
+                    # generate block offsets
+                    block_offsets = np.arange(0, block_size)
+
+                    # compute slot mapping: slot = block_id * block_size + offset
+                    slot_mapping = (
+                        block_offsets.reshape((1, block_size)) + block_ids_array.reshape((num_blocks, 1)) * block_size
+                    ).flatten()[:num_tokens]
+
+                    routed_experts = self.routed_experts_reader.get_routed_experts(indices=slot_mapping)
                 kv_transfer_params = self._free_request(request)
                 if status_before_stop == RequestStatus.RUNNING:
                     stopped_running_reqs.add(request)
@@ -165,7 +195,13 @@ class OmniARScheduler(VLLMScheduler):
                 struct_output_request = request.structured_output_request
                 assert struct_output_request is not None
                 assert struct_output_request.grammar is not None
-                struct_output_request.grammar.accept_tokens(req_id, new_token_ids)
+                ok = struct_output_request.grammar.accept_tokens(req_id, new_token_ids)
+                if not ok:
+                    logger.warning(
+                        "Unexpected: grammar rejected tokens %s for request %s.",
+                        new_token_ids,
+                        req_id,
+                    )
 
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
                 request.num_nans_in_logits = num_nans_in_logits[req_id]
@@ -187,6 +223,7 @@ class OmniARScheduler(VLLMScheduler):
                         kv_transfer_params=kv_transfer_params,
                         trace_headers=request.trace_headers,
                         num_cached_tokens=request.num_cached_tokens,
+                        routed_experts=routed_experts,
                         num_nans_in_logits=request.num_nans_in_logits,
                     )
                 )
@@ -201,6 +238,20 @@ class OmniARScheduler(VLLMScheduler):
             # This is a rare case and unlikely to impact performance.
             self.waiting.remove_requests(stopped_preempted_reqs)
 
+        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
+            requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
+            self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
+            for request in requests:
+                outputs[request.client_index].append(
+                    EngineCoreOutput(
+                        request_id=request.request_id,
+                        new_token_ids=[],
+                        finish_reason=request.get_finished_reason(),
+                        events=request.take_events(),
+                        trace_headers=request.trace_headers,
+                        num_cached_tokens=request.num_cached_tokens,
+                    )
+                )
         # KV Connector: update state for finished KV Transfers.
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
@@ -238,7 +289,7 @@ class OmniARScheduler(VLLMScheduler):
                     engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
             finished_req_ids.clear()
 
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
+        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None:
             # Return stats to only one of the front-ends.
             if (eco := next(iter(engine_core_outputs.values()), None)) is None:
                 # We must return the stats even if there are no request

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -90,7 +90,7 @@ class OmniGenerationScheduler(VLLMScheduler):
             any_request = self.running[0]
             num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request.request_id)
 
-        # Assemble SchedulerOutput (align with v0.12.0)
+        # Assemble SchedulerOutput (align with v0.14.0)
         if self.use_v2_model_runner:
             # No resumed reqs in fast path; pass prefill_token_ids for new reqs.
             new_reqs_data = [
@@ -129,7 +129,7 @@ class OmniGenerationScheduler(VLLMScheduler):
             preempted_req_ids=set(),
         )
 
-        # Record the request ids scheduled in this step (v0.12.0 behavior).
+        # 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())
 
@@ -176,7 +176,7 @@ class OmniGenerationScheduler(VLLMScheduler):
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
         kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
-        # Merge connector-side stats (align with v0.12.0)
+        # Merge connector-side stats (align with v0.14.0)
         if kv_connector_stats and self.connector:
             kv_stats = self.connector.get_kv_connector_stats()
             if kv_stats:
@@ -294,7 +294,7 @@ class OmniGenerationScheduler(VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Collect and publish KV cache events (align with v0.12.0)
+        # Collect and publish KV cache events (align with v0.14.0)
         events = self.kv_cache_manager.take_events()
         if self.connector is not None:
             connector_events = self.connector.take_events()

vllm_omni/diffusion/forward_context.py

@@ -86,5 +86,14 @@ def set_forward_context(
         attn_metadata=attn_metadata,
         split_text_embed_in_sp=split_text_embed_in_sp,
     )
+    # vLLM CustomOp dispatch (e.g. QKVParallelLinear) requires a global
+    # vLLM config set via set_current_vllm_config().
     with override_forward_context(forward_context):
-        yield
+        if vllm_config is None:
+            yield
+        else:
+            # Local import to avoid importing vllm.config.vllm at module import time.
+            from vllm.config.vllm import set_current_vllm_config
+
+            with set_current_vllm_config(vllm_config):
+                yield

vllm_omni/diffusion/models/glm_image/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM Image diffusion model components."""
+
+from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+    GlmImageKVCache,
+    GlmImageTransformer2DModel,
+)
+from vllm_omni.diffusion.models.glm_image.pipeline_glm_image import (
+    GlmImagePipeline,
+    get_glm_image_post_process_func,
+    get_glm_image_pre_process_func,
+)
+
+__all__ = [
+    "GlmImageKVCache",
+    "GlmImagePipeline",
+    "GlmImageTransformer2DModel",
+    "get_glm_image_post_process_func",
+    "get_glm_image_pre_process_func",
+]

vllm_omni/diffusion/models/glm_image/glm_image_transformer.py

@@ -0,0 +1,792 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+from enum import Enum
+from typing import Any
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import FeedForward
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.transformers.transformer_glm_image import GlmImageCombinedTimestepSizeEmbeddings
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import QKVParallelLinear
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.base import CachedTransformer
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+logger = init_logger(__name__)
+
+
+class GlmImageImageProjector(nn.Module):
+    """Projects latent image patches to transformer hidden dimension."""
+
+    def __init__(
+        self,
+        in_channels: int = 16,
+        hidden_size: int = 2560,
+        patch_size: int = 2,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, channel, height, width = hidden_states.shape
+        post_patch_height = height // self.patch_size
+        post_patch_width = width // self.patch_size
+
+        # Reshape: [B, C, H, W] -> [B, H', W', C*p*p] -> [B, H'*W', C*p*p]
+        hidden_states = hidden_states.reshape(
+            batch_size, channel, post_patch_height, self.patch_size, post_patch_width, self.patch_size
+        )
+        hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
+        hidden_states = self.proj(hidden_states)
+        return hidden_states
+
+
+class GlmImageRotaryPosEmbed(nn.Module):
+    """Rotary positional embedding for 2D image patches."""
+
+    def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.theta = theta
+
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, num_channels, height, width = hidden_states.shape
+        height, width = height // self.patch_size, width // self.patch_size
+
+        dim_h, dim_w = self.dim // 2, self.dim // 2
+        h_inv_freq = 1.0 / (
+            self.theta ** (torch.arange(0, dim_h, 2, dtype=torch.float32)[: (dim_h // 2)].float() / dim_h)
+        )
+        w_inv_freq = 1.0 / (
+            self.theta ** (torch.arange(0, dim_w, 2, dtype=torch.float32)[: (dim_w // 2)].float() / dim_w)
+        )
+        h_seq = torch.arange(height, device=hidden_states.device)
+        w_seq = torch.arange(width, device=hidden_states.device)
+        h_inv_freq = h_inv_freq.to(hidden_states.device)
+        w_inv_freq = w_inv_freq.to(hidden_states.device)
+
+        freqs_h = torch.outer(h_seq, h_inv_freq)
+        freqs_w = torch.outer(w_seq, w_inv_freq)
+
+        # Create position matrices: [height, 1, dim//4] and [1, width, dim//4]
+        freqs_h = freqs_h.unsqueeze(1).expand(height, width, -1)
+        freqs_w = freqs_w.unsqueeze(0).expand(height, width, -1)
+
+        # Concatenate: [height, width, dim//2] -> [height, width, dim]
+        freqs = torch.cat([freqs_h, freqs_w], dim=-1)
+        freqs = torch.cat([freqs, freqs], dim=-1)
+        freqs = freqs.reshape(height * width, -1)
+        return (freqs.cos(), freqs.sin())
+
+
+class GlmImageAdaLayerNormZero(nn.Module):
+    """Adaptive LayerNorm with zero initialization for both image and text streams."""
+
+    def __init__(self, embedding_dim: int, dim: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.linear = nn.Linear(embedding_dim, 12 * dim, bias=True)
+
+    def forward(
+        self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
+    ) -> tuple[torch.Tensor, ...]:
+        dtype = hidden_states.dtype
+        norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
+        norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype)
+
+        emb = self.linear(temb)
+        (
+            shift_msa,
+            c_shift_msa,
+            scale_msa,
+            c_scale_msa,
+            gate_msa,
+            c_gate_msa,
+            shift_mlp,
+            c_shift_mlp,
+            scale_mlp,
+            c_scale_mlp,
+            gate_mlp,
+            c_gate_mlp,
+        ) = emb.chunk(12, dim=1)
+
+        hidden_states = norm_hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+        encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_msa.unsqueeze(1)) + c_shift_msa.unsqueeze(1)
+
+        return (
+            hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        )
+
+
+class GlmImageAdaLayerNormContinuous(nn.Module):
+    """Final AdaLN for output projection (no activation before Linear)."""
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        elementwise_affine: bool = True,
+        eps: float = 1e-5,
+        bias: bool = True,
+    ):
+        super().__init__()
+        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
+        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
+
+    def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
+        # NO SiLU here
+        emb = self.linear(conditioning_embedding.to(x.dtype))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+class KVCacheMode(Enum):
+    """Mode for KV cache operations.
+
+    - WRITE: Store the K/V tensors from condition images
+    - READ: Concatenate cached K/V with current K/V
+    - SKIP: Do not use cache (pass-through)
+    """
+
+    WRITE = "write"
+    READ = "read"
+    SKIP = "skip"
+
+
+class GlmImageLayerKVCache:
+    """KV cache for a single attention layer.
+
+    Stores key and value tensors for image editing. The cache accumulates
+    KV pairs during write mode and provides them during read mode.
+
+    Shape convention (vllm-omni):
+        key/value: [batch_size, seq_length, num_heads, head_dim]
+    """
+
+    def __init__(self):
+        self.k_cache: torch.Tensor | None = None
+        self.v_cache: torch.Tensor | None = None
+
+    def store(self, key: torch.Tensor, value: torch.Tensor) -> None:
+        """Store or accumulate KV tensors.
+
+        If cache is empty, stores the tensors directly.
+        If cache is not empty, concatenates new tensors along seq_length dim.
+
+        Args:
+            key: Key tensor of shape [B, S, H, D]
+            value: Value tensor of shape [B, S, H, D]
+        """
+        if self.k_cache is None:
+            self.k_cache = key
+            self.v_cache = value
+        else:
+            # Concatenate along sequence dimension (dim=1 for [B, S, H, D])
+            self.k_cache = torch.cat([self.k_cache, key], dim=1)
+            self.v_cache = torch.cat([self.v_cache, value], dim=1)
+
+    def get(self) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        """Get cached KV tensors.
+
+        Returns:
+            Tuple of (k_cache, v_cache), both may be None if cache is empty.
+        """
+        return self.k_cache, self.v_cache
+
+    def clear(self) -> None:
+        """Clear the cache."""
+        self.k_cache = None
+        self.v_cache = None
+
+    @property
+    def is_empty(self) -> bool:
+        """Check if cache is empty."""
+        return self.k_cache is None
+
+    def __repr__(self) -> str:
+        if self.is_empty:
+            return "GlmImageLayerKVCache(empty)"
+        return f"GlmImageLayerKVCache(k_shape={self.k_cache.shape}, v_shape={self.v_cache.shape})"
+
+
+class GlmImageKVCache:
+    """Container for all layers' KV caches.
+
+    Manages KV cache for all transformer layers in GLM-Image model.
+    Provides a unified interface for setting mode and clearing cache.
+
+    Args:
+        num_layers: Number of transformer layers in the model.
+
+    Example:
+        kv_cache = GlmImageKVCache(num_layers=28)
+        kv_cache.set_mode(KVCacheMode.WRITE)
+        # ... process condition image ...
+        kv_cache.set_mode(KVCacheMode.READ)
+        # ... process target image ...
+        kv_cache.clear()
+    """
+
+    def __init__(self, num_layers: int):
+        self.num_layers = num_layers
+        self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)]
+        self._mode: KVCacheMode | None = None
+
+    def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
+        """Get cache for a specific layer.
+
+        Args:
+            layer_idx: Index of the layer (0-indexed).
+
+        Returns:
+            GlmImageLayerKVCache for the specified layer.
+
+        Raises:
+            IndexError: If layer_idx is out of range.
+        """
+        if layer_idx < 0 or layer_idx >= self.num_layers:
+            raise IndexError(f"Layer index {layer_idx} out of range [0, {self.num_layers})")
+        return self.caches[layer_idx]
+
+    def __len__(self) -> int:
+        """Return number of layers."""
+        return self.num_layers
+
+    @property
+    def mode(self) -> KVCacheMode | None:
+        """Get current cache mode."""
+        return self._mode
+
+    def set_mode(self, mode: KVCacheMode | str | None) -> None:
+        """Set cache mode for all layers.
+
+        Args:
+            mode: Cache mode (WRITE, READ, SKIP) or string ("write", "read", "skip").
+                  Use None to disable cache operations.
+
+        Raises:
+            ValueError: If mode is an invalid string.
+        """
+        if mode is None:
+            self._mode = None
+        elif isinstance(mode, str):
+            try:
+                self._mode = KVCacheMode(mode.lower())
+            except ValueError:
+                raise ValueError(f"Invalid mode: '{mode}', must be one of 'write', 'read', 'skip'")
+        else:
+            self._mode = mode
+
+    def clear(self) -> None:
+        """Clear cache for all layers and reset mode."""
+        for cache in self.caches:
+            cache.clear()
+        self._mode = None
+
+    @property
+    def is_empty(self) -> bool:
+        """Check if all layer caches are empty."""
+        return all(cache.is_empty for cache in self.caches)
+
+    def __repr__(self) -> str:
+        mode_str = self._mode.value if self._mode else "None"
+        return f"GlmImageKVCache(num_layers={self.num_layers}, mode={mode_str}, is_empty={self.is_empty})"
+
+
+class GlmImageAttention(nn.Module):
+    """
+    Joint attention for GLM-Image model using vllm-omni's optimized attention.
+
+    This combines text and image streams for joint attention computation.
+    Supports KV caching for image editing workflows via external cache.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        head_dim: int,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.inner_dim = num_heads * head_dim
+
+        # QKV projection (fused for efficiency)
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=head_dim,
+            total_num_heads=num_heads,
+            disable_tp=True,
+            bias=True,
+        )
+
+        # QK normalization (LayerNorm, not RMSNorm for GLM-Image)
+        self.norm_q = nn.LayerNorm(head_dim, elementwise_affine=False, eps=eps)
+        self.norm_k = nn.LayerNorm(head_dim, elementwise_affine=False, eps=eps)
+
+        # Output projection
+        self.to_out = nn.Sequential(
+            nn.Linear(self.inner_dim, dim, bias=out_bias),
+            nn.Dropout(0.0),
+        )
+
+        # RoPE and attention
+        self.rope = RotaryEmbedding(is_neox_style=False)
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=head_dim,
+            softmax_scale=1.0 / (head_dim**0.5),
+            causal=False,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        kv_cache: GlmImageLayerKVCache | None = None,
+        kv_cache_mode: KVCacheMode | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass for joint attention.
+
+        Args:
+            hidden_states: Image hidden states [B, img_seq_len, D]
+            encoder_hidden_states: Text hidden states [B, text_seq_len, D]
+            image_rotary_emb: Tuple of (cos, sin) for RoPE
+            attention_mask: Optional attention mask for text tokens
+            kv_cache: Optional layer KV cache for image editing
+            kv_cache_mode: Cache mode (WRITE, READ, SKIP)
+
+        Returns:
+            Tuple of (image_hidden_states, text_hidden_states)
+        """
+        dtype = encoder_hidden_states.dtype
+        batch_size, text_seq_length, _ = encoder_hidden_states.shape
+
+        # Concatenate text and image: [text, image]
+        hidden_states_combined = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        # QKV projection
+        qkv, _ = self.to_qkv(hidden_states_combined)
+        query, key, value = qkv.chunk(3, dim=-1)
+
+        # Reshape: [B, S, H*D] -> [B, S, H, D]
+        query = query.unflatten(-1, (self.num_heads, -1))
+        key = key.unflatten(-1, (self.num_heads, -1))
+        value = value.unflatten(-1, (self.num_heads, -1))
+
+        # QK normalization
+        query = self.norm_q(query).to(dtype=dtype)
+        key = self.norm_k(key).to(dtype=dtype)
+
+        # Apply RoPE only to image tokens (not text tokens)
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+            cos = cos.to(query.dtype)
+            sin = sin.to(query.dtype)
+            # Only apply RoPE to image part (after text_seq_length)
+            query_img = query[:, text_seq_length:, :, :]
+            key_img = key[:, text_seq_length:, :, :]
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            query_img = apply_rotary_emb(query_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
+            # key_img = self.rope(key_img, cos, sin)
+            key_img = apply_rotary_emb(key_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
+            query = torch.cat([query[:, :text_seq_length, :, :], query_img], dim=1)
+            key = torch.cat([key[:, :text_seq_length, :, :], key_img], dim=1)
+
+        # Handle KV cache for image editing
+        if kv_cache is not None and kv_cache_mode is not None:
+            if kv_cache_mode == KVCacheMode.WRITE:
+                kv_cache.store(key, value)
+            elif kv_cache_mode == KVCacheMode.READ:
+                k_cached, v_cached = kv_cache.get()
+                if k_cached is not None:
+                    key = torch.cat([k_cached, key], dim=1)
+                    value = torch.cat([v_cached, value], dim=1)
+            # KVCacheMode.SKIP: do nothing
+
+        # Attention computation
+        hidden_states_out = self.attn(query, key, value)
+        hidden_states_out = hidden_states_out.flatten(2, 3)
+        hidden_states_out = hidden_states_out.to(dtype)
+
+        # Output projection
+        hidden_states_out = self.to_out(hidden_states_out)
+
+        # Split back to text and image
+        encoder_hidden_states_out = hidden_states_out[:, :text_seq_length, :]
+        hidden_states_out = hidden_states_out[:, text_seq_length:, :]
+
+        return hidden_states_out, encoder_hidden_states_out
+
+
+class GlmImageTransformerBlock(nn.Module):
+    """Single transformer block for GLM-Image."""
+
+    def __init__(
+        self,
+        dim: int = 2560,
+        num_attention_heads: int = 64,
+        attention_head_dim: int = 40,
+        time_embed_dim: int = 512,
+    ) -> None:
+        super().__init__()
+
+        # 1. Attention with AdaLN
+        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
+        self.attn1 = GlmImageAttention(
+            dim=dim,
+            num_heads=num_attention_heads,
+            head_dim=attention_head_dim,
+        )
+
+        # 2. Feedforward
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        attention_kwargs: dict[str, Any] | None = None,
+        kv_cache: GlmImageLayerKVCache | None = None,
+        kv_cache_mode: KVCacheMode | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass for transformer block.
+
+        Args:
+            hidden_states: Image hidden states
+            encoder_hidden_states: Text hidden states
+            temb: Timestep embedding
+            image_rotary_emb: RoPE embeddings
+            attention_mask: Text attention mask
+            attention_kwargs: Additional attention arguments
+            kv_cache: Layer-specific KV cache for image editing
+            kv_cache_mode: Cache mode (WRITE, READ, SKIP)
+
+        Returns:
+            Tuple of (image_hidden_states, text_hidden_states)
+        """
+        # 1. Timestep conditioning via AdaLN
+        (
+            norm_hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            norm_encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        ) = self.norm1(hidden_states, encoder_hidden_states, temb)
+
+        # 2. Attention
+        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            attention_mask=attention_mask,
+            kv_cache=kv_cache,
+            kv_cache_mode=kv_cache_mode,
+        )
+        hidden_states = hidden_states + attn_hidden_states * gate_msa.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + attn_encoder_hidden_states * c_gate_msa.unsqueeze(1)
+
+        # 3. Feedforward
+        norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) * (
+            1 + c_scale_mlp.unsqueeze(1)
+        ) + c_shift_mlp.unsqueeze(1)
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output_context = self.ff(norm_encoder_hidden_states)
+        hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1)
+
+        return hidden_states, encoder_hidden_states
+
+
+class GlmImageTransformer2DModel(CachedTransformer):
+    """
+    GLM-Image Transformer model for 2D image generation.
+
+    This is the vllm-omni optimized version of the GLM-Image DiT model.
+
+    Args:
+        od_config: OmniDiffusionConfig containing model configuration. The
+            transformer hyper-parameters (e.g. patch size / channels / heads)
+            are read from `od_config.tf_model_config`.
+    """
+
+    def __init__(
+        self,
+        od_config: OmniDiffusionConfig,
+    ):
+        super().__init__()
+
+        patch_size = od_config.tf_model_config.patch_size
+        in_channels = od_config.tf_model_config.in_channels
+        out_channels = od_config.tf_model_config.out_channels
+        num_attention_heads = od_config.tf_model_config.num_attention_heads
+        attention_head_dim = od_config.tf_model_config.attention_head_dim
+        time_embed_dim = od_config.tf_model_config.time_embed_dim
+        condition_dim = od_config.tf_model_config.condition_dim
+        prior_vq_quantizer_codebook_size = od_config.tf_model_config.prior_vq_quantizer_codebook_size
+        text_embed_dim = od_config.tf_model_config.text_embed_dim
+
+        # Get num_layers from config if available
+        model_config = od_config.tf_model_config
+        if model_config is not None and hasattr(model_config, "num_layers"):
+            num_layers = model_config.num_layers
+
+        self.od_config = od_config
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords
+        pooled_projection_dim = 2 * 2 * condition_dim
+        inner_dim = num_attention_heads * attention_head_dim
+
+        # 1. RoPE
+        self.rope = GlmImageRotaryPosEmbed(attention_head_dim, patch_size, theta=10000.0)
+
+        # 2. Patch & Text-timestep embedding
+        self.image_projector = GlmImageImageProjector(in_channels, inner_dim, patch_size)
+        self.glyph_projector = FeedForward(text_embed_dim, inner_dim, inner_dim=inner_dim, activation_fn="gelu")
+        self.prior_token_embedding = nn.Embedding(prior_vq_quantizer_codebook_size, inner_dim)
+        self.prior_projector = FeedForward(inner_dim, inner_dim, inner_dim=inner_dim, activation_fn="linear-silu")
+
+        self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings(
+            embedding_dim=time_embed_dim,
+            condition_dim=condition_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            timesteps_dim=time_embed_dim,
+        )
+
+        # 3. Transformer blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                GlmImageTransformerBlock(inner_dim, num_attention_heads, attention_head_dim, time_embed_dim)
+                for _ in range(num_layers)
+            ]
+        )
+
+        # 4. Output projection
+        self.norm_out = GlmImageAdaLayerNormContinuous(inner_dim, time_embed_dim, elementwise_affine=False)
+        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels, bias=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        prior_token_id: torch.Tensor,
+        prior_token_drop: torch.Tensor,
+        timestep: torch.LongTensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        kv_cache: GlmImageKVCache | None = None,
+    ) -> torch.Tensor | Transformer2DModelOutput:
+        """
+        Forward pass of the GLM-Image Transformer.
+
+        Args:
+            hidden_states: Input latent tensor of shape [B, C, H, W].
+            encoder_hidden_states: Text embeddings of shape [B, S, D].
+            prior_token_id: Prior VQ token IDs.
+            prior_token_drop: Mask for dropping prior tokens (CFG).
+            timestep: Diffusion timestep.
+            target_size: Target image size for conditioning.
+            crop_coords: Crop coordinates for conditioning.
+            attention_kwargs: Additional attention arguments.
+            return_dict: Whether to return a dataclass.
+            attention_mask: Optional attention mask for text tokens.
+            image_rotary_emb: Pre-computed rotary embeddings.
+            kv_cache: Optional KV cache for image editing. When provided,
+                      the cache's mode determines behavior:
+                      - WRITE: Store KV from condition images
+                      - READ: Use cached KV during generation
+                      - SKIP: No caching (same as None)
+
+        Returns:
+            Output tensor or Transformer2DModelOutput.
+        """
+        batch_size, num_channels, height, width = hidden_states.shape
+
+        # Get KV cache mode
+        kv_cache_mode = kv_cache.mode if kv_cache is not None else None
+
+        # 1. RoPE
+        if image_rotary_emb is None:
+            image_rotary_emb = self.rope(hidden_states)
+            # Move to correct device
+            image_rotary_emb = (
+                image_rotary_emb[0].to(hidden_states.device),
+                image_rotary_emb[1].to(hidden_states.device),
+            )
+
+        # 2. Patch & Timestep embeddings
+        p = self.patch_size
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        hidden_states = self.image_projector(hidden_states)
+        encoder_hidden_states = self.glyph_projector(encoder_hidden_states)
+
+        # Prior embedding with dropout
+        prior_embedding = self.prior_token_embedding(prior_token_id)
+        prior_embedding[prior_token_drop] *= 0.0
+        prior_hidden_states = self.prior_projector(prior_embedding)
+        hidden_states = hidden_states + prior_hidden_states
+
+        # Timestep conditioning
+        temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype)
+
+        # 3. Transformer blocks
+        for layer_idx, block in enumerate(self.transformer_blocks):
+            # Get layer-specific KV cache if available
+            layer_kv_cache = kv_cache[layer_idx] if kv_cache is not None else None
+
+            hidden_states, encoder_hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                temb,
+                image_rotary_emb,
+                attention_mask,
+                attention_kwargs,
+                kv_cache=layer_kv_cache,
+                kv_cache_mode=kv_cache_mode,
+            )
+
+        # 4. Output norm & projection
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # 5. Unpatchify: [B, H'*W', C*p*p] -> [B, C, H, W]
+        hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p)
+        output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """
+        Load weights from pretrained checkpoint.
+
+        This method handles the mapping from diffusers weight names to vllm-omni weight names,
+        especially for fused QKV projections.
+        """
+        stacked_params_mapping = [
+            # Fused QKV projection: to_q, to_k, to_v -> to_qkv
+            (".to_qkv", ".to_q", "q"),
+            (".to_qkv", ".to_k", "k"),
+            (".to_qkv", ".to_v", "v"),
+        ]
+
+        params_dict = dict(self.named_parameters())
+
+        # Also include buffers (for any beta/eps parameters)
+        for name, buffer in self.named_buffers():
+            params_dict[name] = buffer
+
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            # Handle fused QKV projections
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+
+                # Map diffusers name to vllm-omni name
+                name = name.replace(weight_name, param_name)
+
+                if name not in params_dict:
+                    logger.warning(f"Skipping weight {name} - not found in model")
+                    break
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Standard weight loading (not fused)
+                if name not in params_dict:
+                    logger.warning(f"Skipping weight {name} - not found in model")
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add(name)
+
+        return loaded_params
+
+    def create_kv_cache(self) -> GlmImageKVCache:
+        """
+        Create a KV cache for image editing.
+
+        Returns a new GlmImageKVCache instance sized for this model's
+        number of transformer layers. Use this for image editing workflows.
+
+        Example:
+            kv_cache = transformer.create_kv_cache()
+            kv_cache.set_mode("write")
+            transformer(condition_image, kv_cache=kv_cache)
+            kv_cache.set_mode("read")
+            for t in timesteps:
+                transformer(noisy_target, kv_cache=kv_cache)
+            kv_cache.clear()
+
+        Returns:
+            GlmImageKVCache instance with correct number of layers.
+        """
+        return GlmImageKVCache(num_layers=len(self.transformer_blocks))
+
+    @property
+    def num_layers(self) -> int:
+        """Return number of transformer layers."""
+        return len(self.transformer_blocks)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Return dtype of model parameters."""
+        return next(self.parameters()).dtype

vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py

@@ -0,0 +1,981 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+GlmImagePipeline implementation for vLLM-Omni.
+
+This pipeline implements GLM-Image text-to-image generation with:
+- AR stage: GlmImageForConditionalGeneration generates prior tokens
+- DiT stage: GlmImageTransformer2DModel performs diffusion denoising
+- VAE: AutoencoderKL decodes latents to images
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+import logging
+import os
+import re
+from collections.abc import Iterable
+
+import numpy as np
+import PIL.Image
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import (
+    ByT5Tokenizer,
+    GlmImageForConditionalGeneration,
+    GlmImageProcessor,
+    T5EncoderModel,
+)
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+    GlmImageKVCache,
+    GlmImageTransformer2DModel,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_glm_image_pre_process_func(od_config: OmniDiffusionConfig):
+    """Get pre-processing function for GLM-Image pipeline.
+
+    Pre-processes condition images before they are sent to the pipeline.
+    This is called by DiffusionEngine before batching requests.
+    """
+    model_name = od_config.model
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+
+    vae_config_path = os.path.join(model_path, "vae/config.json")
+    with open(vae_config_path) as f:
+        vae_config = json.load(f)
+        block_out_channels = vae_config.get("block_out_channels", [128, 256, 512, 512])
+        vae_scale_factor = 2 ** (len(block_out_channels) - 1)
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
+    # GLM-Image uses patch_size=2 for transformer
+    patch_size = 2
+
+    def pre_process_func(requests: list[OmniDiffusionRequest]):
+        """Pre-process condition images for Image Edit mode."""
+        for req in requests:
+            images = req.pil_image
+            if images is None:
+                # Text-to-image mode, no preprocessing needed
+                continue
+
+            if not isinstance(images, list):
+                images = [images]
+
+            preprocessed = []
+            height, width = None, None
+
+            for img in images:
+                if isinstance(img, PIL.Image.Image):
+                    img_h, img_w = img.size[::-1]  # PIL is (width, height)
+                else:
+                    img_h, img_w = img.shape[:2]
+
+                # Align to multiple of vae_scale_factor * patch_size
+                multiple_of = vae_scale_factor * patch_size
+                img_h = (img_h // multiple_of) * multiple_of
+                img_w = (img_w // multiple_of) * multiple_of
+
+                processed = image_processor.preprocess(img, height=img_h, width=img_w)
+                preprocessed.append(processed)
+
+                # Use first image dimensions as default
+                if height is None:
+                    height, width = img_h, img_w
+
+            # Store in request
+            req.preprocessed_image = preprocessed
+            req.prompt_image = images  # Keep original PIL images
+            if req.height is None:
+                req.height = height
+            if req.width is None:
+                req.width = width
+
+        return requests
+
+    return pre_process_func
+
+
+def get_glm_image_post_process_func(od_config: OmniDiffusionConfig):
+    """Get post-processing function for GLM-Image pipeline."""
+    model_name = od_config.model
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+
+    vae_config_path = os.path.join(model_path, "vae/config.json")
+    with open(vae_config_path) as f:
+        vae_config = json.load(f)
+        block_out_channels = vae_config.get("block_out_channels", [128, 256, 512, 512])
+        vae_scale_factor = 2 ** (len(block_out_channels) - 1)
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
+
+    def post_process_func(images: torch.Tensor) -> list[PIL.Image.Image]:
+        return image_processor.postprocess(images, output_type="pil")
+
+    return post_process_func
+
+
+def calculate_shift(
+    image_seq_len: int,
+    base_seq_len: int = 256,
+    base_shift: float = 0.25,
+    max_shift: float = 0.75,
+) -> float:
+    """Calculate timestep shift based on image sequence length."""
+    m = (image_seq_len / base_seq_len) ** 0.5
+    mu = m * max_shift + base_shift
+    return mu
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, int]:
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps.
+    Handles custom timesteps and sigmas schedules.
+    """
+    accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+    accepts_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+
+    if timesteps is not None and sigmas is not None:
+        # Both provided - check if scheduler supports both
+        if not accepts_timesteps and not accepts_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep or sigma schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif timesteps is not None:
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        if not accepts_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigma schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+
+    return timesteps, num_inference_steps
+
+
+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")
+
+
+class GlmImagePipeline(nn.Module):
+    """
+    GLM-Image Pipeline for text-to-image and image-to-image generation.
+
+    This pipeline integrates:
+    - AR model (GlmImageForConditionalGeneration): Generates prior image tokens
+    - Text encoder (T5EncoderModel): Encodes glyph/text embeddings
+    - DiT model (GlmImageTransformer2DModel): Diffusion transformer
+    - VAE (AutoencoderKL): Encodes/decodes images to/from latent space
+
+    The pipeline flow:
+    1. AR generates prior_token_ids from text prompt
+    2. T5 encodes glyph text for text rendering
+    3. DiT performs iterative denoising conditioned on prior tokens
+    4. VAE decodes final latents to image
+    """
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+        self.device = get_local_device()
+
+        model = od_config.model
+        local_files_only = os.path.exists(model)
+
+        if local_files_only:
+            model_path = model
+        else:
+            model_path = download_weights_from_hf_specific(model, od_config.revision, ["*"])
+
+        # Load scheduler
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model_path, subfolder="scheduler", local_files_only=True
+        )
+
+        # Load AR model (vision_language_encoder)
+        logger.info("Loading GlmImageForConditionalGeneration (AR model)...")
+        self.vision_language_encoder = GlmImageForConditionalGeneration.from_pretrained(
+            model_path,
+            subfolder="vision_language_encoder",
+            local_files_only=True,
+            torch_dtype=torch.bfloat16,
+        ).to(self.device)
+        self.vision_language_encoder.eval()
+
+        # Load processor for AR model
+        self.processor = GlmImageProcessor.from_pretrained(model_path, subfolder="processor", local_files_only=True)
+
+        # Load text encoder (T5 for glyph embeddings)
+        logger.info("Loading T5EncoderModel (glyph encoder)...")
+        self.text_encoder = T5EncoderModel.from_pretrained(
+            model_path,
+            subfolder="text_encoder",
+            local_files_only=True,
+            torch_dtype=torch.bfloat16,
+        ).to(self.device)
+        self.text_encoder.eval()
+
+        # Load tokenizer for glyph encoding
+        self.tokenizer = ByT5Tokenizer.from_pretrained(model_path, subfolder="tokenizer", local_files_only=True)
+
+        # Load VAE
+        logger.info("Loading AutoencoderKL (VAE)...")
+        self.vae = AutoencoderKL.from_pretrained(
+            model_path, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16
+        ).to(self.device)
+        self.vae.eval()
+
+        # Load transformer (DiT)
+        logger.info("Loading GlmImageTransformer2DModel (DiT)...")
+        self.transformer = GlmImageTransformer2DModel(od_config=od_config)
+
+        # Weight sources for DiT loading
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=od_config.revision,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+
+        # Configure scale factors
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.default_sample_size = 128
+
+        # Get transformer config for patch size
+        self._patch_size = getattr(self.transformer, "patch_size", 2)
+
+    # ==================== Input Validation ====================
+
+    def check_inputs(
+        self,
+        prompt: str | list[str] | None,
+        height: int | None,
+        width: int | None,
+        prompt_embeds: torch.Tensor | None = None,
+    ) -> None:
+        """Validate input arguments before generation."""
+        # Check dimension alignment
+        multiple_of = self.vae_scale_factor * self._patch_size
+        if height is not None and height % multiple_of != 0:
+            logger.warning(
+                f"`height` should be divisible by {multiple_of} but is {height}. "
+                "Dimensions will be adjusted accordingly."
+            )
+        if width is not None and width % multiple_of != 0:
+            logger.warning(
+                f"`width` should be divisible by {multiple_of} but is {width}. Dimensions will be adjusted accordingly."
+            )
+
+        # Check prompt/prompt_embeds mutual exclusivity
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. "
+                "Please provide only one of the two."
+            )
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Provide either `prompt` or `prompt_embeds`. Cannot leave both undefined.")
+
+        # Check prompt type
+        if prompt is not None and not isinstance(prompt, (str, list)):
+            raise ValueError(f"`prompt` must be of type `str` or `list` but is {type(prompt)}")
+
+    # ==================== AR Stage Methods ====================
+
+    @staticmethod
+    def _compute_generation_params(
+        image_grid_thw: torch.Tensor,
+        is_text_to_image: bool,
+    ) -> tuple[int, int, int, int]:
+        """
+        Compute AR generation parameters from image grid.
+
+        Args:
+            image_grid_thw: Image grid tensor of shape [N, 3] where each row is [t, h, w]
+            is_text_to_image: Whether this is text-to-image (vs image-to-image)
+
+        Returns:
+            Tuple of (max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w)
+        """
+        grid_sizes = []
+        grid_hw = []
+
+        for i in range(image_grid_thw.shape[0]):
+            t, h, w = image_grid_thw[i].tolist()
+            grid_sizes.append(int(h * w))
+            grid_hw.append((int(h), int(w)))
+
+        if not is_text_to_image:
+            # Image-to-image: only generate target image tokens
+            max_new_tokens = grid_sizes[-1] + 1
+            large_image_start_offset = 0
+            target_grid_h, target_grid_w = grid_hw[-1]
+        else:
+            # Text-to-image: generate both small preview and large target
+            total_tokens = sum(grid_sizes)
+            max_new_tokens = total_tokens + 1
+            large_image_start_offset = sum(grid_sizes[1:])
+            target_grid_h, target_grid_w = grid_hw[0]
+
+        return max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w
+
+    @staticmethod
+    def _extract_large_image_tokens(
+        outputs: torch.Tensor, input_length: int, large_image_start_offset: int, large_image_tokens: int
+    ) -> torch.Tensor:
+        """Extract large image tokens from AR output."""
+        generated_tokens = outputs[0][input_length:]
+        large_image_start = large_image_start_offset
+        large_image_end = large_image_start + large_image_tokens
+        return generated_tokens[large_image_start:large_image_end]
+
+    @staticmethod
+    def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) -> torch.Tensor:
+        """Upsample token IDs by 2x using nearest neighbor interpolation."""
+        token_ids = token_ids.view(1, 1, token_h, token_w)
+        token_ids = torch.nn.functional.interpolate(token_ids.float(), scale_factor=2, mode="nearest").to(
+            dtype=torch.long
+        )
+        token_ids = token_ids.view(1, -1)
+        return token_ids
+
+    @torch.inference_mode()
+    def generate_prior_tokens(
+        self,
+        prompt: str,
+        height: int,
+        width: int,
+        image: list[PIL.Image.Image] | None = None,
+        factor: int = 32,
+    ) -> tuple[torch.Tensor, list[torch.Tensor] | None]:
+        """
+        Generate prior tokens using the AR model.
+
+        Args:
+            prompt: Text prompt for generation
+            height: Target image height
+            width: Target image width
+            image: Optional condition images for image-to-image
+            factor: Token factor (default 32)
+
+        Returns:
+            Tuple of (prior_token_ids, prior_token_image_ids)
+            prior_token_image_ids is a list of tensors, one per condition image
+        """
+        device = self.vision_language_encoder.device
+        height = (height // factor) * factor
+        width = (width // factor) * factor
+        is_text_to_image = image is None or len(image) == 0
+
+        # Build message content
+        content = []
+        if image is not None:
+            for img in image:
+                content.append({"type": "image", "image": img})
+        content.append({"type": "text", "text": prompt})
+        messages = [{"role": "user", "content": content}]
+
+        # Apply chat template - processor will handle target dimensions and build grid
+        inputs = self.processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            target_h=height,
+            target_w=width,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(device)
+
+        image_grid_thw = inputs.get("image_grid_thw")
+
+        # Compute generation parameters from the full grid
+        max_new_tokens, large_image_offset, token_h, token_w = self._compute_generation_params(
+            image_grid_thw=image_grid_thw, is_text_to_image=is_text_to_image
+        )
+
+        # Process condition images if provided
+        # Use image_grid_thw[:-1] to exclude the target image grid (last entry)
+        prior_token_image_ids = None
+        if image is not None and image_grid_thw is not None and len(image_grid_thw) > 1:
+            # Get features only for condition images (exclude target image grid)
+            condition_grid = image_grid_thw[:-1]
+            prior_token_image_embed = self.vision_language_encoder.get_image_features(
+                inputs["pixel_values"], condition_grid
+            )
+            prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0)
+            flat_prior_token_image_ids = self.vision_language_encoder.get_image_tokens(
+                prior_token_image_embed, condition_grid
+            )
+            # Split by image grid sizes and convert to list
+            split_sizes = (condition_grid.prod(dim=-1)).tolist()
+            prior_token_image_ids_list = torch.split(flat_prior_token_image_ids, split_sizes, dim=0)
+            # Convert to list with upsampling
+            prior_token_image_ids = []
+            for i, token_ids in enumerate(prior_token_image_ids_list):
+                grid_t, grid_h, grid_w = condition_grid[i].tolist()
+                token_ids = token_ids.view(1, -1)
+                # Upsample 2x (from d32 to d64)
+                token_ids_upsampled = self._upsample_token_ids(token_ids, grid_h, grid_w)
+                prior_token_image_ids.append(token_ids_upsampled)
+
+        # Generate with AR model
+        outputs = self.vision_language_encoder.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+        )
+
+        # Extract and upsample tokens
+        large_image_tokens = token_h * token_w
+        prior_token_ids_d32 = self._extract_large_image_tokens(
+            outputs, inputs["input_ids"].shape[-1], large_image_offset, large_image_tokens
+        )
+        prior_token_ids = self._upsample_token_ids(prior_token_ids_d32, token_h, token_w)
+
+        return prior_token_ids, prior_token_image_ids
+
+    # ==================== Text Encoding Methods ====================
+
+    def get_glyph_texts(self, prompt: str | list[str]) -> list[str]:
+        """Extract text within quotes for glyph rendering."""
+        prompt = prompt[0] if isinstance(prompt, list) else prompt
+        ocr_texts = (
+            re.findall(r"'([^']*)'", prompt)
+            + re.findall(r"“([^“”]*)”", prompt)
+            + re.findall(r'"([^"]*)"', prompt)
+            + re.findall(r"「([^「」]*)」", prompt)
+        )
+        return ocr_texts
+
+    def _get_glyph_embeds(
+        self,
+        prompt: str | list[str],
+        max_sequence_length: int = 2048,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
+        """Get glyph embeddings from T5 encoder for text rendering."""
+        device = device or self.device
+        dtype = dtype or self.text_encoder.dtype
+
+        glyph_texts = self.get_glyph_texts(prompt)
+        input_ids = self.tokenizer(
+            glyph_texts if len(glyph_texts) > 0 else [""],
+            max_length=max_sequence_length,
+            truncation=True,
+        ).input_ids
+
+        # Pad to even length
+        input_ids = [[self.tokenizer.pad_token_id] * ((len(ids) + 1) % 2) + ids for ids in input_ids]
+        max_length = max(len(ids) for ids in input_ids)
+
+        attention_mask = torch.tensor(
+            [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids],
+            device=device,
+        )
+        input_ids = torch.tensor(
+            [ids + [self.tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids],
+            device=device,
+        )
+
+        outputs = self.text_encoder(input_ids, attention_mask=attention_mask)
+        glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0)
+
+        return glyph_embeds.to(device=device, dtype=dtype)
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        do_classifier_free_guidance: bool = True,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.Tensor | None = None,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        max_sequence_length: int = 2048,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Encode prompt into glyph embeddings for text rendering."""
+        device = device or self.device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds = self._get_glyph_embeds(prompt, max_sequence_length, device, dtype)
+
+        seq_len = prompt_embeds.size(1)
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        negative_prompt_embeds = None
+        if do_classifier_free_guidance:
+            negative_prompt = [""] * batch_size
+            negative_prompt_embeds = self._get_glyph_embeds(negative_prompt, max_sequence_length, device, dtype)
+            seq_len = negative_prompt_embeds.size(1)
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds, negative_prompt_embeds
+
+    # ==================== Latent Preparation ====================
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: torch.Generator | None,
+        latents: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Prepare random noise latents."""
+        if latents is not None:
+            return latents.to(device)
+
+        shape = (
+            batch_size,
+            num_channels_latents,
+            int(height) // self.vae_scale_factor,
+            int(width) // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(f"Passed {len(generator)} generators but batch size is {batch_size}.")
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        return latents
+
+    def diffuse(
+        self,
+        latents: torch.Tensor,
+        prior_token_id: torch.Tensor,
+        prompt_embeds: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        timesteps: torch.Tensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        guidance_scale: float,
+        do_classifier_free_guidance: bool,
+        kv_caches: GlmImageKVCache | None = None,
+    ) -> torch.Tensor:
+        """
+        Denoising loop for diffusion process with CFG-Parallel support.
+
+        Args:
+            latents: Initial noise latents
+            prior_token_id: Prior tokens generated by AR model
+            prompt_embeds: Encoded positive prompt embeddings (glyph embeddings)
+            negative_prompt_embeds: Encoded negative prompt embeddings
+            timesteps: Denoising timesteps
+            target_size: Target image size tensor [[height, width]]
+            crop_coords: Crop coordinates tensor
+            guidance_scale: CFG scale
+            do_classifier_free_guidance: Whether to apply CFG
+            kv_caches: Optional KV cache for Image Edit mode
+
+        Returns:
+            Denoised latents ready for VAE decode
+        """
+        # Prepare conditional/unconditional drop flags
+        prior_token_drop_cond = torch.full_like(prior_token_id, False, dtype=torch.bool)
+        prior_token_drop_uncond = torch.full_like(prior_token_id, True, dtype=torch.bool)
+
+        transformer_dtype = self.transformer.dtype
+
+        # Enable CFG-parallel: rank0 computes positive, rank1 computes negative
+        cfg_parallel_ready = do_classifier_free_guidance and get_classifier_free_guidance_world_size() > 1
+
+        for i, t in enumerate(timesteps):
+            latent_model_input = latents.to(transformer_dtype)
+            timestep = t.expand(latents.shape[0]) - 1
+
+            if cfg_parallel_ready:
+                cfg_group = get_cfg_group()
+                cfg_rank = get_classifier_free_guidance_rank()
+
+                if cfg_rank == 0:
+                    # Rank 0: Compute positive (conditional) prediction
+                    local_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        encoder_hidden_states=prompt_embeds,
+                        prior_token_id=prior_token_id,
+                        prior_token_drop=prior_token_drop_cond,
+                        timestep=timestep,
+                        target_size=target_size,
+                        crop_coords=crop_coords,
+                        kv_cache=kv_caches,
+                        return_dict=False,
+                    )[0].float()
+                else:
+                    # Rank 1: Compute negative (unconditional) prediction
+                    local_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        prior_token_id=prior_token_id,
+                        prior_token_drop=prior_token_drop_uncond,
+                        timestep=timestep,
+                        target_size=target_size,
+                        crop_coords=crop_coords,
+                        kv_cache=kv_caches,
+                        return_dict=False,
+                    )[0].float()
+
+                # All-gather predictions from all ranks
+                gathered = cfg_group.all_gather(local_pred, separate_tensors=True)
+
+                if cfg_rank == 0:
+                    # Rank 0: Combine predictions and apply CFG
+                    noise_pred_cond = gathered[0]
+                    noise_pred_uncond = gathered[1]
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+                    # Scheduler step
+                    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                # Broadcast updated latents to all ranks
+                cfg_group.broadcast(latents, src=0)
+
+            else:
+                # Sequential CFG (single GPU or no CFG)
+                # Conditional forward pass
+                noise_pred_cond = self.transformer(
+                    hidden_states=latent_model_input,
+                    encoder_hidden_states=prompt_embeds,
+                    prior_token_id=prior_token_id,
+                    prior_token_drop=prior_token_drop_cond,
+                    timestep=timestep,
+                    target_size=target_size,
+                    crop_coords=crop_coords,
+                    kv_cache=kv_caches,
+                    return_dict=False,
+                )[0].float()
+
+                if do_classifier_free_guidance:
+                    # Unconditional forward pass
+                    noise_pred_uncond = self.transformer(
+                        hidden_states=latent_model_input,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        prior_token_id=prior_token_id,
+                        prior_token_drop=prior_token_drop_uncond,
+                        timestep=timestep,
+                        target_size=target_size,
+                        crop_coords=crop_coords,
+                        kv_cache=kv_caches,
+                        return_dict=False,
+                    )[0].float()
+
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+                else:
+                    noise_pred = noise_pred_cond
+
+                # Scheduler step
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+        return latents
+
+    # ==================== Main Forward Pass ====================
+
+    def _prepare_condition_image_kv_cache(
+        self,
+        condition_images: list[torch.Tensor],
+        prior_token_image_ids: list[torch.Tensor],
+        prompt_embeds: torch.Tensor,
+        generator: torch.Generator | None = None,
+    ) -> GlmImageKVCache:
+        """
+        Prepare KV cache by running condition images through transformer at timestep 0.
+
+        This is used for Image Edit mode where we need to cache the condition image's
+        KV states for cross-attention during denoising.
+
+        Args:
+            condition_images: List of preprocessed condition images
+            prior_token_image_ids: Prior token IDs for each condition image from AR model
+            prompt_embeds: Prompt embeddings (used to get dtype)
+            generator: Optional random generator
+
+        Returns:
+            GlmImageKVCache with cached KV states from condition images
+        """
+        kv_caches = self.transformer.create_kv_cache()
+        kv_caches.set_mode("write")
+
+        # Prepare VAE normalization parameters
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.latent_channels, 1, 1)
+            .to(device=self.device, dtype=prompt_embeds.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.vae.config.latents_std)
+            .view(1, self.vae.config.latent_channels, 1, 1)
+            .to(device=self.device, dtype=prompt_embeds.dtype)
+        )
+
+        # Process each condition image through transformer to populate KV cache
+        for condition_image, condition_prior_token_id in zip(condition_images, prior_token_image_ids):
+            condition_image = condition_image.to(device=self.device, dtype=prompt_embeds.dtype)
+
+            # Encode condition image to latent space
+            # Use argmax (mode) for deterministic encoding of condition images
+            condition_latent = retrieve_latents(
+                self.vae.encode(condition_image), generator=generator, sample_mode="argmax"
+            )
+            condition_latent = (condition_latent - latents_mean) / latents_std
+
+            # Run forward pass at timestep 0 to cache KV states
+            # Empty encoder_hidden_states since we only want to cache image features
+            _ = self.transformer(
+                hidden_states=condition_latent,
+                encoder_hidden_states=torch.zeros_like(prompt_embeds)[:1, :0, ...],
+                prior_token_id=condition_prior_token_id,
+                prior_token_drop=torch.full_like(condition_prior_token_id, False, dtype=torch.bool),
+                timestep=torch.zeros((1,), device=self.device),
+                target_size=torch.tensor([condition_image.shape[-2:]], device=self.device, dtype=prompt_embeds.dtype),
+                crop_coords=torch.zeros((1, 2), device=self.device, dtype=prompt_embeds.dtype),
+                kv_cache=kv_caches,
+                return_dict=False,
+            )
+
+        return kv_caches
+
+    @torch.inference_mode()
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        """
+        Main generation forward pass.
+
+        Args:
+            req: OmniDiffusionRequest with generation parameters
+
+        Returns:
+            DiffusionOutput containing generated image
+        """
+        prompt = req.prompt or ""
+        if isinstance(prompt, list):
+            prompt = prompt[0] if prompt else ""
+
+        # Get pre-computed prompt embeddings if provided
+        prompt_embeds = req.prompt_embeds if isinstance(req.prompt_embeds, torch.Tensor) else None
+
+        # Get condition images for Image Edit mode
+        # Use pre-processed images from pre_process_func
+        preprocessed_images = req.preprocessed_image
+        condition_images = getattr(req, "prompt_image", None)
+        img_height = req.height
+        img_width = req.width
+
+        is_image_edit = preprocessed_images is not None
+
+        # Use image dimensions as default if available
+        height = req.height or img_height or self.default_sample_size * self.vae_scale_factor
+        width = req.width or img_width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.num_inference_steps or 50
+        guidance_scale = req.guidance_scale or 1.5
+
+        # 0. Validate inputs
+        self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds)
+
+        batch_size = 1
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # Set seed if provided
+        generator = None
+        if req.seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(req.seed)
+
+        # 1. Get prior tokens - either from external source (multistage) or generate internally
+        # Check if prior_token_ids are provided externally (from AR stage in multistage mode)
+        external_prior_tokens = req.extra.get("prior_token_ids") if req.extra else None
+        external_prior_image_ids = req.extra.get("prior_token_image_ids") if req.extra else None
+
+        if external_prior_tokens is not None:
+            # Multistage mode: use externally provided prior tokens from vLLM AR stage
+            logger.info("Using externally provided prior tokens from AR stage...")
+            prior_token_id = external_prior_tokens
+            if isinstance(prior_token_id, list):
+                prior_token_id = torch.tensor(prior_token_id, dtype=torch.long, device=self.device)
+            elif isinstance(prior_token_id, torch.Tensor):
+                prior_token_id = prior_token_id.to(device=self.device, dtype=torch.long)
+            # Ensure shape is [1, num_tokens] for batch processing
+            if prior_token_id.dim() == 1:
+                prior_token_id = prior_token_id.unsqueeze(0)
+            prior_token_image_ids = external_prior_image_ids
+        else:
+            # Single-stage mode: generate prior tokens with internal AR model
+            logger.info("Generating prior tokens with AR model...")
+            prior_token_id, prior_token_image_ids = self.generate_prior_tokens(
+                prompt=prompt,
+                image=condition_images,
+                height=height,
+                width=width,
+            )
+
+        # 2. Encode prompt for glyph embeddings
+        logger.info("Encoding prompt...")
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            num_images_per_prompt=1,
+            prompt_embeds=prompt_embeds,
+            device=self.device,
+            dtype=self.transformer.dtype,
+        )
+
+        # 3. Prepare KV cache for Image Edit mode
+        kv_caches = None
+        if is_image_edit and prior_token_image_ids is not None:
+            logger.info("Preparing KV cache for Image Edit mode...")
+            kv_caches = self._prepare_condition_image_kv_cache(
+                condition_images=preprocessed_images,
+                prior_token_image_ids=prior_token_image_ids,
+                prompt_embeds=prompt_embeds,
+                generator=generator,
+            )
+            # Switch to read mode for denoising
+            kv_caches.set_mode("read")
+
+        # 4. Prepare latents
+        latent_channels = self.transformer.in_channels
+        latents = self.prepare_latents(
+            batch_size=batch_size,
+            num_channels_latents=latent_channels,
+            height=height,
+            width=width,
+            dtype=prompt_embeds.dtype,
+            device=self.device,
+            generator=generator,
+        )
+
+        # 5. Prepare timesteps
+        image_seq_len = ((height // self.vae_scale_factor) * (width // self.vae_scale_factor)) // (self._patch_size**2)
+        timesteps_array = np.linspace(self.scheduler.config.num_train_timesteps, 1.0, num_inference_steps + 1)[:-1]
+        timesteps_array = timesteps_array.astype(np.int64).astype(np.float32)
+        sigmas = timesteps_array / self.scheduler.config.num_train_timesteps
+
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("base_shift", 0.25),
+            self.scheduler.config.get("max_shift", 0.75),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, self.device, timesteps_array.tolist(), sigmas.tolist(), mu=mu
+        )
+
+        # 6. Prepare conditioning tensors
+        target_size = torch.tensor([[height, width]], dtype=prompt_embeds.dtype, device=self.device)
+        crop_coords = torch.zeros((1, 2), dtype=prompt_embeds.dtype, device=self.device)
+
+        # 7. Denoising loop with CFG-parallel support
+        logger.info(f"Starting denoising loop with {num_inference_steps} steps...")
+        latents = self.diffuse(
+            latents=latents,
+            prior_token_id=prior_token_id,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            timesteps=timesteps,
+            target_size=target_size,
+            crop_coords=crop_coords,
+            guidance_scale=guidance_scale,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            kv_caches=kv_caches,
+        )
+
+        # 8. VAE decode
+        logger.info("Decoding latents with VAE...")
+        latents = latents.to(self.vae.dtype)
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.latent_channels, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.vae.config.latents_std)
+            .view(1, self.vae.config.latent_channels, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents = latents * latents_std + latents_mean
+        image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
+
+        # 9. Leave post-process to vllm-omni pipeline
+
+        return DiffusionOutput(output=image)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load transformer weights."""
+        # Filter weights for transformer only
+        transformer_weights = (
+            (name.replace("transformer.", ""), weight) for name, weight in weights if name.startswith("transformer.")
+        )
+        return self.transformer.load_weights(transformer_weights)

vllm_omni/diffusion/registry.py

@@ -29,6 +29,11 @@ _DIFFUSION_MODELS = {
         "pipeline_qwen_image_layered",
         "QwenImageLayeredPipeline",
     ),
+    "GlmImagePipeline": (
+        "glm_image",
+        "pipeline_glm_image",
+        "GlmImagePipeline",
+    ),
     "ZImagePipeline": (
         "z_image",
         "pipeline_z_image",
@@ -117,6 +122,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "QwenImagePipeline": "get_qwen_image_post_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
     "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func",
+    "GlmImagePipeline": "get_glm_image_post_process_func",
     "ZImagePipeline": "get_post_process_func",
     "OvisImagePipeline": "get_ovis_image_post_process_func",
     "WanPipeline": "get_wan22_post_process_func",
@@ -133,6 +139,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     # arch: pre_process_func
     # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
     # where mod_folder and mod_relname are  defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
+    "GlmImagePipeline": "get_glm_image_pre_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
     "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func",

vllm_omni/engine/__init__.py

@@ -2,20 +2,14 @@
 Engine components for vLLM-Omni.
 """
 
-import time
-from collections.abc import Mapping
 from typing import Any
 
 import msgspec
 import torch
 from vllm.v1.engine import (
-    EngineCoreEvent,
+    EngineCoreOutput,
+    EngineCoreOutputs,
     EngineCoreRequest,
-    FinishReason,
-    LogprobsLists,
-    LogprobsTensors,
-    SchedulerStats,
-    UtilityOutput,
 )
 
 
@@ -79,64 +73,9 @@ class OmniEngineCoreRequest(EngineCoreRequest):
     additional_information: AdditionalInformationPayload | None = None
 
 
-class OmniEngineCoreOutput(
-    msgspec.Struct,
-    array_like=True,  # type: ignore[call-arg]
-    omit_defaults=True,  # type: ignore[call-arg]
-    gc=False,
-):  # type: ignore[call-arg]
-    request_id: str
-    new_token_ids: list[int]
-
-    new_logprobs: LogprobsLists | None = None
-    new_prompt_logprobs_tensors: LogprobsTensors | None = None
-
+class OmniEngineCoreOutput(EngineCoreOutput):
     pooling_output: dict[str, torch.Tensor] | None = None
 
-    finish_reason: FinishReason | None = None
-    stop_reason: int | str | None = None
-    events: list[EngineCoreEvent] | None = None
-    kv_transfer_params: dict[str, Any] | None = None
 
-    trace_headers: Mapping[str, str] | None = None
-    # The number of tokens with prefix cache hits.
-    num_cached_tokens: int = 0
-
-    # The number of NaNs in logits.
-    # A value greater than 0 indicates that the output is corrupted.
-    num_nans_in_logits: int = 0
-
-    @property
-    def finished(self) -> bool:
-        return self.finish_reason is not None
-
-
-class OmniEngineCoreOutputs(
-    msgspec.Struct,
-    array_like=True,  # type: ignore[call-arg]
-    omit_defaults=True,  # type: ignore[call-arg]
-    gc=False,
-):  # type: ignore[call-arg]
-    # NOTE(Nick): We could consider ways to make this more compact,
-    # e.g. columnwise layout
-
-    engine_index: int = 0
-
-    # [num_reqs]
+class OmniEngineCoreOutputs(EngineCoreOutputs):
     outputs: list[OmniEngineCoreOutput] = []
-    scheduler_stats: SchedulerStats | None = None
-    timestamp: float = 0.0
-
-    utility_output: UtilityOutput | None = None
-    finished_requests: set[str] | None = None
-
-    # In DP case, used to signal that the current wave of requests
-    # has finished and the engines are paused.
-    wave_complete: int | None = None
-    # In DP case, used to signal that a request was received for an
-    # "old" wave, so the next wave needs to be started in other engines.
-    start_wave: int | None = None
-
-    def __post_init__(self):
-        if self.timestamp == 0.0:
-            self.timestamp = time.monotonic()

vllm_omni/engine/output_processor.py

@@ -8,8 +8,6 @@ from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.tokenizers import TokenizerLike
 from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
-from vllm.v1.engine.detokenizer import IncrementalDetokenizer
-from vllm.v1.engine.logprobs import LogprobsProcessor
 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
 from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestOutputCollector, RequestState
 from vllm.v1.engine.parallel_sampling import ParentRequest
@@ -37,65 +35,6 @@ class OmniRequestState(RequestState):
         self.mm_type: str | None = None
         self.mm_accumulated: Dict[str, Any] | None = None
 
-    @classmethod
-    def from_new_request(
-        cls,
-        tokenizer: TokenizerLike,
-        request: EngineCoreRequest,
-        prompt: str | None,
-        parent_req: ParentRequest | None,
-        request_index: int,
-        queue: Any | None,
-        log_stats: bool,
-        stream_interval: int,
-    ) -> "OmniRequestState":
-        if sampling_params := request.sampling_params:
-            if not sampling_params.detokenize:
-                tokenizer = None
-            output_kind = sampling_params.output_kind
-            logprobs_processor = LogprobsProcessor.from_new_request(
-                tokenizer=tokenizer,
-                request=request,
-            )
-            detokenizer = IncrementalDetokenizer.from_new_request(
-                tokenizer=tokenizer,
-                request=request,
-            )
-            max_tokens_param = sampling_params.max_tokens
-            top_p = sampling_params.top_p
-            n = sampling_params.n
-            temperature = sampling_params.temperature
-        else:
-            logprobs_processor = None
-            detokenizer = None
-            max_tokens_param = None
-            top_p = None
-            n = None
-            temperature = None
-            assert request.pooling_params is not None
-            output_kind = request.pooling_params.output_kind
-
-        return cls(
-            request_id=request.request_id,
-            parent_req=parent_req,
-            request_index=request_index,
-            lora_name=(request.lora_request.name if request.lora_request is not None else None),
-            output_kind=output_kind,
-            prompt=prompt,
-            prompt_token_ids=request.prompt_token_ids,
-            prompt_embeds=request.prompt_embeds,
-            logprobs_processor=logprobs_processor,
-            detokenizer=detokenizer,
-            max_tokens_param=max_tokens_param,
-            top_p=top_p,
-            n=n,
-            temperature=temperature,
-            arrival_time=request.arrival_time,
-            queue=queue,
-            log_stats=log_stats,
-            stream_interval=stream_interval,
-        )
-
     def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
         if payload is None:
             return

vllm_omni/entrypoints/async_omni_llm.py

@@ -10,7 +10,7 @@ import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.tokenizers import init_tokenizer_from_config
+from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.tracing import init_tracer
 from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
 from vllm.usage.usage_lib import UsageContext
@@ -111,7 +111,7 @@ class AsyncOmniLLM(AsyncLLM):
             tokenizer = None
         else:
             # Tokenizer (+ ensure liveness if running in another process).
-            tokenizer = init_tokenizer_from_config(model_config=vllm_config.model_config)
+            tokenizer = cached_tokenizer_from_config(model_config=vllm_config.model_config)
 
         # InputProcessor (converts Inputs --> EngineCoreRequests).
         self.input_processor = OmniInputProcessor(

vllm_omni/entrypoints/chat_utils.py

@@ -22,7 +22,6 @@ from vllm.entrypoints.chat_utils import (
     _postprocess_messages,
     _ToolParser,
 )
-from vllm.transformers_utils.tokenizer import AnyTokenizer
 
 
 class OmniAsyncMultiModalItemTracker(AsyncMultiModalItemTracker):
@@ -129,7 +128,6 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser):
 def parse_chat_messages_futures(
     messages: list[ChatCompletionMessageParam],
     model_config: ModelConfig,
-    tokenizer: AnyTokenizer,
     content_format: _ChatTemplateContentFormat,
     mm_processor_kwargs: dict[str, Any] | None = None,
 ) -> tuple[
@@ -138,7 +136,7 @@ def parse_chat_messages_futures(
     MultiModalUUIDDict | None,
 ]:
     conversation: list[ConversationMessage] = []
-    mm_tracker = OmniAsyncMultiModalItemTracker(model_config, tokenizer)
+    mm_tracker = OmniAsyncMultiModalItemTracker(model_config)
 
     for msg in messages:
         sub_messages = _parse_chat_message_content(

vllm_omni/entrypoints/omni_llm.py

@@ -1,12 +1,15 @@
+from collections.abc import Callable
 from typing import Any
 
 import cloudpickle
 from pydantic import ValidationError
+from tqdm import tqdm
 
 # External library imports (vLLM)
 from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
 from vllm.entrypoints.llm import LLM
 from vllm.logger import init_logger
+from vllm.outputs import PoolingRequestOutput, RequestOutput
 from vllm.plugins.io_processors import get_io_processor
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.counter import Counter
@@ -190,3 +193,47 @@ class OmniLLM(LLM):
             self.close()
         except Exception as e:
             logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
+
+    def _run_engine(self, *, use_tqdm: bool | Callable[..., tqdm] = True) -> list[RequestOutput | PoolingRequestOutput]:
+        # Initialize tqdm.
+        if use_tqdm:
+            num_requests = self.llm_engine.get_num_unfinished_requests()
+            tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
+            pbar = tqdm_func(
+                total=num_requests,
+                desc="Processed prompts",
+                dynamic_ncols=True,
+                postfix=(f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s"),
+            )
+
+        # Run the engine.
+        outputs: list[RequestOutput | PoolingRequestOutput] = []
+        total_in_toks = 0
+        total_out_toks = 0
+        while self.llm_engine.has_unfinished_requests():
+            step_outputs = self.llm_engine.step()
+            for output in step_outputs:
+                if output.finished:
+                    outputs.append(output)
+                    if use_tqdm:
+                        if isinstance(output, RequestOutput):
+                            # Calculate tokens only for RequestOutput
+                            n = len(output.outputs)
+                            assert output.prompt_token_ids is not None
+                            total_in_toks += len(output.prompt_token_ids) * n
+                            in_spd = total_in_toks / pbar.format_dict["elapsed"]
+                            total_out_toks += sum(len(stp.token_ids) for stp in output.outputs)
+                            out_spd = total_out_toks / pbar.format_dict["elapsed"]
+                            pbar.postfix = f"est. speed input: {in_spd:.2f} toks/s, output: {out_spd:.2f} toks/s"
+                            pbar.update(n)
+                        else:
+                            pbar.update(1)
+                        if pbar.n == num_requests:
+                            pbar.refresh()
+
+        if use_tqdm:
+            pbar.close()
+        # Sort the outputs by the int part of request ID which is in format of 'int-uuid'.
+        # This is necessary because some requests may be finished earlier than
+        # its previous requests.
+        return sorted(outputs, key=lambda x: int(x.request_id.split("-")[0]))

vllm_omni/entrypoints/openai/api_server.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import multiprocessing
 import multiprocessing.forkserver as forkserver
 import os
@@ -14,31 +16,48 @@ import vllm.envs as envs
 from fastapi import Depends, HTTPException, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 from starlette.datastructures import State
-from vllm.config import VllmConfig
+from starlette.routing import Route
 from vllm.engine.protocol import EngineClient
-from vllm.entrypoints.chat_utils import load_chat_template, resolve_hf_chat_template, resolve_mistral_chat_template
+from vllm.entrypoints.anthropic.serving_messages import AnthropicServingMessages
 from vllm.entrypoints.launcher import serve_http
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.openai.api_server import (
     base,
     build_app,
     load_log_config,
-    maybe_register_tokenizer_info_endpoint,
     router,
     setup_server,
-    validate_json_request,
 )
-from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse
-from vllm.entrypoints.openai.serving_models import BaseModelPath, LoRAModulePath, OpenAIServingModels
-from vllm.entrypoints.openai.tool_parsers import ToolParserManager
-
-# yapf conflicts with isort for this block
-# yapf: disable
-# yapf: enable
+from vllm.entrypoints.openai.orca_metrics import metrics_header
+from vllm.entrypoints.openai.protocol import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ErrorResponse,
+)
+from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
+from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
+from vllm.entrypoints.openai.serving_responses import OpenAIServingResponses
+from vllm.entrypoints.openai.serving_transcription import (
+    OpenAIServingTranscription,
+    OpenAIServingTranslation,
+)
+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.pooling.serving import OpenAIServingPooling
+from vllm.entrypoints.pooling.score.serving import ServingScores
+from vllm.entrypoints.serve.disagg.serving import ServingTokens
+from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
 from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer
-from vllm.entrypoints.utils import load_aware_call, with_cancellation
+from vllm.entrypoints.utils import (
+    load_aware_call,
+    process_chat_template,
+    process_lora_modules,
+    with_cancellation,
+)
 from vllm.logger import init_logger
-from vllm.tokenizers import MistralTokenizer
+from vllm.tasks import POOLING_TASKS
+from vllm.tool_parsers import ToolParserManager
 from vllm.utils.system_utils import decorate_logs
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
@@ -57,6 +76,29 @@ from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 
 logger = init_logger(__name__)
 
+ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format"
+
+
+def _remove_route_from_router(router_obj, path: str, methods: set[str] | None = None):
+    """Remove a route from the router by path and optionally by methods.
+
+    This is needed because vllm's api_server registers routes when imported,
+    and we need to override some routes (like /v1/chat/completions) with
+    omni-specific implementations.
+    """
+    routes_to_remove = []
+    for route in router_obj.routes:
+        if isinstance(route, Route) and route.path == path:
+            if methods is None or (hasattr(route, "methods") and route.methods & methods):
+                routes_to_remove.append(route)
+
+    for route in routes_to_remove:
+        router_obj.routes.remove(route)
+
+
+# Remove vllm's /v1/chat/completions route so we can register our own omni version
+_remove_route_from_router(router, "/v1/chat/completions", {"POST"})
+
 
 # Server entry points
 
@@ -88,6 +130,10 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
     if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
         ToolParserManager.import_tool_parser(args.tool_parser_plugin)
+    if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3:
+        from vllm.reasoning import ReasoningParserManager
+
+        ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin)
 
     # Load logging config for uvicorn if specified
     log_config = load_log_config(args.log_config_file)
@@ -98,11 +144,11 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
         args,
         client_config=client_config,
     ) as engine_client:
-        maybe_register_tokenizer_info_endpoint(args)
         app = build_app(args)
 
+        await omni_init_app_state(engine_client, app.state, args)
+
         vllm_config = await engine_client.get_vllm_config()
-        await omni_init_app_state(engine_client, vllm_config, app.state, args)
 
         # Check if pure diffusion mode (vllm_config will be None)
         is_pure_diffusion = vllm_config is None
@@ -233,7 +279,6 @@ async def build_async_omni_from_stage_config(
 
 async def omni_init_app_state(
     engine_client: EngineClient,
-    vllm_config: VllmConfig | None,
     state: State,
     args: Namespace,
 ) -> None:
@@ -246,10 +291,12 @@ async def omni_init_app_state(
 
     Args:
         engine_client: Engine client instance (AsyncOmni)
-        vllm_config: vLLM configuration object (may be None for pure diffusion)
         state: FastAPI application state object to initialize
         args: Parsed command-line arguments
     """
+    # Get vllm_config from engine_client (following 0.14.0 pattern)
+    vllm_config = await engine_client.get_vllm_config()
+
     # Detect if it's pure Diffusion mode (single stage and is Diffusion)
     is_pure_diffusion = False
     if hasattr(engine_client, "stage_configs") and engine_client.stage_configs:
@@ -273,6 +320,7 @@ async def omni_init_app_state(
     base_model_paths = [BaseModelPath(name=name, model_path=args.model) for name in served_model_names]
     state.engine_client = engine_client
     state.log_stats = not args.disable_log_stats
+    state.args = args
 
     # For omni models
     state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
@@ -302,32 +350,18 @@ async def omni_init_app_state(
             logger.warning("vllm_config is None, some features may not work correctly")
 
     state.vllm_config = vllm_config
-    if vllm_config is not None:
-        _model_config = vllm_config.model_config
 
-    resolved_chat_template = load_chat_template(args.chat_template)
-    if resolved_chat_template is not None and vllm_config is not None:
-        # Get the tokenizer to check official template
-        tokenizer = await engine_client.get_tokenizer()
+    # Get supported tasks
+    supported_tasks: set[str] = {"generate"}
+    if hasattr(engine_client, "get_supported_tasks"):
+        supported_tasks = set(await engine_client.get_supported_tasks())
+    logger.info("Supported tasks: %s", supported_tasks)
 
-        if tokenizer is not None:
-            if isinstance(tokenizer, MistralTokenizer):
-                # The warning is logged in resolve_mistral_chat_template.
-                resolved_chat_template = resolve_mistral_chat_template(chat_template=resolved_chat_template)
-            else:
-                hf_chat_template = resolve_hf_chat_template(
-                    tokenizer=tokenizer,
-                    chat_template=None,
-                    tools=None,
-                    model_config=vllm_config.model_config,
-                )
-
-                if hf_chat_template != resolved_chat_template:
-                    logger.warning(
-                        "Using supplied chat template: %s\nIt is different from official chat template '%s'. This discrepancy may lead to performance degradation.",  # noqa: E501
-                        resolved_chat_template,
-                        args.model,
-                    )
+    resolved_chat_template = await process_chat_template(
+        args.chat_template,
+        engine_client,
+        vllm_config.model_config if vllm_config is not None else None,
+    )
 
     if args.tool_server == "demo":
         tool_server: ToolServer | None = DemoToolServer()
@@ -340,23 +374,12 @@ async def omni_init_app_state(
         tool_server = None
 
     # Merge default_mm_loras into the static lora_modules
-    default_mm_loras = {}
-    if vllm_config is not None and vllm_config.lora_config is not None:
-        default_mm_loras = vllm_config.lora_config.default_mm_loras
-
-    lora_modules = args.lora_modules
-    if default_mm_loras:
-        default_mm_lora_paths = [
-            LoRAModulePath(
-                name=modality,
-                path=lora_path,
-            )
-            for modality, lora_path in default_mm_loras.items()
-        ]
-        if args.lora_modules is None:
-            lora_modules = default_mm_lora_paths
-        else:
-            lora_modules += default_mm_lora_paths
+    default_mm_loras = (
+        vllm_config.lora_config.default_mm_loras
+        if vllm_config is not None and vllm_config.lora_config is not None
+        else {}
+    )
+    lora_modules = process_lora_modules(args.lora_modules, default_mm_loras)
 
     # Ensure input_processor, io_processor, and model_config exist for OpenAIServingModels compatibility
     if (
@@ -415,24 +438,182 @@ async def omni_init_app_state(
         lora_modules=lora_modules,
     )
     await state.openai_serving_models.init_static_loras()
-    state.openai_serving_chat = OmniOpenAIServingChat(
+
+    state.openai_serving_responses = (
+        OpenAIServingResponses(
+            engine_client,
+            state.openai_serving_models,
+            request_logger=request_logger,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+            enable_auto_tools=args.enable_auto_tool_choice,
+            tool_parser=args.tool_call_parser,
+            tool_server=tool_server,
+            reasoning_parser=args.structured_outputs_config.reasoning_parser,
+            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
+    )
+    state.openai_serving_chat = (
+        OmniOpenAIServingChat(
+            engine_client,
+            state.openai_serving_models,
+            args.response_role,
+            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,
+            return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+            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,
+            reasoning_parser=args.structured_outputs_config.reasoning_parser,
+            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,
+            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_completion = (
+        OpenAIServingCompletion(
+            engine_client,
+            state.openai_serving_models,
+            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
+    )
+    state.openai_serving_pooling = (
+        OpenAIServingPooling(
+            engine_client,
+            state.openai_serving_models,
+            supported_tasks=supported_tasks,
+            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,
+            log_error_stack=args.log_error_stack,
+        )
+        if any(task in POOLING_TASKS for task in supported_tasks)
+        else None
+    )
+    state.openai_serving_embedding = (
+        OpenAIServingEmbedding(
+            engine_client,
+            state.openai_serving_models,
+            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,
+            log_error_stack=args.log_error_stack,
+        )
+        if "embed" in supported_tasks
+        else None
+    )
+    state.openai_serving_classification = (
+        ServingClassification(
+            engine_client,
+            state.openai_serving_models,
+            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,
+            log_error_stack=args.log_error_stack,
+        )
+        if "classify" in supported_tasks
+        else None
+    )
+    state.openai_serving_scores = (
+        ServingScores(
+            engine_client,
+            state.openai_serving_models,
+            request_logger=request_logger,
+            score_template=resolved_chat_template,
+            log_error_stack=args.log_error_stack,
+        )
+        if ("embed" in supported_tasks or "score" in supported_tasks)
+        else None
+    )
+    state.openai_serving_tokenization = OpenAIServingTokenization(
         engine_client,
         state.openai_serving_models,
-        args.response_role,
         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,
-        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
-        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,
-        reasoning_parser=args.structured_outputs_config.reasoning_parser,
-        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,
     )
+    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
+        else None
+    )
+    state.openai_serving_translation = (
+        OpenAIServingTranslation(
+            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
+        else None
+    )
+    state.anthropic_serving_messages = (
+        AnthropicServingMessages(
+            engine_client,
+            state.openai_serving_models,
+            args.response_role,
+            request_logger=request_logger,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+            enable_auto_tools=args.enable_auto_tool_choice,
+            tool_parser=args.tool_call_parser,
+            reasoning_parser=args.structured_outputs_config.reasoning_parser,
+            enable_prompt_tokens_details=args.enable_prompt_tokens_details,
+            enable_force_include_usage=args.enable_force_include_usage,
+        )
+        if "generate" in supported_tasks
+        else None
+    )
+    state.serving_tokens = (
+        ServingTokens(
+            engine_client,
+            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,
+        )
+        if "generate" in supported_tasks
+        else None
+    )
 
     state.openai_serving_speech = OmniOpenAIServingSpeech(
         engine_client, state.openai_serving_models, request_logger=request_logger
@@ -463,6 +644,7 @@ def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
 @with_cancellation
 @load_aware_call
 async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
+    metrics_header_format = raw_request.headers.get(ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, "")
     handler = Omnichat(raw_request)
     if handler is None:
         return base(raw_request).create_error_response(message="The model does not support Chat Completions API")
@@ -474,7 +656,8 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
 
     if isinstance(generator, ErrorResponse):
         return JSONResponse(
-            content=generator.model_dump(), status_code=generator.code if hasattr(generator, "code") else 400
+            content=generator.model_dump(),
+            status_code=generator.error.code if generator.error else 400,
         )
 
     elif isinstance(generator, ChatCompletionResponse):
@@ -490,18 +673,27 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
             try:
                 # Use serialize_as_any=True to bypass type checking
                 response_dict = generator.model_dump(mode="json", serialize_as_any=True, warnings="none")
-                return JSONResponse(content=response_dict)
+                return JSONResponse(
+                    content=response_dict,
+                    headers=metrics_header(metrics_header_format),
+                )
             except Exception:
                 # Fallback: convert to JSON string and parse back to avoid any serialization issues
                 try:
                     response_json = generator.model_dump_json(warnings="none", serialize_as_any=True)
                     response_dict = json_lib.loads(response_json)
-                    return JSONResponse(content=response_dict)
+                    return JSONResponse(
+                        content=response_dict,
+                        headers=metrics_header(metrics_header_format),
+                    )
                 except Exception:
                     # Last resort: regular dump with warnings suppressed
                     with warnings_module.catch_warnings():
                         warnings_module.filterwarnings("ignore", category=UserWarning)
-                        return JSONResponse(content=generator.model_dump(mode="json", warnings="none"))
+                        return JSONResponse(
+                            content=generator.model_dump(mode="json", warnings="none"),
+                            headers=metrics_header(metrics_header_format),
+                        )
 
     return StreamingResponse(content=generator, media_type="text/event-stream")
 

vllm_omni/entrypoints/openai/protocol/chat_completion.py

@@ -1,4 +1,6 @@
-from vllm.entrypoints.openai.protocol import ChatCompletionStreamResponse
+from vllm.entrypoints.openai.protocol import (
+    ChatCompletionStreamResponse,
+)
 
 
 class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse):

vllm_omni/entrypoints/openai/serving_chat.py

@@ -29,7 +29,10 @@ from vllm.entrypoints.chat_utils import (
     make_tool_call_id,
     resolve_chat_template_content_format,
 )
-from vllm.entrypoints.harmony_utils import get_streamable_parser_for_assistant, parse_chat_output
+from vllm.entrypoints.openai.parser.harmony_utils import (
+    get_streamable_parser_for_assistant,
+    parse_chat_output,
+)
 from vllm.entrypoints.openai.protocol import (
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionRequest,
@@ -41,29 +44,24 @@ from vllm.entrypoints.openai.protocol import (
     DeltaFunctionCall,
     DeltaMessage,
     DeltaToolCall,
+    ErrorInfo,
     ErrorResponse,
     FunctionCall,
     FunctionDefinition,
     PromptTokenUsageInfo,
     RequestResponseMetadata,
+    ResponsesRequest,
     ToolCall,
     UsageInfo,
 )
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 from vllm.entrypoints.openai.serving_engine import (
     ChatLikeRequest,
-    EngineTokensPrompt,
-    RequestPrompt,
-    ResponsesRequest,
-    TextTokensPrompt,
     clamp_prompt_logprobs,
-    is_list_of,
 )
-from vllm.entrypoints.openai.tool_parsers import ToolParser
-from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
 from vllm.entrypoints.utils import should_include_usage
-from vllm.inputs.data import PromptType
+from vllm.inputs.data import PromptType, TokensPrompt
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.outputs import RequestOutput
@@ -75,8 +73,10 @@ from vllm.tokenizers.mistral import (
     truncate_tool_call_ids,
     validate_request_params,
 )
+from vllm.tool_parsers import ToolParser
+from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.transformers_utils.tokenizer import AnyTokenizer
-from vllm.utils.collection_utils import as_list
+from vllm.utils.collection_utils import as_list, is_list_of
 
 from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
@@ -177,52 +177,68 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 truncate_tool_call_ids(request)
                 validate_request_params(request)
 
-            if (
-                request.tool_choice == "auto"
-                and not (self.enable_auto_tools and tool_parser is not None)
-                and not isinstance(tokenizer, MistralTokenizer)
-                and not self.use_harmony
+            # Check if tool parsing is unavailable (common condition)
+            tool_parsing_unavailable = (
+                tool_parser is None and not isinstance(tokenizer, MistralTokenizer) and not self.use_harmony
+            )
+
+            # Validate tool_choice when tool parsing is required but unavailable
+            if tool_parsing_unavailable and request.tool_choice not in (
+                None,
+                "none",
             ):
-                # for hf tokenizers, "auto" tools requires
-                # --enable-auto-tool-choice and --tool-call-parser
-                return self.create_error_response(
-                    '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set'
-                )
+                if request.tool_choice == "auto" and not self.enable_auto_tools:
+                    # for hf tokenizers, "auto" tools requires
+                    # --enable-auto-tool-choice and --tool-call-parser
+                    return self.create_error_response(
+                        '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set'
+                    )
+                elif request.tool_choice != "auto":
+                    # "required" or named tool requires tool parser
+                    return self.create_error_response(
+                        f'tool_choice="{request.tool_choice}" requires --tool-call-parser to be set'
+                    )
 
             if request.tools is None or (request.tool_choice == "none" and self.exclude_tools_when_tool_choice_none):
                 tool_dicts = None
             else:
                 tool_dicts = [tool.model_dump() for tool in request.tools]
 
-            # Common case.
-            request_chat_template = request.chat_template
-            chat_template_kwargs = request.chat_template_kwargs
-            if not self.trust_request_chat_template and (
-                request_chat_template is not None
-                or (chat_template_kwargs and chat_template_kwargs.get("chat_template") is not None)
-            ):
-                return self.create_error_response(
-                    "Chat template is passed with request, but --trust-request-chat-template is not set. "
-                    "Refused request with untrusted chat template."
+            if not self.use_harmony:
+                error_check_ret = self._validate_chat_template(
+                    request_chat_template=request.chat_template,
+                    chat_template_kwargs=request.chat_template_kwargs,
+                    trust_request_chat_template=self.trust_request_chat_template,
                 )
-            (
-                conversation,
-                request_prompts,
-                engine_prompts,
-            ) = await self._preprocess_chat(
-                request,
-                tokenizer,
-                request.messages,
-                chat_template=request_chat_template or self.chat_template,
-                chat_template_content_format=self.chat_template_content_format,
-                add_generation_prompt=request.add_generation_prompt,
-                continue_final_message=request.continue_final_message,
-                tool_dicts=tool_dicts,
-                documents=request.documents,
-                chat_template_kwargs=request.chat_template_kwargs,
-                tool_parser=tool_parser,
-                add_special_tokens=request.add_special_tokens,
-            )
+                if error_check_ret is not None:
+                    return error_check_ret
+
+                chat_template_kwargs = request.chat_template_kwargs or {}
+                chat_template_kwargs.update(reasoning_effort=request.reasoning_effort)
+
+                (
+                    conversation,
+                    request_prompts,
+                    engine_prompts,
+                ) = await self._preprocess_chat(
+                    request,
+                    tokenizer,
+                    request.messages,
+                    chat_template=request.chat_template or self.chat_template,
+                    chat_template_content_format=self.chat_template_content_format,
+                    add_generation_prompt=request.add_generation_prompt,
+                    continue_final_message=request.continue_final_message,
+                    tool_dicts=tool_dicts,
+                    documents=getattr(request, "documents", None),
+                    chat_template_kwargs=chat_template_kwargs,
+                    default_chat_template_kwargs=self.default_chat_template_kwargs,
+                    tool_parser=tool_parser,
+                    add_special_tokens=request.add_special_tokens,
+                )
+            else:
+                should_include_tools = tool_dicts is not None
+                conversation, engine_prompts = self._make_request_with_harmony(request, should_include_tools)
+                request_prompts = [engine_prompt.get("prompt_token_ids", []) for engine_prompt in engine_prompts]
 
         except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
             logger.exception("Error in preprocessing prompt inputs")
@@ -314,12 +330,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         tool_dicts: list[dict[str, Any]] | None = None,
         documents: list[dict[str, str]] | None = None,
         chat_template_kwargs: dict[str, Any] | None = None,
+        default_chat_template_kwargs: dict[str, Any] | None = None,
         tool_parser: Callable[[TokenizerLike], ToolParser] | None = None,
         add_special_tokens: bool = False,
     ) -> tuple[
         list[ConversationMessage],
-        Sequence[RequestPrompt],
-        list[EngineTokensPrompt],
+        Sequence[PromptType],
+        list[TokensPrompt],
     ]:
         model_config = self.model_config
 
@@ -333,11 +350,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         conversation, mm_data_future, mm_uuids = parse_chat_messages_futures(
             messages,
             model_config,
-            tokenizer,
             content_format=resolved_content_format,
             mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
         )
 
+        # Merge default_chat_template_kwargs with request-provided kwargs
+        # Request kwargs take precedence over defaults
+        merged_kwargs = self._prepare_extra_chat_template_kwargs(
+            chat_template_kwargs,
+            default_chat_template_kwargs,
+        )
+
         _chat_template_kwargs: dict[str, Any] = dict(
             chat_template=chat_template,
             add_generation_prompt=add_generation_prompt,
@@ -345,7 +368,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             tools=tool_dicts,
             documents=documents,
         )
-        _chat_template_kwargs.update(chat_template_kwargs or {})
+        _chat_template_kwargs.update(merged_kwargs)
 
         request_prompt: str | list[int]
 
@@ -388,7 +411,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 "Prompt has to be a string",
                 "when the tokenizer is not initialised",
             )
-            prompt_inputs = TextTokensPrompt(prompt=request_prompt, prompt_token_ids=[1])
+            prompt_inputs = TokensPrompt(prompt=request_prompt, prompt_token_ids=[1])
         elif isinstance(request_prompt, str):
             prompt_inputs = await self._tokenize_prompt_input_async(
                 request,
@@ -399,20 +422,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         else:
             # For MistralTokenizer
             assert is_list_of(request_prompt, int), "Prompt has to be either a string or a list of token ids"
-            prompt_inputs = TextTokensPrompt(
+            prompt_inputs = TokensPrompt(
                 prompt=tokenizer.decode(request_prompt),
                 prompt_token_ids=request_prompt,
             )
 
-        engine_prompt = EngineTokensPrompt(prompt_token_ids=prompt_inputs["prompt_token_ids"])
+        engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["prompt_token_ids"])
         if mm_data is not None:
             engine_prompt["multi_modal_data"] = mm_data
 
         if mm_uuids is not None:
             engine_prompt["multi_modal_uuids"] = mm_uuids
 
-        if request.mm_processor_kwargs is not None:
-            engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
+        mm_processor_kwargs = getattr(request, "mm_processor_kwargs", None)
+        if mm_processor_kwargs is not None:
+            engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs
 
         if hasattr(request, "cache_salt") and request.cache_salt is not None:
             engine_prompt["cache_salt"] = request.cache_salt
@@ -513,7 +537,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     def _log_inputs(
         self,
         request_id: str,
-        inputs: RequestPrompt | PromptType,
+        inputs: PromptType,
         params_list: list[SamplingParams] | None,
         lora_request: LoRARequest | None,
     ) -> None:
@@ -599,9 +623,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         try:
             if self.reasoning_parser:
+                chat_template_kwargs = self._prepare_extra_chat_template_kwargs(
+                    request.chat_template_kwargs,
+                    self.default_chat_template_kwargs,
+                )
                 reasoning_parser = self.reasoning_parser(
                     tokenizer,
-                    chat_template_kwargs=request.chat_template_kwargs,  # type: ignore
+                    chat_template_kwargs=chat_template_kwargs,  # type: ignore
                 )
         except RuntimeError as e:
             logger.exception("Error in reasoning parser creation.")
@@ -1052,10 +1080,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         # wasn't ready to send a token, then
                         #   get the next token without streaming a chunk
                         if delta_message is None:
-                            if output.finish_reason is None:
+                            if output.finish_reason is None and not request.return_token_ids:
                                 continue
-                            else:
-                                delta_message = DeltaMessage()
+                            delta_message = DeltaMessage()
 
                         # Log streaming delta if output logging is enabled
                         if self.enable_log_outputs and self.request_logger:
@@ -1438,13 +1465,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             if self.reasoning_parser:
                 try:
-                    reasoning_parser = self.reasoning_parser(tokenizer)
+                    chat_template_kwargs = self._prepare_extra_chat_template_kwargs(
+                        request.chat_template_kwargs,
+                        self.default_chat_template_kwargs,
+                    )
+                    reasoning_parser = self.reasoning_parser(
+                        tokenizer,
+                        chat_template_kwargs=chat_template_kwargs,  # type: ignore
+                    )
                 except RuntimeError as e:
                     logger.exception("Error in reasoning parser creation.")
                     return self.create_error_response(str(e))
                 # If the reasoning parser is enabled,
                 # tool calls are extracted exclusively from the content.
-                reasoning_content, content = reasoning_parser.extract_reasoning_content(output.text, request=request)
+                reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request)
                 if not request.include_reasoning:
                     reasoning_content = None
             else:
@@ -2049,7 +2083,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     ) -> ErrorResponse:
         """Create an error response following OpenAI error format."""
         return ErrorResponse(
-            message=message,
-            type=err_type,
-            code=status_code,
+            error=ErrorInfo(
+                message=message,
+                type=err_type,
+                code=status_code,
+            )
         )

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -107,6 +107,7 @@ class Qwen2_5OmniForConditionalGeneration(
             )
             if t2w_token_end_id:
                 self.model.set_suppress_start_id(t2w_token_end_id + 1)
+            self.requires_raw_input_tokens = True
 
         elif self.model_stage == "code2wav":
             self.thinker = None
@@ -125,6 +126,7 @@ class Qwen2_5OmniForConditionalGeneration(
             self._token2wav_conds: dict[str, torch.Tensor] = {}
             self._token2wav_ref_mels: dict[str, torch.Tensor] = {}
             self.model = self.token2wav
+            self.requires_raw_input_tokens = True
         else:
             raise ValueError("Invalid model stage")
 
@@ -248,14 +250,27 @@ class Qwen2_5OmniForConditionalGeneration(
             if is_npu():
                 # TODO: remove this hack when NPU supports batched inputs properly
                 thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids
-                thinker_positions = positions[0] if positions.ndim > 1 else positions
+                # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it
+                if positions.ndim == 2 and positions.shape[0] == 3:
+                    thinker_positions = positions  # MRoPE positions, keep as is
+                else:
+                    thinker_positions = positions[0] if positions.ndim > 1 else positions
                 thinker_inputs_embeds = (
                     inputs_embeds[0] if inputs_embeds is not None and added_batch_dim else inputs_embeds
                 )
             else:
-                thinker_input_ids = input_ids
-                thinker_positions = positions[0]
-                thinker_inputs_embeds = inputs_embeds
+                # Squeeze back if we added batch dim earlier
+                thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids
+                # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it
+                if positions.ndim == 2 and positions.shape[0] == 3:
+                    thinker_positions = positions  # MRoPE positions, keep as is
+                elif added_batch_dim:
+                    thinker_positions = positions[0]
+                else:
+                    thinker_positions = positions
+                thinker_inputs_embeds = (
+                    inputs_embeds[0] if inputs_embeds is not None and added_batch_dim else inputs_embeds
+                )
 
             # Run thinker
             thinker_output = self.thinker(
@@ -288,10 +303,16 @@ class Qwen2_5OmniForConditionalGeneration(
             if not hasattr(self, "voice_type"):
                 self.voice_type = voice_type
 
+            # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it
+            if positions.ndim == 2 and positions.shape[0] == 3:
+                talker_positions = positions  # MRoPE positions, keep as is
+            else:
+                talker_positions = positions[0]
+
             with torch.inference_mode():
                 talker_hidden = self.talker(
                     input_ids=input_ids,
-                    positions=positions[0],
+                    positions=talker_positions,
                     inputs_embeds=inputs_embeds,
                 )
 

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -1,4 +1,4 @@
-"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker (v0.12) with minimal overrides."""
+"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker (v0.14) with minimal overrides."""
 
 from collections.abc import Iterable
 from typing import Any
@@ -12,6 +12,7 @@ from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
     Qwen2_5OmniAudioEncoder,
 )
 from vllm.config import VllmConfig
+from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
@@ -26,8 +27,6 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
     Qwen2_5OmniThinkerProcessingInfo,
-    get_llm_pos_ids_for_vision,
-    split_list_into_ranges,
 )
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase,
@@ -46,7 +45,9 @@ from vllm.model_executor.models.utils import (
     WeightsMapper,
     init_vllm_registered_model,
     maybe_prefix,
+    split_list_into_ranges,
 )
+from vllm.model_executor.models.vision import get_llm_pos_ids_for_vision
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
@@ -166,6 +167,43 @@ class Qwen2_5OmniConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixi
                 video_grid_thw=video_grid_thw,
             )
 
+    def _process_image_input(self, image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:
+        if image_input["type"] == "image_embeds":
+            return image_input["image_embeds"].type(self.visual.dtype)
+
+        grid_thw = image_input["image_grid_thw"]
+        assert grid_thw.ndim == 2
+
+        pixel_values = image_input["pixel_values"].type(self.visual.dtype)
+        with set_forward_context(None, self.vllm_config):
+            image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
+        # Split concatenated embeddings for each image item.
+        merge_size = self.visual.spatial_merge_size
+        sizes = grid_thw.prod(-1) // merge_size // merge_size
+
+        return image_embeds.split(sizes.tolist())
+
+    def _process_video_input(
+        self,
+        video_input: Qwen2_5_VLVideoInputs,
+        video_hashes: list[str] = None,
+        cached_video_embeds: torch.Tensor = None,
+    ) -> torch.Tensor:
+        if video_input["type"] == "video_embeds":
+            return video_input["video_embeds"].type(self.visual.dtype)
+
+        grid_thw = video_input["video_grid_thw"]
+        assert grid_thw.ndim == 2
+
+        pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype)
+        with set_forward_context(None, self.vllm_config):
+            video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
+        # Split concatenated embeddings for each video item.
+        merge_size = self.visual.spatial_merge_size
+        sizes = grid_thw.prod(-1) // merge_size // merge_size
+
+        return video_embeds.split(sizes.tolist())
+
 
 @MULTIMODAL_REGISTRY.register_processor(
     Qwen2_5OmniThinkerMultiModalProcessor,
@@ -180,8 +218,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
     SupportsMRoPE,
     Qwen2_5OmniConditionalGenerationMixin,
 ):
-    merge_by_field_config = True
-
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "thinker.lm_head.": "language_model.lm_head.",
@@ -250,6 +286,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                 norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
                 quant_config=quant_config,
                 prefix=maybe_prefix(prefix, "visual"),
+                multimodal_config=multimodal_config,
             )
         else:
             self.visual = None

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py

@@ -3,9 +3,6 @@ from collections.abc import Iterable
 import torch
 from torch import nn
 from transformers import Qwen2Config
-from vllm.attention.backends.abstract import (
-    AttentionType,
-)
 from vllm.attention.layer import Attention
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -15,7 +12,6 @@ from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.pooler import Pooler, PoolingType
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding
@@ -24,15 +20,14 @@ from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
-    WeightsMapper,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
 )
 from vllm.sequence import IntermediateTensors
-from vllm.v1.outputs import PoolerOutput, SamplerOutput
-from vllm.v1.pool.metadata import PoolingMetadata
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
@@ -130,7 +125,6 @@ class Qwen2Attention(nn.Module):
 
         self.rotary_pos_emb = get_rope(
             head_size=self.head_dim,
-            rotary_dim=self.head_dim,
             max_position=max_position,
             is_neox_style=True,
             rope_parameters={
@@ -460,66 +454,3 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
             skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
         )
         return loader.load_weights(weights)
-
-
-class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
-    packed_modules_mapping = {
-        "qkv_proj": [
-            "q_proj",
-            "k_proj",
-            "v_proj",
-        ],
-        "gate_up_proj": [
-            "gate_proj",
-            "up_proj",
-        ],
-    }
-
-    hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-        config = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-        lora_config = vllm_config.lora_config
-        pooler_config = vllm_config.model_config.pooler_config
-
-        self.config = config
-        self.lora_config = lora_config
-
-        self.quant_config = quant_config
-        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
-
-        # TODO: Replace this model class with as_embedding_model(
-        # Qwen2ForCausalLM) after changing the default pooling method
-        if pooler_config.pooling_type is None:
-            logger.warning(
-                "This embedding model will default to last-token pooling in "
-                "an upcoming version. To avoid breaking changes, you should "
-                'pass `--override-pooler-config \'{"pooling_type": "MEAN"}\'`'
-                " explicitly."
-            )
-
-        self._pooler = Pooler.from_config_with_defaults(
-            pooler_config, pooling_type=PoolingType.MEAN, normalize=True, softmax=False
-        )
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None = None,
-    ) -> torch.Tensor:
-        return self.model(input_ids, positions, intermediate_tensors)
-
-    def pooler(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> PoolerOutput | None:
-        return self._pooler(hidden_states, pooling_metadata)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        weights = self.hf_to_vllm_mapper.apply(weights)
-        weights = ((name, data) for name, data in weights if not name.startswith("lm_head."))
-        self.model.load_weights(weights)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -154,6 +154,7 @@ class Qwen3OmniMoeForConditionalGeneration(
 
             # for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights
             self._init_special_tokens_embeddings()
+            self.requires_raw_input_tokens = True
 
         elif self.model_stage == "code2wav":
             self.thinker = None
@@ -168,6 +169,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 architectures=["Qwen3OmniMoeCode2Wav"],
             )
             self.model = self.code2wav
+            self.requires_raw_input_tokens = True
         else:
             raise ValueError(
                 f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'talker', 'code2wav'"
@@ -354,18 +356,24 @@ class Qwen3OmniMoeForConditionalGeneration(
         elif self.model_stage == "code2wav":
             # Extract codec codes from input
             codes = []
-            if input_ids is not None:
+            if input_ids.shape[0] % 16 == 0:
                 codes.append(input_ids.reshape(1, 16, -1))
-
             else:
-                # for profile, we use max length from inputs_embeds
-                codes.append(
-                    torch.zeros(
-                        (1, 16, inputs_embeds.shape[1]),
-                        dtype=torch.long,
-                        device=inputs_embeds.device,
-                    )
+                logger.warning(
+                    (
+                        "Input_ids length: %s is not divisible by 16, padding "
+                        "with zeros. This should only happen in warm up."
+                    ),
+                    input_ids.shape[0],
                 )
+                input_ids_flatten = input_ids.reshape(-1)
+                input_ids_flatten = torch.cat(
+                    [
+                        input_ids_flatten,
+                        torch.zeros(16 - input_ids.shape[0] % 16, dtype=torch.long, device=input_ids.device),
+                    ]
+                )
+                codes.append(input_ids_flatten.reshape(1, 16, -1))
 
             # Generate audio from codec codes
             audio_tensors = []

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -37,9 +37,6 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeConfig,
     Qwen3OmniMoeThinkerConfig,
 )
-from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
-    Qwen3OmniMoeAudioEncoder,
-)
 from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (
     Qwen3OmniMoeProcessor,
 )
@@ -56,6 +53,7 @@ from vllm.model_executor.models.interfaces import (
     SupportsMultiModal,
     SupportsPP,
 )
+from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
@@ -69,6 +67,7 @@ from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeModel as _Qwen3MoeLLMModel
 from vllm.model_executor.models.qwen3_omni_moe_thinker import (
     Qwen3Omni_VisionTransformer,
+    Qwen3OmniMoeAudioEncoder,
     _get_feat_extract_output_lengths,
 )
 from vllm.model_executor.models.utils import (
@@ -248,11 +247,36 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
             # https://github.com/huggingface/transformers/pull/41473
             mm_kwargs = dict(mm_kwargs)
             tok_kwargs = dict(tok_kwargs)
+            mm_kwargs["audio_kwargs"] = dict(mm_kwargs.get("audio_kwargs") or {})
+            mm_kwargs["text_kwargs"] = dict(mm_kwargs.get("text_kwargs") or {})
             if Version(TRANSFORMERS_VERSION) < Version("4.58.0"):
+                # Extract audio_sample_rate before restructuring
+                audio_sample_rate = mm_kwargs.pop("audio_sample_rate", None)
+
                 # move truncation to audio_kwargs level to avoid conflict
                 # with tok_kwargs
-                mm_kwargs["audio_kwargs"] = {"truncation": mm_kwargs.pop("truncation", False)}
-                mm_kwargs["text_kwargs"] = {"truncation": tok_kwargs.pop("truncation", False)}
+                mm_kwargs["audio_kwargs"].setdefault("truncation", mm_kwargs.pop("truncation", False))
+                mm_kwargs["text_kwargs"].setdefault("truncation", tok_kwargs.pop("truncation", False))
+
+                # Validate and conditionally pass audio_sample_rate
+                # WhisperFeatureExtractor has a fixed sampling rate, and vLLM's
+                # audio loader already resamples audio to the target rate.
+                # Only pass the value if it matches to avoid unexpected behavior.
+                if audio_sample_rate is not None:
+                    expected_sr = feature_extractor.sampling_rate
+                    if audio_sample_rate != expected_sr:
+                        logger.warning(
+                            "[%s] audio_sample_rate mismatch: user provided %dHz "
+                            "but model expects %dHz. Ignoring user value. "
+                            "vLLM's audio loader already resampled to %dHz.",
+                            self.__class__.__name__,
+                            audio_sample_rate,
+                            expected_sr,
+                            expected_sr,
+                        )
+                    else:
+                        # Sample rate matches, safe to pass
+                        mm_kwargs["audio_kwargs"]["audio_sample_rate"] = audio_sample_rate
 
         hf_inputs = super()._call_hf_processor(
             prompt=prompt,
@@ -398,11 +422,11 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
         if audio_feature_lengths is None and feature_attention_mask is None:
             audio_output_lengths = []
         elif audio_feature_lengths is not None:
-            _, audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths)
+            audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths)
             audio_output_lengths = audio_output_lens.tolist()
         elif feature_attention_mask is not None:
             assert isinstance(feature_attention_mask, torch.Tensor)
-            _, audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1))
+            audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1))
             audio_output_lengths = audio_output_lens.tolist()
 
         # number of audios read from video.
@@ -569,59 +593,21 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
 
 class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixin):
-    def _parse_and_validate_audio_input(self, **kwargs: object) -> Qwen2_5OmniAudioFeatureInputs | None:
-        input_audio_features = kwargs.pop("input_audio_features", None)
-        audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
-        feature_attention_mask = kwargs.pop("feature_attention_mask", None)
-        if input_audio_features is None:
-            return None
-        if (
-            input_audio_features is not None
-            and isinstance(input_audio_features, torch.Tensor)
-            and input_audio_features.ndim == 3
-        ):
-            # (batch_size, feature_dim, chunk_size) -> (feature_dim, batch_size * chunk_size)
-            input_audio_features = input_audio_features.permute(1, 0, 2).flatten(1)
-        elif input_audio_features is not None and isinstance(input_audio_features, list):
-            input_audio_features = torch.cat(input_audio_features, dim=-1)
-        if (
-            audio_feature_lengths is not None
-            and isinstance(audio_feature_lengths, torch.Tensor)
-            and audio_feature_lengths.ndim == 2
-        ):
-            audio_feature_lengths = audio_feature_lengths.reshape(-1)
-        elif audio_feature_lengths is not None and isinstance(audio_feature_lengths, list):
-            audio_feature_lengths = torch.cat(audio_feature_lengths, dim=-1)
-        if (
-            feature_attention_mask is not None
-            and isinstance(feature_attention_mask, torch.Tensor)
-            and feature_attention_mask.ndim == 3
-        ):
-            feature_attention_mask = feature_attention_mask.reshape(-1, feature_attention_mask.shape[-1])
-        elif feature_attention_mask is not None and isinstance(feature_attention_mask, list):
-            for i in range(len(feature_attention_mask)):
-                feature_attention_mask[i] = feature_attention_mask[i].reshape(-1)
-        return Qwen2_5OmniAudioFeatureInputs(
-            type="audio_features",
-            input_features=input_audio_features,
-            audio_feature_lengths=audio_feature_lengths,
-            feature_attention_mask=feature_attention_mask,
-        )
-
     def _process_audio_input(
         self,
         audio_input: Qwen2_5OmniAudioFeatureInputs,
         audio_hashes: list[str] | None = None,
         cached_audio_features: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, ...]:
         input_features = audio_input["input_features"]
         audio_feature_lengths = audio_input["audio_feature_lengths"]
 
-        audio_feat_lengths, audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)
+        audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)
 
         audio_outputs = self.audio_tower(
             input_features.to(self.audio_tower.dtype),
             feature_lens=audio_feature_lengths,
+            aftercnn_lens=audio_output_lengths,
         )
         audio_features = audio_outputs.last_hidden_state
         return audio_features.split(audio_output_lengths.tolist())
@@ -639,8 +625,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
     SupportsMRoPE,
     Qwen3OmniMoeConditionalGenerationMixin,
 ):
-    merge_by_field_config = True
-
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "thinker.lm_head.": "language_model.lm_head.",
@@ -649,6 +633,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         }
     )
 
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": [
+            "gate_proj",
+            "up_proj",
+        ],
+    }
+
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> str | None:
         if modality.startswith("image"):
@@ -669,28 +665,16 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         self.config = thinker_config
         self.multimodal_config = multimodal_config
 
-        # force "use_flash_attention_2=True" to audio tower to align
-        # the results.
-        if flash_attn is not None:
-            audio_config = thinker_config.audio_config
-            audio_config._attn_implementation_autoset = True
-            audio_config._attn_implementation = "flash_attention_2"
-        else:
-            logger.warning(
-                "flash_attn is not available, the model may not yield the "
-                "exactly same result as the transformers implementation "
-                "in the audio tower part."
-            )
+        self.audio_tower = Qwen3OmniMoeAudioEncoder(
+            thinker_config.audio_config,
+        )
 
-        self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config)
-
-        attn_backend_override = multimodal_config.mm_encoder_attn_backend if multimodal_config is not None else None
         self.visual = Qwen3Omni_VisionTransformer(
             vision_config=thinker_config.vision_config,
             norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
             quant_config=quant_config,
             prefix=maybe_prefix(prefix, "visual"),
-            attn_backend_override=attn_backend_override,
+            multimodal_config=multimodal_config,
         )
         self.quant_config = quant_config
 
@@ -811,8 +795,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             return inputs_embeds
 
         deepstack_input_embeds = None
-        # TODO (ywang96): support overlapping modalitiy embeddings so that
-        # `use_audio_in_video` will work on V1.
         # split the feat dim to obtain multi-scale visual feature
         has_vision_embeddings = [
             embeddings.shape[-1] != self.config.text_config.hidden_size for embeddings in multimodal_embeddings
@@ -1006,7 +988,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
                 bos_len = 1
                 llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
                 st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                _, audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx])
+                audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx])
                 llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
                 llm_pos_ids_list.append(llm_pos_ids)
                 st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
@@ -1078,7 +1060,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
                 llm_pos_ids_list.append(bos_block)
                 llm_pos_ids_list.append(bos_block)
                 st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                _, audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx])
+                audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx])
                 audio_llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
                 grid_t = video_grid_thw[video_idx][0]
                 grid_hs = video_grid_thw[:, 1]
@@ -1122,3 +1104,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
 
         mrope_position_delta = llm_positions.max() + 1 - seq_len
         return llm_positions, mrope_position_delta
+
+    def get_mm_mapping(self) -> MultiModelKeys:
+        """
+        Get the module prefix in multimodal models
+        """
+        return MultiModelKeys.from_string_field(
+            language_model="language_model",
+            connector="visual.merger",
+            tower_model=["visual.", "audio_tower."],
+        )

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -76,6 +76,7 @@ stage_args:
       trust_remote_code: true
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
+      async_scheduling: false
       engine_output_type: audio
     engine_input_source: [1]
     final_output: true

vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml

@@ -16,7 +16,7 @@ stage_args:
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
-      enforce_eager: false
+      enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
@@ -46,8 +46,8 @@ stage_args:
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.3
-      enforce_eager: false
+      gpu_memory_utilization: 0.35
+      enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output codec codes for code2wav
       # tensor_parallel_size: 2
@@ -80,6 +80,7 @@ stage_args:
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
+      async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: audio  # Final output: audio waveform
       gpu_memory_utilization: 0.1

vllm_omni/worker/gpu_ar_model_runner.py

@@ -12,10 +12,15 @@ from typing import Any, NamedTuple
 import numpy as np
 import torch
 from vllm.config import CUDAGraphMode
+from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
+from vllm.distributed.kv_transfer import get_kv_transfer_group
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+)
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
-from vllm.v1.outputs import AsyncModelRunnerOutput
+from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.structured_output.utils import apply_grammar_bitmask
 from vllm.v1.utils import record_function_or_nullcontext
@@ -27,6 +32,7 @@ from vllm.v1.worker.gpu_model_runner import (
     get_tp_group,
     has_kv_transfer_group,
 )
+from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
 from vllm_omni.outputs import OmniModelRunnerOutput
@@ -44,6 +50,7 @@ class ExecuteModelState(NamedTuple):
     sample_hidden_states: torch.Tensor
     aux_hidden_states: list[torch.Tensor] | None
     ec_connector_output: Any
+    cudagraph_stats: Any
     multimodal_outputs: Any
 
 
@@ -82,67 +89,138 @@ class GPUARModelRunner(OmniGPUModelRunner):
         scheduler_output: SchedulerOutput,
         intermediate_tensors: IntermediateTensors | None = None,
     ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None:
-        with record_function_or_nullcontext("Preprocess"):
-            with self.synchronize_input_prep():
-                self._update_states(scheduler_output)
-                self._decode_and_store_request_payloads(scheduler_output)
+        if self.execute_model_state is not None:
+            raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
-                if not scheduler_output.total_num_scheduled_tokens:
-                    if not has_kv_transfer_group():
-                        return EMPTY_MODEL_RUNNER_OUTPUT
-                    return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
-                if self.cache_config.kv_sharing_fast_prefill:
-                    assert not self.input_batch.num_prompt_logprobs, (
-                        "--kv-sharing-fast-prefill produces incorrect "
-                        "logprobs for prompt tokens, tokens, please disable "
-                        "it when the requests need prompt logprobs"
-                    )
+        if self.vllm_config.model_config.enable_return_routed_experts:
+            capturer = RoutedExpertsCapturer.get_instance()
+            if capturer is not None:
+                capturer.clear_buffer()  # noqa
+            else:
+                logger.error("RoutedExpertsCapturer not initialized.")
 
-                num_reqs = self.input_batch.num_reqs
-                req_ids = self.input_batch.req_ids
-                tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
-                num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
-                max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
-                num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
+        if scheduler_output.preempted_req_ids and has_kv_transfer_group():
+            get_kv_transfer_group().handle_preemptions(scheduler_output.preempted_req_ids)
 
-                logits_indices, spec_decode_metadata = self._prepare_inputs(
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        with (
+            record_function_or_nullcontext("gpu_model_runner: preprocess"),
+            self.synchronize_input_prep(),
+        ):
+            # Update persistent batch states.
+            self._update_states(scheduler_output)
+
+            if has_ec_transfer() and get_ec_transfer().is_producer:
+                with self.maybe_get_ec_connector_output(
                     scheduler_output,
+                    encoder_cache=self.encoder_cache,
+                ) as ec_connector_output:
+                    self._execute_mm_encoder(scheduler_output)
+                    return make_empty_encoder_model_runner_output(scheduler_output)
+
+            if not num_scheduled_tokens:
+                if (
+                    self.parallel_config.distributed_executor_backend == "external_launcher"
+                    and self.parallel_config.data_parallel_size > 1
+                ):
+                    # this is a corner case when both external launcher
+                    # and DP are enabled, num_scheduled_tokens could be
+                    # 0, and has_unfinished_requests in the outer loop
+                    # returns True. before returning early here we call
+                    # dummy run to ensure coordinate_batch_across_dp
+                    # is called into to avoid out of sync issues.
+                    self._dummy_run(1)
+                if not has_kv_transfer_group():
+                    # Return empty ModelRunnerOutput if no work to do.
+                    return EMPTY_MODEL_RUNNER_OUTPUT
+                return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
+
+            if self.cache_config.kv_sharing_fast_prefill:
+                assert not self.num_prompt_logprobs, (
+                    "--kv-sharing-fast-prefill produces incorrect "
+                    "logprobs for prompt tokens, tokens, please disable "
+                    "it when the requests need prompt logprobs"
+                )
+
+            num_reqs = self.input_batch.num_reqs
+            req_ids = self.input_batch.req_ids
+            tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
+            num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
+            max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
+            num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
+
+            logits_indices, spec_decode_metadata = self._prepare_inputs(
+                scheduler_output,
+                num_scheduled_tokens_np,
+            )
+
+            cascade_attn_prefix_lens = None
+            # Disable cascade attention when using microbatching (DBO)
+            if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
+                # Pre-compute cascade attention prefix lengths
+                cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
                     num_scheduled_tokens_np,
+                    self.input_batch.num_computed_tokens_cpu[:num_reqs],
+                    scheduler_output.num_common_prefix_blocks,
                 )
 
-                (
-                    cudagraph_mode,
-                    batch_desc,
-                    ubatch_slices,
-                    num_tokens_across_dp,
-                ) = self._determine_batch_execution_and_padding(
-                    num_tokens=num_tokens_unpadded,
-                    num_reqs=num_reqs,
-                    num_scheduled_tokens_np=num_scheduled_tokens_np,
-                    max_num_scheduled_tokens=max_num_scheduled_tokens,
-                    use_cascade_attn=False,
-                )
+            (
+                cudagraph_mode,
+                batch_desc,
+                should_ubatch,
+                num_tokens_across_dp,
+                cudagraph_stats,
+            ) = self._determine_batch_execution_and_padding(
+                num_tokens=num_tokens_unpadded,
+                num_reqs=num_reqs,
+                num_scheduled_tokens_np=num_scheduled_tokens_np,
+                max_num_scheduled_tokens=max_num_scheduled_tokens,
+                use_cascade_attn=cascade_attn_prefix_lens is not None,
+                num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
+            )
 
-                num_tokens_padded = batch_desc.num_tokens
-                num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
-                use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
-                pad_attn = cudagraph_mode == CUDAGraphMode.FULL
+            logger.debug(
+                "Running batch with cudagraph_mode: %s, batch_descriptor: %s, "
+                "should_ubatch: %s, num_tokens_across_dp: %s",
+                cudagraph_mode,
+                batch_desc,
+                should_ubatch,
+                num_tokens_across_dp,
+            )
 
-                (
-                    attn_metadata,
-                    spec_decode_common_attn_metadata,
-                ) = self._build_attention_metadata(
-                    num_tokens=num_tokens_unpadded,
-                    num_tokens_padded=num_tokens_padded if pad_attn else None,
-                    num_reqs=num_reqs,
-                    num_reqs_padded=num_reqs_padded if pad_attn else None,
-                    max_query_len=max_num_scheduled_tokens,
-                    ubatch_slices=ubatch_slices,
-                    logits_indices=logits_indices,
-                    use_spec_decode=use_spec_decode,
-                    num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
-                    cascade_attn_prefix_lens=None,
-                )
+            num_tokens_padded = batch_desc.num_tokens
+            num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
+            ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
+                should_ubatch,
+                num_scheduled_tokens_np,
+                num_tokens_padded,
+                num_reqs_padded,
+                self.parallel_config.num_ubatches,
+            )
+
+            logger.debug(
+                "ubatch_slices: %s, ubatch_slices_padded: %s",
+                ubatch_slices,
+                ubatch_slices_padded,
+            )
+
+            pad_attn = cudagraph_mode == CUDAGraphMode.FULL
+
+            use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
+            ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
+
+            attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
+                num_tokens=num_tokens_unpadded,
+                num_tokens_padded=num_tokens_padded if pad_attn else None,
+                num_reqs=num_reqs,
+                num_reqs_padded=num_reqs_padded if pad_attn else None,
+                max_query_len=max_num_scheduled_tokens,
+                ubatch_slices=ubatch_slices_attn,
+                logits_indices=logits_indices,
+                use_spec_decode=use_spec_decode,
+                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
+                cascade_attn_prefix_lens=cascade_attn_prefix_lens,
+            )
 
             (
                 input_ids,
@@ -151,16 +229,18 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 intermediate_tensors,
                 model_kwargs,
                 ec_connector_output,
-            ) = self._preprocess(
-                scheduler_output,
-                num_tokens_padded,
-                intermediate_tensors,
-            )
+            ) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors)
 
+        # Set cudagraph mode to none if calc_kv_scales is true.
+        # KV scales calculation involves dynamic operations that are incompatible
+        # with CUDA graph capture.
         if self.calculate_kv_scales:
             cudagraph_mode = CUDAGraphMode.NONE
+            # Mark KV scales as calculated after the first forward pass
             self.calculate_kv_scales = False
 
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
         with (
             set_forward_context(
                 attn_metadata,
@@ -169,9 +249,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 num_tokens_across_dp=num_tokens_across_dp,
                 cudagraph_runtime_mode=cudagraph_mode,
                 batch_descriptor=batch_desc,
-                ubatch_slices=ubatch_slices,
+                ubatch_slices=ubatch_slices_padded,
             ),
-            record_function_or_nullcontext("Forward"),
+            record_function_or_nullcontext("gpu_model_runner: forward"),
             self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
         ):
             model_output = self._model_forward(
@@ -208,27 +288,32 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 logger.debug("[AR] execute_model: multimodal_outputs is None")
 
             if not self.broadcast_pp_output:
+                # Common case.
                 if not get_pp_group().is_last_rank:
+                    # Return the intermediate tensors.
                     assert isinstance(hidden_states, IntermediateTensors)
                     hidden_states.kv_connector_output = kv_connector_output
+                    self.kv_connector_output = kv_connector_output
                     return hidden_states
 
                 if self.is_pooling_model:
-                    output = self._pool(
+                    # Return the pooling output.
+                    return self._pool(
                         hidden_states,
-                        num_tokens_padded,
+                        num_scheduled_tokens,
                         num_scheduled_tokens_np,
+                        kv_connector_output,
                     )
-                    output.kv_connector_output = kv_connector_output
-                    return output
 
                 sample_hidden_states = hidden_states[logits_indices]
                 logits = self.model.compute_logits(
                     sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
                 )
             else:
+                # Rare case.
                 assert not self.is_pooling_model
 
+                sample_hidden_states = hidden_states[logits_indices]
                 if not get_pp_group().is_last_rank:
                     all_gather_tensors = {
                         "residual": not is_residual_scattered_for_sp(self.vllm_config, num_tokens_padded)
@@ -240,7 +325,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     )
                     logits = None
                 else:
-                    sample_hidden_states = hidden_states[logits_indices]
                     logits = self.model.compute_logits(
                         sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
                     )
@@ -264,6 +348,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             sample_hidden_states,
             aux_hidden_states,
             ec_connector_output,
+            cudagraph_stats,
             multimodal_outputs,
         )
         self.kv_connector_output = kv_connector_output
@@ -278,14 +363,20 @@ class GPUARModelRunner(OmniGPUModelRunner):
         self.kv_connector_output = None
 
         if self.execute_model_state is None:
+            # Nothing to do (PP non-final rank case), output isn't used.
             if not kv_connector_output:
                 return None  # type: ignore[return-value]
+
+            # In case of PP with kv transfer, we need to pass through the
+            # kv_connector_output
             if kv_connector_output.is_empty():
                 return EMPTY_MODEL_RUNNER_OUTPUT
+
             output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
             output.kv_connector_output = kv_connector_output
             return output
 
+        # Unpack ephemeral state.
         (
             scheduler_output,
             logits,
@@ -295,16 +386,20 @@ class GPUARModelRunner(OmniGPUModelRunner):
             sample_hidden_states,
             aux_hidden_states,
             ec_connector_output,
+            cudagraph_stats,
             multimodal_outputs,
         ) = self.execute_model_state
         self.execute_model_state = None
 
+        # Apply structured output bitmasks if present.
         if grammar_output is not None:
             apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits)
 
         with record_function_or_nullcontext("gpu_model_runner: sample"):
             sampler_output = self._sample(logits, spec_decode_metadata)
 
+        self._draft_token_ids = None
+        self._draft_token_req_ids = None
         self.input_batch.prev_sampled_token_ids = None
 
         def propose_draft_token_ids(sampled_token_ids):
@@ -320,39 +415,40 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     spec_decode_metadata,
                     spec_decode_common_attn_metadata,
                 )
+                self._copy_draft_token_ids_to_cpu(scheduler_output)
 
         spec_config = self.speculative_config
-        use_padded_batch_for_eagle = (
-            spec_config is not None and spec_config.use_eagle() and not spec_config.disable_padded_drafter_batch
-        )
-        effective_drafter_max_model_len = self.max_model_len
-        if effective_drafter_max_model_len is None:
-            effective_drafter_max_model_len = self.model_config.max_model_len
-        if (
-            spec_config is not None
-            and spec_config.draft_model_config is not None
-            and spec_config.draft_model_config.max_model_len is not None
-        ):
-            effective_drafter_max_model_len = spec_config.draft_model_config.max_model_len
-        input_fits_in_drafter = spec_decode_common_attn_metadata and (
-            spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens <= effective_drafter_max_model_len
-        )
-        if use_padded_batch_for_eagle:
-            assert self.speculative_config is not None
-            assert isinstance(self.drafter, EagleProposer)
-            sampled_token_ids = sampler_output.sampled_token_ids
-            if input_fits_in_drafter:
-                propose_draft_token_ids(sampled_token_ids)
-            elif self.valid_sampled_token_count_event is not None:
-                assert spec_decode_common_attn_metadata is not None
-                next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded(
-                    spec_decode_common_attn_metadata,
-                    sampled_token_ids,
-                    self.requests,
-                    self.input_batch,
-                    self.discard_request_mask.gpu,
-                )
-                self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
+        propose_drafts_after_bookkeeping = False
+        if spec_config is not None:
+            input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
+                spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
+                <= self.effective_drafter_max_model_len
+            )
+            if spec_config.use_eagle() and not spec_config.disable_padded_drafter_batch:
+                # EAGLE speculative decoding can use the GPU sampled tokens
+                # as inputs, and does not need to wait for bookkeeping to finish.
+                assert isinstance(self.drafter, EagleProposer)
+                sampled_token_ids = sampler_output.sampled_token_ids
+                if input_fits_in_drafter:
+                    propose_draft_token_ids(sampled_token_ids)
+                elif self.valid_sampled_token_count_event is not None:
+                    assert spec_decode_common_attn_metadata is not None
+                    next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded(
+                        spec_decode_common_attn_metadata,
+                        sampled_token_ids,
+                        self.requests,
+                        self.input_batch,
+                        self.discard_request_mask.gpu,
+                    )
+                    self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
+                    # Since we couldn't run the drafter,
+                    # just use zeros for the draft tokens.
+                    self._draft_token_ids = torch.zeros(1, device=self.device, dtype=torch.int32).expand(
+                        len(self.input_batch.req_ids), self.num_spec_tokens
+                    )
+                    self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
+            else:
+                propose_drafts_after_bookkeeping = input_fits_in_drafter
 
         with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
             (
@@ -372,7 +468,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 spec_decode_metadata,
             )
 
-        if self.speculative_config and not use_padded_batch_for_eagle and input_fits_in_drafter:
+        if propose_drafts_after_bookkeeping:
+            # ngram and other speculative decoding methods use the sampled
+            # tokens on the CPU, so they are run after bookkeeping.
             propose_draft_token_ids(valid_sampled_token_ids)
 
         with record_function_or_nullcontext("gpu_model_runner: eplb"):
@@ -421,6 +519,12 @@ 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:
+                capturer = RoutedExpertsCapturer.get_instance()
+                if capturer is not None:
+                    capturer.save_captured_experts(indices=self.slot_mapping)  # noqa
+                else:
+                    logger.error("RoutedExpertsCapturer not initialized.")
             output = OmniModelRunnerOutput(
                 req_ids=req_ids_output_copy,
                 req_id_to_index=req_id_to_index_output_copy,
@@ -431,6 +535,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 kv_connector_output=kv_connector_output,
                 ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
                 num_nans_in_logits=num_nans_in_logits,
+                cudagraph_stats=cudagraph_stats,
             )
 
         if not self.use_async_scheduling:

vllm_omni/worker/gpu_ar_worker.py

@@ -3,13 +3,14 @@ import os
 
 import torch
 from vllm.logger import init_logger
-from vllm.model_executor import set_random_seed
 from vllm.platforms import current_platform
-from vllm.utils.mem_constants import GiB_bytes
-from vllm.utils.mem_utils import MemorySnapshot
+from vllm.utils.mem_utils import MemorySnapshot, format_gib
+from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_worker import Worker as GPUWorker
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm.v1.worker.utils import request_memory
+from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
 
@@ -24,21 +25,19 @@ class GPUARWorker(GPUWorker):
     """
 
     def init_device(self):
-        device = self.device_config.device
-        if isinstance(device, torch.device) and device.type == "cuda":
+        if self.device_config.device_type == "cuda":
             # This env var set by Ray causes exceptions with graph building.
             os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+            parallel_config = self.parallel_config
             if (
-                self.parallel_config.data_parallel_size > 1
-                and self.parallel_config.data_parallel_size_local > 0
-                and self.parallel_config.distributed_executor_backend not in ["ray", "external_launcher"]
-                and self.vllm_config.parallel_config.data_parallel_backend != "ray"
-                and self.vllm_config.parallel_config.nnodes_within_dp == 1
+                parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
+                and parallel_config.data_parallel_backend != "ray"
+                and parallel_config.nnodes_within_dp == 1
             ):
                 # Use local DP rank if available, otherwise use global DP rank.
                 dp_local_rank = self.parallel_config.data_parallel_rank_local
                 if dp_local_rank is None:
-                    dp_local_rank = self.parallel_config.data_parallel_rank
+                    dp_local_rank = self.parallel_config.data_parallel_index
 
                 tp_pp_world_size = (
                     self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
@@ -80,25 +79,17 @@ class GPUARWorker(GPUWorker):
             torch.cuda.empty_cache()
 
             # take current memory snapshot
-            self.init_snapshot = MemorySnapshot()
-            self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
-            if self.init_snapshot.free_memory < self.requested_memory:
-
-                def gib(bytes_val: float) -> float:
-                    return round(bytes_val / GiB_bytes, 2)
-
-                raise ValueError(
-                    f"Free memory on device "
-                    f"({gib(self.init_snapshot.free_memory)}/"
-                    f"{gib(self.init_snapshot.total_memory)} GiB) on startup "
-                    f"is less than desired GPU memory utilization "
-                    f"({self.cache_config.gpu_memory_utilization}, "
-                    f"{gib(self.requested_memory)} GiB). Decrease GPU memory "
-                    f"utilization or reduce GPU memory used by other processes."
-                )
+            self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
+            self.requested_memory = request_memory(init_snapshot, self.cache_config)
+            logger.debug("worker init memory snapshot: %r", self.init_snapshot)
+            logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
         else:
             raise RuntimeError(f"Not support device type: {self.device_config.device}")
 
+        # Initialize workspace manager
+        num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
+        init_workspace_manager(self.device, num_ubatches)
+
         # Construct the model runner
         self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
 

vllm_omni/worker/gpu_generation_model_runner.py

@@ -8,13 +8,20 @@ from __future__ import annotations
 
 import gc
 import logging
+from copy import copy
 
 import numpy as np
 import torch
 from vllm.config import CUDAGraphMode
-from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
+from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+)
+from vllm.model_executor.models.interfaces import supports_mm_encoder_only
 from vllm.utils.math_utils import cdiv
-from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
+from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.utils import record_function_or_nullcontext
 from vllm.v1.worker.gpu_model_runner import (
@@ -25,9 +32,11 @@ from vllm.v1.worker.gpu_model_runner import (
     get_pp_group,
     set_forward_context,
 )
+from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
 from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 logger = logging.getLogger(__name__)
@@ -47,57 +56,139 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         scheduler_output: SchedulerOutput,
         intermediate_tensors: IntermediateTensors | None = None,
     ) -> OmniModelRunnerOutput | IntermediateTensors:
-        with record_function_or_nullcontext("Preprocess"):
-            with self.synchronize_input_prep():
-                self._update_states(scheduler_output)
-                if not scheduler_output.total_num_scheduled_tokens:
+        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:
+            capturer = RoutedExpertsCapturer.get_instance()
+            if capturer is not None:
+                capturer.clear_buffer()  # noqa
+            else:
+                logger.error("RoutedExpertsCapturer not initialized.")
+
+        if scheduler_output.preempted_req_ids and has_kv_transfer_group():
+            get_kv_transfer_group().handle_preemptions(scheduler_output.preempted_req_ids)
+
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        with (
+            record_function_or_nullcontext("gpu_model_runner: preprocess"),
+            self.synchronize_input_prep(),
+        ):
+            self._update_states(scheduler_output)
+            if not scheduler_output.total_num_scheduled_tokens:
+                return EMPTY_MODEL_RUNNER_OUTPUT
+
+            if has_ec_transfer() and get_ec_transfer().is_producer:
+                with self.maybe_get_ec_connector_output(
+                    scheduler_output,
+                    encoder_cache=self.encoder_cache,
+                ) as ec_connector_output:
+                    self._execute_mm_encoder(scheduler_output)
+                    return make_empty_encoder_model_runner_output(scheduler_output)
+
+            if not num_scheduled_tokens:
+                if (
+                    self.parallel_config.distributed_executor_backend == "external_launcher"
+                    and self.parallel_config.data_parallel_size > 1
+                ):
+                    # this is a corner case when both external launcher
+                    # and DP are enabled, num_scheduled_tokens could be
+                    # 0, and has_unfinished_requests in the outer loop
+                    # returns True. before returning early here we call
+                    # dummy run to ensure coordinate_batch_across_dp
+                    # is called into to avoid out of sync issues.
+                    self._dummy_run(1)
+                if not has_kv_transfer_group():
+                    # Return empty ModelRunnerOutput if no work to do.
                     return EMPTY_MODEL_RUNNER_OUTPUT
 
-                num_reqs = self.input_batch.num_reqs
-                req_ids = self.input_batch.req_ids
-                tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
-                num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
-                max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
-                num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
+                return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
 
-                logits_indices, spec_decode_metadata = self._prepare_inputs(
-                    scheduler_output,
+            if self.cache_config.kv_sharing_fast_prefill:
+                assert not self.num_prompt_logprobs, (
+                    "--kv-sharing-fast-prefill produces incorrect "
+                    "logprobs for prompt tokens, tokens, please disable "
+                    "it when the requests need prompt logprobs"
+                )
+            num_reqs = self.input_batch.num_reqs
+            req_ids = self.input_batch.req_ids
+            tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
+            num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
+            max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
+            num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
+
+            logits_indices, spec_decode_metadata = self._prepare_inputs(
+                scheduler_output,
+                num_scheduled_tokens_np,
+            )
+
+            cascade_attn_prefix_lens = None
+            # Disable cascade attention when using microbatching (DBO)
+            if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
+                # Pre-compute cascade attention prefix lengths
+                cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
                     num_scheduled_tokens_np,
+                    self.input_batch.num_computed_tokens_cpu[:num_reqs],
+                    scheduler_output.num_common_prefix_blocks,
                 )
 
-                (
-                    cudagraph_mode,
-                    batch_desc,
-                    ubatch_slices,
-                    num_tokens_across_dp,
-                ) = self._determine_batch_execution_and_padding(
-                    num_tokens=num_tokens_unpadded,
-                    num_reqs=num_reqs,
-                    num_scheduled_tokens_np=num_scheduled_tokens_np,
-                    max_num_scheduled_tokens=max_num_scheduled_tokens,
-                    use_cascade_attn=False,
-                )
+            (
+                cudagraph_mode,
+                batch_desc,
+                should_ubatch,
+                num_tokens_across_dp,
+                cudagraph_stats,
+            ) = self._determine_batch_execution_and_padding(
+                num_tokens=num_tokens_unpadded,
+                num_reqs=num_reqs,
+                num_scheduled_tokens_np=num_scheduled_tokens_np,
+                max_num_scheduled_tokens=max_num_scheduled_tokens,
+                use_cascade_attn=cascade_attn_prefix_lens is not None,
+                num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
+            )
 
-                num_tokens_padded = batch_desc.num_tokens
-                num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
-                use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
-                pad_attn = cudagraph_mode == CUDAGraphMode.FULL
+            logger.debug(
+                "Running batch with cudagraph_mode: %s, batch_descriptor: %s, "
+                "should_ubatch: %s, num_tokens_across_dp: %s",
+                cudagraph_mode,
+                batch_desc,
+                should_ubatch,
+                num_tokens_across_dp,
+            )
 
-                (
-                    attn_metadata,
-                    spec_decode_common_attn_metadata,
-                ) = self._build_attention_metadata(
-                    num_tokens=num_tokens_unpadded,
-                    num_tokens_padded=num_tokens_padded if pad_attn else None,
-                    num_reqs=num_reqs,
-                    num_reqs_padded=num_reqs_padded if pad_attn else None,
-                    max_query_len=max_num_scheduled_tokens,
-                    ubatch_slices=ubatch_slices,
-                    logits_indices=logits_indices,
-                    use_spec_decode=use_spec_decode,
-                    num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
-                    cascade_attn_prefix_lens=None,
-                )
+            num_tokens_padded = batch_desc.num_tokens
+            num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
+            ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
+                should_ubatch,
+                num_scheduled_tokens_np,
+                num_tokens_padded,
+                num_reqs_padded,
+                self.parallel_config.num_ubatches,
+            )
+
+            logger.debug(
+                "ubatch_slices: %s, ubatch_slices_padded: %s",
+                ubatch_slices,
+                ubatch_slices_padded,
+            )
+
+            pad_attn = cudagraph_mode == CUDAGraphMode.FULL
+
+            use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
+            ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
+
+            attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
+                num_tokens=num_tokens_unpadded,
+                num_tokens_padded=num_tokens_padded if pad_attn else None,
+                num_reqs=num_reqs,
+                num_reqs_padded=num_reqs_padded if pad_attn else None,
+                max_query_len=max_num_scheduled_tokens,
+                ubatch_slices=ubatch_slices_attn,
+                logits_indices=logits_indices,
+                use_spec_decode=use_spec_decode,
+                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
+                cascade_attn_prefix_lens=cascade_attn_prefix_lens,
+            )
 
             (
                 input_ids,
@@ -112,10 +203,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 intermediate_tensors,
             )
 
+        # Set cudagraph mode to none if calc_kv_scales is true.
+        # KV scales calculation involves dynamic operations that are incompatible
+        # with CUDA graph capture.
         if self.calculate_kv_scales:
             cudagraph_mode = CUDAGraphMode.NONE
+            # Mark KV scales as calculated after the first forward pass
             self.calculate_kv_scales = False
 
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
         with (
             set_forward_context(
                 attn_metadata,
@@ -124,7 +221,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 num_tokens_across_dp=num_tokens_across_dp,
                 cudagraph_runtime_mode=cudagraph_mode,
                 batch_descriptor=batch_desc,
-                ubatch_slices=ubatch_slices,
+                ubatch_slices=ubatch_slices_padded,
             ),
             record_function_or_nullcontext("Forward"),
             self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
@@ -139,6 +236,61 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             )
 
         _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
+        self.execute_model_state = ExecuteModelState(
+            scheduler_output,
+            None,
+            spec_decode_metadata,
+            spec_decode_common_attn_metadata,
+            None,
+            None,
+            None,
+            ec_connector_output,
+            cudagraph_stats,
+            multimodal_outputs,
+        )
+        self.kv_connector_output = kv_connector_output
+        return None
+
+    @torch.inference_mode()
+    def sample_tokens(
+        self,
+        grammar_output: GrammarOutput | None = None,
+    ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
+        # NOTE: Even though the model is non-autoregressive, we still need
+        # this function to match the interface of the engine core.
+        # In this case, this function
+        kv_connector_output = self.kv_connector_output
+        self.kv_connector_output = None
+
+        if self.execute_model_state is None:
+            # Nothing to do (PP non-final rank case), output isn't used.
+            if not kv_connector_output:
+                return None  # type: ignore[return-value]
+
+            # In case of PP with kv transfer, we need to pass through the
+            # kv_connector_output
+            if kv_connector_output.is_empty():
+                return EMPTY_MODEL_RUNNER_OUTPUT
+
+            output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
+            output.kv_connector_output = kv_connector_output
+            return output
+
+        # Unpack ephemeral state.
+        (
+            scheduler_output,
+            logits,
+            spec_decode_metadata,
+            spec_decode_common_attn_metadata,
+            hidden_states,
+            sample_hidden_states,
+            aux_hidden_states,
+            ec_connector_output,
+            cudagraph_stats,
+            multimodal_outputs,
+        ) = self.execute_model_state
+        self.execute_model_state = None
+
         pooler_output: list[object] = []
         if isinstance(multimodal_outputs, torch.Tensor):
             assert multimodal_outputs.shape[0] == 1, (
@@ -169,6 +321,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             pooler_output=pooler_output,
             kv_connector_output=kv_connector_output,
             num_nans_in_logits={},
+            cudagraph_stats=cudagraph_stats,
+            ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
         )
 
         if not self.use_async_scheduling:
@@ -176,9 +330,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
 
         return AsyncGPUModelRunnerOutput(
             model_runner_output=output,
-            sampled_token_ids=[],
+            sampled_token_ids=torch.tensor([], device=self.device),
             invalid_req_indices=[],
             async_output_copy_stream=self.async_output_copy_stream,
+            vocab_size=self.input_batch.vocab_size,
+            logprobs_tensors=None,
         )
 
     def _run_generation_model(
@@ -206,7 +362,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             positions=positions,
             intermediate_tensors=intermediate_tensors,
             inputs_embeds=inputs_embeds,
-            **MultiModalKwargs.as_kwargs(model_kwargs, device=self.device),
+            **model_kwargs,
             sampling_metadata=self.input_batch.sampling_metadata,
             logits_index=logits_indices,
             sampler=self.sampler,
@@ -263,6 +419,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             remove_lora: If False, dummy LoRAs are not destroyed after the run
             activate_lora: If False, dummy_run is performed without LoRAs.
         """
+        if supports_mm_encoder_only(self.model):
+            # The current dummy run only covers LM execution, so we can skip it.
+            # mm encoder dummy run may need to add in the future.
+            return torch.tensor([]), torch.tensor([])
+
         assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
 
         # If cudagraph_mode.decode_mode() == FULL and
@@ -316,23 +477,25 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
 
         num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
 
-        _cudagraph_mode, batch_desc, ubatch_slices, num_tokens_across_dp = self._determine_batch_execution_and_padding(
-            num_tokens=num_tokens_unpadded,
-            num_reqs=num_reqs,
-            num_scheduled_tokens_np=num_scheduled_tokens,
-            max_num_scheduled_tokens=max_query_len,
-            use_cascade_attn=False,
-            allow_microbatching=allow_microbatching,
-            force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE),
-            # `force_uniform_decode` is used for cudagraph capture; because for
-            # capturing mixed prefill-decode batches, we sometimes use
-            # num_tokens == num_reqs which looks like a uniform decode batch to the
-            # dispatcher; but we actually want to capture a piecewise cudagraph
-            force_uniform_decode=uniform_decode,
-            # `force_has_lora` is used for cudagraph capture; because LoRA is
-            # activated later in the context manager, but we need to know the
-            # LoRA state when determining the batch descriptor for capture
-            force_has_lora=activate_lora,
+        _cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = (
+            self._determine_batch_execution_and_padding(
+                num_tokens=num_tokens_unpadded,
+                num_reqs=num_reqs,
+                num_scheduled_tokens_np=num_scheduled_tokens,
+                max_num_scheduled_tokens=max_query_len,
+                use_cascade_attn=False,
+                allow_microbatching=allow_microbatching,
+                force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE),
+                # `force_uniform_decode` is used for cudagraph capture; because for
+                # capturing mixed prefill-decode batches, we sometimes use
+                # num_tokens == num_reqs which looks like a uniform decode batch to the
+                # dispatcher; but we actually want to capture a piecewise cudagraph
+                force_uniform_decode=uniform_decode,
+                # `force_has_lora` is used for cudagraph capture; because LoRA is
+                # activated later in the context manager, but we need to know the
+                # LoRA state when determining the batch descriptor for capture
+                force_has_lora=activate_lora,
+            )
         )
 
         if cudagraph_runtime_mode is None:
@@ -345,6 +508,18 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
 
         num_tokens_padded = batch_desc.num_tokens
         num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
+        ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
+            should_ubatch,
+            num_scheduled_tokens,
+            num_tokens_padded,
+            num_reqs_padded,
+            self.vllm_config.parallel_config.num_ubatches,
+        )
+        logger.debug(
+            "ubatch_slices: %s, ubatch_slices_padded: %s",
+            ubatch_slices,
+            ubatch_slices_padded,
+        )
 
         attn_metadata: PerLayerAttnMetadata | None = None
 
@@ -366,11 +541,12 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
             self.query_start_loc.copy_to_gpu()
 
+            pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
             attn_metadata, _ = self._build_attention_metadata(
                 num_tokens=num_tokens_unpadded,
                 num_reqs=num_reqs_padded,
                 max_query_len=max_query_len,
-                ubatch_slices=ubatch_slices,
+                ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
                 for_cudagraph_capture=is_graph_capturing,
             )
 
@@ -383,10 +559,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         ):
             # Make sure padding doesn't exceed max_num_tokens
             assert num_tokens_padded <= self.max_num_tokens
-            model_kwargs = self._init_model_kwargs(num_tokens_padded)
+            model_kwargs = self._init_model_kwargs()
             if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
+                input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)
+
                 model_kwargs = {
                     **model_kwargs,
                     **self._dummy_mm_kwargs(num_reqs),
@@ -394,7 +570,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             elif self.enable_prompt_embeds:
                 input_ids = None
                 inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
-                model_kwargs = self._init_model_kwargs(num_tokens_padded)
+                model_kwargs = self._init_model_kwargs()
             else:
                 input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
@@ -418,16 +594,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
 
                 intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False)
 
-            if ubatch_slices is not None:
+            if ubatch_slices_padded is not None:
                 # Adjust values to reflect a single ubatch.
                 # TODO(sage,lucas): this is cruft that should be addressed in
                 #  the padding refactor.
-                num_tokens_padded = ubatch_slices[0].num_tokens
+                num_tokens_padded = ubatch_slices_padded[0].num_tokens
                 if num_tokens_across_dp is not None:
                     num_tokens_across_dp[:] = num_tokens_padded
 
             with (
-                self.maybe_randomize_inputs(input_ids),
+                self.maybe_randomize_inputs(input_ids, inputs_embeds),
                 set_forward_context(
                     attn_metadata,
                     self.vllm_config,
@@ -435,7 +611,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                     num_tokens_across_dp=num_tokens_across_dp,
                     cudagraph_runtime_mode=cudagraph_runtime_mode,
                     batch_descriptor=batch_desc,
-                    ubatch_slices=ubatch_slices,
+                    ubatch_slices=ubatch_slices_padded,
                 ),
             ):
                 outputs = self.model(
@@ -453,10 +629,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
             if self.speculative_config and self.speculative_config.use_eagle():
                 assert isinstance(self.drafter, EagleProposer)
+                # Eagle currently only supports PIECEWISE cudagraphs.
+                # Therefore only use cudagraphs if the main model uses PIECEWISE
+                # NOTE(lucas): this is a hack, need to clean up.
                 use_cudagraphs = (
-                    cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE)
-                    and not self.speculative_config.enforce_eager
-                )
+                    (is_graph_capturing and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE)
+                    or (not is_graph_capturing and cudagraph_runtime_mode != CUDAGraphMode.NONE)
+                ) and not self.speculative_config.enforce_eager
 
                 # Note(gnovack) - We need to disable cudagraphs for one of the two
                 # lora cases when cudagraph_specialize_lora is enabled. This is a
@@ -471,6 +650,17 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                     is_graph_capturing=is_graph_capturing,
                 )
 
+        # We register layerwise NVTX hooks here after the first dynamo tracing is
+        # done to avoid nvtx operations in hook functions being traced by
+        # torch dynamo and causing graph breaks.
+        # Note that for DYNAMO_ONCE and VLLM_COMPILE mode,
+        # compiled model's dynamo tracing is only done once and the compiled model's
+        # __call__ function is replaced by calling the compiled function.
+        # So it's safe to register hooks here. Hooks will be registered to
+        # both compiled and uncompiled models but they will never
+        # be called on the compiled model execution path.
+        self._register_layerwise_nvtx_hooks()
+
         # This is necessary to avoid blocking DP.
         # For dummy runs, we typically skip EPLB since we don't have any real
         # requests to process.

vllm_omni/worker/gpu_generation_worker.py

@@ -2,16 +2,20 @@ import gc
 import os
 
 import torch
-from vllm.model_executor import set_random_seed
+from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from vllm.utils.mem_constants import GiB_bytes
-from vllm.utils.mem_utils import MemorySnapshot
+from vllm.utils.mem_utils import MemorySnapshot, format_gib
+from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_worker import Worker as GPUWorker
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm.v1.worker.utils import request_memory
+from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
 
+logger = init_logger(__name__)
+
 
 class GPUGenerationWorker(GPUWorker):
     """GPU Worker for Generation model (non-autoregressive waveform generation).
@@ -21,21 +25,19 @@ class GPUGenerationWorker(GPUWorker):
     """
 
     def init_device(self):
-        device = self.device_config.device
-        if isinstance(device, torch.device) and device.type == "cuda":
+        if self.device_config.device_type == "cuda":
             # This env var set by Ray causes exceptions with graph building.
             os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+            parallel_config = self.parallel_config
             if (
-                self.parallel_config.data_parallel_size > 1
-                and self.parallel_config.data_parallel_size_local > 0
-                and self.parallel_config.distributed_executor_backend not in ["ray", "external_launcher"]
-                and self.vllm_config.parallel_config.data_parallel_backend != "ray"
-                and self.vllm_config.parallel_config.nnodes_within_dp == 1
+                parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
+                and parallel_config.data_parallel_backend != "ray"
+                and parallel_config.nnodes_within_dp == 1
             ):
                 # Use local DP rank if available, otherwise use global DP rank.
                 dp_local_rank = self.parallel_config.data_parallel_rank_local
                 if dp_local_rank is None:
-                    dp_local_rank = self.parallel_config.data_parallel_rank
+                    dp_local_rank = self.parallel_config.data_parallel_index
 
                 tp_pp_world_size = (
                     self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
@@ -77,26 +79,17 @@ class GPUGenerationWorker(GPUWorker):
             torch.cuda.empty_cache()
 
             # take current memory snapshot
-            self.init_snapshot = MemorySnapshot()
-            self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
-            if self.init_snapshot.free_memory < self.requested_memory:
-
-                def gib(bytes_val: float) -> float:
-                    return round(bytes_val / GiB_bytes, 2)
-
-                raise ValueError(
-                    f"Free memory on device "
-                    f"({gib(self.init_snapshot.free_memory)}/"
-                    f"{gib(self.init_snapshot.total_memory)} GiB) on startup "
-                    f"is less than desired GPU memory utilization "
-                    f"({self.cache_config.gpu_memory_utilization}, "
-                    f"{gib(self.requested_memory)} GiB). Decrease GPU memory "
-                    f"utilization or reduce GPU memory used by other processes."
-                )
+            self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
+            self.requested_memory = request_memory(init_snapshot, self.cache_config)
+            logger.debug("worker init memory snapshot: %r", self.init_snapshot)
+            logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
         else:
             raise RuntimeError(f"Not support device type: {self.device_config.device}")
 
-        # Construct the model runner
+        # Initialize workspace manager
+        num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
+        init_workspace_manager(self.device, num_ubatches)
+
         self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device)
 
         if self.rank == 0:

vllm_omni/worker/gpu_model_runner.py

@@ -8,7 +8,7 @@ from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
-from vllm.model_executor.models.interfaces import supports_mrope
+from vllm.model_executor.models.interfaces import supports_mm_encoder_only, supports_mrope
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
 from vllm.utils.import_utils import LazyLoader
@@ -16,6 +16,7 @@ from vllm.utils.math_utils import cdiv
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.worker.gpu_input_batch import CachedRequestState
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata
+from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 
@@ -129,6 +130,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         # Remove finished requests from the cached states.
         for req_id in scheduler_output.finished_req_ids:
             self.requests.pop(req_id, None)
+            self.num_prompt_logprobs.pop(req_id, None)
         # 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
@@ -149,7 +151,14 @@ class OmniGPUModelRunner(GPUModelRunner):
         # they will be scheduled again sometime in the future.
         scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
         cached_req_ids = self.input_batch.req_id_to_index.keys()
-        unscheduled_req_ids = cached_req_ids - scheduled_req_ids
+        resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
+        # NOTE(zhuohan): cached_req_ids and resumed_req_ids are usually disjoint,
+        # so `(scheduled_req_ids - resumed_req_ids) == scheduled_req_ids` holds
+        # apart from the forced-preemption case in reset_prefix_cache. And in
+        # that case we include the resumed_req_ids in the unscheduled set so
+        # that they get cleared from the persistent batch before being re-scheduled
+        # in the normal resumed request path.
+        unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids)
         # NOTE(woosuk): The persistent batch optimization assumes that
         # consecutive batches contain mostly the same requests. If batches
         # have low request overlap (e.g., alternating between two distinct
@@ -241,20 +250,62 @@ class OmniGPUModelRunner(GPUModelRunner):
                 logger.error(f"Error decoding additional information: {e}")
                 pass
 
+            if sampling_params and sampling_params.prompt_logprobs is not None:
+                self.num_prompt_logprobs[req_id] = (
+                    self.input_batch.vocab_size
+                    if sampling_params.prompt_logprobs == -1
+                    else sampling_params.prompt_logprobs
+                )
             # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
             if self.uses_mrope:
                 self._init_mrope_positions(req_state)
 
+            # Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
+            if self.uses_xdrope_dim > 0:
+                self._init_xdrope_positions(req_state)
+
             reqs_to_add.append(self.requests[req_id])
 
         # Update the states of the running/resumed requests.
         is_last_rank = get_pp_group().is_last_rank
         req_data = scheduler_output.scheduled_cached_reqs
+        scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
+
+        # Wait until valid_sampled_tokens_count is copied to cpu,
+        # then use it to update actual num_computed_tokens of each request.
+        valid_sampled_token_count = self._get_valid_sampled_token_count()
+
         for i, req_id in enumerate(req_data.req_ids):
             req_state = self.requests[req_id]
             num_computed_tokens = req_data.num_computed_tokens[i]
             new_block_ids = req_data.new_block_ids[i]
-            resumed_from_preemption = req_data.resumed_from_preemption[i]
+            resumed_from_preemption = req_id in req_data.resumed_req_ids
+            num_output_tokens = req_data.num_output_tokens[i]
+            req_index = self.input_batch.req_id_to_index.get(req_id)
+
+            if req_state.prev_num_draft_len and self.use_async_scheduling:
+                # 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 = [],
+                # 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
+                # 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.
+                if req_index is None:
+                    req_state.prev_num_draft_len = 0
+                else:
+                    assert self.input_batch.prev_req_id_to_index is not None
+                    prev_req_index = self.input_batch.prev_req_id_to_index[req_id]
+                    num_accepted = valid_sampled_token_count[prev_req_index] - 1
+                    num_rejected = req_state.prev_num_draft_len - num_accepted
+                    num_computed_tokens -= num_rejected
+                    req_state.output_token_ids.extend([-1] * num_accepted)
 
             # Update the cached states.
             req_state.num_computed_tokens = num_computed_tokens
@@ -272,6 +323,13 @@ class OmniGPUModelRunner(GPUModelRunner):
                     req_state.output_token_ids.append(new_token_ids[-1])
                 elif num_new_tokens > 0:
                     req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
+            elif num_output_tokens < len(req_state.output_token_ids):
+                # Some output tokens were discarded due to a sync-KV-load
+                # failure. Align the cached state.
+                del req_state.output_token_ids[num_output_tokens:]
+                if req_index is not None:
+                    end_idx = self.input_batch.num_prompt_tokens[req_index] + num_output_tokens
+                    self.input_batch.num_tokens_no_spec[req_index] = end_idx
 
             # Update the block IDs.
             if not resumed_from_preemption:
@@ -280,6 +338,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                     for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
                         block_ids.extend(new_ids)
             else:
+                assert req_index is None
                 assert new_block_ids is not None
                 # The request is resumed from preemption.
                 # Replace the existing block IDs with the new ones.
@@ -290,6 +349,13 @@ class OmniGPUModelRunner(GPUModelRunner):
                 # The request is not in the persistent batch.
                 # The request was either preempted and resumed later, or was not
                 # scheduled in the previous step and needs to be added again.
+
+                if self.use_async_scheduling and num_output_tokens > 0:
+                    # We must recover the output token ids for resumed requests in the
+                    # async scheduling case, so that correct input_ids are obtained.
+                    resumed_token_ids = req_data.all_token_ids[req_id]
+                    req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]
+
                 reqs_to_add.append(req_state)
                 continue
 
@@ -306,22 +372,15 @@ class OmniGPUModelRunner(GPUModelRunner):
                 end_token_index = num_computed_tokens + len(new_token_ids)
                 self.input_batch.token_ids_cpu[req_index, start_token_index:end_token_index] = new_token_ids
                 self.input_batch.num_tokens_no_spec[req_index] = end_token_index
-                self.input_batch.num_tokens[req_index] = end_token_index
 
             # Add spec_token_ids to token_ids_cpu.
-            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())
-            if spec_token_ids:
-                num_spec_tokens = len(spec_token_ids)
-                start_index = self.input_batch.num_tokens_no_spec[req_index]
-                end_token_index = start_index + num_spec_tokens
-                self.input_batch.token_ids_cpu[req_index, start_index:end_token_index] = spec_token_ids
-                # NOTE(woosuk): `num_tokens` here may include spec tokens.
-                self.input_batch.num_tokens[req_index] += num_spec_tokens
+            self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens)
 
         # Add the new or resumed requests to the persistent batch.
         # The smaller empty indices are filled first.
         for request in reqs_to_add:
             self.input_batch.add_request(request)
+            self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens)
 
         # Condense the batched states if there are gaps left by removed requests
         self.input_batch.condense()
@@ -388,6 +447,11 @@ class OmniGPUModelRunner(GPUModelRunner):
             remove_lora: If False, dummy LoRAs are not destroyed after the run
             activate_lora: If False, dummy_run is performed without LoRAs.
         """
+        if supports_mm_encoder_only(self.model):
+            # The current dummy run only covers LM execution, so we can skip it.
+            # mm encoder dummy run may need to add in the future.
+            return torch.tensor([]), torch.tensor([])
+
         assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
 
         # If cudagraph_mode.decode_mode() == FULL and
@@ -441,23 +505,25 @@ class OmniGPUModelRunner(GPUModelRunner):
 
         num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
 
-        _cudagraph_mode, batch_desc, ubatch_slices, num_tokens_across_dp = self._determine_batch_execution_and_padding(
-            num_tokens=num_tokens_unpadded,
-            num_reqs=num_reqs,
-            num_scheduled_tokens_np=num_scheduled_tokens,
-            max_num_scheduled_tokens=max_query_len,
-            use_cascade_attn=False,
-            allow_microbatching=allow_microbatching,
-            force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE),
-            # `force_uniform_decode` is used for cudagraph capture; because for
-            # capturing mixed prefill-decode batches, we sometimes use
-            # num_tokens == num_reqs which looks like a uniform decode batch to the
-            # dispatcher; but we actually want to capture a piecewise cudagraph
-            force_uniform_decode=uniform_decode,
-            # `force_has_lora` is used for cudagraph capture; because LoRA is
-            # activated later in the context manager, but we need to know the
-            # LoRA state when determining the batch descriptor for capture
-            force_has_lora=activate_lora,
+        _cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = (
+            self._determine_batch_execution_and_padding(
+                num_tokens=num_tokens_unpadded,
+                num_reqs=num_reqs,
+                num_scheduled_tokens_np=num_scheduled_tokens,
+                max_num_scheduled_tokens=max_query_len,
+                use_cascade_attn=False,
+                allow_microbatching=allow_microbatching,
+                force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE),
+                # `force_uniform_decode` is used for cudagraph capture; because for
+                # capturing mixed prefill-decode batches, we sometimes use
+                # num_tokens == num_reqs which looks like a uniform decode batch to the
+                # dispatcher; but we actually want to capture a piecewise cudagraph
+                force_uniform_decode=uniform_decode,
+                # `force_has_lora` is used for cudagraph capture; because LoRA is
+                # activated later in the context manager, but we need to know the
+                # LoRA state when determining the batch descriptor for capture
+                force_has_lora=activate_lora,
+            )
         )
 
         if cudagraph_runtime_mode is None:
@@ -470,6 +536,18 @@ class OmniGPUModelRunner(GPUModelRunner):
 
         num_tokens_padded = batch_desc.num_tokens
         num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
+        ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
+            should_ubatch,
+            num_scheduled_tokens,
+            num_tokens_padded,
+            num_reqs_padded,
+            self.vllm_config.parallel_config.num_ubatches,
+        )
+        logger.debug(
+            "ubatch_slices: %s, ubatch_slices_padded: %s",
+            ubatch_slices,
+            ubatch_slices_padded,
+        )
 
         attn_metadata: PerLayerAttnMetadata | None = None
 
@@ -491,11 +569,12 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
             self.query_start_loc.copy_to_gpu()
 
+            pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
             attn_metadata, _ = self._build_attention_metadata(
                 num_tokens=num_tokens_unpadded,
                 num_reqs=num_reqs_padded,
                 max_query_len=max_query_len,
-                ubatch_slices=ubatch_slices,
+                ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
                 for_cudagraph_capture=is_graph_capturing,
             )
 
@@ -508,10 +587,10 @@ class OmniGPUModelRunner(GPUModelRunner):
         ):
             # Make sure padding doesn't exceed max_num_tokens
             assert num_tokens_padded <= self.max_num_tokens
-            model_kwargs = self._init_model_kwargs(num_tokens_padded)
+            model_kwargs = self._init_model_kwargs()
             if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
+                input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)
+
                 model_kwargs = {
                     **model_kwargs,
                     **self._dummy_mm_kwargs(num_reqs),
@@ -519,7 +598,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             elif self.enable_prompt_embeds:
                 input_ids = None
                 inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
-                model_kwargs = self._init_model_kwargs(num_tokens_padded)
+                model_kwargs = self._init_model_kwargs()
             else:
                 input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
@@ -543,16 +622,16 @@ class OmniGPUModelRunner(GPUModelRunner):
 
                 intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False)
 
-            if ubatch_slices is not None:
+            if ubatch_slices_padded is not None:
                 # Adjust values to reflect a single ubatch.
                 # TODO(sage,lucas): this is cruft that should be addressed in
                 #  the padding refactor.
-                num_tokens_padded = ubatch_slices[0].num_tokens
+                num_tokens_padded = ubatch_slices_padded[0].num_tokens
                 if num_tokens_across_dp is not None:
                     num_tokens_across_dp[:] = num_tokens_padded
 
             with (
-                self.maybe_randomize_inputs(input_ids),
+                self.maybe_randomize_inputs(input_ids, inputs_embeds),
                 set_forward_context(
                     attn_metadata,
                     self.vllm_config,
@@ -560,7 +639,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                     num_tokens_across_dp=num_tokens_across_dp,
                     cudagraph_runtime_mode=cudagraph_runtime_mode,
                     batch_descriptor=batch_desc,
-                    ubatch_slices=ubatch_slices,
+                    ubatch_slices=ubatch_slices_padded,
                 ),
             ):
                 if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
@@ -585,10 +664,13 @@ class OmniGPUModelRunner(GPUModelRunner):
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
             if self.speculative_config and self.speculative_config.use_eagle():
                 assert isinstance(self.drafter, EagleProposer)
+                # Eagle currently only supports PIECEWISE cudagraphs.
+                # Therefore only use cudagraphs if the main model uses PIECEWISE
+                # NOTE(lucas): this is a hack, need to clean up.
                 use_cudagraphs = (
-                    cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE)
-                    and not self.speculative_config.enforce_eager
-                )
+                    (is_graph_capturing and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE)
+                    or (not is_graph_capturing and cudagraph_runtime_mode != CUDAGraphMode.NONE)
+                ) and not self.speculative_config.enforce_eager
 
                 # Note(gnovack) - We need to disable cudagraphs for one of the two
                 # lora cases when cudagraph_specialize_lora is enabled. This is a
@@ -603,6 +685,17 @@ class OmniGPUModelRunner(GPUModelRunner):
                     is_graph_capturing=is_graph_capturing,
                 )
 
+        # We register layerwise NVTX hooks here after the first dynamo tracing is
+        # done to avoid nvtx operations in hook functions being traced by
+        # torch dynamo and causing graph breaks.
+        # Note that for DYNAMO_ONCE and VLLM_COMPILE mode,
+        # compiled model's dynamo tracing is only done once and the compiled model's
+        # __call__ function is replaced by calling the compiled function.
+        # So it's safe to register hooks here. Hooks will be registered to
+        # both compiled and uncompiled models but they will never
+        # be called on the compiled model execution path.
+        self._register_layerwise_nvtx_hooks()
+
         # This is necessary to avoid blocking DP.
         # For dummy runs, we typically skip EPLB since we don't have any real
         # requests to process.
@@ -772,7 +865,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         num_input_tokens: int,
         intermediate_tensors: IntermediateTensors | None = None,
     ):
-        """Align with v0.12 preprocess and omni's additional information handling."""
+        """Align with v0.14.0 preprocess and omni's additional information handling."""
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         is_first_rank = get_pp_group().is_first_rank
         is_encoder_decoder = self.model_config.is_encoder_decoder
@@ -802,10 +895,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             # TODO(woosuk): Avoid the copy. Optimize.
             self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
 
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens)
             model_kwargs = {
-                **self._init_model_kwargs(num_scheduled_tokens),
+                **self._init_model_kwargs(),
                 **self._extract_mm_kwargs(scheduler_output),
             }
         elif self.enable_prompt_embeds and is_first_rank:
@@ -829,7 +921,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
 
             inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            model_kwargs = self._init_model_kwargs()
             input_ids = self.input_ids.gpu[:num_input_tokens]
         else:
             # For text-only models, we use token ids as input.
@@ -838,7 +930,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             # then the embedding layer is not included in the CUDA graph.
             input_ids = self.input_ids.gpu[:num_input_tokens]
             inputs_embeds = None
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            model_kwargs = self._init_model_kwargs()
 
         if self.uses_mrope:
             positions = self.mrope_positions.gpu[:, :num_input_tokens]
@@ -929,7 +1021,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         decode_batch_size = len(decode_req_ids)
         if decode_batch_size == 0:
             return
-        _cudagraph_mode, batch_desc, _, _ = self._determine_batch_execution_and_padding(
+        _cudagraph_mode, batch_desc, _, _, _ = self._determine_batch_execution_and_padding(
             num_tokens=decode_batch_size,
             num_reqs=decode_batch_size,
             num_scheduled_tokens_np=np.ones(decode_batch_size, dtype=np.int32),