diff --git a/vllm_omni/core/sched/__init__.py b/vllm_omni/core/sched/__init__.py index 9b3a34ca5..ecf18d07a 100644 --- a/vllm_omni/core/sched/__init__.py +++ b/vllm_omni/core/sched/__init__.py @@ -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", ] diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py deleted file mode 100644 index 56b7be68c..000000000 --- a/vllm_omni/core/sched/diffusion_scheduler.py +++ /dev/null @@ -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 diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py similarity index 99% rename from vllm_omni/core/sched/scheduler.py rename to vllm_omni/core/sched/omni_ar_scheduler.py index 1e480c6e1..99bde7638 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -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. diff --git a/vllm_omni/core/sched/generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py similarity index 99% rename from vllm_omni/core/sched/generation_scheduler.py rename to vllm_omni/core/sched/omni_generation_scheduler.py index 9a0c07266..59f954a8a 100644 --- a/vllm_omni/core/sched/generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index 8441db598..6192c0510 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml index 0c9fe7ccb..e73ad0205 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -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 diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu_diffusion_model_runner.py deleted file mode 100644 index 401ff91bd..000000000 --- a/vllm_omni/worker/gpu_diffusion_model_runner.py +++ /dev/null @@ -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() diff --git a/vllm_omni/worker/gpu_diffusion_worker.py b/vllm_omni/worker/gpu_diffusion_worker.py deleted file mode 100644 index 4731e26f1..000000000 --- a/vllm_omni/worker/gpu_diffusion_worker.py +++ /dev/null @@ -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)