diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index c5d847345..0b74ba0ba 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -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// +# 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]" diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 0b484c507..9e8e8e6c3 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -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 = ( diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index dad07dc22..3199e432a 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -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 diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 2af96e718..4e903cafc 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -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 diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 9abcae94b..61aa70970 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -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) diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index 933ae7629..b86a97534 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -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), ) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index d550cf555..9d12c25e7 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -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 diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index a1e467a88..7dd3df2b0 100644 --- a/vllm_omni/engine/input_processor.py +++ b/vllm_omni/engine/input_processor.py @@ -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, diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index e39dc9fec..6a494f376 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -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 diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index b557c07dd..aadbe5756 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -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: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 97357dc3b..33355ed3f 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -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 diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index a2070f89d..8bcdd24ea 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -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: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index fb52c7e46..4ceea5516 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index cfee79157..c7854442f 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -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" diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index a4d0432c5..d993994ed 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -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 diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index d04c66331..1c200557d 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -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, ) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index bb2de4f14..58a4b25d2 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -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, diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index f54182eed..ee1944a77 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -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 ] | for transcription + # Translate this audio [from into ] | 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: """ diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 9332abeb1..7b23b3b91 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -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) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 2c82b5cd4..af69bd207 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -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