import itertools from dataclasses import dataclass from typing import Any, Dict, List, Optional import torch from vllm.lora.layers import LoRAMapping from vllm.sequence import MultiModalKwargs from vllm.utils import async_tensor_h2d, flatten_2d_lists from vllm.worker.model_runner import ModelInputForGPU, ModelInputForGPUBuilder @dataclass(frozen=True) class OmniModelInputForGPU(ModelInputForGPU): # New: optional additional information dict for current scheduled tokens additional_information: Optional[Dict[str, Any]] = None class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder): def build(self) -> OmniModelInputForGPU: """Finalize the builder intermediate data and create on-device tensors. """ # Combine and flatten intermediate data. input_tokens = list[int]() inputs_embeds_list = list[torch.Tensor]() token_types = list[int]() for inter_data in self.inter_data_list: for cur_input_tokens in inter_data.input_tokens: input_tokens.extend(cur_input_tokens) for cur_token_types in inter_data.token_types: token_types.extend(cur_token_types) if inter_data.inputs_embeds is not None: inputs_embeds_list.append( inter_data.inputs_embeds.to( dtype=self.runner.model_config.dtype, device=self.runner.device ) ) # Support v1 direct-transfer prompt embeds attached on request state if inter_data.inputs_embeds is None: try: # Locate req_id -> request state to fetch decoded CPU embeds req_id = inter_data.request_id # type: ignore[attr-defined] req_state = getattr(self.runner, "requests", {}).get(req_id) pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) if pe_cpu is not None: inter_data.inputs_embeds = pe_cpu.to( dtype=self.runner.model_config.dtype, device=self.runner.device, ) inputs_embeds_list.append(inter_data.inputs_embeds) except Exception: pass inputs_embeds: Optional[torch.Tensor] if len(inputs_embeds_list) == 0: inputs_embeds = None else: inputs_embeds = torch.cat(inputs_embeds_list, dim=0).to( dtype=self.runner.model_config.dtype, device=self.runner.device ) assert len(inputs_embeds) == len(input_tokens) if not input_tokens and inputs_embeds is None: # This may happen when all prefill requests hit # prefix caching and there is no decode request. return self.model_input_cls() mrope_input_positions: Optional[List[List[int]]] = None if any( inter_data.mrope_input_positions is not None for inter_data in self.inter_data_list ): mrope_input_positions = [[] for _ in range(3)] for idx in range(3): for inter_data in self.inter_data_list: msections = inter_data.mrope_input_positions if msections is None: for _seq_input_positions in inter_data.input_positions: mrope_input_positions[idx].extend(_seq_input_positions) else: for _seq_mrope_input_positions in msections: mrope_input_positions[idx].extend( _seq_mrope_input_positions[idx] ) input_positions = None else: input_positions = [] for inter_data in self.inter_data_list: for cur_input_positions in inter_data.input_positions: input_positions.extend(cur_input_positions) seq_lens = [] query_lens = [] max_decode_seq_len = 0 max_encoder_seq_len = 0 for inter_data in self.inter_data_list: seq_lens.extend(inter_data.seq_lens) query_lens.extend(inter_data.query_lens) if not inter_data.is_prompt: max_decode_seq_len = max(max_decode_seq_len, max(inter_data.seq_lens)) if self.runner.model_config.is_encoder_decoder: max_encoder_seq_len = max( max_encoder_seq_len, inter_data.encoder_seq_len ) # Mapping from request IDs to sequence IDs. Used for Jamba models # that manages the cache by itself. request_ids_to_seq_ids = { data.request_id: data.seq_ids for data in self.inter_data_list } cuda_graph_pad_size = self._get_cuda_graph_pad_size( num_seqs=len(seq_lens), max_decode_seq_len=max_decode_seq_len, max_encoder_seq_len=max_encoder_seq_len, ) batch_size = len(input_tokens) if cuda_graph_pad_size != -1: # If cuda graph can be used, pad tensors accordingly. # See `capture_model` API for more details. # vLLM uses cuda graph only for decoding requests. batch_size += cuda_graph_pad_size # Tokens and positions. if cuda_graph_pad_size: input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size)) assert self.runner.device is not None input_tokens_tensor = async_tensor_h2d( input_tokens, torch.long, self.runner.device, self.runner.pin_memory ) token_types_tensor = ( async_tensor_h2d( token_types, torch.long, self.runner.device, self.runner.pin_memory ) if token_types else None ) if mrope_input_positions is not None: for idx in range(3): mrope_input_positions[idx].extend( itertools.repeat(0, cuda_graph_pad_size) ) input_positions_tensor = async_tensor_h2d( mrope_input_positions, torch.long, self.runner.device, self.runner.pin_memory, ) else: input_positions.extend(itertools.repeat(0, cuda_graph_pad_size)) input_positions_tensor = async_tensor_h2d( input_positions, torch.long, self.runner.device, self.runner.pin_memory ) # Sequence and query lengths. if cuda_graph_pad_size: seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size)) # Attention metadata. attn_metadata = self.attn_metadata_builder.build( seq_lens, query_lens, cuda_graph_pad_size, batch_size ) # LoRA data. lora_requests = set() lora_mapping = None if self.enable_lora: lora_requests = set( r for data in self.inter_data_list for r in data.lora_requests ) lora_index_mapping = flatten_2d_lists( [ flatten_2d_lists(inter_data.lora_index_mapping) for inter_data in self.inter_data_list ] ) if cuda_graph_pad_size: lora_index_mapping.extend(itertools.repeat(0, cuda_graph_pad_size)) lora_prompt_mapping = flatten_2d_lists( [ flatten_2d_lists(inter_data.lora_prompt_mapping) for inter_data in self.inter_data_list ] ) lora_mapping = LoRAMapping( **dict( index_mapping=lora_index_mapping, prompt_mapping=lora_prompt_mapping, is_prefill=not self.decode_only, ) ) # Multi-modal data. multi_modal_kwargs_list = [ data.multi_modal_kwargs for data in self.inter_data_list if data.multi_modal_kwargs is not None ] multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list) return self.model_input_cls( input_tokens=input_tokens_tensor, inputs_embeds=inputs_embeds, input_positions=input_positions_tensor, token_types=token_types_tensor, attn_metadata=attn_metadata, seq_lens=seq_lens, query_lens=query_lens, lora_mapping=lora_mapping, lora_requests=lora_requests, multi_modal_kwargs=multi_modal_kwargs, request_ids_to_seq_ids=request_ids_to_seq_ids, finished_requests_ids=self.finished_requests_ids, )