diff --git a/README.md b/README.md index 47820fd5b..7b1f8a294 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ vLLM-omni is built on a modular architecture that extends vLLM's core functional ## 🛠️ Installation -### Installation of vLLM +### Environment setup Use Docker to keep consistent basic environment (Optional, Recommended) ```bash @@ -51,23 +51,39 @@ pip install uv uv venv --python 3.12 --seed source .venv/bin/activate ``` -Install certain version of vllm with commitid: 808a7b69df479b6b3a16181711cac7ca28a9b941 + +### Installation of vLLM (for users) + +Now we build it based on vLLM v0.11.0. Please install it with command below. +```bash +uv pip install vllm==0.11.0 --torch-backend=auto +``` + +### Installation of vLLM (for developers) + +Install stable release version of vllm with 0.11.0 ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout 808a7b69df479b6b3a16181711cac7ca28a9b941 +git checkout v0.11.0 ``` Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set VLLM_PRECOMPILED_WHEEL_LOCATION as your local absolute path of whl file. ```bash -export VLLM_COMMIT=808a7b69df479b6b3a16181711cac7ca28a9b941 -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://wheels.vllm.ai/${VLLM_COMMIT}/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.11.0/vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl ``` -Install vllm with command below. +Install vllm with command below (If you have no existing PyTorch). ```bash uv pip install --editable . ``` +Install vllm with command below (If you already have PyTorch). +```bash +python use_existing_torch.py +uv pip install -r requirements/build.txt +uv pip install --no-build-isolation --editable . +``` + ### Installation of vLLM-omni Install additional requirements for vllm-omni ```bash @@ -90,7 +106,7 @@ The output audio is saved in ./output_audio ## To-do list - [x] Offline inference example for Qwen2.5-omni with single request -- [ ] Adaptation from current vllm branch to stable vllm v0.11.0 +- [x] Adaptation from current vllm branch to stable vllm v0.11.0 - [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests - [ ] Online inference support - [ ] Support for other models diff --git a/docs/architecture/high level arch design.md b/docs/architecture/high_level_arch_design.md similarity index 100% rename from docs/architecture/high level arch design.md rename to docs/architecture/high_level_arch_design.md diff --git a/examples/offline_inference/qwen_2_5_omni/README.md b/examples/offline_inference/qwen_2_5_omni/README.md index cbbf5c0e2..f6b3a4fdb 100644 --- a/examples/offline_inference/qwen_2_5_omni/README.md +++ b/examples/offline_inference/qwen_2_5_omni/README.md @@ -2,7 +2,7 @@ ## 🛠️ Installation -### Installation of vLLM +### Environment setup Use Docker to keep consistent basic environment (Optional, Recommended) ```bash @@ -15,23 +15,39 @@ pip install uv uv venv --python 3.12 --seed source .venv/bin/activate ``` -Install certain version of vllm with commitid: 808a7b69df479b6b3a16181711cac7ca28a9b941 + +### Installation of vLLM (for users) + +Now we build it based on vLLM v0.11.0. Please install it with command below. +```bash +uv pip install vllm==0.11.0 --torch-backend=auto +``` + +### Installation of vLLM (for developers) + +Install stable release version of vllm with 0.11.0 ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout 808a7b69df479b6b3a16181711cac7ca28a9b941 +git checkout v0.11.0 ``` Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set VLLM_PRECOMPILED_WHEEL_LOCATION as your local absolute path of whl file. ```bash -export VLLM_COMMIT=808a7b69df479b6b3a16181711cac7ca28a9b941 -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://wheels.vllm.ai/${VLLM_COMMIT}/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.11.0/vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl ``` -Install vllm with command below. +Install vllm with command below (If you have no existing PyTorch). ```bash uv pip install --editable . ``` +Install vllm with command below (If you already have PyTorch). +```bash +python use_existing_torch.py +uv pip install -r requirements/build.txt +uv pip install --no-build-isolation --editable . +``` + ### Installation of vLLM-omni Install additional requirements for vllm-omni ```bash @@ -54,7 +70,7 @@ The output audio is saved in ./output_audio ## To-do list - [x] Offline inference example for Qwen2.5-omni with single request -- [ ] Adaptation from current vllm branch to stable vllm v0.11.0 +- [x] Adaptation from current vllm branch to stable vllm v0.11.0 - [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests - [ ] Online inference support - [ ] Support for other models diff --git a/examples/offline_inference/qwen_2_5_omni/end2end.py b/examples/offline_inference/qwen_2_5_omni/end2end.py index 6dd06abb0..fa5df3e0b 100644 --- a/examples/offline_inference/qwen_2_5_omni/end2end.py +++ b/examples/offline_inference/qwen_2_5_omni/end2end.py @@ -7,8 +7,8 @@ import numpy as np import soundfile as sf import torch from utils import make_omni_prompt -from vllm.sampling_params import SamplingParams +from vllm.sampling_params import SamplingParams from vllm_omni.entrypoints.omni_llm import OmniLLM _os_env_toggle.environ["VLLM_USE_V1"] = "1" diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh index c2f8e8ea1..7f02a672c 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -1,5 +1,4 @@ export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH -export HF_ENDPOINT=https://hf-mirror.com python end2end.py --model Qwen/Qwen2.5-Omni-7B \ --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \ --voice-type "m02" \ diff --git a/examples/offline_inference/qwen_2_5_omni/utils.py b/examples/offline_inference/qwen_2_5_omni/utils.py index 221f86119..730d2a077 100644 --- a/examples/offline_inference/qwen_2_5_omni/utils.py +++ b/examples/offline_inference/qwen_2_5_omni/utils.py @@ -9,8 +9,8 @@ import soundfile as sf import torch import torchvision.io from processing_omni import fetch_image, fetch_video -from vllm.inputs import TextPrompt +from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt diff --git a/requirements.txt b/requirements.txt index 4a24e9b14..4a7116da4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,14 +1,4 @@ -cloudpickle==3.1.1 -msgspec==0.19.0 -numpy==2.3.4 omegaconf==2.3.0 -pydantic==2.12.3 -transformers==4.57.0 -typing_extensions==4.15.0 decord==0.6.0 librosa==0.11.0 -packaging==25.0 -Pillow==12.0.0 -Requests==2.32.5 resampy==0.4.3 -soundfile==0.13.1 diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index c2c0a5482..0fae07e7a 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -2,33 +2,7 @@ Configuration module for vLLM-omni. """ -from typing import Optional - -from pydantic import ConfigDict -from pydantic.dataclasses import dataclass -from vllm.config import ModelConfig, config - -import vllm_omni.model_executor.models as me_models - - -@config -@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) -class OmniModelConfig(ModelConfig): - """Configuration for Omni models, extending the base ModelConfig.""" - - stage_id: int = 0 - model_stage: str = "thinker" - model_arch: str = "Qwen2_5OmniForConditionalGeneration" - engine_output_type: Optional[str] = None - - @property - def registry(self): - return me_models.OmniModelRegistry - - @property - def architectures(self) -> list[str]: - return [self.model_arch] - +from vllm_omni.config.model import OmniModelConfig __all__ = [ "OmniModelConfig", diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py new file mode 100644 index 000000000..00e56813d --- /dev/null +++ b/vllm_omni/config/model.py @@ -0,0 +1,26 @@ +from typing import Optional + +from pydantic import ConfigDict +from pydantic.dataclasses import dataclass + +import vllm_omni.model_executor.models as me_models +from vllm.config import ModelConfig, config + + +@config +@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +class OmniModelConfig(ModelConfig): + """Configuration for Omni models, extending the base ModelConfig.""" + + stage_id: int = 0 + model_stage: str = "thinker" + model_arch: str = "Qwen2_5OmniForConditionalGeneration" + engine_output_type: Optional[str] = None + + @property + def registry(self): + return me_models.OmniModelRegistry + + @property + def architectures(self) -> list[str]: + return [self.model_arch] diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py index d4da2c2d1..ee878a540 100644 --- a/vllm_omni/core/sched/diffusion_scheduler.py +++ b/vllm_omni/core/sched/diffusion_scheduler.py @@ -3,6 +3,7 @@ from collections import defaultdict from typing import Optional from vllm.distributed.kv_events import KVEventBatch +from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.request_queue import create_request_queue from vllm.v1.core.sched.scheduler import ( EngineCoreOutputs, @@ -11,8 +12,8 @@ from vllm.v1.core.sched.scheduler import ( SchedulerOutput, SpecDecodingStats, ) +from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput - from vllm_omni.core.sched.output import OmniNewRequestData from vllm_omni.core.sched.scheduler import OmniScheduler from vllm_omni.outputs import OmniModelRunnerOutput @@ -37,11 +38,10 @@ class DiffusionScheduler(OmniScheduler): scheduled_resumed_reqs: list[Request] = [] scheduled_running_reqs: list[Request] = [] - req_to_new_block_ids: dict[str, tuple[list[int], ...]] = {} + req_to_new_blocks: dict[str, KVCacheBlocks] = {} num_scheduled_tokens: dict[str, int] = {} scheduled_spec_decode_tokens: dict[str, list[int]] = {} scheduled_encoder_inputs: dict[str, list[int]] = {} - structured_output_request_ids: dict[str, int] = {} # Temporary queue: preserve waiting order, do not disturb non-diffusion requests skipped_waiting_requests = create_request_queue(self.policy) @@ -80,7 +80,7 @@ class DiffusionScheduler(OmniScheduler): if self.log_stats: request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) - req_to_new_block_ids[request.request_id] = new_blocks.get_block_ids() + req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens token_budget -= num_new_tokens capacity -= 1 @@ -104,15 +104,11 @@ class DiffusionScheduler(OmniScheduler): ) ) - grammar_bitmask = self.structured_output_manager.grammar_bitmask( - self.requests, - structured_output_request_ids, - scheduled_spec_decode_tokens, - ) - # Assemble SchedulerOutput new_reqs_data = [ - OmniNewRequestData.from_request(req, req_to_new_block_ids[req.request_id]) + OmniNewRequestData.from_request( + req, req_to_new_blocks[req.request_id].get_block_ids() + ) for req in scheduled_new_reqs ] cached_reqs_data = self._make_cached_request_data( @@ -120,7 +116,13 @@ class DiffusionScheduler(OmniScheduler): scheduled_resumed_reqs, num_scheduled_tokens, scheduled_spec_decode_tokens, - req_to_new_block_ids, + req_to_new_blocks, + ) + scheduled_requests = ( + scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs + ) + structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask( + scheduled_requests, scheduled_spec_decode_tokens ) total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) @@ -133,7 +135,7 @@ class DiffusionScheduler(OmniScheduler): scheduled_encoder_inputs=scheduled_encoder_inputs, num_common_prefix_blocks=num_common_prefix_blocks, finished_req_ids=self.finished_req_ids, - free_encoder_input_ids=self.encoder_cache_manager.get_freed_ids(), + free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), structured_output_request_ids=structured_output_request_ids, grammar_bitmask=grammar_bitmask, ) @@ -145,6 +147,17 @@ class DiffusionScheduler(OmniScheduler): # Publish KV events (aligned with v1) events = self.kv_cache_manager.take_events() + + # collect KV cache events from connector + if self.connector is not None: + connector_events = self.connector.take_events() + if connector_events: + if events is None: + events = list(connector_events) + else: + events.extend(connector_events) + + # publish collected KV cache events if events: batch = KVEventBatch(ts=time.time(), events=events) self.kv_event_publisher.publish(batch) @@ -173,15 +186,18 @@ class DiffusionScheduler(OmniScheduler): This method is modified to stop the request immediately for the diffusion model. """ sampled_token_ids = model_runner_output.sampled_token_ids - spec_token_ids = model_runner_output.spec_token_ids logprobs = model_runner_output.logprobs prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict num_scheduled_tokens = scheduler_output.num_scheduled_tokens pooler_outputs = model_runner_output.pooler_output num_nans_in_logits = model_runner_output.num_nans_in_logits + kv_connector_output = model_runner_output.kv_connector_output outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: Optional[SpecDecodingStats] = None + kv_connector_stats = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, # the below loop can be a performance bottleneck. We should do our best @@ -206,20 +222,19 @@ class DiffusionScheduler(OmniScheduler): scheduler_output.scheduled_spec_decode_tokens.get(req_id) ) if scheduled_spec_token_ids: + num_draft_tokens = len(scheduled_spec_token_ids) + num_accepted = len(generated_token_ids) - 1 + num_rejected = num_draft_tokens - num_accepted # num_computed_tokens represents the number of tokens # processed in the current step, considering scheduled # tokens and rejections. If some tokens are rejected, # num_computed_tokens is decreased by the number of rejected - # tokens, where is given by: - # len(scheduled_spec_token_ids) + 1 - len(generated_token_ids). - num_tokens_rejected = ( - len(scheduled_spec_token_ids) + 1 - len(generated_token_ids) - ) - request.num_computed_tokens -= num_tokens_rejected + # tokens. + request.num_computed_tokens -= num_rejected spec_decoding_stats = self.make_spec_decoding_stats( spec_decoding_stats, - num_draft_tokens=len(scheduled_spec_token_ids), - num_accepted_tokens=len(generated_token_ids) - 1, + num_draft_tokens=num_draft_tokens, + num_accepted_tokens=num_accepted, ) new_logprobs = None @@ -263,17 +278,6 @@ class DiffusionScheduler(OmniScheduler): if num_nans_in_logits is not None and req_id in num_nans_in_logits: request.num_nans_in_logits = num_nans_in_logits[req_id] - # Add newly generated spec token ids to the request. - if spec_token_ids is not None: - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - # Needs to happen after new_token_ids are accepted. - request.spec_token_ids = metadata.grammar.validate_tokens( # type: ignore[union-attr] # noqa: E501 - spec_token_ids[req_index] - ) - else: - request.spec_token_ids = spec_token_ids[req_index] - # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) if new_token_ids or pooler_output is not None or kv_transfer_params: @@ -290,19 +294,18 @@ class DiffusionScheduler(OmniScheduler): stop_reason=request.stop_reason, events=request.take_events(), kv_transfer_params=kv_transfer_params, + trace_headers=request.trace_headers, num_cached_tokens=request.num_cached_tokens, ) ) - else: # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors # Remove the stopped requests from the running and waiting queues. if stopped_running_reqs: - self.running = [ - req for req in self.running if req not in stopped_running_reqs - ] + if stopped_running_reqs: + self.running = remove_all(self.running, stopped_running_reqs) if stopped_preempted_reqs: # This is a rare case and unlikely to impact performance. self.waiting.remove_requests(stopped_preempted_reqs) @@ -332,10 +335,14 @@ class DiffusionScheduler(OmniScheduler): ) finished_req_ids.clear() - if engine_core_outputs: + if ( + stats := self.make_stats(spec_decoding_stats, kv_connector_stats) + ) is not None: # Return stats to only one of the front-ends. - next(iter(engine_core_outputs.values())).scheduler_stats = self.make_stats( - spec_decoding_stats - ) + if (eco := next(iter(engine_core_outputs.values()), None)) is None: + # We must return the stats even if there are no request + # outputs this step. + engine_core_outputs[0] = eco = EngineCoreOutputs() + eco.scheduler_stats = stats return engine_core_outputs diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index 30c3d2db0..45eb6bee0 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -3,7 +3,6 @@ from typing import Optional from vllm.v1.core.sched.output import NewRequestData from vllm.v1.request import Request - from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload @@ -23,9 +22,7 @@ class OmniNewRequestData(NewRequestData): return cls( req_id=request.request_id, prompt_token_ids=request.prompt_token_ids, - mm_inputs=request.mm_inputs, - mm_hashes=request.mm_hashes, - mm_positions=request.mm_positions, + mm_features=request.mm_features, sampling_params=request.sampling_params, pooling_params=request.pooling_params, block_ids=block_ids, diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py index 4e925010b..fbd02c4b7 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/scheduler.py @@ -32,9 +32,7 @@ class OmniScheduler(VLLMScheduler): omni_nr = OmniNewRequestData( req_id=nr.req_id, prompt_token_ids=nr.prompt_token_ids, - mm_inputs=nr.mm_inputs, - mm_hashes=nr.mm_hashes, - mm_positions=nr.mm_positions, + mm_features=nr.mm_features, sampling_params=nr.sampling_params, pooling_params=nr.pooling_params, block_ids=nr.block_ids, diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 2de831388..dd3a9a710 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -5,6 +5,7 @@ Engine components for vLLM-omni. from typing import Any, Optional import msgspec + from vllm.v1.engine import EngineCoreRequest diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 9277a9d82..92036ae7c 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -2,8 +2,6 @@ from dataclasses import dataclass from typing import Optional from vllm.engine.arg_utils import EngineArgs -from vllm.utils import FlexibleArgumentParser - from vllm_omni.config import OmniModelConfig @@ -14,31 +12,6 @@ class OmniEngineArgs(EngineArgs): 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=OmniEngineArgs.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=OmniEngineArgs.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() diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 85be92fac..f870a0f16 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,10 +1,11 @@ from typing import Any, Callable, Dict, Optional, Union import torch + from vllm.logger import init_logger +from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind from vllm.transformers_utils.tokenizer import AnyTokenizer -from vllm.transformers_utils.tokenizer_group import TokenizerGroup from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason from vllm.v1.engine.detokenizer import IncrementalDetokenizer from vllm.v1.engine.logprobs import LogprobsProcessor @@ -16,6 +17,7 @@ from vllm.v1.engine.output_processor import ( ) from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.metrics.stats import IterationStats +from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -56,10 +58,16 @@ class OmniRequestState(RequestState): request=request, ) max_tokens_param = sampling_params.max_tokens + top_p = sampling_params.top_p + n = sampling_params.n + temperature = sampling_params.temperature else: logprobs_processor = None detokenizer = None max_tokens_param = None + top_p = None + n = None + temperature = None assert request.pooling_params is not None output_kind = request.pooling_params.output_kind @@ -73,9 +81,13 @@ class OmniRequestState(RequestState): output_kind=output_kind, prompt=prompt, prompt_token_ids=request.prompt_token_ids, + prompt_embeds=request.prompt_embeds, logprobs_processor=logprobs_processor, detokenizer=detokenizer, max_tokens_param=max_tokens_param, + top_p=top_p, + n=n, + temperature=temperature, arrival_time=request.arrival_time, queue=queue, log_stats=log_stats, @@ -112,21 +124,13 @@ class OmniRequestState(RequestState): finish_reason: Optional[FinishReason], stop_reason: Optional[Union[int, str]], kv_transfer_params: Optional[dict[str, Any]] = None, - num_cached_tokens: Optional[int] = None, - ) -> Optional[Any]: + ) -> Optional[Union[OmniRequestOutput, PoolingRequestOutput]]: finished = finish_reason is not None final_only = self.output_kind == RequestOutputKind.FINAL_ONLY if not finished and final_only: return None - if num_cached_tokens is not None: - # Keep num_cached_tokens in RequestOutput for compatibility - try: - self.num_cached_tokens = num_cached_tokens # type: ignore[attr-defined] - except Exception: - pass - request_id = self.request_id output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) @@ -187,7 +191,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): def __init__( self, - tokenizer: TokenizerGroup, + tokenizer: AnyTokenizer, log_stats: bool, engine_core_output_type: Optional[str] = None, ): @@ -214,14 +218,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): if request_id in self.request_states: raise ValueError(f"Request id {request_id} already running.") - tokenizer = ( - None - if not self.tokenizer - else self.tokenizer.get_lora_tokenizer(request.lora_request) - ) - req_state = OmniRequestState.from_new_request( - tokenizer=tokenizer, + tokenizer=self.tokenizer, request=request, prompt=prompt, parent_req=parent_req, @@ -266,7 +264,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): finish_reason = eco.finish_reason stop_reason = eco.stop_reason kv_transfer_params = eco.kv_transfer_params - num_cached_tokens = eco.num_cached_tokens + req_state.num_cached_tokens = eco.num_cached_tokens req_state.is_prefilling = False # 2) Detokenize and logprobs when text path @@ -308,7 +306,6 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): finish_reason, stop_reason, kv_transfer_params, - num_cached_tokens, ) if ro: # Attach accumulated multimodal payload if any @@ -341,6 +338,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): self._update_stats_from_finished( req_state, finish_reason, iteration_stats ) + if self.tracer: + self.do_tracing(eco, req_state, iteration_stats) # Cleanup per-request mm state if isinstance(req_state, OmniRequestState): req_state.mm_accumulated = None diff --git a/vllm_omni/engine/processor.py b/vllm_omni/engine/processor.py index 450acee9c..d9fb00c13 100644 --- a/vllm_omni/engine/processor.py +++ b/vllm_omni/engine/processor.py @@ -1,21 +1,22 @@ import time -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from typing import Any, Optional, Union import torch + from vllm.config import VllmConfig 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 MULTIMODAL_REGISTRY, MultiModalKwargs, MultiModalRegistry -from vllm.multimodal.inputs import PlaceholderRange -from vllm.multimodal.utils import merge_and_sort_multimodal_metadata +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.multimodal.utils import argsort_mm_positions from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams -from vllm.transformers_utils.tokenizer_group import TokenizerGroup +from vllm.transformers_utils.tokenizer import AnyTokenizer +from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.v1.engine.processor import Processor - from vllm_omni.engine import ( AdditionalInformationEntry, AdditionalInformationPayload, @@ -51,12 +52,15 @@ class OmniProcessor(Processor): def __init__( self, vllm_config: VllmConfig, - tokenizer: TokenizerGroup, + tokenizer: AnyTokenizer, mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, ): super().__init__(vllm_config, tokenizer, mm_registry) self.input_preprocessor = OmniInputPreprocessor( - self.model_config, self.tokenizer, mm_registry + self.model_config, + self.tokenizer, + mm_registry, + mm_processor_cache=self.mm_processor_cache, ) def process_inputs( @@ -75,9 +79,7 @@ class OmniProcessor(Processor): # 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.") + self._validate_params(params) data_parallel_size = self.vllm_config.parallel_config.data_parallel_size if data_parallel_rank is not None and not ( @@ -91,33 +93,65 @@ class OmniProcessor(Processor): if arrival_time is None: arrival_time = time.time() + # Optionally generate multimodal hash overrides to avoid hashing + # multimodal data items by their content as their identifiers. + + # NOTE: when users explicitly turn off BOTH prefix caching and input + # processing caching, no multimodal features or embeddings will be + # reused across requests, therefore identifying multimodal data items + # by their content is no longer necessary, and we create uuids with + # request id-modality-index as multimodal hash overrides. + if ( + self.model_config.multimodal_config + and self.model_config.multimodal_config.mm_processor_cache_gb == 0 + and not self.cache_config.enable_prefix_caching + ): + mm_uuids = self._maybe_build_mm_uuids(request_id, prompt) + else: + # Otherwise, use user-provided uuids as multimodal hash overrides + # if provided. + self._validate_multi_modal_uuids(prompt) + if isinstance(prompt, dict): + mm_uuids = prompt.get("multi_modal_uuids") + else: + mm_uuids = None + # 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, + mm_uuids=mm_uuids, ) 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) + eos_token_id = self.input_preprocessor.get_eos_token_id() encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs) + self._validate_model_inputs(encoder_inputs, decoder_inputs) - # TODO: Impl encoder-decoder - if encoder_inputs is not None: - raise NotImplementedError + # Mypy does not always properly infer the types of some elements of + # discriminated unions of TypedDicts, because of how it handles + # inheritance of TypedDict. If we explicitly extract the items we want + # we can avoid type errors from using `dict.get` later in the method. + prompt_str: Optional[str] = ( + None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt") + ) + prompt_token_ids = ( + decoder_inputs["prompt_token_ids"] + if decoder_inputs["type"] != "embeds" + else None + ) + prompt_embeds = ( + decoder_inputs["prompt_embeds"] + if decoder_inputs["type"] == "embeds" + else None + ) sampling_params = None pooling_params = None @@ -126,66 +160,41 @@ class OmniProcessor(Processor): 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"] + seq_len = length_from_prompt_token_ids_or_embeds( + prompt_token_ids, prompt_embeds ) + sampling_params.max_tokens = self.model_config.max_model_len - seq_len 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) - ) + sampling_params.update_from_tokenizer(self.tokenizer) 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 + mm_features: Optional[list[MultiModalFeatureSpec]] = None + if decoder_inputs["type"] == "multimodal": decoder_mm_inputs = decoder_inputs["mm_kwargs"] + decoder_mm_positions = decoder_inputs["mm_placeholders"] + decoder_mm_hashes = decoder_inputs["mm_hashes"] # 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, - ) + sorted_mm_idxs = argsort_mm_positions(decoder_mm_positions) - # 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 + mm_features = [] + for modality, idx in sorted_mm_idxs: + mm_features.append( + MultiModalFeatureSpec( + data=decoder_mm_inputs[modality][idx], + modality=modality, + identifier=decoder_mm_hashes[modality][idx], + mm_position=decoder_mm_positions[modality][idx], + ) ) - else: - sorted_mm_inputs = orig_sorted_mm_inputs # Serialize prompt_embeds and additional_information if provided # (direct-transfer path) @@ -232,12 +241,10 @@ class OmniProcessor(Processor): entries=entries ) - return decoder_inputs.get("prompt"), OmniEngineCoreRequest( + return prompt_str, 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, + prompt_token_ids=prompt_token_ids, + mm_features=mm_features, sampling_params=sampling_params, pooling_params=pooling_params, eos_token_id=eos_token_id, @@ -246,6 +253,7 @@ class OmniProcessor(Processor): cache_salt=decoder_inputs.get("cache_salt"), priority=priority, data_parallel_rank=data_parallel_rank, + trace_headers=trace_headers, prompt_embeds=prompt_embeds_payload, additional_information=additional_information_payload, ) diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 97f445235..57bd12958 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,18 +1,18 @@ from typing import Any, Optional, Sequence, Union import cloudpickle -import vllm.envs as envs from pydantic import ValidationError -from vllm.config import CompilationConfig, is_init_field + +from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field from vllm.engine.arg_utils import HfOverrides from vllm.entrypoints.llm import LLM from vllm.inputs import PromptType from vllm.logger import init_logger +from vllm.plugins.io_processors import get_io_processor from vllm.sampling_params import SamplingParams from vllm.usage.usage_lib import UsageContext from vllm.utils import Counter from vllm.v1.engine.llm_engine import LLMEngine - from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.engine.processor import OmniProcessor @@ -98,6 +98,9 @@ class OmniStageLLM(LLM): Union[int, dict[str, Any], CompilationConfig] ] = None, hf_overrides: Optional[HfOverrides] = None, + structured_outputs_config: Optional[ + Union[dict[str, Any], StructuredOutputsConfig] + ] = None, **kwargs, ): """LLM constructor.""" @@ -114,7 +117,7 @@ class OmniStageLLM(LLM): if "kv_transfer_config" in kwargs and isinstance( kwargs["kv_transfer_config"], dict ): - from vllm.config import KVTransferConfig + from vllm.config.kv_transfer import KVTransferConfig raw_config_dict = kwargs["kv_transfer_config"] try: @@ -139,22 +142,37 @@ class OmniStageLLM(LLM): level=compilation_config ) elif isinstance(compilation_config, dict): - - def predicate(x): - return is_init_field(CompilationConfig, x[0]) - compilation_config_instance = CompilationConfig( - **dict(filter(predicate, compilation_config.items())) + **{ + k: v + for k, v in compilation_config.items() + if is_init_field(CompilationConfig, k) + } ) else: compilation_config_instance = compilation_config else: compilation_config_instance = CompilationConfig() + if structured_outputs_config is not None: + if isinstance(structured_outputs_config, dict): + structured_outputs_instance = StructuredOutputsConfig( + **{ + k: v + for k, v in structured_outputs_config.items() + if is_init_field(StructuredOutputsConfig, k) + } + ) + else: + structured_outputs_instance = structured_outputs_config + else: + structured_outputs_instance = StructuredOutputsConfig() + engine_args = OmniEngineArgs( model=model, hf_overrides=hf_overrides, compilation_config=compilation_config_instance, + structured_outputs_config=structured_outputs_instance, **kwargs, ) @@ -175,11 +193,14 @@ class OmniStageLLM(LLM): self.request_counter = Counter() self.default_sampling_params: Union[dict[str, Any], None] = None - if envs.VLLM_USE_V1: - supported_tasks = self.llm_engine.get_supported_tasks() # type: ignore - else: - supported_tasks = self.llm_engine.model_config.supported_tasks + supported_tasks = self.llm_engine.get_supported_tasks() # type: ignore logger.info("Supported_tasks: %s", supported_tasks) self.supported_tasks = supported_tasks + + # Load the Input/Output processor plugin if any + io_processor_plugin = self.llm_engine.model_config.io_processor_plugin + self.io_processor = get_io_processor( + self.llm_engine.vllm_config, io_processor_plugin + ) diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index d0451cd81..06b1f1b4c 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -9,7 +9,6 @@ from vllm.inputs import TextPrompt from vllm.v1.engine import EngineCoreOutput from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.llm_engine import LLMEngine - from vllm_omni.inputs.data import OmniTokensPrompt diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 1e96dfe60..383c9f47d 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -2,6 +2,7 @@ import os from pathlib import Path from omegaconf import OmegaConf + from vllm.transformers_utils.config import get_config # Get the project root directory (2 levels up from this file) diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 34dd2fd62..abadb919c 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -1,6 +1,7 @@ from typing import Any, NotRequired, Optional import torch + from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt @@ -33,7 +34,6 @@ class OmniEmbedsPrompt(EmbedsPrompt): 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, @@ -45,8 +45,6 @@ def token_inputs_omni( 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: diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 349b21685..b1943318b 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -1,11 +1,10 @@ from typing import Any, Optional, Union from typing_extensions import assert_never + from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt from vllm.inputs.preprocess import InputPreprocessor -from vllm.lora.request import LoRARequest -from vllm.multimodal.inputs import MultiModalInputs - +from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni from vllm_omni.inputs.parse import parse_singleton_prompt_omni @@ -15,11 +14,12 @@ class OmniInputPreprocessor(InputPreprocessor): self, parsed_content: TokensPrompt, tokenization_kwargs: Optional[dict[str, Any]] = None, - lora_request: Optional[LoRARequest] = None, - return_mm_hashes: bool = False, + *, + mm_uuids: Optional[MultiModalUUIDDict] = None, ) -> Union[OmniTokenInputs, MultiModalInputs]: - prompt_token_ids = parsed_content["prompt_token_ids"] - token_type_ids = parsed_content.get("token_type_ids") + prompt_token_ids = self._truncate_inputs( + parsed_content["prompt_token_ids"], tokenization_kwargs + ) prompt_embeds = parsed_content.get("prompt_embeds") additional_information = parsed_content.get("additional_information") @@ -30,13 +30,11 @@ class OmniInputPreprocessor(InputPreprocessor): multi_modal_data, parsed_content.get("mm_processor_kwargs"), tokenization_kwargs=tokenization_kwargs, - lora_request=lora_request, - return_mm_hashes=return_mm_hashes, + mm_uuids=mm_uuids, ) 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, ) @@ -50,8 +48,8 @@ class OmniInputPreprocessor(InputPreprocessor): self, prompt: SingletonPrompt, tokenization_kwargs: Optional[dict[str, Any]] = None, - lora_request: Optional[LoRARequest] = None, - return_mm_hashes: bool = False, + *, + mm_uuids: Optional[MultiModalUUIDDict] = None, ) -> SingletonInputs: """ Extract the singleton inputs from a prompt. @@ -73,95 +71,19 @@ class OmniInputPreprocessor(InputPreprocessor): if parsed["type"] == "tokens": return self._process_tokens( parsed["content"], - lora_request=lora_request, - return_mm_hashes=return_mm_hashes, + mm_uuids=mm_uuids, ) if parsed["type"] == "text": return self._process_text( parsed["content"], tokenization_kwargs=tokenization_kwargs, - lora_request=lora_request, - return_mm_hashes=return_mm_hashes, + mm_uuids=mm_uuids, ) 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, + mm_uuids=mm_uuids, ) assert_never(parsed) diff --git a/vllm_omni/model_executor/layers/mrope.py b/vllm_omni/model_executor/layers/mrope.py index 5f35b64d1..e986c2b8d 100644 --- a/vllm_omni/model_executor/layers/mrope.py +++ b/vllm_omni/model_executor/layers/mrope.py @@ -4,6 +4,7 @@ from typing import Optional, Union import numpy as np import torch from transformers import PretrainedConfig + from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding @@ -48,6 +49,14 @@ class MRotaryEmbedding(RotaryEmbedding): is_neox_style: bool, dtype: torch.dtype, mrope_section: Optional[list[int]] = None, + mrope_interleaved: bool = False, + # YaRN parameters. + *, + scaling_factor: Optional[float] = None, + extrapolation_factor: float = 1, + attn_factor: float = 1, + beta_fast: int = 32, + beta_slow: int = 1, ) -> None: # In Qwen2.5-VL, the maximum index value is related to the duration of # the input video. We enlarge max_position_embeddings to 4 times to get diff --git a/vllm_omni/model_executor/model_loader/weight_utils.py b/vllm_omni/model_executor/model_loader/weight_utils.py index ee5e5a6d0..949c858dc 100644 --- a/vllm_omni/model_executor/model_loader/weight_utils.py +++ b/vllm_omni/model_executor/model_loader/weight_utils.py @@ -1,8 +1,10 @@ import time +from pathlib import Path from typing import Optional, Union import huggingface_hub from huggingface_hub import snapshot_download + from vllm.logger import init_logger from vllm.model_executor.model_loader.weight_utils import DisabledTqdm, get_lock @@ -16,7 +18,8 @@ def download_weights_from_hf_specific( revision: Optional[str] = None, ignore_patterns: Optional[Union[str, list[str]]] = None, ) -> str: - """Download model weights from Hugging Face Hub. + """Download model weights from Hugging Face Hub. Users can specify the + allow_patterns to download only the necessary weights. Args: model_name_or_path (str): The model name or path. @@ -33,6 +36,7 @@ def download_weights_from_hf_specific( Returns: str: The path to the downloaded model weights. """ + assert len(allow_patterns) > 0 local_only = huggingface_hub.constants.HF_HUB_OFFLINE logger.info("Using model weights format %s", allow_patterns) @@ -40,15 +44,20 @@ def download_weights_from_hf_specific( # downloading the same model weights at the same time. with get_lock(model_name_or_path, cache_dir): start_time = time.perf_counter() - hf_folder = snapshot_download( - model_name_or_path, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - cache_dir=cache_dir, - tqdm_class=DisabledTqdm, - revision=revision, - local_files_only=local_only, - ) + for allow_pattern in allow_patterns: + hf_folder = snapshot_download( + model_name_or_path, + allow_patterns=allow_pattern, + ignore_patterns=ignore_patterns, + cache_dir=cache_dir, + tqdm_class=DisabledTqdm, + revision=revision, + local_files_only=local_only, + ) + # If we have downloaded weights for this allow_pattern, + # we don't need to check the rest. + if any(Path(hf_folder).glob(allow_pattern)): + break time_taken = time.perf_counter() - start_time if time_taken > 0.5: logger.info( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index 894cd6084..1c600fbdd 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni.py @@ -11,9 +11,9 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( Qwen2_5OmniTalkerConfig, Qwen2_5OmniThinkerConfig, ) + from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, @@ -22,12 +22,13 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerProcessingInfo, ) from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix -from vllm.model_executor.sampling_metadata import SamplingMetadata # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors - +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -167,7 +168,7 @@ class Qwen2_5OmniForConditionalGeneration( def sampler(self): if hasattr(self.model, "sampler"): return self.model.sampler - return get_sampler() + return Sampler() def get_input_embeddings( self, @@ -556,16 +557,14 @@ class Qwen2_5OmniForConditionalGeneration( return output_token_ids, processed_output_token_embeds def compute_logits( - self, - hidden_states: Union[torch.Tensor, OmniOutput], - sampling_metadata: SamplingMetadata, + self, hidden_states: Union[torch.Tensor, OmniOutput] ) -> Optional[torch.Tensor]: # Handle OmniOutput type if isinstance(hidden_states, OmniOutput): hidden_states = hidden_states.text_hidden_states # Use thinker model for logits computation - return self.model.compute_logits(hidden_states, sampling_metadata) + return self.model.compute_logits(hidden_states) def sample( self, diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni_talker.py index 2451f04d2..6e9527cdd 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni_talker.py @@ -14,7 +14,6 @@ from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import ( from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.linear import ColumnParallelLinear -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsMultiModal, @@ -34,9 +33,11 @@ from vllm.model_executor.models.utils import ( maybe_prefix, merge_multimodal_embeddings, ) -from vllm.model_executor.sampling_metadata import SamplingMetadata from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler @MULTIMODAL_REGISTRY.register_processor( @@ -112,7 +113,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration( if hasattr(self.language_model, "sampler"): return self.language_model.sampler - return get_sampler() + return Sampler() def get_input_embeddings( self, @@ -165,12 +166,8 @@ class Qwen2_5OmniTalkerForConditionalGeneration( ) return hidden_states - def compute_logits( - self, - hidden_states: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> Optional[torch.Tensor]: - return self.language_model.compute_logits(hidden_states, sampling_metadata) + def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: + return self.language_model.compute_logits(hidden_states) def sample( self, diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py index 1fbb58a84..f1b20cfeb 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py @@ -17,6 +17,7 @@ from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import ( Qwen2_5OmniProcessor, ) from transformers.models.whisper import WhisperFeatureExtractor + from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( @@ -47,7 +48,6 @@ from vllm.model_executor.models.utils import ( maybe_prefix, merge_multimodal_embeddings, ) -from vllm.model_executor.sampling_metadata import SamplingMetadata from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( ImageItem, @@ -73,7 +73,6 @@ from vllm.multimodal.processing import ( from vllm.multimodal.profiling import BaseDummyInputsBuilder from vllm.sequence import IntermediateTensors from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens - from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding try: @@ -939,12 +938,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration( ) return text_inputs_embeds, hidden_states.unsqueeze(0) # (1, S, D) - def compute_logits( - self, - hidden_states: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> Optional[torch.Tensor]: - return self.language_model.compute_logits(hidden_states, sampling_metadata) + def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: + return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py index fb612470d..9af9dcdbe 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py @@ -22,9 +22,9 @@ from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import ( # Bring in HF base classes, configs and utilities used below from transformers.utils.logging import get_logger as _hf_get_logger + from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( AutoWeightsLoader as _Vllm_AutoWeightsLoader, @@ -34,8 +34,10 @@ from vllm.model_executor.models.utils import ( init_vllm_registered_model as _Vllm_init_vllm_registered_model, ) from vllm.model_executor.models.utils import maybe_prefix as _Vllm_maybe_prefix -from vllm.model_executor.sampling_metadata import SamplingMetadata from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler # Provide a no-op auto_docstring decorator to satisfy annotations if missing @@ -1743,7 +1745,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): @property def sampler(self): # Token2Wav does not use sampler; return vLLM default for API parity - return get_sampler() + return Sampler() def forward( self, @@ -1767,11 +1769,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): **kwargs, ) - def compute_logits( - self, - hidden_states: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> Optional[torch.Tensor]: + def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: # Token2Wav outputs waveform; logits are not applicable return hidden_states diff --git a/vllm_omni/model_executor/models/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_old.py index bde5740eb..0ffe22b22 100644 --- a/vllm_omni/model_executor/models/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_old.py @@ -3,6 +3,7 @@ from typing import Iterable, Optional, Set, Tuple, Union import torch from torch import nn from transformers import Qwen2Config + from vllm.attention import Attention, AttentionType from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -19,7 +20,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.pooler import Pooler, PoolingType from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -38,9 +38,11 @@ from vllm.model_executor.models.utils import ( make_layers, maybe_prefix, ) -from vllm.model_executor.sampling_metadata import SamplingMetadata -from vllm.sequence import IntermediateTensors, PoolerOutput +from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import PoolerOutput, SamplerOutput from vllm.v1.pool.metadata import PoolingMetadata +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler logger = init_logger(__name__) @@ -449,7 +451,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) - self.sampler = get_sampler() + self.sampler = Sampler() self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors @@ -473,9 +475,8 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): def compute_logits( self, hidden_states: torch.Tensor, - sampling_metadata: SamplingMetadata, ) -> Optional[torch.Tensor]: - logits = self.logits_processor(self.lm_head, hidden_states, sampling_metadata) + logits = self.logits_processor(self.lm_head, hidden_states) return logits def sample( diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 0216eb408..4dc10d4fd 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -1,8 +1,8 @@ from typing import Union import torch -from vllm.inputs import TextPrompt +from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 454e7b423..cf97114eb 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from typing import Optional import torch + from vllm.outputs import RequestOutput from vllm.v1.outputs import ModelRunnerOutput diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index bbf24c3f8..b8de3010b 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -6,7 +6,6 @@ from vllm.model_executor.layers.rotary_embedding import ( ) from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest from vllm.v1.request import Request as _OriginalRequest - from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 4c15dbb21..cb5e012f5 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -1,10 +1,11 @@ -from typing import Optional +from typing import TYPE_CHECKING, Callable, Optional -from vllm.multimodal.inputs import MultiModalKwargs -from vllm.utils import is_list_of from vllm.v1.request import Request from vllm.v1.structured_output.request import StructuredOutputRequest +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_utils import BlockHash + from vllm_omni.engine import ( AdditionalInformationPayload, OmniEngineCoreRequest, @@ -29,20 +30,17 @@ class OmniRequest(Request): ) @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" + def from_engine_core_request( + cls, + request: OmniEngineCoreRequest, + block_hasher: Optional[Callable[["Request"], list["BlockHash"]]], + ) -> "Request": 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, + mm_features=request.mm_features, sampling_params=request.sampling_params, pooling_params=request.pooling_params, eos_token_id=request.eos_token_id, @@ -55,6 +53,8 @@ class OmniRequest(Request): ), cache_salt=request.cache_salt, priority=request.priority, + trace_headers=request.trace_headers, + block_hasher=block_hasher, prompt_embeds=request.prompt_embeds, additional_information=request.additional_information, ) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 1aeea3855..20a98e55a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -6,28 +6,34 @@ and also outputs sampled tokens. from __future__ import annotations -from typing import Any, List, Optional, Union +from typing import Any, Optional, Union import numpy as np import torch -from vllm import envs -from vllm.multimodal.inputs import MultiModalKwargs -from vllm.v1.attention.backends.utils import CommonAttentionMetadata + +from vllm.forward_context import BatchDescriptor +from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.outputs import AsyncModelRunnerOutput +from vllm.v1.structured_output.utils import apply_grammar_bitmask +from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, + AsyncGPUModelRunnerOutput, IntermediateTensors, get_pp_group, get_tp_group, has_kv_transfer_group, set_forward_context, ) - +from vllm.v1.worker.ubatch_utils import UBatchSlices +from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner +logger = init_logger(__name__) + class GPUARModelRunner(OmniGPUModelRunner): """Autoregressive GPU model runner that returns hidden states per request. @@ -41,129 +47,57 @@ class GPUARModelRunner(OmniGPUModelRunner): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.input_ids = torch.zeros( - self.max_num_tokens, dtype=torch.int64, device=self.device - ) + self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) - @torch.inference_mode() - def execute_model( + def _preprocess( self, scheduler_output: "SchedulerOutput", + num_scheduled_tokens_np: np.ndarray, intermediate_tensors: Optional[IntermediateTensors] = None, - ) -> Union[OmniModelRunnerOutput, IntermediateTensors]: - # Update internal state with the new schedule - self._update_states(scheduler_output) + ubatch_slices: Optional[UBatchSlices] = None, + num_tokens_after_padding: Optional[torch.Tensor] = None, + ) -> tuple[ + int, + int, + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + torch.Tensor, + Optional[IntermediateTensors], + dict[str, Any], + ]: - # Decode per-request prompt_embeds / additional_hidden_states payloads - # (if present) into CPU tensors - try: - new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) - if new_reqs: - for nr in new_reqs: - req_id = getattr(nr, "req_id", None) or getattr( - nr, "request_id", None - ) - if req_id is None: - continue - # prompt_embeds - payload_pe = getattr(nr, "prompt_embeds", None) - if payload_pe is not None: - if isinstance(payload_pe, torch.Tensor): - pe_cpu = payload_pe.detach().to("cpu").contiguous() - elif isinstance(payload_pe, PromptEmbedsPayload): - dt = np.dtype(getattr(payload_pe, "dtype", "float32")) - arr = np.frombuffer(payload_pe.data, dtype=dt) - arr = arr.reshape(payload_pe.shape) - pe_cpu = torch.from_numpy(arr) - else: - pe_cpu = None - if pe_cpu is not None and req_id in self.requests: - setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) - # additional_information - payload_info = getattr(nr, "additional_information", None) - if payload_info is not None: - info_dict = {} - if isinstance(payload_info, dict): - # Already decoded - info_dict = payload_info - elif isinstance(payload_info, AdditionalInformationPayload): - for k, entry in payload_info.entries.items(): - if entry.tensor_data is not None: - dt = np.dtype( - getattr(entry, "tensor_dtype", "float32") - ) - arr = np.frombuffer(entry.tensor_data, dtype=dt) - arr = arr.reshape(entry.tensor_shape) - info_dict[k] = torch.from_numpy(arr) - else: - info_dict[k] = entry.list_data - if info_dict and req_id in self.requests: - setattr( - self.requests[req_id], - "additional_information_cpu", - info_dict, - ) - except Exception: - pass - - # If there's no work to do, either return empty output or kv-only path - if not scheduler_output.total_num_scheduled_tokens: - if not has_kv_transfer_group(): - return EMPTY_MODEL_RUNNER_OUTPUT - return self.kv_connector_no_forward(scheduler_output, self.vllm_config) - - # Prepare decoder inputs and attention metadata - ( - attn_metadata, - attention_cuda_graphs, - logits_indices, - spec_decode_metadata, - num_scheduled_tokens_np, - spec_decode_common_attn_metadata, - ) = self._prepare_inputs(scheduler_output) - - # Determine number of input tokens for this iteration num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + if ubatch_slices: + assert num_tokens_after_padding is not None + num_input_tokens = int(num_tokens_after_padding[0].item() * 2) + self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens) + elif ubatch_slices is None: + num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens) + num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens) + num_input_tokens += num_pad + + # _prepare_inputs may reorder the batch, so we must gather multi + # modal outputs after that to ensure the correct order if ( - self.use_cuda_graph - and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1] + self.supports_mm_inputs + and get_pp_group().is_first_rank + and not self.model_config.is_encoder_decoder ): - num_input_tokens = self.vllm_config.pad_for_cudagraph(num_scheduled_tokens) - else: - tp_size = self.vllm_config.parallel_config.tensor_parallel_size - if ( - self.compilation_config.pass_config.enable_sequence_parallelism - and tp_size > 1 - ): - from vllm.utils import round_up # lazy local import - - num_input_tokens = round_up(num_scheduled_tokens, tp_size) - else: - num_input_tokens = num_scheduled_tokens - - # DP padding - num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens) - num_input_tokens += num_pad - - # Multimodal handling (encode and gather embeddings if needed) - if self.is_multimodal_model: + # Run the multimodal encoder if any. self._execute_mm_encoder(scheduler_output) mm_embeds = self._gather_mm_embeddings(scheduler_output) - else: - mm_embeds = [] - # Always assemble inputs_embeds on first PP rank; overlay per-request - # prompt_embeds and collect additional_hidden_states for prefill only - if get_pp_group().is_first_rank: + # NOTE(woosuk): To unify token ids and soft tokens (vision + # embeddings), we always use embeddings (rather than token ids) + # as input to the multimodal model, even when the input is text. inputs_embeds_scheduled = self.model.get_input_embeddings( - input_ids=self.input_ids[:num_scheduled_tokens], - multimodal_embeddings=( - (mm_embeds or None) if self.is_multimodal_model else None - ), + input_ids=self.input_ids.gpu[:num_scheduled_tokens], + multimodal_embeddings=mm_embeds or None, ) - # Copy into persistent buffer to enable CUDA Graph capture - self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) + # TODO(woosuk): Avoid the copy. Optimize. + self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) # Reset per-step additional information collector if hasattr(self, "_forward_additional_information"): @@ -237,25 +171,52 @@ class GPUARModelRunner(OmniGPUModelRunner): # Mixed types: wrap existing into list self._forward_additional_information[k] = [prev_val, v] - input_ids = self.input_ids[:num_input_tokens] # preserved for APIs - inputs_embeds = self.inputs_embeds[:num_input_tokens] - model_mm_kwargs = ( - self._extract_mm_kwargs(scheduler_output) - if self.is_multimodal_model - else {} + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = { + **self._init_model_kwargs(num_scheduled_tokens), + **self._extract_mm_kwargs(scheduler_output), + } + elif self.enable_prompt_embeds and get_pp_group().is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the CUDA graph all the time. The v0 + # engine avoids this by "double compiling" the CUDA graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the CUDA graph will be more performant (like in the else case + # below). + token_ids_idx = ( + self.is_token_ids.gpu[:num_scheduled_tokens] + .nonzero(as_tuple=False) + .squeeze(1) ) + # Some tokens ids may need to become embeds + if token_ids_idx.numel() > 0: + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds + + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = self._init_model_kwargs(num_input_tokens) + input_ids = None else: - input_ids = self.input_ids[:num_input_tokens] + # For text-only models, we use token ids as input. + # While it is possible to use embeddings as input just like the + # multimodal models, it is not desirable for performance since + # then the embedding layer is not included in the CUDA graph. + input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = None - model_mm_kwargs = {} - - # Positions/mRoPE + model_kwargs = self._init_model_kwargs(num_input_tokens) if self.uses_mrope: - positions = self.mrope_positions[:, :num_input_tokens] + positions = self.mrope_positions.gpu[:, :num_input_tokens] else: - positions = self.positions[:num_input_tokens] + positions = self.positions.gpu[:num_input_tokens] - # Handle pipeline-parallel intermediate tensors if get_pp_group().is_first_rank: intermediate_tensors = None else: @@ -263,17 +224,175 @@ class GPUARModelRunner(OmniGPUModelRunner): num_input_tokens, intermediate_tensors, True ) - # Some attention backends only support CUDA Graphs in pure decode. - skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs + if ( + self.model_config.is_encoder_decoder + and scheduler_output.scheduled_encoder_inputs + ): + encoder_inputs = self._extract_encoder_inputs(scheduler_output) + model_kwargs.update(encoder_inputs) - # Forward pass - with set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_input_tokens, - num_tokens_across_dp=num_tokens_across_dp, - skip_cuda_graphs=skip_cuda_graphs, - ), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output: + return ( + num_scheduled_tokens, + num_input_tokens, + num_tokens_after_padding, + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ) + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: "SchedulerOutput", + intermediate_tensors: Optional[IntermediateTensors] = None, + ) -> Union[OmniModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]: + with record_function_or_nullcontext("Preprocess"): + with self.synchronize_input_prep(): + # Update persistent batch states. + self._update_states(scheduler_output) + + # Decode per-request prompt_embeds / additional_hidden_states payloads + # (if present) into CPU tensors + try: + new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) + if new_reqs: + for nr in new_reqs: + req_id = getattr(nr, "req_id", None) or getattr( + nr, "request_id", None + ) + if req_id is None: + continue + # prompt_embeds + payload_pe = getattr(nr, "prompt_embeds", None) + if payload_pe is not None: + if isinstance(payload_pe, torch.Tensor): + pe_cpu = payload_pe.detach().to("cpu").contiguous() + elif isinstance(payload_pe, PromptEmbedsPayload): + dt = np.dtype( + getattr(payload_pe, "dtype", "float32") + ) + arr = np.frombuffer(payload_pe.data, dtype=dt) + arr = arr.reshape(payload_pe.shape) + pe_cpu = torch.from_numpy(arr) + else: + pe_cpu = None + if pe_cpu is not None and req_id in self.requests: + setattr( + self.requests[req_id], + "prompt_embeds_cpu", + pe_cpu, + ) + # additional_information + payload_info = getattr(nr, "additional_information", None) + if payload_info is not None: + info_dict = {} + if isinstance(payload_info, dict): + # Already decoded + info_dict = payload_info + elif isinstance( + payload_info, AdditionalInformationPayload + ): + for k, entry in payload_info.entries.items(): + if entry.tensor_data is not None: + dt = np.dtype( + getattr( + entry, "tensor_dtype", "float32" + ) + ) + arr = np.frombuffer( + entry.tensor_data, dtype=dt + ) + arr = arr.reshape(entry.tensor_shape) + info_dict[k] = torch.from_numpy(arr) + else: + info_dict[k] = entry.list_data + if info_dict and req_id in self.requests: + setattr( + self.requests[req_id], + "additional_information_cpu", + info_dict, + ) + except Exception as e: + logger.error( + f"Error decoding prompt_embeds / additional_information: {e}" + ) + pass + + if not scheduler_output.total_num_scheduled_tokens: + if not has_kv_transfer_group(): + # Return empty ModelRunnerOutput if no work to do. + return EMPTY_MODEL_RUNNER_OUTPUT + return self.kv_connector_no_forward( + scheduler_output, self.vllm_config + ) + if self.cache_config.kv_sharing_fast_prefill: + assert not self.input_batch.num_prompt_logprobs, ( + "--kv-sharing-fast-prefill produces incorrect " + "logprobs for prompt tokens, tokens, please disable " + "it when the requests need prompt logprobs" + ) + + # Prepare the decoder inputs. + ( + attn_metadata, + logits_indices, + spec_decode_metadata, + num_scheduled_tokens_np, + spec_decode_common_attn_metadata, + max_query_len, + ubatch_slices, + num_tokens_after_padding, + ) = self._prepare_inputs(scheduler_output) + + ( + num_scheduled_tokens, + num_input_tokens, + num_tokens_across_dp, + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ) = self._preprocess( + scheduler_output, + num_scheduled_tokens_np, + intermediate_tensors, + ubatch_slices, + num_tokens_after_padding, + ) + + uniform_decode = (max_query_len == self.uniform_decode_query_len) and ( + num_scheduled_tokens == self.input_batch.num_reqs * max_query_len + ) + batch_descriptor = BatchDescriptor( + num_tokens=num_input_tokens, uniform_decode=uniform_decode + ) + cudagraph_runtime_mode, batch_descriptor = ( + self.cudagraph_dispatcher.dispatch(batch_descriptor) + ) + + # This is currently to get around the assert in the DPMetadata + # where it wants `num_tokens_across_dp` to align with `num_tokens` + if ubatch_slices is not None: + num_input_tokens = ubatch_slices[0].num_tokens + + # Run the model. + # Use persistent buffers for CUDA graphs. + with ( + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + ubatch_slices=ubatch_slices, + ), + record_function_or_nullcontext("Forward"), + self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, + ): model_kwargs_extra = {} # Only pass additional_information for the prefill part @@ -290,189 +409,166 @@ class GPUARModelRunner(OmniGPUModelRunner): positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, - **MultiModalKwargs.as_kwargs( - model_mm_kwargs, - device=self.device, - ), + **model_kwargs, sampling_metadata=self.input_batch.sampling_metadata, logits_index=logits_indices, sampler=self.sampler, **model_kwargs_extra, ) - if self.use_aux_hidden_state_outputs: - hidden_states, _aux_hidden_states = model_output - else: - hidden_states = model_output - _aux_hidden_states = None + with record_function_or_nullcontext("Postprocess"): + if self.use_aux_hidden_state_outputs: + # True when EAGLE 3 is used. + hidden_states, aux_hidden_states = model_output + else: + # Common case. + hidden_states = model_output + aux_hidden_states = None - text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs( - hidden_states - ) - - # Mid PP stages return intermediate tensors unmodified - if not get_pp_group().is_last_rank: - assert isinstance(text_hidden_states, IntermediateTensors) - text_hidden_states.kv_connector_output = kv_connector_output - return text_hidden_states - - # Broadcast PP output for external_launcher (torchrun) - broadcast_pp_output = ( - self.parallel_config.distributed_executor_backend == "external_launcher" - and len(get_pp_group().ranks) > 0 - ) - if not get_pp_group().is_last_rank: - assert isinstance(text_hidden_states, IntermediateTensors) - if not broadcast_pp_output: - text_hidden_states.kv_connector_output = kv_connector_output - return text_hidden_states - get_pp_group().send_tensor_dict( - text_hidden_states.tensors, all_gather_group=get_tp_group() + hidden_states, multimodal_outputs = self.extract_multimodal_outputs( + hidden_states ) - logits = None - else: - if self.input_batch.pooling_params: - return self._pool( - text_hidden_states, - num_scheduled_tokens, - num_scheduled_tokens_np, - kv_connector_output, + + if not self.broadcast_pp_output: + # Common case. + if not get_pp_group().is_last_rank: + # Return the intermediate tensors. + assert isinstance(hidden_states, IntermediateTensors) + hidden_states.kv_connector_output = kv_connector_output + return hidden_states + + if self.is_pooling_model: + # Return the pooling output. + output = self._pool( + hidden_states, num_scheduled_tokens, num_scheduled_tokens_np + ) + output.kv_connector_output = kv_connector_output + return output + + sample_hidden_states = hidden_states[logits_indices] + logits = self.model.compute_logits(sample_hidden_states) + else: + # Rare case. + assert not self.is_pooling_model + + if not get_pp_group().is_last_rank: + all_gather_tensors = { + "residual": not is_residual_scattered_for_sp( + self.vllm_config, num_input_tokens + ) + } + get_pp_group().send_tensor_dict( + hidden_states.tensors, + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, + ) + logits = None + else: + sample_hidden_states = hidden_states[logits_indices] + logits = self.model.compute_logits(sample_hidden_states) + + model_output_broadcast_data = {} + if logits is not None: + model_output_broadcast_data["logits"] = logits.contiguous() + + model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( + model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 + ) + assert model_output_broadcast_data is not None + logits = model_output_broadcast_data["logits"] + + # Apply structured output bitmasks if present + if scheduler_output.grammar_bitmask is not None: + apply_grammar_bitmask( + scheduler_output, self.input_batch, logits, self.device ) - sample_hidden_states = text_hidden_states[logits_indices] - logits = self.model.compute_logits(sample_hidden_states, None) - if broadcast_pp_output: - model_output_broadcast_data = ( - { - "logits": logits.contiguous(), - } - if logits is not None - else {} - ) - model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( - model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 - ) - assert model_output_broadcast_data is not None - logits = model_output_broadcast_data["logits"] + with record_function_or_nullcontext("Sample"): + sampler_output = self._sample(logits, spec_decode_metadata) - # Apply structured output bitmasks if present - if scheduler_output.grammar_bitmask is not None: - self.apply_grammar_bitmask(scheduler_output, logits) - - # Sample the next token and get logprobs if needed (with spec decode) - sampling_metadata = self.input_batch.sampling_metadata - if spec_decode_metadata is None: - sampler_output = self.sampler( - logits=logits, - sampling_metadata=sampling_metadata, - ) - else: - assert logits is not None - bonus_logits = logits[spec_decode_metadata.bonus_logits_indices] - sampler_output = self.sampler( - logits=bonus_logits, - sampling_metadata=sampling_metadata, - ) - bonus_token_ids = sampler_output.sampled_token_ids - target_logits = logits[spec_decode_metadata.target_logits_indices] - output_token_ids = self.rejection_sampler( - spec_decode_metadata, - None, - target_logits, - bonus_token_ids, - sampling_metadata, - ) - sampler_output.sampled_token_ids = output_token_ids - - num_nans_in_logits = {} - if envs.VLLM_COMPUTE_NANS_IN_LOGITS: - num_nans_in_logits = self._get_nans_in_logits(logits) - - # Handle partial prefill: discard sampled tokens and rewind RNG - discard_sampled_tokens_req_indices = [] - for i, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests[req_id] - seq_len = ( - req_state.num_computed_tokens - + scheduler_output.num_scheduled_tokens[req_id] - ) - if seq_len < req_state.num_tokens: - generator = self.input_batch.generators.get(i) - if generator is not None: - generator.set_offset(generator.get_offset() - 4) - discard_sampled_tokens_req_indices.append(i) - - # Move CPU sync parts - logprobs_tensors = sampler_output.logprobs_tensors - logprobs_lists = ( - logprobs_tensors.tolists() if logprobs_tensors is not None else None - ) - - # Prompt logprobs if needed - prompt_logprobs_dict = self._get_prompt_logprobs_dict( - text_hidden_states[:num_scheduled_tokens], - scheduler_output, - ) - - # Parse valid sampled tokens - sampled_token_ids = sampler_output.sampled_token_ids - max_gen_len = sampled_token_ids.shape[-1] - if max_gen_len == 1: - valid_sampled_token_ids = sampled_token_ids.tolist() - else: - valid_sampled_token_ids = self.rejection_sampler.parse_output( - sampled_token_ids, - self.input_batch.vocab_size, - ) - for i in discard_sampled_tokens_req_indices: - valid_sampled_token_ids[i].clear() - - # Cache sampled tokens - for req_idx, sampled_ids in enumerate(valid_sampled_token_ids): - if not sampled_ids: - continue - start_idx = self.input_batch.num_tokens_no_spec[req_idx] - end_idx = start_idx + len(sampled_ids) - self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids - self.input_batch.num_tokens_no_spec[req_idx] = end_idx - self.input_batch.num_tokens[req_idx] = end_idx - req_id = self.input_batch.req_ids[req_idx] - req_state = self.requests[req_id] - req_state.output_token_ids.extend(sampled_ids) - - # Speculative decoding draft tokens if configured - if not self.speculative_config: - spec_token_ids = None - else: + def propose_draft_token_ids(sampled_token_ids): assert spec_decode_common_attn_metadata is not None - spec_token_ids = self.propose_draft_token_ids( - scheduler_output, - valid_sampled_token_ids, - sampling_metadata, - text_hidden_states, - sample_hidden_states, - _aux_hidden_states, - spec_decode_metadata, - spec_decode_common_attn_metadata, + with record_function_or_nullcontext("Draft"): + self._draft_token_ids = self.propose_draft_token_ids( + scheduler_output, + sampled_token_ids, + self.input_batch.sampling_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + spec_decode_metadata, + spec_decode_common_attn_metadata, + ) + + use_padded_batch_for_eagle = ( + self.speculative_config + and self.speculative_config.use_eagle() + and not self.speculative_config.disable_padded_drafter_batch + ) + effective_drafter_max_model_len = self.max_model_len + if effective_drafter_max_model_len is None: + effective_drafter_max_model_len = self.model_config.max_model_len + if ( + self.speculative_config + and self.speculative_config.draft_model_config is not None + and self.speculative_config.draft_model_config.max_model_len is not None + ): + effective_drafter_max_model_len = ( + self.speculative_config.draft_model_config.max_model_len ) + input_fits_in_drafter = spec_decode_common_attn_metadata and ( + spec_decode_common_attn_metadata.seq_lens.max() + + self.speculative_config.num_speculative_tokens + <= effective_drafter_max_model_len + ) + if use_padded_batch_for_eagle and input_fits_in_drafter: + # EAGLE speculative decoding can use the GPU sampled tokens + # as inputs, and does not need to wait for bookkeeping to finish. + propose_draft_token_ids(sampler_output.sampled_token_ids) + + with record_function_or_nullcontext("Bookkeep"): + ( + num_nans_in_logits, + logprobs_lists, + valid_sampled_token_ids, + prompt_logprobs_dict, + req_ids_output_copy, + req_id_to_index_output_copy, + invalid_req_indices, + ) = self._bookkeeping_sync( + scheduler_output, + sampler_output, + logits, + hidden_states, + num_scheduled_tokens, + ) + + if ( + self.speculative_config + and not use_padded_batch_for_eagle + and input_fits_in_drafter + ): + # ngram and other speculative decoding methods use the sampled + # tokens on the CPU, so they are run after bookkeeping. + propose_draft_token_ids(valid_sampled_token_ids) + + with record_function_or_nullcontext("EPLB"): + self.eplb_step() # Convert to per-request tensors on CPU - text_hidden_states_cpu = text_hidden_states.detach().to("cpu").contiguous() + hidden_states_cpu = hidden_states.detach().to("cpu").contiguous() pooler_output: list[Optional[torch.Tensor]] = [] prev_logits_index = 0 for logits_index in logits_indices: pooler_output.append( - text_hidden_states_cpu[prev_logits_index : logits_index + 1] + hidden_states_cpu[prev_logits_index : logits_index + 1] ) prev_logits_index = logits_index + 1 - self.eplb_step() - - return OmniModelRunnerOutput( - req_ids=self.input_batch.req_ids, - req_id_to_index=self.input_batch.req_id_to_index, + output = OmniModelRunnerOutput( + req_ids=req_ids_output_copy, + req_id_to_index=req_id_to_index_output_copy, sampled_token_ids=valid_sampled_token_ids, - spec_token_ids=spec_token_ids, logprobs=logprobs_lists, prompt_logprobs_dict=prompt_logprobs_dict, pooler_output=( @@ -484,165 +580,12 @@ class GPUARModelRunner(OmniGPUModelRunner): num_nans_in_logits=num_nans_in_logits, ) - @torch.inference_mode() - def extract_multimodal_outputs( - self, hidden_states: Union[torch.Tensor, List[torch.Tensor]] - ) -> dict: - if ( - hasattr(self.model, "have_multimodal_outputs") - and self.model.have_multimodal_outputs - ): - text_hidden_states = hidden_states.text_hidden_states - multimodal_outputs = hidden_states.multimodal_outputs + if not self.use_async_scheduling: + return output - elif isinstance(hidden_states, torch.Tensor): - text_hidden_states = hidden_states - multimodal_outputs = {} - elif isinstance(hidden_states, List): - text_hidden_states = hidden_states[0] - multimodal_outputs = {} - else: - raise ValueError(f"Invalid hidden states type: {type(hidden_states)}") - return text_hidden_states, multimodal_outputs - - @torch.inference_mode() - def _dummy_run( - self, - num_tokens: int, - capture_attn_cudagraph: bool = False, - skip_eplb: bool = False, - is_profile: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor]: - - # Padding for DP - num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) - num_tokens += num_pad - - # Set num_scheduled_tokens based on num_tokens and max_num_seqs - # for dummy run with LoRA so that the num_reqs collectively - # has num_tokens in total. - assert num_tokens <= self.scheduler_config.max_num_batched_tokens - max_num_reqs = self.scheduler_config.max_num_seqs - num_reqs = min(num_tokens, max_num_reqs) - min_tokens_per_req = num_tokens // num_reqs - num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs - num_scheduled_tokens_list[-1] += num_tokens % num_reqs - assert sum(num_scheduled_tokens_list) == num_tokens - assert len(num_scheduled_tokens_list) == num_reqs - num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) - - attn_metadata: Optional[dict[str, Any]] = None - if capture_attn_cudagraph: - attn_metadata = {} - - # Make sure max_model_len is used at the graph capture time. - self.seq_lens_np[:num_reqs] = self.max_model_len - self.seq_lens_np[num_reqs:] = 0 - self.seq_lens[:num_reqs].copy_( - self.seq_lens_cpu[:num_reqs], non_blocking=True - ) - - for kv_cache_group_id, kv_cache_group_spec in enumerate( - self.kv_cache_config.kv_cache_groups - ): - common_attn_metadata = CommonAttentionMetadata( - query_start_loc=self.query_start_loc[: num_reqs + 1], - query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1], - seq_lens=self.seq_lens[:num_reqs], - seq_lens_cpu=self.seq_lens_cpu[:num_reqs], - num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[ # noqa: E501 - :num_reqs - ], - num_reqs=num_reqs, - num_actual_tokens=num_tokens, - max_query_len=num_tokens, - block_table_tensor=self.input_batch.block_table[ - kv_cache_group_id - ].get_device_tensor()[:num_reqs], - slot_mapping=self.input_batch.block_table[ - kv_cache_group_id - ].slot_mapping[:num_tokens], - causal=True, - ) - - for attn_group in self.attn_groups[kv_cache_group_id]: - attn_metadata_i = ( - attn_group.metadata_builder.build_for_cudagraph_capture( - common_attn_metadata - ) - ) - for layer_name in kv_cache_group_spec.layer_names: - attn_metadata[layer_name] = attn_metadata_i - - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens): - if self.is_multimodal_model: - input_ids = None - inputs_embeds = self.inputs_embeds[:num_tokens] - model_mm_kwargs = self._dummy_mm_kwargs(num_reqs) - else: - input_ids = self.input_ids[:num_tokens] - inputs_embeds = None - model_mm_kwargs = {} - - if self.uses_mrope: - positions = self.mrope_positions[:, :num_tokens] - else: - positions = self.positions[:num_tokens] - - if get_pp_group().is_first_rank: - intermediate_tensors = None - else: - if self.intermediate_tensors is None: - self.intermediate_tensors = ( - self.model.make_empty_intermediate_tensors( - batch_size=self.max_num_tokens, - dtype=self.model_config.dtype, - device=self.device, - ) - ) - - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_tokens, None, False - ) - - with self.maybe_randomize_inputs(input_ids), set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_tokens, - num_tokens_across_dp=num_tokens_across_dp, - ): - outputs = self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **MultiModalKwargs.as_kwargs( - model_mm_kwargs, - device=self.device, - ), - ) - - if self.use_aux_hidden_state_outputs: - hidden_states, _ = outputs - else: - hidden_states = outputs - - if self.speculative_config and self.speculative_config.use_eagle(): - assert isinstance(self.drafter, EagleProposer) - self.drafter.dummy_run(num_tokens) - - # This is necessary to avoid blocking DP. - # For dummy runs, we typically skip EPLB since we don't have any real - # requests to process. - # However, in DP settings, there may be cases when some DP ranks do - # not have any requests to process, so they're executing dummy batches. - # In such cases, we still have to trigger EPLB to make sure - # ranks execute the rearrangement in synchronization. - if not skip_eplb: - self.eplb_step(is_dummy=True, is_profile=is_profile) - - logit_indices = np.cumsum(num_scheduled_tokens) - 1 - hidden_states, multimodal_outputs = self.extract_multimodal_outputs( - hidden_states + return AsyncGPUModelRunnerOutput( + model_runner_output=output, + sampled_token_ids=sampler_output.sampled_token_ids, + invalid_req_indices=invalid_req_indices, + async_output_copy_stream=self.async_output_copy_stream, ) - return hidden_states, hidden_states[logit_indices] diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index ab9754631..838068df7 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -2,34 +2,42 @@ import gc import os import torch + from vllm.model_executor import set_random_seed from vllm.platforms import current_platform from vllm.utils import GiB_bytes, MemorySnapshot from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment -from vllm.worker.worker import _check_if_gpu_supports_dtype - from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner class GPUARWorker(GPUWorker): def init_device(self): if self.device_config.device.type == "cuda": - # torch.distributed.all_reduce does not free the input tensor until - # the synchronization point. This causes the memory usage to grow - # as the number of all_reduce calls increases. This env var disables - # this behavior. - # Related issue: - # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 - os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" - # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) self.device = torch.device(f"cuda:{self.local_rank}") current_platform.set_device(self.device) - _check_if_gpu_supports_dtype(self.model_config.dtype) + current_platform.check_if_supports_dtype(self.model_config.dtype) + + # Initialize the distributed environment BEFORE taking + # memory snapshot + # This ensures NCCL buffers are allocated before we measure + # available memory + init_worker_distributed_environment( + self.vllm_config, + self.rank, + self.distributed_init_method, + self.local_rank, + current_platform.dist_backend, + ) + + # Set random seed. + set_random_seed(self.model_config.seed) + + # Now take memory snapshot after NCCL is initialized gc.collect() torch.cuda.empty_cache() @@ -55,16 +63,6 @@ class GPUARWorker(GPUWorker): ) else: raise RuntimeError(f"Not support device type: {self.device_config.device}") - # Initialize the distributed environment. - init_worker_distributed_environment( - self.vllm_config, - self.rank, - self.distributed_init_method, - self.local_rank, - current_platform.dist_backend, - ) - # Set random seed. - set_random_seed(self.model_config.seed) # Construct the model runner self.model_runner: GPUARModelRunner = GPUARModelRunner( diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu_diffusion_model_runner.py index c694208a9..5d4a5d6fb 100644 --- a/vllm_omni/worker/gpu_diffusion_model_runner.py +++ b/vllm_omni/worker/gpu_diffusion_model_runner.py @@ -2,21 +2,33 @@ from __future__ import annotations import gc import logging -from typing import List, Optional, Union +from typing import Any, List, Optional, Union import numpy as np import torch + +from vllm.config import CUDAGraphMode +from vllm.forward_context import BatchDescriptor from vllm.multimodal.inputs import MultiModalKwargs -from vllm.v1.attention.backends.utils import CommonAttentionMetadata +from vllm.utils import cdiv +from vllm.v1.attention.backends.utils import ( + CommonAttentionMetadata, + split_attn_metadata, +) from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, + AsyncGPUModelRunnerOutput, IntermediateTensors, + PerLayerAttnMetadata, get_pp_group, set_forward_context, ) +from vllm.v1.worker.ubatch_splitting import ubatch_split +from vllm.v1.worker.ubatch_utils import UBatchSlices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs - from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner @@ -31,65 +43,97 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): - Executes diffusion process and returns tensors via `pooler_output`. """ - @torch.inference_mode() - def execute_model( + def _preprocess( self, scheduler_output: "SchedulerOutput", intermediate_tensors: Optional[IntermediateTensors] = None, - ) -> Union[OmniModelRunnerOutput, IntermediateTensors]: - self._update_states(scheduler_output) - - if not scheduler_output.total_num_scheduled_tokens: - - return EMPTY_MODEL_RUNNER_OUTPUT - - # Prepare decoder inputs and attention metadata (for batch/order mapping) - ( - attn_metadata, - attention_cuda_graphs, - logits_indices, - spec_decode_metadata, - num_scheduled_tokens_np, - spec_decode_common_attn_metadata, - ) = self._prepare_inputs(scheduler_output) + ubatch_slices: Optional[UBatchSlices] = None, + num_tokens_after_padding: Optional[torch.Tensor] = None, + ) -> tuple[ + int, + int, + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + torch.Tensor, + Optional[IntermediateTensors], + dict[str, Any], + ]: # Input token count for this iteration (not used by diffusion, but # retained to keep DP padding/ordering consistent) num_input_tokens = scheduler_output.total_num_scheduled_tokens - num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens) + num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens) num_input_tokens += num_pad - # Multimodal conditioning (e.g., text/audio/video encoders) - if self.is_multimodal_model: + # _prepare_inputs may reorder the batch, so we must gather multi + # modal outputs after that to ensure the correct order + if ( + self.supports_mm_inputs + and get_pp_group().is_first_rank + and not self.model_config.is_encoder_decoder + ): + # Run the multimodal encoder if any. self._execute_mm_encoder(scheduler_output) mm_embeds = self._gather_mm_embeddings(scheduler_output) - else: - mm_embeds = [] - # Build inputs to mirror AR runner: input_ids/positions/embeds - if self.is_multimodal_model and get_pp_group().is_first_rank: + # NOTE(woosuk): To unify token ids and soft tokens (vision + # embeddings), we always use embeddings (rather than token ids) + # as input to the multimodal model, even when the input is text. inputs_embeds_scheduled = self.model.get_input_embeddings( - input_ids=self.input_ids[: scheduler_output.total_num_scheduled_tokens], + input_ids=self.input_ids.gpu[:num_input_tokens], multimodal_embeddings=mm_embeds or None, ) - self.inputs_embeds[: scheduler_output.total_num_scheduled_tokens].copy_( - inputs_embeds_scheduled + + # TODO(woosuk): Avoid the copy. Optimize. + self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled) + + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = { + **self._init_model_kwargs(num_input_tokens), + **self._extract_mm_kwargs(scheduler_output), + } + elif self.enable_prompt_embeds and get_pp_group().is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the CUDA graph all the time. The v0 + # engine avoids this by "double compiling" the CUDA graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the CUDA graph will be more performant (like in the else case + # below). + token_ids_idx = ( + self.is_token_ids.gpu[:num_input_tokens] + .nonzero(as_tuple=False) + .squeeze(1) ) - input_ids = self.input_ids[:num_input_tokens] - inputs_embeds = self.inputs_embeds[:num_input_tokens] - model_mm_kwargs = self._extract_mm_kwargs(scheduler_output) + # Some tokens ids may need to become embeds + if token_ids_idx.numel() > 0: + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds + + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = self._init_model_kwargs(num_input_tokens) + input_ids = self.input_ids.gpu[:num_input_tokens] else: - input_ids = self.input_ids[:num_input_tokens] + # For text-only models, we use token ids as input. + # While it is possible to use embeddings as input just like the + # multimodal models, it is not desirable for performance since + # then the embedding layer is not included in the CUDA graph. + input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = None - model_mm_kwargs = {} - - # Positions (mrope or standard) + model_kwargs = self._init_model_kwargs(num_input_tokens) if self.uses_mrope: - positions = self.mrope_positions[:, :num_input_tokens] + positions = self.mrope_positions.gpu[:, :num_input_tokens] else: - positions = self.positions[:num_input_tokens] + positions = self.positions.gpu[:num_input_tokens] - # Intermediate tensors sync for PP (if any) if get_pp_group().is_first_rank: intermediate_tensors = None else: @@ -97,32 +141,94 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): num_input_tokens, intermediate_tensors, True ) - # Set forward context mainly for resource management and kv connector - skip_cuda_graphs = True # diffusion path does not rely on cuda graphs here - with set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_input_tokens, - num_tokens_across_dp=num_tokens_across_dp, - skip_cuda_graphs=skip_cuda_graphs, - ), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output: + if ( + self.model_config.is_encoder_decoder + and scheduler_output.scheduled_encoder_inputs + ): + encoder_inputs = self._extract_encoder_inputs(scheduler_output) + model_kwargs.update(encoder_inputs) - if not get_pp_group().is_last_rank: - # For non-last PP stages, pass through intermediate tensors. - assert intermediate_tensors is not None - intermediate_tensors.kv_connector_output = kv_connector_output - return intermediate_tensors + return ( + num_input_tokens, + num_input_tokens, + num_tokens_after_padding, + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ) + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: "SchedulerOutput", + intermediate_tensors: Optional[IntermediateTensors] = None, + ) -> Union[OmniModelRunnerOutput, IntermediateTensors]: + with record_function_or_nullcontext("Preprocess"): + with self.synchronize_input_prep(): + # Update persistent batch states. + self._update_states(scheduler_output) + if not scheduler_output.total_num_scheduled_tokens: + return EMPTY_MODEL_RUNNER_OUTPUT + + # Prepare the decoder inputs. + ( + attn_metadata, + logits_indices, + spec_decode_metadata, + num_scheduled_tokens_np, + spec_decode_common_attn_metadata, + max_query_len, + ubatch_slices, + num_tokens_after_padding, + ) = self._prepare_inputs(scheduler_output) + + # Input token count for this iteration (not used by diffusion, but + # retained to keep DP padding/ordering consistent) + ( + num_scheduled_tokens, + num_input_tokens, + num_tokens_across_dp, + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ) = self._preprocess( + scheduler_output, + intermediate_tensors, + ubatch_slices, + num_tokens_after_padding, + ) + + cudagraph_runtime_mode = CUDAGraphMode.NONE + # Run the model. + # Use persistent buffers for CUDA graphs. + with ( + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=None, + ubatch_slices=ubatch_slices, + ), + record_function_or_nullcontext("Forward"), + self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, + ): outputs = self._run_diffusion( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, - multimodal_kwargs=model_mm_kwargs, + multimodal_kwargs=model_kwargs, logits_indices=logits_indices, ) - _, multimodal_outputs = self.extract_multimodal_outputs(outputs) + _, multimodal_outputs = self.extract_multimodal_outputs(outputs) # Ensure one tensor per request, map to CPU for output struct pooler_output: List[Optional[torch.Tensor]] = [] if isinstance(multimodal_outputs, torch.Tensor): @@ -145,13 +251,10 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): else: raise RuntimeError("Unsupported diffusion output type") - self.eplb_step() - - return OmniModelRunnerOutput( + output = OmniModelRunnerOutput( req_ids=self.input_batch.req_ids, req_id_to_index=self.input_batch.req_id_to_index, sampled_token_ids=[], - spec_token_ids=None, logprobs=None, prompt_logprobs_dict={}, pooler_output=pooler_output, @@ -159,6 +262,16 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): num_nans_in_logits={}, ) + if not self.use_async_scheduling: + return output + + return AsyncGPUModelRunnerOutput( + model_runner_output=output, + sampled_token_ids=[], + invalid_req_indices=[], + async_output_copy_stream=self.async_output_copy_stream, + ) + def _run_diffusion( self, *, @@ -202,84 +315,217 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): def _dummy_run( self, num_tokens: int, - capture_attn_cudagraph: bool = False, + cudagraph_runtime_mode: Optional[CUDAGraphMode] = None, + force_attention: bool = False, + uniform_decode: bool = False, + allow_microbatching: bool = True, skip_eplb: bool = False, is_profile: bool = False, + create_mixed_batch: bool = False, + remove_lora: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: - # Padding for DP - num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) - num_tokens += num_pad + """ + Run a dummy forward pass to warm up/profile run or capture the + CUDA graph for the model. + + Args: + num_tokens: Number of tokens to run the dummy forward pass. + cudagraph_runtime_mode: used to control the behavior. + - if not set will determine the cudagraph mode based on using + the self.cudagraph_dispatcher. + - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run + - CUDAGraphMode.PIECEWISE: Piecewise cudagraph. + - CUDAGraphMode.FULL: Full cudagraph, attention metadata is + needed. + force_attention: If True, always create attention metadata. Used to + warm up attention backend when mode is NONE. + uniform_decode: If True, the batch is a uniform decode batch. + skip_eplb: If True, skip EPLB state update. + is_profile: If True, this is a profile run. + create_mixed_batch: If True, create a mixed batch with both decode + (1 token) and prefill (multiple tokens) requests. + remove_lora: If False, dummy LoRAs are not destroyed after the run + """ + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in { + CUDAGraphMode.NONE, + CUDAGraphMode.PIECEWISE, + CUDAGraphMode.FULL, + } + + # If cudagraph_mode.decode_mode() == FULL and + # cudagraph_mode.separate_routine(). This means that we are using + # different graphs and/or modes for mixed prefill-decode batches vs. + # uniform decode batches. A uniform decode batch means that all + # requests have identical query length, except a potential virtual + # request (shorter) in the batch account for padding. + # Uniform decode batch could either be common pure decode, where + # max_query_len == 1, or speculative decode, where + # max_query_len == 1 + num_spec_decode_tokens. + + # When setting max_query_len = 1, we switch to and capture the optimized + # routine of FA2 for pure decode, i.e., Flashdecode + an optimization + # for GQA/MQA. + max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens # Set num_scheduled_tokens based on num_tokens and max_num_seqs # for dummy run with LoRA so that the num_reqs collectively # has num_tokens in total. assert num_tokens <= self.scheduler_config.max_num_batched_tokens max_num_reqs = self.scheduler_config.max_num_seqs - num_reqs = min(num_tokens, max_num_reqs) - min_tokens_per_req = num_tokens // num_reqs - num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs - num_scheduled_tokens_list[-1] += num_tokens % num_reqs + if create_mixed_batch: + assert not uniform_decode + # Create mixed batch: + # first half decode tokens, second half one prefill + num_decode_tokens = num_tokens // 2 + num_prefill_tokens = num_tokens - num_decode_tokens + num_reqs = num_decode_tokens + 1 + + # Create decode requests (1 token each) followed by prefill request + num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens] + # Note: Overriding max_query_len to be the prefill tokens + max_query_len = num_prefill_tokens + elif uniform_decode: + assert not create_mixed_batch + num_reqs = cdiv(num_tokens, max_query_len) + num_scheduled_tokens_list = [max_query_len] * num_reqs + if num_tokens % max_query_len != 0: + num_scheduled_tokens_list[-1] = num_tokens % max_query_len + else: + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs + assert sum(num_scheduled_tokens_list) == num_tokens assert len(num_scheduled_tokens_list) == num_reqs num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) - attn_metadata: Optional[dict[str, dict]] = None - if capture_attn_cudagraph: - attn_metadata = {} + total_num_scheduled_tokens = int(num_scheduled_tokens.sum()) - # Make sure max_model_len is used at the graph capture time. - self.seq_lens_np[:num_reqs] = self.max_model_len - self.seq_lens_np[num_reqs:] = 0 - self.seq_lens[:num_reqs].copy_( - self.seq_lens_cpu[:num_reqs], non_blocking=True + ubatch_slices = None + num_tokens_after_padding = None + + # We currently only microbatch if the number of tokens is + # over a certain threshold. + if self.parallel_config.enable_dbo and allow_microbatching: + ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split( + num_scheduled_tokens, + total_num_scheduled_tokens, + total_num_scheduled_tokens, + uniform_decode=uniform_decode, + vllm_config=self.vllm_config, ) + # Currently when DBO is enabled `ubatch_split` returns + # the num_tokens_after_padding for a single ubatch, but we have 2 + # TODO(sage,lucas): this is cruft that should be addressed in the + # padding refactor. + if ubatch_num_tokens_after_padding is not None: + num_tokens_after_padding = ubatch_num_tokens_after_padding * 2 + + # If we failed to microbatch, currently need to resynchronize + # TODO(lucas,sage): we should be able to avoid this second sync by + # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into + # a single `coordinate_batch_across_dp` function. + if num_tokens_after_padding is None: + num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) + num_tokens_after_padding = num_tokens + num_pad + else: + num_tokens_across_dp = num_tokens_after_padding + num_tokens_after_padding = int(num_tokens_after_padding[0].item()) + + attn_metadata: Optional[PerLayerAttnMetadata] = None + + # If force_attention is True, we always capture attention. Otherwise, + # it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + attn_metadata = {} + if ubatch_slices is not None: + attn_metadata = [dict() for _ in range(len(ubatch_slices))] + + if create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + else: + seq_lens = max_query_len + self.seq_lens.np[:num_reqs] = seq_lens + self.seq_lens.np[num_reqs:] = 0 + self.seq_lens.copy_to_gpu() + + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() for kv_cache_group_id, kv_cache_group_spec in enumerate( self.kv_cache_config.kv_cache_groups ): common_attn_metadata = CommonAttentionMetadata( - query_start_loc=self.query_start_loc[: num_reqs + 1], - query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1], - seq_lens=self.seq_lens[:num_reqs], - seq_lens_cpu=self.seq_lens_cpu[:num_reqs], + query_start_loc=self.query_start_loc.gpu[: num_reqs + 1], + query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1], + seq_lens=self.seq_lens.gpu[:num_reqs], + seq_lens_cpu=self.seq_lens.cpu[:num_reqs], num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[ # noqa: E501 :num_reqs ], num_reqs=num_reqs, num_actual_tokens=num_tokens, - max_query_len=num_tokens, + max_query_len=max_query_len, + max_seq_len=self.max_model_len, block_table_tensor=self.input_batch.block_table[ kv_cache_group_id - ].get_device_tensor()[:num_reqs], + ].get_device_tensor(num_reqs), slot_mapping=self.input_batch.block_table[ kv_cache_group_id - ].slot_mapping[:num_tokens], + ].slot_mapping.gpu[:num_tokens], causal=True, ) - for attn_group in self.attn_groups[kv_cache_group_id]: - attn_metadata_i = ( - attn_group.metadata_builder.build_for_cudagraph_capture( + if ubatch_slices is not None: + common_attn_metadata_list = split_attn_metadata( + ubatch_slices, common_attn_metadata + ) + for ubid, common_attn_metadata in enumerate( + common_attn_metadata_list + ): + assert common_attn_metadata.max_query_len == 1 + attn_metadata_i = attn_group.get_metadata_builder( + ubatch_id=ubid + ).build_for_cudagraph_capture(common_attn_metadata) + for layer_name in attn_group.layer_names: + assert type(attn_metadata) is list + attn_metadata[ubid][layer_name] = attn_metadata_i + else: + assert type(attn_metadata) is dict + attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture( # noqa: E501 common_attn_metadata ) - ) - for layer_name in kv_cache_group_spec.layer_names: - attn_metadata[layer_name] = attn_metadata_i + for layer_name in attn_group.layer_names: + attn_metadata[layer_name] = attn_metadata_i - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens): - if self.is_multimodal_model: + with self.maybe_dummy_run_with_lora( + self.lora_config, num_scheduled_tokens, remove_lora + ): + model_kwargs = self._init_model_kwargs(num_tokens) + if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: input_ids = None - inputs_embeds = self.inputs_embeds[:num_tokens] - model_mm_kwargs = self._dummy_mm_kwargs(num_reqs) + inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + model_kwargs = { + **model_kwargs, + **self._dummy_mm_kwargs(num_reqs), + } + elif self.enable_prompt_embeds: + input_ids = None + inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + model_kwargs = self._init_model_kwargs(num_tokens) else: - input_ids = self.input_ids[:num_tokens] + input_ids = self.input_ids.gpu[:num_tokens] inputs_embeds = None - model_mm_kwargs = {} if self.uses_mrope: - positions = self.mrope_positions[:, :num_tokens] + positions = self.mrope_positions.gpu[:, :num_tokens] else: - positions = self.positions[:num_tokens] + positions = self.positions.gpu[:num_tokens] if get_pp_group().is_first_rank: intermediate_tensors = None @@ -297,23 +543,53 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): num_tokens, None, False ) - # Diffusion path: avoid CUDA graphs; we only use context for resource wiring + # filter out the valid batch descriptor + _cg_mode, batch_descriptor = ( + self.cudagraph_dispatcher.dispatch( + BatchDescriptor( + num_tokens=num_tokens_after_padding, + uniform_decode=uniform_decode, + ) + ) + if not is_profile + else (CUDAGraphMode.NONE, None) + ) + if cudagraph_runtime_mode is not None: + # we allow forcing NONE when the dispatcher disagrees to support + # warm ups for cudagraph capture + assert ( + cudagraph_runtime_mode == CUDAGraphMode.NONE + or cudagraph_runtime_mode == _cg_mode + ), ( + f"Cudagraph runtime mode mismatch at dummy_run. " + f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}." + ) + else: + cudagraph_runtime_mode = _cg_mode + + if ubatch_slices is not None: + # Adjust values to reflect a single ubatch. + # TODO(sage,lucas): this is cruft that should be addressed in + # the padding refactor. + num_tokens_after_padding = ubatch_slices[0].num_tokens + if num_tokens_across_dp is not None: + num_tokens_across_dp[:] = num_tokens_after_padding + with self.maybe_randomize_inputs(input_ids), set_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_tokens, + num_tokens=num_tokens_after_padding, num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + ubatch_slices=ubatch_slices, ): outputs = self.model( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, - **MultiModalKwargs.as_kwargs( - model_mm_kwargs, - device=self.device, - ), - sampler=None, + **model_kwargs, ) if self.use_aux_hidden_state_outputs: @@ -321,16 +597,23 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): else: hidden_states = outputs - # Extract multimodal outputs if present; we ignore them here because - # dummy run returns tensors only. The actual diffusion runner returns - # multimodal outputs via pooler_output in execute_model. - text_hidden_states, _ = self.extract_multimodal_outputs(hidden_states) + if self.speculative_config and self.speculative_config.use_eagle(): + assert isinstance(self.drafter, EagleProposer) + self.drafter.dummy_run(num_tokens) - if not skip_eplb: - self.eplb_step(is_dummy=True, is_profile=is_profile) + # This is necessary to avoid blocking DP. + # For dummy runs, we typically skip EPLB since we don't have any real + # requests to process. + # However, in DP settings, there may be cases when some DP ranks do + # not have any requests to process, so they're executing dummy batches. + # In such cases, we still have to trigger EPLB to make sure + # ranks execute the rearrangement in synchronization. + if not skip_eplb: + self.eplb_step(is_dummy=True, is_profile=is_profile) - # logit_indices = np.cumsum(num_scheduled_tokens) - 1 # unused variable - return text_hidden_states, None + # logit_indices = np.cumsum(num_scheduled_tokens) - 1 + hidden_states, _ = self.extract_multimodal_outputs(hidden_states) + return hidden_states, None @torch.inference_mode() def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: @@ -338,42 +621,76 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): return None def profile_run(self) -> None: - # Profile with multimodal encoder & encoder cache, similar to base but - # without any logits/sampler warming. - if self.is_multimodal_model: - mm_budget = self.mm_budget - assert mm_budget is not None - - # TODO: handle encoder-decoder models once supported. - if (mm_budget.get_encoder_budget()) > 0: # encoder_budget unused - ( - dummy_modality, - max_tokens, - ) = mm_budget.get_modality_with_max_tokens() - ( - max_mm_items_per_prompt, - max_mm_items_per_batch, - ) = mm_budget.get_max_items(dummy_modality, max_tokens) - - batched_dummy_mm_inputs = self._get_mm_dummy_batch( - dummy_modality, - max_mm_items_per_batch, + # Profile with multimodal encoder & encoder cache. + if self.supports_mm_inputs: + if self.model_config.multimodal_config.skip_mm_profiling: + logger.info( + "Skipping memory profiling for multimodal encoder and " + "encoder cache." ) + else: + mm_budget = self.mm_budget + assert mm_budget is not None - dummy_encoder_outputs = self.model.get_multimodal_embeddings( - **batched_dummy_mm_inputs - ) + 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 + ] - sanity_check_mm_encoder_outputs( - dummy_encoder_outputs, - expected_num_items=max_mm_items_per_batch, - ) + 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, + ) - self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs)) + # 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.get_multimodal_embeddings( + **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 + # (encode_budget, hidden_size) and scatter encoder + # output into it. + encoder_output_shape = dummy_encoder_outputs[0].shape + if encoder_output_shape[0] < encoder_budget: + expanded_outputs = [] + for output in dummy_encoder_outputs: + expanded = output.new_zeros( + (encoder_budget, encoder_output_shape[-1]) + ) + 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)) + + # Add `is_profile` here to pre-allocate communication buffers hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True) if get_pp_group().is_last_rank: - pass # No sampler/pooler warmup for diffusion + pass self._sync_device() del hidden_states self.encoder_cache.clear() diff --git a/vllm_omni/worker/gpu_diffusion_worker.py b/vllm_omni/worker/gpu_diffusion_worker.py index 831881cf0..46af50190 100644 --- a/vllm_omni/worker/gpu_diffusion_worker.py +++ b/vllm_omni/worker/gpu_diffusion_worker.py @@ -2,34 +2,42 @@ import gc import os import torch + from vllm.model_executor import set_random_seed from vllm.platforms import current_platform from vllm.utils import GiB_bytes, MemorySnapshot from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment -from vllm.worker.worker import _check_if_gpu_supports_dtype - from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner class GPUDiffusionWorker(GPUWorker): def init_device(self): if self.device_config.device.type == "cuda": - # torch.distributed.all_reduce does not free the input tensor until - # the synchronization point. This causes the memory usage to grow - # as the number of all_reduce calls increases. This env var disables - # this behavior. - # Related issue: - # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 - os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" - # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) self.device = torch.device(f"cuda:{self.local_rank}") current_platform.set_device(self.device) - _check_if_gpu_supports_dtype(self.model_config.dtype) + current_platform.check_if_supports_dtype(self.model_config.dtype) + + # Initialize the distributed environment BEFORE taking + # memory snapshot + # This ensures NCCL buffers are allocated before we measure + # available memory + init_worker_distributed_environment( + self.vllm_config, + self.rank, + self.distributed_init_method, + self.local_rank, + current_platform.dist_backend, + ) + + # Set random seed. + set_random_seed(self.model_config.seed) + + # Now take memory snapshot after NCCL is initialized gc.collect() torch.cuda.empty_cache() @@ -55,16 +63,6 @@ class GPUDiffusionWorker(GPUWorker): ) else: raise RuntimeError(f"Not support device type: {self.device_config.device}") - # Initialize the distributed environment. - init_worker_distributed_environment( - self.vllm_config, - self.rank, - self.distributed_init_method, - self.local_rank, - current_platform.dist_backend, - ) - # Set random seed. - set_random_seed(self.model_config.seed) # Construct the model runner self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner( diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 0cba53e93..5369da899 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1,25 +1,23 @@ -from typing import TYPE_CHECKING, Any, Optional, Union, cast +from typing import TYPE_CHECKING, List, Optional, Union, cast import numpy as np import torch -import vllm.envs as envs -from vllm.distributed.kv_transfer import has_kv_transfer_group -from vllm.distributed.parallel_state import get_pp_group, get_tp_group -from vllm.forward_context import set_forward_context + +from vllm.config import CUDAGraphMode +from vllm.distributed.parallel_state import get_pp_group +from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.model_executor.models.interfaces_base import VllmModelForPooling -from vllm.multimodal.inputs import MultiModalKwargs from vllm.sampling_params import SamplingType -from vllm.sequence import IntermediateTensors -from vllm.utils import LazyLoader, round_up -from vllm.v1.attention.backends.utils import CommonAttentionMetadata -from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT +from vllm.utils import LazyLoader, cdiv +from vllm.v1.attention.backends.utils import ( + CommonAttentionMetadata, + split_attn_metadata, +) from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.worker.gpu_input_batch import CachedRequestState -from vllm.v1.worker.gpu_model_runner import GPUModelRunner - -from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding -from vllm_omni.outputs import OmniModelRunnerOutput +from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata +from vllm.v1.worker.ubatch_splitting import ubatch_split if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -48,7 +46,6 @@ class OmniGPUModelRunner(GPUModelRunner): # Remove finished requests from the cached states. for req_id in scheduler_output.finished_req_ids: self.requests.pop(req_id, None) - self.encoder_cache.pop(req_id, None) # Remove the finished requests from the persistent batch. # NOTE(woosuk): There could be an edge case where finished_req_ids and # scheduled_req_ids overlap. This happens when a request is aborted and @@ -59,12 +56,8 @@ class OmniGPUModelRunner(GPUModelRunner): self.input_batch.remove_request(req_id) # Free the cached encoder outputs. - for req_id, input_id in scheduler_output.free_encoder_input_ids: - encoder_outputs = self.encoder_cache.get(req_id) - if encoder_outputs is not None: - encoder_outputs.pop(input_id, None) - if not encoder_outputs: - self.encoder_cache.pop(req_id, None) + for mm_hash in scheduler_output.free_encoder_mm_hashes: + self.encoder_cache.pop(mm_hash, None) # Remove the unscheduled requests from the persistent batch. # NOTE(woosuk): The unscheduled requests are either preempted requests @@ -81,7 +74,7 @@ class OmniGPUModelRunner(GPUModelRunner): for req_id in unscheduled_req_ids: self.input_batch.remove_request(req_id) - req_ids_to_add: list[str] = [] + reqs_to_add: list[CachedRequestState] = [] # Add new requests to the cached states. for new_req_data in scheduler_output.scheduled_new_reqs: req_id = new_req_data.req_id @@ -97,20 +90,20 @@ class OmniGPUModelRunner(GPUModelRunner): else: generator = None - if pooling_params: - assert ( - task := pooling_params.task - ) is not None, "You did not set `task` in the API" + if self.is_pooling_model: + assert pooling_params is not None + task = pooling_params.task + assert task is not None, "You did not set `task` in the API" - model = cast(VllmModelForPooling, self.model) + model = cast(VllmModelForPooling, self.get_model()) to_update = model.pooler.get_pooling_updates(task) to_update.apply(pooling_params) - self.requests[req_id] = CachedRequestState( + req_state = CachedRequestState( req_id=req_id, prompt_token_ids=new_req_data.prompt_token_ids, - mm_inputs=new_req_data.mm_inputs, - mm_positions=new_req_data.mm_positions, + prompt_embeds=new_req_data.prompt_embeds, + mm_features=new_req_data.mm_features, sampling_params=sampling_params, pooling_params=pooling_params, generator=generator, @@ -119,6 +112,7 @@ class OmniGPUModelRunner(GPUModelRunner): output_token_ids=[], lora_request=new_req_data.lora_request, ) + self.requests[req_id] = req_state # If prompt embeddings are provided, decode and attach to inter_data try: @@ -136,8 +130,8 @@ class OmniGPUModelRunner(GPUModelRunner): new_req_data.prompt_embeds = pe_cpu # type: ignore[assignment] except Exception: pass - except Exception: - pass + except Exception as e: + logger.error(f"Error decoding prompt embeds: {e}") # Decode additional_information payloads (dictionary) try: if getattr(new_req_data, "additional_information", None) is not None: @@ -146,7 +140,7 @@ class OmniGPUModelRunner(GPUModelRunner): if isinstance(payload_info, dict): info_dict = payload_info else: - from vllm.v1.engine import AdditionalInformationPayload + from vllm_omni.engine import AdditionalInformationPayload if isinstance(payload_info, AdditionalInformationPayload): for k, entry in payload_info.entries.items(): @@ -165,44 +159,15 @@ class OmniGPUModelRunner(GPUModelRunner): "additional_information_cpu", info_dict, ) - except Exception: + except Exception as e: + logger.error(f"Error decoding additional information: {e}") pass # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: - image_grid_thw = [] - video_grid_thw = [] - second_per_grid_ts = [] - audio_feature_lengths = [] - use_audio_in_video = False - for mm_input in self.requests[req_id].mm_inputs: - if mm_input.get("image_grid_thw") is not None: - image_grid_thw.extend(mm_input["image_grid_thw"].tolist()) - if mm_input.get("video_grid_thw") is not None: - video_grid_thw.extend(mm_input["video_grid_thw"].tolist()) - if mm_input.get("second_per_grid_ts") is not None: - second_per_grid_ts.extend(mm_input["second_per_grid_ts"]) - if mm_input.get("audio_feature_lengths") is not None: - audio_feature_lengths.extend(mm_input["audio_feature_lengths"]) - if mm_input.get("use_audio_in_video") is True: - use_audio_in_video = True + self._init_mrope_positions(req_state) - hf_config = self.model_config.hf_config - - ( - self.requests[req_id].mrope_positions, - self.requests[req_id].mrope_position_delta, - ) = MRotaryEmbedding.get_input_positions_tensor( - self.requests[req_id].prompt_token_ids, - hf_config=hf_config, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - audio_feature_lengths=audio_feature_lengths, - use_audio_in_video=use_audio_in_video, - ) - - req_ids_to_add.append(req_id) + reqs_to_add.append(self.requests[req_id]) # Update the states of the running/resumed requests. is_last_rank = get_pp_group().is_last_rank @@ -234,10 +199,12 @@ class OmniGPUModelRunner(GPUModelRunner): # Update the block IDs. if not resumed_from_preemption: - # Append the new blocks to the existing block IDs. - for block_ids, new_ids in zip(req_state.block_ids, new_block_ids): - block_ids.extend(new_ids) + if new_block_ids is not None: + # Append the new blocks to the existing block IDs. + for block_ids, new_ids in zip(req_state.block_ids, new_block_ids): + block_ids.extend(new_ids) else: + assert new_block_ids is not None # The request is resumed from preemption. # Replace the existing block IDs with the new ones. req_state.block_ids = new_block_ids @@ -247,12 +214,13 @@ class OmniGPUModelRunner(GPUModelRunner): # The request is not in the persistent batch. # The request was either preempted and resumed later, or was not # scheduled in the previous step and needs to be added again. - req_ids_to_add.append(req_id) + reqs_to_add.append(req_state) continue # Update the persistent batch. self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens - self.input_batch.block_table.append_row(new_block_ids, req_index) + if new_block_ids is not None: + self.input_batch.block_table.append_row(new_block_ids, req_index) # For the last rank, we don't need to update the token_ids_cpu # because the sampled tokens are already cached. @@ -282,9 +250,8 @@ class OmniGPUModelRunner(GPUModelRunner): # Add the new or resumed requests to the persistent batch. # The smaller empty indices are filled first. - for req_id in req_ids_to_add: - req_state = self.requests[req_id] - self.input_batch.add_request(req_state) + for request in reqs_to_add: + self.input_batch.add_request(request) # Condense the batched states if there are gaps left by removed requests self.input_batch.condense() @@ -293,356 +260,10 @@ class OmniGPUModelRunner(GPUModelRunner): # Refresh batch metadata with any pending updates. self.input_batch.refresh_metadata() - @torch.inference_mode() - def execute_model( - self, - scheduler_output: "SchedulerOutput", - intermediate_tensors: Optional[IntermediateTensors] = None, - ) -> Union[OmniModelRunnerOutput, IntermediateTensors]: - self._update_states(scheduler_output) - if not scheduler_output.total_num_scheduled_tokens: - if not has_kv_transfer_group(): - # Return empty ModelRunnerOutput if there's no work to do. - return EMPTY_MODEL_RUNNER_OUTPUT - - return self.kv_connector_no_forward(scheduler_output, self.vllm_config) - - # Prepare the decoder inputs. - ( - attn_metadata, - attention_cuda_graphs, - logits_indices, - spec_decode_metadata, - num_scheduled_tokens_np, - spec_decode_common_attn_metadata, - ) = self._prepare_inputs(scheduler_output) - - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens - if ( - self.use_cuda_graph - and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1] - ): - # Use piecewise CUDA graphs. - # Add padding to the batch size. - num_input_tokens = self.vllm_config.pad_for_cudagraph(num_scheduled_tokens) - else: - # Eager mode. - # Pad tokens to multiple of tensor_parallel_size when - # enabled collective fusion for SP - tp_size = self.vllm_config.parallel_config.tensor_parallel_size - if ( - self.compilation_config.pass_config.enable_sequence_parallelism - and tp_size > 1 - ): - num_input_tokens = round_up(num_scheduled_tokens, tp_size) - else: - num_input_tokens = num_scheduled_tokens - - # Padding for DP - num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens) - num_input_tokens += num_pad - - # _prepare_inputs may reorder the batch, so we must gather multi - # modal outputs after that to ensure the correct order - if self.is_multimodal_model: - # Run the multimodal encoder if any. - self._execute_mm_encoder(scheduler_output) - mm_embeds = self._gather_mm_embeddings(scheduler_output) - else: - mm_embeds = [] - - if get_pp_group().is_first_rank: - # Always prepare inputs_embeds on the first PP rank so both - # multimodal and text-only models can consume embeddings. - inputs_embeds_scheduled = self.model.get_input_embeddings( - input_ids=self.input_ids[:num_scheduled_tokens], - multimodal_embeddings=( - (mm_embeds or None) if self.is_multimodal_model else None - ), - ) - - # TODO(woosuk): Avoid the copy. Optimize. - self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) - - # Overlay per-request custom prompt_embeds for the prompt portion - for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests[req_id] - pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) - if pe_cpu is None: - continue - num_computed_tokens = int( - self.input_batch.num_computed_tokens_cpu[req_index] - ) - prompt_len = len(req_state.prompt_token_ids) - prompt_remaining = max(0, prompt_len - num_computed_tokens) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - overlay_len = min(sched_tokens, prompt_remaining) - if overlay_len <= 0: - continue - src = pe_cpu[ - num_computed_tokens : num_computed_tokens + overlay_len - ].to(dtype=self.dtype, device=self.device, non_blocking=True) - start_offset = int(self.query_start_loc_cpu[req_index]) - self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) - - input_ids = self.input_ids[:num_input_tokens] # preserved for APIs - inputs_embeds = self.inputs_embeds[:num_input_tokens] - model_mm_kwargs = ( - self._extract_mm_kwargs(scheduler_output) - if self.is_multimodal_model - else {} - ) - else: - # For non-first PP ranks, use token ids as usual; embeddings are - # only consumed on the first rank. - input_ids = self.input_ids[:num_input_tokens] - inputs_embeds = None - model_mm_kwargs = {} - if self.uses_mrope: - positions = self.mrope_positions[:, :num_input_tokens] - else: - positions = self.positions[:num_input_tokens] - - if get_pp_group().is_first_rank: - intermediate_tensors = None - else: - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_input_tokens, intermediate_tensors, True - ) - - # Some attention backends only support CUDA Graphs in pure decode. - # If attention doesn't support CUDA Graphs for this batch, but we - # compiled with full CUDA graphs, we have to skip them entirely. - skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs - - # Run the model. - # Use persistent buffers for CUDA graphs. - with set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_input_tokens, - num_tokens_across_dp=num_tokens_across_dp, - skip_cuda_graphs=skip_cuda_graphs, - ), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output: - - model_output = self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **MultiModalKwargs.as_kwargs( - model_mm_kwargs, - device=self.device, - ), - sampling_metadata=self.input_batch.sampling_metadata, - logits_index=logits_indices, - sampler=self.sampler, - ) - - if self.use_aux_hidden_state_outputs: - hidden_states, aux_hidden_states = model_output - else: - hidden_states = model_output - aux_hidden_states = None - - text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs( - hidden_states - ) - - # Broadcast PP output for external_launcher (torchrun) - # to make sure we are synced across pp ranks - # TODO: Support overlapping mirco-batches - # https://github.com/vllm-project/vllm/issues/18019 - broadcast_pp_output = ( - self.parallel_config.distributed_executor_backend == "external_launcher" - and len(get_pp_group().ranks) > 0 - ) - if not get_pp_group().is_last_rank: - # For mid-pipeline stages, return the hidden states. - assert isinstance(text_hidden_states, IntermediateTensors) - if not broadcast_pp_output: - text_hidden_states.kv_connector_output = kv_connector_output - return text_hidden_states - get_pp_group().send_tensor_dict( - text_hidden_states.tensors, all_gather_group=get_tp_group() - ) - logits = None - else: - if self.input_batch.pooling_params: - return self._pool( - text_hidden_states, - num_scheduled_tokens, - num_scheduled_tokens_np, - kv_connector_output, - ) - - sample_hidden_states = text_hidden_states[logits_indices] - logits = self.model.compute_logits(sample_hidden_states, None) - if broadcast_pp_output: - model_output_broadcast_data = ( - { - "logits": logits.contiguous(), - } - if logits is not None - else {} - ) - model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( - model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 - ) - assert model_output_broadcast_data is not None - logits = model_output_broadcast_data["logits"] - - # Apply structured output bitmasks if present - if scheduler_output.grammar_bitmask is not None: - self.apply_grammar_bitmask(scheduler_output, logits) - - # Sample the next token and get logprobs if needed. - sampling_metadata = self.input_batch.sampling_metadata - if spec_decode_metadata is None: - sampler_output = self.sampler( - logits=logits, - sampling_metadata=sampling_metadata, - ) - else: - # When indexing with a tensor (bonus_logits_indices), PyTorch - # creates a new tensor with separate storage from the original - # logits tensor. This means any in-place operations on bonus_logits - # won't affect the original logits tensor. - assert logits is not None - bonus_logits = logits[spec_decode_metadata.bonus_logits_indices] - sampler_output = self.sampler( - logits=bonus_logits, - sampling_metadata=sampling_metadata, - ) - bonus_token_ids = sampler_output.sampled_token_ids - - # Just like `bonus_logits`, `target_logits` is a new tensor with - # separate storage from the original `logits` tensor. Therefore, - # it is safe to update `target_logits` in place. - target_logits = logits[spec_decode_metadata.target_logits_indices] - output_token_ids = self.rejection_sampler( - spec_decode_metadata, - None, # draft_probs - target_logits, - bonus_token_ids, - sampling_metadata, - ) - sampler_output.sampled_token_ids = output_token_ids - - num_nans_in_logits = {} - if envs.VLLM_COMPUTE_NANS_IN_LOGITS: - num_nans_in_logits = self._get_nans_in_logits(logits) - - # TODO(woosuk): The following loop can be slow since it iterates over - # the requests one by one. Optimize. - discard_sampled_tokens_req_indices = [] - for i, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests[req_id] - seq_len = ( - req_state.num_computed_tokens - + scheduler_output.num_scheduled_tokens[req_id] - ) - if seq_len < req_state.num_tokens: - # Ignore the sampled token for partial prefills. - # Rewind the generator state as if the token was not sampled. - # This relies on cuda-specific torch-internal impl details - generator = self.input_batch.generators.get(i) - if generator is not None: - generator.set_offset(generator.get_offset() - 4) - # Record the index of the request that should not be sampled, - # so that we could clear the sampled tokens before returning. - discard_sampled_tokens_req_indices.append(i) - - # NOTE: GPU -> CPU Sync happens here. - # Move as many CPU operations as possible before this sync point. - logprobs_tensors = sampler_output.logprobs_tensors - logprobs_lists = ( - logprobs_tensors.tolists() if logprobs_tensors is not None else None - ) - - # Compute prompt logprobs if needed. - prompt_logprobs_dict = self._get_prompt_logprobs_dict( - text_hidden_states[:num_scheduled_tokens], - scheduler_output, - ) - - # Get the valid generated tokens. - # import os # unused import - - sampled_token_ids = sampler_output.sampled_token_ids - max_gen_len = sampled_token_ids.shape[-1] - if max_gen_len == 1: - # No spec decode tokens. - valid_sampled_token_ids = sampled_token_ids.tolist() - else: - # Includes spec decode tokens. - valid_sampled_token_ids = self.rejection_sampler.parse_output( - sampled_token_ids, - self.input_batch.vocab_size, - ) - # Mask out the sampled tokens that should not be sampled. - for i in discard_sampled_tokens_req_indices: - valid_sampled_token_ids[i].clear() - - # Cache the sampled tokens in the model runner, so that the scheduler - # doesn't need to send them back. - # NOTE(woosuk): As an exception, when using PP, the scheduler sends - # the sampled tokens back, because there's no direct communication - # between the first-stage worker and the last-stage worker. - for req_idx, sampled_ids in enumerate(valid_sampled_token_ids): - if not sampled_ids: - continue - - start_idx = self.input_batch.num_tokens_no_spec[req_idx] - end_idx = start_idx + len(sampled_ids) - assert end_idx <= self.max_model_len, ( - "Sampled token IDs exceed the max model length. " - f"Total number of tokens: {end_idx} > max_model_len: " - f"{self.max_model_len}" - ) - - self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids - self.input_batch.num_tokens_no_spec[req_idx] = end_idx - self.input_batch.num_tokens[req_idx] = end_idx - req_id = self.input_batch.req_ids[req_idx] - req_state = self.requests[req_id] - req_state.output_token_ids.extend(sampled_ids) - - if not self.speculative_config: - # Speculative decoding is not enabled. - spec_token_ids = None - else: - assert spec_decode_common_attn_metadata is not None - spec_token_ids = self.propose_draft_token_ids( - scheduler_output, - valid_sampled_token_ids, - sampling_metadata, - text_hidden_states, - sample_hidden_states, - aux_hidden_states, - spec_decode_metadata, - spec_decode_common_attn_metadata, - ) - - self.eplb_step() - - return OmniModelRunnerOutput( - req_ids=self.input_batch.req_ids, - req_id_to_index=self.input_batch.req_id_to_index, - sampled_token_ids=valid_sampled_token_ids, - spec_token_ids=spec_token_ids, - logprobs=logprobs_lists, - prompt_logprobs_dict=prompt_logprobs_dict, - pooler_output=[], - kv_connector_output=kv_connector_output, - num_nans_in_logits=num_nans_in_logits, - multimodal_outputs=multimodal_outputs, - ) - @torch.inference_mode() def extract_multimodal_outputs( - self, hidden_states: torch.Tensor - ) -> tuple[torch.Tensor, dict]: + self, hidden_states: Union[torch.Tensor, List[torch.Tensor]] + ) -> dict: if ( hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs @@ -650,94 +271,231 @@ class OmniGPUModelRunner(GPUModelRunner): text_hidden_states = hidden_states.text_hidden_states multimodal_outputs = hidden_states.multimodal_outputs - else: + elif isinstance(hidden_states, torch.Tensor): text_hidden_states = hidden_states multimodal_outputs = {} + elif isinstance(hidden_states, List): + text_hidden_states = hidden_states[0] + multimodal_outputs = {} + else: + raise ValueError(f"Invalid hidden states type: {type(hidden_states)}") return text_hidden_states, multimodal_outputs @torch.inference_mode() def _dummy_run( self, num_tokens: int, - capture_attn_cudagraph: bool = False, + cudagraph_runtime_mode: Optional[CUDAGraphMode] = None, + force_attention: bool = False, + uniform_decode: bool = False, + allow_microbatching: bool = True, skip_eplb: bool = False, is_profile: bool = False, + create_mixed_batch: bool = False, + remove_lora: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Run a dummy forward pass to warm up/profile run or capture the + CUDA graph for the model. - # Padding for DP - num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) - num_tokens += num_pad + Args: + num_tokens: Number of tokens to run the dummy forward pass. + cudagraph_runtime_mode: used to control the behavior. + - if not set will determine the cudagraph mode based on using + the self.cudagraph_dispatcher. + - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run + - CUDAGraphMode.PIECEWISE: Piecewise cudagraph. + - CUDAGraphMode.FULL: Full cudagraph, attention metadata is + needed. + force_attention: If True, always create attention metadata. Used to + warm up attention backend when mode is NONE. + uniform_decode: If True, the batch is a uniform decode batch. + skip_eplb: If True, skip EPLB state update. + is_profile: If True, this is a profile run. + create_mixed_batch: If True, create a mixed batch with both decode + (1 token) and prefill (multiple tokens) requests. + remove_lora: If False, dummy LoRAs are not destroyed after the run + """ + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in { + CUDAGraphMode.NONE, + CUDAGraphMode.PIECEWISE, + CUDAGraphMode.FULL, + } + + # If cudagraph_mode.decode_mode() == FULL and + # cudagraph_mode.separate_routine(). This means that we are using + # different graphs and/or modes for mixed prefill-decode batches vs. + # uniform decode batches. A uniform decode batch means that all + # requests have identical query length, except a potential virtual + # request (shorter) in the batch account for padding. + # Uniform decode batch could either be common pure decode, where + # max_query_len == 1, or speculative decode, where + # max_query_len == 1 + num_spec_decode_tokens. + + # When setting max_query_len = 1, we switch to and capture the optimized + # routine of FA2 for pure decode, i.e., Flashdecode + an optimization + # for GQA/MQA. + max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens # Set num_scheduled_tokens based on num_tokens and max_num_seqs # for dummy run with LoRA so that the num_reqs collectively # has num_tokens in total. assert num_tokens <= self.scheduler_config.max_num_batched_tokens max_num_reqs = self.scheduler_config.max_num_seqs - num_reqs = min(num_tokens, max_num_reqs) - min_tokens_per_req = num_tokens // num_reqs - num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs - num_scheduled_tokens_list[-1] += num_tokens % num_reqs + if create_mixed_batch: + assert not uniform_decode + # Create mixed batch: + # first half decode tokens, second half one prefill + num_decode_tokens = num_tokens // 2 + num_prefill_tokens = num_tokens - num_decode_tokens + num_reqs = num_decode_tokens + 1 + + # Create decode requests (1 token each) followed by prefill request + num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens] + # Note: Overriding max_query_len to be the prefill tokens + max_query_len = num_prefill_tokens + elif uniform_decode: + assert not create_mixed_batch + num_reqs = cdiv(num_tokens, max_query_len) + num_scheduled_tokens_list = [max_query_len] * num_reqs + if num_tokens % max_query_len != 0: + num_scheduled_tokens_list[-1] = num_tokens % max_query_len + else: + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs + assert sum(num_scheduled_tokens_list) == num_tokens assert len(num_scheduled_tokens_list) == num_reqs num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) - attn_metadata: Optional[dict[str, Any]] = None - if capture_attn_cudagraph: - attn_metadata = {} + total_num_scheduled_tokens = int(num_scheduled_tokens.sum()) - # Make sure max_model_len is used at the graph capture time. - self.seq_lens_np[:num_reqs] = self.max_model_len - self.seq_lens_np[num_reqs:] = 0 - self.seq_lens[:num_reqs].copy_( - self.seq_lens_cpu[:num_reqs], non_blocking=True + ubatch_slices = None + num_tokens_after_padding = None + + # We currently only microbatch if the number of tokens is + # over a certain threshold. + if self.parallel_config.enable_dbo and allow_microbatching: + ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split( + num_scheduled_tokens, + total_num_scheduled_tokens, + total_num_scheduled_tokens, + uniform_decode=uniform_decode, + vllm_config=self.vllm_config, ) + # Currently when DBO is enabled `ubatch_split` returns + # the num_tokens_after_padding for a single ubatch, but we have 2 + # TODO(sage,lucas): this is cruft that should be addressed in the + # padding refactor. + if ubatch_num_tokens_after_padding is not None: + num_tokens_after_padding = ubatch_num_tokens_after_padding * 2 + + # If we failed to microbatch, currently need to resynchronize + # TODO(lucas,sage): we should be able to avoid this second sync by + # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into + # a single `coordinate_batch_across_dp` function. + if num_tokens_after_padding is None: + num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) + num_tokens_after_padding = num_tokens + num_pad + else: + num_tokens_across_dp = num_tokens_after_padding + num_tokens_after_padding = int(num_tokens_after_padding[0].item()) + + attn_metadata: Optional[PerLayerAttnMetadata] = None + + # If force_attention is True, we always capture attention. Otherwise, + # it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + attn_metadata = {} + if ubatch_slices is not None: + attn_metadata = [dict() for _ in range(len(ubatch_slices))] + + if create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + else: + seq_lens = max_query_len + self.seq_lens.np[:num_reqs] = seq_lens + self.seq_lens.np[num_reqs:] = 0 + self.seq_lens.copy_to_gpu() + + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() for kv_cache_group_id, kv_cache_group_spec in enumerate( self.kv_cache_config.kv_cache_groups ): common_attn_metadata = CommonAttentionMetadata( - query_start_loc=self.query_start_loc[: num_reqs + 1], - query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1], - seq_lens=self.seq_lens[:num_reqs], - seq_lens_cpu=self.seq_lens_cpu[:num_reqs], + query_start_loc=self.query_start_loc.gpu[: num_reqs + 1], + query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1], + seq_lens=self.seq_lens.gpu[:num_reqs], + seq_lens_cpu=self.seq_lens.cpu[:num_reqs], num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[ # noqa: E501 :num_reqs ], num_reqs=num_reqs, num_actual_tokens=num_tokens, - max_query_len=num_tokens, + max_query_len=max_query_len, + max_seq_len=self.max_model_len, block_table_tensor=self.input_batch.block_table[ kv_cache_group_id - ].get_device_tensor()[:num_reqs], + ].get_device_tensor(num_reqs), slot_mapping=self.input_batch.block_table[ kv_cache_group_id - ].slot_mapping[:num_tokens], + ].slot_mapping.gpu[:num_tokens], causal=True, ) - for attn_group in self.attn_groups[kv_cache_group_id]: - attn_metadata_i = ( - attn_group.metadata_builder.build_for_cudagraph_capture( + if ubatch_slices is not None: + common_attn_metadata_list = split_attn_metadata( + ubatch_slices, common_attn_metadata + ) + for ubid, common_attn_metadata in enumerate( + common_attn_metadata_list + ): + assert common_attn_metadata.max_query_len == 1 + attn_metadata_i = attn_group.get_metadata_builder( + ubatch_id=ubid + ).build_for_cudagraph_capture(common_attn_metadata) + for layer_name in attn_group.layer_names: + assert type(attn_metadata) is list + attn_metadata[ubid][layer_name] = attn_metadata_i + else: + assert type(attn_metadata) is dict + attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture( # noqa: E501 common_attn_metadata ) - ) - for layer_name in kv_cache_group_spec.layer_names: - attn_metadata[layer_name] = attn_metadata_i + for layer_name in attn_group.layer_names: + attn_metadata[layer_name] = attn_metadata_i - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens): - if self.is_multimodal_model: + with self.maybe_dummy_run_with_lora( + self.lora_config, num_scheduled_tokens, remove_lora + ): + model_kwargs = self._init_model_kwargs(num_tokens) + if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: input_ids = None - inputs_embeds = self.inputs_embeds[:num_tokens] - model_mm_kwargs = self._dummy_mm_kwargs(num_reqs) + inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + model_kwargs = { + **model_kwargs, + **self._dummy_mm_kwargs(num_reqs), + } + elif self.enable_prompt_embeds: + input_ids = None + inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + model_kwargs = self._init_model_kwargs(num_tokens) else: - input_ids = self.input_ids[:num_tokens] + input_ids = self.input_ids.gpu[:num_tokens] inputs_embeds = None - model_mm_kwargs = {} if self.uses_mrope: - positions = self.mrope_positions[:, :num_tokens] + positions = self.mrope_positions.gpu[:, :num_tokens] else: - positions = self.positions[:num_tokens] + positions = self.positions.gpu[:num_tokens] if get_pp_group().is_first_rank: intermediate_tensors = None @@ -755,22 +513,53 @@ class OmniGPUModelRunner(GPUModelRunner): num_tokens, None, False ) + # filter out the valid batch descriptor + _cg_mode, batch_descriptor = ( + self.cudagraph_dispatcher.dispatch( + BatchDescriptor( + num_tokens=num_tokens_after_padding, + uniform_decode=uniform_decode, + ) + ) + if not is_profile + else (CUDAGraphMode.NONE, None) + ) + if cudagraph_runtime_mode is not None: + # we allow forcing NONE when the dispatcher disagrees to support + # warm ups for cudagraph capture + assert ( + cudagraph_runtime_mode == CUDAGraphMode.NONE + or cudagraph_runtime_mode == _cg_mode + ), ( + f"Cudagraph runtime mode mismatch at dummy_run. " + f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}." + ) + else: + cudagraph_runtime_mode = _cg_mode + + if ubatch_slices is not None: + # Adjust values to reflect a single ubatch. + # TODO(sage,lucas): this is cruft that should be addressed in + # the padding refactor. + num_tokens_after_padding = ubatch_slices[0].num_tokens + if num_tokens_across_dp is not None: + num_tokens_across_dp[:] = num_tokens_after_padding + with self.maybe_randomize_inputs(input_ids), set_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_tokens, + num_tokens=num_tokens_after_padding, num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + ubatch_slices=ubatch_slices, ): outputs = self.model( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, - **MultiModalKwargs.as_kwargs( - model_mm_kwargs, - device=self.device, - ), - sampler=None, + **model_kwargs, ) if self.use_aux_hidden_state_outputs: @@ -778,10 +567,6 @@ class OmniGPUModelRunner(GPUModelRunner): else: hidden_states = outputs - text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs( - hidden_states - ) - if self.speculative_config and self.speculative_config.use_eagle(): assert isinstance(self.drafter, EagleProposer) self.drafter.dummy_run(num_tokens) @@ -797,4 +582,7 @@ class OmniGPUModelRunner(GPUModelRunner): self.eplb_step(is_dummy=True, is_profile=is_profile) logit_indices = np.cumsum(num_scheduled_tokens) - 1 - return text_hidden_states, text_hidden_states[logit_indices] + hidden_states, multimodal_outputs = self.extract_multimodal_outputs( + hidden_states + ) + return hidden_states, hidden_states[logit_indices] diff --git a/vllm_omni/worker/model_runner.py b/vllm_omni/worker/model_runner.py deleted file mode 100644 index 4ff3ce5eb..000000000 --- a/vllm_omni/worker/model_runner.py +++ /dev/null @@ -1,216 +0,0 @@ -import itertools -from dataclasses import dataclass -from typing import Any, Dict, List, Optional - -import torch -from vllm.lora.layers import LoRAMapping -from vllm.sequence import MultiModalKwargs -from vllm.utils import async_tensor_h2d, flatten_2d_lists -from vllm.worker.model_runner import ModelInputForGPU, ModelInputForGPUBuilder - - -@dataclass(frozen=True) -class OmniModelInputForGPU(ModelInputForGPU): - # New: optional additional information dict for current scheduled tokens - additional_information: Optional[Dict[str, Any]] = None - - -class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder): - def build(self) -> OmniModelInputForGPU: - """Finalize the builder intermediate data and - create on-device tensors. - """ - # Combine and flatten intermediate data. - input_tokens = list[int]() - inputs_embeds_list = list[torch.Tensor]() - token_types = list[int]() - for inter_data in self.inter_data_list: - for cur_input_tokens in inter_data.input_tokens: - input_tokens.extend(cur_input_tokens) - for cur_token_types in inter_data.token_types: - token_types.extend(cur_token_types) - if inter_data.inputs_embeds is not None: - inputs_embeds_list.append( - inter_data.inputs_embeds.to( - dtype=self.runner.model_config.dtype, device=self.runner.device - ) - ) - # Support v1 direct-transfer prompt embeds attached on request state - if inter_data.inputs_embeds is None: - try: - # Locate req_id -> request state to fetch decoded CPU embeds - req_id = inter_data.request_id # type: ignore[attr-defined] - req_state = getattr(self.runner, "requests", {}).get(req_id) - pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) - if pe_cpu is not None: - inter_data.inputs_embeds = pe_cpu.to( - dtype=self.runner.model_config.dtype, - device=self.runner.device, - ) - inputs_embeds_list.append(inter_data.inputs_embeds) - except Exception: - pass - inputs_embeds: Optional[torch.Tensor] - if len(inputs_embeds_list) == 0: - inputs_embeds = None - else: - inputs_embeds = torch.cat(inputs_embeds_list, dim=0).to( - dtype=self.runner.model_config.dtype, device=self.runner.device - ) - assert len(inputs_embeds) == len(input_tokens) - - if not input_tokens and inputs_embeds is None: - # This may happen when all prefill requests hit - # prefix caching and there is no decode request. - return self.model_input_cls() - - mrope_input_positions: Optional[List[List[int]]] = None - if any( - inter_data.mrope_input_positions is not None - for inter_data in self.inter_data_list - ): - mrope_input_positions = [[] for _ in range(3)] - for idx in range(3): - for inter_data in self.inter_data_list: - msections = inter_data.mrope_input_positions - if msections is None: - for _seq_input_positions in inter_data.input_positions: - mrope_input_positions[idx].extend(_seq_input_positions) - else: - for _seq_mrope_input_positions in msections: - mrope_input_positions[idx].extend( - _seq_mrope_input_positions[idx] - ) - input_positions = None - else: - input_positions = [] - for inter_data in self.inter_data_list: - for cur_input_positions in inter_data.input_positions: - input_positions.extend(cur_input_positions) - - seq_lens = [] - query_lens = [] - max_decode_seq_len = 0 - max_encoder_seq_len = 0 - for inter_data in self.inter_data_list: - seq_lens.extend(inter_data.seq_lens) - query_lens.extend(inter_data.query_lens) - if not inter_data.is_prompt: - max_decode_seq_len = max(max_decode_seq_len, max(inter_data.seq_lens)) - if self.runner.model_config.is_encoder_decoder: - max_encoder_seq_len = max( - max_encoder_seq_len, inter_data.encoder_seq_len - ) - - # Mapping from request IDs to sequence IDs. Used for Jamba models - # that manages the cache by itself. - request_ids_to_seq_ids = { - data.request_id: data.seq_ids for data in self.inter_data_list - } - - cuda_graph_pad_size = self._get_cuda_graph_pad_size( - num_seqs=len(seq_lens), - max_decode_seq_len=max_decode_seq_len, - max_encoder_seq_len=max_encoder_seq_len, - ) - - batch_size = len(input_tokens) - if cuda_graph_pad_size != -1: - # If cuda graph can be used, pad tensors accordingly. - # See `capture_model` API for more details. - # vLLM uses cuda graph only for decoding requests. - batch_size += cuda_graph_pad_size - - # Tokens and positions. - if cuda_graph_pad_size: - input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size)) - assert self.runner.device is not None - input_tokens_tensor = async_tensor_h2d( - input_tokens, torch.long, self.runner.device, self.runner.pin_memory - ) - - token_types_tensor = ( - async_tensor_h2d( - token_types, torch.long, self.runner.device, self.runner.pin_memory - ) - if token_types - else None - ) - - if mrope_input_positions is not None: - for idx in range(3): - mrope_input_positions[idx].extend( - itertools.repeat(0, cuda_graph_pad_size) - ) - input_positions_tensor = async_tensor_h2d( - mrope_input_positions, - torch.long, - self.runner.device, - self.runner.pin_memory, - ) - else: - input_positions.extend(itertools.repeat(0, cuda_graph_pad_size)) - input_positions_tensor = async_tensor_h2d( - input_positions, torch.long, self.runner.device, self.runner.pin_memory - ) - # Sequence and query lengths. - if cuda_graph_pad_size: - seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size)) - - # Attention metadata. - attn_metadata = self.attn_metadata_builder.build( - seq_lens, query_lens, cuda_graph_pad_size, batch_size - ) - - # LoRA data. - lora_requests = set() - lora_mapping = None - if self.enable_lora: - lora_requests = set( - r for data in self.inter_data_list for r in data.lora_requests - ) - lora_index_mapping = flatten_2d_lists( - [ - flatten_2d_lists(inter_data.lora_index_mapping) - for inter_data in self.inter_data_list - ] - ) - if cuda_graph_pad_size: - lora_index_mapping.extend(itertools.repeat(0, cuda_graph_pad_size)) - lora_prompt_mapping = flatten_2d_lists( - [ - flatten_2d_lists(inter_data.lora_prompt_mapping) - for inter_data in self.inter_data_list - ] - ) - - lora_mapping = LoRAMapping( - **dict( - index_mapping=lora_index_mapping, - prompt_mapping=lora_prompt_mapping, - is_prefill=not self.decode_only, - ) - ) - - # Multi-modal data. - multi_modal_kwargs_list = [ - data.multi_modal_kwargs - for data in self.inter_data_list - if data.multi_modal_kwargs is not None - ] - multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list) - - return self.model_input_cls( - input_tokens=input_tokens_tensor, - inputs_embeds=inputs_embeds, - input_positions=input_positions_tensor, - token_types=token_types_tensor, - attn_metadata=attn_metadata, - seq_lens=seq_lens, - query_lens=query_lens, - lora_mapping=lora_mapping, - lora_requests=lora_requests, - multi_modal_kwargs=multi_modal_kwargs, - request_ids_to_seq_ids=request_ids_to_seq_ids, - finished_requests_ids=self.finished_requests_ids, - )