[Rebase] Rebase to vllm 0.12.0 (#335)

Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Rein Yang <ruiruyang2@gmail.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

.buildkite/scripts/simple_test.sh

@@ -47,7 +47,7 @@ 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.11.0
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.12.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/

.gitignore

@@ -221,6 +221,7 @@ datasets/
 outputs/
 results/
 generated/
+output_*/
 
 # Configuration overrides
 configs/local.yaml

docker/Dockerfile.ci

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

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

@@ -18,9 +18,9 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 #### Installation of vLLM
 
-vLLM-Omni is built based on vLLM v0.11.0. Please install it with command below.
+vLLM-Omni is built based on vLLM v0.12.0. Please install it with command below.
 ```bash
-uv pip install vllm==0.11.0 --torch-backend=auto
+uv pip install vllm==0.12.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -34,10 +34,10 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.11.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.12.0 release version of the library
 
 ```bash
-uv pip install vllm==0.11.0 --torch-backend=auto
+uv pip install vllm==0.12.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -54,11 +54,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.11.0
+git checkout v0.12.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.11.0/vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl
+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
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -89,7 +89,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.11.0rc1 \
+    vllm/vllm-omni:v0.12.0rc1 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -19,7 +19,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 #### Get the correct docker environment
 
-The v0.11.0 releases is on Oct 3, 2025, so we will pull a AMD nightly docker image close to the release period. `rocm/vllm-dev:nightly_main_20251005` is picked as it has a correct `transformers` version that supports qwen3_omni.
+The v0.12.0 releases is on Dec 3, 2025, so we will pull a AMD nightly docker image close to the release period. `rocm/vllm-dev:nightly_main_20251205` is picked as it has a correct `transformers` version that supports qwen3_omni.
 
 `docker pull rocm/vllm-dev:nightly_main_20251005`
 
@@ -51,18 +51,18 @@ sudo apt install ffmpeg -y
 ```
 
 #### Installation of vLLM
-Install the v0.11.0 release of vLLM as vLLM-Omni is currently depend on v0.11.0.
+Install the v0.12.0 release of vLLM as vLLM-Omni is currently depend on v0.12.0.
 
 ```bash
 # uninstall existing vllm which is based on the main branch, not on the releases branch
 python3 -m pip uninstall -y vllm
-git clone https://github.com/vllm-project/vllm.git vllm0110
-cd vllm0110/
+git clone https://github.com/vllm-project/vllm.git vllm0120
+cd vllm0120/
 
-# checkout the version tag v0.11.0
-git checkout -b v0.11.0 v0.11.0
+# checkout the version tag v0.12.0
+git checkout -b v0.12.0 v0.12.0
 
-# you should see 0.11.0
+# you should see 0.12.0
 python -c "import setuptools_scm; print(setuptools_scm.get_version())"
 
 # PYTORCH_ROCM_ARCH=<your-gpu-arch> python3 setup.py develop
@@ -92,7 +92,7 @@ pip install vllm-omni
 # clone and checkout to the release version.
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
-git checkout v0.11.0rc1
+git checkout v0.12.0rc1
 pip install -e .
 ```
 

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.11.0rc2
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -17,7 +17,7 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.sampling_params import SamplingParams
-from vllm.utils import FlexibleArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni.entrypoints.omni import Omni
 
@@ -119,7 +119,7 @@ def get_use_audio_in_video_query(
     question = "Describe the content of the video, then convert what the baby say into text."
     prompt = (
         f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>"
+        "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|><|audio_bos|><|AUDIO|><|audio_eos|>"
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
@@ -319,7 +319,6 @@ def main(args):
         query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
     else:
         query_result = query_func()
-
     omni_llm = Omni(
         model=model_name,
         log_stats=args.enable_stats,

examples/offline_inference/qwen3_omni/end2end.py

@@ -17,7 +17,7 @@ from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
-from vllm.utils import FlexibleArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni.entrypoints.omni import Omni
 
@@ -281,7 +281,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="mixed_modalities",
+        default="use_video",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -4,7 +4,7 @@ import os
 import requests
 from openai import OpenAI
 from vllm.assets.audio import AudioAsset
-from vllm.utils import FlexibleArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 # Modify OpenAI's API key and API base to use vLLM's API server.
 openai_api_key = "EMPTY"

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -5,7 +5,7 @@ from typing import NamedTuple
 import requests
 from openai import OpenAI
 from vllm.assets.audio import AudioAsset
-from vllm.utils import FlexibleArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 # Modify OpenAI's API key and API base to use vLLM's API server.
 openai_api_key = "EMPTY"

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vllm-omni"
-version = "0.11.0rc1"
+version = "0.12.0rc1"
 description = "A framework for efficient model inference with omni-modality models"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"
@@ -37,7 +37,7 @@ dependencies = [
     "gradio==5.50",
     "soundfile>=0.13.1",
     "cache-dit==1.1.8",
-    # "vllm==0.11.0",  # TODO: fix the entrypoints overwrite problem
+    # "vllm==0.12.0",  # TODO: fix the entrypoints overwrite problem
 ]
 
 [project.optional-dependencies]

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -44,7 +44,7 @@ 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.5
+       gpu_memory_utilization: 0.32
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -27,7 +27,7 @@ test_params = [(model, stage_config) for model in models for stage_config in sta
 def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None:
     """Test processing video, generating audio output."""
     model, stage_config_path = test_config
-    with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
+    with omni_runner(model, seed=42, stage_configs_path=stage_config_path, init_sleep_seconds=90) as runner:
         # Prepare inputs
         question = "Describe the video briefly."
         video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays

tests/e2e/online_serving/test_qwen3_omni.py

@@ -109,7 +109,7 @@ def omni_server(request):
     Multi-stage initialization can take 10-20+ minutes.
     """
     model, stage_config_path = request.param
-    with OmniServer(model, ["--stage-configs-path", stage_config_path]) as server:
+    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--init-sleep-seconds", "90"]) as server:
         yield server
 
 

tests/entrypoints/test_omni_llm.py

@@ -393,12 +393,12 @@ def mock_get_config(monkeypatch):
     )
     # Also mock in processor module if it's imported
     monkeypatch.setattr(
-        "vllm_omni.engine.processor.length_from_prompt_token_ids_or_embeds",
+        "vllm_omni.engine.input_processor.length_from_prompt_token_ids_or_embeds",
         _mock_length_from_prompt_token_ids_or_embeds,
         raising=False,
     )
     # If processor module is already imported, patch it directly
-    processor_module_path = "vllm_omni.engine.processor"
+    processor_module_path = "vllm_omni.engine.input_processor"
     if processor_module_path in sys.modules:
         processor_module = sys.modules[processor_module_path]
         setattr(

vllm_omni/config/model.py

@@ -1,17 +1,17 @@
-import json
 import warnings
 from importlib.util import find_spec
-from typing import Any, Literal
+from typing import Any
 
 import torch
 import vllm.envs as envs
-from omegaconf import DictConfig, OmegaConf
 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,
@@ -27,9 +27,14 @@ from vllm.transformers_utils.config import (
     get_hf_image_processor_config,
     get_hf_text_config,
     get_pooling_config,
-    is_interleaved,
 )
-from vllm.transformers_utils.utils import maybe_model_redirect
+from vllm.transformers_utils.gguf_utils import (
+    maybe_patch_hf_config_from_gguf,
+)
+from vllm.transformers_utils.utils import (
+    is_gguf,
+    maybe_model_redirect,
+)
 
 import vllm_omni.model_executor.models as me_models
 
@@ -87,29 +92,19 @@ class OmniModelConfig(ModelConfig):
     def __post_init__(
         self,
         # Multimodal config init vars
-        limit_mm_per_prompt: dict[str, int] | None,
+        limit_mm_per_prompt: dict[str, int | dict[str, int]] | None,
+        enable_mm_embeds: bool | None,
         media_io_kwargs: dict[str, dict[str, Any]] | None,
         mm_processor_kwargs: dict[str, Any] | None,
         mm_processor_cache_gb: float | None,
         mm_processor_cache_type: MMCacheType | None,
         mm_shm_cache_max_object_size_mb: int | None,
         mm_encoder_tp_mode: MMEncoderTPMode | None,
+        mm_encoder_attn_backend: AttentionBackendEnum | str | None,
         interleave_mm_strings: bool | None,
         skip_mm_profiling: bool | None,
         video_pruning_rate: float | None,
     ) -> None:
-        # Set the default seed to 0 in V1.
-        if envs.VLLM_USE_V1 and self.seed is None:
-            self.seed = 0
-            if not envs.VLLM_ENABLE_V1_MULTIPROCESSING:
-                logger.warning(
-                    "The global random seed is set to %d. Since "
-                    "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may "
-                    "affect the random state of the Python process that "
-                    "launched vLLM.",
-                    self.seed,
-                )
-
         # Keep set served_model_name before maybe_model_redirect(self.model)
         self.served_model_name = get_served_model_name(self.model, self.served_model_name)
         self.model = maybe_model_redirect(self.model)
@@ -124,53 +119,36 @@ class OmniModelConfig(ModelConfig):
             self.hf_config_path = maybe_model_redirect(self.hf_config_path)
 
         if callable(self.hf_overrides):
-            hf_overrides_kw: dict[str, Any] = {}
+            hf_overrides_kw = {}
             hf_overrides_fn = self.hf_overrides
             dict_overrides: dict[str, Any] = {}
         else:
-            # Separate dict overrides from flat ones; convert OmegaConf to dict
+            # Separate dict overrides from flat ones
+            # We'll determine how to apply dict overrides after loading the config
             hf_overrides_kw = {}
             dict_overrides = {}
             for key, value in self.hf_overrides.items():
-                if isinstance(value, DictConfig):
-                    value = OmegaConf.to_container(value, resolve=True, structured_config_mode=None)
                 if isinstance(value, dict):
                     dict_overrides[key] = value
                 else:
                     hf_overrides_kw[key] = value
             hf_overrides_fn = None
 
-        if self.rope_scaling:
-            hf_override: dict[str, Any] = {"rope_scaling": self.rope_scaling}
-            hf_overrides_kw.update(hf_override)
-            hf_overrides_str = json.dumps(hf_overrides_kw)
-            msg = (
-                "`--rope-scaling` will be removed in a future release. "
-                f"'Please instead use `--hf-overrides '{hf_overrides_str}'`"
-            )
-            warnings.warn(DeprecationWarning(msg), stacklevel=2)
-        if self.rope_theta is not None:
-            hf_override = {"rope_theta": self.rope_theta}
-            hf_overrides_kw.update(hf_override)
-            hf_overrides_str = json.dumps(hf_overrides_kw)
-            msg = (
-                "`--rope-theta` will be removed in a future release. "
-                f"'Please instead use `--hf-overrides '{hf_overrides_str}'`"
-            )
-            warnings.warn(DeprecationWarning(msg), stacklevel=2)
-
         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 "
+                "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", stacklevel=2)
+            warnings.warn(
+                "override-attention-dtype is set but not using ROCm platform",
+                stacklevel=2,
+            )
 
         if self.enable_sleep_mode and not current_platform.is_sleep_mode_available():
             raise ValueError("Sleep mode is not supported on current platform.")
@@ -184,11 +162,14 @@ class OmniModelConfig(ModelConfig):
             hf_overrides_kw=hf_overrides_kw,
             hf_overrides_fn=hf_overrides_fn,
         )
-
-        if dict_overrides:
-            self._apply_dict_overrides(hf_config, dict_overrides)
+        hf_config = maybe_patch_hf_config_from_gguf(
+            self.model,
+            hf_config,
+        )
 
         self.hf_config = hf_config
+        if dict_overrides:
+            self._apply_dict_overrides(hf_config, dict_overrides)
         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()
@@ -216,7 +197,7 @@ class OmniModelConfig(ModelConfig):
 
         if self.task is not None:
             runner: RunnerOption = "auto"
-            convert: ConvertType | Literal["auto"] = "auto"  # type: ignore[name-defined]
+            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."
             )
@@ -228,7 +209,7 @@ class OmniModelConfig(ModelConfig):
             if is_generative_model and is_pooling_model:
                 if is_generative_task:
                     runner = "generate"
-                    convert = "auto"  # type: ignore[assignment]
+                    convert = "auto"
                     msg_hint = (
                         "Please replace this option with `--runner "
                         "generate` to continue using this model "
@@ -236,7 +217,7 @@ class OmniModelConfig(ModelConfig):
                     )
                 elif is_pooling_task:
                     runner = "pooling"
-                    convert = "auto"  # type: ignore[assignment]
+                    convert = "auto"
                     msg_hint = (
                         "Please replace this option with `--runner "
                         "pooling` to continue using this model "
@@ -247,7 +228,7 @@ class OmniModelConfig(ModelConfig):
             elif is_generative_model or is_pooling_model:
                 if is_generative_task:
                     runner = "generate"
-                    convert = "auto"  # type: ignore[assignment]
+                    convert = "auto"
                     msg_hint = "Please remove this option"
                 elif is_pooling_task:
                     runner = "pooling"
@@ -260,12 +241,29 @@ class OmniModelConfig(ModelConfig):
                 else:  # task == "auto"
                     pass
             else:
-                raise AssertionError(
-                    f"The model should be a generative or pooling model when task is set to {self.task!r}."
-                )
+                # 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  # type: ignore[assignment]
+            self.convert = convert
 
             msg = f"{msg_prefix} {msg_hint}"
             warnings.warn(msg, DeprecationWarning, stacklevel=2)
@@ -276,6 +274,7 @@ class OmniModelConfig(ModelConfig):
         if self.runner_type == "generate" and not is_generative_model:
             generate_converts = _RUNNER_CONVERTS["generate"]
             if self.convert_type not in generate_converts:
+                # Currently we don't have any converters for generative models
                 raise ValueError("This model does not support `--runner generate`.")
         if self.runner_type == "pooling" and not is_pooling_model:
             pooling_converts = _RUNNER_CONVERTS["pooling"]
@@ -296,23 +295,12 @@ class OmniModelConfig(ModelConfig):
 
         # Init pooler config if needed
         if self.runner_type == "pooling":
-            if self.override_pooler_config is not None:
-                logger.warning_once(
-                    "`override_pooler_config` is deprecated and will be "
-                    "removed in v0.12.0 or v1.0.0, whichever is sooner. "
-                    "Please use `pooler_config` instead."
-                )
-
-                if isinstance(self.override_pooler_config, dict):
-                    self.pooler_config = PoolerConfig(**self.override_pooler_config)
-                else:
-                    self.pooler_config = self.override_pooler_config
-
             if self.pooler_config is None:
                 self.pooler_config = PoolerConfig()
 
             base_config = get_pooling_config(self.model, self.revision)
             if base_config is not None:
+                # Only set values that are not overridden by the user
                 for k, v in base_config.items():
                     if getattr(self.pooler_config, k) is None:
                         setattr(self.pooler_config, k, v)
@@ -329,23 +317,6 @@ class OmniModelConfig(ModelConfig):
             revision=self.revision,
         )
 
-        # Interleaved attention is not supported by some backends in V0
-        if (
-            not self.disable_sliding_window
-            and is_interleaved(self.hf_text_config)
-            and not envs.VLLM_USE_V1
-            and (backend := envs.VLLM_ATTENTION_BACKEND) in ("XFORMERS", "FLASHINFER")
-        ):
-            logger.warning_once(
-                "%s has interleaved attention, which is currently not "
-                "supported by the %s backend. Disabling sliding window and "
-                "capping the max length to the sliding window size (%d).",
-                self.hf_text_config.model_type,
-                backend,
-                self.hf_text_config.sliding_window,
-            )
-            self.disable_sliding_window = True
-
         self.original_max_model_len = self.max_model_len
         self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
         # Init multimodal config if needed
@@ -359,12 +330,14 @@ class OmniModelConfig(ModelConfig):
 
             mm_config_kwargs = dict(
                 limit_per_prompt=limit_mm_per_prompt,
+                enable_mm_embeds=enable_mm_embeds,
                 media_io_kwargs=media_io_kwargs,
                 mm_processor_kwargs=mm_processor_kwargs,
                 mm_processor_cache_gb=mm_processor_cache_gb,
                 mm_processor_cache_type=mm_processor_cache_type,
                 mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
                 mm_encoder_tp_mode=mm_encoder_tp_mode,
+                mm_encoder_attn_backend=mm_encoder_attn_backend,
                 interleave_mm_strings=interleave_mm_strings,
                 skip_mm_profiling=skip_mm_profiling,
                 video_pruning_rate=video_pruning_rate,
@@ -374,12 +347,20 @@ class OmniModelConfig(ModelConfig):
 
             self.multimodal_config = MultiModalConfig(**mm_config_kwargs)
 
+        # Multimodal GGUF models must use original repo for mm processing
+        if is_gguf(self.tokenizer) and self.is_multimodal_model:
+            raise ValueError(
+                "Loading a multimodal GGUF model needs to use original "
+                "tokenizer. Please specify the unquantized hf model's "
+                "repo name or path using the --tokenizer argument."
+            )
+
         if self.disable_sliding_window:
+            # Set after get_and_verify_max_len to ensure that max_model_len
+            # can be correctly capped to sliding window size
             self.hf_text_config.sliding_window = None
 
-        if not self.skip_tokenizer_init:
-            self._verify_tokenizer_mode()
-
+        # Avoid running try_verify_and_update_config multiple times
         self.config_updated = False
 
         self._verify_quantization()

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from time import time
 
+from vllm.distributed.kv_events import KVEventBatch
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
@@ -63,14 +66,6 @@ class OmniARScheduler(VLLMScheduler):
         scheduler_output: SchedulerOutput,
         model_runner_output: ModelRunnerOutput,
     ) -> dict[int, EngineCoreOutputs]:
-        """Override update_from_output to fix check_stop bug when output_token_ids is empty.
-
-        The original vLLM implementation calls check_stop when pooler_outputs exist,
-        but doesn't check if output_token_ids is empty. This causes IndexError during
-        prefill phase when pooler_outputs exist but output_token_ids is still empty.
-
-        This override adds a safety check before calling check_stop.
-        """
         sampled_token_ids = model_runner_output.sampled_token_ids
         logprobs = model_runner_output.logprobs
         prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
@@ -81,7 +76,20 @@ class OmniARScheduler(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
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        if kv_connector_stats and self.connector:
+            kv_stats = self.connector.get_kv_connector_stats()
+            if kv_stats:
+                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
+
+        failed_kv_load_req_ids = None
+        if kv_connector_output and kv_connector_output.invalid_block_ids:
+            # These blocks contain externally computed tokens that failed to
+            # load. Identify affected requests and adjust their computed token
+            # count to trigger recomputation of the invalid blocks.
+            failed_kv_load_req_ids = self._handle_invalid_blocks(kv_connector_output.invalid_block_ids)
 
         # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
         # the below loop can be a performance bottleneck. We should do our best
@@ -90,6 +98,9 @@ class OmniARScheduler(VLLMScheduler):
         stopped_preempted_reqs: set[Request] = set()
         for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
             assert num_tokens_scheduled > 0
+            if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
+                # Skip requests that were recovered from KV load failure
+                continue
             request = self.requests.get(req_id)
             if request is None:
                 # The request is already finished. This can happen if the
@@ -110,9 +121,16 @@ class OmniARScheduler(VLLMScheduler):
                 # tokens and rejections. If some tokens are rejected,
                 # num_computed_tokens is decreased by the number of rejected
                 # tokens.
-                request.num_computed_tokens -= num_rejected
+                if request.num_computed_tokens > 0:
+                    request.num_computed_tokens -= num_rejected
+                # If async scheduling, num_output_placeholders also includes
+                # the scheduled spec tokens count and so is similarly adjusted.
+                if request.num_output_placeholders > 0:
+                    request.num_output_placeholders -= num_rejected
                 spec_decoding_stats = self.make_spec_decoding_stats(
-                    spec_decoding_stats, num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted
+                    spec_decoding_stats,
+                    num_draft_tokens=num_draft_tokens,
+                    num_accepted_tokens=num_accepted,
                 )
 
             stopped = False
@@ -125,28 +143,12 @@ 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]
-                stopped = check_stop(request, self.max_model_len,
-                                     pooler_output)
-            """
-
-            # Stop checking for pooler models.
-            # FIX: Check if output_token_ids is empty before calling check_stop
-            # to prevent IndexError during prefill phase when pooler_outputs exist
-            # but output_token_ids is still empty.
-            pooler_output = None
-            if pooler_outputs:
-                pooler_output = pooler_outputs[req_index]
-                # Only call check_stop if output_token_ids is not empty
-                # This prevents IndexError when accessing output_token_ids[-1]
                 if request.output_token_ids:
                     stopped = check_stop(request, self.max_model_len, pooler_output)
-                # If output_token_ids is empty, we don't call check_stop
-                # and keep stopped=False (don't stop during prefill phase)
 
             if stopped:
                 kv_transfer_params = self._free_request(request)
@@ -157,17 +159,13 @@ class OmniARScheduler(VLLMScheduler):
 
             # Extract sample logprobs if needed.
             if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
-                # NOTE: once we support N tokens per step (spec decode),
-                # the outer lists can be of length > 1.
-                new_logprobs = logprobs.slice(req_index, req_index + 1)
+                new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
 
             if new_token_ids and self.structured_output_manager.should_advance(request):
-                # NOTE: structured_output_request
-                # should not be None if use_structured_output, we have
-                # checked above, so safe to ignore type warning
-                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]
-                    req_id, new_token_ids
-                )
+                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)
 
             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]
@@ -189,6 +187,7 @@ class OmniARScheduler(VLLMScheduler):
                         kv_transfer_params=kv_transfer_params,
                         trace_headers=request.trace_headers,
                         num_cached_tokens=request.num_cached_tokens,
+                        num_nans_in_logits=request.num_nans_in_logits,
                     )
                 )
             else:
@@ -203,8 +202,25 @@ class OmniARScheduler(VLLMScheduler):
             self.waiting.remove_requests(stopped_preempted_reqs)
 
         # KV Connector: update state for finished KV Transfers.
-        if model_runner_output.kv_connector_output:
-            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
+        if kv_connector_output:
+            self._update_from_kv_xfer_finished(kv_connector_output)
+
+        # collect KV cache events from KV cache manager
+        events = self.kv_cache_manager.take_events()
+
+        # collect KV cache events from connector
+        if self.connector is not None:
+            connector_events = self.connector.take_events()
+            if connector_events:
+                if events is None:
+                    events = list(connector_events)
+                else:
+                    events.extend(connector_events)
+
+        # publish collected KV cache events
+        if events:
+            batch = KVEventBatch(ts=time.time(), events=events)
+            self.kv_event_publisher.publish(batch)
 
         # Create EngineCoreOutputs for all clients that have requests with
         # outputs in this step.

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -1,17 +1,15 @@
 import time
 from collections import defaultdict
 
+from vllm.distributed.kv_events import KVEventBatch
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.request_queue import create_request_queue
-from vllm.v1.core.sched.scheduler import (
-    Request,
-    RequestStatus,
-    SchedulerOutput,
-    SpecDecodingStats,
-)
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.request import Request, RequestStatus
+from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.outputs import OmniModelRunnerOutput
@@ -26,15 +24,10 @@ class OmniGenerationScheduler(VLLMScheduler):
           default vLLM scheduling.
         """
 
-        # Select requests with zero prompt and using pooling
-        # (diffusion results returned via pooler_output)
         token_budget = self.max_num_scheduled_tokens
-        capacity = self.max_num_running_reqs - len(self.running)
         scheduled_timestamp = time.monotonic()
 
         scheduled_new_reqs: list[Request] = []
-        scheduled_resumed_reqs: list[Request] = []
-        scheduled_running_reqs: list[Request] = []
 
         req_to_new_blocks: dict[str, KVCacheBlocks] = {}
         num_scheduled_tokens: dict[str, int] = {}
@@ -46,7 +39,7 @@ class OmniGenerationScheduler(VLLMScheduler):
 
         # Fast path selection and scheduling (treat all as diffusion requests,
         # independent of pooling_params)
-        while self.waiting and token_budget > 0 and capacity > 0:
+        while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs:
             request = self.waiting.peek_request()
             # Uniformly treat as diffusion. A feature flag can be added later
             # via config or request tag.
@@ -81,7 +74,6 @@ class OmniGenerationScheduler(VLLMScheduler):
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens
             token_budget -= num_new_tokens
-            capacity -= 1
             scheduled_new_reqs.append(request)
 
         # Return skipped waiting requests
@@ -96,25 +88,31 @@ class OmniGenerationScheduler(VLLMScheduler):
         num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
         if self.running:
             any_request = self.running[0]
-            num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
-                any_request, len(self.running)
-            )
+            num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request.request_id)
 
-        # Assemble SchedulerOutput
-        new_reqs_data = [
-            OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
-            for req in scheduled_new_reqs
-        ]
+        # Assemble SchedulerOutput (align with v0.12.0)
+        if self.use_v2_model_runner:
+            # No resumed reqs in fast path; pass prefill_token_ids for new reqs.
+            new_reqs_data = [
+                OmniNewRequestData.from_request(
+                    req,
+                    req_to_new_blocks[req.request_id].get_block_ids(),
+                    getattr(req, "_all_token_ids", None),
+                )
+                for req in scheduled_new_reqs
+            ]
+        else:
+            new_reqs_data = [
+                OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
+                for req in scheduled_new_reqs
+            ]
+        # No running/resumed reqs scheduled in our fast path
         cached_reqs_data = self._make_cached_request_data(
-            scheduled_running_reqs,
-            scheduled_resumed_reqs,
-            num_scheduled_tokens,
-            scheduled_spec_decode_tokens,
-            req_to_new_blocks,
-        )
-        scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
-        structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
-            scheduled_requests, scheduled_spec_decode_tokens
+            running_reqs=[],
+            resumed_reqs=[],
+            num_scheduled_tokens=num_scheduled_tokens,
+            spec_decode_tokens=scheduled_spec_decode_tokens,
+            req_to_new_blocks=req_to_new_blocks,
         )
 
         total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
@@ -128,14 +126,21 @@ class OmniGenerationScheduler(VLLMScheduler):
             num_common_prefix_blocks=num_common_prefix_blocks,
             finished_req_ids=self.finished_req_ids,
             free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
-            structured_output_request_ids=structured_output_request_ids,
-            grammar_bitmask=grammar_bitmask,
+            preempted_req_ids=set(),
         )
 
+        # Record the request ids scheduled in this step (v0.12.0 behavior).
+        self.prev_step_scheduled_req_ids.clear()
+        self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
+
         # KVTransfer: package metadata
         if self.connector is not None:
             meta = self.connector.build_connector_meta(scheduler_output)
             scheduler_output.kv_connector_metadata = meta
+        # EC Connector: package metadata
+        if self.ec_connector is not None:
+            ec_meta = self.ec_connector.build_connector_meta(scheduler_output)
+            scheduler_output.ec_connector_metadata = ec_meta
 
         # Update internal state (advance num_computed_tokens, free encoder inputs,
         # etc.)
@@ -171,6 +176,15 @@ 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)
+        if kv_connector_stats and self.connector:
+            kv_stats = self.connector.get_kv_connector_stats()
+            if kv_stats:
+                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
+
+        failed_kv_load_req_ids = None
+        if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None):
+            failed_kv_load_req_ids = self._handle_invalid_blocks(kv_connector_output.invalid_block_ids)
 
         # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
         # the below loop can be a performance bottleneck. We should do our best
@@ -179,6 +193,9 @@ class OmniGenerationScheduler(VLLMScheduler):
         stopped_preempted_reqs: set[Request] = set()
         for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
             assert num_tokens_scheduled > 0
+            if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
+                # Skip requests that were recovered from KV load failure
+                continue
             request = self.requests.get(req_id)
             if request is None:
                 # The request is already finished. This can happen if the
@@ -199,7 +216,8 @@ class OmniGenerationScheduler(VLLMScheduler):
                 # tokens and rejections. If some tokens are rejected,
                 # num_computed_tokens is decreased by the number of rejected
                 # tokens.
-                request.num_computed_tokens -= num_rejected
+                if request.num_computed_tokens > 0:
+                    request.num_computed_tokens -= num_rejected
                 spec_decoding_stats = self.make_spec_decoding_stats(
                     spec_decoding_stats,
                     num_draft_tokens=num_draft_tokens,
@@ -227,9 +245,7 @@ class OmniGenerationScheduler(VLLMScheduler):
 
             # Extract sample logprobs if needed.
             if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
-                # NOTE: once we support N tokens per step (spec decode),
-                # the outer lists can be of length > 1.
-                new_logprobs = logprobs.slice(req_index, req_index + 1)
+                new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
 
             if new_token_ids and self.structured_output_manager.should_advance(request):
                 # NOTE: structured_output_request should not be None if
@@ -245,12 +261,6 @@ class OmniGenerationScheduler(VLLMScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
-            pooling_output_dict = None
-            if pooler_output is not None:
-                # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
-                # which expects Optional[dict[str, torch.Tensor]]
-                pooling_output_dict = {"model_outputs": pooler_output}
             if new_token_ids or pooler_output is not None or kv_transfer_params:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
@@ -260,12 +270,13 @@ class OmniGenerationScheduler(VLLMScheduler):
                         finish_reason=request.get_finished_reason(),
                         new_logprobs=new_logprobs,
                         new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                        pooling_output=pooling_output_dict,
+                        pooling_output=pooler_output,
                         stop_reason=request.stop_reason,
                         events=request.take_events(),
                         kv_transfer_params=kv_transfer_params,
                         trace_headers=request.trace_headers,
                         num_cached_tokens=request.num_cached_tokens,
+                        num_nans_in_logits=request.num_nans_in_logits,
                     )
                 )
             else:
@@ -280,8 +291,21 @@ class OmniGenerationScheduler(VLLMScheduler):
             self.waiting.remove_requests(stopped_preempted_reqs)
 
         # KV Connector: update state for finished KV Transfers.
-        if model_runner_output.kv_connector_output:
-            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
+        if kv_connector_output:
+            self._update_from_kv_xfer_finished(kv_connector_output)
+
+        # Collect and publish KV cache events (align with v0.12.0)
+        events = self.kv_cache_manager.take_events()
+        if self.connector is not None:
+            connector_events = self.connector.take_events()
+            if connector_events:
+                if events is None:
+                    events = list(connector_events)
+                else:
+                    events.extend(connector_events)
+        if events:
+            batch = KVEventBatch(ts=time.time(), events=events)
+            self.kv_event_publisher.publish(batch)
 
         # Create EngineCoreOutputs for all clients that have requests with
         # outputs in this step.

vllm_omni/core/sched/output.py

@@ -29,6 +29,7 @@ class OmniNewRequestData(NewRequestData):
         cls,
         request: Request,
         block_ids: tuple[list[int], ...],
+        prefill_token_ids: list[int] | None = None,
     ) -> "OmniNewRequestData":
         """Create OmniNewRequestData from a Request object.
 
@@ -49,5 +50,6 @@ class OmniNewRequestData(NewRequestData):
             num_computed_tokens=request.num_computed_tokens,
             lora_request=request.lora_request,
             prompt_embeds=request.prompt_embeds,
+            prefill_token_ids=prefill_token_ids,
             additional_information=request.additional_information,
         )

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -13,7 +13,6 @@ from torch import nn
 from vllm.config import ModelConfig
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
     download_weights_from_hf,
@@ -22,6 +21,7 @@ from vllm.model_executor.model_loader.weight_utils import (
     maybe_download_from_modelscope,
     safetensors_weights_iterator,
 )
+from vllm.utils.torch_utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.registry import initialize_model

vllm_omni/diffusion/worker/gpu_worker.py

@@ -9,7 +9,7 @@ import zmq
 from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
-from vllm.utils import DeviceMemoryProfiler, GiB_bytes
+from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
 
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.data import (

vllm_omni/diffusion/worker/npu/npu_worker.py

@@ -14,7 +14,7 @@ from vllm.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm.logger import init_logger
-from vllm.utils import DeviceMemoryProfiler, GiB_bytes
+from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
 
 from vllm_omni.diffusion.data import (
     SHUTDOWN_MESSAGE,

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -292,7 +292,7 @@ def build_stage_connectors(
     from .config import ConnectorSpec
 
     connectors: dict[tuple[str, str], Any] = {}
-    # 将字典格式的配置转换为ConnectorSpec对象
+    # Convert dictionary-formatted config to ConnectorSpec objects
     stage_connector_specs = {}
     for input_key, config in connectors_config.items():
         if not input_key.startswith("from_stage_"):
@@ -310,7 +310,7 @@ def build_stage_connectors(
         stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec
 
     try:
-        # 使用统一的连接器创建逻辑
+        # Use unified connector creation logic
         connectors = create_connectors_from_config(stage_connector_specs)
     except Exception as exc:  # pragma: no cover - defensive logging
         # Fail fast so the stage does not start with missing connectors.

vllm_omni/engine/__init__.py

@@ -102,6 +102,10 @@ class OmniEngineCoreOutput(
     # 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

vllm_omni/engine/input_processor.py

@@ -1,6 +1,6 @@
 import time
 from collections.abc import Mapping
-from typing import Any
+from typing import Any, cast
 
 import torch
 from vllm.config import VllmConfig
@@ -9,14 +9,14 @@ from vllm.inputs.parse import split_enc_dec_inputs
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict
 from vllm.multimodal.utils import argsort_mm_positions
 from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
-from vllm.transformers_utils.tokenizer import AnyTokenizer
+from vllm.tokenizers import TokenizerLike
 from vllm.utils import length_from_prompt_token_ids_or_embeds
-from vllm.v1.engine.processor import Processor
+from vllm.v1.engine.input_processor import InputProcessor
 
 from vllm_omni.engine import (
     AdditionalInformationEntry,
@@ -29,7 +29,7 @@ from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 logger = init_logger(__name__)
 
 
-class OmniProcessor(Processor):
+class OmniInputProcessor(InputProcessor):
     """Processor for omni models, handling multimodal inputs and embeddings.
 
     Extends the base vLLM Processor with support for processing prompt
@@ -75,7 +75,7 @@ class OmniProcessor(Processor):
     def __init__(
         self,
         vllm_config: VllmConfig,
-        tokenizer: AnyTokenizer,
+        tokenizer: TokenizerLike,
         mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
     ):
         super().__init__(vllm_config, tokenizer, mm_registry)
@@ -127,8 +127,6 @@ class OmniProcessor(Processor):
             ValueError: If data_parallel_rank is out of range or prompt_embeds
                 has incorrect shape
         """
-        # TODO(woosuk): Support pooling models.
-        # TODO(woosuk): Support encoder-decoder models.
         self._validate_lora(lora_request)
         self._validate_params(params)
 
@@ -158,7 +156,7 @@ class OmniProcessor(Processor):
             # if provided.
             self._validate_multi_modal_uuids(prompt)
             if isinstance(prompt, dict):
-                mm_uuids = prompt.get("multi_modal_uuids")
+                mm_uuids = cast(MultiModalUUIDDict | None, prompt.get("multi_modal_uuids"))
             else:
                 mm_uuids = None
 
@@ -171,11 +169,13 @@ class OmniProcessor(Processor):
             tokenization_kwargs=tokenization_kwargs,
             mm_uuids=mm_uuids,
         )
+
         current_platform.validate_request(
             prompt=prompt,
             params=params,
             processed_inputs=processed_inputs,
         )
+
         eos_token_id = self.input_preprocessor.get_eos_token_id()
 
         encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
@@ -185,7 +185,7 @@ class OmniProcessor(Processor):
         # discriminated unions of TypedDicts, because of how it handles
         # inheritance of TypedDict. If we explicitly extract the items we want
         # we can avoid type errors from using `dict.get` later in the method.
-        prompt_str: str | None = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
+        _prompt_str: str | None = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
         prompt_token_ids = decoder_inputs["prompt_token_ids"] if decoder_inputs["type"] != "embeds" else None
         prompt_embeds = decoder_inputs["prompt_embeds"] if decoder_inputs["type"] == "embeds" else None
 
@@ -266,7 +266,7 @@ class OmniProcessor(Processor):
                 entries[key] = entry
             additional_information_payload = AdditionalInformationPayload(entries=entries)
 
-        return prompt_str, OmniEngineCoreRequest(
+        return OmniEngineCoreRequest(
             request_id=request_id,
             prompt_token_ids=prompt_token_ids,
             mm_features=mm_features,

vllm_omni/engine/output_processor.py

@@ -6,7 +6,7 @@ import torch
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
-from vllm.transformers_utils.tokenizer import AnyTokenizer
+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
@@ -40,13 +40,14 @@ class OmniRequestState(RequestState):
     @classmethod
     def from_new_request(
         cls,
-        tokenizer: AnyTokenizer,
+        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:
@@ -92,6 +93,7 @@ class OmniRequestState(RequestState):
             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:
@@ -196,6 +198,26 @@ class OmniRequestState(RequestState):
         if not finished and final_only:
             return None
 
+        if self.stream_interval > 1:
+            assert self.detokenizer is not None
+
+            # Send output request only when
+            # 1. It has finished, or
+            # 2. It is the first token, or
+            # 3. It has reached the stream interval number of tokens
+            if not (
+                finished
+                or self.sent_tokens_offset == 0
+                or len(self.detokenizer.output_token_ids) - self.sent_tokens_offset >= self.stream_interval
+            ):
+                return None
+
+            if self.output_kind == RequestOutputKind.DELTA:
+                # Send tokens from the offset in DELTA mode, otherwise all
+                # tokens are sent.
+                new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :]
+                self.sent_tokens_offset = len(self.detokenizer.output_token_ids)
+
         request_id = self.request_id
         output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
 
@@ -247,7 +269,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
     def __init__(
         self,
-        tokenizer: AnyTokenizer,
+        tokenizer: TokenizerLike,
         log_stats: bool,
         engine_core_output_type: str | None = None,
     ):
@@ -314,9 +336,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             request_index=request_index,
             queue=queue,
             log_stats=self.log_stats,
+            stream_interval=self.stream_interval,
         )
         self.request_states[request_id] = req_state
-        self.lora_states.add_request(req_state)
         if parent_req:
             self.parent_requests[parent_req.request_id] = parent_req
 
@@ -422,6 +444,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 parent_req = req_state.parent_req
                 if parent_req and not parent_req.child_requests:
                     self.parent_requests.pop(parent_req.request_id, None)
+                if not self.request_states:
+                    self._requests_drained.set()
                 if not eco.finished:
                     reqs_to_abort.append(req_id)
                 self._update_stats_from_finished(req_state, finish_reason, iteration_stats)

vllm_omni/entrypoints/async_omni.py

@@ -19,12 +19,13 @@ from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
+from vllm.plugins.io_processors import get_io_processor
 from vllm.sampling_params import SamplingParams
+from vllm.tokenizers import TokenizerLike, init_tokenizer_from_config
 from vllm.tracing import init_tracer
 from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
-from vllm.transformers_utils.tokenizer import AnyTokenizer, init_tokenizer_from_configs
 from vllm.usage.usage_lib import UsageContext
-from vllm.utils import Device, deprecate_kwargs
+from vllm.utils.func_utils import deprecate_kwargs
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.core_client import EngineCoreClient
 from vllm.v1.engine.exceptions import EngineDeadError
@@ -44,8 +45,8 @@ from vllm_omni.distributed.ray_utils.utils import (
     try_close_ray,
 )
 from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
+from vllm_omni.engine.input_processor import OmniInputProcessor
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
-from vllm_omni.engine.processor import OmniProcessor
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.log_utils import (
     OrchestratorMetrics,
@@ -114,6 +115,10 @@ class AsyncOmni(EngineClient):
         self.batch_timeout = cli_args.batch_timeout
         self._enable_stats: bool = bool(cli_args.log_stats)
 
+        # Pause / resume state for async RL workflows.
+        self._pause_cond = asyncio.Condition()
+        self._paused = False
+
         base_engine_args = AsyncOmniEngineArgs.from_cli_args(cli_args).__dict__.copy()
 
         if cli_args.stage_configs_path is None:
@@ -143,6 +148,23 @@ class AsyncOmni(EngineClient):
         self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
         self._initialize_stages(model, cli_args.init_sleep_seconds, cli_args.shm_threshold_bytes, cli_args.init_timeout)
 
+        # Here is a duplicated init for the first stage to get the tokenizer,
+        # input_processor, and io_processor
+        # TODO: Refactor the code to avoid the duplicated init, we need to pass
+        # the tokenizer, input_processor, and io_processor back here from the
+        # first stage.
+        vllm_config = self.stage_list[0].vllm_config
+        self.tokenizer = init_tokenizer_from_config(model_config=vllm_config.model_config)
+        self.input_processor = OmniInputProcessor(
+            vllm_config=vllm_config,
+            tokenizer=self.tokenizer,
+        )
+        self.model_config = vllm_config.model_config
+        self.io_processor = get_io_processor(
+            vllm_config,
+            self.model_config.io_processor_plugin,
+        )
+
     def _initialize_stages(
         self,
         model: str,
@@ -304,6 +326,10 @@ class AsyncOmni(EngineClient):
         Raises:
             ValueError: If sampling_params_list has incorrect length.
         """
+        # Wait until generation is resumed if the engine is paused.
+        async with self._pause_cond:
+            await self._pause_cond.wait_for(lambda: not self._paused)
+
         logger.debug("[Orchestrator] generate() called")
 
         # Start output handler on the first call to generate()
@@ -634,7 +660,7 @@ class AsyncOmni(EngineClient):
     async def get_input_preprocessor(self) -> InputPreprocessor:
         return None
 
-    async def get_tokenizer(self) -> AnyTokenizer:
+    async def get_tokenizer(self) -> TokenizerLike:
         for stage in self.stage_list:
             if stage.is_comprehension:
                 return stage.tokenizer
@@ -655,7 +681,7 @@ class AsyncOmni(EngineClient):
     async def reset_mm_cache(self) -> None:
         pass
 
-    async def reset_prefix_cache(self, device: Device | None = None) -> None:
+    async def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool:
         pass
 
     async def sleep(self, level: int = 1) -> None:
@@ -686,6 +712,54 @@ class AsyncOmni(EngineClient):
     async def stop_profile(self) -> None:
         raise NotImplementedError("stop_profile() is not implemented for AsyncOmni")
 
+    async def pause_generation(
+        self,
+        *,
+        wait_for_inflight_requests: bool = False,
+        clear_cache: bool = True,
+    ) -> None:
+        """
+        Pause generation to allow model weight updates.
+
+        New generation/encoding requests are blocked until resume.
+
+        Args:
+            wait_for_inflight_requests: When ``True`` waits for in-flight
+                requests to finish before pausing. When ``False`` (default),
+                immediately aborts any in-flight requests.
+            clear_cache: Whether to clear KV cache and prefix cache after
+                draining. Set to ``False`` to preserve cache for faster resume.
+                Default is ``True`` (clear caches).
+        """
+
+        async with self._pause_cond:
+            if self._paused:
+                return
+            self._paused = True
+
+        # Note: AsyncOmni uses a stage-based architecture without a central
+        # output_processor. For now, we simply set the pause flag and let
+        # new requests wait. In-flight requests will complete naturally.
+        # TODO: Implement request abortion for stages if needed.
+
+        # Clear cache if requested
+        if clear_cache:
+            await self.reset_prefix_cache()
+            await self.reset_mm_cache()
+
+    async def resume_generation(self) -> None:
+        """Resume generation after :meth:`pause_generation`."""
+
+        async with self._pause_cond:
+            self._paused = False
+            self._pause_cond.notify_all()  # Wake up all waiting requests
+
+    async def is_paused(self) -> bool:
+        """Return whether the engine is currently paused."""
+
+        async with self._pause_cond:
+            return self._paused
+
 
 class AsyncOmniStageLLM(AsyncLLM):
     """Async single-stage LLM engine for use within a stage worker process.
@@ -749,14 +823,6 @@ class AsyncOmniStageLLM(AsyncLLM):
         Returns:
             None
         """
-        if not envs.VLLM_USE_V1:
-            raise ValueError(
-                "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. "
-                "This should not happen. As a workaround, try using "
-                "AsyncLLMEngine.from_vllm_config(...) or explicitly set "
-                "VLLM_USE_V1=0 or 1 and report this issue on Github."
-            )
-
         # Ensure we can serialize custom transformer configs
         maybe_register_config_serialize_by_value()
 
@@ -773,28 +839,33 @@ class AsyncOmniStageLLM(AsyncLLM):
             )
 
         if self.model_config.skip_tokenizer_init:
-            self.tokenizer = None
+            tokenizer = None
         else:
             # Tokenizer (+ ensure liveness if running in another process).
-            self.tokenizer = init_tokenizer_from_configs(model_config=vllm_config.model_config)
+            tokenizer = init_tokenizer_from_config(model_config=vllm_config.model_config)
 
-        # Processor (converts Inputs --> EngineCoreRequests).
-        self.processor = OmniProcessor(
+        # InputProcessor (converts Inputs --> EngineCoreRequests).
+        self.input_processor = OmniInputProcessor(
             vllm_config=vllm_config,
-            tokenizer=self.tokenizer,
+            tokenizer=tokenizer,
             mm_registry=mm_registry,
         )
 
         # OutputProcessor (converts EngineCoreOutputs --> RequestOutput).
         self.output_processor = MultimodalOutputProcessor(
-            tokenizer=self.tokenizer,
+            tokenizer=tokenizer,
             log_stats=self.log_stats,
             engine_core_output_type=engine_args.engine_output_type,
         )
+
         if self.observability_config.otlp_traces_endpoint is not None:
             tracer = init_tracer("vllm.llm_engine", self.observability_config.otlp_traces_endpoint)
             self.output_processor.tracer = tracer
 
+        # Pause / resume state for async RL workflows.
+        self._pause_cond = asyncio.Condition()
+        self._paused = False
+
         # EngineCore (starts the engine in background process).
         self.engine_core = EngineCoreClient.make_async_mp_client(
             vllm_config=vllm_config,
@@ -862,14 +933,6 @@ class AsyncOmniStageLLM(AsyncLLM):
         client_index: int = 0,
         disable_log_requests: bool = True,  # Deprecated, will be removed
     ) -> "AsyncLLM":
-        if not envs.VLLM_USE_V1:
-            raise ValueError(
-                "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. "
-                "This should not happen. As a workaround, try using "
-                "AsyncLLMEngine.from_vllm_config(...) or explicitly set "
-                "VLLM_USE_V1=0 or 1 and report this issue on Github."
-            )
-
         # Create the LLMEngine.
         return cls(
             vllm_config=vllm_config,

vllm_omni/entrypoints/cli/main.py

@@ -16,7 +16,7 @@ def main():
         return
     else:
         from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup
-        from vllm.utils import FlexibleArgumentParser
+        from vllm.utils.argparse_utils import FlexibleArgumentParser
 
         import vllm_omni.entrypoints.cli.serve
 

vllm_omni/entrypoints/cli/serve.py

@@ -12,7 +12,7 @@ from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger
-from vllm.utils import FlexibleArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni.entrypoints.openai.api_server import omni_run_server
 

vllm_omni/entrypoints/omni_llm.py

@@ -17,7 +17,7 @@ from vllm.logger import init_logger
 from vllm.plugins.io_processors import get_io_processor
 from vllm.sampling_params import SamplingParams
 from vllm.usage.usage_lib import UsageContext
-from vllm.utils import Counter
+from vllm.utils.counter import Counter
 from vllm.v1.engine.llm_engine import LLMEngine
 
 from vllm_omni.distributed.omni_connectors import (
@@ -33,8 +33,8 @@ from vllm_omni.distributed.ray_utils.utils import (
     try_close_ray,
 )
 from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.input_processor import OmniInputProcessor
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
-from vllm_omni.engine.processor import OmniProcessor
 from vllm_omni.entrypoints.log_utils import (
     OrchestratorMetrics,
     configure_orchestrator_logger,
@@ -624,7 +624,7 @@ class OmniStageLLM(LLM):
             log_stats=self.llm_engine.log_stats,
             engine_core_output_type=engine_args.engine_output_type,
         )
-        self.llm_engine.processor = OmniProcessor(
+        self.llm_engine.input_processor = OmniInputProcessor(
             vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer
         )
         self.engine_class = type(self.llm_engine)
@@ -641,3 +641,5 @@ class OmniStageLLM(LLM):
         # Load the Input/Output processor plugin if any
         io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
         self.io_processor = get_io_processor(self.llm_engine.vllm_config, io_processor_plugin)
+        self.model_config = self.llm_engine.model_config
+        self.input_processor = self.llm_engine.input_processor

vllm_omni/entrypoints/omni_stage.py

@@ -22,7 +22,7 @@ from vllm.inputs import TextPrompt
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
 from vllm.sampling_params import SamplingParams
-from vllm.transformers_utils.tokenizer import AnyTokenizer
+from vllm.tokenizers import TokenizerLike
 from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine import EngineCoreOutput
 from vllm.v1.engine.async_llm import AsyncLLM
@@ -112,7 +112,7 @@ class OmniStage:
         """
         self.vllm_config = vllm_config
 
-    def set_tokenizer(self, tokenizer: AnyTokenizer) -> None:
+    def set_tokenizer(self, tokenizer: TokenizerLike) -> None:
         """Set the tokenizer for this stage.
 
         Args:

vllm_omni/entrypoints/openai/api_server.py

@@ -37,8 +37,8 @@ from vllm.entrypoints.openai.tool_parsers import ToolParserManager
 from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.utils import load_aware_call, with_cancellation
 from vllm.logger import init_logger
-from vllm.transformers_utils.tokenizer import MistralTokenizer
-from vllm.utils import decorate_logs
+from vllm.tokenizers import MistralTokenizer
+from vllm.utils.system_utils import decorate_logs
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
@@ -305,13 +305,8 @@ async def build_async_omni_from_stage_config(
     """
 
     # V1 AsyncLLM.
-    assert envs.VLLM_USE_V1
-
     if disable_frontend_multiprocessing:
-        logger.warning(
-            "V1 is enabled, but got --disable-frontend-multiprocessing. "
-            "To disable frontend multiprocessing, set VLLM_USE_V1=0."
-        )
+        logger.warning("V1 is enabled, but got --disable-frontend-multiprocessing.")
 
     async_omni: EngineClient | None = None
 
@@ -358,7 +353,7 @@ async def omni_init_app_state(
     state.engine_client = engine_client
     state.log_stats = not args.disable_log_stats
     state.vllm_config = vllm_config
-    model_config = vllm_config.model_config
+    _model_config = vllm_config.model_config
     state.log_stats = not args.disable_log_stats
 
     # For omni models
@@ -416,14 +411,12 @@ async def omni_init_app_state(
 
     state.openai_serving_models = OpenAIServingModels(
         engine_client=engine_client,
-        model_config=model_config,
         base_model_paths=base_model_paths,
         lora_modules=lora_modules,
     )
     await state.openai_serving_models.init_static_loras()
     state.openai_serving_chat = OmniOpenAIServingChat(
         engine_client,
-        model_config,
         state.openai_serving_models,
         args.response_role,
         request_logger=request_logger,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -61,13 +61,14 @@ from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import SamplingParams
-from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
-from vllm.transformers_utils.tokenizers import (
+from vllm.tokenizers import TokenizerLike
+from vllm.tokenizers.mistral import (
+    MistralTokenizer,
     maybe_serialize_tool_calls,
     truncate_tool_call_ids,
     validate_request_params,
 )
-from vllm.utils import as_list
+from vllm.utils.collection_utils import as_list
 
 from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures
 from vllm_omni.outputs import OmniRequestOutput
@@ -290,7 +291,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
     async def _preprocess_chat(
         self,
         request: ChatLikeRequest | ResponsesRequest,
-        tokenizer: AnyTokenizer,
+        tokenizer: TokenizerLike,
         messages: list[ChatCompletionMessageParam],
         chat_template: str | None,
         chat_template_content_format: ChatTemplateContentFormatOption,
@@ -299,7 +300,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
         tool_dicts: list[dict[str, Any]] | None = None,
         documents: list[dict[str, str]] | None = None,
         chat_template_kwargs: dict[str, Any] | None = None,
-        tool_parser: Callable[[AnyTokenizer], ToolParser] | None = None,
+        tool_parser: Callable[[TokenizerLike], ToolParser] | None = None,
         add_special_tokens: bool = False,
     ) -> tuple[
         list[ConversationMessage],
@@ -450,7 +451,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
         request_id: str,
         model_name: str,
         conversation: list[ConversationMessage],
-        tokenizer: AnyTokenizer,
+        tokenizer: TokenizerLike,
         request_metadata: RequestResponseMetadata,
     ) -> ErrorResponse | ChatCompletionResponse:
         created_time = int(time.time())
@@ -542,7 +543,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
         self,
         request: ChatCompletionRequest,
         omni_outputs: OmniRequestOutput,
-        tokenizer: AnyTokenizer,
+        tokenizer: TokenizerLike,
         conversation: list[ConversationMessage],
         role: str,
     ):

vllm_omni/inputs/preprocess.py

@@ -36,7 +36,7 @@ class OmniInputPreprocessor(InputPreprocessor):
             inputs = self._process_multimodal(
                 prompt_token_ids,
                 multi_modal_data,
-                parsed_content.get("mm_processor_kwargs"),
+                parsed_content.get("mm_processor_kwargs") or {},
                 tokenization_kwargs=tokenization_kwargs,
                 mm_uuids=mm_uuids,
             )

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -15,6 +15,12 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
+)
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 
 # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
@@ -26,12 +32,6 @@ from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
-    Qwen2_5OmniConditionalGenerationMixin,
-    Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerMultiModalProcessor,
-    Qwen2_5OmniThinkerProcessingInfo,
-)
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
 from vllm_omni.utils.platform_utils import is_npu
@@ -161,18 +161,21 @@ class Qwen2_5OmniForConditionalGeneration(
             return self.model.sampler
         return Sampler()
 
-    def get_input_embeddings(
+    def embed_input_ids(
         self,
         input_ids: torch.Tensor,
         multimodal_embeddings=None,
+        is_multimodal=None,
     ) -> torch.Tensor:
         if self.model_stage == "code2wav":
             return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
-        return self.model.get_input_embeddings(input_ids, multimodal_embeddings)
+        return self.model.embed_input_ids(
+            input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal
+        )
 
-    def get_multimodal_embeddings(self, **kwargs):
+    def embed_multimodal(self, **kwargs):
         # Delegate to thinker model for multimodal processing
-        return self.model.get_multimodal_embeddings(**kwargs)
+        return self.model.embed_multimodal(**kwargs)
 
     def last_index_of(self, list, value):
         return len(list) - 1 - list[::-1].index(value)
@@ -258,7 +261,7 @@ class Qwen2_5OmniForConditionalGeneration(
 
             # Text-only path
             return OmniOutput(
-                text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states),
+                text_hidden_states=(text_hidden_states.reshape(-1, text_hidden_states.shape[-1])),
                 multimodal_outputs=None,
             )
 
@@ -288,7 +291,7 @@ class Qwen2_5OmniForConditionalGeneration(
 
             # Ensure we have base embeddings when only ids are provided
             if inputs_embeds is None and input_ids is not None:
-                inputs_embeds = self.talker.get_input_embeddings(input_ids)
+                inputs_embeds = self.talker.embed_input_ids(input_ids)
 
             # ------- Request-scoped additional information (no cross-request concat) -------
             request_ids: list[str] | None = kwargs.get("request_ids")  # ordered
@@ -456,7 +459,7 @@ class Qwen2_5OmniForConditionalGeneration(
                         dtype=torch.bfloat16,
                     ).to(self._module_device(self.model)),
                     self.talker.thinker_to_talker_proj(
-                        self.talker.get_input_embeddings(
+                        self.talker.embed_input_ids(
                             torch.tensor([TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID])
                             .to(torch.bfloat16)
                             .to(self._module_device(self.model))
@@ -790,7 +793,7 @@ class Qwen2_5OmniForConditionalGeneration(
         output_prompt_embeds,
         output_token_ids,
     ):
-        processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings(
+        processed_output_token_embeds = output_prompt_embeds + self.talker.embed_input_ids(
             output_token_ids
         )  # for decode
         return output_token_ids, processed_output_token_embeds

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py

@@ -12,7 +12,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
-    Qwen2_5OmniConditionalGenerationMixin,
     Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
     Qwen2_5OmniThinkerProcessingInfo,
@@ -23,7 +22,6 @@ from vllm.model_executor.models.utils import (
     WeightsMapper,
     init_vllm_registered_model,
     maybe_prefix,
-    merge_multimodal_embeddings,
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
@@ -31,6 +29,8 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import Qwen2_5OmniConditionalGenerationMixin
+
 
 @MULTIMODAL_REGISTRY.register_processor(
     Qwen2_5OmniThinkerMultiModalProcessor,
@@ -58,7 +58,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         super().__init__()
         config: Qwen2_5OmniTalkerConfig = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
-
+        self.vllm_config = vllm_config
         self.prefix = prefix
         self.quant_config = quant_config
 
@@ -103,26 +103,24 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
 
         return Sampler()
 
-    def get_input_embeddings(
+    def embed_input_ids(
         self,
         input_ids: torch.Tensor,
         multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+        handle_oov_mm_token: bool = False,
     ) -> torch.Tensor:
-        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
-        if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
-            # TODO (ywang96): support overlapping modalitiy embeddings so that
-            # `use_audio_in_video` will work on V1.
-            inputs_embeds = merge_multimodal_embeddings(
-                input_ids,
-                inputs_embeds,
-                multimodal_embeddings,
-                [
-                    self.config.image_token_index,
-                    self.config.video_token_index,
-                    self.config.audio_token_index,
-                ],
-            )
-        return inputs_embeds
+        # This is to satisfy the type checker for each overload
+        if multimodal_embeddings is None or is_multimodal is None:
+            return super().embed_input_ids(input_ids)
+
+        return super().embed_input_ids(
+            input_ids,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+            handle_oov_mm_token=handle_oov_mm_token,
+        )
 
     def forward(
         self,
@@ -139,7 +137,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
             inputs_embeds = None
         elif inputs_embeds is None:
             # for profile_run:
-            inputs_embeds = self.get_input_embeddings(input_ids)
+            inputs_embeds = self.embed_input_ids(input_ids)
 
         input_ids = None
 
@@ -213,7 +211,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
                 mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs)
         return mm_input_by_modality
 
-    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
         mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py

@@ -3,7 +3,10 @@ from collections.abc import Iterable
 import torch
 from torch import nn
 from transformers import Qwen2Config
-from vllm.attention import Attention, AttentionType
+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
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
@@ -125,12 +128,15 @@ class Qwen2Attention(nn.Module):
             prefix=f"{prefix}.o_proj",
         )
 
-        self.rotary_emb = get_rope(
-            self.head_dim,
+        self.rotary_pos_emb = get_rope(
+            head_size=self.head_dim,
             rotary_dim=self.head_dim,
             max_position=max_position,
-            base=self.rope_theta,
-            rope_scaling=rope_scaling,
+            is_neox_style=True,
+            rope_parameters={
+                "base": self.rope_theta,
+                **rope_scaling,
+            },
         )
         self.attn = Attention(
             self.num_heads,
@@ -150,7 +156,7 @@ class Qwen2Attention(nn.Module):
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = self.rotary_emb(positions, q, k)
+        q, k = self.rotary_pos_emb(positions, q, k)
         attn_output = self.attn(q, k, v)
         output, _ = self.o_proj(attn_output)
         return output
@@ -294,7 +300,7 @@ class Qwen2Model(nn.Module):
         else:
             self.norm = PPMissingLayer()
 
-    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
 
     def forward(
@@ -308,7 +314,7 @@ class Qwen2Model(nn.Module):
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
             else:
-                hidden_states = self.get_input_embeddings(input_ids)
+                hidden_states = self.embed_input_ids(input_ids)
             residual = None
         else:
             assert intermediate_tensors is not None
@@ -420,8 +426,8 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
 
         self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
 
-    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.model.get_input_embeddings(input_ids)
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
 
     def forward(
         self,

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -18,6 +18,12 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.qwen3_omni_moe_thinker import (
+    Qwen3OmniMoeConditionalGenerationMixin,
+    Qwen3OmniMoeThinkerDummyInputsBuilder,
+    Qwen3OmniMoeThinkerMultiModalProcessor,
+    Qwen3OmniMoeThinkerProcessingInfo,
+)
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
@@ -28,13 +34,6 @@ from vllm.v1.sample.sampler import Sampler
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 
-from .qwen3_omni_moe_thinker import (
-    Qwen3OmniMoeConditionalGenerationMixin,
-    Qwen3OmniMoeThinkerDummyInputsBuilder,
-    Qwen3OmniMoeThinkerMultiModalProcessor,
-    Qwen3OmniMoeThinkerProcessingInfo,
-)
-
 # Special token IDs for Qwen3 Omni MoE
 # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
 
@@ -226,20 +225,21 @@ class Qwen3OmniMoeForConditionalGeneration(
             return self.model.sampler
         return Sampler()
 
-    def get_input_embeddings(
+    def embed_input_ids(
         self,
         input_ids: torch.Tensor,
         multimodal_embeddings=None,
+        is_multimodal=None,
     ) -> torch.Tensor:
-        """Get input embeddings for the active model stage."""
         if self.model_stage == "code2wav":
-            # Code2wav doesn't use text embeddings
             return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
-        return self.model.get_input_embeddings(input_ids, multimodal_embeddings)
+        return self.model.embed_input_ids(
+            input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal
+        )
 
-    def get_multimodal_embeddings(self, **kwargs):
+    def embed_multimodal(self, **kwargs):
         """Delegate to active model for multimodal processing."""
-        return self.model.get_multimodal_embeddings(**kwargs)
+        return self.model.embed_multimodal(**kwargs)
 
     # ==================== Forward Pass ====================
     def _get_talker_suppressed_tokens(self):
@@ -281,16 +281,16 @@ class Qwen3OmniMoeForConditionalGeneration(
         # ========== Stage 1: Thinker ==========
         if self.model_stage == "thinker":
             # Normalize to batched inputs if needed
-            added_batch_dim = False
+            _added_batch_dim = False
             if input_ids is not None and input_ids.ndim == 1:
                 input_ids = input_ids.unsqueeze(0)
-                added_batch_dim = True
+                _added_batch_dim = True
             if positions is not None and positions.ndim == 1:
                 positions = positions.unsqueeze(0)
-                added_batch_dim = True
+                _added_batch_dim = True
             if inputs_embeds is not None and inputs_embeds.ndim == 2:
                 inputs_embeds = inputs_embeds.unsqueeze(0)
-                added_batch_dim = True
+                _added_batch_dim = True
 
             thinker_dev = self._module_device(self.thinker)
 
@@ -301,7 +301,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     dtype=torch.long,
                     device=thinker_dev,
                 ).unsqueeze(0)
-                added_batch_dim = True
+                _added_batch_dim = True
 
             # Move to thinker device
             if input_ids is not None and input_ids.device != thinker_dev:
@@ -337,7 +337,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     device=self._module_device(self.thinker),
                     dtype=torch.long,
                 )
-                thinker_tts_embeds = self.thinker.get_input_embeddings(tts_tokens)  # [1,3,thinker_hidden]
+                thinker_tts_embeds = self.thinker.embed_input_ids(tts_tokens)  # [1,3,thinker_hidden]
                 if (
                     isinstance(thinker_tts_embeds, torch.Tensor)
                     and thinker_tts_embeds.ndim == 3
@@ -353,7 +353,7 @@ class Qwen3OmniMoeForConditionalGeneration(
 
             # Return text-only output (with multimodal sidecar)
             return OmniOutput(
-                text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states),
+                text_hidden_states=(text_hidden_states.reshape(-1, text_hidden_states.shape[-1])),
                 multimodal_outputs=multimodal_outputs,
             )
 
@@ -379,7 +379,7 @@ class Qwen3OmniMoeForConditionalGeneration(
 
             # Ensure we have base embeddings when only ids are provided
             if inputs_embeds is None and input_ids is not None:
-                inputs_embeds = self.talker.get_input_embeddings(input_ids)
+                inputs_embeds = self.talker.embed_input_ids(input_ids)
 
             # ------- Request-scoped additional information (no cross-request concat) -------
             request_ids: list[str] | None = kwargs.get("request_ids")  # ordered
@@ -748,8 +748,6 @@ class Qwen3OmniMoeForConditionalGeneration(
         - hidden_projection: Used at runtime for multimodal hidden states (audio/image/video)
           from thinker's last layer, not needed for special token initialization
         """
-        # Get embeddings from both models
-        # self.thinker_embedding = self.thinker.model.get_input_embeddings()
         self.talker_embedding = self._load_talker_embedding()
 
         # Get configuration
@@ -960,9 +958,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     device=tts_pad_embed.device,
                     dtype=torch.bfloat16,
                 ),
-                self.talker.get_input_embeddings(codec_special_tokens).to(
-                    device=tts_pad_embed.device, dtype=torch.bfloat16
-                ),
+                self.talker.embed_input_ids(codec_special_tokens).to(device=tts_pad_embed.device, dtype=torch.bfloat16),
             ),
             dim=0,
         )
@@ -1059,7 +1055,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 logitsprocs=LogitsProcessors(),
             )
         # Use active model for logits computation
-        logits = self.model.compute_logits(hidden_states, sampling_metadata)  # V, d
+        logits = self.model.compute_logits(hidden_states)  # V, d
         # Talker: suppress tokens by setting their probability to ~1e-9 (finite very small),
         # implemented by assigning their logits to log(1e-9).
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -23,6 +23,7 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeMLP
+from vllm.model_executor.models.qwen3_omni_moe_thinker import Qwen3Omni_VisionTransformer
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     WeightsMapper,
@@ -31,14 +32,12 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
-from vllm.v1.sample.metadata import SamplingMetadata
 
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import (
     Qwen3OmniMoeTalkerCodePredictor,
 )
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3MoeLLMForCausalLM,
-    Qwen3Omni_VisionTransformer,
     Qwen3OmniMoeConditionalGenerationMixin,
     Qwen3OmniMoeThinkerMultiModalProcessor,
     Qwen3OmniMoeThinkerProcessingInfo,
@@ -106,9 +105,11 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config
+        talker_config.text_config.rope_parameters = talker_config.text_config.rope_scaling
+        talker_config.text_config.rope_parameters["rope_theta"] = talker_config.text_config.rope_theta
         self.quant_config = vllm_config.quant_config
         self.prefix = prefix
-
+        self.vllm_config = vllm_config
         self.config = talker_config
         self.vocab_size = talker_config.text_config.vocab_size
         self.router_aux_loss_coef = talker_config.text_config.router_aux_loss_coef
@@ -193,7 +194,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
             pos_codes = [layer0_code]  # Start with layer 0: [batch, 1]
 
             # Initial input: [last_talker_hidden, layer0_embed]
-            layer0_embed = self.get_input_embeddings(layer0_code)
+            layer0_embed = self.embed_input_ids(layer0_code)
             prev_embed = layer0_embed  # Track previous layer embedding
             try:
                 current_input = torch.cat([last_talker_hidden, prev_embed], dim=1)  # [batch, 2, hidden_size]
@@ -260,7 +261,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
         for pos in range(seq_len):
             # Layer 0 embedding
             layer0_code = result_codes[:, 0, pos : pos + 1]  # [batch, 1]
-            layer0_embed = self.get_input_embeddings(layer0_code)  # [batch, 1, hidden_size]
+            layer0_embed = self.embed_input_ids(layer0_code)  # [batch, 1, hidden_size]
 
             # mid layers hidden states (from CodePredictor)
             mid_residual_hiddens = middle_hidden_states[pos]  # [batch, num_code_groups-2, hidden_size]
@@ -440,7 +441,6 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
     def compute_logits(
         self,
         hidden_states: torch.Tensor,
-        sampling_metadata: SamplingMetadata = None,
     ) -> torch.Tensor | None:
         """Compute logits for audio codec codes (layer 0 of RVQ).
 
@@ -471,7 +471,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
                 mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs)
         return mm_input_by_modality
 
-    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings | None:
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
         mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []
@@ -504,9 +504,14 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
                 multimodal_embeddings += tuple(audio_embeddings)
         return multimodal_embeddings
 
-    def get_input_embeddings(self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None):
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        is_multimodal: bool = False,
+    ):
         """Get the input embedding layer (for codec tokens)."""
-        return self.language_model.get_input_embeddings(input_ids, multimodal_embeddings)
+        return self.language_model.embed_input_ids(input_ids)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights for the talker model.
@@ -695,10 +700,9 @@ class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM):
 
         return forward_with_shared_expert
 
-    def get_input_embeddings(
+    def embed_input_ids(
         self,
         input_ids: torch.Tensor,
-        multimodal_embeddings: MultiModalEmbeddings | None = None,
-        generation_steps=None,
+        **kwargs: object,
     ) -> torch.Tensor:
         return self.model.codec_embedding(input_ids)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -22,15 +22,13 @@
 # limitations under the License.
 """Inference-only Qwen3-Omni-Moe model (thinker part)."""
 
-import math
-from collections.abc import Callable, Iterable, Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from functools import partial
 from typing import Any
 
 import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from packaging.version import Version
 from transformers import PretrainedConfig
 from transformers import __version__ as TRANSFORMERS_VERSION
@@ -46,52 +44,44 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (
     Qwen3OmniMoeProcessor,
 )
 from transformers.models.whisper import WhisperFeatureExtractor
-from vllm.attention.utils.fa_utils import get_flash_attn_version
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
 from vllm.logger import init_logger
-from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    ReplicatedLinear,
-    RowParallelLinear,
-)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
+    SupportsMRoPE,
     SupportsMultiModal,
     SupportsPP,
 )
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
-    Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
 )
 from vllm.model_executor.models.qwen2_5_vl import (
-    Qwen2_5_VisionAttention,
-    Qwen2_5_VisionRotaryEmbedding,
     Qwen2_5_VLProcessingInfo,
 )
-from vllm.model_executor.models.qwen2_audio import Qwen2AudioInputs, Qwen2AudioProcessingInfo
+from vllm.model_executor.models.qwen2_audio import Qwen2AudioProcessingInfo
 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,
+    _get_feat_extract_output_lengths,
+)
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     WeightsMapper,
     _merge_multimodal_embeddings,
-    flatten_bn,
     maybe_prefix,
 )
 from vllm.model_executor.models.vision import (
-    get_vit_attn_backend,
+    get_llm_pos_ids_for_vision,
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import (
-    MultiModalKwargsItems,
-)
+from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems
 from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems
 from vllm.multimodal.processing import (
     MultiModalPromptUpdates,
@@ -100,504 +90,20 @@ from vllm.multimodal.processing import (
     PromptUpdate,
     PromptUpdateDetails,
 )
-from vllm.platforms import _Backend
 from vllm.sequence import IntermediateTensors
-from vllm.v1.sample.metadata import SamplingMetadata
 
-from .qwen3_moe import Qwen3MoeModel
+from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
+    Qwen2_5OmniConditionalGenerationMixin,
+)
 
 try:
     import flash_attn
 except (ImportError, ModuleNotFoundError):
     flash_attn = None
 
-from torch import Tensor
-
 logger = init_logger(__name__)
 
 
-# Local helper functions (not available in vLLM)
-def conv3d_to_linear_weight(weight: torch.Tensor) -> torch.Tensor:
-    """Convert 3D convolutional weight to linear weight for patch embedding."""
-    # Weight shape: (out_channels, in_channels, t, h, w)
-    # Need to flatten spatial dimensions: (out_channels, in_channels * t * h * w)
-    flattened_weight = weight.flatten(1)  # Flatten all dims except out_channels
-    return flattened_weight
-
-
-def get_llm_pos_ids_for_vision(
-    start_idx: int,
-    vision_idx: int,
-    spatial_merge_size: int,
-    t_index: list[torch.Tensor],
-    grid_hs: list[torch.Tensor],
-    grid_ws: list[torch.Tensor],
-) -> torch.Tensor:
-    """Generate 3D position IDs for vision tokens (temporal, height, width)."""
-    llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
-    llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
-    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten().float()
-    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float()
-    t_index_tensor = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float()
-    llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
-    llm_pos_ids = llm_pos_ids + start_idx
-    return llm_pos_ids
-
-
-# Type aliases for audio inputs
-Qwen2_5OmniAudioFeatureInputs = Qwen2AudioInputs
-
-
-def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):
-    input_lengths_leave = input_lengths % 100
-    feat_lengths = (input_lengths_leave - 1) // 2 + 1
-    output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
-    return feat_lengths, output_lengths
-
-
-class Qwen3_VisionPatchEmbed(nn.Module):
-    def __init__(
-        self,
-        patch_size: int = 14,
-        temporal_patch_size: int = 2,
-        in_channels: int = 3,
-        hidden_size: int = 1152,
-    ) -> None:
-        super().__init__()
-        self.patch_size = patch_size
-        self.temporal_patch_size = temporal_patch_size
-        self.hidden_size = hidden_size
-
-        kernel_size = (temporal_patch_size, patch_size, patch_size)
-        self.proj = ReplicatedLinear(
-            in_channels * math.prod(kernel_size),
-            hidden_size,
-            bias=True,
-            return_bias=False,
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        L, C = x.shape
-        x = self.proj(x)
-        return x
-
-
-class Qwen3_VisionMLP(nn.Module):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: int,
-        bias: bool = False,
-        act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.linear_fc1 = ColumnParallelLinear(
-            in_features,
-            hidden_features,
-            bias=bias,
-            quant_config=quant_config,
-            return_bias=False,
-            prefix=f"{prefix}.linear_fc1",
-        )
-        self.linear_fc2 = RowParallelLinear(
-            hidden_features,
-            in_features,
-            bias=bias,
-            quant_config=quant_config,
-            return_bias=False,
-            prefix=f"{prefix}.linear_fc2",
-        )
-        self.act_fn = act_fn
-
-    def forward(self, x: torch.Tensor):
-        mlp_output = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
-        return mlp_output
-
-
-class Qwen3_VisionBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        mlp_hidden_dim: int,
-        act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
-        norm_layer: Callable[[int], nn.Module] | None = None,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        if norm_layer is None:
-            norm_layer = partial(nn.LayerNorm, eps=1e-6)
-        self.norm1 = norm_layer(dim)
-        self.norm2 = norm_layer(dim)
-        self.attn = Qwen2_5_VisionAttention(
-            embed_dim=dim,
-            num_heads=num_heads,
-            projection_size=dim,
-            quant_config=quant_config,
-            prefix=f"{prefix}.attn",
-        )
-        self.mlp = Qwen3_VisionMLP(
-            dim,
-            mlp_hidden_dim,
-            act_fn=act_fn,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.mlp",
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        rotary_pos_emb: torch.Tensor,
-        max_seqlen: int | None = None,  # Only used for Flash Attention
-        seqlens: list[int] | None = None,  # Only used for xFormers
-    ) -> torch.Tensor:
-        x = x + self.attn(
-            self.norm1(x),
-            cu_seqlens=cu_seqlens,
-            rotary_pos_emb=rotary_pos_emb,
-            max_seqlen=max_seqlen,
-            seqlens=seqlens,
-        )
-
-        x = x + self.mlp(self.norm2(x))
-        return x
-
-
-class Qwen3_VisionPatchMerger(nn.Module):
-    def __init__(
-        self,
-        d_model: int,
-        context_dim: int,
-        norm_layer: Callable[[int], nn.Module] | None = None,
-        spatial_merge_size: int = 2,
-        use_postshuffle_norm: bool = False,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.hidden_size = context_dim * (spatial_merge_size**2)
-
-        self.use_postshuffle_norm = use_postshuffle_norm
-        if self.use_postshuffle_norm:
-            context_dim = self.hidden_size
-
-        if norm_layer is None:
-            norm_layer = partial(nn.LayerNorm, eps=1e-6)
-        self.use_postshuffle_norm = use_postshuffle_norm
-        self.ln_q = norm_layer(self.hidden_size if use_postshuffle_norm else context_dim)
-        self.mlp = nn.ModuleList(
-            [
-                ColumnParallelLinear(
-                    self.hidden_size,
-                    self.hidden_size,
-                    bias=True,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.mlp.0",
-                ),
-                nn.GELU(),
-                RowParallelLinear(
-                    self.hidden_size,
-                    d_model,
-                    bias=True,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.mlp.2",
-                ),
-            ]
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.use_postshuffle_norm:
-            x = self.ln_q(x.view(-1, self.hidden_size))
-        else:
-            x = self.ln_q(x).view(-1, self.hidden_size)
-
-        mlp_fc1, mlp_act, mlp_fc2 = self.mlp
-        x_parallel, _ = mlp_fc1(x)
-        x_parallel = mlp_act(x_parallel)
-        out, _ = mlp_fc2(x_parallel)
-        return out
-
-
-class Qwen3Omni_VisionTransformer(nn.Module):
-    def __init__(
-        self,
-        vision_config,
-        norm_eps: float = 1e-6,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-        attn_backend_override: _Backend | None = None,
-    ) -> None:
-        super().__init__()
-        self.hidden_size = vision_config.hidden_size
-        self.num_heads = vision_config.num_heads
-        self.image_size = vision_config.image_size
-        self.patch_size = vision_config.patch_size
-        self.spatial_merge_size = vision_config.spatial_merge_size
-        self.spatial_merge_unit = self.spatial_merge_size**2
-        self.temporal_patch_size = vision_config.temporal_patch_size
-        self.num_grid_per_side = self.image_size // self.patch_size
-        self.apply_vit_abs_pos_embed = vision_config.apply_vit_abs_pos_embed
-        self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes
-
-        self.patch_embed = Qwen3_VisionPatchEmbed(
-            patch_size=self.patch_size,
-            temporal_patch_size=self.temporal_patch_size,
-            in_channels=vision_config.in_channels,
-            hidden_size=self.hidden_size,
-        )
-
-        # vit pos embedding, TODO: spatial_patch_size vs patch_size
-        if self.apply_vit_abs_pos_embed:
-            self.pos_embed = nn.Embedding(self.num_grid_per_side**2, self.hidden_size)
-        else:
-            self.pos_embed = nn.Parameter(torch.empty([1, self.num_grid_per_side**2, self.hidden_size]))
-
-        norm_layer = partial(nn.LayerNorm, eps=norm_eps)
-        head_dim = self.hidden_size // self.num_heads
-        self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
-
-        self.blocks = nn.ModuleList(
-            [
-                Qwen3_VisionBlock(
-                    dim=self.hidden_size,
-                    num_heads=self.num_heads,
-                    mlp_hidden_dim=vision_config.intermediate_size,
-                    act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act],
-                    norm_layer=norm_layer,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.blocks.{layer_idx}",
-                )
-                for layer_idx in range(vision_config.depth)
-            ]
-        )
-        self.merger = Qwen3_VisionPatchMerger(
-            d_model=vision_config.out_hidden_size,
-            context_dim=self.hidden_size,
-            norm_layer=norm_layer,
-            spatial_merge_size=self.spatial_merge_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.merger",
-        )
-        if self.deepstack_visual_indexes is not None:
-            self.merger_list = nn.ModuleList(
-                [
-                    Qwen3_VisionPatchMerger(
-                        d_model=vision_config.out_hidden_size,
-                        context_dim=self.hidden_size,
-                        spatial_merge_size=self.spatial_merge_size,
-                        use_postshuffle_norm=True,
-                        norm_layer=norm_layer,
-                        quant_config=quant_config,
-                        prefix=f"{prefix}.merger_list.{layer_idx}",
-                    )
-                    for layer_idx in range(len(self.deepstack_visual_indexes))
-                ]
-            )
-
-        self.attn_backend = get_vit_attn_backend(
-            head_size=head_dim,
-            dtype=torch.get_default_dtype(),
-            # attn_backend_override=attn_backend_override,
-        )
-        if self.attn_backend != _Backend.FLASH_ATTN and get_flash_attn_version() is not None:
-            self.attn_backend = _Backend.FLASH_ATTN
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self.patch_embed.proj.weight.dtype
-
-    @property
-    def device(self) -> torch.device:
-        return self.patch_embed.proj.weight.device
-
-    def rot_pos_emb(self, grid_thw):
-        pos_ids = []
-        for t, h, w in grid_thw:
-            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
-            hpos_ids = hpos_ids.reshape(
-                h // self.spatial_merge_size,
-                self.spatial_merge_size,
-                w // self.spatial_merge_size,
-                self.spatial_merge_size,
-            )
-            hpos_ids = hpos_ids.permute(0, 2, 1, 3)
-            hpos_ids = hpos_ids.flatten()
-
-            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
-            wpos_ids = wpos_ids.reshape(
-                h // self.spatial_merge_size,
-                self.spatial_merge_size,
-                w // self.spatial_merge_size,
-                self.spatial_merge_size,
-            )
-            wpos_ids = wpos_ids.permute(0, 2, 1, 3)
-            wpos_ids = wpos_ids.flatten()
-            pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
-        pos_ids = torch.cat(pos_ids, dim=0)
-        max_grid_size = grid_thw[:, 1:].max()
-        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
-        rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
-        return rotary_pos_emb
-
-    def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:
-        num_grid_per_side = self.num_grid_per_side
-        m_size = self.spatial_merge_size
-        hidden_dim = self.pos_embed.embedding_dim
-
-        outputs = []
-        for t, h, w in grid_thw:
-            h_idxs = torch.linspace(0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device)
-            w_idxs = torch.linspace(0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device)
-
-            h_floor = h_idxs.to(torch.long)
-            w_floor = w_idxs.to(torch.long)
-            h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)
-            w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)
-
-            dh = h_idxs - h_floor
-            dw = w_idxs - w_floor
-
-            # Create meshgrid view for all h, w vars
-            dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")
-            h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij")
-            h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij")
-            h_floor_grid_idx = h_floor_grid * num_grid_per_side
-            h_ceil_grid_idx = h_ceil_grid * num_grid_per_side
-
-            # original computation of weights
-            # w00 = (1 - dh_grid) * (1 - dw_grid)
-            # w01 = (1 - dh_grid) * dw_grid
-            # w10 = dh_grid * (1 - dw_grid)
-            # w11 = dh_grid * dw_grid
-            # we reuse w11 here to avoid duplicate
-            # dh_grid * dw_grid computation
-            w11 = dh_grid * dw_grid
-            w10 = dh_grid - w11
-            w01 = dw_grid - w11
-            w00 = 1 - dh_grid - dw_grid + w11
-
-            idx00 = h_floor_grid_idx + w_floor_grid
-            idx01 = h_floor_grid_idx + w_ceil_grid
-            idx10 = h_ceil_grid_idx + w_floor_grid
-            idx11 = h_ceil_grid_idx + w_ceil_grid
-
-            indices = torch.stack([idx00, idx01, idx10, idx11], dim=0).reshape(4, -1)
-            weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1)
-            weights = weights.to(dtype=self.dtype, device=self.device)
-
-            embeds = self.pos_embed(indices)
-            weighted_embeds = embeds * weights
-            p0, p1, p2, p3 = weighted_embeds.unbind(dim=0)
-            combined = p0 + p1 + p2 + p3
-
-            combined = combined.view(h * w, hidden_dim)
-            repeated = combined.unsqueeze(0).expand(t, -1, -1).contiguous()
-            repeated = repeated.view(t, h // m_size, m_size, w // m_size, m_size, hidden_dim)
-            repeated = repeated.permute(0, 1, 3, 2, 4, 5).reshape(-1, hidden_dim)
-            outputs.append(repeated)
-
-        return torch.cat(outputs, dim=0)
-
-    def compute_attn_mask_seqlen(
-        self,
-        cu_seqlens: torch.Tensor,
-    ) -> tuple[int | None, list[int] | None]:
-        max_seqlen, seqlens = None, None
-        if self.attn_backend == _Backend.FLASH_ATTN:
-            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
-        elif self.attn_backend == _Backend.XFORMERS:
-            seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
-        return max_seqlen, seqlens
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        grid_thw: list[list[int]],
-    ) -> torch.Tensor:
-        hidden_states = x.to(device=self.device, dtype=self.dtype)
-        hidden_states = self.patch_embed(hidden_states)
-
-        if self.apply_vit_abs_pos_embed:
-            pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
-            hidden_states = hidden_states + pos_embeds
-        rotary_pos_emb = self.rot_pos_emb(grid_thw)
-
-        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
-            dim=0,
-            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
-        )
-        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
-
-        hidden_states = hidden_states.unsqueeze(1)
-        rotary_pos_emb = rotary_pos_emb.to(hidden_states.device)
-        max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
-
-        hidden_states_list = []
-        deepstack_visual_indexes = self.deepstack_visual_indexes
-
-        for layer_num, blk in enumerate(self.blocks):
-            hidden_states = blk(
-                hidden_states,
-                cu_seqlens=cu_seqlens,
-                rotary_pos_emb=rotary_pos_emb,
-                max_seqlen=max_seqlen,
-                seqlens=seqlens,
-            )
-            if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes:
-                hidden_states_list.append(hidden_states)
-
-        hidden_states = self.merger(hidden_states)
-
-        # processing deepstack
-        if deepstack_visual_indexes is not None:
-            processed_hidden_states_list = [hidden_states]
-            for idx, x in enumerate(hidden_states_list):
-                x = self.merger_list[idx](x)
-                processed_hidden_states_list.append(x)
-            # we cat the original visual features and deepstack features
-            # along the feature dim
-            hidden_states = torch.cat(
-                processed_hidden_states_list, dim=1
-            )  # [seq_len, hidden_size * (1 + depth_of_deepstack)]
-
-        return hidden_states
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("attn.qkv.", "attn.q.", "q"),
-            ("attn.qkv.", "attn.k.", "k"),
-            ("attn.qkv.", "attn.v.", "v"),
-        ]
-        params_dict = dict(self.named_parameters(remove_duplicate=False))
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            if name.endswith("patch_embed.proj.weight"):
-                loaded_weight = conv3d_to_linear_weight(loaded_weight)
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                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
-
-
 @support_torch_compile(
     dynamic_arg_dims={
         "input_ids": 0,
@@ -607,44 +113,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):
         "deepstack_input_embeds": 0,
     }
 )
-class Qwen3MoeLLMModel(Qwen3MoeModel):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__(vllm_config=vllm_config, prefix=prefix)
-        # from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeTextRMSNorm
-        # config = vllm_config.model_config.hf_config
-        # self.norm = Qwen3OmniMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.deepstack_multiscale_layer_start = 1
-
-    def _deepstack_process(self, hidden_states: torch.Tensor, visual_embeds: torch.Tensor) -> torch.Tensor:
-        """
-        Process deepstack visual embeddings by selectively adding them only to visual positions.
-        This matches the transformers implementation which uses visual_pos_masks.
-
-        Args:
-            hidden_states: Hidden states tensor [num_tokens, hidden_size]
-            visual_embeds: Visual embeddings tensor [num_tokens, hidden_size] (zero-padded for non-visual positions)
-
-        Returns:
-            Modified hidden_states with visual embeddings added only to visual positions
-        """
-        # Derive visual position mask from visual_embeds (non-zero positions are visual)
-        # Use a small threshold to handle floating point precision issues
-        visual_pos_masks = torch.any(torch.abs(visual_embeds) > 1e-8, dim=-1)
-
-        if visual_pos_masks.any():
-            visual_pos_masks = visual_pos_masks.to(hidden_states.device)
-            visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
-            # Match transformers implementation EXACTLY (lines 2966-2969):
-            # In transformers: local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds
-            # Note: visual_embeds in transformers is NOT indexed, meaning it's already shaped
-            # [num_visual_tokens, hidden_size]. In vllm-omni, visual_embeds is [num_tokens, hidden_size]
-            # zero-padded, so we extract only visual positions to match transformers behavior.
-            visual_embeds_only = visual_embeds[visual_pos_masks, :]  # Extract only visual positions
-            local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds_only
-            hidden_states[visual_pos_masks, :] = local_this
-
-        return hidden_states
-
+class Qwen3MoeLLMModel(_Qwen3MoeLLMModel):
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -660,7 +129,7 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
             else:
-                hidden_states = self.get_input_embeddings(input_ids)
+                hidden_states = self.embed_input_ids(input_ids)
             residual = None
         else:
             assert intermediate_tensors is not None
@@ -683,11 +152,7 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
             )
 
             if deepstack_input_embeds is not None and layer_idx in range(0, len(deepstack_input_embeds)):
-                # Use _deepstack_process to match transformers implementation
-                # This selectively adds visual embeddings only to visual positions
-                hidden_states = self._deepstack_process(
-                    hidden_states, deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]
-                )
+                hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
@@ -856,7 +321,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
                 # Derive audio placeholders from video placeholders
                 mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
             else:
-                prompt_ids, prompt, mm_placeholders = self._apply_prompt_updates(
+                prompt_ids, mm_placeholders = self._apply_prompt_updates(
                     prompt_ids,
                     mm_prompt_updates,
                 )
@@ -866,7 +331,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
                 mm_item_counts,
             )
 
-        return prompt_ids, prompt, mm_placeholders
+        return prompt_ids, mm_placeholders
 
     def get_updates_use_audio_in_video(
         self,
@@ -962,7 +427,6 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
         def get_replacement_qwen2_vision(item_idx: int, modality: str):
             grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx]
-            # grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx]
             assert isinstance(grid_thw, torch.Tensor)
             merge_length = image_processor.merge_size**2
 
@@ -974,8 +438,6 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
         def get_replacement_qwen2_use_audio_in_video(item_idx: int):
             nonlocal audio_in_video_item_idx
-            # audio_num_features = audio_output_lengths[audio_item_idx + item_idx]
-            # video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
             audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx]
             video_grid_thw = out_mm_data["video_grid_thw"][item_idx]
 
@@ -1022,7 +484,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
     def _derive_audio_from_video_placeholders(
         self,
         placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
-        mm_prompt_updates: Mapping[str, list[PromptUpdateDetails]],
+        mm_prompt_updates: MultiModalPromptUpdates,
     ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
         """
         Helper to derive audio placeholders from video placeholders when
@@ -1116,8 +578,6 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
         input_features = audio_input["input_features"]
         audio_feature_lengths = audio_input["audio_feature_lengths"]
 
-        audio_feature_lengths = flatten_bn(audio_feature_lengths, concat=True)
-
         audio_feat_lengths, audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)
 
         audio_outputs = self.audio_tower(
@@ -1138,6 +598,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
     nn.Module,
     SupportsMultiModal,
     SupportsPP,
+    SupportsMRoPE,
     Qwen3OmniMoeConditionalGenerationMixin,
 ):
     merge_by_field_config = True
@@ -1163,6 +624,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
+        self.vllm_config = vllm_config  # needed for torch compile forward context
         thinker_config: Qwen3OmniMoeThinkerConfig = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
         multimodal_config = vllm_config.model_config.multimodal_config
@@ -1184,24 +646,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
 
         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
-        # )
+        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,
+            attn_backend_override=attn_backend_override,
         )
         self.quant_config = quant_config
 
-        lm_vllm_config = vllm_config.with_hf_config(thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"])
-
         self.language_model = Qwen3MoeLLMForCausalLM(
-            vllm_config=lm_vllm_config,
+            vllm_config=vllm_config.with_hf_config(thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"]),
             prefix=maybe_prefix(prefix, "language_model"),
         )
 
@@ -1274,7 +730,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
     def get_language_model(self) -> torch.nn.Module:
         return self.language_model
 
-    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings | None:
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
         mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []
@@ -1298,27 +754,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
                 multimodal_embeddings += tuple(audio_embeddings)
         return multimodal_embeddings
 
-    def _get_text_embeddings(
-        self,
-        input_ids: Tensor,
-        embed_input_ids: Callable[[Tensor], Tensor],
-        *,
-        is_multimodal: Tensor | None,
-        handle_oov_mm_token: bool,
-    ) -> Tensor:
-        if handle_oov_mm_token and is_multimodal is not None:
-            is_text = ~is_multimodal
-            text_embeds = embed_input_ids(input_ids[is_text])
-
-            return torch.empty(
-                (input_ids.shape[0], text_embeds.shape[1]),
-                dtype=text_embeds.dtype,
-                device=text_embeds.device,
-            ).masked_scatter_(is_text.unsqueeze_(-1), text_embeds)
-
-        return embed_input_ids(input_ids)
-
-    def get_input_embeddings(
+    def embed_input_ids(
         self,
         input_ids: torch.Tensor,
         multimodal_embeddings: MultiModalEmbeddings | None = None,
@@ -1326,18 +762,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         is_multimodal: torch.Tensor | None = None,
         handle_oov_mm_token: bool = False,
     ) -> torch.Tensor:
-        # Compute is_multimodal if not provided and we have multimodal embeddings
-        if is_multimodal is None and multimodal_embeddings is not None and len(multimodal_embeddings) > 0:
-            # Create a boolean tensor marking which positions are multimodal tokens
-            is_multimodal = (
-                (input_ids == self.config.image_token_id)
-                | (input_ids == self.config.video_token_id)
-                | (input_ids == self.config.audio_token_id)
-            )
-
-        inputs_embeds = self._get_text_embeddings(
+        inputs_embeds = self._embed_text_input_ids(
             input_ids,
-            self.language_model.get_input_embeddings,
+            self.language_model.embed_input_ids,
             is_multimodal=is_multimodal,
             handle_oov_mm_token=handle_oov_mm_token,
         )
@@ -1406,7 +833,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
-        *,
         capture_layer_indices: Sequence[int] | None = None,
         return_hidden_states: bool = False,
         **kwargs: object,
@@ -1415,7 +841,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             inputs_embeds = None
 
         if self.use_deepstack and inputs_embeds is not None and get_pp_group().is_first_rank:
-            deepstack_input_embeds = self._get_deepstack_input_embeds(input_ids.size(1))
+            deepstack_input_embeds = self._get_deepstack_input_embeds(inputs_embeds.size(0))
         else:
             deepstack_input_embeds = None
 
@@ -1431,14 +857,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         )
 
         if inputs_embeds is not None and get_pp_group().is_first_rank:
-            self._clear_deepstack_input_embeds(input_ids.size(1))
+            self._clear_deepstack_input_embeds(inputs_embeds.size(0))
 
-        return hidden_states.unsqueeze(0), captured_hidden_states
+        return hidden_states, captured_hidden_states
 
     def compute_logits(
         self,
         hidden_states: torch.Tensor,
-        sampling_metadata: SamplingMetadata = None,
     ) -> torch.Tensor | None:
         return self.language_model.compute_logits(hidden_states)
 
@@ -1449,55 +874,48 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         )
         loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
-        # Log load summary
-        try:
-            total_bytes = 0
-            for name, param in self.named_parameters():
-                if param is not None and param.data is not None:
-                    total_bytes += param.data.numel() * param.data.element_size()
-            device = next(self.parameters()).device
-            logger.info(
-                "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
-                self.__class__.__name__,
-                True,
-                total_bytes / (1024**2),
-                str(device),
-            )
-        except Exception:
-            pass
-
         return loaded_weights
 
     def get_mrope_input_positions(
         self,
         input_tokens: list[int],
-        hf_config: PretrainedConfig,
-        image_grid_thw: list[list[int]] | torch.Tensor | None,
-        video_grid_thw: list[list[int]] | torch.Tensor | None,
-        second_per_grid_ts: list[float] | None = None,
-        context_len: int = 0,
-        seq_len: int | None = None,
-        audio_feature_lengths: torch.Tensor | None = None,
-        use_audio_in_video: bool = False,
+        mm_features: list[MultiModalFeatureSpec],
     ) -> tuple[torch.Tensor, int]:
-        config = hf_config.thinker_config
-        if isinstance(image_grid_thw, list):
-            image_grid_thw = torch.tensor(image_grid_thw)
-        if isinstance(video_grid_thw, list):
-            video_grid_thw = torch.tensor(video_grid_thw)
+        kwargs = MultiModalFeatureSpec.gather_kwargs(
+            mm_features,
+            {
+                "image_grid_thw",
+                "video_grid_thw",
+                "second_per_grid_ts",
+                "audio_feature_lengths",
+                "use_audio_in_video",
+            },
+        )
+        image_grid_thw = kwargs.get("image_grid_thw", [])
+        video_grid_thw = kwargs.get("video_grid_thw", [])
+        second_per_grid_ts = kwargs.get("second_per_grid_ts", [])
+        audio_feature_lengths = kwargs.get("audio_feature_lengths", [])
+        use_audio_in_video = any(kwargs.get("use_audio_in_video", []))
+
+        image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(image_grid_thw)
+        video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(video_grid_thw)
+
         input_ids = torch.tensor(input_tokens)
         if input_ids is None or input_ids.ndim != 1:
             raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids")
 
         seq_len = input_ids.shape[0]
-        if audio_feature_lengths is not None and not isinstance(audio_feature_lengths, torch.Tensor):
-            audio_feature_lengths = torch.as_tensor(audio_feature_lengths, dtype=torch.long)
+
+        if isinstance(audio_feature_lengths, list):
+            audio_feature_lengths = torch.tensor(audio_feature_lengths, dtype=torch.long)
+
         if not len(second_per_grid_ts) and len(video_grid_thw):
             second_per_grid_ts = 2.0
             second_per_grids = torch.ones(len(video_grid_thw), dtype=torch.float32) * second_per_grid_ts
         else:
             second_per_grids = torch.tensor(second_per_grid_ts, dtype=torch.float32)
 
+        config = self.config
         spatial_merge_size = config.vision_config.spatial_merge_size
         image_token_id = config.image_token_id
         video_token_id = config.video_token_id

vllm_omni/model_executor/models/registry.py

@@ -10,12 +10,12 @@ _OMNI_MODELS = {
         "qwen2_5_omni",
         "qwen2_5_omni_thinker",
         "Qwen2_5OmniThinkerForConditionalGeneration",
-    ),  # noqa: E501
+    ),
     "Qwen2_5OmniTalkerModel": (
         "qwen2_5_omni",
         "qwen2_5_omni_talker",
         "Qwen2_5OmniTalkerForConditionalGeneration",
-    ),  # noqa: E501
+    ),
     "Qwen2_5OmniToken2WavModel": (
         "qwen2_5_omni",
         "qwen2_5_omni_token2wav",

vllm_omni/request.py

@@ -2,7 +2,6 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING
 
 from vllm.v1.request import Request
-from vllm.v1.structured_output.request import StructuredOutputRequest
 
 if TYPE_CHECKING:
     from vllm.v1.core.kv_cache_utils import BlockHash
@@ -57,19 +56,16 @@ class OmniRequest(Request):
             request_id=request.request_id,
             client_index=request.client_index,
             prompt_token_ids=request.prompt_token_ids,
+            prompt_embeds=request.prompt_embeds,
             mm_features=request.mm_features,
             sampling_params=request.sampling_params,
             pooling_params=request.pooling_params,
             eos_token_id=request.eos_token_id,
             arrival_time=request.arrival_time,
             lora_request=request.lora_request,
-            structured_output_request=(
-                StructuredOutputRequest(sampling_params=request.sampling_params) if request.sampling_params else None
-            ),
             cache_salt=request.cache_salt,
             priority=request.priority,
             trace_headers=request.trace_headers,
             block_hasher=block_hasher,
-            prompt_embeds=request.prompt_embeds,
             additional_information=request.additional_information,
         )

vllm_omni/worker/gpu_ar_model_runner.py

@@ -6,11 +6,17 @@ and also outputs sampled tokens.
 
 from __future__ import annotations
 
+from copy import copy
+from typing import Any, NamedTuple
+
+import numpy as np
 import torch
-from vllm.forward_context import BatchDescriptor
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
-from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import AsyncModelRunnerOutput
+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
 from vllm.v1.worker.gpu_model_runner import (
@@ -20,7 +26,6 @@ from vllm.v1.worker.gpu_model_runner import (
     get_pp_group,
     get_tp_group,
     has_kv_transfer_group,
-    set_forward_context,
 )
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
@@ -30,14 +35,25 @@ from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 logger = init_logger(__name__)
 
 
+class ExecuteModelState(NamedTuple):
+    scheduler_output: SchedulerOutput
+    logits: torch.Tensor | None
+    spec_decode_metadata: Any
+    spec_decode_common_attn_metadata: Any
+    hidden_states: torch.Tensor
+    sample_hidden_states: torch.Tensor
+    aux_hidden_states: list[torch.Tensor] | None
+    ec_connector_output: Any
+    multimodal_outputs: Any
+
+
 class GPUARModelRunner(OmniGPUModelRunner):
     """Autoregressive GPU model runner that returns hidden states per request.
 
-    This runner follows the same preparation and forward path as GPUModelRunner
-    (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
-    and additionally performs lightweight sampling so that sampled tokens are
-    available in outputs. Hidden representations are taken at the same indices
-    that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
+    Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and
+    reuses Omni hooks for additional_information / multimodal outputs. This
+    class only overrides sample_tokens to expose hidden states + multimodal
+    outputs per request while keeping Async output semantics.
     """
 
     def __init__(self, *args, **kwargs):
@@ -65,19 +81,14 @@ class GPUARModelRunner(OmniGPUModelRunner):
         self,
         scheduler_output: SchedulerOutput,
         intermediate_tensors: IntermediateTensors | None = None,
-    ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
+    ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None:
         with record_function_or_nullcontext("Preprocess"):
             with self.synchronize_input_prep():
-                # Update persistent batch states.
                 self._update_states(scheduler_output)
-
-                # Decode per-request prompt_embeds / additional_information payloads
-                # (if present) into CPU tensors
                 self._decode_and_store_request_payloads(scheduler_output)
 
                 if not scheduler_output.total_num_scheduled_tokens:
                     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:
@@ -87,64 +98,83 @@ class GPUARModelRunner(OmniGPUModelRunner):
                         "it when the requests need prompt logprobs"
                     )
 
-                # Prepare the decoder inputs.
+                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,
+                )
+
+                (
+                    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,
+                )
+
+                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
+
                 (
                     attn_metadata,
-                    logits_indices,
-                    spec_decode_metadata,
-                    num_scheduled_tokens_np,
                     spec_decode_common_attn_metadata,
-                    max_query_len,
-                    ubatch_slices,
-                    num_tokens_after_padding,
-                ) = self._prepare_inputs(scheduler_output)
+                ) = 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_scheduled_tokens,
-                num_input_tokens,
-                num_tokens_across_dp,
                 input_ids,
                 inputs_embeds,
                 positions,
                 intermediate_tensors,
                 model_kwargs,
-                per_req_additional_information,
+                ec_connector_output,
             ) = self._preprocess(
                 scheduler_output,
-                num_scheduled_tokens_np,
+                num_tokens_padded,
                 intermediate_tensors,
-                ubatch_slices,
-                num_tokens_after_padding,
             )
 
-            uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
-                num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
-            )
-            batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=uniform_decode)
-            cudagraph_runtime_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(batch_descriptor)
+        if self.calculate_kv_scales:
+            cudagraph_mode = CUDAGraphMode.NONE
+            self.calculate_kv_scales = False
 
-        # This is currently to get around the assert in the DPMetadata
-        # where it wants `num_tokens_across_dp` to align with `num_tokens`
-        if ubatch_slices is not None:
-            num_input_tokens = ubatch_slices[0].num_tokens
-
-        # Run the model.
-        # Use persistent buffers for CUDA graphs.
         with (
             set_forward_context(
                 attn_metadata,
                 self.vllm_config,
-                num_tokens=num_input_tokens,
+                num_tokens=num_tokens_padded,
                 num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=batch_descriptor,
+                cudagraph_runtime_mode=cudagraph_mode,
+                batch_descriptor=batch_desc,
                 ubatch_slices=ubatch_slices,
             ),
             record_function_or_nullcontext("Forward"),
             self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
         ):
-            model_kwargs_extra = self._build_model_kwargs_extra(per_req_additional_information, num_scheduled_tokens_np)
-            model_output = self.model(
+            model_output = self._model_forward(
                 input_ids=input_ids,
                 positions=positions,
                 intermediate_tensors=intermediate_tensors,
@@ -153,10 +183,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 sampling_metadata=self.input_batch.sampling_metadata,
                 logits_index=logits_indices,
                 sampler=self.sampler,
-                **model_kwargs_extra,
             )
 
-        with record_function_or_nullcontext("Postprocess"):
+        with record_function_or_nullcontext("gpu_model_runner: postprocess"):
             if self.use_aux_hidden_state_outputs:
                 # True when EAGLE 3 is used.
                 hidden_states, aux_hidden_states = model_output
@@ -165,38 +194,42 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 hidden_states = model_output
                 aux_hidden_states = None
 
-            # hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
-            #     hidden_states
-            # )
             multimodal_outputs = model_output.multimodal_outputs
             hidden_states = model_output.text_hidden_states
-            # The model side may return per-request additional_information updates (model-agnostic channel).
-            # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
-            # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
-            self._process_additional_information_updates(multimodal_outputs)
+
+            if multimodal_outputs is not None:
+                keys_or_type = (
+                    list(multimodal_outputs.keys())
+                    if isinstance(multimodal_outputs, dict)
+                    else type(multimodal_outputs)
+                )
+                logger.debug(f"[AR] execute_model: multimodal_outputs keys = {keys_or_type}")
+            else:
+                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
                     return hidden_states
 
                 if self.is_pooling_model:
-                    # Return the pooling output.
-                    output = self._pool(hidden_states, num_scheduled_tokens, num_scheduled_tokens_np)
+                    output = self._pool(
+                        hidden_states,
+                        num_tokens_padded,
+                        num_scheduled_tokens_np,
+                    )
                     output.kv_connector_output = kv_connector_output
                     return output
 
                 sample_hidden_states = hidden_states[logits_indices]
                 logits = self.model.compute_logits(sample_hidden_states)
             else:
-                # Rare case.
                 assert not self.is_pooling_model
 
                 if not get_pp_group().is_last_rank:
                     all_gather_tensors = {
-                        "residual": not is_residual_scattered_for_sp(self.vllm_config, num_input_tokens)
+                        "residual": not is_residual_scattered_for_sp(self.vllm_config, num_tokens_padded)
                     }
                     get_pp_group().send_tensor_dict(
                         hidden_states.tensors,
@@ -208,26 +241,71 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     sample_hidden_states = hidden_states[logits_indices]
                     logits = self.model.compute_logits(sample_hidden_states)
 
-                model_output_broadcast_data = {}
+                model_output_broadcast_data: dict[str, Any] = {}
                 if logits is not None:
                     model_output_broadcast_data["logits"] = logits.contiguous()
 
-                model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
+                broadcasted = get_pp_group().broadcast_tensor_dict(
                     model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
                 )
-                assert model_output_broadcast_data is not None
-                logits = model_output_broadcast_data["logits"]
+                assert broadcasted is not None
+                logits = broadcasted["logits"]
 
-            # Apply structured output bitmasks if present
-            if scheduler_output.grammar_bitmask is not None:
-                apply_grammar_bitmask(scheduler_output, self.input_batch, logits, self.device)
+        self.execute_model_state = ExecuteModelState(
+            scheduler_output,
+            logits,
+            spec_decode_metadata,
+            spec_decode_common_attn_metadata,
+            hidden_states,
+            sample_hidden_states,
+            aux_hidden_states,
+            ec_connector_output,
+            multimodal_outputs,
+        )
+        self.kv_connector_output = kv_connector_output
+        return None
 
-        with record_function_or_nullcontext("Sample"):
+    @torch.inference_mode()
+    def sample_tokens(
+        self,
+        grammar_output: GrammarOutput | None,
+    ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
+        kv_connector_output = self.kv_connector_output
+        self.kv_connector_output = None
+
+        if self.execute_model_state is None:
+            if not kv_connector_output:
+                return None  # type: ignore[return-value]
+            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
+
+        (
+            scheduler_output,
+            logits,
+            spec_decode_metadata,
+            spec_decode_common_attn_metadata,
+            hidden_states,
+            sample_hidden_states,
+            aux_hidden_states,
+            ec_connector_output,
+            multimodal_outputs,
+        ) = self.execute_model_state
+        self.execute_model_state = None
+
+        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.input_batch.prev_sampled_token_ids = None
+
         def propose_draft_token_ids(sampled_token_ids):
             assert spec_decode_common_attn_metadata is not None
-            with record_function_or_nullcontext("Draft"):
+            with record_function_or_nullcontext("gpu_model_runner: draft"):
                 self._draft_token_ids = self.propose_draft_token_ids(
                     scheduler_output,
                     sampled_token_ids,
@@ -239,30 +317,40 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     spec_decode_common_attn_metadata,
                 )
 
+        spec_config = self.speculative_config
         use_padded_batch_for_eagle = (
-            self.speculative_config
-            and self.speculative_config.use_eagle()
-            and not self.speculative_config.disable_padded_drafter_batch
+            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 (
-            self.speculative_config
-            and self.speculative_config.draft_model_config is not None
-            and self.speculative_config.draft_model_config.max_model_len is not None
+            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 = self.speculative_config.draft_model_config.max_model_len
+            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.seq_lens.max() + self.speculative_config.num_speculative_tokens
-            <= effective_drafter_max_model_len
+            spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens <= effective_drafter_max_model_len
         )
-        if use_padded_batch_for_eagle and input_fits_in_drafter:
-            # EAGLE speculative decoding can use the GPU sampled tokens
-            # as inputs, and does not need to wait for bookkeeping to finish.
-            propose_draft_token_ids(sampler_output.sampled_token_ids)
+        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)
 
-        with record_function_or_nullcontext("Bookkeep"):
+        with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
             (
                 num_nans_in_logits,
                 logprobs_lists,
@@ -276,82 +364,91 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 sampler_output,
                 logits,
                 hidden_states,
-                num_scheduled_tokens,
+                scheduler_output.total_num_scheduled_tokens,
+                spec_decode_metadata,
             )
 
         if self.speculative_config and not use_padded_batch_for_eagle and input_fits_in_drafter:
-            # 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("EPLB"):
+        with record_function_or_nullcontext("gpu_model_runner: eplb"):
             self.eplb_step()
 
-        # Convert to per-request tensors on CPU
+        self._process_additional_information_updates(multimodal_outputs)
+
         hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
-        # pooler_output: list[torch.Tensor | None] = []
-        pooler_output: list[torch.Tensor | None] = []
-        prev_logits_index = 0
-        for rid, logits_index in zip(req_ids_output_copy, logits_indices):
-            # Base payload: hidden slice for this request in this iteration
-            hidden_slice = hidden_states_cpu[prev_logits_index : logits_index + 1]
+        num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
+        if num_scheduled_tokens_np is None:
+            req_ids = self.input_batch.req_ids
+            num_scheduled_tokens_np = np.array(
+                [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids],
+                dtype=np.int32,
+            )
+
+        pooler_output: list[dict[str, object]] = []
+        for rid in req_ids_output_copy:
+            idx = req_id_to_index_output_copy[rid]
+            start = int(self.query_start_loc.cpu[idx])
+            sched = int(num_scheduled_tokens_np[idx])
+            end = start + sched
+            hidden_slice = hidden_states_cpu[start:end]
             payload: dict[str, object] = {"hidden": hidden_slice}
-            # Merge multimodal_outputs if present
             if isinstance(multimodal_outputs, dict) and multimodal_outputs:
                 mm_payload: dict[str, object] = {}
                 for k, v in multimodal_outputs.items():
                     try:
-                        # Case 1: tensor aligned on token dimension
                         if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
-                            mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
-                        elif isinstance(v, torch.Tensor) and v.shape[0] != hidden_states_cpu.shape[0]:
-                            logger.error(
-                                f"Error in merge multimodal outputs: Tensor dimension mismatch, \
-                                          {v.shape} != {hidden_states_cpu.shape} for {k}"
-                            )
-                        # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers)
+                            mm_payload[k] = v.detach().to("cpu")[start:end].contiguous()
                         elif isinstance(v, dict):
                             sub_dict: dict[str, torch.Tensor] = {}
                             for sk, sv in v.items():
                                 if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
-                                    sub_dict[str(sk)] = (
-                                        sv.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
-                                    )
+                                    sub_dict[str(sk)] = sv.detach().to("cpu")[start:end].contiguous()
                             if sub_dict:
                                 mm_payload[k] = sub_dict
                         elif isinstance(v, list):
                             element = v[0]
                             if isinstance(element, torch.Tensor):
                                 element = element.detach().to("cpu").contiguous()
-                            multimodal_outputs[k] = v[1:] if len(v) > 1 else v
                             mm_payload[k] = element
                     except Exception as e:
-                        # Best-effort; skip malformed entries
                         logger.error(f"Error in merge multimodal outputs: {e}")
                 if mm_payload:
                     payload.update(mm_payload)
-            pooler_output.append(payload)  # type: ignore[arg-type]
-            prev_logits_index = logits_index + 1
-        output = OmniModelRunnerOutput(
-            req_ids=req_ids_output_copy,
-            req_id_to_index=req_id_to_index_output_copy,
-            sampled_token_ids=valid_sampled_token_ids,
-            logprobs=logprobs_lists,
-            prompt_logprobs_dict=prompt_logprobs_dict,
-            pooler_output=(pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None),
-            kv_connector_output=kv_connector_output,
-            num_nans_in_logits=num_nans_in_logits,
-        )
+            pooler_output.append(payload)
+        with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
+            output = OmniModelRunnerOutput(
+                req_ids=req_ids_output_copy,
+                req_id_to_index=req_id_to_index_output_copy,
+                sampled_token_ids=valid_sampled_token_ids,
+                logprobs=logprobs_lists,
+                prompt_logprobs_dict=prompt_logprobs_dict,
+                pooler_output=(pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None),
+                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,
+            )
 
         if not self.use_async_scheduling:
             return output
+        with record_function_or_nullcontext("gpu_model_runner: AsyncGPUModelRunnerOutput"):
+            async_output = AsyncGPUModelRunnerOutput(
+                model_runner_output=output,
+                sampled_token_ids=sampler_output.sampled_token_ids,
+                logprobs_tensors=sampler_output.logprobs_tensors,
+                invalid_req_indices=invalid_req_indices,
+                async_output_copy_stream=self.async_output_copy_stream,
+                vocab_size=self.input_batch.vocab_size,
+            )
+        with record_function_or_nullcontext("gpu_model_runner: set_async_sampled_token_ids"):
+            # Save ref of sampled_token_ids CPU tensor if the batch contains
+            # any requests with sampling params that require output ids.
+            self.input_batch.set_async_sampled_token_ids(
+                async_output.sampled_token_ids_cpu,
+                async_output.async_copy_ready_event,
+            )
 
-        return AsyncGPUModelRunnerOutput(
-            model_runner_output=output,
-            sampled_token_ids=sampler_output.sampled_token_ids,
-            invalid_req_indices=invalid_req_indices,
-            async_output_copy_stream=self.async_output_copy_stream,
-        )
+        return async_output
 
     def _merge_additional_information_update(self, req_id: str, upd: dict) -> None:
         req_state = self.requests.get(req_id)
@@ -365,13 +462,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
             if isinstance(v, torch.Tensor):
                 merged[k] = v.detach().to("cpu").contiguous()
             elif isinstance(v, list):
-                new_list = []
-                for item in v:
-                    if isinstance(item, torch.Tensor):
-                        new_list.append(item.detach().to("cpu").contiguous())
-                    else:
-                        new_list.append(item)
-                merged[k] = new_list
+                merged[k] = [
+                    (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v
+                ]
             else:
                 merged[k] = v
         setattr(req_state, "additional_information_cpu", merged)

vllm_omni/worker/gpu_ar_worker.py

@@ -4,7 +4,8 @@ import os
 import torch
 from vllm.model_executor import set_random_seed
 from vllm.platforms import current_platform
-from vllm.utils import GiB_bytes, MemorySnapshot
+from vllm.utils.mem_constants import GiB_bytes
+from vllm.utils.mem_utils import MemorySnapshot
 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
@@ -20,19 +21,37 @@ class GPUARWorker(GPUWorker):
     """
 
     def init_device(self):
-        """Initialize the GPU device and autoregressive model runner.
-
-        Sets up CUDA device, initializes distributed environment,
-        sets random seed, and creates a GPUARModelRunner instance.
-
-        Raises:
-            ValueError: If free GPU memory is insufficient for the
-                requested memory utilization
-            RuntimeError: If device type is not supported
-        """
-        if self.device_config.device.type == "cuda":
+        device = self.device_config.device
+        if isinstance(device, torch.device) and device.type == "cuda":
             # This env var set by Ray causes exceptions with graph building.
             os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+            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
+            ):
+                # 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
+
+                tp_pp_world_size = (
+                    self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
+                )
+
+                # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
+                self.local_rank += dp_local_rank * tp_pp_world_size
+                assert self.local_rank < torch.cuda.device_count(), (
+                    f"DP adjusted local rank {self.local_rank} is out of bounds. "
+                )
+                visible_device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
+                assert self.parallel_config.local_world_size <= visible_device_count, (
+                    f"local_world_size ({self.parallel_config.local_world_size}) must "
+                    f"be less than or equal to the number of visible devices "
+                    f"({visible_device_count})."
+                )
             self.device = torch.device(f"cuda:{self.local_rank}")
             current_platform.set_device(self.device)
 
@@ -62,23 +81,23 @@ class GPUARWorker(GPUWorker):
             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(b):
-                    return round(b / GiB_bytes, 2)
+                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"({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"{gib(self.requested_memory)} GiB). Decrease GPU memory "
                     f"utilization or reduce GPU memory used by other processes."
                 )
         else:
             raise RuntimeError(f"Not support device type: {self.device_config.device}")
 
         # Construct the model runner
-        self.model_runner: GPUARModelRunner = GPUARModelRunner(self.vllm_config, self.device)
+        self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.

vllm_omni/worker/gpu_generation_model_runner.py

@@ -8,14 +8,12 @@ from __future__ import annotations
 
 import gc
 import logging
-from typing import Any
 
 import numpy as np
 import torch
 from vllm.config import CUDAGraphMode
-from vllm.forward_context import BatchDescriptor
 from vllm.multimodal.inputs import MultiModalKwargs
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.utils import record_function_or_nullcontext
@@ -27,7 +25,6 @@ from vllm.v1.worker.gpu_model_runner import (
     get_pp_group,
     set_forward_context,
 )
-from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
 from vllm_omni.outputs import OmniModelRunnerOutput
@@ -44,110 +41,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
     - Executes generation process and returns tensors via `pooler_output`.
     """
 
-    def _preprocess(
-        self,
-        scheduler_output: SchedulerOutput,
-        intermediate_tensors: IntermediateTensors | None = None,
-        ubatch_slices: UBatchSlices | None = None,
-        num_tokens_after_padding: torch.Tensor | None = None,
-    ) -> tuple[
-        int,
-        int,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor,
-        IntermediateTensors | None,
-        dict[str, Any],
-    ]:
-        # Input token count for this iteration (not used by diffusion, but
-        # retained to keep DP padding/ordering consistent)
-        num_input_tokens = scheduler_output.total_num_scheduled_tokens
-        num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
-        num_input_tokens += num_pad
-
-        # _prepare_inputs may reorder the batch, so we must gather multi
-        # modal outputs after that to ensure the correct order
-        if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
-            # Run the multimodal encoder if any.
-            self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
-
-            # NOTE(woosuk): To unify token ids and soft tokens (vision
-            # embeddings), we always use embeddings (rather than token ids)
-            # as input to the multimodal model, even when the input is text.
-            inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids.gpu[:num_input_tokens],
-                multimodal_embeddings=mm_embeds or None,
-            )
-
-            # TODO(woosuk): Avoid the copy. Optimize.
-            self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled)
-
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = {
-                **self._init_model_kwargs(num_input_tokens),
-                **self._extract_mm_kwargs(scheduler_output),
-            }
-        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
-            # Get the input embeddings for the tokens that are not input embeds,
-            # then put them into the appropriate positions.
-            # TODO(qthequartermasterman): Since even when prompt embeds are
-            # enabled, (a) not all requests will use prompt embeds, and (b)
-            # after the initial prompt is processed, the rest of the generated
-            # tokens will be token ids, it is not desirable to have the
-            # embedding layer outside of the CUDA graph all the time. The v0
-            # engine avoids this by "double compiling" the CUDA graph, once
-            # with input_ids and again with inputs_embeds, for all num_tokens.
-            # If a batch only has token ids, then including the embedding layer
-            # in the CUDA graph will be more performant (like in the else case
-            # below).
-            token_ids_idx = self.is_token_ids.gpu[:num_input_tokens].nonzero(as_tuple=False).squeeze(1)
-            # Some tokens ids may need to become embeds
-            if token_ids_idx.numel() > 0:
-                token_ids = self.input_ids.gpu[token_ids_idx]
-                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
-                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)
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-        else:
-            # For text-only models, we use token ids as input.
-            # While it is possible to use embeddings as input just like the
-            # multimodal models, it is not desirable for performance since
-            # 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)
-        if self.uses_mrope:
-            positions = self.mrope_positions.gpu[:, :num_input_tokens]
-        else:
-            positions = self.positions.gpu[:num_input_tokens]
-
-        if get_pp_group().is_first_rank:
-            intermediate_tensors = None
-        else:
-            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True
-            )
-
-        if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
-            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
-            model_kwargs.update(encoder_inputs)
-
-        return (
-            num_input_tokens,
-            num_input_tokens,
-            num_tokens_after_padding,
-            input_ids,
-            inputs_embeds,
-            positions,
-            intermediate_tensors,
-            model_kwargs,
-        )
-
     @torch.inference_mode()
     def execute_model(
         self,
@@ -156,52 +49,81 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
     ) -> OmniModelRunnerOutput | IntermediateTensors:
         with record_function_or_nullcontext("Preprocess"):
             with self.synchronize_input_prep():
-                # Update persistent batch states.
                 self._update_states(scheduler_output)
                 if not scheduler_output.total_num_scheduled_tokens:
                     return EMPTY_MODEL_RUNNER_OUTPUT
 
-                # Prepare the decoder inputs.
+                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,
+                )
+
+                (
+                    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,
+                )
+
+                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
+
                 (
                     attn_metadata,
-                    logits_indices,
-                    spec_decode_metadata,
-                    num_scheduled_tokens_np,
                     spec_decode_common_attn_metadata,
-                    max_query_len,
-                    ubatch_slices,
-                    num_tokens_after_padding,
-                ) = self._prepare_inputs(scheduler_output)
+                ) = 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,
+                )
 
-            # Input token count for this iteration (not used by diffusion, but
-            # retained to keep DP padding/ordering consistent)
             (
-                num_scheduled_tokens,
-                num_input_tokens,
-                num_tokens_across_dp,
                 input_ids,
                 inputs_embeds,
                 positions,
                 intermediate_tensors,
                 model_kwargs,
+                ec_connector_output,
             ) = self._preprocess(
                 scheduler_output,
+                num_tokens_padded,
                 intermediate_tensors,
-                ubatch_slices,
-                num_tokens_after_padding,
             )
 
-        cudagraph_runtime_mode = CUDAGraphMode.NONE
-        # Run the model.
-        # Use persistent buffers for CUDA graphs.
+        if self.calculate_kv_scales:
+            cudagraph_mode = CUDAGraphMode.NONE
+            self.calculate_kv_scales = False
+
         with (
             set_forward_context(
                 attn_metadata,
                 self.vllm_config,
-                num_tokens=num_input_tokens,
+                num_tokens=num_tokens_padded,
                 num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=None,
+                cudagraph_runtime_mode=cudagraph_mode,
+                batch_descriptor=batch_desc,
                 ubatch_slices=ubatch_slices,
             ),
             record_function_or_nullcontext("Forward"),
@@ -212,27 +134,32 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 positions=positions,
                 intermediate_tensors=intermediate_tensors,
                 inputs_embeds=inputs_embeds,
-                multimodal_kwargs=model_kwargs,
+                model_kwargs=model_kwargs,
                 logits_indices=logits_indices,
             )
 
         _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
-        # Ensure one tensor per request, map to CPU for output struct
-        pooler_output: list[torch.Tensor | None] = []
+        pooler_output: list[object] = []
         if isinstance(multimodal_outputs, torch.Tensor):
-            # If model returned a single stacked tensor, split by requests
+            assert multimodal_outputs.shape[0] == 1, (
+                "model should return a single tensor, to return multiple tensors, use a dict"
+            )
             assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
             for i in range(self.input_batch.num_reqs):
-                pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
+                pooler_output.append({"model_outputs": multimodal_outputs[i].detach().to("cpu").contiguous()})
         elif isinstance(multimodal_outputs, list):
+            assert len(multimodal_outputs) == 1, (
+                "model should return a single list, to return multiple lists, use a dict"
+            )
             for out in multimodal_outputs:
-                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+                pooler_output.append(
+                    {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
+                )
         elif isinstance(multimodal_outputs, dict):
-            for out in multimodal_outputs.values():
-                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+            for key, out in multimodal_outputs.items():
+                pooler_output.append({key: out.detach().to("cpu").contiguous() if out is not None else None})
         else:
             raise RuntimeError("Unsupported diffusion output type")
-
         output = OmniModelRunnerOutput(
             req_ids=self.input_batch.req_ids,
             req_id_to_index=self.input_batch.req_id_to_index,
@@ -257,11 +184,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
     def _run_generation_model(
         self,
         *,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor | None,
         intermediate_tensors: IntermediateTensors | None,
         inputs_embeds: torch.Tensor | None,
-        multimodal_kwargs: dict,
+        model_kwargs: dict,
         logits_indices: torch.Tensor,
     ) -> torch.Tensor | list[torch.Tensor]:
         """Run generation from codec codes to waveforms.
@@ -279,7 +206,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             positions=positions,
             intermediate_tensors=intermediate_tensors,
             inputs_embeds=inputs_embeds,
-            **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
+            **MultiModalKwargs.as_kwargs(model_kwargs, device=self.device),
             sampling_metadata=self.input_batch.sampling_metadata,
             logits_index=logits_indices,
             sampler=self.sampler,
@@ -305,10 +232,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         cudagraph_runtime_mode: CUDAGraphMode | None = None,
         force_attention: bool = False,
         uniform_decode: bool = False,
+        allow_microbatching: bool = True,
         skip_eplb: bool = False,
         is_profile: bool = False,
         create_mixed_batch: bool = False,
         remove_lora: bool = True,
+        activate_lora: bool = False,
+        is_graph_capturing: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -331,12 +261,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             create_mixed_batch: If True, create a mixed batch with both decode
                 (1 token) and prefill (multiple tokens) requests.
             remove_lora: If False, dummy LoRAs are not destroyed after the run
+            activate_lora: If False, dummy_run is performed without LoRAs.
         """
-        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
-            CUDAGraphMode.NONE,
-            CUDAGraphMode.PIECEWISE,
-            CUDAGraphMode.FULL,
-        }
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
 
         # If cudagraph_mode.decode_mode() == FULL and
         # cudagraph_mode.separate_routine(). This means that we are using
@@ -358,37 +285,79 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         # has num_tokens in total.
         assert num_tokens <= self.scheduler_config.max_num_batched_tokens
         max_num_reqs = self.scheduler_config.max_num_seqs
+        if create_mixed_batch:
+            assert not uniform_decode
+            # Create mixed batch:
+            # first half decode tokens, second half one prefill
+            num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2)
+            num_prefill_tokens = num_tokens - num_decode_tokens
+            num_reqs = num_decode_tokens + 1
 
-        num_reqs = min(num_tokens, max_num_reqs)
-        min_tokens_per_req = num_tokens // num_reqs
-        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
-        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+            # Create decode requests (1 token each) followed by prefill request
+            num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]
+            # Note: Overriding max_query_len to be the prefill tokens
+            max_query_len = num_prefill_tokens
+        elif uniform_decode:
+            assert not create_mixed_batch
+            num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len))
+            num_scheduled_tokens_list = [max_query_len] * num_reqs
+            if num_tokens % max_query_len != 0:
+                num_scheduled_tokens_list[-1] = num_tokens % max_query_len
+        else:
+            num_reqs = min(num_tokens, max_num_reqs)
+            min_tokens_per_req = num_tokens // num_reqs
+            num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+            num_scheduled_tokens_list[-1] += num_tokens % num_reqs
 
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
         num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+        num_tokens_unpadded = int(num_scheduled_tokens.sum())
 
-        num_tokens_after_padding = None
+        num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
 
-        # If we failed to microbatch, currently need to resynchronize
-        # TODO(lucas,sage): we should be able to avoid this second sync by
-        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
-        #  a single `coordinate_batch_across_dp` function.
-        if num_tokens_after_padding is None:
-            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-            num_tokens_after_padding = num_tokens + num_pad
+        _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,
+        )
+
+        if cudagraph_runtime_mode is None:
+            cudagraph_runtime_mode = _cudagraph_mode
         else:
-            num_tokens_across_dp = num_tokens_after_padding
-            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
+            assert cudagraph_runtime_mode == _cudagraph_mode, (
+                f"Cudagraph runtime mode mismatch in dummy_run. "
+                f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}."
+            )
+
+        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
 
         attn_metadata: PerLayerAttnMetadata | None = None
 
         # If force_attention is True, we always capture attention. Otherwise,
         # it only happens for cudagraph_runtime_mode=FULL.
         if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-            attn_metadata = {}
-
-            seq_lens = max_query_len
+            if create_mixed_batch:
+                # In the mixed batch mode (used for FI warmup), we use
+                # shorter sequence lengths to run faster.
+                # TODO(luka) better system for describing dummy batches
+                seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
+            else:
+                seq_lens = max_query_len  # type: ignore[assignment]
             self.seq_lens.np[:num_reqs] = seq_lens
             self.seq_lens.np[num_reqs:] = 0
             self.seq_lens.copy_to_gpu()
@@ -397,50 +366,45 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
             self.query_start_loc.copy_to_gpu()
 
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
-                common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
-                    seq_lens=self.seq_lens.gpu[:num_reqs],
-                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],  # noqa: E501
-                    num_reqs=num_reqs,
-                    num_actual_tokens=num_tokens,
-                    max_query_len=max_query_len,
-                    max_seq_len=self.max_model_len,
-                    block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs),
-                    slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
-                    causal=True,
-                )
-                for attn_group in self.attn_groups[kv_cache_group_id]:
-                    assert type(attn_metadata) is dict
-                    attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
-                        common_attn_metadata
-                    )  # noqa: E501
-                    for layer_name in attn_group.layer_names:
-                        attn_metadata[layer_name] = attn_metadata_i
+            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,
+                for_cudagraph_capture=is_graph_capturing,
+            )
 
-        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
-            model_kwargs = self._init_model_kwargs(num_tokens)
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config,
+            num_scheduled_tokens,
+            num_sampled_tokens,
+            activate_lora,
+            remove_lora,
+        ):
+            # 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)
             if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
                 input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
                 model_kwargs = {
                     **model_kwargs,
                     **self._dummy_mm_kwargs(num_reqs),
                 }
             elif self.enable_prompt_embeds:
                 input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
-                model_kwargs = self._init_model_kwargs(num_tokens)
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
+                model_kwargs = self._init_model_kwargs(num_tokens_padded)
             else:
-                input_ids = self.input_ids.gpu[:num_tokens]
+                input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
 
             if self.uses_mrope:
-                positions = self.mrope_positions.gpu[:, :num_tokens]
+                positions = self.mrope_positions.gpu[:, :num_tokens_padded]
+            elif self.uses_xdrope_dim > 0:
+                positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
             else:
-                positions = self.positions.gpu[:num_tokens]
+                positions = self.positions.gpu[:num_tokens_padded]
 
             if get_pp_group().is_first_rank:
                 intermediate_tensors = None
@@ -452,38 +416,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                         device=self.device,
                     )
 
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False)
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False)
 
-            # filter out the valid batch descriptor
-            _cg_mode, batch_descriptor = (
-                self.cudagraph_dispatcher.dispatch(
-                    BatchDescriptor(
-                        num_tokens=num_tokens_after_padding,
-                        uniform_decode=uniform_decode,
-                    )
-                )
-                if not is_profile
-                else (CUDAGraphMode.NONE, None)
-            )
-            if cudagraph_runtime_mode is not None:
-                # we allow forcing NONE when the dispatcher disagrees to support
-                # warm ups for cudagraph capture
-                assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
-                    f"Cudagraph runtime mode mismatch at dummy_run. "
-                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
-                )
-            else:
-                cudagraph_runtime_mode = _cg_mode
+            if ubatch_slices 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
+                if num_tokens_across_dp is not None:
+                    num_tokens_across_dp[:] = num_tokens_padded
 
             with (
                 self.maybe_randomize_inputs(input_ids),
                 set_forward_context(
                     attn_metadata,
                     self.vllm_config,
-                    num_tokens=num_tokens_after_padding,
+                    num_tokens=num_tokens_padded,
                     num_tokens_across_dp=num_tokens_across_dp,
                     cudagraph_runtime_mode=cudagraph_runtime_mode,
-                    batch_descriptor=batch_descriptor,
+                    batch_descriptor=batch_desc,
+                    ubatch_slices=ubatch_slices,
                 ),
             ):
                 outputs = self.model(
@@ -498,10 +450,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 hidden_states, _ = outputs
             else:
                 hidden_states = outputs
-
+            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)
-                self.drafter.dummy_run(num_tokens)
+                use_cudagraphs = (
+                    cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE)
+                    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
+                # short term mitigation for issue mentioned in
+                # https://github.com/vllm-project/vllm/issues/28334
+                if self.compilation_config.cudagraph_specialize_lora and activate_lora:
+                    use_cudagraphs = False
+
+                self.drafter.dummy_run(
+                    num_tokens,
+                    use_cudagraphs=use_cudagraphs,
+                    is_graph_capturing=is_graph_capturing,
+                )
 
         # This is necessary to avoid blocking DP.
         # For dummy runs, we typically skip EPLB since we don't have any real
@@ -510,15 +478,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         # not have any requests to process, so they're executing dummy batches.
         # In such cases, we still have to trigger EPLB to make sure
         # ranks execute the rearrangement in synchronization.
+        if not skip_eplb:
+            self.eplb_step(is_dummy=True, is_profile=is_profile)
 
-        # logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
         return hidden_states, None
 
     def profile_run(self) -> None:
         # Profile with multimodal encoder & encoder cache.
         if self.supports_mm_inputs:
-            if self.model_config.multimodal_config.skip_mm_profiling:
+            mm_config = self.model_config.multimodal_config
+            if mm_config is not None and mm_config.skip_mm_profiling:
                 logger.info("Skipping memory profiling for multimodal encoder and encoder cache.")
             else:
                 mm_budget = self.mm_budget
@@ -532,8 +501,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                     max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[dummy_modality]
 
                     logger.info(
-                        "Encoder cache will be initialized with a budget "
-                        "of %s tokens, and profiled with %s %s items of the maximum feature size.",
+                        "Encoder cache will be initialized with a budget of "
+                        "%s tokens, and profiled with %s %s items of the "
+                        "maximum feature size.",
                         encoder_budget,
                         max_mm_items_per_batch,
                         dummy_modality,
@@ -546,7 +516,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                     )
 
                     # Run multimodal encoder.
-                    dummy_encoder_outputs = self.model.get_multimodal_embeddings(**batched_dummy_mm_inputs)
+                    dummy_encoder_outputs = self.model.embed_multimodal(**batched_dummy_mm_inputs)
 
                     sanity_check_mm_encoder_outputs(
                         dummy_encoder_outputs,
@@ -556,13 +526,15 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                     # NOTE: This happens when encoder cache needs to store
                     # the embeddings that encoder outputs are scattered onto.
                     # In this case we create dummy embeddings of size
-                    # (encode_budget, hidden_size) and scatter encoder
-                    # output into it.
+                    # (max_tokens_for_modality, hidden_size) and scatter
+                    # encoder output into it.
                     encoder_output_shape = dummy_encoder_outputs[0].shape
-                    if encoder_output_shape[0] < encoder_budget:
+                    max_mm_tokens_per_item = mm_budget.max_tokens_by_modality[dummy_modality]
+                    if encoder_output_shape[0] < max_mm_tokens_per_item:
+                        encoder_hidden_size = encoder_output_shape[-1]
                         expanded_outputs = []
                         for output in dummy_encoder_outputs:
-                            expanded = output.new_zeros((encoder_budget, encoder_output_shape[-1]))
+                            expanded = output.new_zeros((max_mm_tokens_per_item, encoder_hidden_size))
                             num_tokens = output.shape[0]
                             expanded[:num_tokens].copy_(output)
                             expanded_outputs.append(expanded)
@@ -574,8 +546,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
 
         # Add `is_profile` here to pre-allocate communication buffers
         hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
-        if get_pp_group().is_last_rank:
-            pass
+        output = None
         self._sync_device()
         del hidden_states
         self.encoder_cache.clear()

vllm_omni/worker/gpu_generation_worker.py

@@ -4,7 +4,8 @@ import os
 import torch
 from vllm.model_executor import set_random_seed
 from vllm.platforms import current_platform
-from vllm.utils import GiB_bytes, MemorySnapshot
+from vllm.utils.mem_constants import GiB_bytes
+from vllm.utils.mem_utils import MemorySnapshot
 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
@@ -20,18 +21,37 @@ class GPUGenerationWorker(GPUWorker):
     """
 
     def init_device(self):
-        """Initialize CUDA device and distributed environment."""
-        if self.device_config.device.type == "cuda":
-            # torch.distributed.all_reduce does not free the input tensor until
-            # the synchronization point. This causes the memory usage to grow
-            # as the number of all_reduce calls increases. This env var disables
-            # this behavior.
-            # Related issue:
-            # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
-            os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
-
+        device = self.device_config.device
+        if isinstance(device, torch.device) and device.type == "cuda":
             # This env var set by Ray causes exceptions with graph building.
             os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+            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
+            ):
+                # 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
+
+                tp_pp_world_size = (
+                    self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
+                )
+
+                # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
+                self.local_rank += dp_local_rank * tp_pp_world_size
+                assert self.local_rank < torch.cuda.device_count(), (
+                    f"DP adjusted local rank {self.local_rank} is out of bounds. "
+                )
+                visible_device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
+                assert self.parallel_config.local_world_size <= visible_device_count, (
+                    f"local_world_size ({self.parallel_config.local_world_size}) must "
+                    f"be less than or equal to the number of visible devices "
+                    f"({visible_device_count})."
+                )
             self.device = torch.device(f"cuda:{self.local_rank}")
             current_platform.set_device(self.device)
 
@@ -61,23 +81,23 @@ class GPUGenerationWorker(GPUWorker):
             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(b):
-                    return round(b / GiB_bytes, 2)
+                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"({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"{gib(self.requested_memory)} GiB). Decrease GPU memory "
                     f"utilization or reduce GPU memory used by other processes."
                 )
         else:
             raise RuntimeError(f"Not support device type: {self.device_config.device}")
 
         # Construct the model runner
-        self.model_runner: GPUGenerationModelRunner = GPUGenerationModelRunner(self.vllm_config, self.device)
+        self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.

vllm_omni/worker/gpu_model_runner.py

@@ -4,22 +4,17 @@ import numpy as np
 import torch
 from vllm.config import CUDAGraphMode
 from vllm.distributed.parallel_state import get_pp_group
-from vllm.forward_context import BatchDescriptor, set_forward_context
+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_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
-from vllm.utils import LazyLoader, cdiv
-from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
-    split_attn_metadata,
-)
+from vllm.utils.import_utils import LazyLoader
+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_splitting import ubatch_split
-from vllm.v1.worker.ubatch_utils import UBatchSlices
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
@@ -35,6 +30,12 @@ logger = init_logger(__name__)
 
 
 class OmniGPUModelRunner(GPUModelRunner):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._omni_per_req_additional_information: dict[str, dict] | None = None
+        self._omni_num_scheduled_tokens_np: np.ndarray | None = None
+        self._omni_last_model_output: object | None = None
+
     def _init_mrope_positions(self, req_state: CachedRequestState):
         """Initialize M-RoPE positions for multimodal inputs.
 
@@ -334,6 +335,8 @@ class OmniGPUModelRunner(GPUModelRunner):
         is_profile: bool = False,
         create_mixed_batch: bool = False,
         remove_lora: bool = True,
+        activate_lora: bool = False,
+        is_graph_capturing: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -356,12 +359,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             create_mixed_batch: If True, create a mixed batch with both decode
                 (1 token) and prefill (multiple tokens) requests.
             remove_lora: If False, dummy LoRAs are not destroyed after the run
+            activate_lora: If False, dummy_run is performed without LoRAs.
         """
-        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
-            CUDAGraphMode.NONE,
-            CUDAGraphMode.PIECEWISE,
-            CUDAGraphMode.FULL,
-        }
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
 
         # If cudagraph_mode.decode_mode() == FULL and
         # cudagraph_mode.separate_routine(). This means that we are using
@@ -387,7 +387,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             assert not uniform_decode
             # Create mixed batch:
             # first half decode tokens, second half one prefill
-            num_decode_tokens = num_tokens // 2
+            num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2)
             num_prefill_tokens = num_tokens - num_decode_tokens
             num_reqs = num_decode_tokens + 1
 
@@ -397,7 +397,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             max_query_len = num_prefill_tokens
         elif uniform_decode:
             assert not create_mixed_batch
-            num_reqs = cdiv(num_tokens, max_query_len)
+            num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len))
             num_scheduled_tokens_list = [max_query_len] * num_reqs
             if num_tokens % max_query_len != 0:
                 num_scheduled_tokens_list[-1] = num_tokens % max_query_len
@@ -410,56 +410,52 @@ class OmniGPUModelRunner(GPUModelRunner):
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
         num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+        num_tokens_unpadded = int(num_scheduled_tokens.sum())
 
-        total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
+        num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
 
-        ubatch_slices = None
-        num_tokens_after_padding = None
+        _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,
+        )
 
-        # We currently only microbatch if the number of tokens is
-        # over a certain threshold.
-        if self.parallel_config.enable_dbo and allow_microbatching:
-            ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split(
-                num_scheduled_tokens,
-                total_num_scheduled_tokens,
-                total_num_scheduled_tokens,
-                uniform_decode=uniform_decode,
-                vllm_config=self.vllm_config,
-            )
-            # Currently when DBO is enabled `ubatch_split` returns
-            # the num_tokens_after_padding for a single ubatch, but we have 2
-            # TODO(sage,lucas): this is cruft that should be addressed in the
-            # padding refactor.
-            if ubatch_num_tokens_after_padding is not None:
-                num_tokens_after_padding = ubatch_num_tokens_after_padding * 2
-
-        # If we failed to microbatch, currently need to resynchronize
-        # TODO(lucas,sage): we should be able to avoid this second sync by
-        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
-        #  a single `coordinate_batch_across_dp` function.
-        if num_tokens_after_padding is None:
-            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-            num_tokens_after_padding = num_tokens + num_pad
+        if cudagraph_runtime_mode is None:
+            cudagraph_runtime_mode = _cudagraph_mode
         else:
-            num_tokens_across_dp = num_tokens_after_padding
-            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
+            assert cudagraph_runtime_mode == _cudagraph_mode, (
+                f"Cudagraph runtime mode mismatch in dummy_run. "
+                f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}."
+            )
+
+        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
 
         attn_metadata: PerLayerAttnMetadata | None = None
 
         # If force_attention is True, we always capture attention. Otherwise,
         # it only happens for cudagraph_runtime_mode=FULL.
         if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-            attn_metadata = {}
-            if ubatch_slices is not None:
-                attn_metadata = [dict() for _ in range(len(ubatch_slices))]
-
             if create_mixed_batch:
                 # In the mixed batch mode (used for FI warmup), we use
                 # shorter sequence lengths to run faster.
                 # TODO(luka) better system for describing dummy batches
                 seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
             else:
-                seq_lens = max_query_len
+                seq_lens = max_query_len  # type: ignore[assignment]
             self.seq_lens.np[:num_reqs] = seq_lens
             self.seq_lens.np[num_reqs:] = 0
             self.seq_lens.copy_to_gpu()
@@ -468,61 +464,45 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
             self.query_start_loc.copy_to_gpu()
 
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
-                common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
-                    seq_lens=self.seq_lens.gpu[:num_reqs],
-                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],  # noqa: E501
-                    num_reqs=num_reqs,
-                    num_actual_tokens=num_tokens,
-                    max_query_len=max_query_len,
-                    max_seq_len=self.max_model_len,
-                    block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs),
-                    slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
-                    causal=True,
-                )
-                for attn_group in self.attn_groups[kv_cache_group_id]:
-                    if ubatch_slices is not None:
-                        common_attn_metadata_list = split_attn_metadata(ubatch_slices, common_attn_metadata)
-                        for ubid, common_attn_metadata in enumerate(common_attn_metadata_list):
-                            assert common_attn_metadata.max_query_len == 1
-                            attn_metadata_i = attn_group.get_metadata_builder(
-                                ubatch_id=ubid
-                            ).build_for_cudagraph_capture(common_attn_metadata)
-                            for layer_name in attn_group.layer_names:
-                                assert type(attn_metadata) is list
-                                attn_metadata[ubid][layer_name] = attn_metadata_i
-                    else:
-                        assert type(attn_metadata) is dict
-                        attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
-                            common_attn_metadata
-                        )  # noqa: E501
-                        for layer_name in attn_group.layer_names:
-                            attn_metadata[layer_name] = attn_metadata_i
+            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,
+                for_cudagraph_capture=is_graph_capturing,
+            )
 
-        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
-            model_kwargs = self._init_model_kwargs(num_tokens)
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config,
+            num_scheduled_tokens,
+            num_sampled_tokens,
+            activate_lora,
+            remove_lora,
+        ):
+            # 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)
             if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
                 input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
                 model_kwargs = {
                     **model_kwargs,
                     **self._dummy_mm_kwargs(num_reqs),
                 }
             elif self.enable_prompt_embeds:
                 input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
-                model_kwargs = self._init_model_kwargs(num_tokens)
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
+                model_kwargs = self._init_model_kwargs(num_tokens_padded)
             else:
-                input_ids = self.input_ids.gpu[:num_tokens]
+                input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
 
             if self.uses_mrope:
-                positions = self.mrope_positions.gpu[:, :num_tokens]
+                positions = self.mrope_positions.gpu[:, :num_tokens_padded]
+            elif self.uses_xdrope_dim > 0:
+                positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
             else:
-                positions = self.positions.gpu[:num_tokens]
+                positions = self.positions.gpu[:num_tokens_padded]
 
             if get_pp_group().is_first_rank:
                 intermediate_tensors = None
@@ -534,46 +514,25 @@ class OmniGPUModelRunner(GPUModelRunner):
                         device=self.device,
                     )
 
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False)
-
-            # filter out the valid batch descriptor
-            _cg_mode, batch_descriptor = (
-                self.cudagraph_dispatcher.dispatch(
-                    BatchDescriptor(
-                        num_tokens=num_tokens_after_padding,
-                        uniform_decode=uniform_decode,
-                    )
-                )
-                if not is_profile
-                else (CUDAGraphMode.NONE, None)
-            )
-            if cudagraph_runtime_mode is not None:
-                # we allow forcing NONE when the dispatcher disagrees to support
-                # warm ups for cudagraph capture
-                assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
-                    f"Cudagraph runtime mode mismatch at dummy_run. "
-                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
-                )
-            else:
-                cudagraph_runtime_mode = _cg_mode
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False)
 
             if ubatch_slices 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_after_padding = ubatch_slices[0].num_tokens
+                num_tokens_padded = ubatch_slices[0].num_tokens
                 if num_tokens_across_dp is not None:
-                    num_tokens_across_dp[:] = num_tokens_after_padding
+                    num_tokens_across_dp[:] = num_tokens_padded
 
             with (
                 self.maybe_randomize_inputs(input_ids),
                 set_forward_context(
                     attn_metadata,
                     self.vllm_config,
-                    num_tokens=num_tokens_after_padding,
+                    num_tokens=num_tokens_padded,
                     num_tokens_across_dp=num_tokens_across_dp,
                     cudagraph_runtime_mode=cudagraph_runtime_mode,
-                    batch_descriptor=batch_descriptor,
+                    batch_descriptor=batch_desc,
                     ubatch_slices=ubatch_slices,
                 ),
             ):
@@ -589,10 +548,26 @@ class OmniGPUModelRunner(GPUModelRunner):
                 hidden_states, _ = outputs
             else:
                 hidden_states = outputs
-
+            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)
-                self.drafter.dummy_run(num_tokens)
+                use_cudagraphs = (
+                    cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE)
+                    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
+                # short term mitigation for issue mentioned in
+                # https://github.com/vllm-project/vllm/issues/28334
+                if self.compilation_config.cudagraph_specialize_lora and activate_lora:
+                    use_cudagraphs = False
+
+                self.drafter.dummy_run(
+                    num_tokens,
+                    use_cudagraphs=use_cudagraphs,
+                    is_graph_capturing=is_graph_capturing,
+                )
 
         # This is necessary to avoid blocking DP.
         # For dummy runs, we typically skip EPLB since we don't have any real
@@ -605,110 +580,8 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.eplb_step(is_dummy=True, is_profile=is_profile)
 
         logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
-        return hidden_states, hidden_states[logit_indices]
-
-    def _preprocess(
-        self,
-        scheduler_output: "SchedulerOutput",
-        num_scheduled_tokens_np: np.ndarray,
-        intermediate_tensors: IntermediateTensors | None = None,
-        ubatch_slices: UBatchSlices | None = None,
-        num_tokens_after_padding: torch.Tensor | None = None,
-    ) -> tuple[
-        int,
-        int,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor,
-        IntermediateTensors | None,
-        dict[str, Any],
-        dict[str, dict] | None,
-    ]:
-        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if ubatch_slices:
-            assert num_tokens_after_padding is not None
-            num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
-            self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
-        elif ubatch_slices is None:
-            num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
-            num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
-            num_input_tokens += num_pad
-
-        # _prepare_inputs may reorder the batch, so we must gather multi
-        # modal outputs after that to ensure the correct order
-        per_req_additional_information: dict[str, dict] | None = None
-        if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
-            # Build multimodal inputs and overlay prompt embeds; collect per-request info
-            per_req_additional_information = self._build_mm_inputs_and_overlays(
-                scheduler_output, num_scheduled_tokens, num_scheduled_tokens_np
-            )
-
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = {
-                **self._init_model_kwargs(num_scheduled_tokens),
-                **self._extract_mm_kwargs(scheduler_output),
-            }
-        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
-            # Get the input embeddings for the tokens that are not input embeds,
-            # then put them into the appropriate positions.
-            # TODO(qthequartermasterman): Since even when prompt embeds are
-            # enabled, (a) not all requests will use prompt embeds, and (b)
-            # after the initial prompt is processed, the rest of the generated
-            # tokens will be token ids, it is not desirable to have the
-            # embedding layer outside of the CUDA graph all the time. The v0
-            # engine avoids this by "double compiling" the CUDA graph, once
-            # with input_ids and again with inputs_embeds, for all num_tokens.
-            # If a batch only has token ids, then including the embedding layer
-            # in the CUDA graph will be more performant (like in the else case
-            # below).
-            token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
-            # Some tokens ids may need to become embeds
-            if token_ids_idx.numel() > 0:
-                token_ids = self.input_ids.gpu[token_ids_idx]
-                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
-                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)
-            input_ids = None
-        else:
-            # For text-only models, we use token ids as input.
-            # While it is possible to use embeddings as input just like the
-            # multimodal models, it is not desirable for performance since
-            # 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)
-        if self.uses_mrope:
-            positions = self.mrope_positions.gpu[:, :num_input_tokens]
-        else:
-            positions = self.positions.gpu[:num_input_tokens]
-
-        if get_pp_group().is_first_rank:
-            intermediate_tensors = None
-        else:
-            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True
-            )
-
-        if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
-            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
-            model_kwargs.update(encoder_inputs)
-
-        return (
-            num_scheduled_tokens,
-            num_input_tokens,
-            num_tokens_after_padding,
-            input_ids,
-            inputs_embeds,
-            positions,
-            intermediate_tensors,
-            model_kwargs,
-            per_req_additional_information,
-        )
+        logit_indices_device = torch.from_numpy(logit_indices).to(self.device, non_blocking=True)
+        return hidden_states, hidden_states[logit_indices_device]
 
     def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None:
         """Decode per-request prompt_embeds and additional_information for newly
@@ -770,7 +643,14 @@ class OmniGPUModelRunner(GPUModelRunner):
         for req_id in self.input_batch.req_ids:
             req_state = self.requests.get(req_id)
             info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
-            per_req_runtime_info.append(info if isinstance(info, dict) else {})
+            if info and isinstance(info, dict):
+                per_req_runtime_info.append(info)
+                if "thinker_reply_part_per_request" in info:
+                    q = info["thinker_reply_part_per_request"]
+                    if hasattr(q, "shape"):
+                        logger.debug(f"[OMNI] req={req_id} has thinker_reply_part_per_request queue shape: {q.shape}")
+            else:
+                per_req_runtime_info.append({})
         return per_req_runtime_info
 
     def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]:
@@ -798,24 +678,34 @@ class OmniGPUModelRunner(GPUModelRunner):
             model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
         try:
             model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information()
-            model_kwargs_extra["request_ids"] = self.input_batch.req_ids
-            model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(num_scheduled_tokens_np)
-        except Exception:
-            pass
+            model_kwargs_extra["request_ids"] = list(self.input_batch.req_ids)
+            if num_scheduled_tokens_np is not None:
+                model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(num_scheduled_tokens_np)
+        except Exception as e:
+            logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}")
+            import traceback
+
+            traceback.print_exc()
         return model_kwargs_extra
 
     def _process_additional_information_updates(self, multimodal_outputs: object) -> None:
         """Process model-provided per-request additional_information updates and merge into request state."""
         try:
+            if multimodal_outputs is None:
+                logger.debug("[OMNI] multimodal_outputs is None, no updates to process")
+                return
             if not isinstance(multimodal_outputs, dict):
+                logger.debug(f"[OMNI] multimodal_outputs is not dict: {type(multimodal_outputs)}")
                 return
             if (
                 "additional_information_update" not in multimodal_outputs
                 and "additional_information_update_by_req_id" not in multimodal_outputs
             ):
+                logger.debug(f"[OMNI] multimodal_outputs has no update keys: {list(multimodal_outputs.keys())}")
                 return
             updates_list = multimodal_outputs.get("additional_information_update")
             if isinstance(updates_list, list):
+                logger.debug(f"[OMNI] Processing additional_information_update list with {len(updates_list)} items")
                 for idx, upd in enumerate(updates_list):
                     if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
                         continue
@@ -823,18 +713,33 @@ class OmniGPUModelRunner(GPUModelRunner):
                     self._merge_additional_information_update(req_id, upd)
             updates_map = multimodal_outputs.get("additional_information_update_by_req_id")
             if isinstance(updates_map, dict):
+                logger.debug(
+                    f"[OMNI] Processing additional_information_update_by_req_id for {list(updates_map.keys())}"
+                )
                 for req_id, upd in updates_map.items():
                     if not isinstance(upd, dict):
                         continue
                     if req_id not in self.requests:
+                        logger.warning(f"[OMNI] req_id {req_id} not found in requests")
                         continue
                     self._merge_additional_information_update(req_id, upd)
+
+                    req_state = self.requests.get(req_id)
+                    if req_state:
+                        info = getattr(req_state, "additional_information_cpu", {})
+                        if "thinker_reply_part_per_request" in info:
+                            q = info["thinker_reply_part_per_request"]
+                            if hasattr(q, "shape"):
+                                logger.debug(f"[OMNI] After update, req={req_id} queue shape: {q.shape}")
         except Exception as e:
             logger.error(
                 f"Error merging for requests:{self.input_batch.req_ids} "
                 f"additional information update: {e}, with the multimodal_outputs "
                 f"as {multimodal_outputs}"
             )
+            import traceback
+
+            traceback.print_exc()
 
     def _build_mm_inputs_and_overlays(
         self,
@@ -895,3 +800,164 @@ class OmniGPUModelRunner(GPUModelRunner):
                 if req_info:
                     per_req_additional_information[req_id] = req_info
         return per_req_additional_information
+
+    def _preprocess(
+        self,
+        scheduler_output: "SchedulerOutput",
+        num_input_tokens: int,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ):
+        """Align with v0.12 preprocess and omni's additional information handling."""
+        # Decode payload first, ensure request state has prompt_embeds / additional_information
+        self._decode_and_store_request_payloads(scheduler_output)
+
+        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
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        ec_connector_output = None
+
+        if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder:
+            # Run the multimodal encoder if any.
+            with self.maybe_get_ec_connector_output(
+                scheduler_output,
+                encoder_cache=self.encoder_cache,
+            ) as ec_connector_output:
+                self._execute_mm_encoder(scheduler_output)
+                mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output)
+
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            inputs_embeds_scheduled = self.model.embed_input_ids(
+                self.input_ids.gpu[:num_scheduled_tokens],
+                multimodal_embeddings=mm_embeds,
+                is_multimodal=is_mm_embed,
+            )
+
+            # 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]
+            model_kwargs = {
+                **self._init_model_kwargs(num_scheduled_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+        elif self.enable_prompt_embeds and is_first_rank:
+            # Get the input embeddings for the tokens that are not input embeds,
+            # then put them into the appropriate positions.
+            # TODO(qthequartermasterman): Since even when prompt embeds are
+            # enabled, (a) not all requests will use prompt embeds, and (b)
+            # after the initial prompt is processed, the rest of the generated
+            # tokens will be token ids, it is not desirable to have the
+            # embedding layer outside of the CUDA graph all the time. The v0
+            # engine avoids this by "double compiling" the CUDA graph, once
+            # with input_ids and again with inputs_embeds, for all num_tokens.
+            # If a batch only has token ids, then including the embedding layer
+            # in the CUDA graph will be more performant (like in the else case
+            # below).
+            token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
+            # Some tokens ids may need to become embeds
+            if token_ids_idx.numel() > 0:
+                token_ids = self.input_ids.gpu[token_ids_idx]
+                tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids)
+                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)
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # 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)
+
+        if self.uses_mrope:
+            positions = self.mrope_positions.gpu[:, :num_input_tokens]
+        elif self.uses_xdrope_dim > 0:
+            positions = self.xdrope_positions.gpu[:, :num_input_tokens]
+        else:
+            positions = self.positions.gpu[:num_input_tokens]
+
+        if is_first_rank:
+            intermediate_tensors = None
+        else:
+            assert intermediate_tensors is not None
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True
+            )
+
+        if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
+            # Run the encoder, just like we do with other multimodal inputs.
+            # For an encoder-decoder model, our processing here is a bit
+            # simpler, because the outputs are just passed to the decoder.
+            # We are not doing any prompt replacement. We also will only
+            # ever have a single encoder input.
+            encoder_outputs = self._execute_mm_encoder(scheduler_output)
+            model_kwargs.update({"encoder_outputs": encoder_outputs})
+
+        req_ids = self.input_batch.req_ids
+        num_scheduled_tokens_np = np.array(
+            [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids],
+            dtype=np.int32,
+        )
+        self._omni_num_scheduled_tokens_np = num_scheduled_tokens_np
+
+        # Always collect additional_information, not only in prefill.
+        # _collect_additional_information_for_prefill handles prompt_embeds overlay during prefill,
+        # but even in decode we need per_req_additional_information so that
+        # _build_model_kwargs_extra can pass runtime_additional_information correctly.
+        per_req_additional_information: dict[str, dict] = {}
+        if inputs_embeds is not None:
+            # Prefill: overlay prompt_embeds and collect additional_information
+            per_req_additional_information = self._collect_additional_information_for_prefill(num_scheduled_tokens_np)
+        # Save for _model_forward regardless of prefill; runtime_additional_information
+        # will be fetched from request state inside _build_model_kwargs_extra
+        self._omni_per_req_additional_information = per_req_additional_information
+
+        return (
+            input_ids,
+            inputs_embeds,
+            positions,
+            intermediate_tensors,
+            model_kwargs,
+            ec_connector_output,
+        )
+
+    def _model_forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **model_kwargs: dict[str, Any],
+    ):
+        """Inject omni-specific kwargs into forward and cache model output"""
+        model_kwargs_extra = self._build_model_kwargs_extra(
+            self._omni_per_req_additional_information,
+            self._omni_num_scheduled_tokens_np,
+        )
+
+        runtime_info = model_kwargs_extra.get("runtime_additional_information", [])
+        if runtime_info:
+            for i, info in enumerate(runtime_info):
+                if info:
+                    logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}")
+
+        model_output = super()._model_forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **model_kwargs,
+            **model_kwargs_extra,
+        )
+        # Cache model output so later sample_tokens can consume multimodal results.
+        self._omni_last_model_output = model_output
+        return model_output