217 lines
8.6 KiB
Python
217 lines
8.6 KiB
Python
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,
|
|
)
|