From 7e4a604d4d39024f461af0d89a2f71ad840806ef Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 9 Feb 2026 03:07:43 +0000 Subject: [PATCH] reabse 2025/2/9 Signed-off-by: tzhouam --- .../test_gpu_generation_model_runner.py | 1 + vllm_omni/config/model.py | 9 ++ .../connectors/shm_connector.py | 62 ++++++++-- vllm_omni/engine/input_processor.py | 7 +- vllm_omni/entrypoints/openai/api_server.py | 6 +- vllm_omni/entrypoints/openai/serving_chat.py | 109 ++++++++---------- vllm_omni/inputs/data.py | 4 +- vllm_omni/inputs/parse.py | 61 +++++++--- vllm_omni/outputs.py | 7 +- vllm_omni/request.py | 1 + vllm_omni/worker/gpu_ar_worker.py | 2 + .../worker/gpu_generation_model_runner.py | 81 ++++++------- vllm_omni/worker/gpu_generation_worker.py | 2 + vllm_omni/worker/gpu_model_runner.py | 2 + 14 files changed, 208 insertions(+), 146 deletions(-) diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index 25ed1ae86..f08df8754 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -24,6 +24,7 @@ def _make_runner(multimodal_outputs): None, None, multimodal_outputs, + None, ) runner.kv_connector_output = None runner.input_batch = _DummyInputBatch() diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 3199e432a..d19931172 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -65,6 +65,15 @@ class OmniModelConfig(ModelConfig): def architectures(self) -> list[str]: return [self.model_arch] + @property + def embedding_size(self): + if self.hf_config_name is not None: + stage_config = getattr(self.hf_config, self.hf_config_name, None) + override = getattr(stage_config, "embedding_size", None) + if override is not None: + return override + return super().embedding_size + def draw_hf_text_config(self): # 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, diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 9ad32e1c7..a11f38832 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -42,7 +42,13 @@ class SharedMemoryConnector(OmniConnectorBase): "inline_writes": 0, } - def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: + def put( + self, + from_stage: str, + to_stage: str, + put_key: str, + data: Any, + ) -> tuple[bool, int, dict[str, Any] | None]: try: # Always serialize first to check size (and for SHM writing) # Note: For extremely large objects in "inline" mode (e.g. Ray), @@ -51,9 +57,7 @@ class SharedMemoryConnector(OmniConnectorBase): payload = self.serialize_obj(data) size = len(payload) - metadata = {} - # if size > self.threshold: - if True: # TODO: correct put & get logic + if size > self.threshold: # Use Shared Memory lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock" with open(lock_file, "w") as lockf: @@ -62,13 +66,13 @@ class SharedMemoryConnector(OmniConnectorBase): fcntl.flock(lockf, fcntl.LOCK_UN) # meta contains {'name': ..., 'size': ...} - metadata[put_key] = {"shm": meta, "size": size} + metadata = {"shm": meta, "size": size} self._metrics["shm_writes"] += 1 else: # Inline - pass bytes directly to avoid double serialization of the object # We already serialized it to check size, so we pass the bytes. # The Queue will pickle these bytes (fast), avoiding re-serializing the complex object. - metadata[put_key] = {"inline_bytes": payload, "size": size} + metadata = {"inline_bytes": payload, "size": size} self._metrics["inline_writes"] += 1 self._metrics["puts"] += 1 @@ -80,7 +84,47 @@ class SharedMemoryConnector(OmniConnectorBase): logger.error(f"SharedMemoryConnector put failed for req {put_key}: {e}") return False, 0, None - def get(self, from_stage: str, to_stage: str, get_key: str, metadata=None) -> tuple[Any, int] | None: + def get( + self, + from_stage: str, + to_stage: str, + get_key: str, + metadata=None, + ) -> tuple[Any, int] | None: + if metadata is not None: + # Some callers may wrap metadata by request id. + if isinstance(metadata, dict) and get_key in metadata: + metadata = metadata.get(get_key) + + if not isinstance(metadata, dict): + return None + + if "inline_bytes" in metadata: + try: + obj = self.deserialize_obj(metadata["inline_bytes"]) + return obj, int(metadata.get("size", 0)) + except Exception as e: + logger.error(f"SharedMemoryConnector inline get failed for req {get_key}: {e}") + return None + + if "shm" in metadata: + try: + shm_handle = metadata["shm"] + lock_file = f"/dev/shm/shm_{shm_handle['name']}_lockfile.lock" + with open(lock_file, "w") as lockf: + fcntl.flock(lockf, fcntl.LOCK_SH) + data_bytes = shm_read_bytes(shm_handle) + fcntl.flock(lockf, fcntl.LOCK_UN) + if os.path.exists(lock_file): + os.remove(lock_file) + obj = self.deserialize_obj(data_bytes) + return obj, int(metadata.get("size", 0)) + except Exception as e: + logger.error(f"SharedMemoryConnector shm get failed for req {get_key}: {e}") + return None + + return None + from multiprocessing import shared_memory as shm_pkg # Wait for shared memory to be available (with retry logic) @@ -98,10 +142,10 @@ class SharedMemoryConnector(OmniConnectorBase): else: # Max retries reached, return None logger.warning(f"Shared memory '{get_key}' not found after {max_retries} retries") - return None, 0 + return None if shm is None: - return None, 0 + return None try: lock_file = f"/dev/shm/shm_{get_key}_lockfile.lock" diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index 7dd3df2b0..edd2bba02 100644 --- a/vllm_omni/engine/input_processor.py +++ b/vllm_omni/engine/input_processor.py @@ -14,7 +14,9 @@ from vllm.multimodal.processing.context import set_request_id from vllm.multimodal.utils import argsort_mm_positions from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams +from vllm.renderers.inputs import DictPrompt, TokPrompt from vllm.sampling_params import SamplingParams +from vllm.tasks import SupportedTask from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.torch_utils import set_default_torch_num_threads from vllm.v1.engine.input_processor import InputProcessor @@ -89,7 +91,7 @@ class OmniInputProcessor(InputProcessor): def process_inputs( self, request_id: str, - prompt: PromptType, + prompt: PromptType | DictPrompt | TokPrompt, params: SamplingParams | PoolingParams, arrival_time: float | None = None, lora_request: LoRARequest | None = None, @@ -97,6 +99,7 @@ class OmniInputProcessor(InputProcessor): trace_headers: Mapping[str, str] | None = None, priority: int = 0, data_parallel_rank: int | None = None, + supported_tasks: tuple[SupportedTask, ...] | None = None, resumable: bool = False, ) -> OmniEngineCoreRequest: """Process input prompt into an engine core request. @@ -129,7 +132,7 @@ class OmniInputProcessor(InputProcessor): has incorrect shape """ self._validate_lora(lora_request) - self._validate_params(params) + self._validate_params(params, supported_tasks) parallel_config = self.vllm_config.parallel_config dp_size = parallel_config.data_parallel_size diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 5a211b528..5901aec05 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -53,8 +53,8 @@ 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 ( +from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config +from vllm.entrypoints.openai.speech_to_text.serving import ( OpenAIServingTranscription, OpenAIServingTranslation, ) @@ -191,7 +191,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) # Load logging config for uvicorn if specified - log_config = load_log_config(getattr(args, "log_config_file", None)) + log_config = get_uvicorn_log_config(args) if log_config is not None: uvicorn_kwargs["log_config"] = log_config diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index c7854442f..da3c07027 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -3,7 +3,7 @@ import base64 import json import time import uuid -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable from datetime import datetime, timedelta, timezone from io import BytesIO from typing import TYPE_CHECKING, Any, Final, Optional, cast @@ -66,12 +66,14 @@ from vllm.entrypoints.utils import should_include_usage from vllm.inputs.data import PromptType, TokensPrompt from vllm.logger import init_logger from vllm.outputs import RequestOutput +from vllm.reasoning import ReasoningParser from vllm.renderers.hf import ( resolve_chat_template_content_format, ) from vllm.renderers.hf import ( safe_apply_chat_template as apply_hf_chat_template, ) +from vllm.renderers.inputs import TokPrompt from vllm.renderers.mistral import ( safe_apply_chat_template as apply_mistral_chat_template, ) @@ -182,6 +184,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if tokenizer is None: tokenizer = await self.engine_client.get_tokenizer() + reasoning_parser: ReasoningParser | None = None + if self.reasoning_parser_cls: + chat_template_kwargs = self._prepare_extra_chat_template_kwargs( + request.chat_template_kwargs, + self.default_chat_template_kwargs, + ) + reasoning_parser = self.reasoning_parser_cls( + tokenizer, + chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg] + ) + tool_parser = self.tool_parser if isinstance(tokenizer, MistralTokenizer): @@ -236,11 +249,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): chat_template_kwargs, self.default_chat_template_kwargs, ) - ( - conversation, - request_prompts, - engine_prompts, - ) = await self._preprocess_chat( + conversation, engine_prompts = await self._preprocess_chat( request, request.messages, default_template=request.chat_template or self.chat_template, @@ -258,7 +267,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): else: should_include_tools = tool_dicts is not None conversation, engine_prompts = self._make_request_with_harmony(request, should_include_tools) - request_prompts = [engine_prompt.get("prompt_token_ids", []) for engine_prompt in engine_prompts] except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e: logger.exception("Error in preprocessing prompt inputs") @@ -287,7 +295,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): self._log_inputs( request_id, - request_prompts[i], + engine_prompt, params_list=sampling_params_list, lora_request=lora_request, ) @@ -316,6 +324,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation, tokenizer, request_metadata, + reasoning_parser, ) try: @@ -327,6 +336,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation, tokenizer, request_metadata, + reasoning_parser, ) except ValueError as e: return self.create_error_response(e) @@ -346,11 +356,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): continue_final_message: bool = False, documents: list[dict[str, str]] | None = None, add_special_tokens: bool = False, - ) -> tuple[ - list[ConversationMessage], - Sequence[PromptType], - list[TokensPrompt], - ]: + ) -> tuple[list[ConversationMessage], list[TokPrompt]]: model_config = self.model_config # OMNI: Use self.renderer if renderer not provided (vLLM new API) if renderer is None: @@ -468,7 +474,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if hasattr(request, "cache_salt") and request.cache_salt is not None: engine_prompt["cache_salt"] = request.cache_salt - return conversation, [request_prompt], [engine_prompt] + return conversation, [engine_prompt] def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]: final_sampling_params_list = [] @@ -564,20 +570,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): def _log_inputs( self, request_id: str, - inputs: PromptType, + inputs: PromptType | TokPrompt, params_list: list[SamplingParams] | None, lora_request: LoRARequest | None, ) -> None: if self.request_logger is None: return - prompt, prompt_token_ids, prompt_embeds = None, None, None - if isinstance(inputs, str): - prompt = inputs - elif isinstance(inputs, list): - prompt_token_ids = inputs - else: - prompt = getattr(inputs, "prompt", None) - prompt_token_ids = getattr(inputs, "prompt_token_ids", None) + components = self._extract_prompt_components(inputs) + prompt = components.text + prompt_token_ids = components.token_ids + prompt_embeds = components.embeds logger.info( "Received request %s: prompt: %r, params_list: %s, prompt_token_ids: %s, prompt_embeds shape: %s, lora_request: %s.", # noqa: E501 @@ -598,6 +600,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation: list[ConversationMessage], tokenizer: AnyTokenizer, request_metadata: RequestResponseMetadata, + reasoning_parser: ReasoningParser | None = None, ): created_time = int(time.time()) chunk_object_type: Final = "chat.completion.chunk" @@ -639,7 +642,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Only one of these will be used, thus previous_texts and # all_previous_token_ids will not be used twice in the same iteration. - if tool_choice_auto or self.reasoning_parser: + if tool_choice_auto or reasoning_parser: # These are only required in "auto" tool choice case all_previous_token_ids = [[]] * num_choices # For reasoning parser and tool call all enabled @@ -647,23 +650,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): reasoning_end_arr = [False] * num_choices else: all_previous_token_ids = None - - try: - if self.reasoning_parser: - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) - reasoning_parser = self.reasoning_parser( - tokenizer, - chat_template_kwargs=chat_template_kwargs, # type: ignore - ) - except RuntimeError as e: - logger.exception("Error in reasoning parser creation.") - data = self.create_streaming_error_response(e) - yield f"data: {data}\n\n" - yield "data: [DONE]\n\n" - return # Prepare the tool parser if it's needed try: if tool_choice_auto and self.tool_parser: @@ -815,7 +801,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): delta_message: DeltaMessage | None # just update previous_texts and previous_token_ids - if tool_choice_auto or self.reasoning_parser: + if tool_choice_auto or reasoning_parser: assert previous_texts is not None assert all_previous_token_ids is not None previous_text = previous_texts[i] @@ -882,7 +868,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # handle streaming deltas for tools with named tool_choice elif tool_choice_function_name: if ( - self.reasoning_parser + reasoning_parser and not reasoning_end_arr[i] and not reasoning_parser.is_reasoning_end(previous_token_ids) ): @@ -912,7 +898,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): current_text = "" else: # Just to add remaining `content` - if self.reasoning_parser: + if reasoning_parser: delta_text = previous_text + delta_text current_text = "" @@ -948,14 +934,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): output_token_ids = as_list(output.token_ids) if ( - self.reasoning_parser is not None + reasoning_parser is not None and not reasoning_end_arr[i] and res.prompt_token_ids and reasoning_parser.is_reasoning_end(res.prompt_token_ids) ): reasoning_end_arr[i] = True - if self.reasoning_parser and not reasoning_end_arr[i]: + if reasoning_parser and not reasoning_end_arr[i]: delta_message = reasoning_parser.extract_reasoning_streaming( previous_text, current_text, @@ -994,7 +980,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # handle streaming deltas for tools with "auto" tool choice # and reasoning parser - elif tool_choice_auto and self.reasoning_parser: + elif tool_choice_auto and reasoning_parser: assert tool_parser is not None assert reasoning_parser is not None assert added_content_delta_arr is not None @@ -1075,7 +1061,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tools_streamed[i] = True # when only reasoning - elif self.reasoning_parser: + elif reasoning_parser: delta_message = reasoning_parser.extract_reasoning_streaming( previous_text, current_text, @@ -1089,7 +1075,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): delta_message = DeltaMessage(content=delta_text) # update the previous values for the next iteration - if (tool_choice_auto or self.reasoning_parser) and not self.use_harmony: + if (tool_choice_auto or reasoning_parser) and not self.use_harmony: assert previous_texts is not None assert all_previous_token_ids is not None previous_texts[i] = current_text @@ -1327,6 +1313,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation: list[ConversationMessage], tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, + reasoning_parser: ReasoningParser | None = None, ) -> ErrorResponse | ChatCompletionResponse: created_time = int(time.time()) final_res: RequestOutput | None = None @@ -1372,7 +1359,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): prompt_logprobs, prompt_token_ids, kv_transfer_params, - ) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role) + ) = self._create_text_choice( + request, + omni_outputs, + tokenizer, + conversation, + role, + reasoning_parser, + ) elif omni_outputs.final_output_type == "audio": choices_data = self._create_audio_choice(omni_outputs, role, request, stream=False) elif omni_outputs.final_output_type == "image": @@ -1432,6 +1426,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer: TokenizerLike, conversation: list[ConversationMessage], role: str, + reasoning_parser: ReasoningParser | None = None, ): final_res = omni_outputs.request_output if self.tool_call_id_type == "kimi_k2": @@ -1499,19 +1494,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): choices.append(choice_data) continue - if self.reasoning_parser: - try: - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) - reasoning_parser = self.reasoning_parser( - tokenizer, - chat_template_kwargs=chat_template_kwargs, # type: ignore - ) - except RuntimeError as e: - logger.exception("Error in reasoning parser creation.") - return self.create_error_response(e) + if reasoning_parser: # If the reasoning parser is enabled, # tool calls are extracted exclusively from the content. reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request) diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index d993994ed..7455af6ab 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -106,8 +106,8 @@ class OmniEmbedsPrompt(EmbedsPrompt): # Must ensure that all additional prompt types are inherited from vLLM prompt types # Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime. # Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate() -OmniSingletonPrompt: TypeAlias = str | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt -"""Omni singleton prompt type extending vLLM's SingletonPrompt with additional fields.""" +OmniSingletonPrompt: TypeAlias = str | list[int] | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt +"""Omni singleton prompt type extending vLLM's SingletonPrompt.""" OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt diff --git a/vllm_omni/inputs/parse.py b/vllm_omni/inputs/parse.py index 694fd44bf..8b4f28596 100644 --- a/vllm_omni/inputs/parse.py +++ b/vllm_omni/inputs/parse.py @@ -1,11 +1,29 @@ -from vllm.inputs.parse import ( - ParsedEmbedsPrompt, - ParsedSingletonPrompt, - ParsedStrPrompt, - ParsedTextPrompt, - ParsedTokensPrompt, - SingletonPrompt, -) +from typing import Literal, TypedDict + +from vllm.inputs import EmbedsPrompt, SingletonPrompt, TextPrompt, TokensPrompt + + +class ParsedStrPrompt(TypedDict): + type: Literal["str"] + content: str + + +class ParsedTextPrompt(TypedDict): + type: Literal["text"] + content: TextPrompt + + +class ParsedTokensPrompt(TypedDict): + type: Literal["tokens"] + content: TokensPrompt + + +class ParsedEmbedsPrompt(TypedDict): + type: Literal["embeds"] + content: EmbedsPrompt + + +ParsedSingletonPrompt = ParsedStrPrompt | ParsedTextPrompt | ParsedTokensPrompt | ParsedEmbedsPrompt def parse_singleton_prompt_omni( @@ -29,16 +47,21 @@ def parse_singleton_prompt_omni( TypeError: If the prompt type is not supported """ if isinstance(prompt, str): - return ParsedStrPrompt(type="str", content=prompt) - elif isinstance(prompt, dict): - # Type ignores are because mypy does not correctly infer the TypedDicts - # Pyright does succeed. + return {"type": "str", "content": prompt} + if isinstance(prompt, list): + if not all(isinstance(token, int) for token in prompt): + raise TypeError("Token prompt should be a list of integers") + return { + "type": "tokens", + "content": TokensPrompt(prompt_token_ids=prompt), + } + if isinstance(prompt, dict): # Priority tokens: When both tokens and embeds exist, keep both and - # follow the tokens path + # follow the tokens path. if "prompt_token_ids" in prompt: - return ParsedTokensPrompt(type="tokens", content=prompt) # type: ignore[typeddict-item] - elif "prompt_embeds" in prompt: - return ParsedEmbedsPrompt(type="embeds", content=prompt) # type: ignore[typeddict-item] - elif "prompt" in prompt: - return ParsedTextPrompt(type="text", content=prompt) - raise TypeError("inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt") + return {"type": "tokens", "content": prompt} + if "prompt_embeds" in prompt: + return {"type": "embeds", "content": prompt} + if "prompt" in prompt: + return {"type": "text", "content": prompt} + raise TypeError("inputs must be a string, list of token IDs, TextPrompt, TokensPrompt, or EmbedsPrompt") diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 64569846a..3708bc8f9 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -128,9 +128,10 @@ class OmniRequestOutput: For diffusion outputs, this returns the local _multimodal_output field. """ if self.request_output is not None: - # Check completion outputs first (where multimodal_output is attached) - if self.request_output.outputs: - for output in self.request_output.outputs: + # Check completion outputs first (where multimodal_output is attached). + outputs = getattr(self.request_output, "outputs", None) + if isinstance(outputs, list) and outputs: + for output in outputs: mm = getattr(output, "multimodal_output", None) if mm: return mm diff --git a/vllm_omni/request.py b/vllm_omni/request.py index c992816de..a869fa23c 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -91,4 +91,5 @@ class OmniRequest(Request): block_hasher=block_hasher, additional_information=request.additional_information, resumable=request.resumable, + reasoning_ended=request.reasoning_ended, ) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 30edba643..198d42c8a 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -4,6 +4,7 @@ import os import torch from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.tracing import instrument from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats @@ -25,6 +26,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): model runners for text generation stages (e.g., thinker stages). """ + @instrument(span_name="Init device") def init_device(self): if self.device_config.device_type == "cuda": # This env var set by Ray causes exceptions with graph building. diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index af69bd207..80999e9b7 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -735,55 +735,46 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): assert mm_budget is not None if (encoder_budget := mm_budget.get_encoder_budget()) > 0: - # NOTE: Currently model is profiled with a single non-text - # modality with the max possible input tokens even when - # it supports multiple. - dummy_modality = mm_budget.get_modality_with_max_tokens() - max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[dummy_modality] + if not mm_budget.mm_max_toks_per_item: + # All modality limits are 0 — embedding-only mode. + # Budget is non-zero for embedding storage, but + # there's no encoder to profile. + logger.info( + "Skipping encoder profiling for embedding-only " + "mode (all modality limits=0 with " + "enable_mm_embeds=True).", + ) + else: + # NOTE: Currently model is profiled with a single non-text + # modality with the max possible input tokens even when + # it supports multiple. + dummy_modality = mm_budget.get_modality_with_max_tokens() + max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[dummy_modality] - logger.info( - "Encoder cache will be initialized with a budget of " - "%s tokens, and profiled with %s %s items of the " - "maximum feature size.", - encoder_budget, - max_mm_items_per_batch, - dummy_modality, - ) + logger.info( + "Encoder cache will be initialized with a budget of " + "%s tokens, and profiled with %s %s items of the " + "maximum feature size.", + encoder_budget, + max_mm_items_per_batch, + dummy_modality, + ) - # Create dummy batch of multimodal inputs. - batched_dummy_mm_inputs = self._get_mm_dummy_batch( - dummy_modality, - max_mm_items_per_batch, - ) + # Create dummy batch of multimodal inputs. + batched_dummy_mm_inputs = self._get_mm_dummy_batch( + dummy_modality, + max_mm_items_per_batch, + ) - # Run multimodal encoder. - dummy_encoder_outputs = self.model.embed_multimodal(**batched_dummy_mm_inputs) + # Run multimodal encoder. + dummy_encoder_outputs = self.model.embed_multimodal(**batched_dummy_mm_inputs) - sanity_check_mm_encoder_outputs( - dummy_encoder_outputs, - expected_num_items=max_mm_items_per_batch, - ) - - # NOTE: This happens when encoder cache needs to store - # the embeddings that encoder outputs are scattered onto. - # In this case we create dummy embeddings of size - # (max_tokens_for_modality, hidden_size) and scatter - # encoder output into it. - encoder_output_shape = dummy_encoder_outputs[0].shape - max_mm_tokens_per_item = mm_budget.max_tokens_by_modality[dummy_modality] - if encoder_output_shape[0] < max_mm_tokens_per_item: - encoder_hidden_size = encoder_output_shape[-1] - expanded_outputs = [] - for output in dummy_encoder_outputs: - expanded = output.new_zeros((max_mm_tokens_per_item, encoder_hidden_size)) - num_tokens = output.shape[0] - expanded[:num_tokens].copy_(output) - expanded_outputs.append(expanded) - - dummy_encoder_outputs = expanded_outputs - - # Cache the dummy encoder outputs. - self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs)) + sanity_check_mm_encoder_outputs( + dummy_encoder_outputs, + expected_num_items=max_mm_items_per_batch, + ) + for i, output in enumerate(dummy_encoder_outputs): + self.encoder_cache[f"tmp_{i}"] = output # Add `is_profile` here to pre-allocate communication buffers hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True) diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 8ba4efa3d..7de9a7922 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -4,6 +4,7 @@ import os import torch from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.tracing import instrument from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats @@ -25,6 +26,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): worker_cls: "vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner" """ + @instrument(span_name="Init device") def init_device(self): if self.device_config.device_type == "cuda": # This env var set by Ray causes exceptions with graph building. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index d52a2014f..5735ef1e9 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.models.interfaces import supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.sampling_params import SamplingType +from vllm.tracing import instrument from vllm.utils.import_utils import LazyLoader from vllm.utils.math_utils import cdiv from vllm.v1.spec_decode.eagle import EagleProposer @@ -40,6 +41,7 @@ class OmniGPUModelRunner(GPUModelRunner): self._omni_num_scheduled_tokens_np: np.ndarray | None = None self._omni_last_model_output: object | None = None + @instrument(span_name="Loading (GPU)") def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) # TODO move this model specific logic to a separate class