diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 406c26b39..f0a906709 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -1,9 +1,55 @@ """ Engine components for vLLM-omni. """ +import msgspec +from typing import Any -from .output_processor import MultimodalOutputProcessor +from vllm.v1.engine import EngineCoreRequest +from typing import Optional -__all__ = [ - "MultimodalOutputProcessor", -] +class PromptEmbedsPayload(msgspec.Struct): + """Serialized prompt embeddings payload for direct transfer. + + data: raw bytes of the tensor in row-major order + shape: [seq_len, hidden_size] + dtype: torch dtype name (e.g., "float16", "float32") + """ + + data: bytes + shape: list[int] + dtype: str + + +class AdditionalInformationEntry(msgspec.Struct): + """One entry of additional_information. + + Two supported forms are encoded: + - tensor: data/shape/dtype + - list: a Python list (msgspec-serializable) + Exactly one of (tensor_data, list_data) should be non-None. + """ + + # Tensor form + tensor_data: Optional[bytes] = None + tensor_shape: Optional[list[int]] = None + tensor_dtype: Optional[str] = None + + # List form + list_data: Optional[list[Any]] = None + + +class AdditionalInformationPayload(msgspec.Struct): + """Serialized dictionary payload for additional_information. + + Keys are strings; values are encoded as AdditionalInformationEntry. + """ + + entries: dict[str, AdditionalInformationEntry] + + +class OmniEngineCoreRequest(EngineCoreRequest): + # Optional prompt embeddings (direct-transfer version) + prompt_embeds: Optional[PromptEmbedsPayload] = None + # Optional additional information dictionary (serialized) + additional_information: Optional[AdditionalInformationPayload] = None + \ No newline at end of file diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py new file mode 100644 index 000000000..52253c118 --- /dev/null +++ b/vllm_omni/engine/arg_utils.py @@ -0,0 +1,50 @@ +from vllm.engine.arg_utils import EngineArgs +from typing import Literal, Optional +from dataclasses import dataclass +from vllm_omni.config import OmniModelConfig +from vllm.utils import FlexibleArgumentParser + + +@dataclass +class OmniEngineArgs(EngineArgs): + stage_id: int = 0 + model_stage: str = "thinker" + model_arch: str = "Qwen2_5OmniForConditionalGeneration" + engine_output_type: Optional[str] = None + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + """Shared CLI arguments for vLLM engine.""" + parser.add_argument( + "--engine-output-type", + type=str, + default=EngineArgs.engine_output_type, + help=( + "Declare EngineCoreOutput.output_type (e.g., 'text', 'image', " + "'text+image', 'latent'). This will be written into " + "model_config.engine_output_type for schedulers to use." + ), + ) + parser.add_argument("--model-stage", type=str, default=EngineArgs.model_stage, + help="Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). This will be written into model_config.model_stage for schedulers to use.") + return parser + + def create_model_config(self) -> OmniModelConfig: + # 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['model_stage'] = self.model_stage + config_dict['model_arch'] = self.model_arch + config_dict['engine_output_type'] = self.engine_output_type + + # Create and return the OmniModelConfig instance + omni_config = OmniModelConfig(**config_dict) + omni_config.hf_config.architectures = omni_config.architectures + + return omni_config \ No newline at end of file diff --git a/vllm_omni/engine/diffusion_engine.py b/vllm_omni/engine/diffusion_engine.py deleted file mode 100644 index ef8a5b51d..000000000 --- a/vllm_omni/engine/diffusion_engine.py +++ /dev/null @@ -1,133 +0,0 @@ -"""Diffusion engine adapters for diffusers-backed DiT stages.""" - -from __future__ import annotations - -from typing import Any, Optional - -from ..config import DiTConfig -from ..worker.gpu_diffusion_worker import ( - DiffusionGPUWorker, - DiffusionRunnerOutput, -) - - -class DiffusionEngine: - """Base diffusion engine storing shared DiT configuration.""" - - def __init__( - self, - dit_config: DiTConfig, - model_path: Optional[str] = None, - log_stats: bool = False, - multiprocess_mode: bool = False, - ) -> None: - self.dit_config = dit_config - - model_cfg = getattr(dit_config, "model_config", None) - config_model_path = None - if model_cfg is not None: - if isinstance(model_cfg, dict): - config_model_path = ( - model_cfg.get("model") - or model_cfg.get("model_path") - ) - else: - config_model_path = ( - getattr(model_cfg, "model", None) - or getattr(model_cfg, "model_path", None) - ) - - self.model_path = model_path or config_model_path - self.model_config = dit_config.model_config - self.cache_config = dit_config.dit_cache_config - - self.height = dit_config.height - self.width = dit_config.width - self.num_inference_steps = dit_config.num_inference_steps - self.guidance_scale = dit_config.guidance_scale - - self.log_stats = log_stats - self.multiprocess_mode = multiprocess_mode - - def generate(self, *args, **kwargs): # pragma: no cover - placeholder - raise NotImplementedError("This method should be implemented in subclasses.") - - -class DiffusersPipelineEngine(DiffusionEngine): - """Adapter that invokes a diffusers pipeline via DiffusionGPUWorker.""" - - def __init__( - self, - dit_config: DiTConfig, - model_path: Optional[str] = None, - log_stats: bool = False, - multiprocess_mode: bool = False, - ) -> None: - super().__init__(dit_config, model_path, log_stats, multiprocess_mode) - - resolved_model_path = self.model_path - if resolved_model_path is None: - raise ValueError("Model path must be provided for DiffusersPipelineEngine.") - - device_cfg = getattr(dit_config, "device_config", None) - model_cfg = getattr(dit_config, "model_config", None) - - device = None - dtype = None - - if device_cfg: - if isinstance(device_cfg, dict): - device = device_cfg.get("device") - dtype = device_cfg.get("dtype") - else: - device = getattr(device_cfg, "device", None) - dtype = getattr(device_cfg, "dtype", None) - - if dtype is None and model_cfg: - if isinstance(model_cfg, dict): - dtype = model_cfg.get("dtype") - else: - dtype = getattr(model_cfg, "dtype", None) - - self.worker = DiffusionGPUWorker( - resolved_model_path, - pipeline_name=dit_config.diffusers_pipeline, - device=device, - dtype=dtype, - ) - - def generate( - self, - prompt: str, - *, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: Optional[int] = None, - guidance_scale: Optional[float] = None, - negative_prompt: Optional[str] = None, - seed: Optional[int] = None, - image: Optional[Any] = None, - ) -> DiffusionRunnerOutput: - height = int(height) if height is not None else self.height - width = int(width) if width is not None else self.width - num_steps = ( - int(num_inference_steps) - if num_inference_steps is not None - else self.num_inference_steps - ) - guidance = ( - float(guidance_scale) - if guidance_scale is not None - else self.guidance_scale - ) - - return self.worker.generate( - prompt=prompt, - height=height, - width=width, - num_inference_steps=num_steps, - guidance_scale=guidance, - negative_prompt=negative_prompt, - seed=seed, - image=image, - ) diff --git a/vllm_omni/engine/processor.py b/vllm_omni/engine/processor.py index e69de29bb..1c5891265 100644 --- a/vllm_omni/engine/processor.py +++ b/vllm_omni/engine/processor.py @@ -0,0 +1,216 @@ +import time +from collections.abc import Mapping, Sequence +from typing import Any, Optional, Union + +from vllm.inputs import ProcessorInputs, PromptType +from vllm.inputs.parse import split_enc_dec_inputs +from vllm.lora.request import LoRARequest +from vllm.multimodal import MultiModalKwargs +from vllm.multimodal.inputs import PlaceholderRange +from vllm.multimodal.utils import merge_and_sort_multimodal_metadata +from vllm.pooling_params import PoolingParams +from vllm.sampling_params import SamplingParams +from vllm.config import VllmConfig +from vllm.transformers_utils.tokenizer_group import TokenizerGroup +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm_omni.inputs.preprocess import OmniInputPreprocessor + +from vllm.v1.engine.processor import Processor +from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload, AdditionalInformationEntry, OmniEngineCoreRequest + + +class OmniProcessor(Processor): + def __init__(self, + vllm_config: VllmConfig, + tokenizer: TokenizerGroup, + mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, + ): + super().__init__(vllm_config, tokenizer, mm_registry) + self.input_preprocessor = OmniInputPreprocessor(self.model_config, + self.tokenizer, + mm_registry) + + def process_inputs( + self, + request_id: str, + prompt: PromptType, + params: Union[SamplingParams, PoolingParams], + arrival_time: Optional[float] = None, + lora_request: Optional[LoRARequest] = None, + tokenization_kwargs: Optional[dict[str, Any]] = None, + trace_headers: Optional[Mapping[str, str]] = None, + priority: int = 0, + data_parallel_rank: Optional[int] = None, + ) -> tuple[Optional[str], OmniEngineCoreRequest]: + + # TODO(woosuk): Support pooling models. + # TODO(woosuk): Support encoder-decoder models. + self._validate_lora(lora_request) + self._validate_params(params, lora_request) + if trace_headers is not None: + raise ValueError("V1 does not support tracing yet.") + + data_parallel_size = self.vllm_config.parallel_config.data_parallel_size + if data_parallel_rank is not None and not (0 <= data_parallel_rank < + data_parallel_size): + raise ValueError(f"data_parallel_rank {data_parallel_rank} " + f"is out of range [0, {data_parallel_size}).") + + if arrival_time is None: + arrival_time = time.time() + + # Process inputs, which includes: + # 1. Tokenize text prompt, with LoRA request if one exists. + # 2. For multimodal models with a merged preprocessor, preprocess + # multimodal data and expand prompt token ids accordingly. + return_mm_hashes = (self.model_config.processor_return_mm_hashes + or bool(self.cache_config.enable_prefix_caching)) + processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess( + prompt, + tokenization_kwargs=tokenization_kwargs, + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + from vllm.platforms import current_platform + current_platform.validate_request( + prompt=prompt, + params=params, + processed_inputs=processed_inputs, + ) + eos_token_id = self.input_preprocessor.get_eos_token_id(lora_request) + + self._validate_model_inputs(processed_inputs, lora_request) + + encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs) + + # TODO: Impl encoder-decoder + if encoder_inputs is not None: + raise NotImplementedError + + sampling_params = None + pooling_params = None + if isinstance(params, SamplingParams): + # TODO: can we avoid cloning here in multiproc case? + sampling_params = params.clone() + # If unset max tokens, then generate up to the max_model_len. + if sampling_params.max_tokens is None: + sampling_params.max_tokens = ( + self.model_config.max_model_len - + len(decoder_inputs["prompt_token_ids"])) + sampling_params.update_from_generation_config( + self.generation_config_fields, eos_token_id) + if self.tokenizer is not None: + sampling_params.update_from_tokenizer( + self.tokenizer.get_lora_tokenizer(lora_request)) + else: + pooling_params = params.clone() + + # Multimodal related. + sorted_mm_inputs: Optional[Sequence[Optional[MultiModalKwargs]]] = None + sorted_mm_positions: Optional[list[PlaceholderRange]] = None + sorted_mm_hashes: Optional[list[str]] = None + if decoder_inputs["type"] == "multimodal": + decoder_mm_inputs = decoder_inputs["mm_kwargs"] + + # Merge and flatten multimodal placeholders, hashes and inputs + # from dictionaries to lists, and sort them by each item's position + # in the input sequence. + ( + sorted_item_modalities, + sorted_mm_positions, + sorted_mm_hashes, + ) = merge_and_sort_multimodal_metadata( + decoder_inputs["mm_placeholders"], + decoder_inputs["mm_hashes"] if return_mm_hashes else None, + ) + + # The output of merged multi-modal processor (`decoder_mm_inputs`) + # is a single MultiModalKwargs for all items from all modalities. + # This code flattens kwargs for individual items in a list and + # sorts them by each item's position in the input sequence if there + # are multiple modalities. + unique_modalities = set(sorted_item_modalities) + if len(unique_modalities) > 1: + orig_sorted_mm_inputs = [] + used_indices = {modality: 0 for modality in unique_modalities} + + for modality in sorted_item_modalities: + items = decoder_mm_inputs.get_items(modality) + item = items[used_indices[modality]] + + orig_sorted_mm_inputs.append( + MultiModalKwargs.from_items([item])) + used_indices[modality] += 1 + else: + orig_sorted_mm_inputs = [ + MultiModalKwargs.from_items([item]) for item in + decoder_mm_inputs.get_items(sorted_item_modalities[0]) + ] + + if sorted_mm_hashes is not None: + sorted_mm_inputs = self.mm_input_cache_client.get_and_update( + orig_sorted_mm_inputs, sorted_mm_hashes) + else: + sorted_mm_inputs = orig_sorted_mm_inputs + + # Serialize prompt_embeds and additional_information if provided (direct-transfer path) + prompt_embeds_payload: Optional[PromptEmbedsPayload] = None + additional_information_payload: Optional[AdditionalInformationPayload] = None + if "prompt_embeds" in decoder_inputs: # type: ignore[operator] + import numpy as np + import torch + pe: torch.Tensor = decoder_inputs["prompt_embeds"] # type: ignore[index] + if pe.ndim != 2: + raise ValueError( + "prompt_embeds must be of shape (seq_len, hidden_size)") + # Move to CPU and ensure contiguous memory for stable serialization + pe_cpu = pe.detach().to("cpu").contiguous() + seq_len, hidden_size = pe_cpu.shape + dtype_str = str(pe_cpu.dtype).replace("torch.", "") + data_bytes = pe_cpu.numpy().tobytes() + prompt_embeds_payload = PromptEmbedsPayload( + data=data_bytes, + shape=[int(seq_len), int(hidden_size)], + dtype=dtype_str, + ) + if "additional_information" in decoder_inputs: # type: ignore[operator] + import numpy as np + import torch + raw_info: dict[str, Any] = decoder_inputs["additional_information"] # type: ignore[index] + entries: dict[str, AdditionalInformationEntry] = {} + for key, value in raw_info.items(): + if isinstance(value, torch.Tensor): + v_cpu = value.detach().to("cpu").contiguous() + dtype_str = str(v_cpu.dtype).replace("torch.", "") + data_bytes = v_cpu.numpy().tobytes() + entry = AdditionalInformationEntry( + tensor_data=data_bytes, + tensor_shape=[int(x) for x in list(v_cpu.shape)], + tensor_dtype=dtype_str, + ) + elif isinstance(value, list): + entry = AdditionalInformationEntry(list_data=value) + else: + raise ValueError( + "additional_information values must be Tensor or list") + entries[key] = entry + additional_information_payload = AdditionalInformationPayload( + entries=entries) + + return decoder_inputs.get("prompt"), OmniEngineCoreRequest( + request_id=request_id, + prompt_token_ids=decoder_inputs["prompt_token_ids"], + mm_inputs=sorted_mm_inputs, + mm_hashes=sorted_mm_hashes, + mm_placeholders=sorted_mm_positions, + sampling_params=sampling_params, + pooling_params=pooling_params, + eos_token_id=eos_token_id, + arrival_time=arrival_time, + lora_request=lora_request, + cache_salt=decoder_inputs.get("cache_salt"), + priority=priority, + data_parallel_rank=data_parallel_rank, + prompt_embeds=prompt_embeds_payload, + additional_information=additional_information_payload, + ) \ No newline at end of file diff --git a/vllm_omni/inputs/__init__.py b/vllm_omni/inputs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py new file mode 100644 index 000000000..5682b9dab --- /dev/null +++ b/vllm_omni/inputs/data.py @@ -0,0 +1,56 @@ +from vllm.inputs.data import TokensPrompt, EmbedsPrompt, TokenInputs +from typing import Any, NotRequired, Optional +import torch + + +class OmniTokensPrompt(TokensPrompt): + prompt_embeds: NotRequired[torch.Tensor] + """The embeddings of the prompt.""" + + # New: optional additional information dictionary + # Values may be torch.Tensor or list + additional_information: NotRequired[dict[str, Any]] + + +class OmniTokenInputs(TokenInputs): + # New: optional prompt embeddings aligned with token ids + prompt_embeds: NotRequired[torch.Tensor] + + # New: optional additional information dictionary + # Values may be torch.Tensor or list + additional_information: NotRequired[dict[str, Any]] + + +class OmniEmbedsPrompt(EmbedsPrompt): + # New: optional prompt embeddings aligned with token ids + prompt_embeds: NotRequired[torch.Tensor] + + # New: optional additional information dictionary + # Values may be torch.Tensor or list + additional_information: NotRequired[dict[str, Any]] + + +def token_inputs_omni( + prompt_token_ids: list[int], + token_type_ids: Optional[list[int]] = None, + prompt: Optional[str] = None, + cache_salt: Optional[str] = None, + prompt_embeds: Optional[torch.Tensor] = None, + additional_information: Optional[dict[str, Any]] = None, +) -> OmniTokenInputs: + """Construct [`TokenInputs`][vllm.inputs.data.TokenInputs] from optional + values.""" + inputs = OmniTokenInputs(type="token", prompt_token_ids=prompt_token_ids) + + if prompt is not None: + inputs["prompt"] = prompt + if token_type_ids is not None: + inputs["token_type_ids"] = token_type_ids + if cache_salt is not None: + inputs["cache_salt"] = cache_salt + if prompt_embeds is not None: + inputs["prompt_embeds"] = prompt_embeds + if additional_information is not None: + inputs["additional_information"] = additional_information + + return inputs \ No newline at end of file diff --git a/vllm_omni/inputs/parse.py b/vllm_omni/inputs/parse.py new file mode 100644 index 000000000..a7c72b9a9 --- /dev/null +++ b/vllm_omni/inputs/parse.py @@ -0,0 +1,21 @@ +from vllm.inputs.parse import ParsedSingletonPrompt, ParsedStrPrompt, ParsedEmbedsPrompt, ParsedTokensPrompt, ParsedTextPrompt +from vllm.inputs.data import SingletonPrompt + + +def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPrompt: + 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. + # 优先 tokens：当 tokens 与 embeds 同在时，保留两者并走 tokens 路径 + 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") \ No newline at end of file diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py new file mode 100644 index 000000000..88e97b58a --- /dev/null +++ b/vllm_omni/inputs/preprocess.py @@ -0,0 +1,166 @@ +from typing import Optional, Any, Union +from typing_extensions import assert_never + +from vllm.lora.request import LoRARequest +from vllm.inputs.preprocess import InputPreprocessor +from vllm.inputs.data import TokensPrompt, SingletonPrompt, SingletonInputs, TextPrompt +from vllm.multimodal.inputs import MultiModalInputs +from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni +from vllm_omni.inputs.parse import parse_singleton_prompt_omni + + +class OmniInputPreprocessor(InputPreprocessor): + def _process_tokens( + self, + parsed_content: TokensPrompt, + tokenization_kwargs: Optional[dict[str, Any]] = None, + lora_request: Optional[LoRARequest] = None, + return_mm_hashes: bool = False, + ) -> Union[OmniTokenInputs, MultiModalInputs]: + prompt_token_ids = parsed_content["prompt_token_ids"] + token_type_ids = parsed_content.get("token_type_ids") + prompt_embeds = parsed_content.get("prompt_embeds") + additional_information = parsed_content.get("additional_information") + + inputs: Union[OmniTokenInputs, MultiModalInputs] + if multi_modal_data := parsed_content.get("multi_modal_data"): + inputs = self._process_multimodal( + prompt_token_ids, + multi_modal_data, + parsed_content.get("mm_processor_kwargs"), + tokenization_kwargs=tokenization_kwargs, + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + else: + inputs = token_inputs_omni( + prompt_token_ids=prompt_token_ids, + token_type_ids=token_type_ids, + prompt_embeds=prompt_embeds, + additional_information=additional_information, + ) + + if cache_salt := parsed_content.get("cache_salt"): + inputs["cache_salt"] = cache_salt + + return inputs + + def _prompt_to_llm_inputs( + self, + prompt: SingletonPrompt, + tokenization_kwargs: Optional[dict[str, Any]] = None, + lora_request: Optional[LoRARequest] = None, + return_mm_hashes: bool = False, + ) -> SingletonInputs: + """ + Extract the singleton inputs from a prompt. + + 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: + + * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance + """ + parsed = parse_singleton_prompt_omni(prompt) + + if parsed["type"] == "embeds": + return self._process_embeds(parsed["content"]) + if parsed["type"] == "tokens": + return self._process_tokens( + parsed["content"], + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + if parsed["type"] == "text": + return self._process_text( + parsed["content"], + tokenization_kwargs=tokenization_kwargs, + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + if parsed["type"] == "str": + return self._process_text( + TextPrompt(prompt=parsed["content"]), + tokenization_kwargs=tokenization_kwargs, + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + + assert_never(parsed) + + async def _process_tokens_async( + self, + parsed_content: TokensPrompt, + tokenization_kwargs: Optional[dict[str, Any]] = None, + lora_request: Optional[LoRARequest] = None, + return_mm_hashes: bool = False, + ) -> Union[OmniTokenInputs, MultiModalInputs]: + prompt_token_ids = parsed_content["prompt_token_ids"] + token_type_ids = parsed_content.get("token_type_ids") + prompt_embeds = parsed_content.get("prompt_embeds") + additional_information = parsed_content.get("additional_information") + + inputs: Union[OmniTokenInputs, MultiModalInputs] + if multi_modal_data := parsed_content.get("multi_modal_data"): + inputs = await self._process_multimodal_async( + prompt_token_ids, + multi_modal_data, + parsed_content.get("mm_processor_kwargs"), + tokenization_kwargs=tokenization_kwargs, + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + else: + inputs = token_inputs_omni( + prompt_token_ids=prompt_token_ids, + token_type_ids=token_type_ids, + prompt_embeds=prompt_embeds, + additional_information=additional_information, + ) + + if cache_salt := parsed_content.get("cache_salt"): + inputs["cache_salt"] = cache_salt + + return inputs + + async def _prompt_to_llm_inputs_async( + self, + prompt: SingletonPrompt, + tokenization_kwargs: Optional[dict[str, Any]] = None, + lora_request: Optional[LoRARequest] = None, + return_mm_hashes: bool = False, + ) -> SingletonInputs: + """ + Async version of + [`_prompt_to_llm_inputs`][vllm.inputs.preprocess.InputPreprocessor._prompt_to_llm_inputs]. + """ + parsed = parse_singleton_prompt_omni(prompt) + + if parsed["type"] == "embeds": + return await self._process_embeds_async(parsed["content"]) + if parsed["type"] == "tokens": + return await self._process_tokens_async( + parsed["content"], + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + if parsed["type"] == "text": + return await self._process_text_async( + parsed["content"], + tokenization_kwargs=tokenization_kwargs, + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + if parsed["type"] == "str": + return await self._process_text_async( + TextPrompt(prompt=parsed["content"]), + tokenization_kwargs=tokenization_kwargs, + lora_request=lora_request, + return_mm_hashes=return_mm_hashes, + ) + + assert_never(parsed) \ No newline at end of file diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py new file mode 100644 index 000000000..5d2891847 --- /dev/null +++ b/vllm_omni/outputs.py @@ -0,0 +1,15 @@ +from vllm.v1.outputs import ModelRunnerOutput +from vllm.outputs import RequestOutput +from typing import Optional +from dataclasses import dataclass +import torch + + +class OmniModelRunnerOutput(ModelRunnerOutput): + multimodal_outputs: Optional[dict[str, torch.Tensor]] = None + +@dataclass +class OmniRequestOutput(RequestOutput): + stage_id: int + final_output_type: str + request_output: RequestOutput \ No newline at end of file diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py new file mode 100644 index 000000000..783796ca7 --- /dev/null +++ b/vllm_omni/patch.py @@ -0,0 +1,21 @@ +import sys + +from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt +from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding as _OriginalMRotaryEmbedding +from vllm.v1.request import Request as _OriginalRequest +from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest + +from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding +from vllm_omni.request import OmniRequest +from vllm_omni.engine import OmniEngineCoreRequest + +for module_name, module in sys.modules.items(): + if hasattr(module, 'TokensPrompt') and module.TokensPrompt == _OriginalTokensPrompt: + module.TokensPrompt = OmniTokensPrompt + if hasattr(module, 'MRotaryEmbedding') and module.MRotaryEmbedding == _OriginalMRotaryEmbedding: + module.MRotaryEmbedding = MRotaryEmbedding + if hasattr(module, 'Request') and module.Request == _OriginalRequest: + module.Request = OmniRequest + if hasattr(module, 'EngineCoreRequest') and module.EngineCoreRequest == _OriginalEngineCoreRequest: + module.EngineCoreRequest = OmniEngineCoreRequest \ No newline at end of file diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 4bdc62b25..b0da21e63 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -1,299 +1,47 @@ -""" -OmniRequest: Extended request class for vLLM-omni multimodal processing. - -This class extends vLLM's Request to support multimodal and non-autoregressive -processing with additional fields and methods specific to vLLM-omni. -""" - -import enum -import time -from typing import Optional, Dict, Any, List, Union -from dataclasses import dataclass, field -from vllm.v1.request import Request as vLLMRequest +from vllm.v1.request import Request +from vllm.v1.structured_output.request import StructuredOutputRequest +from vllm.multimodal.inputs import MultiModalKwargs +from vllm.utils import is_list_of +from typing import Optional +from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload, OmniEngineCoreRequest -class RequestType(enum.Enum): - """Types of requests supported by vLLM-omni.""" - TEXT = "text" - IMAGE = "image" - AUDIO = "audio" - VIDEO = "video" - MULTIMODAL = "multimodal" - DIFFUSION = "diffusion" - - -class ProcessingStage(enum.Enum): - """Processing stages for multi-stage models.""" - PREPROCESSING = "preprocessing" - AR_GENERATION = "ar_generation" - DIFFUSION_GENERATION = "diffusion_generation" - POSTPROCESSING = "postprocessing" - COMPLETED = "completed" - - -@dataclass -class MultimodalData: - """Container for multimodal input data.""" - data_type: str # "image", "audio", "video", etc. - data: Any # The actual data (numpy array, bytes, etc.) - metadata: Dict[str, Any] = field(default_factory=dict) - processed: bool = False - - -@dataclass -class DiffusionParams: - """Parameters specific to diffusion model generation.""" - num_inference_steps: int = 50 - guidance_scale: float = 7.5 - seed: Optional[int] = None - scheduler: str = "ddpm" - strength: float = 1.0 - eta: float = 0.0 - - -class OmniRequest(vLLMRequest): - """ - Extended request class for vLLM-omni multimodal and non-autoregressive processing. - - This class extends vLLM's Request with additional fields and methods to support: - - Multimodal input processing - - Non-autoregressive generation (diffusion models) - - Multi-stage processing pipelines - - Enhanced caching for different model types - """ - +class OmniRequest(Request): def __init__( - self, - request_id: str, - prompt: Optional[str] = None, - prompt_token_ids: Optional[List[int]] = None, - sampling_params: Optional[Any] = None, - arrival_time: Optional[float] = None, - lora_request: Optional[Any] = None, - multi_modal_data: Optional[Dict[str, Any]] = None, - multi_modal_placeholders: Optional[Dict[str, str]] = None, - priority: int = 0, - # vLLM-omni specific parameters - request_type: RequestType = RequestType.TEXT, - processing_stage: ProcessingStage = ProcessingStage.PREPROCESSING, - multimodal_inputs: Optional[List[MultimodalData]] = None, - diffusion_params: Optional[DiffusionParams] = None, - output_format: str = "text", - cache_key: Optional[str] = None, - stage_configs: Optional[Dict[str, Any]] = None, - **kwargs - ): - # Initialize parent class - super().__init__( - request_id=request_id, - prompt=prompt, - prompt_token_ids=prompt_token_ids, - sampling_params=sampling_params, - arrival_time=arrival_time or time.time(), - lora_request=lora_request, - multi_modal_data=multi_modal_data, - multi_modal_placeholders=multi_modal_placeholders, - priority=priority, - **kwargs - ) - - # vLLM-omni specific attributes - self.request_type = request_type - self.processing_stage = processing_stage - self.multimodal_inputs = multimodal_inputs or [] - self.diffusion_params = diffusion_params or DiffusionParams() - self.output_format = output_format - self.cache_key = cache_key - self.stage_configs = stage_configs or {} - - # Processing state - self.current_stage = 0 - self.stage_results = {} - self.hidden_states = None - self.intermediate_outputs = [] - - # Timing and metrics - self.stage_timings = {} - self.total_processing_time = 0.0 - - # Error handling - self.errors = [] - self.retry_count = 0 - self.max_retries = 3 + self, + prompt_embeds: Optional[PromptEmbedsPayload] = None, + additional_information: Optional[AdditionalInformationPayload] = None, + *args, **kwargs): + super().__init__(*args, **kwargs) + # Serialized prompt embeddings payload (optional) + self.prompt_embeds: Optional[PromptEmbedsPayload] = prompt_embeds + # Serialized additional information payload (optional) + self.additional_information: Optional[AdditionalInformationPayload] = additional_information - def add_multimodal_input(self, data: MultimodalData) -> None: - """Add a multimodal input to the request.""" - self.multimodal_inputs.append(data) - - def get_multimodal_inputs_by_type(self, data_type: str) -> List[MultimodalData]: - """Get all multimodal inputs of a specific type.""" - return [inp for inp in self.multimodal_inputs if inp.data_type == data_type] - - def update_processing_stage(self, stage: ProcessingStage) -> None: - """Update the current processing stage.""" - self.processing_stage = stage - self.stage_timings[stage.value] = time.time() - - def add_stage_result(self, stage: str, result: Any) -> None: - """Add a result from a processing stage.""" - self.stage_results[stage] = result - self.intermediate_outputs.append({ - 'stage': stage, - 'result': result, - 'timestamp': time.time() - }) - - def get_stage_result(self, stage: str) -> Any: - """Get a result from a specific processing stage.""" - return self.stage_results.get(stage) - - def set_hidden_states(self, hidden_states: Any) -> None: - """Set hidden states for the request.""" - self.hidden_states = hidden_states - - def get_hidden_states(self) -> Any: - """Get hidden states for the request.""" - return self.hidden_states - - def add_error(self, error: str) -> None: - """Add an error to the request.""" - self.errors.append({ - 'error': error, - 'timestamp': time.time(), - 'stage': self.processing_stage.value - }) - - def has_errors(self) -> bool: - """Check if the request has any errors.""" - return len(self.errors) > 0 - - def can_retry(self) -> bool: - """Check if the request can be retried.""" - return self.retry_count < self.max_retries - - def increment_retry(self) -> None: - """Increment the retry count.""" - self.retry_count += 1 - - def generate_cache_key(self) -> str: - """Generate a cache key for the request.""" - if self.cache_key: - return self.cache_key - - # Generate cache key based on request content - key_parts = [ - self.request_id, - str(self.request_type.value), - str(self.prompt_token_ids) if self.prompt_token_ids else str(self.prompt), - str(self.diffusion_params.num_inference_steps) if self.diffusion_params else "0", - str(len(self.multimodal_inputs)) - ] - return "_".join(key_parts) - - def get_processing_time(self) -> float: - """Get the total processing time for the request.""" - if self.stage_timings: - return time.time() - min(self.stage_timings.values()) - return 0.0 - - def is_completed(self) -> bool: - """Check if the request is completed.""" - return self.processing_stage == ProcessingStage.COMPLETED - - def to_dict(self) -> Dict[str, Any]: - """Convert the request to a dictionary for serialization.""" - return { - 'request_id': self.request_id, - 'request_type': self.request_type.value, - 'processing_stage': self.processing_stage.value, - 'prompt': self.prompt, - 'prompt_token_ids': self.prompt_token_ids, - 'output_format': self.output_format, - 'multimodal_inputs_count': len(self.multimodal_inputs), - 'diffusion_params': self.diffusion_params.__dict__ if self.diffusion_params else None, - 'stage_results': list(self.stage_results.keys()), - 'has_errors': self.has_errors(), - 'retry_count': self.retry_count, - 'processing_time': self.get_processing_time() - } - - def __repr__(self) -> str: - """String representation of the request.""" - return (f"OmniRequest(id={self.request_id}, " - f"type={self.request_type.value}, " - f"stage={self.processing_stage.value}, " - f"multimodal_inputs={len(self.multimodal_inputs)})") + @classmethod + def from_engine_core_request(cls, request: OmniEngineCoreRequest) -> "Request": + if request.mm_inputs is not None: + assert isinstance(request.mm_inputs, list) + assert is_list_of(request.mm_inputs, MultiModalKwargs), ( + "mm_inputs was not updated in EngineCore.add_request") - -# Factory functions for creating different types of requests -def create_text_request( - request_id: str, - prompt: str, - sampling_params: Optional[Any] = None, - **kwargs -) -> OmniRequest: - """Create a text-only request.""" - return OmniRequest( - request_id=request_id, - prompt=prompt, - request_type=RequestType.TEXT, - sampling_params=sampling_params, - **kwargs - ) - - -def create_image_request( - request_id: str, - prompt: str, - image_data: Any, - diffusion_params: Optional[DiffusionParams] = None, - **kwargs -) -> OmniRequest: - """Create an image generation request.""" - multimodal_input = MultimodalData( - data_type="image", - data=image_data, - metadata={"is_input": True} - ) - - return OmniRequest( - request_id=request_id, - prompt=prompt, - request_type=RequestType.IMAGE, - multimodal_inputs=[multimodal_input], - diffusion_params=diffusion_params, - output_format="image", - **kwargs - ) - - -def create_multimodal_request( - request_id: str, - prompt: str, - multimodal_inputs: List[MultimodalData], - **kwargs -) -> OmniRequest: - """Create a multimodal request.""" - return OmniRequest( - request_id=request_id, - prompt=prompt, - request_type=RequestType.MULTIMODAL, - multimodal_inputs=multimodal_inputs, - **kwargs - ) - - -def create_diffusion_request( - request_id: str, - prompt: str, - diffusion_params: DiffusionParams, - **kwargs -) -> OmniRequest: - """Create a diffusion model request.""" - return OmniRequest( - request_id=request_id, - prompt=prompt, - request_type=RequestType.DIFFUSION, - diffusion_params=diffusion_params, - **kwargs - ) \ No newline at end of file + return cls( + request_id=request.request_id, + client_index=request.client_index, + prompt_token_ids=request.prompt_token_ids, + multi_modal_inputs=request.mm_inputs, + multi_modal_hashes=request.mm_hashes, + multi_modal_placeholders=request.mm_placeholders, + sampling_params=request.sampling_params, + pooling_params=request.pooling_params, + eos_token_id=request.eos_token_id, + arrival_time=request.arrival_time, + lora_request=request.lora_request, + structured_output_request=StructuredOutputRequest( + sampling_params=request.sampling_params) \ + if request.sampling_params else None, + cache_salt=request.cache_salt, + priority=request.priority, + prompt_embeds=request.prompt_embeds, + additional_information=request.additional_information, + ) \ No newline at end of file