Add Omni model components and input processing for hidden states support
- Introduced `OmniModelRunnerOutput` and `OmniRequestOutput` classes for handling multimodal outputs. - Implemented `OmniEngineCoreRequest` and `OmniInputPreprocessor` to manage prompt embeddings and additional information. - Created `OmniTokensPrompt`, `OmniTokenInputs`, and `OmniEmbedsPrompt` for enhanced token handling. - Added parsing and preprocessing functions to support new input structures. - Removed the obsolete `diffusion_engine.py` file as part of the refactor.
This commit is contained in:
tzhouam committed
1 parent
7a85034d9b
commit
83e7f6ba12
11 files changed
+637
-431
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
Whitespace-only changes.
@@ -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
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
+42
-294
@@ -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
|
||||
)
|
||||
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,
|
||||
)
|
||||
Reference in new issue
Block a user