Dev/rebase v0170 (#1639)

Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

docker/Dockerfile.ci

@@ -1,9 +1,8 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.16.0
+ARG VLLM_BASE_TAG=v0.17.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
-
 COPY . .
 
 # Install system dependencies

docker/Dockerfile.rocm

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

tests/diffusion/attention/test_attention_sp.py

@@ -414,7 +414,8 @@ def ulysses_attention_on_test_model(
             # Generate and save full input data with fixed seed
             # Reinitialize RNG to ensure reproducibility
             torch.manual_seed(42)
-            current_omni_platform.seed_everything(42)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(42)
             full_hidden_states = torch.randn(
                 (batch_size, seq_len, hidden_size),
                 dtype=dtype,

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -66,14 +66,13 @@ def shm_connector():
 
 
 def test_put_get_inline(shm_connector):
-    """Test inline transfer for small data."""
+    """Test transfer for small data (currently always uses SHM)."""
     data = {"small": "data"}
-    # Ensure data is smaller than threshold (100 bytes)
 
     success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
     assert success is True
-    assert "inline_bytes" in metadata
-    assert "shm" not in metadata
+    assert "shm" in metadata
+    assert "size" in metadata
 
     # Retrieve
     retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
@@ -87,7 +86,7 @@ def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.M
     data = {"large": "x" * 200}
 
     # Mock SHM return values
-    mock_handle = {"name": "test_shm", "size": 200}
+    mock_handle = {"name": "req_2", "size": 200}
     mock_write = mocker.MagicMock(return_value=mock_handle)
     monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write)
 

tests/e2e/online_serving/test_qwen3_tts.py

@@ -28,7 +28,7 @@ def get_stage_config(name: str = "qwen3_tts.yaml"):
     return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="class")
 def omni_server():
     """Start vLLM-Omni server with CustomVoice model."""
     stage_config_path = get_stage_config()
@@ -225,7 +225,7 @@ class TestQwen3TTSAPIEndpoints:
         assert len(data["data"]) > 0
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="class")
 def omni_server_no_async_chunk():
     """Start vLLM-Omni server with non-async-chunk config."""
     stage_config_path = get_stage_config("qwen3_tts_no_async_chunk.yaml")

tests/e2e/stage_configs/qwen2_5_omni_ci.yaml

@@ -74,12 +74,13 @@ stage_args:
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.15
+      gpu_memory_utilization: 0.9 #increase the gpu memory utilization to enable the test on H800
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
       engine_output_type: audio
-      max_num_batched_tokens: 100000
+      max_num_batched_tokens: 4096
+      max_model_len: 4096
     engine_input_source: [1]
     final_output: true
     final_output_type: audio

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -48,7 +48,7 @@ stage_args:
     model_arch: Qwen3OmniMoeForConditionalGeneration
     worker_type: ar
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    gpu_memory_utilization: 0.6
+    gpu_memory_utilization: 0.5
     enforce_eager: false
     trust_remote_code: true
     engine_output_type: latent  # Output codec codes for code2wav

tests/entrypoints/test_omni_llm.py

@@ -1,3 +1,4 @@
+import inspect
 import uuid
 import warnings
 from queue import Empty, Queue
@@ -1194,3 +1195,72 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo
 
     # Verify stop_stage_worker was called (process should be set)
     assert omni.stage_list[0]._proc is not None
+
+
+# ---------------------------------------------------------------------------
+# Signature compatibility tests — catch upstream API drift early
+# ---------------------------------------------------------------------------
+
+_INHERITANCE_PAIRS: list[tuple[str, str, str, str]] = [
+    (
+        "vllm.entrypoints.llm",
+        "LLM",
+        "vllm_omni.entrypoints.omni_llm",
+        "OmniLLM",
+    ),
+]
+
+
+def _import_class(module_path: str, class_name: str):
+    mod = __import__(module_path, fromlist=[class_name])
+    return getattr(mod, class_name)
+
+
+@pytest.mark.parametrize(
+    "up_mod,up_cls,omni_mod,omni_cls",
+    _INHERITANCE_PAIRS,
+    ids=[f"{pair[3]}({pair[1]})" for pair in _INHERITANCE_PAIRS],
+)
+def test_overridden_method_signatures_compatible(up_mod: str, up_cls: str, omni_mod: str, omni_cls: str):
+    """All params accepted by a base-class method must also be accepted by
+    the overriding method in the vllm-omni subclass.  This prevents
+    TypeError at runtime when upstream callers pass new arguments."""
+    BaseCls = _import_class(up_mod, up_cls)
+    OmniCls = _import_class(omni_mod, omni_cls)
+
+    failures: list[str] = []
+    for name in dir(OmniCls):
+        if name.startswith("__") and name != "__init__":
+            continue
+        omni_method = getattr(OmniCls, name, None)
+        base_method = getattr(BaseCls, name, None)
+        if not (callable(omni_method) and callable(base_method)):
+            continue
+        if omni_method is base_method:
+            continue
+        try:
+            base_sig = inspect.signature(base_method)
+            omni_sig = inspect.signature(omni_method)
+        except (ValueError, TypeError):
+            continue
+
+        omni_has_var_keyword = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in omni_sig.parameters.values())
+
+        base_params = base_sig.parameters
+        omni_params = set(omni_sig.parameters.keys())
+
+        missing = []
+        for pname, param in base_params.items():
+            if pname in omni_params:
+                continue
+            if omni_has_var_keyword and param.kind in (
+                inspect.Parameter.KEYWORD_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            ):
+                continue
+            missing.append(pname)
+
+        if missing:
+            failures.append(f"{omni_cls}.{name}() missing params {sorted(missing)}; base={base_sig}, omni={omni_sig}")
+
+    assert not failures, f"Signature mismatches found ({len(failures)}):\n" + "\n".join(f"  - {f}" for f in failures)

tests/perf/stage_configs/qwen3_omni.yaml

@@ -85,7 +85,7 @@ stage_args:
       engine_output_type: audio  # Final output: audio waveform
       gpu_memory_utilization: 0.1
       distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
+      max_num_batched_tokens: 100000
       hf_config_name: thinker_config
     engine_input_source: [1]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav

tests/worker/test_gpu_generation_model_runner.py

@@ -34,6 +34,7 @@ def _make_runner(multimodal_outputs):
     runner.use_async_scheduling = False
     runner.device = torch.device("cpu")
     runner.supports_mm_inputs = False
+    runner.speculative_config = None
     return runner
 
 

vllm_omni/benchmarks/metrics/metrics.py

@@ -2,10 +2,10 @@ import warnings
 from dataclasses import dataclass
 
 import numpy as np
-from transformers import PreTrainedTokenizerBase
 from vllm.benchmarks.datasets import SampleRequest
 from vllm.benchmarks.lib.endpoint_request_func import RequestFuncOutput
 from vllm.benchmarks.serve import MILLISECONDS_TO_SECONDS_CONVERSION, TERM_PLOTLIB_AVAILABLE, BenchmarkMetrics, TaskType
+from vllm.tokenizers import TokenizerLike
 
 
 @dataclass
@@ -134,7 +134,7 @@ def calculate_metrics(
     input_requests: list[SampleRequest],
     outputs: list[RequestFuncOutput],
     dur_s: float,
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: TokenizerLike | None,
     selected_percentiles: list[float],
     goodput_config_dict: dict[str, float],
     task_type,
@@ -169,17 +169,21 @@ def calculate_metrics(
     audio_rtfs: list[float] = []
     audio_duration: list[float] = []
     audio_frames: list[int] = []
+    input_audio_duration = 0.0
     for i in range(len(outputs)):
         if outputs[i].success:
             output_len = outputs[i].output_tokens
 
             if not output_len:
-                # We use the tokenizer to count the number of output tokens
-                # for some serving backends instead of looking at
-                # len(outputs[i].itl) since multiple output tokens may be
-                # bundled together
-                # Note : this may inflate the output token count slightly
-                output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
+                if tokenizer is None:
+                    output_len = 1
+                else:
+                    # We use the tokenizer to count the number of output tokens
+                    # for some serving backends instead of looking at
+                    # len(outputs[i].itl) since multiple output tokens may be
+                    # bundled together
+                    # Note : this may inflate the output token count slightly
+                    output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
             actual_output_lens.append(output_len)
             total_input += input_requests[i].prompt_len
             tpot = 0
@@ -196,6 +200,7 @@ def calculate_metrics(
             audio_duration.append(getattr(outputs[i], "audio_duration", 0.0))
             audio_frames.append(getattr(outputs[i], "audio_frames", 0.0))
             e2els.append(outputs[i].latency)
+            input_audio_duration += outputs[i].input_audio_duration
             completed += 1
         else:
             actual_output_lens.append(0)
@@ -335,6 +340,7 @@ def calculate_metrics(
         percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles],
         max_output_tokens_per_s=max_output_tokens_per_s,
         max_concurrent_requests=max_concurrent_requests,
+        rtfx=input_audio_duration / dur_s,
     )
     print_metrics(
         task_type,

vllm_omni/benchmarks/patch/patch.py

@@ -16,7 +16,6 @@ from typing import Literal
 import aiohttp
 from pydub import AudioSegment
 from tqdm.asyncio import tqdm
-from transformers import PreTrainedTokenizerBase
 from vllm.benchmarks import datasets
 from vllm.benchmarks.datasets import SampleRequest
 from vllm.benchmarks.lib.endpoint_request_func import (
@@ -31,6 +30,7 @@ from vllm.benchmarks.lib.endpoint_request_func import (
     _validate_api_url,
 )
 from vllm.logger import init_logger
+from vllm.tokenizers import TokenizerLike
 
 logger = init_logger(__name__)
 from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset
@@ -134,6 +134,8 @@ async def async_request_openai_chat_omni_completions(
 
                     messages = handler.add_chunk(chunk_bytes)
                     for message in messages:
+                        if type(message) is bytes:
+                            message = message.decode("utf-8")
                         # NOTE: SSE comments (often used as pings) start with
                         # a colon. These are not JSON data payload and should
                         # be skipped.
@@ -311,7 +313,7 @@ async def benchmark(
     base_url: str,
     model_id: str,
     model_name: str,
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: TokenizerLike | None,
     input_requests: list[SampleRequest],
     logprobs: int | None,
     request_rate: float,
@@ -576,6 +578,7 @@ async def benchmark(
             "errors": [output.error for output in outputs],
             "max_output_tokens_per_s": metrics.max_output_tokens_per_s,
             "max_concurrent_requests": metrics.max_concurrent_requests,
+            "rtfx": metrics.rtfx,
         }
     else:
         result = {

vllm_omni/config/model.py

@@ -101,6 +101,7 @@ class OmniModelConfig(ModelConfig):
     def __post_init__(
         self,
         # Multimodal config init vars
+        language_model_only: bool,
         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,
@@ -117,6 +118,7 @@ class OmniModelConfig(ModelConfig):
     ) -> None:
         # Call parent's __post_init__ to handle all standard ModelConfig initialization
         super().__post_init__(
+            language_model_only=language_model_only,
             limit_mm_per_prompt=limit_mm_per_prompt,
             enable_mm_embeds=enable_mm_embeds,
             media_io_kwargs=media_io_kwargs,

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -246,7 +246,7 @@ class OmniARScheduler(VLLMScheduler):
             generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
 
             scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
-            if scheduled_spec_token_ids:
+            if scheduled_spec_token_ids and generated_token_ids:
                 num_draft_tokens = len(scheduled_spec_token_ids)
                 num_accepted = len(generated_token_ids) - 1
                 num_rejected = num_draft_tokens - num_accepted
@@ -481,10 +481,7 @@ class OmniARScheduler(VLLMScheduler):
         assert request.is_finished()
 
         # 1. Standard cleanup parts from base _free_request
-        connector_delay_free_blocks = False
-        kv_xfer_params = None
-        if self.connector is not None:
-            connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
+        connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
 
         self.encoder_cache_manager.free(request)
         request_id = request.request_id

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -6,6 +6,7 @@ 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.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.sched.interface import PauseState
 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 Scheduler as VLLMScheduler
@@ -41,8 +42,12 @@ class OmniGenerationScheduler(VLLMScheduler):
         """
 
         token_budget = self.max_num_scheduled_tokens
+        if self._pause_state == PauseState.PAUSED_ALL:
+            token_budget = 0
         scheduled_timestamp = time.monotonic()
 
+        self.kv_cache_manager.new_step_starts()
+
         scheduled_new_reqs: list[Request] = []
 
         req_to_new_blocks: dict[str, KVCacheBlocks] = {}
@@ -108,6 +113,8 @@ class OmniGenerationScheduler(VLLMScheduler):
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens
             cached_prompt_token_ids[request.request_id] = request.prompt_token_ids
+            if request.num_cached_tokens < 0:
+                request.num_cached_tokens = num_computed_tokens
             cached_additional_information[request.request_id] = getattr(request, "additional_information", None)
             token_budget -= num_new_tokens
             scheduled_running_reqs.append(request)
@@ -119,7 +126,12 @@ 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 len(self.running) < self.max_num_running_reqs:
+        while (
+            self.waiting
+            and token_budget > 0
+            and len(self.running) < self.max_num_running_reqs
+            and self._pause_state == PauseState.UNPAUSED
+        ):
             request = self.waiting.peek_request()
             # OMNI: Skip requests that are not in self.requests
             if request.request_id not in self.requests or (
@@ -172,6 +184,8 @@ class OmniGenerationScheduler(VLLMScheduler):
 
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens
+            if request.num_cached_tokens < 0:
+                request.num_cached_tokens = 0
             token_budget -= num_new_tokens
             scheduled_new_reqs.append(request)
 

vllm_omni/core/sched/output.py

@@ -3,30 +3,26 @@ from dataclasses import dataclass, field
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
 from vllm.v1.request import Request
 
-from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
+from vllm_omni.engine import AdditionalInformationPayload
 
 
 @dataclass
 class OmniNewRequestData(NewRequestData):
     """New request data for omni models with embeddings support.
 
-    Extends NewRequestData to include prompt embeddings and additional
-    information for direct transfer between pipeline stages.
+    Extends NewRequestData to include additional information for direct
+    transfer between pipeline stages.
+
+    Note: prompt_embeds is inherited from NewRequestData
+    (torch.Tensor | None).
 
     Args:
-        prompt_embeds: Optional serialized prompt embeddings payload
-            (overrides parent's torch.Tensor type with PromptEmbedsPayload
-            for cross-process serialization)
         external_req_id: Optional external request ID for tracking
         additional_information: Optional serialized additional information
             dictionary containing tensors or lists
     """
 
-    # Optional serialized prompt embeddings (override parent type for serialization)
-    prompt_embeds: PromptEmbedsPayload | None = None  # type: ignore[assignment]
-    # Optional external request ID for tracking
     external_req_id: str | None = None
-    # Optional serialized additional information
     additional_information: AdditionalInformationPayload | None = None
 
     @classmethod

vllm_omni/diffusion/utils/hf_utils.py

@@ -53,10 +53,8 @@ def is_diffusion_model(model_name: str) -> bool:
 
     # Strategy 2: Check using vllm's utility (works for both local and remote models)
     try:
-        # Try to get model_index.json using vllm's utility
         config_dict = get_hf_file_to_dict("model_index.json", model_name)
-        # Verify it has the required fields for a diffusers model
-        if config_dict.get("_class_name") and config_dict.get("_diffusers_version"):
+        if config_dict is not None and config_dict.get("_class_name") and config_dict.get("_diffusers_version"):
             logger.debug("Detected diffusion model via model_index.json")
             return True
     except Exception as e:

vllm_omni/distributed/omni_connectors/adapter.py

@@ -35,12 +35,26 @@ def try_send_via_connector(
     try:
         t0 = time.time()
 
+        # Strip non-serializable multimodal feature fields from original_prompt
+        # before including it in metadata.  After stage-0 runs, the TokPrompt
+        # returned by render_chat_async may carry processed multimodal features
+        # (mm_kwargs, mm_placeholders, mm_hashes) that contain MultiModalKwargsItems
+        # objects, which are not supported by OmniMsgpackEncoder.  The receiving
+        # side (try_recv_via_connector) only extracts "engine_inputs" from the
+        # payload and never uses "original_prompt", so stripping these fields
+        # only affects debug metadata and is safe.
+        _MM_FEATURE_KEYS = frozenset({"mm_kwargs", "mm_placeholders", "mm_hashes"})
+        if isinstance(original_prompt, dict) and any(k in original_prompt for k in _MM_FEATURE_KEYS):
+            safe_prompt = {k: v for k, v in original_prompt.items() if k not in _MM_FEATURE_KEYS}
+        else:
+            safe_prompt = original_prompt
+
         # Prepare data for connector
         payload_data = {
             "engine_inputs": next_inputs,
             "sampling_params": sampling_params,
             "metadata": {
-                "original_prompt": original_prompt,
+                "original_prompt": safe_prompt,
                 "stage_transition": f"{stage_id}->{next_stage_id}",
                 "timestamp": time.time(),
             },

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -48,6 +48,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             payload = self.serialize_obj(data)
             size = len(payload)
 
+            # Currently, we always use SHM.
             if True:
                 # Use Shared Memory
                 lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock"

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -39,7 +39,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs
         self.connector = self.create_connector(model_config)
         super().__init__(model_config)
-        self.model_mode = getattr(model_config, "worker_type", "ar")
+        self.model_mode = getattr(model_config, "worker_type", None) or "ar"
         # State specific to Chunk management
         self.custom_process_next_stage_input_func = None
         custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None)

vllm_omni/distributed/omni_connectors/utils/serialization.py

@@ -145,7 +145,7 @@ class OmniMsgpackEncoder:
             "encoder_prompt": obj.encoder_prompt,
             "encoder_prompt_token_ids": obj.encoder_prompt_token_ids,
             "num_cached_tokens": obj.num_cached_tokens,
-            "multi_modal_placeholders": obj.multi_modal_placeholders,
+            "multi_modal_placeholders": getattr(obj, "multi_modal_placeholders", None),
             "kv_transfer_params": obj.kv_transfer_params,
         }
         # Handle dynamically added multimodal_output attribute
@@ -301,15 +301,22 @@ class OmniMsgpackDecoder:
         """Decode dict to RequestOutput.
 
         RequestOutput is not a dataclass, so msgspec.convert doesn't work.
-        We construct it manually, passing all known fields via **kwargs.
+        We construct it manually using only the known __init__ parameters to
+        avoid triggering the "Ignoring extra arguments" warning in vllm.
+        Fields that are not part of RequestOutput.__init__ (e.g.
+        multi_modal_placeholders, multimodal_output) are extracted first and
+        then restored as dynamic attributes after construction.
         """
-        # Extract multimodal_output before constructing (it's dynamically added)
+        # Extract dynamically-added / non-init fields before constructing so
+        # they are not passed as unknown **kwargs to RequestOutput.__init__.
         mm_output = obj.pop("multimodal_output", None)
+        multi_modal_placeholders = obj.pop("multi_modal_placeholders", None)
 
-        # RequestOutput.__init__ accepts **kwargs for forward compatibility
         ro = RequestOutput(**obj)
 
-        # Restore dynamically added multimodal_output attribute
+        # Restore dynamic attributes that are not part of __init__.
+        if multi_modal_placeholders is not None:
+            setattr(ro, "multi_modal_placeholders", multi_modal_placeholders)
         if mm_output is not None:
             setattr(ro, "multimodal_output", mm_output)
         return ro

vllm_omni/engine/__init__.py

@@ -56,19 +56,19 @@ class AdditionalInformationPayload(msgspec.Struct):
 class OmniEngineCoreRequest(EngineCoreRequest):
     """Engine core request for omni models with embeddings support.
 
-    Extends the base EngineCoreRequest with support for prompt embeddings
-    and additional information payloads, enabling direct transfer of
-    pre-computed embeddings between pipeline stages.
+    Extends the base EngineCoreRequest with support for additional
+    information payloads, enabling direct transfer of pre-computed data
+    between pipeline stages.
+
+    Note: prompt_embeds is inherited from EngineCoreRequest
+    (torch.Tensor | None). PromptEmbedsPayload should be decoded to
+    torch.Tensor before constructing this request.
 
     Attributes:
-        prompt_embeds: Optional serialized prompt embeddings payload for
-            direct transfer between stages
         additional_information: Optional serialized additional information
             dictionary containing tensors or lists to pass along with the request
     """
 
-    # Optional prompt embeddings (direct-transfer version)
-    prompt_embeds: PromptEmbedsPayload | None = None
     # Optional additional information dictionary (serialized)
     additional_information: AdditionalInformationPayload | None = None
 

vllm_omni/engine/arg_utils.py

@@ -116,6 +116,7 @@ class OmniEngineArgs(EngineArgs):
 
         # Keep compatibility when async args are constructed from partial payloads.
         limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
+        language_model_only = getattr(self, "language_model_only", False)
         enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
         interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
         media_io_kwargs = getattr(self, "media_io_kwargs", {})
@@ -169,6 +170,7 @@ class OmniEngineArgs(EngineArgs):
             skip_tokenizer_init=self.skip_tokenizer_init,
             enable_prompt_embeds=self.enable_prompt_embeds,
             served_model_name=self.served_model_name,
+            language_model_only=language_model_only,
             limit_mm_per_prompt=limit_mm_per_prompt,
             enable_mm_embeds=enable_mm_embeds,
             interleave_mm_strings=interleave_mm_strings,
@@ -276,6 +278,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
 
         # Keep compatibility when async args are constructed from partial payloads.
         limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
+        language_model_only = getattr(self, "language_model_only", False)
         enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
         interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
         media_io_kwargs = getattr(self, "media_io_kwargs", {})
@@ -329,6 +332,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
             skip_tokenizer_init=self.skip_tokenizer_init,
             enable_prompt_embeds=self.enable_prompt_embeds,
             served_model_name=self.served_model_name,
+            language_model_only=language_model_only,
             limit_mm_per_prompt=limit_mm_per_prompt,
             enable_mm_embeds=enable_mm_embeds,
             interleave_mm_strings=interleave_mm_strings,

vllm_omni/engine/input_processor.py

@@ -1,6 +1,6 @@
 import time
 from collections.abc import Mapping
-from typing import Any, cast
+from typing import Any
 
 import numpy as np
 import torch
@@ -9,26 +9,14 @@ from vllm.inputs import ProcessorInputs, PromptType
 from vllm.inputs.parse import split_enc_dec_inputs
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict
-
-try:
-    from vllm.multimodal.processing import set_request_id
-except ImportError:  # vllm without set_request_id (older releases)
-    from contextlib import contextmanager
-
-    @contextmanager
-    def set_request_id(_request_id: str):
-        yield
-
-
+from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.multimodal.utils import argsort_mm_positions
-from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
-from vllm.renderers.inputs import DictPrompt, TokPrompt
+from vllm.renderers import BaseRenderer
 from vllm.sampling_params import SamplingParams
 from vllm.tasks import SupportedTask
 from vllm.utils import length_from_prompt_token_ids_or_embeds
-from vllm.utils.torch_utils import set_default_torch_num_threads
+from vllm.utils.jsontree import json_iter_leaves
 from vllm.v1.engine.input_processor import InputProcessor
 
 from vllm_omni.engine import (
@@ -42,6 +30,35 @@ from vllm_omni.lora.request import LoRARequest
 
 logger = init_logger(__name__)
 
+_OMNI_EXTRA_KEYS = (
+    "additional_information",
+    "prompt_embeds",
+    "negative_prompt",
+    "negative_prompt_embeds",
+)
+
+
+def reinject_omni_fields(
+    results: list[ProcessorInputs],
+    original_prompts: list[dict],
+) -> None:
+    """Re-inject omni-specific fields that the upstream renderer discards.
+
+    The upstream renderer's ``process_for_engine`` creates new dicts that only
+    copy standard vLLM fields (prompt_token_ids, multi_modal_data, …).
+    Omni-specific fields such as ``additional_information`` and
+    ``prompt_embeds`` are silently dropped.  This helper copies them back from
+    the *original* parsed prompts into the renderer outputs so they survive
+    into ``OmniInputProcessor.process_inputs()``.
+    """
+    for result, orig in zip(results, original_prompts):
+        if not isinstance(orig, dict):
+            continue
+        for key in _OMNI_EXTRA_KEYS:
+            val = orig.get(key)
+            if val is not None and key not in result:
+                result[key] = val
+
 
 class OmniInputProcessor(InputProcessor):
     """Processor for omni models, handling multimodal inputs and embeddings.
@@ -88,28 +105,29 @@ class OmniInputProcessor(InputProcessor):
     def __init__(
         self,
         vllm_config: VllmConfig,
+        renderer: BaseRenderer | None = None,
+        *,
         mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
     ):
-        super().__init__(vllm_config, mm_registry)
+        super().__init__(vllm_config, renderer=renderer, mm_registry=mm_registry)
         self.input_preprocessor = OmniInputPreprocessor(
-            self.model_config,
-            vllm_config.observability_config,
-            mm_registry,
-            mm_processor_cache=self.mm_processor_cache,
+            vllm_config=vllm_config,
+            renderer=self.renderer,
+            mm_registry=mm_registry,
         )
 
     def process_inputs(
         self,
         request_id: str,
-        prompt: PromptType | DictPrompt | TokPrompt,
+        prompt: PromptType | ProcessorInputs,
         params: SamplingParams | PoolingParams,
+        supported_tasks: tuple[SupportedTask, ...] = ("generate",),
         arrival_time: float | None = None,
         lora_request: LoRARequest | None = None,
         tokenization_kwargs: dict[str, Any] | None = None,
         trace_headers: Mapping[str, str] | None = None,
         priority: int = 0,
         data_parallel_rank: int | None = None,
-        supported_tasks: tuple[SupportedTask, ...] | None = None,
         resumable: bool = False,
     ) -> OmniEngineCoreRequest:
         """Process input prompt into an engine core request.
@@ -123,6 +141,7 @@ class OmniInputProcessor(InputProcessor):
             request_id: Unique identifier for this request
             prompt: Input prompt (text, token IDs, embeddings, or multimodal)
             params: Sampling or pooling parameters for generation
+            supported_tasks: Tuple of supported tasks for validation
             arrival_time: Optional arrival timestamp (defaults to current time)
             lora_request: Optional LoRA adapter request
             tokenization_kwargs: Optional additional tokenization arguments
@@ -130,19 +149,17 @@ class OmniInputProcessor(InputProcessor):
             priority: Request priority (higher values processed first)
             data_parallel_rank: Optional data parallel rank for distributed
                 inference
+            resumable: Whether the request supports streaming input
 
         Returns:
-            Tuple of (prompt_string, OmniEngineCoreRequest) where:
-                - prompt_string: The original prompt as a string, or None if
-                  using embeddings
-                - OmniEngineCoreRequest: Processed request ready for the engine
+            OmniEngineCoreRequest ready for the engine
 
         Raises:
             ValueError: If data_parallel_rank is out of range or prompt_embeds
                 has incorrect shape
         """
-        self._validate_lora(lora_request)
         self._validate_params(params, supported_tasks)
+        self._validate_lora(lora_request)
 
         parallel_config = self.vllm_config.parallel_config
         dp_size = parallel_config.data_parallel_size
@@ -151,50 +168,23 @@ class OmniInputProcessor(InputProcessor):
         if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks):
             raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).")
 
-        if arrival_time is None:
-            arrival_time = time.time()
-
-        # Optionally generate multimodal hash overrides to avoid hashing
-        # multimodal data items by their content as their identifiers.
-
-        # NOTE: when users explicitly turn off BOTH prefix caching and input
-        # processing caching, no multimodal features or embeddings will be
-        # reused across requests, therefore identifying multimodal data items
-        # by their content is no longer necessary, and we create uuids with
-        # request id-modality-index as multimodal hash overrides.
-        if (
-            self.model_config.multimodal_config
-            and self.model_config.multimodal_config.mm_processor_cache_gb == 0
-            and not self.cache_config.enable_prefix_caching
-        ):
-            mm_uuids = self._maybe_build_mm_uuids(request_id, prompt)
+        # Short-circuit for prompts already processed by the renderer
+        # (they carry a "type" key).  Raw prompts must still go through the
+        # omni preprocessor which preserves additional_information, etc.
+        if isinstance(prompt, dict) and "type" in prompt:
+            if arrival_time is None:
+                arrival_time = prompt.get("arrival_time", time.time())
+            processed_inputs: ProcessorInputs = prompt  # type: ignore[assignment]
         else:
-            # Otherwise, use user-provided uuids as multimodal hash overrides
-            # if provided.
-            self._validate_mm_uuids(prompt)
-            if isinstance(prompt, dict):
-                mm_uuids = cast(MultiModalUUIDDict | None, prompt.get("multi_modal_uuids"))
-            else:
-                mm_uuids = None
+            if arrival_time is None:
+                arrival_time = time.time()
 
-        # Process inputs, which includes:
-        # 1. Tokenize text prompt, with LoRA request if one exists.
-        # 2. For multimodal models with a merged preprocessor, preprocess
-        #   multimodal data and expand prompt token ids accordingly.
-        with set_request_id(request_id), set_default_torch_num_threads():
-            processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
+            processed_inputs = self.input_preprocessor.preprocess(
                 prompt,
                 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()
+        self._platform_validate_request(processed_inputs, params)
 
         encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
         self._validate_model_inputs(encoder_inputs, decoder_inputs)
@@ -216,7 +206,10 @@ class OmniInputProcessor(InputProcessor):
             if sampling_params.max_tokens is None:
                 seq_len = length_from_prompt_token_ids_or_embeds(prompt_token_ids, prompt_embeds)
                 sampling_params.max_tokens = self.model_config.max_model_len - seq_len
-            sampling_params.update_from_generation_config(self.generation_config_fields, eos_token_id)
+            sampling_params.update_from_generation_config(
+                self.generation_config_fields,
+                self.renderer.get_eos_token_id(),
+            )
             if self.tokenizer is not None:
                 sampling_params.update_from_tokenizer(self.tokenizer)
         else:
@@ -230,6 +223,13 @@ class OmniInputProcessor(InputProcessor):
             decoder_mm_positions = decoder_inputs["mm_placeholders"]
             decoder_mm_hashes = decoder_inputs["mm_hashes"]
 
+            if not all(isinstance(leaf, str) for leaf in json_iter_leaves(decoder_mm_hashes)):
+                raise ValueError(
+                    f"mm_hashes must contain only strings, got: {decoder_mm_hashes}. "
+                    "This is likely due to an incorrect custom implementation of "
+                    "MultiModalProcessor.apply method."
+                )
+
             # Merge and flatten multimodal placeholders, hashes and inputs
             # from dictionaries to lists, and sort them by each item's position
             # in the input sequence.
@@ -281,7 +281,6 @@ class OmniInputProcessor(InputProcessor):
             mm_features=mm_features,
             sampling_params=sampling_params,
             pooling_params=pooling_params,
-            eos_token_id=eos_token_id,
             arrival_time=arrival_time,
             lora_request=lora_request,
             cache_salt=decoder_inputs.get("cache_salt"),

vllm_omni/engine/output_processor.py

@@ -199,7 +199,7 @@ class OmniRequestState(RequestState):
             if not (
                 finished
                 or self.sent_tokens_offset == 0
-                or len(self.detokenizer.output_token_ids) - self.sent_tokens_offset >= self.stream_interval
+                or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval
             ):
                 return None
 
@@ -207,7 +207,7 @@ class OmniRequestState(RequestState):
                 # 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)
+                self.sent_tokens_offset = self.detokenizer.num_output_tokens()
 
         external_req_id = self.external_req_id
         output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts)
@@ -264,8 +264,10 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     def __init__(
         self,
         tokenizer: TokenizerLike | None,
+        *,
         log_stats: bool,
         stream_interval: int = 1,
+        tracing_enabled: bool = False,
         engine_core_output_type: str | None = None,
     ):
         """Initialize the multimodal output processor.
@@ -278,7 +280,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 (e.g., "image", "audio", "latent"). Used to tag multimodal
                 outputs with the correct modality key.
         """
-        super().__init__(tokenizer=tokenizer, log_stats=log_stats, stream_interval=stream_interval)
+        super().__init__(
+            tokenizer=tokenizer,
+            log_stats=log_stats,
+            stream_interval=stream_interval,
+            tracing_enabled=tracing_enabled,
+        )
         self.engine_core_output_type = engine_core_output_type
 
     def add_request(
@@ -320,8 +327,6 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             log_stats=self.log_stats,
             stream_interval=self.stream_interval,
         )
-        if self._requests_drained.is_set():
-            self._requests_drained.clear()
         self.request_states[request_id] = req_state
         if parent_req:
             self.parent_requests[parent_req.request_id] = parent_req

vllm_omni/entrypoints/async_omni.py

@@ -432,7 +432,10 @@ class AsyncOmni(OmniBase):
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)}
         submit_flag = True
+        _loop_iter = 0
+        _last_progress_ts = time.time()
         while not all(all_stages_finished.values()):
+            _loop_iter += 1
             for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
                 if all_stages_finished[stage_id]:
                     continue
@@ -441,6 +444,7 @@ class AsyncOmni(OmniBase):
                 except asyncio.QueueEmpty:
                     await asyncio.sleep(0.001)
                     continue
+                _last_progress_ts = time.time()
                 engine_outputs, finished, output_to_yield = self._process_single_result(
                     result,
                     stage,
@@ -449,11 +453,20 @@ class AsyncOmni(OmniBase):
                 )
                 if submit_flag and stage_id == 0:
                     submit_flag = False
-                    prompt_token_ids = engine_outputs.prompt_token_ids
+                    prompt_token_ids = getattr(engine_outputs, "prompt_token_ids", None)
+                    if prompt_token_ids is None:
+                        prompt_token_ids = []
                     engine_input = copy.deepcopy(prompt)
-                    next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+                    try:
+                        next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+                    except Exception:
+                        raise
                     engine_input["prompt_token_ids"] = [0] * next_prompt_len
                     engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None
+                    for _mm_key in ("mm_kwargs", "mm_hashes", "mm_placeholders", "multi_modal_uuids"):
+                        engine_input.pop(_mm_key, None)
+                    if engine_input.get("type") == "multimodal":
+                        engine_input["type"] = "token"
                     for i in range(1, len(self.stage_list)):
                         task = {
                             "request_id": request_id,
@@ -706,6 +719,16 @@ class AsyncOmni(OmniBase):
             stage.submit(abort_task)
         return None
 
+    async def get_supported_tasks(self) -> set[str]:
+        tasks: set[str] = set()
+        has_comprehension = any(stage.is_comprehension for stage in self.stage_list)
+        if has_comprehension:
+            tasks.add("generate")
+        for stage in self.stage_list:
+            if getattr(stage, "final_output_type", None) == "audio":
+                tasks.add("speech")
+        return tasks if tasks else {"generate"}
+
     async def get_vllm_config(self) -> VllmConfig:
         for stage in self.stage_list:
             if stage.is_comprehension:
@@ -744,6 +767,8 @@ class AsyncOmni(OmniBase):
         OMNI: Required by upstream OpenAIServingModels.__init__ which
         accesses engine_client.renderer.
         """
+        if self.input_processor is None:
+            return None
         return self.input_processor.renderer
 
     async def do_log_stats(self) -> None:

vllm_omni/entrypoints/async_omni_llm.py

@@ -9,7 +9,7 @@ import torch
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.tokenizers import cached_tokenizer_from_config
+from vllm.renderers import renderer_from_config
 from vllm.tracing import init_tracer
 from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
 from vllm.usage.usage_lib import UsageContext
@@ -106,21 +106,17 @@ class AsyncOmniLLM(AsyncLLM):
                 "enabling logging without default stat loggers"
             )
 
-        if self.model_config.skip_tokenizer_init:
-            tokenizer = None
-        else:
-            # Tokenizer (+ ensure liveness if running in another process).
-            tokenizer = cached_tokenizer_from_config(model_config=vllm_config.model_config)
-
         # InputProcessor (converts Inputs --> EngineCoreRequests).
         self.input_processor = OmniInputProcessor(
             vllm_config=vllm_config,
             mm_registry=mm_registry,
         )
 
+        self.renderer = renderer_from_config(self.vllm_config)
+
         # OutputProcessor (converts EngineCoreOutputs --> RequestOutput).
         self.output_processor = MultimodalOutputProcessor(
-            tokenizer=tokenizer,
+            tokenizer=self.renderer.tokenizer,
             log_stats=self.log_stats,
             engine_core_output_type=engine_args.engine_output_type,
         )

vllm_omni/entrypoints/omni_llm.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import Any
 
 import cloudpickle
@@ -8,9 +8,11 @@ from tqdm import tqdm
 # External library imports (vLLM)
 from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
 from vllm.entrypoints.llm import LLM
+from vllm.inputs import ProcessorInputs, PromptType
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput, RequestOutput
 from vllm.plugins.io_processors import get_io_processor
+from vllm.renderers.inputs.preprocess import parse_model_prompt
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.counter import Counter
 from vllm.v1.engine.llm_engine import LLMEngine
@@ -19,7 +21,7 @@ from vllm_omni.distributed.omni_connectors import initialize_orchestrator_connec
 
 # Internal imports (our code)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
-from vllm_omni.engine.input_processor import OmniInputProcessor
+from vllm_omni.engine.input_processor import OmniInputProcessor, reinject_omni_fields
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.entrypoints.utils import (
     filter_dataclass_kwargs,
@@ -174,12 +176,50 @@ class OmniLLM(LLM):
 
         self.supported_tasks = supported_tasks
 
+        # Keep parity with vLLM's LLM initialization fields so inherited
+        # generate/chat preprocessing paths work as expected.
+        self.renderer = self.llm_engine.renderer
         # 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
 
+        # Parity with upstream LLM for pooling/classify entrypoints
+        chat_template = kwargs.get("chat_template", None)
+        from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
+        from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors
+
+        self.chat_template = load_chat_template(chat_template)
+        self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template)
+        self.init_pooling_io_processors = init_pooling_io_processors(
+            supported_tasks=supported_tasks,
+            model_config=self.model_config,
+            renderer=self.renderer,
+            chat_template_config=self.chat_template_config,
+        )
+
+    # ------------------------------------------------------------------
+    # Override upstream _preprocess_cmpl so that omni-specific fields
+    # (additional_information, prompt_embeds, …) survive the renderer's
+    # process_for_engine step which only copies standard vLLM keys.
+    # ------------------------------------------------------------------
+
+    def _preprocess_cmpl(
+        self,
+        prompts: Sequence[PromptType],
+        tokenization_kwargs: dict[str, Any] | None = None,
+    ) -> Sequence[ProcessorInputs]:
+        renderer = self.renderer
+        model_config = self.model_config
+
+        parsed_prompts = [parse_model_prompt(model_config, prompt) for prompt in prompts]
+        tok_params = renderer.default_cmpl_tok_params.with_kwargs(**(tokenization_kwargs or {}))
+        results = renderer.render_cmpl(parsed_prompts, tok_params)
+
+        reinject_omni_fields(results, parsed_prompts)
+        return results
+
     def close(self) -> None:
         """Close resources.
 
@@ -197,7 +237,9 @@ class OmniLLM(LLM):
         except Exception as e:
             logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
 
-    def _run_engine(self, *, use_tqdm: bool | Callable[..., tqdm] = True) -> list[RequestOutput | PoolingRequestOutput]:
+    def _run_engine(
+        self, output_type=None, *, use_tqdm: bool | Callable[..., tqdm] = True
+    ) -> list[RequestOutput | PoolingRequestOutput]:
         # Initialize tqdm.
         if use_tqdm:
             num_requests = self.llm_engine.get_num_unfinished_requests()

vllm_omni/entrypoints/omni_stage.py

@@ -76,20 +76,16 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int =
     """
     from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
 
-    nvml_available = is_process_scoped_memory_available()
-    pid_host = detect_pid_host()
-
-    if nvml_available and pid_host:
-        logger.info(
+    if is_process_scoped_memory_available() and detect_pid_host():
+        logger.debug(
             "NVML process-scoped memory available and PID host is available — concurrent init is safe, skipping locks"
         )
         yield
         return
     else:
-        logger.info(
-            "Using sequential init locks (nvml_available=%s, pid_host=%s)",
-            nvml_available,
-            pid_host,
+        logger.debug(
+            "NVML unavailable or PID host is not available (usually inside a container, "
+            "--pid=host is not set in docker run command) — using sequential init locks"
         )
 
     from vllm_omni.platforms import current_omni_platform
@@ -1356,7 +1352,7 @@ async def _stage_worker_async(
                 usage_context=usage_context,
                 engine_args=omni_engine_args,
                 disable_log_stats=bool(
-                    engine_args.get("disable_log_stats", True) or getattr(omni_engine_args, "disable_log_stats", True)
+                    engine_args.get("disable_log_stats", False) or getattr(omni_engine_args, "disable_log_stats", False)
                 ),
             )
     if hasattr(stage_engine, "log_stats") and stage_engine.log_stats:

vllm_omni/entrypoints/openai/api_server.py

@@ -287,6 +287,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
             ssl_certfile=args.ssl_certfile,
             ssl_ca_certs=args.ssl_ca_certs,
             ssl_cert_reqs=args.ssl_cert_reqs,
+            ssl_ciphers=args.ssl_ciphers,
             h11_max_incomplete_event_size=args.h11_max_incomplete_event_size,
             h11_max_header_count=args.h11_max_header_count,
             **uvicorn_kwargs,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -453,6 +453,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             default_template_content_format,
         ).with_defaults(default_template_kwargs)
 
+        # OMNI: When use_audio_in_video=True, the qwen2_5_omni_thinker mm
+        # processor asserts that audio items are present alongside video items
+        # (inside render_chat_async).  We must inject audio_url items into the
+        # messages BEFORE calling render_chat_async so the mm processor can
+        # count them correctly during tokenisation.
+        mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {}
+        if mm_proc_kw.get("use_audio_in_video", False):
+            messages = await self._inject_audio_from_video_urls(messages)
+
         (conversation,), (engine_prompt,) = await renderer.render_chat_async(
             [messages],
             chat_params,
@@ -462,28 +471,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             },
         )
 
-        # OMNI: When use_audio_in_video=True, the upstream renderer does not
-        # extract audio from video.  We do it here after rendering so that the
-        # audio data is present in multi_modal_data before the engine processes
-        # the request.
-        mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {}
-        if mm_proc_kw.get("use_audio_in_video", False) and isinstance(engine_prompt, dict):
-            mm_data = engine_prompt.get("multi_modal_data")
-            if mm_data is not None and "video" in mm_data and "audio" not in mm_data:
-                from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async
-
-                video_urls: list[str] = []
-                for msg in messages:
-                    for part in msg.get("content") or []:
-                        if isinstance(part, dict) and part.get("type") == "video_url":
-                            url = part.get("video_url", {}).get("url")
-                            if url:
-                                video_urls.append(url)
-
-                if video_urls:
-                    audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls))
-                    engine_prompt.setdefault("multi_modal_data", {})["audio"] = list(audios)
-
         tokenizer = renderer.get_tokenizer()
 
         # tool parsing is done only if a tool_parser has been set and if
@@ -533,6 +520,95 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return conversation, [engine_prompt]
 
+    async def _inject_audio_from_video_urls(
+        self,
+        messages: list[ChatCompletionMessageParam],
+    ) -> list[ChatCompletionMessageParam]:
+        """Pre-extract audio from video URLs and inject as audio_url content items.
+
+        When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal
+        processor requires that the number of audio items equals the number of
+        video items (it subtracts mm_counts["video"] from mm_counts["audio"]).
+        The client only sends video_url items; this method adds the matching
+        audio_url items on the server side before the renderer processes them.
+        """
+        import io
+
+        from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async
+
+        new_messages: list[ChatCompletionMessageParam] = []
+        for msg in messages:
+            content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
+            if not isinstance(content, list):
+                new_messages.append(msg)
+                continue
+
+            video_urls = [
+                part.get("video_url", {}).get("url")
+                for part in content
+                if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url")
+            ]
+
+            if not video_urls:
+                new_messages.append(msg)
+                continue
+
+            audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls))
+
+            audio_items: list[dict] = []
+            for audio_array, sample_rate in audios:
+                buf = io.BytesIO()
+                if soundfile is not None:
+                    soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV")
+                else:
+                    import struct
+
+                    import numpy as np
+
+                    audio_np = np.asarray(audio_array, dtype=np.float32)
+                    sr = int(sample_rate)
+                    num_channels = 1
+                    bits_per_sample = 32
+                    num_frames = len(audio_np)
+                    data_size = num_frames * num_channels * (bits_per_sample // 8)
+                    # Write minimal RIFF/WAV header
+                    buf.write(b"RIFF")
+                    buf.write(struct.pack("<I", 36 + data_size))
+                    buf.write(b"WAVE")
+                    buf.write(b"fmt ")
+                    buf.write(
+                        struct.pack(
+                            "<IHHIIHH",
+                            16,
+                            3,
+                            num_channels,
+                            sr,
+                            sr * num_channels * (bits_per_sample // 8),
+                            num_channels * (bits_per_sample // 8),
+                            bits_per_sample,
+                        )
+                    )
+                    buf.write(b"data")
+                    buf.write(struct.pack("<I", data_size))
+                    buf.write(audio_np.tobytes())
+
+                audio_b64 = base64.b64encode(buf.getvalue()).decode()
+                audio_items.append(
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"},
+                    }
+                )
+
+            new_content = list(content) + audio_items
+            if isinstance(msg, dict):
+                new_msg = {**msg, "content": new_content}
+            else:
+                new_msg = msg.model_copy(update={"content": new_content})
+            new_messages.append(new_msg)
+
+        return new_messages
+
     def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]:
         final_sampling_params_list = []
         for sampling_params in sampling_params_list:
@@ -1213,11 +1289,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                                     latest_delta_len = len(delta_message.tool_calls[0].function.arguments)
 
                                 # get the expected call based on partial JSON
-                                # parsing which "autocompletes" the JSON
-                                expected_call = json.dumps(
-                                    tool_parser.prev_tool_call_arr[index].get("arguments", {}),
-                                    ensure_ascii=False,
-                                )
+                                # parsing which "autocompletes" the JSON.
+                                # Tool parsers (e.g. Qwen3Coder) store
+                                # arguments as a JSON string in
+                                # prev_tool_call_arr. Calling json.dumps()
+                                # on an already-serialized string would
+                                # double-serialize it (e.g. '{"k":1}' becomes
+                                # '"{\\"k\\":1}"'), which then causes the
+                                # replace() below to fail and append the
+                                # entire double-serialized string as a
+                                # spurious final delta.
+                                args = tool_parser.prev_tool_call_arr[index].get("arguments", {})
+                                if isinstance(args, str):
+                                    expected_call = args
+                                else:
+                                    expected_call = json.dumps(args, ensure_ascii=False)
 
                                 # get what we've streamed so far for arguments
                                 # for the current tool

vllm_omni/inputs/preprocess.py

@@ -4,7 +4,7 @@ from typing_extensions import assert_never
 from vllm.inputs.data import EmbedsInputs, SingletonInputs
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
-from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict
+from vllm.multimodal.inputs import MultiModalInputs
 from vllm.renderers.inputs import SingletonDictPrompt
 
 from vllm_omni.inputs.data import (
@@ -30,8 +30,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         parsed_content: OmniTextPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> OmniTokenInputs | MultiModalInputs:
         """Process text prompts with support for mm_processor_kwargs.
 
@@ -49,9 +47,7 @@ class OmniInputPreprocessor(InputPreprocessor):
                 multi_modal_data,
                 mm_processor_kwargs,
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
-            # Preserve prompt_embeds and additional_information
             prompt_embeds = parsed_content.get("prompt_embeds")
             if prompt_embeds is not None:
                 inputs["prompt_embeds"] = prompt_embeds
@@ -59,14 +55,11 @@ class OmniInputPreprocessor(InputPreprocessor):
             if additional_information is not None:
                 inputs["additional_information"] = additional_information
         elif mm_processor_kwargs:
-            # Support mm_processor_kwargs without multi_modal_data
-            # Use case: GLM-Image text-to-image needs processor to generate grid tokens
             inputs = self._process_multimodal(
                 prompt_text,
-                {},  # Empty multi_modal_data
+                {},
                 mm_processor_kwargs,
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
         else:
             prompt_token_ids = self._tokenize_prompt(
@@ -78,7 +71,12 @@ class OmniInputPreprocessor(InputPreprocessor):
                 prompt_embeds=parsed_content.get("prompt_embeds"),
                 additional_information=parsed_content.get("additional_information"),
             )
-
+        prompt_embeds = parsed_content.get("prompt_embeds")
+        if prompt_embeds is not None:
+            inputs["prompt_embeds"] = prompt_embeds
+        additional_information = parsed_content.get("additional_information")
+        if additional_information is not None:
+            inputs["additional_information"] = additional_information
         if cache_salt := parsed_content.get("cache_salt"):
             inputs["cache_salt"] = cache_salt
 
@@ -88,8 +86,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         parsed_content: OmniTokensPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> OmniTokenInputs | MultiModalInputs:
         prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
         prompt_embeds = parsed_content.get("prompt_embeds")
@@ -102,21 +98,23 @@ class OmniInputPreprocessor(InputPreprocessor):
             inputs = self._process_multimodal(
                 prompt_token_ids,
                 multi_modal_data,
-                parsed_content.get("mm_processor_kwargs") or {},
+                parsed_content.get("mm_processor_kwargs"),
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
+                mm_uuids=parsed_content.get("multi_modal_uuids"),
             )
-            if prompt_embeds is not None:
-                inputs["prompt_embeds"] = prompt_embeds
-            if additional_information is not None:
-                inputs["additional_information"] = additional_information
+
         else:
             inputs = token_inputs_omni(
                 prompt_token_ids=prompt_token_ids,
                 prompt_embeds=prompt_embeds,
                 additional_information=additional_information,
             )
-
+        if prompt_embeds is not None:
+            inputs["prompt_embeds"] = prompt_embeds
+        if additional_information is not None:
+            inputs["additional_information"] = additional_information
+        if prompt_text := parsed_content.get("prompt"):
+            inputs["prompt"] = prompt_text
         if cache_salt := parsed_content.get("cache_salt"):
             inputs["cache_salt"] = cache_salt
 
@@ -145,8 +143,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         prompt: SingletonDictPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> SingletonInputs:
         """
         Extract the singleton inputs from a prompt.
@@ -159,20 +155,18 @@ class OmniInputPreprocessor(InputPreprocessor):
 
         * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance
         """
+        if "prompt_embeds" in prompt:
+            return self._process_embeds(prompt)  # type: ignore[arg-type]
+
         if "prompt_token_ids" in prompt:
             return self._process_tokens(
                 prompt,  # type: ignore[arg-type]
-                mm_uuids=mm_uuids,
             )
 
-        if "prompt_embeds" in prompt:
-            return self._process_embeds(prompt)  # type: ignore[arg-type]
-
         if "prompt" in prompt:
             return self._process_text(
                 prompt,  # type: ignore[arg-type]
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
 
         assert_never(prompt)  # type: ignore[arg-type]

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -19,7 +19,6 @@ from vllm.model_executor.models.interfaces import (
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerDummyInputsBuilder,
 )
-from vllm.model_executor.models.qwen3_omni_moe_thinker import Qwen3Omni_VisionTransformer
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     WeightsMapper,
@@ -33,6 +32,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mt
 )
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3MoeLLMForCausalLM,
+    Qwen3Omni_VisionTransformer,
     Qwen3OmniMoeConditionalGenerationMixin,
     Qwen3OmniMoeThinkerMultiModalProcessor,
     Qwen3OmniMoeThinkerProcessingInfo,

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -43,9 +43,12 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (
 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
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.inputs.data import PromptType
 from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
+    MMEncoderAttention,
+)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.models.interfaces import (
@@ -70,7 +73,9 @@ 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,
+    Qwen3Omni_VisionTransformer as _Qwen3Omni_VisionTransformer,
+)
+from vllm.model_executor.models.qwen3_omni_moe_thinker import (
     Qwen3OmniMoeAudioEncoder,
     _get_feat_extract_output_lengths,
 )
@@ -104,6 +109,90 @@ except (ImportError, ModuleNotFoundError):
 
 logger = init_logger(__name__)
 
+
+class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer):
+    """Subclass that fixes Qwen2_5_VisionAttention.forward() compatibility.
+
+    The upstream Qwen3_VisionBlock.forward() does not pass the
+    ``sequence_lengths`` argument required by the updated
+    Qwen2_5_VisionAttention.forward() signature.  This subclass overrides
+    ``forward()`` to compute ``sequence_lengths`` via MMEncoderAttention
+    (following the pattern in qwen3_vl.py) and calls block internals
+    directly so the argument is forwarded correctly.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        grid_thw,
+    ) -> 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_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw)
+
+        if isinstance(grid_thw, torch.Tensor):
+            grid_thw_np = grid_thw.cpu().numpy().astype(np.int32)
+        else:
+            grid_thw_np = np.array(grid_thw, dtype=np.int32)
+
+        cu_seqlens = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(axis=0, dtype=np.int32)
+        cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
+
+        sequence_lengths = MMEncoderAttention.maybe_compute_sequence_lengths(self.attn_backend, cu_seqlens)
+        if sequence_lengths is not None:
+            sequence_lengths = torch.from_numpy(sequence_lengths).to(self.device, non_blocking=True)
+        max_seqlen = torch.tensor(
+            MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
+            self.attn_backend,
+            cu_seqlens,
+            self.hidden_size,
+            self.tp_size,
+        )
+        cu_seqlens = torch.from_numpy(cu_seqlens).to(self.device, non_blocking=True)
+
+        hidden_states = hidden_states.unsqueeze(1)
+
+        hidden_states_list = []
+        deepstack_visual_indexes = self.deepstack_visual_indexes
+
+        for layer_num, blk in enumerate(self.blocks):
+            hidden_states = hidden_states + blk.attn(
+                blk.norm1(hidden_states),
+                cu_seqlens=cu_seqlens,
+                rotary_pos_emb_cos=rotary_pos_emb_cos,
+                rotary_pos_emb_sin=rotary_pos_emb_sin,
+                max_seqlen=max_seqlen,
+                sequence_lengths=sequence_lengths,
+            )
+            hidden_states = hidden_states + blk.mlp(blk.norm2(hidden_states))
+
+            if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes:
+                hidden_states_list.append(hidden_states)
+
+        hidden_states = self.merger(hidden_states)
+
+        if deepstack_visual_indexes is not None:
+            processed_hidden_states_list = [hidden_states]
+            for idx, x_ds in enumerate(hidden_states_list):
+                x_ds = self.merger_list[idx](x_ds)
+                processed_hidden_states_list.append(x_ds)
+            hidden_states = torch.cat(processed_hidden_states_list, dim=1)
+
+        return hidden_states
+
+
 # Speech input languages supported by Qwen3-Omni
 # From: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct
 ISO639_1_SUPPORTED_LANGS = {
@@ -228,6 +317,38 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"audio": None, "image": None, "video": None}
 
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int] | None = None,
+    ) -> Mapping[str, int] | None:
+        mm_counts = mm_counts or {}
+        requested_modalities = {m for m, c in mm_counts.items() if c > 0}
+        mm_max_tokens: dict[str, int] = {}
+
+        if requested_modalities & {"image", "video"}:
+            vl_tokens = Qwen2_5_VLProcessingInfo.get_mm_max_tokens_per_item(
+                self,
+                seq_len=seq_len,
+                mm_counts=mm_counts,
+            )
+            mm_max_tokens.update({m: vl_tokens[m] for m in ["image", "video"] if m in requested_modalities})
+
+        if "audio" in requested_modalities:
+            audio_tokens = Qwen2AudioProcessingInfo.get_mm_max_tokens_per_item(
+                self,
+                seq_len=seq_len,
+                mm_counts=mm_counts,
+            )
+            if audio_tokens is None:
+                feature_extractor = self.get_feature_extractor()
+                max_audio_samples = feature_extractor.chunk_length * feature_extractor.sampling_rate
+                max_audio_tokens = int(_get_feat_extract_output_lengths(torch.tensor([max_audio_samples])).item())
+                audio_tokens = {"audio": max_audio_tokens}
+            mm_max_tokens["audio"] = audio_tokens["audio"]
+
+        return mm_max_tokens
+
 
 Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder
 
@@ -253,7 +374,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
             return x
 
         # NOTE: WhisperFeatureExtractor cannot handle empty list of audios
-        feature_extractor = self.info.get_feature_extractor()
+        feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
         hop_length = feature_extractor.hop_length
         if audios:
             # NOTE: Qwen3-Omni processor accept "audio"
@@ -653,8 +774,7 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
             feature_lens=audio_feature_lengths,
             aftercnn_lens=audio_output_lengths,
         )
-        # OMNI: audio_tower.forward() returns hidden_states tensor directly
-        audio_features = audio_outputs
+        audio_features = audio_outputs if isinstance(audio_outputs, torch.Tensor) else audio_outputs.last_hidden_state
         return audio_features.split(audio_output_lengths.tolist())
 
 

vllm_omni/request.py

@@ -82,7 +82,6 @@ class OmniRequest(Request):
             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,
             cache_salt=request.cache_salt,

vllm_omni/worker/gpu_ar_model_runner.py

@@ -22,7 +22,9 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
 )
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
+from vllm.v1.spec_decode.draft_model import DraftModelProposer
 from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
 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 (
@@ -258,6 +260,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         # Run the model.
         # Use persistent buffers for CUDA graphs.
+        # When spec decode is enabled, delay clearing connector metadata
+        # until after draft model runs in sample_tokens.
+        clear_kv_metadata = self.speculative_config is None
         with (
             set_forward_context(
                 attn_metadata,
@@ -270,7 +275,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
             ),
             record_function_or_nullcontext("gpu_model_runner: forward"),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+            self.maybe_get_kv_connector_output(
+                scheduler_output, clear_metadata=clear_kv_metadata
+            ) as kv_connector_output,
         ):
             model_output = self._model_forward(
                 input_ids=input_ids,
@@ -376,13 +383,12 @@ class GPUARModelRunner(OmniGPUModelRunner):
         self,
         grammar_output: GrammarOutput | None,
     ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
-        kv_connector_output = self.kv_connector_output
-        self.kv_connector_output = None
-
         kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None)
         self.kv_extracted_req_ids = None
 
         if self.execute_model_state is None:
+            kv_connector_output = self.kv_connector_output
+            self.kv_connector_output = None
             # Nothing to do (PP non-final rank case), output isn't used.
             if not kv_connector_output:
                 return None  # type: ignore[return-value]
@@ -454,10 +460,14 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
                 <= self.effective_drafter_max_model_len
             )
-            if spec_config.use_eagle() and not spec_config.disable_padded_drafter_batch:
-                # EAGLE speculative decoding can use the GPU sampled tokens
-                # as inputs, and does not need to wait for bookkeeping to finish.
-                assert isinstance(self.drafter, EagleProposer)
+            use_gpu_toks = (
+                spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states()
+            ) and not spec_config.disable_padded_drafter_batch
+            if use_gpu_toks:
+                assert isinstance(
+                    self.drafter,
+                    EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,
+                )
                 sampled_token_ids = sampler_output.sampled_token_ids
                 if input_fits_in_drafter:
                     propose_draft_token_ids(sampled_token_ids)
@@ -503,9 +513,19 @@ class GPUARModelRunner(OmniGPUModelRunner):
             # tokens on the CPU, so they are run after bookkeeping.
             propose_draft_token_ids(valid_sampled_token_ids)
 
+        # Clear KV connector metadata after draft model runs (if spec decode).
+        # This was deferred from target model forward to allow draft model
+        # to also save its KV cache.
+        if self.speculative_config is not None:
+            self.clear_kv_connector_metadata()
+
         with record_function_or_nullcontext("gpu_model_runner: eplb"):
             self.eplb_step()
 
+        # kv_connector_output may be modified during drafting
+        kv_connector_output = self.kv_connector_output
+        self.kv_connector_output = None
+
         hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
         num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
         if num_scheduled_tokens_np is None:

vllm_omni/worker/gpu_generation_model_runner.py

@@ -23,7 +23,9 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
+from vllm.v1.spec_decode.draft_model import DraftModelProposer
 from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
 from vllm.v1.utils import record_function_or_nullcontext
 from vllm.v1.worker.gpu_model_runner import (
     EMPTY_MODEL_RUNNER_OUTPUT,
@@ -263,6 +265,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
 
         # Run the model.
         # Use persistent buffers for CUDA graphs.
+        # When spec decode is enabled, delay clearing connector metadata
+        # until after draft model runs in sample_tokens.
+        clear_kv_metadata = self.speculative_config is None
         with (
             set_forward_context(
                 attn_metadata,
@@ -275,7 +280,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
             ),
             record_function_or_nullcontext("Forward"),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+            self.maybe_get_kv_connector_output(
+                scheduler_output, clear_metadata=clear_kv_metadata
+            ) as kv_connector_output,
         ):
             outputs = self._run_generation_model(
                 input_ids=input_ids,
@@ -344,6 +351,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         ) = self.execute_model_state
         self.execute_model_state = None
 
+        # Clear KV connector metadata after draft model runs (if spec decode).
+        if self.speculative_config is not None:
+            self.clear_kv_connector_metadata()
+
         pooler_output: list[object] = []
         if isinstance(multimodal_outputs, torch.Tensor):
             assert multimodal_outputs.shape[0] == 1, (
@@ -487,7 +498,8 @@ 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
-            num_active_loras: Number of active LoRAs to capture for.
+            num_active_loras: Number of distinct active LoRAs to capture for.
+                LoRA is activated when num_active_loras > 0.
         """
         mm_config = self.vllm_config.model_config.multimodal_config
         if mm_config and mm_config.mm_encoder_only:
@@ -515,7 +527,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         # Set num_scheduled_tokens based on num_tokens and max_num_seqs
         # for dummy run with LoRA so that the num_reqs collectively
         # has num_tokens in total.
-        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        assert num_tokens <= self.max_num_tokens
         max_num_reqs = self.scheduler_config.max_num_seqs
         if create_mixed_batch:
             assert not uniform_decode
@@ -566,7 +578,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 # activated later in the context manager, but we need to know the
                 # LoRA state when determining the batch descriptor for capture
                 force_has_lora=num_active_loras > 0,
-                # Capture shape specialization for specific active LoRA counts.
+                # `force_num_active_loras` is used for cudagraph capture; because we
+                # need to capture graphs for specific num_active_loras counts
                 force_num_active_loras=num_active_loras,
             )
         )
@@ -604,33 +617,36 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             ubatch_slices=ubatch_slices_padded,
         )
 
-        # If force_attention is True, we always capture attention. Otherwise,
-        # it only happens for cudagraph_runtime_mode=FULL.
-        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-            if create_mixed_batch:
-                # 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()
+        with self.synchronize_input_prep():
+            # If force_attention is True, we always capture attention.
+            # Otherwise, it only happens for cudagraph_runtime_mode=FULL.
+            if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
+                if create_mixed_batch:
+                    # 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()
 
-            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
-            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
-            self.query_start_loc.copy_to_gpu()
+                cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+                self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+                self.query_start_loc.copy_to_gpu()
 
-            pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
-            attn_metadata, _ = self._build_attention_metadata(
-                num_tokens=num_tokens_unpadded,
-                num_reqs=num_reqs_padded,
-                max_query_len=max_query_len,
-                ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
-                for_cudagraph_capture=is_graph_capturing,
-                slot_mappings=slot_mappings_by_group,
-            )
+                pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
+                attn_metadata, _ = self._build_attention_metadata(
+                    num_tokens=num_tokens_unpadded,
+                    num_tokens_padded=num_tokens_padded if pad_attn else None,
+                    num_reqs=num_reqs_padded,
+                    max_query_len=max_query_len,
+                    ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices),
+                    for_cudagraph_capture=is_graph_capturing,
+                    slot_mappings=slot_mappings_by_group,
+                    use_spec_decode=self.speculative_config is not None,
+                )
 
         with self.maybe_dummy_run_with_lora(
             self.lora_config,
@@ -720,8 +736,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             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)
+            if self.speculative_config and (
+                self.speculative_config.use_eagle()
+                or self.speculative_config.uses_draft_model()
+                or self.speculative_config.uses_extract_hidden_states()
+            ):
+                assert isinstance(
+                    self.drafter,
+                    EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,
+                )
+                assert self.speculative_config is not None
                 # Eagle currently only supports PIECEWISE cudagraphs.
                 # Therefore only use cudagraphs if the main model uses PIECEWISE
                 # NOTE(lucas): this is a hack, need to clean up.

vllm_omni/worker/gpu_model_runner.py

@@ -13,7 +13,9 @@ from vllm.sampling_params import SamplingType
 from vllm.tracing import instrument
 from vllm.utils.import_utils import LazyLoader
 from vllm.utils.math_utils import cdiv
+from vllm.v1.spec_decode.draft_model import DraftModelProposer
 from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
 from vllm.v1.worker.gpu_input_batch import CachedRequestState
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
@@ -275,6 +277,11 @@ class OmniGPUModelRunner(GPUModelRunner):
         # Add new requests to the cached states.
         for new_req_data in scheduler_output.scheduled_new_reqs:
             req_id = new_req_data.req_id
+            if req_id in self.requests:
+                req_state = self._update_streaming_request(req_id, new_req_data)
+                reqs_to_add.append(req_state)
+                continue
+
             sampling_params = new_req_data.sampling_params
             pooling_params = new_req_data.pooling_params
 
@@ -316,10 +323,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                     arr = np.frombuffer(payload.data, dtype=dtype)
                     arr = arr.reshape(payload.shape)
                     pe_cpu = torch.from_numpy(arr)
-                    # Store temporarily on CPU; later moved to device in builder
                     setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                    # Also replace payload with Tensor for user visibility in
-                    # scheduler_output
                     try:
                         new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
                     except Exception:
@@ -350,7 +354,6 @@ class OmniGPUModelRunner(GPUModelRunner):
                                     info_dict[k] = entry.list_data
                     if info_dict:
                         self.model_intermediate_buffer[req_id] = info_dict
-                        # Backward compatible: mirror to old setattr location
                         setattr(
                             self.requests[req_id],
                             "additional_information_cpu",
@@ -358,7 +361,6 @@ class OmniGPUModelRunner(GPUModelRunner):
                         )
             except Exception as e:
                 logger.error(f"Error decoding additional information: {e}")
-                pass
 
             if sampling_params and sampling_params.prompt_logprobs is not None:
                 self.num_prompt_logprobs[req_id] = (
@@ -421,18 +423,15 @@ class OmniGPUModelRunner(GPUModelRunner):
             req_state.num_computed_tokens = num_computed_tokens
 
             if not is_last_rank:
-                # When using PP, the scheduler sends the sampled tokens back,
-                # because there's no direct communication between the first-
-                # stage worker and the last-stage worker.
-                new_token_ids = req_data.new_token_ids[i]
-                # Add the sampled token(s) from the previous step (if any).
-                # This doesn't include "unverified" tokens like spec tokens.
-                num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens
-                if num_new_tokens == 1:
-                    # Avoid slicing list in most common case.
-                    req_state.output_token_ids.append(new_token_ids[-1])
-                elif num_new_tokens > 0:
-                    req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
+                if not req_data.new_token_ids:
+                    new_token_ids: list[int] = []
+                else:
+                    new_token_ids = req_data.new_token_ids[i]
+                    num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens
+                    if num_new_tokens == 1:
+                        req_state.output_token_ids.append(new_token_ids[-1])
+                    elif num_new_tokens > 0:
+                        req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
             elif num_output_tokens < len(req_state.output_token_ids):
                 # Some output tokens were discarded due to a sync-KV-load
                 # failure. Align the cached state.
@@ -533,7 +532,6 @@ class OmniGPUModelRunner(GPUModelRunner):
         remove_lora: bool = True,
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
-        activate_lora: bool | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -556,12 +554,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
-            num_active_loras: Number of active LoRAs to capture for.
-            activate_lora: Backward-compatible override for LoRA activation.
+            num_active_loras: Number of distinct active LoRAs to capture for.
+                LoRA is activated when num_active_loras > 0.
         """
-        if activate_lora is None:
-            activate_lora = num_active_loras > 0
-
         mm_config = self.vllm_config.model_config.multimodal_config
         if mm_config and mm_config.mm_encoder_only:
             # The current dummy run only covers LM execution, so we can skip it.
@@ -588,7 +583,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         # Set num_scheduled_tokens based on num_tokens and max_num_seqs
         # for dummy run with LoRA so that the num_reqs collectively
         # has num_tokens in total.
-        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        assert num_tokens <= self.max_num_tokens
         max_num_reqs = self.scheduler_config.max_num_seqs
         if create_mixed_batch:
             assert not uniform_decode
@@ -638,8 +633,9 @@ class OmniGPUModelRunner(GPUModelRunner):
                 # `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,
-                # Capture shape specialization for specific active LoRA counts.
+                force_has_lora=num_active_loras > 0,
+                # `force_num_active_loras` is used for cudagraph capture; because we
+                # need to capture graphs for specific num_active_loras counts
                 force_num_active_loras=num_active_loras,
             )
         )
@@ -676,40 +672,43 @@ class OmniGPUModelRunner(GPUModelRunner):
             ubatch_slices=ubatch_slices_padded,
         )
 
-        # If force_attention is True, we always capture attention. Otherwise,
-        # it only happens for cudagraph_runtime_mode=FULL.
-        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-            if create_mixed_batch:
-                # 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()
+        with self.synchronize_input_prep():
+            # If force_attention is True, we always capture attention.
+            # Otherwise, it only happens for cudagraph_runtime_mode=FULL.
+            if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
+                if create_mixed_batch:
+                    # 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()
 
-            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
-            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
-            self.query_start_loc.copy_to_gpu()
+                cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+                self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+                self.query_start_loc.copy_to_gpu()
 
-            pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
-            attn_metadata, _ = self._build_attention_metadata(
-                num_tokens=num_tokens_unpadded,
-                num_reqs=num_reqs_padded,
-                max_query_len=max_query_len,
-                ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
-                for_cudagraph_capture=is_graph_capturing,
-                slot_mappings=slot_mappings_by_group,
-            )
+                pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
+                attn_metadata, _ = self._build_attention_metadata(
+                    num_tokens=num_tokens_unpadded,
+                    num_tokens_padded=num_tokens_padded if pad_attn else None,
+                    num_reqs=num_reqs_padded,
+                    max_query_len=max_query_len,
+                    ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices),
+                    for_cudagraph_capture=is_graph_capturing,
+                    slot_mappings=slot_mappings_by_group,
+                    use_spec_decode=self.speculative_config is not None,
+                )
 
         with self.maybe_dummy_run_with_lora(
             self.lora_config,
             num_scheduled_tokens,
             num_sampled_tokens,
-            activate_lora,
             remove_lora,
+            num_active_loras,
         ):
             # Make sure padding doesn't exceed max_num_tokens
             assert num_tokens_padded <= self.max_num_tokens
@@ -798,8 +797,16 @@ class OmniGPUModelRunner(GPUModelRunner):
             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)
+            if self.speculative_config and (
+                self.speculative_config.use_eagle()
+                or self.speculative_config.uses_draft_model()
+                or self.speculative_config.uses_extract_hidden_states()
+            ):
+                assert isinstance(
+                    self.drafter,
+                    EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,
+                )
+                assert self.speculative_config is not None
                 # Eagle currently only supports PIECEWISE cudagraphs.
                 # Therefore only use cudagraphs if the main model uses PIECEWISE
                 # NOTE(lucas): this is a hack, need to clean up.
@@ -812,7 +819,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 # 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:
+                if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0:
                     use_cudagraphs = False
 
                 self.drafter.dummy_run(
@@ -847,64 +854,100 @@ class OmniGPUModelRunner(GPUModelRunner):
         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
-        scheduled requests and store them to CPU in the request state.
-        This version avoids hard dependency on payload classes by duck-typing."""
+    # ------------------------------------------------------------------
+    # Payload decoding helpers (torch.Tensor passthrough + legacy
+    # PromptEmbedsPayload / AdditionalInformationPayload support)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _resolve_prompt_embeds_cpu(
+        pe: "torch.Tensor | object | None",
+    ) -> torch.Tensor | None:
+        """Convert *prompt_embeds* to a contiguous CPU tensor.
+
+        Accepts:
+        - ``torch.Tensor`` – moved to CPU as-is (the normal path after
+          upstream added ``prompt_embeds`` to ``EngineCoreRequest``).
+        - Legacy ``PromptEmbedsPayload`` (or any duck-typed object with
+          ``.data``, ``.shape``, ``.dtype``) – decoded via numpy.
+        - ``None`` – returns ``None``.
+        """
+        if pe is None:
+            return None
         try:
-            new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
-            if not new_reqs:
-                return
-            for nr in new_reqs:
-                req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
-                if req_id is None:
-                    continue
-                # prompt_embeds
-                payload_pe = getattr(nr, "prompt_embeds", None)
-                pe_cpu = None
-                if payload_pe is not None:
-                    if isinstance(payload_pe, torch.Tensor):
-                        pe_cpu = payload_pe.detach().to("cpu").contiguous()
-                    else:
-                        # Try duck-typing a payload with data/shape/dtype
-                        data = getattr(payload_pe, "data", None)
-                        shape = getattr(payload_pe, "shape", None)
-                        if data is not None and shape is not None:
-                            dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
-                            arr = np.frombuffer(data, dtype=dt)
-                            arr = arr.reshape(shape)
-                            pe_cpu = torch.from_numpy(arr.copy())
-                if pe_cpu is not None and req_id in self.requests:
-                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                # additional_information
-                payload_info = getattr(nr, "additional_information", None)
-                if payload_info is not None:
-                    logger.warning_once(
-                        "additional_information on request data is deprecated, use model_intermediate_buffer"
-                    )
-                    info_dict = {}
-                    if isinstance(payload_info, dict):
-                        info_dict = payload_info
-                    else:
-                        # Try duck-typing a payload with entries, each entry may have
-                        # tensor_data/tensor_dtype/tensor_shape or list_data
-                        entries = getattr(payload_info, "entries", None)
-                        if isinstance(entries, dict):
-                            for k, entry in entries.items():
-                                tensor_data = getattr(entry, "tensor_data", None)
-                                if tensor_data is not None:
-                                    dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
-                                    arr = np.frombuffer(tensor_data, dtype=dt)
-                                    arr = arr.reshape(getattr(entry, "tensor_shape", ()))
-                                    info_dict[k] = torch.from_numpy(arr.copy())
-                                else:
-                                    info_dict[k] = getattr(entry, "list_data", None)
-                    if info_dict and req_id in self.requests:
-                        self.model_intermediate_buffer[req_id] = info_dict
-                        # Backward compatible: mirror to old setattr location
-                        setattr(self.requests[req_id], "additional_information_cpu", info_dict)
-        except Exception as e:
-            logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
+            if isinstance(pe, torch.Tensor):
+                return pe.detach().cpu().contiguous()
+            data = getattr(pe, "data", None)
+            shape = getattr(pe, "shape", None)
+            if data is not None and shape is not None:
+                dt = np.dtype(getattr(pe, "dtype", "float32"))
+                arr = np.frombuffer(data, dtype=dt).reshape(shape)
+                return torch.from_numpy(arr.copy())
+        except Exception:
+            logger.exception("Failed to decode prompt_embeds payload")
+        return None
+
+    @staticmethod
+    def _resolve_additional_information(
+        payload: "dict | object | None",
+    ) -> dict[str, object]:
+        """Convert an *additional_information* payload to a plain dict.
+
+        Accepts:
+        - ``dict`` – returned as-is.
+        - ``AdditionalInformationPayload`` (or duck-typed with
+          ``.entries``) – decoded entry-by-entry.
+        - ``None`` – returns ``{}``.
+        """
+        if payload is None:
+            return {}
+        if isinstance(payload, dict):
+            return payload
+        try:
+            entries = getattr(payload, "entries", None)
+            if not isinstance(entries, dict):
+                return {}
+            info: dict[str, object] = {}
+            for k, entry in entries.items():
+                tensor_data = getattr(entry, "tensor_data", None)
+                if tensor_data is not None:
+                    dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
+                    arr = np.frombuffer(tensor_data, dtype=dt)
+                    arr = arr.reshape(getattr(entry, "tensor_shape", ()))
+                    info[k] = torch.from_numpy(arr.copy())
+                else:
+                    info[k] = getattr(entry, "list_data", None)
+            return info
+        except Exception:
+            logger.exception("Failed to decode additional_information payload")
+        return {}
+
+    def _decode_and_store_request_payloads(
+        self,
+        scheduler_output: "SchedulerOutput",
+    ) -> None:
+        """Decode per-request prompt_embeds and additional_information for
+        newly scheduled requests and store them on CPU in the request state.
+        """
+        new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
+        if not new_reqs:
+            return
+        for nr in new_reqs:
+            req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
+            if req_id is None or req_id not in self.requests:
+                continue
+            pe_cpu = self._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None))
+            if pe_cpu is not None:
+                setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
+            info_payload = getattr(nr, "additional_information", None)
+            if info_payload is not None:
+                logger.warning_once(
+                    "additional_information on request data is deprecated, use model_intermediate_buffer"
+                )
+            info_dict = self._resolve_additional_information(info_payload)
+            if info_dict:
+                self.model_intermediate_buffer[req_id] = info_dict
+                setattr(self.requests[req_id], "additional_information_cpu", info_dict)
 
     def _gather_runtime_additional_information(self) -> list[dict]:
         """Gather per-request model_intermediate_buffer in batch order."""