[Engine] Remove Diffusion_XX which duplicates with Generation_XX (#163)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Co-authored-by: Gao Han <gaohan19@huawei.com>
This commit is contained in:
Zhou TaichangandGao Han authored and GitHub committed
1 parent
f3c69df419
commit
c70b167f70
8 files changed
+15
-992
No files matched your search
@@ -2,12 +2,12 @@
|
||||
Scheduling components for vLLM-Omni.
|
||||
"""
|
||||
|
||||
from .diffusion_scheduler import DiffusionScheduler
|
||||
from .omni_ar_scheduler import OmniARScheduler
|
||||
from .omni_generation_scheduler import OmniGenerationScheduler
|
||||
from .output import OmniNewRequestData
|
||||
from .scheduler import OmniScheduler
|
||||
|
||||
__all__ = [
|
||||
"OmniScheduler",
|
||||
"DiffusionScheduler",
|
||||
"OmniARScheduler",
|
||||
"OmniGenerationScheduler",
|
||||
"OmniNewRequestData",
|
||||
]
|
||||
@@ -1,315 +0,0 @@
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Optional
|
||||
|
||||
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 Request, RequestStatus, SchedulerOutput, SpecDecodingStats
|
||||
from vllm.v1.core.sched.utils import remove_all
|
||||
from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
|
||||
|
||||
from vllm_omni.core.sched.output import OmniNewRequestData
|
||||
from vllm_omni.core.sched.scheduler import OmniScheduler
|
||||
from vllm_omni.outputs import OmniModelRunnerOutput
|
||||
|
||||
|
||||
class DiffusionScheduler(OmniScheduler):
|
||||
"""Scheduler for diffusion-based non-autoregressive stages.
|
||||
|
||||
Implements a fast-path scheduling strategy for diffusion models where
|
||||
all input tokens are fed at once. Falls back to standard vLLM scheduling
|
||||
if the token budget cannot accommodate all tokens in a single iteration.
|
||||
"""
|
||||
|
||||
def schedule(self) -> SchedulerOutput:
|
||||
"""Schedule requests for diffusion model execution.
|
||||
|
||||
Uses a fast-path approach that feeds all input tokens at once
|
||||
(or 1 placeholder token if input is empty). Falls back to
|
||||
standard scheduling if token budget is insufficient.
|
||||
|
||||
Returns:
|
||||
SchedulerOutput containing scheduled requests and token allocations
|
||||
"""
|
||||
|
||||
# Select requests with zero prompt and using pooling
|
||||
# (diffusion results returned via pooler_output)
|
||||
token_budget = self.max_num_scheduled_tokens
|
||||
capacity = self.max_num_running_reqs - len(self.running)
|
||||
scheduled_timestamp = time.monotonic()
|
||||
|
||||
scheduled_new_reqs: list[Request] = []
|
||||
scheduled_resumed_reqs: list[Request] = []
|
||||
scheduled_running_reqs: list[Request] = []
|
||||
|
||||
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]] = {}
|
||||
|
||||
# Temporary queue: preserve waiting order, do not disturb non-diffusion requests
|
||||
skipped_waiting_requests = create_request_queue(self.policy)
|
||||
|
||||
# Fast path selection and scheduling (treat all as diffusion requests,
|
||||
# independent of pooling_params)
|
||||
while self.waiting and token_budget > 0 and capacity > 0:
|
||||
request = self.waiting.peek_request()
|
||||
# Uniformly treat as diffusion. A feature flag can be added later
|
||||
# via config or request tag.
|
||||
|
||||
# Allocate all input tokens for the request in one shot
|
||||
# (allocate 1 placeholder if zero)
|
||||
required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
|
||||
if required_tokens > token_budget:
|
||||
# Insufficient budget to process all inputs at once;
|
||||
# stop fast path attempt
|
||||
break
|
||||
num_new_tokens = required_tokens
|
||||
new_blocks = self.kv_cache_manager.allocate_slots(
|
||||
request,
|
||||
num_new_tokens,
|
||||
num_lookahead_tokens=self.num_lookahead_tokens,
|
||||
)
|
||||
if new_blocks is None:
|
||||
# Allocation failed (e.g., VRAM pressure); stop fast path and
|
||||
# fall back to default scheduling
|
||||
# Put the current request back to the head of the waiting queue
|
||||
# Note: the original queue order is preserved
|
||||
break
|
||||
|
||||
# Officially schedule this request
|
||||
request = self.waiting.pop_request()
|
||||
self.running.append(request)
|
||||
request.status = RequestStatus.RUNNING
|
||||
if self.log_stats:
|
||||
request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
|
||||
|
||||
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
|
||||
scheduled_new_reqs.append(request)
|
||||
|
||||
# Return skipped waiting requests
|
||||
if skipped_waiting_requests:
|
||||
self.waiting.prepend_requests(skipped_waiting_requests)
|
||||
|
||||
# If fast path scheduled none, fall back to the original scheduling
|
||||
if not num_scheduled_tokens:
|
||||
return super().schedule()
|
||||
|
||||
# Compute common prefix blocks (aligned with v1)
|
||||
num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
|
||||
if self.running:
|
||||
any_request = self.running[0]
|
||||
num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
|
||||
any_request, len(self.running)
|
||||
)
|
||||
|
||||
# Assemble SchedulerOutput
|
||||
new_reqs_data = [
|
||||
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(
|
||||
scheduled_running_reqs,
|
||||
scheduled_resumed_reqs,
|
||||
num_scheduled_tokens,
|
||||
scheduled_spec_decode_tokens,
|
||||
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())
|
||||
scheduler_output = SchedulerOutput(
|
||||
scheduled_new_reqs=new_reqs_data,
|
||||
scheduled_cached_reqs=cached_reqs_data,
|
||||
num_scheduled_tokens=num_scheduled_tokens,
|
||||
total_num_scheduled_tokens=total_num_scheduled_tokens,
|
||||
scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
|
||||
scheduled_encoder_inputs=scheduled_encoder_inputs,
|
||||
num_common_prefix_blocks=num_common_prefix_blocks,
|
||||
finished_req_ids=self.finished_req_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,
|
||||
)
|
||||
|
||||
# KVTransfer: package metadata
|
||||
if self.connector is not None:
|
||||
meta = self.connector.build_connector_meta(scheduler_output)
|
||||
scheduler_output.kv_connector_metadata = meta
|
||||
|
||||
# Update internal state (advance num_computed_tokens, free encoder inputs,
|
||||
# etc.)
|
||||
self._update_after_schedule(scheduler_output)
|
||||
return scheduler_output
|
||||
|
||||
"""
|
||||
Scheduler for the diffusion model.
|
||||
This scheduler is modified to stop the request immediately for the diffusion model.
|
||||
This is because the diffusion model can generate the final image/audio in one step.
|
||||
Note: This is just a minimal modification to the original scheduler,
|
||||
and there should be some further efforts to optimize the scheduler.
|
||||
The original scheduler is still used for the AR model.
|
||||
"""
|
||||
|
||||
def update_from_output(
|
||||
self,
|
||||
scheduler_output: SchedulerOutput,
|
||||
model_runner_output: OmniModelRunnerOutput,
|
||||
) -> dict[int, EngineCoreOutputs]:
|
||||
"""Update the scheduler state based on the model runner output.
|
||||
|
||||
This method is modified to stop the request immediately for the diffusion model.
|
||||
"""
|
||||
sampled_token_ids = model_runner_output.sampled_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
|
||||
# to avoid expensive operations inside the loop.
|
||||
stopped_running_reqs: set[Request] = set()
|
||||
stopped_preempted_reqs: set[Request] = set()
|
||||
for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
|
||||
assert num_tokens_scheduled > 0
|
||||
request = self.requests.get(req_id)
|
||||
if request is None:
|
||||
# The request is already finished. This can happen if the
|
||||
# request is aborted while the model is executing it (e.g.,
|
||||
# in pipeline parallelism).
|
||||
continue
|
||||
|
||||
req_index = model_runner_output.req_id_to_index[req_id]
|
||||
generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
|
||||
|
||||
scheduled_spec_token_ids = 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.
|
||||
request.num_computed_tokens -= num_rejected
|
||||
spec_decoding_stats = self.make_spec_decoding_stats(
|
||||
spec_decoding_stats,
|
||||
num_draft_tokens=num_draft_tokens,
|
||||
num_accepted_tokens=num_accepted,
|
||||
)
|
||||
|
||||
new_logprobs = None
|
||||
new_token_ids = generated_token_ids
|
||||
kv_transfer_params = None
|
||||
status_before_stop = request.status
|
||||
pooler_output = None
|
||||
if pooler_outputs:
|
||||
pooler_output = pooler_outputs[req_index]
|
||||
|
||||
# Diffusion request: completes in one step; mark finished and free resources
|
||||
request.status = RequestStatus.FINISHED_STOPPED
|
||||
# Optional: set a stop_reason for front-end clarity
|
||||
# (does not affect protocol)
|
||||
request.stop_reason = request.stop_reason or "diffusion_done"
|
||||
kv_transfer_params = self._free_request(request)
|
||||
if status_before_stop == RequestStatus.RUNNING:
|
||||
stopped_running_reqs.add(request)
|
||||
else:
|
||||
stopped_preempted_reqs.add(request)
|
||||
|
||||
# Extract sample logprobs if needed.
|
||||
if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
|
||||
# NOTE: once we support N tokens per step (spec decode),
|
||||
# the outer lists can be of length > 1.
|
||||
new_logprobs = logprobs.slice(req_index, req_index + 1)
|
||||
|
||||
if new_token_ids and self.structured_output_manager.should_advance(request):
|
||||
# NOTE: structured_output_request should not be None if
|
||||
# use_structured_output, we have check above, so safe to ignore
|
||||
# type warning
|
||||
request.structured_output_request.grammar.accept_tokens(req_id, new_token_ids) # type: ignore[union-attr] # noqa: E501
|
||||
|
||||
# spec_token_ids comes from the model runner output
|
||||
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]
|
||||
|
||||
# Get prompt logprobs for this request.
|
||||
prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
|
||||
# Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
|
||||
pooling_output_dict = None
|
||||
if pooler_output is not None:
|
||||
# Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
|
||||
# which expects Optional[dict[str, torch.Tensor]]
|
||||
pooling_output_dict = {"model_outputs": pooler_output}
|
||||
if new_token_ids or pooler_output is not None or kv_transfer_params:
|
||||
# Add EngineCoreOutput for this Request.
|
||||
outputs[request.client_index].append(
|
||||
EngineCoreOutput(
|
||||
request_id=req_id,
|
||||
new_token_ids=new_token_ids,
|
||||
finish_reason=request.get_finished_reason(),
|
||||
new_logprobs=new_logprobs,
|
||||
new_prompt_logprobs_tensors=prompt_logprobs_tensors,
|
||||
pooling_output=pooling_output_dict,
|
||||
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:
|
||||
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)
|
||||
|
||||
# KV Connector: update state for finished KV Transfers.
|
||||
if model_runner_output.kv_connector_output:
|
||||
self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
|
||||
|
||||
# Create EngineCoreOutputs for all clients that have requests with
|
||||
# outputs in this step.
|
||||
engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
|
||||
|
||||
finished_req_ids = self.finished_req_ids_dict
|
||||
if finished_req_ids:
|
||||
# Include ids of requests that finished since last outputs
|
||||
# were sent.
|
||||
for client_index, finished_set in finished_req_ids.items():
|
||||
# Set finished request set in EngineCoreOutputs for this client.
|
||||
if (eco := engine_core_outputs.get(client_index)) is not None:
|
||||
eco.finished_requests = finished_set
|
||||
else:
|
||||
engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
|
||||
finished_req_ids.clear()
|
||||
|
||||
if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
|
||||
# Return stats to only one of the front-ends.
|
||||
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
|
||||
@@ -12,7 +12,7 @@ from vllm.v1.request import Request, RequestStatus
|
||||
from vllm.v1.spec_decode.metrics import SpecDecodingStats
|
||||
|
||||
|
||||
class OmniScheduler(VLLMScheduler):
|
||||
class OmniARScheduler(VLLMScheduler):
|
||||
"""
|
||||
OmniScheduler: Scheduler for vLLM-Omni multimodal processing.
|
||||
|
||||
+2
-2
@@ -10,15 +10,15 @@ from vllm.v1.core.sched.scheduler import (
|
||||
SchedulerOutput,
|
||||
SpecDecodingStats,
|
||||
)
|
||||
from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
|
||||
from vllm.v1.core.sched.utils import remove_all
|
||||
from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
|
||||
|
||||
from vllm_omni.core.sched.output import OmniNewRequestData
|
||||
from vllm_omni.core.sched.scheduler import OmniScheduler
|
||||
from vllm_omni.outputs import OmniModelRunnerOutput
|
||||
|
||||
|
||||
class GenerationScheduler(OmniScheduler):
|
||||
class OmniGenerationScheduler(VLLMScheduler):
|
||||
def schedule(self) -> SchedulerOutput:
|
||||
"""Diffusion fast path:
|
||||
- Feed all input tokens of the request at once
|
||||
@@ -11,7 +11,7 @@ stage_args:
|
||||
model_stage: thinker
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: true # Now we only support eager mode
|
||||
trust_remote_code: true
|
||||
@@ -37,7 +37,7 @@ stage_args:
|
||||
model_stage: talker
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
@@ -62,8 +62,8 @@ stage_args:
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
|
||||
scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
|
||||
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
gpu_memory_utilization: 0.15
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
|
||||
@@ -13,7 +13,7 @@ stage_args:
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.6
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
@@ -42,8 +42,8 @@ stage_args:
|
||||
model_stage: talker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
|
||||
gpu_memory_utilization: 0.3
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.5
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent # Output codec codes for code2wav
|
||||
@@ -72,7 +72,7 @@ stage_args:
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
|
||||
scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
|
||||
@@ -1,576 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
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.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_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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GPUDiffusionModelRunner(OmniGPUModelRunner):
|
||||
"""Diffusion model runner for vLLM-Omni (non-autoregressive).
|
||||
|
||||
- Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
|
||||
- Does not compute logits or perform token sampling.
|
||||
- Executes diffusion process and returns tensors via `pooler_output`.
|
||||
"""
|
||||
|
||||
def _preprocess(
|
||||
self,
|
||||
scheduler_output: SchedulerOutput,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
ubatch_slices: UBatchSlices | None = None,
|
||||
num_tokens_after_padding: torch.Tensor | None = None,
|
||||
) -> tuple[
|
||||
int,
|
||||
int,
|
||||
torch.Tensor | None,
|
||||
torch.Tensor | None,
|
||||
torch.Tensor | None,
|
||||
torch.Tensor,
|
||||
IntermediateTensors | None,
|
||||
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_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.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)
|
||||
|
||||
# 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.gpu[:num_input_tokens],
|
||||
multimodal_embeddings=mm_embeds or None,
|
||||
)
|
||||
|
||||
# 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)
|
||||
# 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:
|
||||
# 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_kwargs = self._init_model_kwargs(num_input_tokens)
|
||||
if self.uses_mrope:
|
||||
positions = self.mrope_positions.gpu[:, :num_input_tokens]
|
||||
else:
|
||||
positions = self.positions.gpu[: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
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
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: IntermediateTensors | None = None,
|
||||
) -> 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_kwargs,
|
||||
logits_indices=logits_indices,
|
||||
)
|
||||
|
||||
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
|
||||
# Ensure one tensor per request, map to CPU for output struct
|
||||
pooler_output: list[torch.Tensor | None] = []
|
||||
if isinstance(multimodal_outputs, torch.Tensor):
|
||||
# If model returned a single stacked tensor, split by requests
|
||||
assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
|
||||
for i in range(self.input_batch.num_reqs):
|
||||
pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
|
||||
elif isinstance(multimodal_outputs, list):
|
||||
for out in multimodal_outputs:
|
||||
pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
|
||||
elif isinstance(multimodal_outputs, dict):
|
||||
for out in multimodal_outputs.values():
|
||||
pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
|
||||
else:
|
||||
raise RuntimeError("Unsupported diffusion output type")
|
||||
|
||||
output = OmniModelRunnerOutput(
|
||||
req_ids=self.input_batch.req_ids,
|
||||
req_id_to_index=self.input_batch.req_id_to_index,
|
||||
sampled_token_ids=[],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=pooler_output,
|
||||
kv_connector_output=kv_connector_output,
|
||||
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,
|
||||
*,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None,
|
||||
inputs_embeds: torch.Tensor | None,
|
||||
multimodal_kwargs: dict,
|
||||
logits_indices: torch.Tensor,
|
||||
) -> torch.Tensor | list[torch.Tensor]:
|
||||
"""Runs the diffusion process and returns per-request tensors.
|
||||
|
||||
Tries model interfaces in the following order for maximal compatibility:
|
||||
1) model.sample(condition=..., **kwargs)
|
||||
2) model.forward(condition=..., **kwargs)
|
||||
3) model.diffuse(condition=..., **kwargs)
|
||||
"""
|
||||
# Keep inputs identical to AR runner
|
||||
kwargs = dict(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
**MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
|
||||
sampling_metadata=self.input_batch.sampling_metadata,
|
||||
logits_index=logits_indices,
|
||||
sampler=self.sampler,
|
||||
)
|
||||
|
||||
# For Qwen 2.5 Omni's current implementation, we only support the forward method
|
||||
if hasattr(self.model, "forward"):
|
||||
return self.model.forward(**kwargs)
|
||||
# TODO: add the diffuse method for other models
|
||||
|
||||
raise RuntimeError(
|
||||
"The loaded model does not expose diffusion interfaces 'sample', 'forward', or 'diffuse'. "
|
||||
"Please implement one of them or adapt the runner."
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def _dummy_run(
|
||||
self,
|
||||
num_tokens: int,
|
||||
cudagraph_runtime_mode: CUDAGraphMode | None = None,
|
||||
force_attention: bool = False,
|
||||
uniform_decode: bool = False,
|
||||
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.
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
num_tokens_after_padding = None
|
||||
|
||||
# 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: PerLayerAttnMetadata | None = 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 = {}
|
||||
|
||||
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.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[:num_reqs], # noqa: E501
|
||||
num_reqs=num_reqs,
|
||||
num_actual_tokens=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),
|
||||
slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
|
||||
causal=True,
|
||||
)
|
||||
for attn_group in self.attn_groups[kv_cache_group_id]:
|
||||
assert type(attn_metadata) is dict
|
||||
attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
|
||||
common_attn_metadata
|
||||
) # noqa: E501
|
||||
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, 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.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.gpu[:num_tokens]
|
||||
inputs_embeds = None
|
||||
|
||||
if self.uses_mrope:
|
||||
positions = self.mrope_positions.gpu[:, :num_tokens]
|
||||
else:
|
||||
positions = self.positions.gpu[: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)
|
||||
|
||||
# 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
|
||||
|
||||
with (
|
||||
self.maybe_randomize_inputs(input_ids),
|
||||
set_forward_context(
|
||||
attn_metadata,
|
||||
self.vllm_config,
|
||||
num_tokens=num_tokens_after_padding,
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
||||
batch_descriptor=batch_descriptor,
|
||||
),
|
||||
):
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
**model_kwargs,
|
||||
)
|
||||
|
||||
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.
|
||||
|
||||
# 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:
|
||||
logger.warning("Dummy sampler run is not implemented for diffusion model")
|
||||
return None
|
||||
|
||||
def profile_run(self) -> None:
|
||||
# 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
|
||||
|
||||
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]
|
||||
|
||||
logger.info(
|
||||
"Encoder cache will be initialized with a budget "
|
||||
"of %s tokens, and profiled with %s %s items of the maximum feature size.",
|
||||
encoder_budget,
|
||||
max_mm_items_per_batch,
|
||||
dummy_modality,
|
||||
)
|
||||
|
||||
# Create dummy batch of multimodal inputs.
|
||||
batched_dummy_mm_inputs = self._get_mm_dummy_batch(
|
||||
dummy_modality,
|
||||
max_mm_items_per_batch,
|
||||
)
|
||||
|
||||
# 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
|
||||
self._sync_device()
|
||||
del hidden_states
|
||||
self.encoder_cache.clear()
|
||||
gc.collect()
|
||||
@@ -1,86 +0,0 @@
|
||||
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_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
|
||||
|
||||
|
||||
class GPUDiffusionWorker(GPUWorker):
|
||||
"""GPU worker for diffusion-based omni model stages.
|
||||
|
||||
Extends the base GPUWorker to initialize and manage diffusion model
|
||||
runners for non-autoregressive generation stages (e.g., image or
|
||||
audio generation stages).
|
||||
"""
|
||||
|
||||
def init_device(self):
|
||||
"""Initialize the GPU device and diffusion model runner.
|
||||
|
||||
Sets up CUDA device, initializes distributed environment,
|
||||
sets random seed, and creates a GPUDiffusionModelRunner instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If free GPU memory is insufficient for the
|
||||
requested memory utilization
|
||||
RuntimeError: If device type is not supported
|
||||
"""
|
||||
if self.device_config.device.type == "cuda":
|
||||
# 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)
|
||||
|
||||
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()
|
||||
|
||||
# take current memory snapshot
|
||||
self.init_snapshot = MemorySnapshot()
|
||||
self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
|
||||
if self.init_snapshot.free_memory < self.requested_memory:
|
||||
|
||||
def GiB(b):
|
||||
return round(b / GiB_bytes, 2)
|
||||
|
||||
raise ValueError(
|
||||
f"Free memory on device "
|
||||
f"({GiB(self.init_snapshot.free_memory)}/"
|
||||
f"{GiB(self.init_snapshot.total_memory)} GiB) on startup "
|
||||
f"is less than desired GPU memory utilization "
|
||||
f"({self.cache_config.gpu_memory_utilization}, "
|
||||
f"{GiB(self.requested_memory)} GiB). Decrease GPU memory "
|
||||
f"utilization or reduce GPU memory used by other processes."
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(f"Not support device type: {self.device_config.device}")
|
||||
|
||||
# Construct the model runner
|
||||
self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(self.vllm_config, self.device)
|
||||
|
||||
if self.rank == 0:
|
||||
# If usage stat is enabled, collect relevant info.
|
||||
report_usage_stats(self.vllm_config)
|
||||
Reference in new issue
Block a user