from __future__ import annotations import gc import logging from typing import List, Optional, Union import numpy as np import torch 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.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, IntermediateTensors, get_pp_group, set_forward_context, ) 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`. """ @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: 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) # 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_input_tokens += num_pad # Multimodal conditioning (e.g., text/audio/video encoders) if self.is_multimodal_model: 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: inputs_embeds_scheduled = self.model.get_input_embeddings( input_ids=self.input_ids[: scheduler_output.total_num_scheduled_tokens], multimodal_embeddings=mm_embeds or None, ) self.inputs_embeds[: scheduler_output.total_num_scheduled_tokens].copy_( inputs_embeds_scheduled ) 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) else: input_ids = self.input_ids[:num_input_tokens] inputs_embeds = None model_mm_kwargs = {} # Positions (mrope or standard) if self.uses_mrope: positions = self.mrope_positions[:, :num_input_tokens] else: positions = self.positions[:num_input_tokens] # Intermediate tensors sync for PP (if any) 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 ) # 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 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 outputs = self._run_diffusion( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, multimodal_kwargs=model_mm_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[Optional[torch.Tensor]] = [] 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") 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=[], spec_token_ids=None, logprobs=None, prompt_logprobs_dict={}, pooler_output=pooler_output, kv_connector_output=kv_connector_output, num_nans_in_logits={}, ) def _run_diffusion( self, *, input_ids: torch.Tensor, positions: torch.Tensor, intermediate_tensors: Optional[IntermediateTensors], inputs_embeds: Optional[torch.Tensor], multimodal_kwargs: dict, logits_indices: torch.Tensor, ) -> Union[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, 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, dict]] = 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 ) # Diffusion path: avoid CUDA graphs; we only use context for resource wiring 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, ), sampler=None, ) if self.use_aux_hidden_state_outputs: hidden_states, _ = outputs 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 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 @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, 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, ) 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, ) self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs)) 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 self._sync_device() del hidden_states self.encoder_cache.clear() gc.collect()