[Rebase] Rebase to vllm v0.19.0 (#2475)

.buildkite/pipeline-intel.yaml

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

.buildkite/test-nightly.yml

@@ -362,7 +362,7 @@ steps:
       - kubernetes:
           podSpec:
             containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:c392ce21e9cf9ea65c52b866447793db10e0261c
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                 resources:
                   limits:
                     nvidia.com/gpu: 2

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.18.0
+ARG VLLM_BASE_TAG=v0.19.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -11,9 +11,7 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Install vllm-omni into the same uv-managed Python environment used by the base image.
-# Use bash -c so that $(python3 -c ...) is expanded inside the container.
-RUN uv pip install --system --no-cache-dir ".[dev]"
+RUN uv pip install --system ".[dev]"
 
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docker/Dockerfile.xpu

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

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

@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
 
 ```bash
-uv pip install vllm==0.18.0 --torch-backend=auto
+uv pip install vllm==0.19.0 --torch-backend=auto
 ```
 
-The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,11 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.18.0
+git checkout v0.19.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.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
+#For CUDA 13.0
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash

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

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
 
 ```bash
-uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
 ```
 
-The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.18.0
+git checkout v0.19.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.18.0 \
+  vllm/vllm-omni-rocm:v0.19.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.18.0
+  vllm/vllm-omni-rocm:v0.19.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

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

docs/mkdocs/hooks/generate_argparse.py

@@ -120,6 +120,7 @@ def extract_omni_serve_subparser_init():
                         "_FlexibleArgumentParser": _FlexibleArgumentParser,
                         "FlexibleArgumentParser": _FlexibleArgumentParser,
                         "make_arg_parser": lambda parser: parser,  # no-op for doc
+                        "_ensure_vllm_platform": lambda: None,  # no-op for doc
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,

tests/diffusion/quantization/test_quantization_quality.py

@@ -118,8 +118,12 @@ def _generate_image(omni, config: QualityTestConfig):
 
     peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
     first = outputs[0]
-    req_out = first.request_output[0] if hasattr(first, "request_output") else first
-    return req_out.images[0], peak_mem
+    if hasattr(first, "images") and first.images:
+        return first.images[0], peak_mem
+    inner = first.request_output
+    if inner is not None and hasattr(inner, "images") and inner.images:
+        return inner.images[0], peak_mem
+    raise ValueError("Could not extract image from output.")
 
 
 def _generate_video(omni, config: QualityTestConfig):

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -229,6 +229,24 @@ def _wait_for_port(host: str, port: int, timeout: int = 30) -> bool:
     return False
 
 
+def _is_mooncake_master_available() -> bool:
+    """Check if mooncake_master binary is present and can actually execute."""
+    import shutil
+
+    binary = shutil.which("mooncake_master")
+    if binary is None:
+        return False
+    try:
+        result = subprocess.run(
+            [binary, "--help"],
+            capture_output=True,
+            timeout=5,
+        )
+        return result.returncode != 127
+    except (subprocess.TimeoutExpired, OSError):
+        return True
+
+
 def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None:
     """Clean up any existing mooncake_master processes.
 
@@ -292,6 +310,8 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
 @hardware_test(res={"cuda": "H100"})
 def test_bagel_text2img_mooncake_connector(run_level):
     """Test Bagel text2img with Mooncake connector for inter-stage communication."""
+    if not _is_mooncake_master_available():
+        pytest.skip("mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)")
     MOONCAKE_HOST = "127.0.0.1"
     MOONCAKE_RPC_PORT = _find_free_port()
     MOONCAKE_HTTP_PORT = _find_free_port()

tests/e2e/online_serving/test_mimo_audio.py

@@ -63,20 +63,27 @@ def download_tokenizer():
 
 
 # CI stage config for H100 / MI325
-stage_configs = [get_chunk_config()]
-tokenizer_path = download_tokenizer()
-os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path
+# Guard module-level setup so test collection doesn't fail in environments
+# where the model cache is read-only or models aren't available.
+try:
+    stage_configs = [get_chunk_config()]
+    tokenizer_path = download_tokenizer()
+    os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path
 
-# Create parameter combinations for model and stage config
-test_params = [
-    OmniServerParams(
-        model=model,
-        stage_config_path=stage_config,
-        server_args=["--chat-template", CHAT_TEMPLATE_PATH],
+    test_params = [
+        OmniServerParams(
+            model=model,
+            stage_config_path=stage_config,
+            server_args=["--chat-template", CHAT_TEMPLATE_PATH],
+        )
+        for model in models
+        for stage_config in stage_configs
+    ]
+except Exception as exc:
+    pytest.skip(
+        f"MiMo-Audio online serving tests skipped: module setup failed ({type(exc).__name__}: {exc})",
+        allow_module_level=True,
     )
-    for model in models
-    for stage_config in stage_configs
-]
 
 
 def get_prompt(prompt_type="text_only"):

tests/engine/test_async_omni_engine_abort.py

@@ -60,6 +60,7 @@ async def generate(
 
 @pytest.mark.core_model
 @pytest.mark.omni
+@pytest.mark.real_hf_config
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.asyncio
 async def test_abort():

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1664,6 +1664,7 @@ class TestFishSpeechServing:
         assert all(allowed_special is None for _, _, allowed_special in tokenizer.calls)
 
     def test_build_fish_clone_prompt_normalizes_text_fields(self, fish_speech_server):
+        fish_speech_server._fish_speech_tokenizer = _FakeFishTokenizer()
         fish_speech_server._estimate_fish_prompt_len = MagicMock(return_value=123)
 
         request = OpenAICreateSpeechRequest(

tests/entrypoints/test_async_omni_abort.py

@@ -13,8 +13,8 @@ def test_generate_accepts_request_after_repeated_cancellations():
         submitted_request_ids = []
         aborted_request_batches = []
 
-        async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id):
-            del prompt, sampling_params_list, final_stage_id
+        async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
+            del prompt, sampling_params_list, final_stage_id, kwargs
             submitted_request_ids.append(request_id)
 
         async def fake_abort_async(request_ids):

tests/entrypoints/test_omni_entrypoints.py

@@ -113,6 +113,7 @@ class FakeAsyncOmniEngine:
         sampling_params_list: list[Any] | None = None,
         final_stage_id: int = 0,
         arrival_time: float | None = None,
+        **kwargs: Any,
     ) -> None:
         msg = {
             "request_id": request_id,

tests/examples/conftest.py

@@ -14,7 +14,6 @@ from collections.abc import Callable
 from pathlib import Path
 from typing import Any, NamedTuple, cast
 
-import mistune
 import pytest
 import torch
 from safetensors.torch import save_file
@@ -62,6 +61,8 @@ class ReadmeSnippet(NamedTuple):
         readme_path: Path,
         skipif: ReadmeSnippetExtractionSkipPredicate | None = None,
     ) -> list["ReadmeSnippet"]:
+        import mistune
+
         markdown = mistune.create_markdown(renderer="ast")
         tokens = markdown(readme_path.read_text(encoding="utf-8"))
         tokens = cast(list[dict[str, Any]], tokens)  # mistune's AST renderer always produces a list, not a str

tests/model_executor/models/test_omni_processing.py

@@ -21,9 +21,10 @@ from vllm.config.multimodal import (
     ImageDummyOptions,
     VideoDummyOptions,
 )
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict
+from vllm.inputs import MultiModalDataDict, MultiModalInput
+from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.cache import MultiModalProcessorOnlyCache
-from vllm.multimodal.inputs import MultiModalInputs, batched_tensors_equal
+from vllm.multimodal.inputs import batched_tensors_equal
 from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
 from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
 
@@ -302,8 +303,8 @@ def test_omni_processing_correctness(
 
 
 def _assert_inputs_equal(
-    a: MultiModalInputs,
-    b: MultiModalInputs,
+    a: MultiModalInput,
+    b: MultiModalInput,
     *,
     ignore_mm_keys: set[str] | None = None,
     msg: str = "",

vllm_omni/benchmarks/patch/patch.py

@@ -358,6 +358,7 @@ async def benchmark(
     lora_modules: Iterable[str] | None,
     extra_headers: dict | None,
     extra_body: dict | None,
+    lora_assignment: Literal["random", "round-robin"] = "random",
     ramp_up_strategy: Literal["linear", "exponential"] | None = None,
     ramp_up_start_rps: int | None = None,
     ramp_up_end_rps: int | None = None,
@@ -454,8 +455,11 @@ async def benchmark(
     print("Starting main benchmark run...")
 
     if lora_modules:
-        # For each input request, choose a LoRA module at random.
-        lora_modules = iter([random.choice(lora_modules) for _ in range(len(input_requests))])
+        lora_modules_list = list(lora_modules)
+        if lora_assignment == "round-robin":
+            lora_modules = iter([lora_modules_list[i % len(lora_modules_list)] for i in range(len(input_requests))])
+        else:
+            lora_modules = iter([random.choice(lora_modules_list) for _ in range(len(input_requests))])
 
     if profile:
         print("Starting profiler...")

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -235,7 +235,10 @@ class OmniARScheduler(VLLMScheduler):
             # 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)
+            failed_kv_load_req_ids = self._handle_invalid_blocks(
+                kv_connector_output.invalid_block_ids,
+                num_scheduled_tokens,
+            )
 
         # 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

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -272,7 +272,7 @@ class OmniGenerationScheduler(VLLMScheduler):
 
         # KVTransfer: package metadata
         if self.connector is not None:
-            meta = self.connector.build_connector_meta(scheduler_output)
+            meta = self._build_kv_connector_meta(self.connector, scheduler_output)
             scheduler_output.kv_connector_metadata = meta
         # EC Connector: package metadata
         if self.ec_connector is not None:
@@ -368,7 +368,10 @@ class OmniGenerationScheduler(VLLMScheduler):
 
         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)
+            failed_kv_load_req_ids = self._handle_invalid_blocks(
+                kv_connector_output.invalid_block_ids,
+                num_scheduled_tokens,
+            )
 
         # 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

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -32,7 +32,9 @@ if current_omni_platform.is_rocm():
         pass
 elif current_omni_platform.is_xpu():
     try:
-        from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func  # noqa: F401
+        from vllm._xpu_ops import xpu_ops  # noqa: F401
+
+        flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func
     except (ImportError, ModuleNotFoundError):
         pass
 elif current_omni_platform.is_musa():

vllm_omni/diffusion/diffusion_engine.py

@@ -361,8 +361,8 @@ class DiffusionEngine:
     def _dummy_run(self):
         """A dummy run to warm up the model."""
         num_inference_steps = 1
-        height = 1024
-        width = 1024
+        height = 512
+        width = 512
         if supports_image_input(self.od_config.model_class_name):
             # Provide a dummy image input if the model supports it
             color_format = image_color_format(self.od_config.model_class_name)

vllm_omni/diffusion/stage_diffusion_client.py

@@ -244,7 +244,20 @@ class StageDiffusionClient:
             return self._output_queue.get_nowait()
         except asyncio.QueueEmpty:
             if not self._shutting_down and self._proc is not None and not self._proc.is_alive():
-                raise RuntimeError(f"StageDiffusionProc died unexpectedly (exit code {self._proc.exitcode})")
+                exitcode = self._proc.exitcode
+                # One final drain – the last ZMQ frame may have arrived
+                # between the first drain and the is_alive() check.
+                self._drain_responses()
+                try:
+                    return self._output_queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    pass
+                if exitcode is not None and exitcode > 128:
+                    sig = exitcode - 128
+                    logger.warning("StageDiffusionProc was killed by signal %d; treating as external shutdown.", sig)
+                    self._shutting_down = True
+                    return None
+                raise RuntimeError(f"StageDiffusionProc died unexpectedly (exit code {exitcode})")
             return None
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -495,7 +495,7 @@ class StageDiffusionProc:
             nonlocal shutdown_requested
             if not shutdown_requested:
                 shutdown_requested = True
-                raise SystemExit()
+                raise SystemExit(128 + signum)
 
         signal.signal(signal.SIGTERM, signal_handler)
         signal.signal(signal.SIGINT, signal_handler)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -17,7 +17,7 @@ from typing import Any
 
 import torch
 import zmq
-from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
+from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
 from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -113,8 +113,12 @@ class DiffusionWorker:
         self.device = current_omni_platform.get_torch_device(rank)
         current_omni_platform.set_device(self.device)
 
-        # Create vllm_config for parallel configuration
-        vllm_config = VllmConfig(compilation_config=CompilationConfig())
+        # Create vllm_config for parallel configuration. Pass explicit device_config
+        # so DeviceConfig does not rely on current_platform in worker subprocesses.
+        vllm_config = VllmConfig(
+            compilation_config=CompilationConfig(),
+            device_config=DeviceConfig(device=self.device),
+        )
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
         vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
         vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel

vllm_omni/engine/arg_utils.py

@@ -1,6 +1,8 @@
 import argparse
 import dataclasses
+import json
 import os
+import tempfile
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -135,6 +137,30 @@ class OmniEngineArgs(EngineArgs):
         self._omni_models_registered = True
         return True
 
+    def _patch_empty_hf_config(self, model_type: str) -> None:
+        """For models with empty config.json (e.g. CosyVoice3), create a
+        patched config in a temp directory with model_type set so that
+        transformers AutoConfig.from_pretrained can resolve the config class.
+        Sets self.hf_config_path to point to the patched directory."""
+        try:
+            from transformers import PretrainedConfig
+
+            config_dict, _ = PretrainedConfig.get_config_dict(self.model)
+            if config_dict.get("model_type"):
+                return  # config.json already has model_type, no patching needed
+        except Exception:
+            return  # can't load config, let vLLM handle the error
+
+        # Create a temp dir with a patched config.json
+        temp_dir = tempfile.mkdtemp(prefix="omni_hf_config_")
+        config_dict["model_type"] = model_type
+        config_dict.setdefault("architectures", [self.model_arch])
+        with open(os.path.join(temp_dir, "config.json"), "w") as f:
+            json.dump(config_dict, f)
+        self.hf_config_path = temp_dir
+        self._temp_config_dir = temp_dir
+        logger.info("Patched empty HF config with model_type=%s at %s", model_type, temp_dir)
+
     def create_model_config(self) -> OmniModelConfig:
         """Create an OmniModelConfig from these engine arguments.
         Returns:
@@ -164,6 +190,18 @@ class OmniEngineArgs(EngineArgs):
                     if model_type is not None:
                         self.hf_overrides.setdefault("model_type", model_type)
 
+            # For models whose HF config.json is empty or lacks model_type
+            # (e.g. CosyVoice3), AutoConfig.from_pretrained fails because it
+            # cannot determine which config class to use from the empty dict.
+            # hf_overrides alone is not enough since transformers reads
+            # model_type from config_dict before applying overrides.
+            # Workaround: create a patched config.json in a temp directory
+            # and point hf_config_path to it so vLLM reads model_type from it.
+            if not self.hf_config_path:
+                model_type = _ARCH_TO_MODEL_TYPE.get(self.model_arch)
+                if model_type is not None:
+                    self._patch_empty_hf_config(model_type)
+
         # Auto-detect tokenizer for models that store it in a subdirectory
         # rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/).
         if not self.tokenizer and self.model:
@@ -200,7 +238,15 @@ class OmniEngineArgs(EngineArgs):
                         logger.warning("Failed to download tokenizer subfolder: %s", e)
 
         # Build the vLLM config first, then use it to create the Omni config.
-        model_config = super().create_model_config()
+        try:
+            model_config = super().create_model_config()
+        finally:
+            # Clean up temp config dir if we created one
+            if hasattr(self, "_temp_config_dir"):
+                import shutil
+
+                shutil.rmtree(self._temp_config_dir, ignore_errors=True)
+                del self._temp_config_dir
 
         omni_config = OmniModelConfig.from_vllm_model_config(
             model_config=model_config,

vllm_omni/engine/async_omni_engine.py

@@ -17,7 +17,7 @@ import threading
 import time
 import uuid
 import weakref
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import asdict
 from typing import TYPE_CHECKING, Any
 
@@ -652,6 +652,12 @@ class AsyncOmniEngine:
         sampling_params_list: Sequence[Any] | None = None,
         final_stage_id: int = 0,
         arrival_time: float | None = None,
+        lora_request: Any = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        trace_headers: Mapping[str, str] | None = None,
+        priority: int = 0,
+        data_parallel_rank: int | None = None,
+        reasoning_ended: bool | None = None,
         *,
         resumable: bool = False,
         message_type: str = "add_request",
@@ -686,12 +692,20 @@ class AsyncOmniEngine:
                 params=params,
                 supported_tasks=self.supported_tasks,
                 arrival_time=arrival_time,
+                lora_request=lora_request,
+                tokenization_kwargs=tokenization_kwargs,
+                trace_headers=trace_headers,
+                priority=priority,
+                data_parallel_rank=data_parallel_rank,
                 resumable=resumable,
             )
             # TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have
             # additional_information field in the prompt.
             request = _upgrade_to_omni_request(request, prompt)
 
+            if reasoning_ended is not None:
+                request.reasoning_ended = reasoning_ended
+
             # Restore external_req_id to the original user-facing request_id.
             # InputProcessor.process_inputs() renames request_id to an internal
             # UUID (saving the original in external_req_id), but then overwrites
@@ -973,6 +987,12 @@ class AsyncOmniEngine:
         sampling_params_list: Sequence[Any] | None = None,
         final_stage_id: int = 0,
         arrival_time: float | None = None,
+        lora_request: Any = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        trace_headers: Mapping[str, str] | None = None,
+        priority: int = 0,
+        data_parallel_rank: int | None = None,
+        reasoning_ended: bool | None = None,
         *,
         resumable: bool = False,
     ) -> None:
@@ -990,6 +1010,12 @@ class AsyncOmniEngine:
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
             arrival_time=arrival_time,
+            lora_request=lora_request,
+            tokenization_kwargs=tokenization_kwargs,
+            trace_headers=trace_headers,
+            priority=priority,
+            data_parallel_rank=data_parallel_rank,
+            reasoning_ended=reasoning_ended,
             resumable=resumable,
         )
         if self.request_queue is None:
@@ -1013,6 +1039,12 @@ class AsyncOmniEngine:
         sampling_params_list: Sequence[Any] | None = None,
         final_stage_id: int = 0,
         arrival_time: float | None = None,
+        lora_request: Any = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        trace_headers: Mapping[str, str] | None = None,
+        priority: int = 0,
+        data_parallel_rank: int | None = None,
+        reasoning_ended: bool | None = None,
         *,
         resumable: bool = False,
     ) -> None:
@@ -1024,6 +1056,12 @@ class AsyncOmniEngine:
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
             arrival_time=arrival_time,
+            lora_request=lora_request,
+            tokenization_kwargs=tokenization_kwargs,
+            trace_headers=trace_headers,
+            priority=priority,
+            data_parallel_rank=data_parallel_rank,
+            reasoning_ended=reasoning_ended,
             resumable=resumable,
         )
 

vllm_omni/engine/stage_init_utils.py

@@ -467,6 +467,35 @@ def initialize_diffusion_stage(
     return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
 
 
+def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int) -> None:
+    """vLLM CoreEngineProcManager / coordinators use ``shutdown()``, not ``close()``."""
+    if resource is None:
+        return
+    shutdown = getattr(resource, "shutdown", None)
+    if callable(shutdown):
+        try:
+            shutdown()
+        except Exception as cleanup_error:
+            logger.warning(
+                "[stage_init] Failed to shutdown launched %s for stage %s: %s",
+                resource_name,
+                stage_id,
+                cleanup_error,
+            )
+        return
+    close = getattr(resource, "close", None)
+    if callable(close):
+        try:
+            close()
+        except Exception as cleanup_error:
+            logger.warning(
+                "[stage_init] Failed to close launched %s for stage %s: %s",
+                resource_name,
+                stage_id,
+                cleanup_error,
+            )
+
+
 def close_started_llm_stage(started: StartedLlmStage) -> None:
     """Terminate the subprocess owned by a launched stage that never attached."""
     if started.proc is None:

vllm_omni/entrypoints/async_omni.py

@@ -9,7 +9,7 @@ from __future__ import annotations
 
 import asyncio
 import time
-from collections.abc import AsyncGenerator, Iterable, Sequence
+from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
 from vllm import TokensPrompt
@@ -88,7 +88,12 @@ class AsyncOmni(EngineClient, OmniBase):
         else:
             vllm_config = self.engine.stage_vllm_configs[stage_index]
             io_processor_plugin = vllm_config.model_config.io_processor_plugin
-            self.io_processor = get_io_processor(vllm_config, io_processor_plugin)
+            renderer = self.renderer
+            if renderer is None:
+                from vllm.renderers import renderer_from_config
+
+                renderer = renderer_from_config(vllm_config)
+            self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin)
 
     def _get_comprehension_stage_index(self) -> int | None:
         fallback_idx: int | None = None
@@ -159,6 +164,10 @@ class AsyncOmni(EngineClient, OmniBase):
         tokenization_kwargs: dict[str, Any] | None = None,
         sampling_params_list: Sequence[OmniSamplingParams] | None = None,
         output_modalities: list[str] | None = None,
+        trace_headers: Mapping[str, str] | None = None,
+        priority: int = 0,
+        data_parallel_rank: int | None = None,
+        reasoning_ended: bool | None = None,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         """Generate outputs for the given prompt(s) asynchronously.
 

vllm_omni/entrypoints/cli/main.py

@@ -28,6 +28,10 @@ def main():
 
         cli_env_setup()
 
+        from vllm_omni.entrypoints.cli.serve import _ensure_vllm_platform
+
+        _ensure_vllm_platform()
+
         parser = FlexibleArgumentParser(
             description="vLLM OMNI CLI",
             epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),

vllm_omni/entrypoints/cli/serve.py

@@ -42,6 +42,37 @@ Search by using: `--help=<ConfigGroup>` to explore options by section (e.g.,
 """
 
 
+def _ensure_vllm_platform():
+    """Ensure vLLM's current_platform is valid before arg parsing.
+
+    Upstream vLLM's argument parser now instantiates DeviceConfig during
+    ``make_arg_parser``, which requires a resolved platform with a non-empty
+    ``device_type``.  In some environments (e.g. editable installs with
+    broken package metadata), vLLM's own platform auto-detection may fail
+    and fall back to ``UnspecifiedPlatform``.  When that happens, use the
+    Omni platform (which has its own detection logic) as a drop-in
+    replacement so that argument parsing succeeds.
+    """
+    from vllm import platforms as vllm_platforms
+
+    if vllm_platforms.current_platform.is_unspecified():
+        from vllm_omni.platforms import current_omni_platform
+
+        if not current_omni_platform.is_unspecified():
+            vllm_platforms.current_platform = current_omni_platform
+            logger.debug(
+                "Replaced vLLM UnspecifiedPlatform with omni platform %s",
+                type(current_omni_platform).__name__,
+            )
+        else:
+            from vllm.platforms.cpu import CpuPlatform
+
+            vllm_platforms.current_platform = CpuPlatform()
+            logger.debug(
+                "Both vLLM and omni platforms are unspecified, falling back to CpuPlatform for arg parsing",
+            )
+
+
 class OmniServeCommand(CLISubcommand):
     """The `serve` subcommand for the vLLM CLI."""
 
@@ -82,6 +113,7 @@ class OmniServeCommand(CLISubcommand):
             usage="vllm serve [model_tag] --omni [options]",
         )
 
+        _ensure_vllm_platform()
         serve_parser = make_arg_parser(serve_parser)
         serve_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name)
 

vllm_omni/entrypoints/openai/api_server.py

@@ -639,6 +639,7 @@ async def omni_init_app_state(
         OpenAIServingResponses(
             engine_client,
             state.openai_serving_models,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -699,7 +700,8 @@ async def omni_init_app_state(
         OpenAIServingPooling(
             engine_client,
             state.openai_serving_models,
-            supported_tasks=supported_tasks,
+            state.openai_serving_render,
+            supported_tasks=tuple(supported_tasks),
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -746,6 +748,7 @@ async def omni_init_app_state(
     state.openai_serving_tokenization = OpenAIServingTokenization(
         engine_client,
         state.openai_serving_models,
+        state.openai_serving_render,
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
@@ -787,6 +790,7 @@ async def omni_init_app_state(
             reasoning_parser=args.structured_outputs_config.reasoning_parser,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_force_include_usage=args.enable_force_include_usage,
+            default_chat_template_kwargs=args.default_chat_template_kwargs,
         )
         if "generate" in supported_tasks
         else None
@@ -795,6 +799,7 @@ async def omni_init_app_state(
         ServingTokens(
             engine_client,
             state.openai_serving_models,
+            state.openai_serving_render,
             request_logger=request_logger,
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -62,7 +62,7 @@ from vllm.entrypoints.openai.parser.harmony_utils import (
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
 from vllm.entrypoints.utils import should_include_usage
-from vllm.inputs.data import PromptType
+from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.outputs import RequestOutput
 from vllm.reasoning import ReasoningParser
@@ -818,7 +818,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         # Prepare the tool parser if it's needed
         try:
             if tool_choice_auto and self.tool_parser:
-                tool_parsers: list[ToolParser | None] = [self.tool_parser(tokenizer)] * num_choices
+                tool_parsers: list[ToolParser | None] = [self.tool_parser(tokenizer, request.tools)] * num_choices
             else:
                 tool_parsers = [None] * num_choices
         except Exception as e:
@@ -1639,12 +1639,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 logprobs = None
 
             if self.use_harmony:
-                reasoning_content, content, _ = parse_chat_output(token_ids)
+                reasoning, content, _ = parse_chat_output(token_ids)
                 if not request.include_reasoning:
-                    reasoning_content = None
+                    reasoning = None
 
                 if self.tool_parser is not None:
-                    tool_parser = self.tool_parser(tokenizer)
+                    tool_parser = self.tool_parser(tokenizer, request.tools)
                     # NOTE: We use token_ids for openai tool parser
                     tool_call_info = tool_parser.extract_tool_calls(
                         "",
@@ -1654,14 +1654,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     content = tool_call_info.content
                     message = ChatMessage(
                         role=role,
-                        reasoning_content=reasoning_content,
+                        reasoning=reasoning,
                         content=content,
                         tool_calls=tool_call_info.tool_calls,
                     )
                 else:
                     message = ChatMessage(
                         role=role,
-                        reasoning_content=reasoning_content,
+                        reasoning=reasoning,
                         content=content,
                     )
 
@@ -1682,11 +1682,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             if reasoning_parser:
                 # If the reasoning parser is enabled,
                 # tool calls are extracted exclusively from the content.
-                reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request)
+                reasoning, content = reasoning_parser.extract_reasoning(output.text, request=request)
                 if not request.include_reasoning:
-                    reasoning_content = None
+                    reasoning = None
             else:
-                reasoning_content = None
+                reasoning = None
                 content = output.text
 
             auto_tools_called = False
@@ -1696,14 +1696,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
                 and request.tool_choice != "required"
             ):
-                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
+                message = ChatMessage(role=role, reasoning=reasoning, content=content)
 
             # if the request uses tools and specified a tool choice
             elif request.tool_choice and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam:
                 tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
                 message = ChatMessage(
                     role=role,
-                    reasoning_content=reasoning_content,
+                    reasoning=reasoning,
                     content="",
                     tool_calls=[
                         tool_call_class(
@@ -1745,13 +1745,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         )
                         for i, tool_call in enumerate(tool_calls)
                     ],
-                    reasoning_content=reasoning_content,
+                    reasoning=reasoning,
                 )
 
             # if the request doesn't use tool choice
             # OR specifies to not use a tool
             elif not request.tool_choice or request.tool_choice == "none":
-                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
+                message = ChatMessage(role=role, reasoning=reasoning, content=content)
 
             # handle when there are tools and tool choice is auto
             elif (
@@ -1761,7 +1761,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 and self.tool_parser
             ):
                 try:
-                    tool_parser = self.tool_parser(tokenizer)
+                    tool_parser = self.tool_parser(tokenizer, request.tools)
                 except RuntimeError as e:
                     logger.exception("Error in tool parser creation.")
                     return self.create_error_response(e)
@@ -1774,7 +1774,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 if tool_call_info.tools_called:
                     message = ChatMessage(
                         role=role,
-                        reasoning_content=reasoning_content,
+                        reasoning=reasoning,
                         content=tool_call_info.content,
                         tool_calls=tool_call_info.tool_calls,
                     )
@@ -1790,7 +1790,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         ret_content = tool_call_info.content
                     message = ChatMessage(
                         role=role,
-                        reasoning_content=reasoning_content,
+                        reasoning=reasoning,
                         content=ret_content,
                     )
 
@@ -1800,7 +1800,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     "Error in chat_completion_full_generator - cannot determine if tools should be extracted. "
                     "Returning a standard chat completion."
                 )
-                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
+                message = ChatMessage(role=role, reasoning=reasoning, content=content)
 
             choice_data = ChatCompletionResponseChoice(
                 index=output.index,

vllm_omni/entrypoints/openai/serving_speech.py

@@ -1642,3 +1642,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             succeeded=succeeded,
             failed=len(final_results) - succeeded,
         )
+
+
+ServingSpeech = OmniOpenAIServingSpeech

vllm_omni/entrypoints/utils.py

@@ -1,7 +1,7 @@
 import os
 import types
 from collections import Counter
-from dataclasses import asdict, fields, is_dataclass
+from dataclasses import fields, is_dataclass
 from pathlib import Path
 from typing import Any, get_args, get_origin
 
@@ -145,12 +145,21 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     if isinstance(obj, set):
         return list(obj)
     # Handle dataclass objects
-    # Note: asdict() recursively converts nested dataclasses but not Counter objects,
-    # so we need to recursively process the result
-    if is_dataclass(obj):
-        result = asdict(obj)
-        # Recursively process the result to convert any Counter objects
-        return _convert_dataclasses_to_dict(result)
+    # Use field iteration instead of asdict() to:
+    # 1. Only include init fields (non-init fields cause "unexpected kwarg" errors)
+    # 2. Skip None values matching field defaults (avoids Pydantic validation
+    #    failures when None is explicitly passed for non-Optional typed fields,
+    #    e.g. CompilationConfig.cudagraph_capture_sizes: list[int] = None)
+    if is_dataclass(obj) and not isinstance(obj, type):
+        result = {}
+        for f in fields(obj):
+            if not f.init:
+                continue
+            value = getattr(obj, f.name)
+            if value is None and f.default is None:
+                continue
+            result[f.name] = _convert_dataclasses_to_dict(value)
+        return result
     # Handle dictionaries (recurse into values) and filter out callables(cause error in OmegaConf.create)
     # Note: This must come AFTER Counter check since Counter is a dict subclass
     if isinstance(obj, dict):

vllm_omni/inputs/data.py

@@ -16,7 +16,7 @@ except ImportError:
 
 
 import torch
-from vllm.inputs.data import EmbedsPrompt, TextPrompt, TokenInputs, TokensPrompt
+from vllm.inputs import EmbedsPrompt, TextPrompt, TokensInput, TokensPrompt
 
 
 class OmniTextPrompt(TextPrompt):
@@ -59,10 +59,10 @@ class OmniTokensPrompt(TokensPrompt):
     additional_information: NotRequired[dict[str, Any]]
 
 
-class OmniTokenInputs(TokenInputs):
+class OmniTokenInputs(TokensInput):
     """Token inputs with optional embeddings and additional information.
 
-    Extends TokenInputs to support prompt embeddings and additional
+    Extends TokensInput to support prompt embeddings and additional
     information payloads for direct transfer between pipeline stages.
 
     Attributes:

vllm_omni/inputs/preprocess.py

@@ -1,10 +1,9 @@
 from typing import Any
 
 from typing_extensions import assert_never
-from vllm.inputs.data import EmbedsInputs, SingletonInputs
+from vllm.inputs import EmbedsInput, MultiModalInput, SingletonInput
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
-from vllm.multimodal.inputs import MultiModalInputs
 from vllm.renderers.inputs import SingletonDictPrompt
 
 from vllm_omni.inputs.data import (
@@ -30,7 +29,7 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         parsed_content: OmniTextPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-    ) -> OmniTokenInputs | MultiModalInputs:
+    ) -> OmniTokenInputs | MultiModalInput:
         """Process text prompts with support for mm_processor_kwargs.
 
         Extends base class to support mm_processor_kwargs without multi_modal_data.
@@ -40,7 +39,7 @@ class OmniInputPreprocessor(InputPreprocessor):
         prompt_text = parsed_content["prompt"]
         mm_processor_kwargs = parsed_content.get("mm_processor_kwargs") or {}
 
-        inputs: OmniTokenInputs | MultiModalInputs
+        inputs: OmniTokenInputs | MultiModalInput
         if multi_modal_data := parsed_content.get("multi_modal_data"):
             inputs = self._process_multimodal(
                 prompt_text,
@@ -86,14 +85,14 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         parsed_content: OmniTokensPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-    ) -> OmniTokenInputs | MultiModalInputs:
+    ) -> OmniTokenInputs | MultiModalInput:
         prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
         prompt_embeds = parsed_content.get("prompt_embeds")
         additional_information = parsed_content.get("additional_information")
 
         multi_modal_data = parsed_content.get("multi_modal_data")
 
-        inputs: OmniTokenInputs | MultiModalInputs
+        inputs: OmniTokenInputs | MultiModalInput
         if multi_modal_data:
             inputs = self._process_multimodal(
                 prompt_token_ids,
@@ -123,7 +122,7 @@ class OmniInputPreprocessor(InputPreprocessor):
     def _process_embeds(
         self,
         parsed_content: OmniEmbedsPrompt,
-    ) -> EmbedsInputs:
+    ) -> EmbedsInput:
         """Process embeddings prompt with omni-specific extensions.
 
         Extends base _process_embeds to handle additional_information payload
@@ -143,7 +142,7 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         prompt: SingletonDictPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-    ) -> SingletonInputs:
+    ) -> SingletonInput:
         """
         Extract the singleton inputs from a prompt.
 
@@ -153,7 +152,7 @@ class OmniInputPreprocessor(InputPreprocessor):
 
         Returns:
 
-        * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance
+        * [`SingletonInput`][vllm.inputs.engine.SingletonInput] instance
         """
         if "prompt_embeds" in prompt:
             return self._process_embeds(prompt)  # type: ignore[arg-type]

vllm_omni/model_executor/models/bagel/bagel.py

@@ -8,6 +8,7 @@ import torch.nn as nn
 from transformers import BatchFeature
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
@@ -20,7 +21,6 @@ from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
-    MultiModalDataDict,
     MultiModalFieldConfig,
     MultiModalKwargsItems,
 )

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -10,10 +10,11 @@ import torch.nn as nn
 from transformers.feature_extraction_utils import BatchFeature
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMultiModal
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
 from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,

vllm_omni/model_executor/models/glm_image/glm_image_ar.py

@@ -40,6 +40,7 @@ from vllm.config import CacheConfig, MultiModalConfig, VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.distributed import utils as dist_utils
+from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention.mm_encoder_attention import (
@@ -73,7 +74,6 @@ from vllm.model_executor.models.utils import (
 from vllm.model_executor.models.vision import get_vit_attn_backend
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
-    MultiModalDataDict,
     MultiModalFeatureSpec,
     MultiModalFieldConfig,
     MultiModalKwargsItems,

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -20,6 +20,7 @@ from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
 from vllm.distributed import get_pp_group
+from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
 from vllm.model_executor.layers.linear import (
@@ -58,7 +59,6 @@ from vllm.model_executor.models.utils import (
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.image import rgba_to_rgb
 from vllm.multimodal.inputs import (
-    MultiModalDataDict,
     MultiModalFeatureSpec,
     MultiModalFieldConfig,
     MultiModalKwargsItems,

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -10,6 +10,7 @@ from torch import nn
 from transformers import BatchFeature, Qwen2Config
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.models import SupportsPP
@@ -18,9 +19,7 @@ from vllm.model_executor.models.utils import init_vllm_registered_model
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     AudioItem,
-    ModalityData,
     MultiModalBatchedField,
-    MultiModalDataDict,
     MultiModalFieldConfig,
     MultiModalFieldElem,
     MultiModalKwargsItem,

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -13,6 +13,7 @@ from transformers.models.qwen2.modeling_qwen2 import (
 )
 from vllm.config import VllmConfig
 from vllm.forward_context import get_forward_context
+from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader, maybe_remap_kv_scale_name
@@ -34,7 +35,6 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
-    MultiModalDataDict,
     MultiModalFieldConfig,
     MultiModalKwargsItems,
 )

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -18,7 +18,7 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeThinkerConfig,
 )
 from vllm.config import ModelConfig, VllmConfig
-from vllm.inputs.data import PromptType, TokensPrompt
+from vllm.inputs import PromptType, TokensPrompt
 from vllm.logger import init_logger
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP, SupportsRealtime

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -46,7 +46,7 @@ from transformers.models.whisper import WhisperFeatureExtractor
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
-from vllm.inputs.data import PromptType
+from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
 from vllm.model_executor.layers.attention.mm_encoder_attention import (

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py

@@ -29,6 +29,7 @@ from mistral_common.tokens.tokenizers.audio import Audio, AudioEncoder
 from transformers import BatchFeature
 from transformers.tokenization_utils_base import TextInput
 from vllm.config import VllmConfig
+from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import SupportsMultiModal
@@ -39,7 +40,6 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
-    MultiModalDataDict,
     MultiModalFieldConfig,
     MultiModalKwargsItems,
     NestedTensors,

vllm_omni/patch.py

@@ -1,7 +1,7 @@
 import sys
 
 from aenum import extend_enum
-from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
+from vllm.inputs import TokensPrompt as _OriginalTokensPrompt
 from vllm.model_executor.layers.rotary_embedding import (
     MRotaryEmbedding as _OriginalMRotaryEmbedding,
 )

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -26,8 +26,8 @@ from vllm.v1.outputs import (
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.structured_output.utils import apply_grammar_bitmask
 from vllm.v1.utils import record_function_or_nullcontext
-from vllm.v1.worker import mamba_utils
 from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata
+from vllm.v1.worker.mamba_utils import preprocess_mamba
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm_ascend.ascend_forward_context import set_ascend_forward_context
 from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
@@ -243,7 +243,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 # '_update_states_after_model_execute', which is not overridden in vLLM-Ascend.
                 # We simply utilize the implementation in vLLM.
                 if self.cache_config.mamba_cache_mode == "align":
-                    mamba_utils.preprocess_mamba(
+                    preprocess_mamba(
                         scheduler_output,
                         self.kv_cache_config,
                         self.cache_config,

vllm_omni/worker/gpu_ar_model_runner.py

@@ -137,8 +137,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
             else:
                 logger.error("RoutedExpertsCapturer not initialized.")
 
-        if scheduler_output.preempted_req_ids and has_kv_transfer_group():
-            get_kv_transfer_group().handle_preemptions(scheduler_output.preempted_req_ids)
+        if has_kv_transfer_group():
+            kv_connector_metadata = scheduler_output.kv_connector_metadata
+            if kv_connector_metadata is not None:
+                get_kv_transfer_group().handle_preemptions(kv_connector_metadata)
 
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         with (
@@ -146,7 +148,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             self.synchronize_input_prep(),
         ):
             # Update persistent batch states.
-            self._update_states(scheduler_output)
+            deferred_state_corrections_fn = self._update_states(scheduler_output)
 
             if has_ec_transfer() and not get_ec_transfer().is_consumer:
                 with self.maybe_get_ec_connector_output(
@@ -417,6 +419,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
         )
         self.kv_connector_output = kv_connector_output
 
+        if deferred_state_corrections_fn:
+            deferred_state_corrections_fn()
+
         return None
 
     @torch.inference_mode()
@@ -474,8 +479,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
         with record_function_or_nullcontext("gpu_model_runner: sample"):
             sampler_output = self._sample(logits, spec_decode_metadata)
 
+        self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output)
+
         self._draft_token_ids = None
         self._draft_token_req_ids = None
+        self.valid_sampled_token_count_gpu = None
         self.input_batch.prev_sampled_token_ids = None
 
         def propose_draft_token_ids(sampled_token_ids):
@@ -515,7 +523,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 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,
+                        self.optimistic_seq_lens_cpu,
                         sampled_token_ids,
                         self.requests,
                         self.input_batch,

vllm_omni/worker/gpu_generation_model_runner.py

@@ -94,8 +94,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             else:
                 logger.error("RoutedExpertsCapturer not initialized.")
 
-        if scheduler_output.preempted_req_ids and has_kv_transfer_group():
-            get_kv_transfer_group().handle_preemptions(scheduler_output.preempted_req_ids)
+        if has_kv_transfer_group():
+            kv_connector_metadata = scheduler_output.kv_connector_metadata
+            if kv_connector_metadata is not None:
+                get_kv_transfer_group().handle_preemptions(kv_connector_metadata)
 
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         with (
@@ -104,7 +106,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         ):
             if self.model_config.async_chunk and num_scheduled_tokens:
                 self._update_request_states(scheduler_output)
-            self._update_states(scheduler_output)
+            deferred_state_corrections_fn = self._update_states(scheduler_output)
             if not scheduler_output.total_num_scheduled_tokens:
                 return EMPTY_MODEL_RUNNER_OUTPUT
 
@@ -309,6 +311,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             slot_mappings,  # OMNI: pass slot_mappings for upstream v1 API compatibility
         )
         self.kv_connector_output = kv_connector_output
+
+        if deferred_state_corrections_fn:
+            deferred_state_corrections_fn()
+
         return None
 
     @torch.inference_mode()
@@ -636,11 +642,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                     seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]  # type: ignore[assignment]
                 else:
                     seq_lens = max_query_len  # type: ignore[assignment]
-                self.seq_lens.np[:num_reqs] = seq_lens
-                self.seq_lens.np[num_reqs:] = 0
-                self.seq_lens.copy_to_gpu()
+                self.seq_lens[:num_reqs] = (
+                    seq_lens
+                    if isinstance(seq_lens, int)
+                    else torch.tensor(seq_lens, dtype=torch.int32, device=self.device)
+                )
+                self.seq_lens[num_reqs:] = 0
 
-                cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+                cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self._arange_scratch)
                 self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
                 self.query_start_loc.copy_to_gpu()
 
@@ -696,7 +705,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             elif self.uses_xdrope_dim > 0:
                 positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
             else:
-                positions = self.positions.gpu[:num_tokens_padded]
+                positions = self.positions[:num_tokens_padded]
 
             if get_pp_group().is_first_rank:
                 intermediate_tensors = None

vllm_omni/worker/gpu_model_runner.py

@@ -243,7 +243,7 @@ class OmniGPUModelRunner(GPUModelRunner):
 
                 mrope_pos_ptr += completion_part_len
 
-    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
+    def _update_states(self, scheduler_output: "SchedulerOutput"):
         """Update the cached states and the persistent batch with the scheduler
         output.
 
@@ -271,7 +271,7 @@ class OmniGPUModelRunner(GPUModelRunner):
 
         # Zero GPU memory for freshly allocated cache blocks to prevent
         # stale NaN/data from corrupting attention or SSM computation.
-        if hasattr(scheduler_output, "new_block_ids_to_zero") and scheduler_output.new_block_ids_to_zero:
+        if scheduler_output.new_block_ids_to_zero:
             self._zero_block_ids(scheduler_output.new_block_ids_to_zero)
 
         # Free the cached encoder outputs.
@@ -300,7 +300,11 @@ class OmniGPUModelRunner(GPUModelRunner):
         for req_id in unscheduled_req_ids:
             self.input_batch.remove_request(req_id)
 
+        if self.use_async_spec_decode:
+            self.prev_num_draft_tokens.np.fill(0)
+
         reqs_to_add: list[CachedRequestState] = []
+        deferred_spec_decode_corrections = []
         # Add new requests to the cached states.
         for new_req_data in scheduler_output.scheduled_new_reqs:
             req_id = new_req_data.req_id
@@ -398,10 +402,6 @@ class OmniGPUModelRunner(GPUModelRunner):
         req_data = scheduler_output.scheduled_cached_reqs
         scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
 
-        # Wait until valid_sampled_tokens_count is copied to cpu,
-        # then use it to update actual num_computed_tokens of each request.
-        valid_sampled_token_count = self._get_valid_sampled_token_count()
-
         for i, req_id in enumerate(req_data.req_ids):
             req_state = self.requests[req_id]
             num_computed_tokens = req_data.num_computed_tokens[i]
@@ -427,12 +427,18 @@ class OmniGPUModelRunner(GPUModelRunner):
                 if req_index is None:
                     req_state.prev_num_draft_len = 0
                 else:
-                    assert self.input_batch.prev_req_id_to_index is not None
-                    prev_req_index = self.input_batch.prev_req_id_to_index[req_id]
-                    num_accepted = valid_sampled_token_count[prev_req_index] - 1
-                    num_rejected = req_state.prev_num_draft_len - num_accepted
-                    num_computed_tokens -= num_rejected
-                    req_state.output_token_ids.extend([-1] * num_accepted)
+                    optimistic_num_accepted = req_state.prev_num_draft_len
+                    req_state.output_token_ids.extend([-1] * optimistic_num_accepted)
+
+                    deferred_spec_decode_corrections.append((req_id, optimistic_num_accepted, req_state))
+
+                    prev_req_index = (
+                        self.input_batch.prev_req_id_to_index.get(req_id)
+                        if self.input_batch.prev_req_id_to_index
+                        else None
+                    )
+                    if prev_req_index is not None:
+                        self.prev_num_draft_tokens.np[prev_req_index] = optimistic_num_accepted
 
             # Update the cached states.
             req_state.num_computed_tokens = num_computed_tokens
@@ -449,7 +455,8 @@ class OmniGPUModelRunner(GPUModelRunner):
                         req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
             elif num_output_tokens < len(req_state.output_token_ids):
                 # Some output tokens were discarded due to a sync-KV-load
-                # failure. Align the cached state.
+                # failure, or output_token_ids was inflated by the optimistic
+                # extend above (async spec decode). Align the cached state.
                 del req_state.output_token_ids[num_output_tokens:]
                 if req_index is not None:
                     end_idx = self.input_batch.num_prompt_tokens[req_index] + num_output_tokens
@@ -513,6 +520,35 @@ class OmniGPUModelRunner(GPUModelRunner):
         # Refresh batch metadata with any pending updates.
         self.input_batch.refresh_metadata()
 
+        if deferred_spec_decode_corrections:
+
+            def correct_spec_decode_token_counts():
+                valid_sampled_token_count = self._get_valid_sampled_token_count()
+                if not valid_sampled_token_count:
+                    return
+                prev_req_id_to_index = self.input_batch.prev_req_id_to_index
+                if not prev_req_id_to_index:
+                    return
+                for (
+                    req_id,
+                    optimistic_num_accepted,
+                    req_state,
+                ) in deferred_spec_decode_corrections:
+                    prev_req_index = prev_req_id_to_index.get(req_id)
+                    if prev_req_index is None:
+                        continue
+                    num_accepted = valid_sampled_token_count[prev_req_index] - 1
+                    correction = optimistic_num_accepted - num_accepted
+                    req_state.num_computed_tokens -= correction
+                    cur_req_index = self.input_batch.req_id_to_index.get(req_id)
+                    if cur_req_index is None:
+                        continue
+                    self.input_batch.num_computed_tokens_cpu[cur_req_index] -= correction
+
+            return correct_spec_decode_token_counts
+        else:
+            return None
+
     @torch.inference_mode()
     def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict:
         if (
@@ -704,11 +740,14 @@ class OmniGPUModelRunner(GPUModelRunner):
                     seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]  # type: ignore[assignment]
                 else:
                     seq_lens = max_query_len  # type: ignore[assignment]
-                self.seq_lens.np[:num_reqs] = seq_lens
-                self.seq_lens.np[num_reqs:] = 0
-                self.seq_lens.copy_to_gpu()
+                self.seq_lens[:num_reqs] = (
+                    seq_lens
+                    if isinstance(seq_lens, int)
+                    else torch.tensor(seq_lens, dtype=torch.int32, device=self.device)
+                )
+                self.seq_lens[num_reqs:] = 0
 
-                cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+                cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self._arange_scratch)
                 self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
                 self.query_start_loc.copy_to_gpu()
 
@@ -759,7 +798,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             elif self.uses_xdrope_dim > 0:
                 positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
             else:
-                positions = self.positions.gpu[:num_tokens_padded]
+                positions = self.positions[:num_tokens_padded]
 
             if get_pp_group().is_first_rank:
                 intermediate_tensors = None
@@ -1164,7 +1203,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         elif self.uses_xdrope_dim > 0:
             positions = self.xdrope_positions.gpu[:, :num_input_tokens]
         else:
-            positions = self.positions.gpu[:num_input_tokens]
+            positions = self.positions[:num_input_tokens]
 
         if is_first_rank:
             intermediate_tensors = None