"""AR GPU Model Runner for vLLM-omni. Exposes per-request hidden representations via ModelRunnerOutput.pooler_output and also outputs sampled tokens. """ from __future__ import annotations from typing import Any, List, Optional, Union import numpy as np import torch from vllm import envs from vllm.multimodal.inputs import MultiModalKwargs from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, IntermediateTensors, get_pp_group, get_tp_group, has_kv_transfer_group, set_forward_context, ) from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner class GPUARModelRunner(OmniGPUModelRunner): """Autoregressive GPU model runner that returns hidden states per request. This runner follows the same preparation and forward path as GPUModelRunner (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs), and additionally performs lightweight sampling so that sampled tokens are available in outputs. Hidden representations are taken at the same indices that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`). """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.input_ids = torch.zeros( self.max_num_tokens, dtype=torch.int64, device=self.device ) @torch.inference_mode() def execute_model( self, scheduler_output: "SchedulerOutput", intermediate_tensors: Optional[IntermediateTensors] = None, ) -> Union[OmniModelRunnerOutput, IntermediateTensors]: # Update internal state with the new schedule self._update_states(scheduler_output) # Decode per-request prompt_embeds / additional_hidden_states payloads # (if present) into CPU tensors try: new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) if new_reqs: for nr in new_reqs: req_id = getattr(nr, "req_id", None) or getattr( nr, "request_id", None ) if req_id is None: continue # prompt_embeds payload_pe = getattr(nr, "prompt_embeds", None) if payload_pe is not None: if isinstance(payload_pe, torch.Tensor): pe_cpu = payload_pe.detach().to("cpu").contiguous() elif isinstance(payload_pe, PromptEmbedsPayload): dt = np.dtype(getattr(payload_pe, "dtype", "float32")) arr = np.frombuffer(payload_pe.data, dtype=dt) arr = arr.reshape(payload_pe.shape) pe_cpu = torch.from_numpy(arr) else: pe_cpu = None if pe_cpu is not None and req_id in self.requests: setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) # additional_information payload_info = getattr(nr, "additional_information", None) if payload_info is not None: info_dict = {} if isinstance(payload_info, dict): # Already decoded info_dict = payload_info elif isinstance(payload_info, AdditionalInformationPayload): for k, entry in payload_info.entries.items(): if entry.tensor_data is not None: dt = np.dtype( getattr(entry, "tensor_dtype", "float32") ) arr = np.frombuffer(entry.tensor_data, dtype=dt) arr = arr.reshape(entry.tensor_shape) info_dict[k] = torch.from_numpy(arr) else: info_dict[k] = entry.list_data if info_dict and req_id in self.requests: setattr( self.requests[req_id], "additional_information_cpu", info_dict, ) except Exception: pass # If there's no work to do, either return empty output or kv-only path if not scheduler_output.total_num_scheduled_tokens: if not has_kv_transfer_group(): return EMPTY_MODEL_RUNNER_OUTPUT return self.kv_connector_no_forward(scheduler_output, self.vllm_config) # Prepare decoder inputs and attention metadata ( attn_metadata, attention_cuda_graphs, logits_indices, spec_decode_metadata, num_scheduled_tokens_np, spec_decode_common_attn_metadata, ) = self._prepare_inputs(scheduler_output) # Determine number of input tokens for this iteration num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens if ( self.use_cuda_graph and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1] ): num_input_tokens = self.vllm_config.pad_for_cudagraph(num_scheduled_tokens) else: tp_size = self.vllm_config.parallel_config.tensor_parallel_size if ( self.compilation_config.pass_config.enable_sequence_parallelism and tp_size > 1 ): from vllm.utils import round_up # lazy local import num_input_tokens = round_up(num_scheduled_tokens, tp_size) else: num_input_tokens = num_scheduled_tokens # DP padding num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens) num_input_tokens += num_pad # Multimodal handling (encode and gather embeddings if needed) if self.is_multimodal_model: self._execute_mm_encoder(scheduler_output) mm_embeds = self._gather_mm_embeddings(scheduler_output) else: mm_embeds = [] # Always assemble inputs_embeds on first PP rank; overlay per-request # prompt_embeds and collect additional_hidden_states for prefill only if get_pp_group().is_first_rank: inputs_embeds_scheduled = self.model.get_input_embeddings( input_ids=self.input_ids[:num_scheduled_tokens], multimodal_embeddings=( (mm_embeds or None) if self.is_multimodal_model else None ), ) # Copy into persistent buffer to enable CUDA Graph capture self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) # Reset per-step additional information collector if hasattr(self, "_forward_additional_information"): self._forward_additional_information = None # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only for req_index, req_id in enumerate(self.input_batch.req_ids): req_state = self.requests[req_id] pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) addi_cpu = getattr(req_state, "additional_information_cpu", None) num_computed_tokens = int( self.input_batch.num_computed_tokens_cpu[req_index] ) prompt_len = len(req_state.prompt_token_ids) prompt_remaining = max(0, prompt_len - num_computed_tokens) sched_tokens = int(num_scheduled_tokens_np[req_index]) overlay_len = min(sched_tokens, prompt_remaining) if overlay_len <= 0: continue if pe_cpu is not None: src = pe_cpu[ num_computed_tokens : num_computed_tokens + overlay_len ].to(dtype=self.dtype, device=self.device, non_blocking=True) start_offset = int(self.query_start_loc_cpu[req_index]) self.inputs_embeds[start_offset : start_offset + overlay_len].copy_( src ) # For additional_information: handle arbitrary keys if addi_cpu is not None and isinstance(addi_cpu, dict): # Lazy init collector dict if ( not hasattr(self, "_forward_additional_information") or self._forward_additional_information is None ): self._forward_additional_information = {} # Process tensors (slice by scheduled prompt range) and # lists (append per-request) for k, v in addi_cpu.items(): if isinstance(v, torch.Tensor): # Slice along token dimension for prefill part try: seg = v[ num_computed_tokens : num_computed_tokens + overlay_len ].to( dtype=self.dtype, device=self.device, non_blocking=True, ) except Exception: # Fallback: move whole tensor if slicing fails seg = v.to( dtype=self.dtype, device=self.device, non_blocking=True, ) prev_val = self._forward_additional_information.get(k) self._forward_additional_information[k] = ( torch.cat([prev_val, seg], dim=0) if isinstance(prev_val, torch.Tensor) else seg.clone() ) elif isinstance(v, list): prev_val = self._forward_additional_information.get(k) if prev_val is None: self._forward_additional_information[k] = [v] elif isinstance(prev_val, list): self._forward_additional_information[k].append(v) else: # Mixed types: wrap existing into list self._forward_additional_information[k] = [prev_val, v] input_ids = self.input_ids[:num_input_tokens] # preserved for APIs inputs_embeds = self.inputs_embeds[:num_input_tokens] model_mm_kwargs = ( self._extract_mm_kwargs(scheduler_output) if self.is_multimodal_model else {} ) else: input_ids = self.input_ids[:num_input_tokens] inputs_embeds = None model_mm_kwargs = {} # Positions/mRoPE if self.uses_mrope: positions = self.mrope_positions[:, :num_input_tokens] else: positions = self.positions[:num_input_tokens] # Handle pipeline-parallel intermediate tensors if get_pp_group().is_first_rank: intermediate_tensors = None else: intermediate_tensors = self.sync_and_slice_intermediate_tensors( num_input_tokens, intermediate_tensors, True ) # Some attention backends only support CUDA Graphs in pure decode. skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs # Forward pass with set_forward_context( attn_metadata, self.vllm_config, num_tokens=num_input_tokens, num_tokens_across_dp=num_tokens_across_dp, skip_cuda_graphs=skip_cuda_graphs, ), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output: model_kwargs_extra = {} # Only pass additional_information for the prefill part if ( hasattr(self, "_forward_additional_information") and self._forward_additional_information is not None and isinstance(self._forward_additional_information, dict) ): model_kwargs_extra["additional_information"] = ( self._forward_additional_information ) model_output = self.model( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, **MultiModalKwargs.as_kwargs( model_mm_kwargs, device=self.device, ), sampling_metadata=self.input_batch.sampling_metadata, logits_index=logits_indices, sampler=self.sampler, **model_kwargs_extra, ) if self.use_aux_hidden_state_outputs: hidden_states, _aux_hidden_states = model_output else: hidden_states = model_output _aux_hidden_states = None text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs( hidden_states ) # Mid PP stages return intermediate tensors unmodified if not get_pp_group().is_last_rank: assert isinstance(text_hidden_states, IntermediateTensors) text_hidden_states.kv_connector_output = kv_connector_output return text_hidden_states # Broadcast PP output for external_launcher (torchrun) broadcast_pp_output = ( self.parallel_config.distributed_executor_backend == "external_launcher" and len(get_pp_group().ranks) > 0 ) if not get_pp_group().is_last_rank: assert isinstance(text_hidden_states, IntermediateTensors) if not broadcast_pp_output: text_hidden_states.kv_connector_output = kv_connector_output return text_hidden_states get_pp_group().send_tensor_dict( text_hidden_states.tensors, all_gather_group=get_tp_group() ) logits = None else: if self.input_batch.pooling_params: return self._pool( text_hidden_states, num_scheduled_tokens, num_scheduled_tokens_np, kv_connector_output, ) sample_hidden_states = text_hidden_states[logits_indices] logits = self.model.compute_logits(sample_hidden_states, None) if broadcast_pp_output: model_output_broadcast_data = ( { "logits": logits.contiguous(), } if logits is not None else {} ) model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 ) assert model_output_broadcast_data is not None logits = model_output_broadcast_data["logits"] # Apply structured output bitmasks if present if scheduler_output.grammar_bitmask is not None: self.apply_grammar_bitmask(scheduler_output, logits) # Sample the next token and get logprobs if needed (with spec decode) sampling_metadata = self.input_batch.sampling_metadata if spec_decode_metadata is None: sampler_output = self.sampler( logits=logits, sampling_metadata=sampling_metadata, ) else: assert logits is not None bonus_logits = logits[spec_decode_metadata.bonus_logits_indices] sampler_output = self.sampler( logits=bonus_logits, sampling_metadata=sampling_metadata, ) bonus_token_ids = sampler_output.sampled_token_ids target_logits = logits[spec_decode_metadata.target_logits_indices] output_token_ids = self.rejection_sampler( spec_decode_metadata, None, target_logits, bonus_token_ids, sampling_metadata, ) sampler_output.sampled_token_ids = output_token_ids num_nans_in_logits = {} if envs.VLLM_COMPUTE_NANS_IN_LOGITS: num_nans_in_logits = self._get_nans_in_logits(logits) # Handle partial prefill: discard sampled tokens and rewind RNG discard_sampled_tokens_req_indices = [] for i, req_id in enumerate(self.input_batch.req_ids): req_state = self.requests[req_id] seq_len = ( req_state.num_computed_tokens + scheduler_output.num_scheduled_tokens[req_id] ) if seq_len < req_state.num_tokens: generator = self.input_batch.generators.get(i) if generator is not None: generator.set_offset(generator.get_offset() - 4) discard_sampled_tokens_req_indices.append(i) # Move CPU sync parts logprobs_tensors = sampler_output.logprobs_tensors logprobs_lists = ( logprobs_tensors.tolists() if logprobs_tensors is not None else None ) # Prompt logprobs if needed prompt_logprobs_dict = self._get_prompt_logprobs_dict( text_hidden_states[:num_scheduled_tokens], scheduler_output, ) # Parse valid sampled tokens sampled_token_ids = sampler_output.sampled_token_ids max_gen_len = sampled_token_ids.shape[-1] if max_gen_len == 1: valid_sampled_token_ids = sampled_token_ids.tolist() else: valid_sampled_token_ids = self.rejection_sampler.parse_output( sampled_token_ids, self.input_batch.vocab_size, ) for i in discard_sampled_tokens_req_indices: valid_sampled_token_ids[i].clear() # Cache sampled tokens for req_idx, sampled_ids in enumerate(valid_sampled_token_ids): if not sampled_ids: continue start_idx = self.input_batch.num_tokens_no_spec[req_idx] end_idx = start_idx + len(sampled_ids) self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids self.input_batch.num_tokens_no_spec[req_idx] = end_idx self.input_batch.num_tokens[req_idx] = end_idx req_id = self.input_batch.req_ids[req_idx] req_state = self.requests[req_id] req_state.output_token_ids.extend(sampled_ids) # Speculative decoding draft tokens if configured if not self.speculative_config: spec_token_ids = None else: assert spec_decode_common_attn_metadata is not None spec_token_ids = self.propose_draft_token_ids( scheduler_output, valid_sampled_token_ids, sampling_metadata, text_hidden_states, sample_hidden_states, _aux_hidden_states, spec_decode_metadata, spec_decode_common_attn_metadata, ) # Convert to per-request tensors on CPU text_hidden_states_cpu = text_hidden_states.detach().to("cpu").contiguous() pooler_output: list[Optional[torch.Tensor]] = [] prev_logits_index = 0 for logits_index in logits_indices: pooler_output.append( text_hidden_states_cpu[prev_logits_index : logits_index + 1] ) prev_logits_index = logits_index + 1 self.eplb_step() return OmniModelRunnerOutput( req_ids=self.input_batch.req_ids, req_id_to_index=self.input_batch.req_id_to_index, sampled_token_ids=valid_sampled_token_ids, spec_token_ids=spec_token_ids, logprobs=logprobs_lists, prompt_logprobs_dict=prompt_logprobs_dict, pooler_output=( pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None ), kv_connector_output=kv_connector_output, num_nans_in_logits=num_nans_in_logits, ) @torch.inference_mode() def extract_multimodal_outputs( self, hidden_states: Union[torch.Tensor, List[torch.Tensor]] ) -> dict: if ( hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs ): text_hidden_states = hidden_states.text_hidden_states multimodal_outputs = hidden_states.multimodal_outputs elif isinstance(hidden_states, torch.Tensor): text_hidden_states = hidden_states multimodal_outputs = {} elif isinstance(hidden_states, List): text_hidden_states = hidden_states[0] multimodal_outputs = {} else: raise ValueError(f"Invalid hidden states type: {type(hidden_states)}") return text_hidden_states, multimodal_outputs @torch.inference_mode() def _dummy_run( self, num_tokens: int, capture_attn_cudagraph: bool = False, skip_eplb: bool = False, is_profile: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # Padding for DP num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) num_tokens += num_pad # Set num_scheduled_tokens based on num_tokens and max_num_seqs # for dummy run with LoRA so that the num_reqs collectively # has num_tokens in total. assert num_tokens <= self.scheduler_config.max_num_batched_tokens max_num_reqs = self.scheduler_config.max_num_seqs num_reqs = min(num_tokens, max_num_reqs) min_tokens_per_req = num_tokens // num_reqs num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs num_scheduled_tokens_list[-1] += num_tokens % num_reqs assert sum(num_scheduled_tokens_list) == num_tokens assert len(num_scheduled_tokens_list) == num_reqs num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) attn_metadata: Optional[dict[str, Any]] = None if capture_attn_cudagraph: attn_metadata = {} # Make sure max_model_len is used at the graph capture time. self.seq_lens_np[:num_reqs] = self.max_model_len self.seq_lens_np[num_reqs:] = 0 self.seq_lens[:num_reqs].copy_( self.seq_lens_cpu[:num_reqs], non_blocking=True ) for kv_cache_group_id, kv_cache_group_spec in enumerate( self.kv_cache_config.kv_cache_groups ): common_attn_metadata = CommonAttentionMetadata( query_start_loc=self.query_start_loc[: num_reqs + 1], query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1], seq_lens=self.seq_lens[:num_reqs], seq_lens_cpu=self.seq_lens_cpu[:num_reqs], num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[ # noqa: E501 :num_reqs ], num_reqs=num_reqs, num_actual_tokens=num_tokens, max_query_len=num_tokens, block_table_tensor=self.input_batch.block_table[ kv_cache_group_id ].get_device_tensor()[:num_reqs], slot_mapping=self.input_batch.block_table[ kv_cache_group_id ].slot_mapping[:num_tokens], causal=True, ) for attn_group in self.attn_groups[kv_cache_group_id]: attn_metadata_i = ( attn_group.metadata_builder.build_for_cudagraph_capture( common_attn_metadata ) ) for layer_name in kv_cache_group_spec.layer_names: attn_metadata[layer_name] = attn_metadata_i with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens): if self.is_multimodal_model: input_ids = None inputs_embeds = self.inputs_embeds[:num_tokens] model_mm_kwargs = self._dummy_mm_kwargs(num_reqs) else: input_ids = self.input_ids[:num_tokens] inputs_embeds = None model_mm_kwargs = {} if self.uses_mrope: positions = self.mrope_positions[:, :num_tokens] else: positions = self.positions[:num_tokens] if get_pp_group().is_first_rank: intermediate_tensors = None else: if self.intermediate_tensors is None: self.intermediate_tensors = ( self.model.make_empty_intermediate_tensors( batch_size=self.max_num_tokens, dtype=self.model_config.dtype, device=self.device, ) ) intermediate_tensors = self.sync_and_slice_intermediate_tensors( num_tokens, None, False ) with self.maybe_randomize_inputs(input_ids), set_forward_context( attn_metadata, self.vllm_config, num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): outputs = self.model( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, **MultiModalKwargs.as_kwargs( model_mm_kwargs, device=self.device, ), ) if self.use_aux_hidden_state_outputs: hidden_states, _ = outputs else: hidden_states = outputs if self.speculative_config and self.speculative_config.use_eagle(): assert isinstance(self.drafter, EagleProposer) self.drafter.dummy_run(num_tokens) # This is necessary to avoid blocking DP. # For dummy runs, we typically skip EPLB since we don't have any real # requests to process. # However, in DP settings, there may be cases when some DP ranks do # not have any requests to process, so they're executing dummy batches. # In such cases, we still have to trigger EPLB to make sure # ranks execute the rearrangement in synchronization. if not skip_eplb: self.eplb_step(is_dummy=True, is_profile=is_profile) logit_indices = np.cumsum(num_scheduled_tokens) - 1 hidden_states, multimodal_outputs = self.extract_multimodal_outputs( hidden_states ) return hidden_states, hidden_states[logit_indices]