anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Dependency refractored to vLLM v0.11.0

Browse Source 
Signed-off-by: Gaohan123 <gaohan19@huawei.com>
This commit is contained in:
Gaohan123 committed 2025-11-04 22:06:41 +08:00
1 parent d4a364f052
commit 6e1478222e
38 files changed
+1500 -1719
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

README.md
+23 -7
View File
@@ -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
docs/architecture/high level arch design.md → docs/architecture/high_level_arch_design.md
View File
File renamed without changes.
examples/offline_inference/qwen_2_5_omni/README.md
+23 -7
View File
@@ -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
examples/offline_inference/qwen_2_5_omni/end2end.py
+1 -1
View File
@@ -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"
examples/offline_inference/qwen_2_5_omni/run.sh
-1
View File
@@ -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" \
examples/offline_inference/qwen_2_5_omni/utils.py
+1 -1
View File
@@ -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
requirements.txt
-10
View File
@@ -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
vllm_omni/config/__init__.py
+1 -27
View File
@@ -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",
vllm_omni/config/model.py
+26
View File
@@ -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]
vllm_omni/core/sched/diffusion_scheduler.py
+48 -41
View File
@@ -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
vllm_omni/core/sched/output.py
+1 -4
View File
@@ -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,
vllm_omni/core/sched/scheduler.py
+1 -3
View File
@@ -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,
vllm_omni/engine/__init__.py
+1
View File
@@ -5,6 +5,7 @@ Engine components for vLLM-omni.
from typing import Any, Optional
import msgspec
from vllm.v1.engine import EngineCoreRequest
vllm_omni/engine/arg_utils.py
-27
View File
@@ -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()
vllm_omni/engine/output_processor.py
+19 -20
View File
@@ -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
vllm_omni/engine/processor.py
+78 -70
View File
@@ -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,
)
vllm_omni/entrypoints/omni_llm.py
+34 -13
View File
@@ -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
)
vllm_omni/entrypoints/omni_stage.py
-1
View File
@@ -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
vllm_omni/entrypoints/utils.py
+1
View File
@@ -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)
vllm_omni/inputs/data.py
+1 -3
View File
@@ -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:
vllm_omni/inputs/preprocess.py
+13 -91
View File
@@ -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)
vllm_omni/model_executor/layers/mrope.py
+9
View File
@@ -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
vllm_omni/model_executor/model_loader/weight_utils.py
+19 -10
View File
@@ -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(
vllm_omni/model_executor/models/qwen2_5_omni.py
+7 -8
View File
@@ -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,
vllm_omni/model_executor/models/qwen2_5_omni_talker.py
+6 -9
View File
@@ -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,
vllm_omni/model_executor/models/qwen2_5_omni_thinker.py
+3 -8
View File
@@ -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(
vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py
+6 -8
View File
@@ -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
vllm_omni/model_executor/models/qwen2_old.py
+7 -6
View File
@@ -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(
vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py
+1 -1
View File
@@ -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
vllm_omni/outputs.py
+1
View File
@@ -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
vllm_omni/patch.py
-1
View File
@@ -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
vllm_omni/request.py
+12 -12
View File
@@ -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,
)
vllm_omni/worker/gpu_ar_model_runner.py
+401 -458
View File
@@ -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]
vllm_omni/worker/gpu_ar_worker.py
+19 -21
View File
@@ -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(
vllm_omni/worker/gpu_diffusion_model_runner.py
+459 -142
View File
@@ -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()
vllm_omni/worker/gpu_diffusion_worker.py
+19 -21
View File
@@ -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(
vllm_omni/worker/gpu_model_runner.py
+259 -471
View File
@@ -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]
vllm_omni/worker/model_runner.py
-216
View File
@@ -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,
)