rebase 2026/2/5

2026-02-05 10:52:28 +00:00
1 parent 0144a424ed
commit 24b04f58a0
@@ -12,6 +12,11 @@ RUN apt-get update && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

+# Install vLLM wheel for the specific commit
+# URL format: https://wheels.vllm.ai/<commit_hash>/
+# Current vLLM commit: a7be77beef5f59d9d349818b4f2860483551b255
+RUN uv pip install --system --upgrade --force-reinstall vllm --extra-index-url https://wheels.vllm.ai/a7be77beef5f59d9d349818b4f2860483551b255
+
 # Install vllm-omni into the same uv-managed Python environment used by the base image.
 RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"

@@ -226,37 +226,6 @@ def get_multi_audios_query() -> QueryResult:
    )


-# def get_use_audio_in_video_query(video_path: str | None = None) -> QueryResult:
-# question = (
-#     "Describe the content of the video in details, then convert what the "
-#     "baby say into text."
-# )
-# prompt = (
-#     f"<|im_start|>system\n{default_system}<|im_end|>\n"
-#     "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
-#     f"{question}<|im_end|>\n"
-#     f"<|im_start|>assistant\n"
-# )
-# if video_path:
-#     if not os.path.exists(video_path):
-#         raise FileNotFoundError(f"Video file not found: {video_path}")
-#     video_frames = video_to_ndarrays(video_path, num_frames=16)
-# else:
-#     video_frames = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays
-# audio = extract_video_audio(video_path, sampling_rate=16000)
-# return QueryResult(
-#     inputs={
-#         "prompt": prompt,
-#         "multi_modal_data": {
-#             "video": video_frames,
-#             "audio": audio,
-#         },
-#         "mm_processor_kwargs": {
-#             "use_audio_in_video": True,
-#         },
-#     },
-#     limit_mm_per_prompt={"audio": 1, "video": 1},
-# )
 def get_use_audio_in_video_query() -> QueryResult:
    question = "Describe the content of the video in details, then convert what the baby say into text."
    prompt = (
@@ -1,28 +1,12 @@
-import warnings
 from dataclasses import field
 from typing import Any

-import torch
 from pydantic import ConfigDict
-from pydantic.dataclasses import dataclass
-from vllm.config import ModelConfig, config
-from vllm.config.model import (
-    _RUNNER_CONVERTS,
-    _get_and_verify_dtype,
-    get_served_model_name,
-)
-from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MultiModalConfig
-from vllm.config.pooler import PoolerConfig
+from vllm.config import ModelConfig
+from vllm.config.multimodal import MMCacheType, MMEncoderTPMode
+from vllm.config.utils import config
 from vllm.logger import init_logger
-from vllm.platforms import current_platform
-from vllm.transformers_utils.config import (
-    get_config,
-    get_hf_image_processor_config,
-    get_hf_text_config,
-    get_pooling_config,
-)
-from vllm.transformers_utils.gguf_utils import is_gguf, maybe_patch_hf_config_from_gguf
-from vllm.transformers_utils.utils import maybe_model_redirect
+from vllm.transformers_utils.config import get_hf_text_config
 from vllm.v1.attention.backends.registry import AttentionBackendEnum

 import vllm_omni.model_executor.models as me_models
@@ -30,8 +14,7 @@ import vllm_omni.model_executor.models as me_models
 logger = init_logger(__name__)


-@config
-@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
+@config(config=ConfigDict(arbitrary_types_allowed=True))
 class OmniModelConfig(ModelConfig):
    """Configuration for Omni models, extending the base ModelConfig.

@@ -117,189 +100,32 @@ class OmniModelConfig(ModelConfig):
        skip_mm_profiling: bool | None,
        video_pruning_rate: float | None,
    ) -> None:
-        # Keep set served_model_name before maybe_model_redirect(self.model)
-        self.served_model_name = get_served_model_name(self.model, self.served_model_name)
-        self.model = maybe_model_redirect(self.model)
-        # The tokenizer is consistent with the model by default.
-        if self.tokenizer is None:
-            self.tokenizer = self.model
-        if self.tokenizer_revision is None:
-            self.tokenizer_revision = self.revision
-        self.tokenizer = maybe_model_redirect(self.tokenizer)
-
-        if isinstance(self.hf_config_path, str):
-            self.hf_config_path = maybe_model_redirect(self.hf_config_path)
-
-        if callable(self.hf_overrides):
-            hf_overrides_kw = {}
-            hf_overrides_fn = self.hf_overrides
-            dict_overrides: dict[str, Any] = {}
-        else:
-            # Separate dict overrides from flat ones
-            # We'll determine how to apply dict overrides after loading the config
-            hf_overrides_kw = {}
-            dict_overrides = {}
-            for key, value in self.hf_overrides.items():
-                if isinstance(value, dict):
-                    dict_overrides[key] = value
-                else:
-                    hf_overrides_kw[key] = value
-            hf_overrides_fn = None
-
-        self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
-
-        if self.override_attention_dtype is not None and not current_platform.is_rocm():
-            warnings.warn(
-                "override-attention-dtype is set but not using ROCm platform",
-                stacklevel=2,
-            )
-
-        if self.enable_sleep_mode and not current_platform.is_sleep_mode_available():
-            raise ValueError("Sleep mode is not supported on current platform.")
-
-        hf_config = get_config(
-            self.hf_config_path or self.model,
-            self.trust_remote_code,
-            self.revision,
-            self.code_revision,
-            self.config_format,
-            hf_overrides_kw=hf_overrides_kw,
-            hf_overrides_fn=hf_overrides_fn,
-        )
-        hf_config = maybe_patch_hf_config_from_gguf(
-            self.model,
-            hf_config,
+        # Call parent's __post_init__ to handle all standard ModelConfig initialization
+        super().__post_init__(
+            limit_mm_per_prompt=limit_mm_per_prompt,
+            enable_mm_embeds=enable_mm_embeds,
+            media_io_kwargs=media_io_kwargs,
+            mm_processor_kwargs=mm_processor_kwargs,
+            mm_processor_cache_gb=mm_processor_cache_gb,
+            mm_processor_cache_type=mm_processor_cache_type,
+            mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
+            mm_encoder_only=mm_encoder_only,
+            mm_encoder_tp_mode=mm_encoder_tp_mode,
+            mm_encoder_attn_backend=mm_encoder_attn_backend,
+            interleave_mm_strings=interleave_mm_strings,
+            skip_mm_profiling=skip_mm_profiling,
+            video_pruning_rate=video_pruning_rate,
        )

-        self.hf_config = hf_config
-        if dict_overrides:
-            self._apply_dict_overrides(hf_config, dict_overrides)
-        self.hf_text_config = self.draw_hf_text_config()
-        self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None)
-        self.encoder_config = self._get_encoder_config()
-        self.hf_image_processor_config = get_hf_image_processor_config(
-            self.model, hf_token=self.hf_token, revision=self.revision
-        )
-        self.model_arch_config = self.get_model_arch_config()
-
-        if self.convert == "mm_encoder_only":
-            logger.warning_once(
-                "`--convert mm_encoder_only` is deprecated and "
-                "will be removed in v0.15. "
-                "Please use --mm-encoder-only` instead."
-            )
-            mm_encoder_only = True
-            self.convert = "none"
-
-        architectures = self.architectures
-        registry = self.registry
-        is_generative_model = registry.is_text_generation_model(architectures, self)
-        is_pooling_model = registry.is_pooling_model(architectures, self)
-
-        self.runner_type = self._get_runner_type(architectures, self.runner)
-        self.convert_type = self._get_convert_type(architectures, self.runner_type, self.convert)
-
-        if self.runner_type == "generate" and not is_generative_model:
-            generate_converts = _RUNNER_CONVERTS["generate"]
-            if self.convert_type not in generate_converts:
-                # Currently we don't have any converters for generative models
-                raise ValueError("This model does not support `--runner generate`.")
-        if self.runner_type == "pooling" and not is_pooling_model:
-            pooling_converts = _RUNNER_CONVERTS["pooling"]
-            if self.convert_type not in pooling_converts:
-                convert_option = "<" + "|".join(pooling_converts) + ">"
-                raise ValueError(
-                    "This model does not support `--runner pooling`. "
-                    f"You can pass `--convert {convert_option} to adapt "
-                    "it into a pooling model."
-                )
-
-        # Note: Initialize these attributes early because transformers fallback
-        # may fail to load dynamic modules in child processes
-        model_info, arch = registry.inspect_model_cls(architectures, self)
-        self._model_info = model_info
-        self._architecture = arch
-        logger.info("Resolved architecture: %s", arch)
-
-        # Init pooler config if needed
-        if self.runner_type == "pooling":
-            if self.pooler_config is None:
-                self.pooler_config = PoolerConfig()
-
-            base_config = get_pooling_config(self.model, self.revision)
-            if base_config is not None:
-                # Only set values that are not overridden by the user
-                for k, v in base_config.items():
-                    if getattr(self.pooler_config, k) is None:
-                        setattr(self.pooler_config, k, v)
-
-            default_seq_pooling_type = self._model_info.default_seq_pooling_type
-            if self.pooler_config.seq_pooling_type is None:
-                self.pooler_config.seq_pooling_type = default_seq_pooling_type
-            default_tok_pooling_type = self._model_info.default_tok_pooling_type
-            if self.pooler_config.tok_pooling_type is None:
-                self.pooler_config.tok_pooling_type = default_tok_pooling_type
-
-        self.dtype: torch.dtype = _get_and_verify_dtype(
-            self.model,
-            self.hf_config,
-            self.dtype,
-            is_pooling_model=self.runner_type == "pooling",
-            revision=self.revision,
-        )
-
-        self.original_max_model_len = self.max_model_len
-        self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
-
-        if self.is_encoder_decoder:
-            self.mm_processor_cache_gb = 0
-            logger.info("Encoder-decoder model detected, disabling mm processor cache.")
-
-        # Init multimodal config if needed
-        if self._model_info.supports_multimodal:
-            if mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data:
-                logger.warning_once(
-                    "This model does not support `--mm-encoder-tp-mode data`. "
-                    "Falling back to `--mm-encoder-tp-mode weights`."
-                )
-                mm_encoder_tp_mode = "weights"
-
-            mm_config_kwargs = dict(
-                limit_per_prompt=limit_mm_per_prompt,
-                enable_mm_embeds=enable_mm_embeds,
-                media_io_kwargs=media_io_kwargs,
-                mm_processor_kwargs=mm_processor_kwargs,
-                mm_processor_cache_gb=mm_processor_cache_gb,
-                mm_processor_cache_type=mm_processor_cache_type,
-                mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
-                mm_encoder_only=mm_encoder_only,
-                mm_encoder_tp_mode=mm_encoder_tp_mode,
-                mm_encoder_attn_backend=mm_encoder_attn_backend,
-                interleave_mm_strings=interleave_mm_strings,
-                skip_mm_profiling=skip_mm_profiling,
-                video_pruning_rate=video_pruning_rate,
-            )
-
-            mm_config_kwargs = {k: v for k, v in mm_config_kwargs.items() if v is not None}
-
-            self.multimodal_config = MultiModalConfig(**mm_config_kwargs)
-
-        # Multimodal GGUF models must use original repo for mm processing
-        if is_gguf(self.tokenizer) and self.is_multimodal_model:
-            raise ValueError(
-                "Loading a multimodal GGUF model needs to use original "
-                "tokenizer. Please specify the unquantized hf model's "
-                "repo name or path using the --tokenizer argument."
-            )
-
-        if self.disable_sliding_window:
-            # Set after get_and_verify_max_len to ensure that max_model_len
-            # can be correctly capped to sliding window size
-            self.hf_text_config.sliding_window = None
-
-        # Avoid running try_verify_and_update_config multiple times
-        self.config_updated = False
-        self._try_verify_and_update_model_config()
-        self._verify_quantization()
-        self._verify_cuda_graph()
-        self._verify_bnb_config()
+        # Override hf_text_config with omni-specific logic for multi-stage models
+        # (e.g., thinker_config, talker_config)
+        new_hf_text_config = self.draw_hf_text_config()
+        if new_hf_text_config is not self.hf_text_config:
+            self.hf_text_config = new_hf_text_config
+            # Recalculate dependent attributes
+            self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None)
+            # Recalculate max_model_len since it depends on hf_text_config
+            self.max_model_len = self.get_and_verify_max_len(self.original_max_model_len)
+            # Reset sliding_window if needed
+            if self.disable_sliding_window:
+                self.hf_text_config.sliding_window = None
@@ -351,6 +351,7 @@ class OmniARScheduler(VLLMScheduler):
                        kv_transfer_params=kv_transfer_params,
                        trace_headers=request.trace_headers,
                        num_cached_tokens=request.num_cached_tokens,
+                        num_external_computed_tokens=request.num_external_computed_tokens,
                        routed_experts=routed_experts,
                        num_nans_in_logits=request.num_nans_in_logits,
                    )
@@ -418,7 +419,7 @@ class OmniARScheduler(VLLMScheduler):

        # publish collected KV cache events
        if events:
-            batch = KVEventBatch(ts=time.time(), events=events)
+            batch = KVEventBatch(ts=time(), events=events)
            self.kv_event_publisher.publish(batch)

        # Create EngineCoreOutputs for all clients that have requests with
@@ -327,7 +327,7 @@ class OmniGenerationScheduler(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
@@ -338,10 +338,16 @@ class OmniGenerationScheduler(VLLMScheduler):
                # tokens.
                if request.num_computed_tokens > 0:
                    request.num_computed_tokens -= num_rejected
+                # If async scheduling, num_output_placeholders also includes
+                # the scheduled spec tokens count and so is similarly adjusted.
+                if request.num_output_placeholders > 0:
+                    request.num_output_placeholders -= num_rejected
                spec_decoding_stats = self.make_spec_decoding_stats(
                    spec_decoding_stats,
                    num_draft_tokens=num_draft_tokens,
                    num_accepted_tokens=num_accepted,
+                    num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
+                    request_id=req_id,
                )

            stopped = False
@@ -407,6 +413,7 @@ class OmniGenerationScheduler(VLLMScheduler):
                        kv_transfer_params=kv_transfer_params,
                        trace_headers=request.trace_headers,
                        num_cached_tokens=request.num_cached_tokens,
+                        num_external_computed_tokens=request.num_external_computed_tokens,
                        routed_experts=routed_experts,
                        num_nans_in_logits=request.num_nans_in_logits,
                    )
@@ -422,6 +429,22 @@ class OmniGenerationScheduler(VLLMScheduler):
            # This is a rare case and unlikely to impact performance.
            self.waiting.remove_requests(stopped_preempted_reqs)

+        # Handle failed KV load requests
+        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
+            requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
+            self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
+            for request in requests:
+                outputs[request.client_index].append(
+                    EngineCoreOutput(
+                        request_id=request.request_id,
+                        new_token_ids=[],
+                        finish_reason=request.get_finished_reason(),
+                        events=request.take_events(),
+                        trace_headers=request.trace_headers,
+                        num_cached_tokens=request.num_cached_tokens,
+                    )
+                )
+
        # KV Connector: update state for finished KV Transfers.
        if kv_connector_output:
            self._update_from_kv_xfer_finished(kv_connector_output)
@@ -15,12 +15,15 @@ class OmniNewRequestData(NewRequestData):

    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
-    prompt_embeds: PromptEmbedsPayload | None = None
+    # 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
@@ -38,13 +41,14 @@ class OmniNewRequestData(NewRequestData):
        Args:
            request: Request object to convert
            block_ids: Tuple of block ID lists for KV cache allocation
+            prefill_token_ids: Optional prefill token IDs for v2 model runner

        Returns:
            OmniNewRequestData instance with data from the request
        """
        return cls(
            req_id=request.request_id,
-            external_req_id=request.external_req_id,
+            external_req_id=getattr(request, "external_req_id", None),
            prompt_token_ids=request.prompt_token_ids,
            mm_features=request.mm_features,
            sampling_params=request.sampling_params,
@@ -52,9 +56,9 @@ class OmniNewRequestData(NewRequestData):
            block_ids=block_ids,
            num_computed_tokens=request.num_computed_tokens,
            lora_request=request.lora_request,
-            prompt_embeds=request.prompt_embeds,
+            prompt_embeds=getattr(request, "prompt_embeds", None),
            prefill_token_ids=prefill_token_ids,
-            additional_information=request.additional_information,
+            additional_information=getattr(request, "additional_information", None),
        )


@@ -1,10 +1,8 @@
 from dataclasses import dataclass, field
 from typing import Any

-from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig
 from vllm.engine.arg_utils import EngineArgs
 from vllm.logger import init_logger
-from vllm.transformers_utils.config import get_hf_text_config
 from vllm.v1.engine.async_llm import AsyncEngineArgs

 from vllm_omni.config import OmniModelConfig
@@ -75,24 +73,6 @@ class OmniEngineArgs(EngineArgs):
    async_chunk: bool = False
    omni_kv_config: dict | None = None

-    def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
-        # transformers' get_text_config method is used to get the text config from thinker_config.
-        # to handle the case that each model stage has their own text config,
-        # we need to draw the text config from the corresponding model stage.
-        hf_config = config_dict["hf_config"]
-        hf_config_name = config_dict["hf_config_name"]
-        try:
-            # Try to get the stage-specific config (e.g., thinker_config, talker_config)
-            stage_config = getattr(hf_config, hf_config_name)
-            return stage_config.get_text_config()
-        except AttributeError:
-            # Fallback: if the attribute doesn't exist, use the default get_hf_text_config
-            logger.warning(
-                f"Config attribute '{hf_config_name}' not found in hf_config, "
-                "falling back to default get_hf_text_config"
-            )
-        return get_hf_text_config(hf_config)
-
    def __post_init__(self) -> None:
        load_omni_general_plugins()
        super().__post_init__()
@@ -112,39 +92,80 @@ class OmniEngineArgs(EngineArgs):
        # register omni models to avoid model not found error
        self._ensure_omni_models_registered()

-        # First, get the base ModelConfig from the parent class
-        base_config = super().create_model_config()
-
-        # Create OmniModelConfig by copying all base config attributes
-        # and adding the new omni-specific fields
-        config_dict = base_config.__dict__.copy()
-        # FIXME(Isotr0py): This is a temporary workaround for multimodal_config
-        config_dict = {
-            **(getattr(mm := config_dict.pop("multimodal_config", None), "__dict__", mm or {})),
-            **config_dict,
-        }
-
-        # Add the new omni-specific fields
-        config_dict["stage_id"] = self.stage_id
-        config_dict["async_chunk"] = self.async_chunk
-        config_dict["model_stage"] = self.model_stage
-        config_dict["model_arch"] = self.model_arch
-        config_dict["engine_output_type"] = self.engine_output_type
        # Build stage_connector_config from stage_connector_spec
        stage_connector_config = {
            "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
            "extra": self.stage_connector_spec.get("extra", {}).copy(),
        }
        stage_connector_config["extra"]["stage_id"] = self.stage_id
-        config_dict["stage_connector_config"] = stage_connector_config

-        config_dict["hf_config_name"] = self.hf_config_name
-        config_dict["custom_process_next_stage_input_func"] = self.custom_process_next_stage_input_func
-        config_dict["omni_kv_config"] = self.omni_kv_config
-        if self.hf_config_name is not None:
-            config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
-        # Create and return the OmniModelConfig instance
-        omni_config = OmniModelConfig(**config_dict)
+        # Create OmniModelConfig directly from engine args
+        # Note: We pass the actual init parameters matching vLLM's EngineArgs.create_model_config()
+        omni_config = OmniModelConfig(
+            # Base ModelConfig fields (matching vLLM's EngineArgs.create_model_config)
+            model=self.model,
+            model_weights=self.model_weights,
+            hf_config_path=self.hf_config_path,
+            runner=self.runner,
+            convert=self.convert,
+            tokenizer=self.tokenizer,
+            tokenizer_mode=self.tokenizer_mode,
+            trust_remote_code=self.trust_remote_code,
+            allowed_local_media_path=self.allowed_local_media_path,
+            allowed_media_domains=self.allowed_media_domains,
+            dtype=self.dtype,
+            seed=self.seed,
+            revision=self.revision,
+            code_revision=self.code_revision,
+            hf_token=self.hf_token,
+            hf_overrides=self.hf_overrides,
+            tokenizer_revision=self.tokenizer_revision,
+            max_model_len=self.max_model_len,
+            quantization=self.quantization,
+            allow_deprecated_quantization=self.allow_deprecated_quantization,
+            enforce_eager=self.enforce_eager,
+            enable_return_routed_experts=self.enable_return_routed_experts,
+            max_logprobs=self.max_logprobs,
+            logprobs_mode=self.logprobs_mode,
+            disable_sliding_window=self.disable_sliding_window,
+            disable_cascade_attn=self.disable_cascade_attn,
+            skip_tokenizer_init=self.skip_tokenizer_init,
+            enable_prompt_embeds=self.enable_prompt_embeds,
+            served_model_name=self.served_model_name,
+            limit_mm_per_prompt=self.limit_mm_per_prompt,
+            enable_mm_embeds=self.enable_mm_embeds,
+            interleave_mm_strings=self.interleave_mm_strings,
+            media_io_kwargs=self.media_io_kwargs,
+            skip_mm_profiling=self.skip_mm_profiling,
+            config_format=self.config_format,
+            mm_processor_kwargs=self.mm_processor_kwargs,
+            mm_processor_cache_gb=self.mm_processor_cache_gb,
+            mm_processor_cache_type=self.mm_processor_cache_type,
+            mm_shm_cache_max_object_size_mb=self.mm_shm_cache_max_object_size_mb,
+            mm_encoder_only=self.mm_encoder_only,
+            mm_encoder_tp_mode=self.mm_encoder_tp_mode,
+            mm_encoder_attn_backend=self.mm_encoder_attn_backend,
+            pooler_config=self.pooler_config,
+            logits_processor_pattern=self.logits_processor_pattern,
+            generation_config=self.generation_config,
+            override_generation_config=self.override_generation_config,
+            enable_sleep_mode=self.enable_sleep_mode,
+            model_impl=self.model_impl,
+            override_attention_dtype=self.override_attention_dtype,
+            logits_processors=self.logits_processors,
+            video_pruning_rate=self.video_pruning_rate,
+            io_processor_plugin=self.io_processor_plugin,
+            # Omni-specific fields
+            stage_id=self.stage_id,
+            async_chunk=self.async_chunk,
+            model_stage=self.model_stage,
+            model_arch=self.model_arch,
+            engine_output_type=self.engine_output_type,
+            hf_config_name=self.hf_config_name,
+            custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
+            stage_connector_config=stage_connector_config,
+            omni_kv_config=self.omni_kv_config,
+        )
        omni_config.hf_config.architectures = omni_config.architectures

        return omni_config
@@ -177,24 +198,6 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
    async_chunk: bool = False
    omni_kv_config: dict | None = None

-    def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
-        # transformers' get_text_config method is used to get the text config from thinker_config.
-        # to handle the case that each model stage has their own text config,
-        # we need to draw the text config from the corresponding model stage.
-        hf_config = config_dict["hf_config"]
-        hf_config_name = config_dict["hf_config_name"]
-        try:
-            # Try to get the stage-specific config (e.g., thinker_config, talker_config)
-            stage_config = getattr(hf_config, hf_config_name)
-            return stage_config.get_text_config()
-        except AttributeError:
-            # Fallback: if the attribute doesn't exist, use the default get_hf_text_config
-            logger.warning(
-                f"Config attribute '{hf_config_name}' not found in hf_config, "
-                "falling back to default get_hf_text_config"
-            )
-        return get_hf_text_config(hf_config)
-
    def __post_init__(self) -> None:
        load_omni_general_plugins()
        super().__post_init__()
@@ -207,35 +210,87 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
        return True

    def create_model_config(self) -> OmniModelConfig:
+        """Create an OmniModelConfig from these engine arguments.
+        Returns:
+            OmniModelConfig instance with all configuration fields set
+        """
        # register omni models to avoid model not found error
        self._ensure_omni_models_registered()
-        # First, get the base ModelConfig from the parent class
-        base_config = super().create_model_config()

-        # Create OmniModelConfig by copying all base config attributes
-        # and adding the new omni-specific fields
-        config_dict = base_config.__dict__.copy()
-
-        # Add the new omni-specific fields
-        config_dict["stage_id"] = self.stage_id
-        config_dict["async_chunk"] = self.async_chunk
-        config_dict["model_stage"] = self.model_stage
-        config_dict["model_arch"] = self.model_arch
-        config_dict["engine_output_type"] = self.engine_output_type
+        # Build stage_connector_config from stage_connector_spec
        stage_connector_config = {
            "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
            "extra": self.stage_connector_spec.get("extra", {}).copy(),
        }
        stage_connector_config["extra"]["stage_id"] = self.stage_id
-        config_dict["stage_connector_config"] = stage_connector_config

-        config_dict["hf_config_name"] = self.hf_config_name
-        config_dict["custom_process_next_stage_input_func"] = self.custom_process_next_stage_input_func
-        config_dict["omni_kv_config"] = self.omni_kv_config
-        if self.hf_config_name is not None:
-            config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
-        # Create and return the OmniModelConfig instance
-        omni_config = OmniModelConfig(**config_dict)
+        # Create OmniModelConfig directly from engine args
+        # Note: We pass the actual init parameters matching vLLM's EngineArgs.create_model_config()
+        omni_config = OmniModelConfig(
+            # Base ModelConfig fields (matching vLLM's EngineArgs.create_model_config)
+            model=self.model,
+            model_weights=self.model_weights,
+            hf_config_path=self.hf_config_path,
+            runner=self.runner,
+            convert=self.convert,
+            tokenizer=self.tokenizer,
+            tokenizer_mode=self.tokenizer_mode,
+            trust_remote_code=self.trust_remote_code,
+            allowed_local_media_path=self.allowed_local_media_path,
+            allowed_media_domains=self.allowed_media_domains,
+            dtype=self.dtype,
+            seed=self.seed,
+            revision=self.revision,
+            code_revision=self.code_revision,
+            hf_token=self.hf_token,
+            hf_overrides=self.hf_overrides,
+            tokenizer_revision=self.tokenizer_revision,
+            max_model_len=self.max_model_len,
+            quantization=self.quantization,
+            allow_deprecated_quantization=self.allow_deprecated_quantization,
+            enforce_eager=self.enforce_eager,
+            enable_return_routed_experts=self.enable_return_routed_experts,
+            max_logprobs=self.max_logprobs,
+            logprobs_mode=self.logprobs_mode,
+            disable_sliding_window=self.disable_sliding_window,
+            disable_cascade_attn=self.disable_cascade_attn,
+            skip_tokenizer_init=self.skip_tokenizer_init,
+            enable_prompt_embeds=self.enable_prompt_embeds,
+            served_model_name=self.served_model_name,
+            limit_mm_per_prompt=self.limit_mm_per_prompt,
+            enable_mm_embeds=self.enable_mm_embeds,
+            interleave_mm_strings=self.interleave_mm_strings,
+            media_io_kwargs=self.media_io_kwargs,
+            skip_mm_profiling=self.skip_mm_profiling,
+            config_format=self.config_format,
+            mm_processor_kwargs=self.mm_processor_kwargs,
+            mm_processor_cache_gb=self.mm_processor_cache_gb,
+            mm_processor_cache_type=self.mm_processor_cache_type,
+            mm_shm_cache_max_object_size_mb=self.mm_shm_cache_max_object_size_mb,
+            mm_encoder_only=self.mm_encoder_only,
+            mm_encoder_tp_mode=self.mm_encoder_tp_mode,
+            mm_encoder_attn_backend=self.mm_encoder_attn_backend,
+            pooler_config=self.pooler_config,
+            logits_processor_pattern=self.logits_processor_pattern,
+            generation_config=self.generation_config,
+            override_generation_config=self.override_generation_config,
+            enable_sleep_mode=self.enable_sleep_mode,
+            model_impl=self.model_impl,
+            override_attention_dtype=self.override_attention_dtype,
+            logits_processors=self.logits_processors,
+            video_pruning_rate=self.video_pruning_rate,
+            io_processor_plugin=self.io_processor_plugin,
+            # Omni-specific fields
+            stage_id=self.stage_id,
+            async_chunk=self.async_chunk,
+            model_stage=self.model_stage,
+            model_arch=self.model_arch,
+            engine_output_type=self.engine_output_type,
+            hf_config_name=self.hf_config_name,
+            custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
+            stage_connector_config=stage_connector_config,
+            omni_kv_config=self.omni_kv_config,
+        )
        omni_config.hf_config.architectures = omni_config.architectures

        return omni_config
@@ -1,4 +1,3 @@
-import os
 import time
 from collections.abc import Mapping
 from typing import Any, cast
@@ -169,14 +168,7 @@ class OmniInputProcessor(InputProcessor):
        # 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.
-        num_threads = int(os.environ.get("OMP_NUM_THREADS", "1"))
-        if "OMP_NUM_THREADS" not in os.environ:
-            logger.debug_once(
-                "OMP_NUM_THREADS is not set; defaulting Torch threads to %d for input preprocessing.",
-                num_threads,
-            )
-
-        with set_request_id(request_id), set_default_torch_num_threads(num_threads):
+        with set_request_id(request_id), set_default_torch_num_threads():
            processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
                prompt,
                tokenization_kwargs=tokenization_kwargs,
@@ -1,4 +1,3 @@
-from ast import Dict
 from collections.abc import Callable
 from typing import Any

@@ -37,8 +36,9 @@ class OmniRequestState(RequestState):
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
+        # Omni-specific: multimodal output accumulation
        self.mm_type: str | None = None
-        self.mm_accumulated: Dict[str, Any] | None = None
+        self.mm_accumulated: dict[str, Any] | None = None

    def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
        if payload is None:
@@ -57,7 +57,7 @@ class OmniRequestState(RequestState):
                return x

            if isinstance(payload, dict):
-                incoming: Dict[str, Any] = {}
+                incoming: dict[str, Any] = {}
                target_key = self.mm_type or "hidden"

                # Iterate directly without unnecessary dict copy
@@ -257,8 +257,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):

    def __init__(
        self,
-        tokenizer: TokenizerLike,
+        tokenizer: TokenizerLike | None,
        log_stats: bool,
+        stream_interval: int = 1,
        engine_core_output_type: str | None = None,
    ):
        """Initialize the multimodal output processor.
@@ -266,11 +267,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
        Args:
            tokenizer: Tokenizer for detokenizing text outputs
            log_stats: Whether to log statistics
+            stream_interval: Stream interval for output generation
            engine_core_output_type: Optional output type specification
                (e.g., "image", "audio", "latents"). Used to route outputs
                to appropriate processors. If None, output type is inferred.
        """
-        super().__init__(tokenizer=tokenizer, log_stats=log_stats)
+        super().__init__(tokenizer=tokenizer, log_stats=log_stats, stream_interval=stream_interval)
        self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
        self._reqid_to_mm_type: dict[str, str] = {}
        self.engine_core_output_type = engine_core_output_type
@@ -6,7 +6,6 @@ import socket
 from typing import TYPE_CHECKING

 import torch
-import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
@@ -164,21 +163,26 @@ class AsyncOmniLLM(AsyncLLM):
        except RuntimeError:
            pass

-        if envs.VLLM_TORCH_PROFILER_DIR and not envs.VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM:
+        # Use profiler_config from vllm_config (new way, aligned with vllm v1)
+        if (
+            vllm_config.profiler_config.profiler == "torch"
+            and not vllm_config.profiler_config.ignore_frontend
+        ):
+            profiler_dir = vllm_config.profiler_config.torch_profiler_dir
            logger.info(
                "Torch profiler enabled. AsyncOmniLLM CPU traces will be collected under %s",
-                envs.VLLM_TORCH_PROFILER_DIR,
+                profiler_dir,
            )
            worker_name = f"{socket.gethostname()}_{os.getpid()}.async_omni_llm"
            self.profiler = torch.profiler.profile(
                activities=[
                    torch.profiler.ProfilerActivity.CPU,
                ],
-                with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
+                with_stack=vllm_config.profiler_config.torch_profiler_with_stack,
                on_trace_ready=torch.profiler.tensorboard_trace_handler(
-                    envs.VLLM_TORCH_PROFILER_DIR,
+                    profiler_dir,
                    worker_name=worker_name,
-                    use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP,
+                    use_gzip=vllm_config.profiler_config.torch_profiler_use_gzip,
                ),
            )
        else:
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import json
 import multiprocessing as mp
 import os
 import time
@@ -626,6 +626,7 @@ def _stage_worker(
        for device_id in devices_to_lock:
            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
            lock_acquired = False
+            next_wait_log = wait_start + 5.0

            while not lock_acquired:
                try:
@@ -654,6 +655,7 @@ def _stage_worker(
                            )
                            break

+
                        # Wait a bit before retrying
                        _time.sleep(0.1)
                except OSError as e:
@@ -1171,6 +1173,7 @@ async def _stage_worker_async(
        for device_id in devices_to_lock:
            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
            lock_acquired = False
+            next_wait_log = wait_start + 5.0

            while not lock_acquired:
                try:
@@ -1200,6 +1203,7 @@ async def _stage_worker_async(
                            )
                            break

+
                        # Wait a bit before retrying
                        _time.sleep(0.1)
                except OSError as e:
@@ -1234,6 +1238,8 @@ async def _stage_worker_async(
            break
        engine_args["stage_connector_spec"] = stage_connector_spec
        engine_args["stage_id"] = stage_id
+    stage_engine = None
+    engine_init_error = None
    try:
        if stage_type == "diffusion":
            # For diffusion, we need to extract diffusion-specific config
@@ -1261,6 +1267,9 @@ async def _stage_worker_async(
                usage_context=usage_context,
                engine_args=omni_engine_args,
            )
+    except Exception as e:
+        engine_init_error = e
+        logger.exception("[Stage-%s] Engine initialization failed: %s", stage_id, e)
    finally:
        # Release all locks by closing file descriptors
        # Locks are automatically released when file descriptors are closed
@@ -1272,6 +1281,11 @@ async def _stage_worker_async(
                logger.debug("Released initialization lock (fd=%s)", lock_fd)
            except (OSError, ValueError):
                pass
+    # If engine init failed, re-raise the exception
+    if engine_init_error is not None:
+        raise engine_init_error
+    if stage_engine is None:
+        raise RuntimeError(f"[Stage-{stage_id}] Engine was not initialized")
    omni_stage.set_async_engine(stage_engine)
    if hasattr(omni_stage.async_engine, "log_stats") and omni_stage.async_engine.log_stats:

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import base64
 import io
-import json
 import multiprocessing
 import multiprocessing.forkserver as forkserver
 import os
@@ -31,9 +30,9 @@ from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.launcher import serve_http
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
-from vllm.entrypoints.openai.api_server import base, load_log_config
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
 from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
+from vllm.entrypoints.openai.basic.api_router import base
 from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionRequest,
    ChatCompletionResponse,
@@ -53,6 +52,7 @@ from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.orca_metrics import metrics_header
 from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
+from vllm.entrypoints.openai.server_utils import load_log_config
 from vllm.entrypoints.openai.translations.serving import (
    OpenAIServingTranscription,
    OpenAIServingTranslation,
@@ -206,7 +206,8 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
        if not supported_tasks:
            supported_tasks = ("generate",)

-        app = build_openai_app(args)
+        # OMNI: Pass supported_tasks to build_app (required by upstream vLLM)
+        app = build_openai_app(args, supported_tasks)
        # OMNI: Remove upstream routes that we override with omni-specific handlers
        _remove_route_from_app(app, "/v1/chat/completions", {"POST"})
        _remove_route_from_app(app, "/v1/models", {"GET"})  # Remove upstream /v1/models to use omni's handler
@@ -13,7 +13,7 @@ import torch
 from fastapi import Request
 from PIL import Image
 from pydantic import TypeAdapter
-from vllm.renderers import RendererLike
+from vllm.renderers.protocol import BaseRenderer

 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
@@ -77,6 +77,7 @@ from vllm.renderers.mistral import (
 )
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import TokenizerLike
+from vllm.tokenizers import TokenizerLike as AnyTokenizer
 from vllm.tokenizers.mistral import (
    MistralTokenizer,
    maybe_serialize_tool_calls,
@@ -85,7 +86,6 @@ from vllm.tokenizers.mistral import (
 )
 from vllm.tool_parsers import ToolParser
 from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
-from vllm.transformers_utils.tokenizer import AnyTokenizer
 from vllm.utils.collection_utils import as_list, is_list_of

 from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures
@@ -231,23 +231,28 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                chat_template_kwargs = request.chat_template_kwargs or {}
                chat_template_kwargs.update(reasoning_effort=request.reasoning_effort)

+                # Merge chat_template_kwargs with defaults
+                merged_template_kwargs = self._prepare_extra_chat_template_kwargs(
+                    chat_template_kwargs,
+                    self.default_chat_template_kwargs,
+                )
                (
                    conversation,
                    request_prompts,
                    engine_prompts,
                ) = await self._preprocess_chat(
                    request,
-                    renderer,
                    request.messages,
-                    chat_template=request.chat_template or self.chat_template,
-                    chat_template_content_format=self.chat_template_content_format,
+                    default_template=request.chat_template or self.chat_template,
+                    default_template_content_format=self.chat_template_content_format,
+                    default_template_kwargs=merged_template_kwargs,
+                    tool_dicts=tool_dicts,
+                    tool_parser=tool_parser,
+                    # OMNI: Additional parameters
+                    renderer=renderer,
                    add_generation_prompt=request.add_generation_prompt,
                    continue_final_message=request.continue_final_message,
-                    tool_dicts=tool_dicts,
                    documents=getattr(request, "documents", None),
-                    chat_template_kwargs=chat_template_kwargs,
-                    default_chat_template_kwargs=self.default_chat_template_kwargs,
-                    tool_parser=tool_parser,
                    add_special_tokens=request.add_special_tokens,
                )
            else:
@@ -329,17 +334,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
    async def _preprocess_chat(
        self,
        request: ChatLikeRequest | ResponsesRequest,
-        renderer: RendererLike,
        messages: list[ChatCompletionMessageParam],
-        chat_template: str | None,
-        chat_template_content_format: ChatTemplateContentFormatOption,
+        default_template: str | None,
+        default_template_content_format: ChatTemplateContentFormatOption,
+        default_template_kwargs: dict[str, Any] | None = None,
+        tool_dicts: list[dict[str, Any]] | None = None,
+        tool_parser: Callable[[TokenizerLike], ToolParser] | None = None,
+        # OMNI: Additional parameters for backward compatibility
+        renderer: BaseRenderer | None = None,
        add_generation_prompt: bool = True,
        continue_final_message: bool = False,
-        tool_dicts: list[dict[str, Any]] | None = None,
        documents: list[dict[str, str]] | None = None,
-        chat_template_kwargs: dict[str, Any] | None = None,
-        default_chat_template_kwargs: dict[str, Any] | None = None,
-        tool_parser: Callable[[TokenizerLike], ToolParser] | None = None,
        add_special_tokens: bool = False,
    ) -> tuple[
        list[ConversationMessage],
@@ -347,8 +352,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
        list[TokensPrompt],
    ]:
        model_config = self.model_config
+        # OMNI: Use self.renderer if renderer not provided (vLLM new API)
+        if renderer is None:
+            renderer = self.renderer
        tokenizer = renderer.get_tokenizer() if renderer is not None else None

+        # Map new parameter names to internal variables
+        chat_template = default_template
+        chat_template_content_format = default_template_content_format
+
        if tokenizer is None or isinstance(tokenizer, MistralTokenizer):
            resolved_content_format = (
                chat_template_content_format if chat_template_content_format != "auto" else "string"
@@ -369,12 +381,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
        )

-        # Merge default_chat_template_kwargs with request-provided kwargs
-        # Request kwargs take precedence over defaults
-        merged_kwargs = self._prepare_extra_chat_template_kwargs(
-            chat_template_kwargs,
-            default_chat_template_kwargs,
-        )
+        # OMNI: Handle merged kwargs - default_template_kwargs replaces the old two-param system
+        merged_kwargs = default_template_kwargs or {}

        _chat_template_kwargs: dict[str, Any] = dict(
            chat_template=chat_template,
@@ -435,12 +443,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            )
            prompt_inputs = TokensPrompt(prompt=request_prompt, prompt_token_ids=[1])
        elif isinstance(request_prompt, str):
-            prompt_inputs = await self._tokenize_prompt_input_async(
-                request,
-                tokenizer,
-                request_prompt,
-                add_special_tokens=add_special_tokens,
-            )
+            # OMNI: Direct tokenization using tokenizer (replaces removed _tokenize_prompt_input_async)
+            prompt_token_ids = tokenizer.encode(request_prompt, add_special_tokens=add_special_tokens)
+            prompt_inputs = TokensPrompt(prompt=request_prompt, prompt_token_ids=prompt_token_ids)
        else:
            # For MistralTokenizer
            assert is_list_of(request_prompt, int), "Prompt has to be either a string or a list of token ids"
@@ -3,7 +3,8 @@ import pprint
 from dataclasses import asdict, dataclass, field
 from typing import Any, TypeAlias

-from vllm import PromptType, SamplingParams
+from vllm.inputs import PromptType
+from vllm.sampling_params import SamplingParams

 from vllm_omni.lora.request import LoRARequest

@@ -1,12 +1,13 @@
 from typing import Any

 from typing_extensions import assert_never
-from vllm.inputs.data import SingletonInputs, SingletonPrompt
+from vllm.inputs.data import EmbedsInputs, SingletonInputs, SingletonPrompt, TextPrompt
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
 from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict

 from vllm_omni.inputs.data import (
+    OmniEmbedsPrompt,
    OmniTextPrompt,
    OmniTokenInputs,
    OmniTokensPrompt,
@@ -101,6 +102,25 @@ class OmniInputPreprocessor(InputPreprocessor):

        return inputs

+    def _process_embeds(
+        self,
+        parsed_content: OmniEmbedsPrompt,
+    ) -> EmbedsInputs:
+        """Process embeddings prompt with omni-specific extensions.
+
+        Extends base _process_embeds to handle additional_information payload
+        for direct transfer between pipeline stages.
+        """
+        # Call parent implementation for base embeds processing
+        inputs = super()._process_embeds(parsed_content)
+
+        # Add omni-specific additional_information if present
+        additional_information = parsed_content.get("additional_information")
+        if additional_information is not None:
+            inputs["additional_information"] = additional_information  # type: ignore[typeddict-unknown-key]
+
+        return inputs
+
    def _prompt_to_llm_inputs(
        self,
        prompt: SingletonPrompt,
@@ -114,15 +134,17 @@ class OmniInputPreprocessor(InputPreprocessor):
        Arguments:

        * prompt: single encoder or decoder input prompt
-        * lora_request: this is only valid for decoder prompts
-        * return_mm_hashes: whether to return multimodal hashes

        Returns:

-        * Input container compatible with vLLM's singleton prompt handling.
+        * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance
        """
        parsed = parse_singleton_prompt_omni(prompt)

+        # Note: omni parsing prioritizes tokens path when both tokens and embeds
+        # exist, keeping both for pipeline stage transfer
+        if parsed["type"] == "embeds":
+            return self._process_embeds(parsed["content"])
        if parsed["type"] == "tokens":
            return self._process_tokens(
                parsed["content"],
@@ -136,7 +158,7 @@ class OmniInputPreprocessor(InputPreprocessor):
            )
        if parsed["type"] == "str":
            return self._process_text(
-                OmniTextPrompt(prompt=parsed["content"]),
+                TextPrompt(prompt=parsed["content"]),
                tokenization_kwargs=tokenization_kwargs,
                mm_uuids=mm_uuids,
            )
@@ -27,6 +27,8 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
    Qwen2_5OmniThinkerDummyInputsBuilder,
    Qwen2_5OmniThinkerMultiModalProcessor,
    Qwen2_5OmniThinkerProcessingInfo,
+    check_interleaved_audio_video,
+    merge_interleaved_embeddings,
 )
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
    Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase,
@@ -43,6 +45,7 @@ from vllm.model_executor.models.qwen2_5_vl import (
 from vllm.model_executor.models.utils import (
    AutoWeightsLoader,
    WeightsMapper,
+    _merge_multimodal_embeddings,
    init_vllm_registered_model,
    maybe_prefix,
    split_list_into_ranges,
@@ -275,28 +278,32 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                "in the audio tower part."
            )

-        if multimodal_config.get_limit_per_prompt("audio"):
-            self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
-        else:
-            self.audio_tower = None
-
-        if multimodal_config.get_limit_per_prompt("image") or multimodal_config.get_limit_per_prompt("video"):
-            self.visual = Qwen2_5_VisionTransformer(
-                vision_config=thinker_config.vision_config,
-                norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
-                quant_config=quant_config,
-                prefix=maybe_prefix(prefix, "visual"),
-            )
-        else:
-            self.visual = None
-
        self.quant_config = quant_config
-        self.language_model = init_vllm_registered_model(
-            vllm_config=vllm_config,
-            prefix=maybe_prefix(prefix, "language_model"),
-            hf_config=thinker_config.text_config,
-            architectures=["Qwen2ForCausalLM"],
-        )
+
+        with self._mark_tower_model(vllm_config, "audio"):
+            if multimodal_config.get_limit_per_prompt("audio"):
+                self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
+            else:
+                self.audio_tower = None
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            if multimodal_config.get_limit_per_prompt("image") or multimodal_config.get_limit_per_prompt("video"):
+                self.visual = Qwen2_5_VisionTransformer(
+                    vision_config=thinker_config.vision_config,
+                    norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "visual"),
+                )
+            else:
+                self.visual = None
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "language_model"),
+                hf_config=thinker_config.text_config,
+                architectures=["Qwen2ForCausalLM"],
+            )

        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors

@@ -503,8 +510,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                multimodal_embeddings += tuple(audio_embeddings)
        return multimodal_embeddings

-    # TODO (ywang96): support overlapping modality embeddings so that
-    # `use_audio_in_video` will work on V1.
    def embed_input_ids(
        self,
        input_ids: torch.Tensor,
@@ -517,13 +522,41 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
        if multimodal_embeddings is None or is_multimodal is None:
            return super().embed_input_ids(input_ids)

-        return super().embed_input_ids(
+        inputs_embeds = self._embed_text_input_ids(
            input_ids,
-            multimodal_embeddings=multimodal_embeddings,
+            self.get_language_model().embed_input_ids,
            is_multimodal=is_multimodal,
            handle_oov_mm_token=handle_oov_mm_token,
        )

+        if len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        # Check for audio-in-video: interleaved video and audio tokens
+        # in the multimodal region.
+        video_token_id = self.config.video_token_index
+        audio_token_id = self.config.audio_token_index
+
+        is_video = is_multimodal & (input_ids == video_token_id)
+        is_audio = is_multimodal & (input_ids == audio_token_id)
+
+        num_video = is_video.sum().item()
+        num_audio = is_audio.sum().item()
+
+        if check_interleaved_audio_video(is_video, is_audio, num_video, num_audio):
+            return merge_interleaved_embeddings(
+                inputs_embeds,
+                multimodal_embeddings,
+                is_video,
+                is_audio,
+                is_multimodal,
+                num_video,
+                num_audio,
+            )
+
+        # Default: standard merge (no interleaving)
+        return _merge_multimodal_embeddings(inputs_embeds, multimodal_embeddings, is_multimodal)
+
    def forward(
        self,
        input_ids: torch.Tensor,
@@ -22,9 +22,9 @@
 # limitations under the License.
 """Inference-only Qwen3-Omni-Moe model (thinker part)."""

-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Iterator, Mapping, Sequence
 from functools import partial
-from typing import Any
+from typing import Any, Literal, cast

 import numpy as np
 import torch
@@ -42,8 +42,9 @@ 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 VllmConfig
+from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
 from vllm.distributed import get_pp_group
+from vllm.inputs.data import PromptType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
@@ -52,12 +53,15 @@ from vllm.model_executor.models.interfaces import (
    SupportsMRoPE,
    SupportsMultiModal,
    SupportsPP,
+    SupportsTranscription,
 )
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
    Qwen2_5OmniAudioFeatureInputs,
    Qwen2_5OmniThinkerDummyInputsBuilder,
    Qwen2_5OmniThinkerMultiModalProcessor,
+    check_interleaved_audio_video,
+    merge_interleaved_embeddings,
 )
 from vllm.model_executor.models.qwen2_5_vl import (
    Qwen2_5_VLProcessingInfo,
@@ -76,9 +80,6 @@ from vllm.model_executor.models.utils import (
    _merge_multimodal_embeddings,
    maybe_prefix,
 )
-from vllm.model_executor.models.vision import (
-    get_llm_pos_ids_for_vision,
-)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems
 from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems
@@ -90,6 +91,7 @@ from vllm.multimodal.processing.processor import (
    PromptUpdateDetails,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.processor import cached_processor_from_config

 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
    Qwen2_5OmniConditionalGenerationMixin,
@@ -102,6 +104,29 @@ except (ImportError, ModuleNotFoundError):

 logger = init_logger(__name__)

+# Speech input languages supported by Qwen3-Omni
+# From: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct
+ISO639_1_SUPPORTED_LANGS = {
+    "en": "English",
+    "zh": "Chinese",
+    "ko": "Korean",
+    "ja": "Japanese",
+    "de": "German",
+    "ru": "Russian",
+    "it": "Italian",
+    "fr": "French",
+    "es": "Spanish",
+    "pt": "Portuguese",
+    "ms": "Malay",
+    "nl": "Dutch",
+    "id": "Indonesian",
+    "tr": "Turkish",
+    "vi": "Vietnamese",
+    "yue": "Cantonese",
+    "ar": "Arabic",
+    "ur": "Urdu",
+}
+

@support_torch_compile(
    dynamic_arg_dims={
@@ -639,6 +664,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
    SupportsPP,
    SupportsMRoPE,
    Qwen3OmniMoeConditionalGenerationMixin,
+    SupportsTranscription,
 ):
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
@@ -660,6 +686,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
        ],
    }

+    supported_languages = ISO639_1_SUPPORTED_LANGS
+
    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
@@ -679,46 +707,57 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
        multimodal_config = vllm_config.model_config.multimodal_config
        self.config = thinker_config
        self.multimodal_config = multimodal_config
-
-        self.audio_tower = Qwen3OmniMoeAudioEncoder(
-            thinker_config.audio_config,
-        )
-
-        self.visual = Qwen3Omni_VisionTransformer(
-            vision_config=thinker_config.vision_config,
-            norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
-            quant_config=quant_config,
-            prefix=maybe_prefix(prefix, "visual"),
-        )
        self.quant_config = quant_config

-        self.language_model = Qwen3MoeLLMForCausalLM(
-            vllm_config=vllm_config.with_hf_config(thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"]),
-            prefix=maybe_prefix(prefix, "language_model"),
-        )
-
-        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+        with self._mark_tower_model(vllm_config, "audio"):
+            self.audio_tower = Qwen3OmniMoeAudioEncoder(
+                thinker_config.audio_config,
+                prefix=maybe_prefix(prefix, "audio_tower"),
+            )

        self.use_deepstack = hasattr(thinker_config.vision_config, "deepstack_visual_indexes")
        self.deepstack_num_level = (
            len(thinker_config.vision_config.deepstack_visual_indexes) if self.use_deepstack else 0
        )
-        # register buffer for deepstack
-        self.deepstack_input_embeds = (
-            [
-                torch.zeros(
-                    vllm_config.scheduler_config.max_num_batched_tokens,
-                    thinker_config.text_config.hidden_size,
-                )
-                for _ in range(self.deepstack_num_level)
-            ]
-            if self.use_deepstack
-            else None
-        )
        self.visual_dim = thinker_config.vision_config.out_hidden_size
        self.multiscale_dim = self.visual_dim * self.deepstack_num_level

-    def _get_deepstack_input_embeds(self, num_tokens: int) -> IntermediateTensors:
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Qwen3Omni_VisionTransformer(
+                vision_config=thinker_config.vision_config,
+                norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+            # register buffer for deepstack
+            if self.use_deepstack:
+                self.deepstack_input_embeds = [
+                    torch.zeros(
+                        vllm_config.scheduler_config.max_num_batched_tokens,
+                        thinker_config.text_config.hidden_size,
+                    )
+                    for _ in range(self.deepstack_num_level)
+                ]
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = Qwen3MoeLLMForCausalLM(
+                vllm_config=vllm_config.with_hf_config(
+                    thinker_config.text_config,
+                    architectures=["Qwen3MoeForCausalLM"],
+                ),
+                prefix=maybe_prefix(prefix, "language_model"),
+            )
+
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+
+    def _get_deepstack_input_embeds(
+        self,
+        num_tokens: int,
+    ) -> IntermediateTensors | None:
+        if not getattr(self, "deepstack_input_embeds", None):
+            return None  # If vision tower is skipped
+
        # get deepstack_input_embeds from buffer, and clear the buffer
        return IntermediateTensors(
            {
@@ -728,6 +767,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
        )

    def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None:
+        if not getattr(self, "deepstack_input_embeds", None):
+            return
+
        # set deepstack_input_embeds to buffer
        num_tokens = deepstack_input_embeds.size(1)
        if num_tokens > self.deepstack_input_embeds[0].size(0):
@@ -744,6 +786,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
            self.deepstack_input_embeds[idx][:num_tokens].copy_(deepstack_input_embeds[idx])

    def _clear_deepstack_input_embeds(self, num_tokens: int) -> None:
+        if not getattr(self, "deepstack_input_embeds", None):
+            return
+
        # clear deepstack_input_embeds in buffer
        if num_tokens > 0:
            for idx in range(self.deepstack_num_level):
@@ -808,6 +853,17 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
            return inputs_embeds

+        # Detect interleaved audio-in-video early, since it affects
+        # both the deepstack path and the final embedding merge.
+        video_token_id = self.config.video_token_id
+        audio_token_id = self.config.audio_token_id
+        is_video = is_multimodal & (input_ids == video_token_id)
+        is_audio = is_multimodal & (input_ids == audio_token_id)
+        num_video = is_video.sum().item()
+        num_audio = is_audio.sum().item()
+
+        is_interleaved = check_interleaved_audio_video(is_video, is_audio, num_video, num_audio)
+
        deepstack_input_embeds = None
        # split the feat dim to obtain multi-scale visual feature
        has_vision_embeddings = [
@@ -816,12 +872,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
        if self.visual.deepstack_visual_indexes is not None and any(has_vision_embeddings):
            multiscale_len = len(self.visual.deepstack_visual_indexes)
            multimodal_embeddings_multiscale = []
-            is_vision = torch.zeros_like(is_multimodal)
-            mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]
-            mm_position_idx = 0
+
+            if is_interleaved:
+                # Use input_ids-based mask for correct vision positions
+                # when audio and video tokens are interleaved.
+                is_vision = is_video.clone()
+            else:
+                is_vision = torch.zeros_like(is_multimodal)
+                mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]
+                mm_position_idx = 0
+
            for index, embeddings in enumerate(multimodal_embeddings):
                num_tokens = embeddings.shape[0]
-                current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens]

                # Vision embeddings
                if embeddings.shape[-1] != self.config.text_config.hidden_size:
@@ -830,13 +892,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
                    embeddings_main, embeddings_multiscale = torch.split(embeddings, [visual_dim, multi_dim], dim=-1)
                    multimodal_embeddings[index] = embeddings_main
                    multimodal_embeddings_multiscale.append(embeddings_multiscale)
-                    is_vision[current_positions] = True
+                    if not is_interleaved:
+                        current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens]
+                        is_vision[current_positions] = True

                # Audio embeddings
                else:
-                    is_vision[current_positions] = False
+                    if not is_interleaved:
+                        current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens]
+                        is_vision[current_positions] = False

-                mm_position_idx += num_tokens
+                if not is_interleaved:
+                    mm_position_idx += num_tokens

            deepstack_input_embeds = inputs_embeds.new_zeros(
                inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1)
@@ -853,6 +920,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
            )
            self._set_deepstack_input_embeds(deepstack_input_embeds)

+        if is_interleaved:
+            return merge_interleaved_embeddings(
+                inputs_embeds,
+                multimodal_embeddings,
+                is_video,
+                is_audio,
+                is_multimodal,
+                num_video,
+                num_audio,
+            )
+
+        # Default: standard merge (no interleaving)
        inputs_embeds = _merge_multimodal_embeddings(
            inputs_embeds=inputs_embeds,
            multimodal_embeddings=multimodal_embeddings,
@@ -910,214 +989,304 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

        return loaded_weights

+    def _compute_audio_token_count(self, audio_feature_length: int) -> int:
+        """Compute audio tokens from feature length using Qwen3-Omni formula."""
+        return _get_feat_extract_output_lengths(torch.tensor([audio_feature_length])).item()
+
+    def _get_audio_for_video_mapping(self, mm_features: list[MultiModalFeatureSpec]) -> tuple[dict[int, int], set[int]]:
+        """
+        Map video offset -> paired audio_feature_length for use_audio_in_video.
+
+        When use_audio_in_video=True, audio is interleaved within video.
+        The pairing is based on feature order in mm_features.
+
+        Returns:
+            Tuple of (video_offset -> audio_feature_length mapping,
+                      set of paired audio offsets to skip)
+        """
+        videos_with_audio = [
+            f
+            for f in mm_features
+            if f.modality == "video" and f.data.get("use_audio_in_video") and f.data["use_audio_in_video"].data.item()
+        ]
+        audios = [f for f in mm_features if f.modality == "audio"]
+
+        mapping: dict[int, int] = {}
+        paired_audio_offsets: set[int] = set()
+        for i, video_f in enumerate(videos_with_audio):
+            if i < len(audios):
+                audio_len = audios[i].data["audio_feature_lengths"].data.item()
+                mapping[video_f.mm_position.offset] = audio_len
+                paired_audio_offsets.add(audios[i].mm_position.offset)
+        return mapping, paired_audio_offsets
+
+    def iter_mm_features(self, mm_features: list[MultiModalFeatureSpec]) -> Iterator[tuple[int, str, dict[str, Any]]]:
+        """
+        Iterate over multimodal features sorted by position offset.
+
+        Yields: (offset, modality, feature_data) where feature_data contains:
+        - image: {"grid_t", "grid_h", "grid_w", "t_factor"}
+        - video: {"grid_t", "grid_h", "grid_w", "t_factor",
+                  "use_audio_in_video", "audio_feature_length"}
+        - audio: {"audio_feature_length"}
+        """
+        config = self.config
+        spatial_merge_size = config.vision_config.spatial_merge_size
+        position_id_per_seconds = config.position_id_per_seconds
+
+        sorted_features = sorted(mm_features, key=lambda f: f.mm_position.offset)
+        audio_for_video, paired_audio_offsets = self._get_audio_for_video_mapping(sorted_features)
+
+        for mm_feature in sorted_features:
+            offset = mm_feature.mm_position.offset
+            modality = mm_feature.modality
+
+            if modality == "image":
+                t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
+                yield (
+                    offset,
+                    "image",
+                    {
+                        "grid_t": t,
+                        "grid_h": h // spatial_merge_size,
+                        "grid_w": w // spatial_merge_size,
+                        "t_factor": position_id_per_seconds,
+                    },
+                )
+            elif modality == "video":
+                t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
+                second_per_grid_ts = 2.0
+                if mm_feature.data.get("second_per_grid_ts"):
+                    second_per_grid_ts = mm_feature.data["second_per_grid_ts"].data.item()
+                use_audio_in_video = bool(
+                    mm_feature.data.get("use_audio_in_video") and mm_feature.data["use_audio_in_video"].data.item()
+                )
+
+                yield (
+                    offset,
+                    "video",
+                    {
+                        "grid_t": t,
+                        "grid_h": h // spatial_merge_size,
+                        "grid_w": w // spatial_merge_size,
+                        "t_factor": second_per_grid_ts * position_id_per_seconds,
+                        "use_audio_in_video": use_audio_in_video,
+                        "audio_feature_length": audio_for_video.get(offset),
+                    },
+                )
+            elif modality == "audio":
+                if offset not in paired_audio_offsets:
+                    audio_len = mm_feature.data["audio_feature_lengths"].data.item()
+                    yield offset, "audio", {"audio_feature_length": audio_len}
+
+    def _compute_interleaved_positions(self, start_idx: int, data: dict[str, Any]) -> tuple[np.ndarray, int]:
+        """
+        Compute positions for interleaved video+audio using Qwen3 token-by-token
+        interleaving logic.
+
+        Returns: (position_ids [3, N], total_token_count)
+        """
+        grid_t = data["grid_t"]
+        grid_h = data["grid_h"]
+        grid_w = data["grid_w"]
+        t_factor = data["t_factor"]
+        audio_feature_length = data["audio_feature_length"]
+
+        audio_len = self._compute_audio_token_count(audio_feature_length)
+
+        h_index = np.tile(np.arange(grid_h).reshape(1, -1, 1), (grid_t, 1, grid_w)).flatten()
+        w_index = np.tile(np.arange(grid_w).reshape(1, 1, -1), (grid_t, grid_h, 1)).flatten()
+        t_index_raw = np.arange(grid_t)
+        t_index_scaled = (t_index_raw * t_factor).astype(np.int64)
+        t_index = np.repeat(t_index_scaled, grid_h * grid_w)
+
+        video_pos = np.stack([t_index, h_index, w_index]) + start_idx
+        audio_pos = np.broadcast_to(np.arange(audio_len), (3, audio_len)) + start_idx
+
+        video_t_values = video_pos[0]
+        audio_t_values = audio_pos[0]
+
+        pos_ids_list: list[np.ndarray] = []
+        video_idx, audio_idx = 0, 0
+        num_video = grid_t * grid_h * grid_w
+
+        while video_idx < num_video and audio_idx < audio_len:
+            if video_t_values[video_idx] <= audio_t_values[audio_idx]:
+                pos_ids_list.append(video_pos[:, video_idx : video_idx + 1])
+                video_idx += 1
+            else:
+                pos_ids_list.append(audio_pos[:, audio_idx : audio_idx + 1])
+                audio_idx += 1
+
+        if video_idx < num_video:
+            pos_ids_list.append(video_pos[:, video_idx:])
+        if audio_idx < audio_len:
+            pos_ids_list.append(audio_pos[:, audio_idx:])
+
+        total_tokens = num_video + audio_len
+        return np.concatenate(pos_ids_list, axis=1), total_tokens
+
+    @classmethod
+    def get_speech_to_text_config(cls, model_config: ModelConfig, task_type: str) -> SpeechToTextConfig:
+        processor = cached_processor_from_config(model_config, processor_cls=Qwen3OmniMoeProcessor)
+        return SpeechToTextConfig(
+            max_audio_clip_s=processor.feature_extractor.chunk_length,
+            sample_rate=processor.feature_extractor.sampling_rate,
+            min_energy_split_window_size=None,
+        )
+
+    @classmethod
+    def get_generation_prompt(
+        cls,
+        audio: np.ndarray,
+        stt_config: SpeechToTextConfig,
+        model_config: ModelConfig,
+        language: str | None,
+        task_type: Literal["transcribe", "translate"],
+        request_prompt: str,
+        to_language: str | None,
+    ) -> PromptType:
+        """
+        Construct a transcription/translation prompt for Qwen3-Omni.
+        """
+        # Transcribe this audio [into <language>] | for transcription
+        # Translate this audio [from <language> into <to_language>] | for translation
+        instruction = "Transcribe" if task_type == "transcribe" else "Translate"
+        instruction += " this audio"
+
+        # Default to_language to English for translation
+        if task_type == "translate" and to_language is None:
+            to_language = "en"
+
+        # Get full language names from supported_languages mapping
+        full_lang_name = cls.supported_languages.get(language, "")
+        full_lang_name_to = cls.supported_languages.get(to_language, "")
+
+        if task_type == "transcribe" and full_lang_name:
+            instruction += f" into {full_lang_name}"
+        elif task_type == "translate":
+            if full_lang_name:
+                instruction += f" from {full_lang_name}"
+            if full_lang_name_to:
+                instruction += f" into {full_lang_name_to}"
+
+        instruction += "."
+
+        if request_prompt:
+            instruction += f" {request_prompt}"
+
+        processor = cached_processor_from_config(model_config, processor_cls=Qwen3OmniMoeProcessor)
+        # Audio placeholder format: <|audio_start|><|audio_pad|><|audio_end|>
+        audio_placeholder = "<|audio_start|><|audio_pad|><|audio_end|>"
+        user_content = f"{audio_placeholder}{instruction}"
+
+        messages = [{"role": "user", "content": user_content}]
+        prompt = processor.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+        audio_data = (audio, stt_config.sample_rate)
+        prompts_dict = {"multi_modal_data": {"audio": audio_data}, "prompt": prompt}
+        return cast(PromptType, prompts_dict)
+
    def get_mrope_input_positions(
        self,
        input_tokens: list[int],
        mm_features: list[MultiModalFeatureSpec],
    ) -> tuple[torch.Tensor, int]:
-        kwargs = MultiModalFeatureSpec.gather_kwargs(
-            mm_features,
-            {
-                "image_grid_thw",
-                "video_grid_thw",
-                "second_per_grid_ts",
-                "audio_feature_lengths",
-                "use_audio_in_video",
-            },
-        )
-        image_grid_thw = kwargs.get("image_grid_thw", [])
-        video_grid_thw = kwargs.get("video_grid_thw", [])
-        second_per_grid_ts = kwargs.get("second_per_grid_ts", [])
-        audio_feature_lengths = kwargs.get("audio_feature_lengths", [])
-        use_audio_in_video = any(kwargs.get("use_audio_in_video", []))
+        """Compute M-RoPE input positions using mm_features directly."""
+        seq_len = len(input_tokens)

-        image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(image_grid_thw)
-        video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(video_grid_thw)
-
-        input_ids = torch.tensor(input_tokens)
-        if input_ids is None or input_ids.ndim != 1:
-            raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids")
-
-        seq_len = input_ids.shape[0]
-
-        if isinstance(audio_feature_lengths, list):
-            audio_feature_lengths = torch.tensor(audio_feature_lengths, dtype=torch.long)
-
-        if not len(second_per_grid_ts) and len(video_grid_thw):
-            second_per_grid_ts = 2.0
-            second_per_grids = torch.ones(len(video_grid_thw), dtype=torch.float32) * second_per_grid_ts
-        else:
-            second_per_grids = torch.tensor(second_per_grid_ts, dtype=torch.float32)
-
-        config = self.config
-        spatial_merge_size = config.vision_config.spatial_merge_size
-        image_token_id = config.image_token_id
-        video_token_id = config.video_token_id
-        audio_token_id = config.audio_token_id
-        vision_start_token_id = config.vision_start_token_id
-        audio_start_token_id = config.audio_start_token_id
-        position_id_per_seconds = config.position_id_per_seconds
-
-        vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
-        if vision_start_indices.numel() > 0:
-            vision_tokens = input_ids[vision_start_indices + 1]
-        else:
-            vision_tokens = input_ids.new_empty((0,), dtype=input_ids.dtype)
-        audio_nums = torch.sum(input_ids == audio_start_token_id)
-        image_nums = (vision_tokens == image_token_id).sum()
-        video_nums = (
-            (vision_tokens == audio_start_token_id).sum()
-            if use_audio_in_video
-            else (vision_tokens == video_token_id).sum()
-        )
-
-        llm_pos_ids_list: list[torch.Tensor] = []
+        llm_pos_ids_list: list[np.ndarray] = []
        st = 0
-        image_idx = 0
-        video_idx = 0
-        audio_idx = 0
-        remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums  # noqa: E501
-        multimodal_nums = image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums  # noqa: E501

-        for _ in range(multimodal_nums):
-            st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-            if (image_token_id in input_tokens or video_token_id in input_tokens) and (
-                remain_videos > 0 or remain_images > 0
-            ):
-                ed_vision_start = input_tokens.index(vision_start_token_id, st)
-            else:
-                ed_vision_start = len(input_tokens) + 1
-            if audio_token_id in input_tokens and remain_audios > 0:
-                ed_audio_start = input_tokens.index(audio_start_token_id, st)
-            else:
-                ed_audio_start = len(input_tokens) + 1
-            min_ed = min(ed_vision_start, ed_audio_start)
+        for offset, modality, data in self.iter_mm_features(mm_features):
+            text_len = offset - st
+            st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0

-            if min_ed == ed_audio_start:
-                text_len = min_ed - st
-                if text_len != 0:
-                    st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                    llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                bos_len = 1
-                llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx])
-                llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
-                llm_pos_ids_list.append(llm_pos_ids)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                eos_len = 1
-                llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st += text_len + bos_len + audio_len + eos_len
-                audio_idx += 1
-                remain_audios -= 1
-            elif min_ed == ed_vision_start and input_ids[ed_vision_start + 1] == image_token_id:
-                text_len = min_ed - st
-                if text_len != 0:
-                    st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                    llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                bos_len = 1
-                llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                grid_t = image_grid_thw[image_idx][0]
-                grid_hs = image_grid_thw[:, 1]
-                grid_ws = image_grid_thw[:, 2]
-                t_index = torch.arange(grid_t) * position_id_per_seconds
-                llm_pos_ids = get_llm_pos_ids_for_vision(
-                    st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
-                llm_pos_ids_list.append(llm_pos_ids)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                eos_len = 1
-                llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st += text_len + bos_len + image_len + eos_len
-                image_idx += 1
-                remain_images -= 1
-            elif (
-                min_ed == ed_vision_start
-                and input_ids[ed_vision_start + 1] == video_token_id
-                and not use_audio_in_video
-            ):
-                text_len = min_ed - st
-                if text_len != 0:
-                    st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                    llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                bos_len = 1
-                llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                grid_t = video_grid_thw[video_idx][0]
-                grid_hs = video_grid_thw[:, 1]
-                grid_ws = video_grid_thw[:, 2]
-                t_index = torch.arange(grid_t) * float(second_per_grids[video_idx].item()) * position_id_per_seconds
-                llm_pos_ids = get_llm_pos_ids_for_vision(
-                    st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
-                llm_pos_ids_list.append(llm_pos_ids)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                eos_len = 1
-                llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st += text_len + bos_len + video_len + eos_len
-                video_idx += 1
-                remain_videos -= 1
-            elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start and use_audio_in_video:
-                text_len = min_ed - st
-                if text_len != 0:
-                    st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                    llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                bos_len = 1
-                bos_block = torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
-                llm_pos_ids_list.append(bos_block)
-                llm_pos_ids_list.append(bos_block)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx])
-                audio_llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
-                grid_t = video_grid_thw[video_idx][0]
-                grid_hs = video_grid_thw[:, 1]
-                grid_ws = video_grid_thw[:, 2]
-                t_index = torch.arange(grid_t) * float(second_per_grids[video_idx].item()) * position_id_per_seconds
-                video_llm_pos_ids = get_llm_pos_ids_for_vision(
-                    st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
-                )
-                video_data_index, audio_data_index = 0, 0
-                while video_data_index < video_llm_pos_ids.shape[-1] and audio_data_index < audio_llm_pos_ids.shape[-1]:
-                    if video_llm_pos_ids[0][video_data_index] <= audio_llm_pos_ids[0][audio_data_index]:
-                        llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_data_index + 1])
-                        video_data_index += 1
-                    else:
-                        llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_data_index + 1])
-                        audio_data_index += 1
-                if video_data_index < video_llm_pos_ids.shape[-1]:
-                    llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_llm_pos_ids.shape[-1]])
-                if audio_data_index < audio_llm_pos_ids.shape[-1]:
-                    llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_llm_pos_ids.shape[-1]])
-                video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
-                st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-                eos_len = 1
-                eos_block = torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
-                llm_pos_ids_list.append(eos_block)
-                llm_pos_ids_list.append(eos_block)
-                st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2  # noqa: E501
-                audio_idx += 1
-                video_idx += 1
-                remain_videos -= 1
-                remain_audios -= 1
+            if text_len > 0:
+                llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx)
+                st_idx += text_len

-        if st < len(input_tokens):
-            st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-            text_len = len(input_tokens) - st
-            llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+            bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
+            llm_pos_ids_list.append(bos_pos)
+            st_idx += 1

-        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+            if modality == "audio":
+                audio_tokens = self._compute_audio_token_count(data["audio_feature_length"])
+                audio_pos = np.broadcast_to(np.arange(audio_tokens), (3, audio_tokens)) + st_idx
+                llm_pos_ids_list.append(audio_pos)
+                st_idx = int(audio_pos.max()) + 1
+
+                eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
+                llm_pos_ids_list.append(eos_pos)
+                st = offset + 1 + audio_tokens + 1
+
+            elif modality == "image":
+                grid_t = data["grid_t"]
+                grid_h = data["grid_h"]
+                grid_w = data["grid_w"]
+                t_factor = data["t_factor"]
+
+                grid_indices = np.indices((grid_t, grid_h, grid_w))
+                if t_factor != 1.0:
+                    grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
+                llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx)
+
+                image_len = grid_t * grid_h * grid_w
+                st_idx = int(llm_pos_ids_list[-1].max()) + 1
+
+                eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
+                llm_pos_ids_list.append(eos_pos)
+                st = offset + 1 + image_len + 1
+
+            elif modality == "video":
+                grid_t = data["grid_t"]
+                grid_h = data["grid_h"]
+                grid_w = data["grid_w"]
+                t_factor = data["t_factor"]
+
+                if not data["use_audio_in_video"]:
+                    grid_indices = np.indices((grid_t, grid_h, grid_w))
+                    if t_factor != 1.0:
+                        grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
+                    llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx)
+
+                    video_len = grid_t * grid_h * grid_w
+                    st_idx = int(llm_pos_ids_list[-1].max()) + 1
+
+                    eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
+                    llm_pos_ids_list.append(eos_pos)
+                    st = offset + 1 + video_len + 1
+                else:
+                    audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1))
+                    llm_pos_ids_list.append(audio_bos_pos)
+
+                    pos_ids, _ = self._compute_interleaved_positions(st_idx, data)
+                    llm_pos_ids_list.append(pos_ids)
+                    st_idx = int(pos_ids.max()) + 1
+
+                    eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
+                    llm_pos_ids_list.append(eos_pos)
+                    llm_pos_ids_list.append(eos_pos)
+
+                    video_len = grid_t * grid_h * grid_w
+                    audio_len = self._compute_audio_token_count(data["audio_feature_length"])
+                    st = offset + 2 + video_len + audio_len + 2
+
+        if st < seq_len:
+            st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0
+            text_len = seq_len - st
+            llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx)
+
+        llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
        if llm_positions.shape[1] != seq_len:
            raise RuntimeError("Position ids length mismatch with input ids length")

-        mrope_position_delta = llm_positions.max() + 1 - seq_len
-        return llm_positions, mrope_position_delta
+        mrope_position_delta = int(llm_positions.max()) + 1 - seq_len
+        return torch.from_numpy(llm_positions), mrope_position_delta

    def get_mm_mapping(self) -> MultiModelKeys:
        """
@@ -50,7 +50,10 @@ class ExecuteModelState(NamedTuple):
    aux_hidden_states: list[torch.Tensor] | None
    ec_connector_output: Any
    cudagraph_stats: Any
+    # OMNI: multimodal_outputs field for omni-specific multimodal handling
    multimodal_outputs: Any
+    # slot_mappings for attention/drafter (aligned with upstream v1 API)
+    slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None


 class GPUARModelRunner(OmniGPUModelRunner):
@@ -386,6 +389,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
            ec_connector_output,
            cudagraph_stats,
            multimodal_outputs,
+            slot_mappings,  # OMNI: pass slot_mappings for drafter
        )
        self.kv_connector_output = kv_connector_output

@@ -428,6 +432,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
            ec_connector_output,
            cudagraph_stats,
            multimodal_outputs,
+            slot_mappings,  # OMNI: unpack slot_mappings for drafter
        ) = self.execute_model_state
        self.execute_model_state = None

@@ -454,6 +459,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                    aux_hidden_states,
                    spec_decode_metadata,
                    spec_decode_common_attn_metadata,
+                    slot_mappings,  # OMNI: pass slot_mappings to drafter (upstream v1 API)
                )
                self._copy_draft_token_ids_to_cpu(scheduler_output)

@@ -281,6 +281,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
            ec_connector_output,
            cudagraph_stats,
            multimodal_outputs,
+            slot_mappings,  # OMNI: pass slot_mappings for upstream v1 API compatibility
        )
        self.kv_connector_output = kv_connector_output
        return None
@@ -322,6 +323,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
            ec_connector_output,
            cudagraph_stats,
            multimodal_outputs,
+            slot_mappings,  # OMNI: unpack slot_mappings for upstream v1 API compatibility
        ) = self.execute_model_state
        self.execute_model_state = None

@@ -696,6 +698,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                    num_tokens,
                    use_cudagraphs=use_cudagraphs,
                    is_graph_capturing=is_graph_capturing,
+                    slot_mappings=slot_mappings,  # OMNI: pass slot_mappings (upstream v1 API)
                )

        # We register layerwise NVTX hooks here after the first dynamo tracing is