| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | from collections import UserDict, defaultdict |
| 6 | from collections.abc import Mapping, Sequence |
| 7 | from dataclasses import dataclass |
| 8 | from functools import cached_property, partial |
| 9 | from itertools import accumulate |
| 10 | from typing import ( |
| 11 | TYPE_CHECKING, |
| 12 | Any, |
| 13 | Literal, |
| 14 | TypeAlias, |
| 15 | TypedDict, |
| 16 | Union, |
| 17 | cast, |
| 18 | ) |
| 19 | |
| 20 | import numpy as np |
| 21 | from PIL.Image import Image |
| 22 | from typing_extensions import TypeVar |
| 23 | |
| 24 | from vllm.utils.collection_utils import is_list_of |
| 25 | from vllm.utils.import_utils import LazyLoader |
| 26 | from vllm.utils.jsontree import json_iter_leaves, json_map_leaves |
| 27 | |
| 28 | from .media import MediaWithBytes |
| 29 | |
| 30 | if TYPE_CHECKING: |
| 31 | import torch |
| 32 | import torch.types |
| 33 | from transformers.feature_extraction_utils import BatchFeature |
| 34 | else: |
| 35 | torch = LazyLoader("torch", globals(), "torch") |
| 36 | |
| 37 | |
| 38 | HfImageItem: TypeAlias = Union["Image", np.ndarray, "torch.Tensor"] |
| 39 | """ |
| 40 | A `transformers.image_utils.ImageInput` representing a single image |
| 41 | item, which can be passed to a HuggingFace `ImageProcessor`. |
| 42 | """ |
| 43 | |
| 44 | HfVideoItem: TypeAlias = Union[ |
| 45 | list["Image"], np.ndarray, "torch.Tensor", list[np.ndarray], list["torch.Tensor"] |
| 46 | ] |
| 47 | """ |
| 48 | A `transformers.image_utils.VideoInput` representing a single video |
| 49 | item, which can be passed to a HuggingFace `VideoProcessor`. |
| 50 | """ |
| 51 | |
| 52 | HfAudioItem: TypeAlias = Union[list[float], np.ndarray, "torch.Tensor"] |
| 53 | """ |
| 54 | Represents a single audio |
| 55 | item, which can be passed to a HuggingFace `AudioProcessor`. |
| 56 | """ |
| 57 | |
| 58 | ImageItem: TypeAlias = Union[HfImageItem, "torch.Tensor", MediaWithBytes[HfImageItem]] |
| 59 | """ |
| 60 | A `transformers.image_utils.ImageInput` representing a single image |
| 61 | item, which can be passed to a HuggingFace `ImageProcessor`. |
| 62 | |
| 63 | Alternatively, a 3-D tensor or batch of 2-D tensors, |
| 64 | which are treated as image embeddings; |
| 65 | these are directly passed to the model without HF processing. |
| 66 | """ |
| 67 | |
| 68 | VideoItem: TypeAlias = Union[ |
| 69 | HfVideoItem, |
| 70 | "torch.Tensor", |
| 71 | tuple[HfVideoItem, dict[str, Any]], |
| 72 | MediaWithBytes[tuple[HfVideoItem, dict[str, Any]]], |
| 73 | ] |
| 74 | """ |
| 75 | A `transformers.video_utils.VideoInput` representing a single video item. |
| 76 | This can be passed to a HuggingFace `VideoProcessor` |
| 77 | with `transformers.video_utils.VideoMetadata`. |
| 78 | |
| 79 | Alternatively, a 3-D tensor or batch of 2-D tensors, |
| 80 | which are treated as video embeddings; |
| 81 | these are directly passed to the model without HF processing. |
| 82 | """ |
| 83 | |
| 84 | AudioItem: TypeAlias = Union[HfAudioItem, tuple[np.ndarray, float], "torch.Tensor"] |
| 85 | """ |
| 86 | Represents a single audio |
| 87 | item, which can be passed to a HuggingFace `AudioProcessor`. |
| 88 | |
| 89 | Alternatively, a tuple `(audio, sampling_rate)`, where the sampling rate |
| 90 | is different from that expected by the model; |
| 91 | these are resampled to the model's sampling rate before being processed by HF. |
| 92 | |
| 93 | Alternatively, a 3-D tensor or batch of 2-D tensors, |
| 94 | which are treated as audio embeddings; |
| 95 | these are directly passed to the model without HF processing. |
| 96 | """ |
| 97 | |
| 98 | |
| 99 | class VisionChunkImage(TypedDict): |
| 100 | """Represents an image wrapped as a vision chunk.""" |
| 101 | |
| 102 | type: Literal["image"] |
| 103 | image: Image |
| 104 | uuid: str | None |
| 105 | |
| 106 | |
| 107 | class VisionChunkVideo(TypedDict): |
| 108 | """Represents a video chunk with metadata.""" |
| 109 | |
| 110 | type: Literal["video_chunk"] |
| 111 | video_chunk: list[Image] |
| 112 | uuid: str | None |
| 113 | prompt: str |
| 114 | video_idx: int |
| 115 | |
| 116 | |
| 117 | VisionChunk: TypeAlias = VisionChunkImage | VisionChunkVideo |
| 118 | """A vision chunk is either an image or a video chunk.""" |
| 119 | |
| 120 | |
| 121 | @dataclass(frozen=True) |
| 122 | class PlaceholderRange: |
| 123 | """Placeholder location information for multi-modal data. |
| 124 | |
| 125 | Example: |
| 126 | Prompt: `AAAA BBBB What is in these images?` |
| 127 | |
| 128 | Images A and B will have: |
| 129 | |
| 130 | ``` |
| 131 | A: PlaceholderRange(offset=0, length=4) |
| 132 | B: PlaceholderRange(offset=5, length=4) |
| 133 | ``` |
| 134 | |
| 135 | """ |
| 136 | |
| 137 | offset: int |
| 138 | """The start index of the placeholder in the prompt.""" |
| 139 | |
| 140 | length: int |
| 141 | """The length of the placeholder.""" |
| 142 | |
| 143 | is_embed: "torch.Tensor | None" = None |
| 144 | """ |
| 145 | A boolean mask of shape `(length,)` indicating which positions |
| 146 | between `offset` and `offset + length` to assign embeddings to. |
| 147 | """ |
| 148 | |
| 149 | @cached_property |
| 150 | def embeds_cumsum(self) -> list[int] | None: |
| 151 | # python list so python indexing avoids torch C++ overhead/conversions/deallocs |
| 152 | return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist() |
| 153 | |
| 154 | def get_num_embeds(self) -> int: |
| 155 | if self.embeds_cumsum is None: |
| 156 | return self.length |
| 157 | |
| 158 | return self.embeds_cumsum[-1] if self.embeds_cumsum else 0 |
| 159 | |
| 160 | def get_embeds_indices_in_range( |
| 161 | self, start_idx: int, end_idx: int |
| 162 | ) -> tuple[int, int]: |
| 163 | """Returns the starting and ending indices of the embeddings of encoder outputs |
| 164 | in the range of [start_idx, end_idx) in the placeholders. |
| 165 | |
| 166 | For example, given: |
| 167 | PlaceholderRange(offset=2, length=5, is_embed=[False, True, False, True, True]) |
| 168 | |
| 169 | If start_idx=3 and end_idx=5, the output is (1, 3) because we want to get |
| 170 | the second and the third embeddings from the encoder output. |
| 171 | """ |
| 172 | if self.embeds_cumsum is None: |
| 173 | return start_idx, end_idx |
| 174 | |
| 175 | embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0 |
| 176 | embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0 |
| 177 | |
| 178 | return embeds_start_idx, embeds_end_idx |
| 179 | |
| 180 | def extract_embeds_range(self) -> list[tuple[int, int]]: |
| 181 | """Extract the start and end indices of the embedded region in prompt. |
| 182 | |
| 183 | For example, given `PlaceholderRange(offset=2, length=5)` and |
| 184 | `is_embed = [False, True, False, True, True]`, the output is |
| 185 | `[(1 + offset, 1 + offset), (3 + offset, 4 + offset)]`. |
| 186 | |
| 187 | Returns: |
| 188 | A tuple `(start, end)` representing the start and end |
| 189 | indices (inclusive) of the embedded region. |
| 190 | Returns full placeholder range if `is_embed` is `None`. |
| 191 | |
| 192 | """ |
| 193 | if self.is_embed is None: |
| 194 | return [(self.offset, self.offset + self.length - 1)] |
| 195 | |
| 196 | mask_i = self.is_embed.int() |
| 197 | starts = torch.nonzero( |
| 198 | torch.diff(mask_i, prepend=mask_i.new_zeros(1)) == 1 |
| 199 | ).flatten() |
| 200 | ends = torch.nonzero( |
| 201 | torch.diff(mask_i, append=mask_i.new_zeros(1)) == -1 |
| 202 | ).flatten() |
| 203 | ranges = torch.stack((starts, ends), dim=1) + self.offset |
| 204 | return [tuple(x) for x in ranges.tolist()] |
| 205 | |
| 206 | def __eq__(self, other: object) -> bool: |
| 207 | if not isinstance(other, self.__class__): |
| 208 | return False |
| 209 | if not (self.offset, self.length) == (other.offset, other.length): |
| 210 | return False |
| 211 | |
| 212 | if self.is_embed is None: |
| 213 | return other.is_embed is None |
| 214 | if other.is_embed is None: |
| 215 | return self.is_embed is None |
| 216 | |
| 217 | return nested_tensors_equal(self.is_embed, other.is_embed) |
| 218 | |
| 219 | |
| 220 | NestedTensors: TypeAlias = Union[ |
| 221 | list["NestedTensors"], |
| 222 | list["torch.Tensor"], |
| 223 | "torch.Tensor", |
| 224 | tuple["torch.Tensor", ...], |
| 225 | ] |
| 226 | """ |
| 227 | Uses a list instead of a tensor if the dimensions of each element do not match. |
| 228 | """ |
| 229 | |
| 230 | |
| 231 | def nested_tensors_equal( |
| 232 | a: NestedTensors, |
| 233 | b: NestedTensors, |
| 234 | check_dtype: bool = True, |
| 235 | ) -> bool: |
| 236 | """Equality check between |
| 237 | [`NestedTensors`][vllm.multimodal.inputs.NestedTensors] objects. |
| 238 | |
| 239 | If `check_dtype` is `True`, the tensors must have the same dtype. |
| 240 | """ |
| 241 | check_dtype_func = ( |
| 242 | lambda a, b, check_dtype: a.dtype == b.dtype if check_dtype else True |
| 243 | ) |
| 244 | if isinstance(a, torch.Tensor): |
| 245 | return ( |
| 246 | isinstance(b, torch.Tensor) |
| 247 | and torch.equal(a, b) |
| 248 | and check_dtype_func(a, b, check_dtype) |
| 249 | ) |
| 250 | elif isinstance(b, torch.Tensor): |
| 251 | return ( |
| 252 | isinstance(a, torch.Tensor) |
| 253 | and torch.equal(b, a) |
| 254 | and check_dtype_func(b, a, check_dtype) |
| 255 | ) |
| 256 | |
| 257 | if isinstance(a, list): |
| 258 | return ( |
| 259 | isinstance(b, list) |
| 260 | and len(a) == len(b) |
| 261 | and all(nested_tensors_equal(a_, b_, check_dtype) for a_, b_ in zip(a, b)) |
| 262 | ) |
| 263 | if isinstance(b, list): |
| 264 | return ( |
| 265 | isinstance(a, list) |
| 266 | and len(b) == len(a) |
| 267 | and all(nested_tensors_equal(b_, a_, check_dtype) for b_, a_ in zip(b, a)) |
| 268 | ) |
| 269 | |
| 270 | if isinstance(a, tuple): |
| 271 | return ( |
| 272 | isinstance(b, tuple) |
| 273 | and len(a) == len(b) |
| 274 | and all(nested_tensors_equal(a_, b_, check_dtype) for a_, b_ in zip(a, b)) |
| 275 | ) |
| 276 | if isinstance(b, tuple): |
| 277 | return ( |
| 278 | isinstance(a, tuple) |
| 279 | and len(b) == len(a) |
| 280 | and all(nested_tensors_equal(b_, a_, check_dtype) for b_, a_ in zip(b, a)) |
| 281 | ) |
| 282 | |
| 283 | # Both a and b are scalars |
| 284 | return a == b |
| 285 | |
| 286 | |
| 287 | def _nested_tensors_are_cpu(tensors: NestedTensors) -> bool: |
| 288 | """Whether every tensor in `tensors` lives in host memory.""" |
| 289 | return not any( |
| 290 | isinstance(x, torch.Tensor) and not x.is_cpu for x in json_iter_leaves(tensors) |
| 291 | ) |
| 292 | |
| 293 | |
| 294 | def _is_dense(t: torch.Tensor) -> bool: |
| 295 | """Whether `t`'s elements fill a contiguous storage range, possibly |
| 296 | permuted (e.g. a transposed matrix). Such layouts can be copied with |
| 297 | pitched cudaMemcpy2D/3DAsync directly from pinned memory, whereas |
| 298 | gapped layouts (e.g. strided slices) stage through a pageable temp.""" |
| 299 | if not t.is_contiguous(): |
| 300 | expected_stride = 1 |
| 301 | for stride, size in sorted(zip(t.stride(), t.shape)): |
| 302 | if size > 1: |
| 303 | if stride != expected_stride: |
| 304 | return False |
| 305 | expected_stride *= size |
| 306 | return True |
| 307 | |
| 308 | |
| 309 | def _nested_tensors_h2d( |
| 310 | tensors: NestedTensors, |
| 311 | device: torch.types.Device, |
| 312 | *, |
| 313 | pin_memory: bool = False, |
| 314 | ) -> NestedTensors: |
| 315 | if device is None: |
| 316 | return tensors |
| 317 | |
| 318 | def _h2d(x: torch.Tensor) -> torch.Tensor: |
| 319 | if not isinstance(x, torch.Tensor): |
| 320 | return x |
| 321 | if pin_memory and x.is_cpu and not (x.is_pinned() and _is_dense(x)): |
| 322 | # Ensure tensor is pinned and dense for non_blocking H2D copy. |
| 323 | x = x.new_empty(x.shape, pin_memory=True).copy_(x) |
| 324 | return x.to(device=device, non_blocking=True) |
| 325 | |
| 326 | return json_map_leaves(_h2d, tensors) |
| 327 | |
| 328 | |
| 329 | BatchedTensorInputs: TypeAlias = dict[str, NestedTensors] |
| 330 | """ |
| 331 | A dictionary containing nested tensors which have been batched via |
| 332 | [`MultiModalKwargsItems.get_data`][vllm.multimodal.inputs.MultiModalKwargsItems.get_data]. |
| 333 | """ |
| 334 | |
| 335 | |
| 336 | def batched_tensors_equal(a: BatchedTensorInputs, b: BatchedTensorInputs) -> bool: |
| 337 | """Equality check between |
| 338 | [`BatchedTensorInputs`][vllm.multimodal.inputs.BatchedTensorInputs] objects. |
| 339 | """ |
| 340 | return all(k in b and nested_tensors_equal(a[k], b[k]) for k in a) |
| 341 | |
| 342 | |
| 343 | @dataclass |
| 344 | class MultiModalFeatureSpec: |
| 345 | """Represents a single multimodal input with its processed data and metadata. |
| 346 | |
| 347 | Used to track multimodal data through processing and caching. |
| 348 | A request containing multiple multimodal items will have one |
| 349 | `MultiModalFeatureSpec` per item. |
| 350 | """ |
| 351 | |
| 352 | data: "MultiModalKwargsItem | None" |
| 353 | """ |
| 354 | Represents multimodal data for this feature. |
| 355 | |
| 356 | Can be `None` if the item is cached, to skip IPC between API server |
| 357 | and engine core processes. |
| 358 | """ |
| 359 | |
| 360 | modality: str |
| 361 | """The input modality, e.g., `"image"`, `"audio"`, `"video"`.""" |
| 362 | |
| 363 | identifier: str |
| 364 | """The hash for caching encoder outputs (with LoRA prefix if applicable).""" |
| 365 | |
| 366 | mm_position: PlaceholderRange |
| 367 | """ |
| 368 | The location of the `modality` tokens corresponding to this item |
| 369 | in the prompt, e.g., `PlaceholderRange(offset=2, length=336)`. |
| 370 | """ |
| 371 | |
| 372 | mm_hash: str | None = None |
| 373 | """The hash for caching processor outputs (without LoRA prefix).""" |
| 374 | |
| 375 | @staticmethod |
| 376 | def gather_kwargs(features: list["MultiModalFeatureSpec"], keys: set[str]): |
| 377 | kwargs = defaultdict[str, list[NestedTensors]](list) |
| 378 | |
| 379 | for f in features: |
| 380 | item = f.data |
| 381 | if item is not None: |
| 382 | for k in keys: |
| 383 | if k in item: |
| 384 | kwargs[k].append(item[k].data) |
| 385 | |
| 386 | return dict(kwargs) |
| 387 | |
| 388 | |
| 389 | @dataclass |
| 390 | class MultiModalFieldElem: |
| 391 | """Represents a processed keyword argument to pass to a model for a |
| 392 | [`MultiModalKwargsItem`][vllm.multimodal.inputs.MultiModalKwargsItem]. |
| 393 | """ |
| 394 | |
| 395 | data: NestedTensors |
| 396 | """ |
| 397 | The tensor data of this field in |
| 398 | [`MultiModalKwargsItem`][vllm.multimodal.inputs.MultiModalKwargsItem], |
| 399 | i.e. the value of the keyword argument to be passed to the model. |
| 400 | |
| 401 | It may be set to `None` if it is determined that the item is cached |
| 402 | in `EngineCore`. |
| 403 | """ |
| 404 | |
| 405 | field: "BaseMultiModalField" |
| 406 | """ |
| 407 | Defines how to combine the tensor data of this field with others |
| 408 | in order to batch multi-modal items together for model inference. |
| 409 | """ |
| 410 | |
| 411 | def __eq__(self, other: object) -> bool: |
| 412 | if not isinstance(other, self.__class__): |
| 413 | return False |
| 414 | |
| 415 | if self.data is None: |
| 416 | data_equal = other.data is None |
| 417 | elif other.data is None: |
| 418 | data_equal = self.data is None |
| 419 | else: |
| 420 | data_equal = nested_tensors_equal(self.data, other.data) |
| 421 | |
| 422 | return data_equal and type(self.field) is type(other.field) # noqa: E721 |
| 423 | |
| 424 | |
| 425 | @dataclass(frozen=True, kw_only=True) |
| 426 | class BaseMultiModalField(ABC): |
| 427 | """Defines how to interpret tensor data belonging to a keyword argument for |
| 428 | [`MultiModalKwargsItems`][vllm.multimodal.inputs.MultiModalKwargsItems], |
| 429 | and vice versa. |
| 430 | """ |
| 431 | |
| 432 | keep_on_cpu: bool = False |
| 433 | """ |
| 434 | If `True`, then this field is excluded from being moved to the accelerator when |
| 435 | [`group_and_batch_mm_items`][vllm.multimodal.utils.group_and_batch_mm_items] |
| 436 | is called to batch the data. |
| 437 | """ |
| 438 | |
| 439 | def _field_factory(self): |
| 440 | f = partial(MultiModalFieldElem, field=self) |
| 441 | |
| 442 | # Allow passing data as positional argument |
| 443 | def factory(data: NestedTensors) -> MultiModalFieldElem: |
| 444 | return f(data=data) |
| 445 | |
| 446 | return factory |
| 447 | |
| 448 | @abstractmethod |
| 449 | def build_elems( |
| 450 | self, |
| 451 | modality: str, |
| 452 | key: str, |
| 453 | data: NestedTensors, |
| 454 | ) -> Sequence[MultiModalFieldElem]: |
| 455 | """Construct |
| 456 | [`MultiModalFieldElem`][vllm.multimodal.inputs.MultiModalFieldElem] |
| 457 | instances to represent the provided data. |
| 458 | |
| 459 | This is the inverse of |
| 460 | [`reduce_data`][vllm.multimodal.inputs.BaseMultiModalField.reduce_data]. |
| 461 | """ |
| 462 | raise NotImplementedError |
| 463 | |
| 464 | @abstractmethod |
| 465 | def _reduce_data( |
| 466 | self, batch: list[NestedTensors], *, pin_memory: bool |
| 467 | ) -> NestedTensors: |
| 468 | raise NotImplementedError |
| 469 | |
| 470 | def reduce_data( |
| 471 | self, |
| 472 | elems: list[MultiModalFieldElem], |
| 473 | *, |
| 474 | device: torch.types.Device = None, |
| 475 | pin_memory: bool = False, |
| 476 | ) -> NestedTensors: |
| 477 | """Merge the data from multiple instances of |
| 478 | [`MultiModalFieldElem`][vllm.multimodal.inputs.MultiModalFieldElem]. |
| 479 | |
| 480 | This is the inverse of |
| 481 | [`build_elems`][vllm.multimodal.inputs.BaseMultiModalField.build_elems]. |
| 482 | """ |
| 483 | field_types = [type(item.field) for item in elems] |
| 484 | if len(set(field_types)) > 1: |
| 485 | raise ValueError(f"Cannot merge different {field_types=}") |
| 486 | |
| 487 | if device is not None and self.keep_on_cpu: |
| 488 | device = "cpu" |
| 489 | if pin_memory and self.keep_on_cpu: |
| 490 | pin_memory = False |
| 491 | if device == "cpu" or device == torch.device("cpu"): |
| 492 | pin_memory = False |
| 493 | |
| 494 | batch = [elem.data for elem in elems] |
| 495 | if pin_memory and not _nested_tensors_are_cpu(batch): |
| 496 | # Data already sits on an accelerator (e.g. a device-side HF |
| 497 | # processor handed it over by IPC). Pinning is a host-memory |
| 498 | # concept, and `pin_memory()` rejects device tensors outright. |
| 499 | pin_memory = False |
| 500 | out = self._reduce_data(batch, pin_memory=pin_memory) |
| 501 | return _nested_tensors_h2d(out, device=device, pin_memory=pin_memory) |
| 502 | |
| 503 | |
| 504 | @dataclass(frozen=True, kw_only=True) |
| 505 | class MultiModalBatchedField(BaseMultiModalField): |
| 506 | """Info: |
| 507 | [`MultiModalFieldConfig.batched`][vllm.multimodal.inputs.MultiModalFieldConfig.batched] |
| 508 | """ |
| 509 | |
| 510 | def build_elems( |
| 511 | self, |
| 512 | modality: str, |
| 513 | key: str, |
| 514 | data: NestedTensors, |
| 515 | ) -> Sequence[MultiModalFieldElem]: |
| 516 | field_factory = self._field_factory() |
| 517 | return [field_factory(item) for item in data] |
| 518 | |
| 519 | def _reduce_data( |
| 520 | self, batch: list[NestedTensors], *, pin_memory: bool |
| 521 | ) -> NestedTensors: |
| 522 | if len(batch) > 0 and is_list_of(batch, torch.Tensor, check="all"): |
| 523 | batch = cast(list[torch.Tensor], batch) |
| 524 | if len(batch) == 1: |
| 525 | # An optimization when `batch` contains only one tensor: |
| 526 | # - produce exactly same result as `torch.stack(batch)` |
| 527 | # - will achieve zero-copy if the tensor is contiguous |
| 528 | return batch[0].unsqueeze(0) |
| 529 | first_shape = batch[0].shape |
| 530 | if all(elem.shape == first_shape for elem in batch): |
| 531 | out = batch[0].new_empty( |
| 532 | (len(batch), *first_shape), pin_memory=pin_memory |
| 533 | ) |
| 534 | return torch.stack(batch, out=out) |
| 535 | |
| 536 | return batch |
| 537 | |
| 538 | |
| 539 | @dataclass(frozen=True, kw_only=True) |
| 540 | class MultiModalFlatField(BaseMultiModalField): |
| 541 | """Info: |
| 542 | [`MultiModalFieldConfig.flat`][vllm.multimodal.inputs.MultiModalFieldConfig.flat] |
| 543 | [`MultiModalFieldConfig.flat_from_sizes`][vllm.multimodal.inputs.MultiModalFieldConfig.flat_from_sizes] |
| 544 | """ |
| 545 | |
| 546 | slices: Sequence[slice] | Sequence[Sequence[slice]] |
| 547 | dim: int = 0 |
| 548 | |
| 549 | def build_elems( |
| 550 | self, |
| 551 | modality: str, |
| 552 | key: str, |
| 553 | data: NestedTensors, |
| 554 | ) -> Sequence[MultiModalFieldElem]: |
| 555 | field_factory = self._field_factory() |
| 556 | if not is_list_of(self.slices, slice, check="all"): |
| 557 | assert isinstance(data, torch.Tensor), ( |
| 558 | "torch.Tensor is required for multiple slices" |
| 559 | ) |
| 560 | return [field_factory(data[cast(slice, s)]) for s in self.slices] |
| 561 | |
| 562 | def _reduce_data( |
| 563 | self, batch: list[NestedTensors], *, pin_memory: bool |
| 564 | ) -> NestedTensors: |
| 565 | if len(batch) > 0 and is_list_of(batch, torch.Tensor, check="all"): |
| 566 | batch = cast(list[torch.Tensor], batch) |
| 567 | if len(batch) == 1: |
| 568 | # An optimization when `batch` contains only one tensor: |
| 569 | # - produce exactly same result as `torch.concat(batch)` |
| 570 | # - will achieve zero-copy if the tensor is contiguous |
| 571 | return batch[0] |
| 572 | |
| 573 | dim = self.dim + (self.dim < 0) * len(batch[0].shape) |
| 574 | |
| 575 | def _shape_before_after(tensor: torch.Tensor): |
| 576 | return tensor.shape[:dim], tensor.shape[dim + 1 :] |
| 577 | |
| 578 | first_shape = _shape_before_after(batch[0]) |
| 579 | |
| 580 | if all(_shape_before_after(elem) == first_shape for elem in batch): |
| 581 | shape_before, shape_after = first_shape |
| 582 | shape_concat = sum(item.shape[dim] for item in batch) |
| 583 | out = batch[0].new_empty( |
| 584 | (*shape_before, shape_concat, *shape_after), pin_memory=pin_memory |
| 585 | ) |
| 586 | return torch.concat(batch, dim=self.dim, out=out) |
| 587 | |
| 588 | # Variable-length case: non-concat dimensions differ |
| 589 | # (e.g., Ultravox with different audio durations). |
| 590 | # Use slice-assign approach (more efficient than padding). |
| 591 | # See: https://github.com/vllm-project/vllm/issues/31658 |
| 592 | |
| 593 | ndim = batch[0].ndim |
| 594 | |
| 595 | # Step 1: Compute output shape |
| 596 | # - Non-concat dims: take max across batch |
| 597 | # - Concat dim: sum across batch |
| 598 | max_sizes: list[int] = [] |
| 599 | for d in range(ndim): |
| 600 | if d == dim: |
| 601 | max_sizes.append(sum(t.shape[d] for t in batch)) |
| 602 | else: |
| 603 | max_sizes.append(max(t.shape[d] for t in batch)) |
| 604 | |
| 605 | # Step 2: Create zero-initialized output tensor |
| 606 | out = batch[0].new_zeros(max_sizes, pin_memory=pin_memory) |
| 607 | |
| 608 | # Step 3: Slice-assign each tensor to its proper position |
| 609 | concat_offset = 0 |
| 610 | for tensor in batch: |
| 611 | slices: list[slice] = [] |
| 612 | for d in range(ndim): |
| 613 | if d == dim: |
| 614 | slices.append( |
| 615 | slice(concat_offset, concat_offset + tensor.shape[d]) |
| 616 | ) |
| 617 | else: |
| 618 | slices.append(slice(0, tensor.shape[d])) |
| 619 | out[tuple(slices)] = tensor |
| 620 | concat_offset += tensor.shape[dim] |
| 621 | |
| 622 | return out |
| 623 | |
| 624 | assert self.dim == 0, "dim == 0 is required for nested list" |
| 625 | return [e for elem in batch for e in elem] |
| 626 | |
| 627 | |
| 628 | @dataclass(frozen=True, kw_only=True) |
| 629 | class MultiModalSharedField(BaseMultiModalField): |
| 630 | """Info: |
| 631 | [`MultiModalFieldConfig.shared`][vllm.multimodal.inputs.MultiModalFieldConfig.shared] |
| 632 | """ |
| 633 | |
| 634 | batch_size: int |
| 635 | |
| 636 | def build_elems( |
| 637 | self, |
| 638 | modality: str, |
| 639 | key: str, |
| 640 | data: NestedTensors, |
| 641 | ) -> Sequence[MultiModalFieldElem]: |
| 642 | field_factory = self._field_factory() |
| 643 | return [field_factory(data)] * self.batch_size |
| 644 | |
| 645 | def _reduce_data( |
| 646 | self, |
| 647 | batch: list[NestedTensors], |
| 648 | *, |
| 649 | pin_memory: bool, |
| 650 | ) -> NestedTensors: |
| 651 | return batch[0] |
| 652 | |
| 653 | |
| 654 | @dataclass(frozen=True) |
| 655 | class MultiModalFieldConfig: |
| 656 | @staticmethod |
| 657 | def batched(modality: str, *, keep_on_cpu: bool = False): |
| 658 | """Defines a field where an element in the batch is obtained by |
| 659 | indexing into the first dimension of the underlying data. |
| 660 | |
| 661 | Args: |
| 662 | modality: The modality of the multi-modal item that uses this |
| 663 | keyword argument. |
| 664 | keep_on_cpu: Whether to keep this field on the CPU for the model inputs. |
| 665 | |
| 666 | Example: |
| 667 | ``` |
| 668 | Input: |
| 669 | Data: [[AAAA] |
| 670 | [BBBB] |
| 671 | [CCCC]] |
| 672 | |
| 673 | Output: |
| 674 | Element 1: [AAAA] |
| 675 | Element 2: [BBBB] |
| 676 | Element 3: [CCCC] |
| 677 | ``` |
| 678 | |
| 679 | """ |
| 680 | return MultiModalFieldConfig( |
| 681 | field=MultiModalBatchedField(keep_on_cpu=keep_on_cpu), |
| 682 | modality=modality, |
| 683 | ) |
| 684 | |
| 685 | @staticmethod |
| 686 | def flat( |
| 687 | modality: str, |
| 688 | slices: Sequence[slice] | Sequence[Sequence[slice]], |
| 689 | dim: int = 0, |
| 690 | *, |
| 691 | keep_on_cpu: bool = False, |
| 692 | ): |
| 693 | """Defines a field where an element in the batch is obtained by |
| 694 | slicing along the first dimension of the underlying data. |
| 695 | |
| 696 | Args: |
| 697 | modality: The modality of the multi-modal item that uses this |
| 698 | keyword argument. |
| 699 | slices: For each multi-modal item, a slice (dim=0) or a tuple of |
| 700 | slices (dim>0) that is used to extract the data corresponding |
| 701 | to it. |
| 702 | dim: The dimension to extract data, default to 0. |
| 703 | keep_on_cpu: Whether to keep this field on the CPU for the model inputs. |
| 704 | |
| 705 | Example: |
| 706 | ``` |
| 707 | Given: |
| 708 | slices: [slice(0, 3), slice(3, 7), slice(7, 9)] |
| 709 | |
| 710 | Input: |
| 711 | Data: [AAABBBBCC] |
| 712 | |
| 713 | Output: |
| 714 | Element 1: [AAA] |
| 715 | Element 2: [BBBB] |
| 716 | Element 3: [CC] |
| 717 | ``` |
| 718 | |
| 719 | ``` |
| 720 | Given: |
| 721 | slices: [ |
| 722 | (slice(None), slice(0, 3)), |
| 723 | (slice(None), slice(3, 7)), |
| 724 | (slice(None), slice(7, 9))] |
| 725 | dim: 1 |
| 726 | |
| 727 | Input: |
| 728 | Data: [[A],[A],[A],[B],[B],[B],[B],[C],[C]] |
| 729 | |
| 730 | Output: |
| 731 | Element 1: [[A],[A],[A]] |
| 732 | Element 2: [[B],[B],[B],[B]] |
| 733 | Element 3: [[C],[C]] |
| 734 | ``` |
| 735 | |
| 736 | """ |
| 737 | return MultiModalFieldConfig( |
| 738 | field=MultiModalFlatField( |
| 739 | slices=slices, |
| 740 | dim=dim, |
| 741 | keep_on_cpu=keep_on_cpu, |
| 742 | ), |
| 743 | modality=modality, |
| 744 | ) |
| 745 | |
| 746 | @staticmethod |
| 747 | def flat_from_sizes( |
| 748 | modality: str, |
| 749 | size_per_item: "torch.Tensor", |
| 750 | dim: int = 0, |
| 751 | *, |
| 752 | keep_on_cpu: bool = False, |
| 753 | ): |
| 754 | """Defines a field where an element in the batch is obtained by |
| 755 | slicing along the first dimension of the underlying data. |
| 756 | |
| 757 | Args: |
| 758 | modality: The modality of the multi-modal item that uses this |
| 759 | keyword argument. |
| 760 | size_per_item: For each multi-modal item, the size of the slice |
| 761 | that is used to extract the data corresponding to it. |
| 762 | dim: The dimension to slice, default to 0. |
| 763 | keep_on_cpu: Whether to keep this field on the CPU for the model inputs. |
| 764 | |
| 765 | Example: |
| 766 | ``` |
| 767 | Given: |
| 768 | size_per_item: [3, 4, 2] |
| 769 | |
| 770 | Input: |
| 771 | Data: [AAABBBBCC] |
| 772 | |
| 773 | Output: |
| 774 | Element 1: [AAA] |
| 775 | Element 2: [BBBB] |
| 776 | Element 3: [CC] |
| 777 | ``` |
| 778 | |
| 779 | ``` |
| 780 | Given: |
| 781 | size_per_item: [3, 4, 2] |
| 782 | dim: 1 |
| 783 | |
| 784 | Input: |
| 785 | Data: [[A],[A],[A],[B],[B],[B],[B],[C],[C]] |
| 786 | |
| 787 | Output: |
| 788 | Element 1: [[A],[A],[A]] |
| 789 | Element 2: [[B],[B],[B],[B]] |
| 790 | Element 3: [[C],[C]] |
| 791 | ``` |
| 792 | |
| 793 | Info: |
| 794 | [`MultiModalFieldConfig.flat`][vllm.multimodal.inputs.MultiModalFieldConfig.flat] |
| 795 | |
| 796 | """ |
| 797 | if size_per_item.ndim != 1: |
| 798 | raise ValueError( |
| 799 | "size_per_item should be a 1-D tensor, " |
| 800 | f"but found shape: {size_per_item.shape}" |
| 801 | ) |
| 802 | |
| 803 | slice_idxs = [0, *accumulate(size_per_item)] |
| 804 | slices = [ |
| 805 | (slice(None, None, None),) * dim |
| 806 | + (slice(slice_idxs[i], slice_idxs[i + 1]),) |
| 807 | for i in range(len(size_per_item)) |
| 808 | ] |
| 809 | |
| 810 | return MultiModalFieldConfig.flat( |
| 811 | modality, |
| 812 | slices, |
| 813 | dim=dim, |
| 814 | keep_on_cpu=keep_on_cpu, |
| 815 | ) |
| 816 | |
| 817 | @staticmethod |
| 818 | def shared( |
| 819 | modality: str, |
| 820 | batch_size: int, |
| 821 | *, |
| 822 | keep_on_cpu: bool = False, |
| 823 | ): |
| 824 | """Defines a field where an element in the batch is obtained by |
| 825 | taking the entirety of the underlying data. |
| 826 | |
| 827 | This means that the data is the same for each element in the batch. |
| 828 | |
| 829 | Args: |
| 830 | modality: The modality of the multi-modal item that uses this |
| 831 | keyword argument. |
| 832 | batch_size: The number of multi-modal items which share this data. |
| 833 | keep_on_cpu: Whether to keep this field on the CPU for the model inputs. |
| 834 | |
| 835 | Example: |
| 836 | ``` |
| 837 | Given: |
| 838 | batch_size: 4 |
| 839 | |
| 840 | Input: |
| 841 | Data: [XYZ] |
| 842 | |
| 843 | Output: |
| 844 | Element 1: [XYZ] |
| 845 | Element 2: [XYZ] |
| 846 | Element 3: [XYZ] |
| 847 | Element 4: [XYZ] |
| 848 | ``` |
| 849 | |
| 850 | """ |
| 851 | return MultiModalFieldConfig( |
| 852 | field=MultiModalSharedField( |
| 853 | batch_size=batch_size, |
| 854 | keep_on_cpu=keep_on_cpu, |
| 855 | ), |
| 856 | modality=modality, |
| 857 | ) |
| 858 | |
| 859 | field: BaseMultiModalField |
| 860 | modality: str |
| 861 | |
| 862 | def build_elems( |
| 863 | self, |
| 864 | key: str, |
| 865 | batch: NestedTensors, |
| 866 | ) -> Sequence[MultiModalFieldElem]: |
| 867 | return self.field.build_elems(self.modality, key, batch) |
| 868 | |
| 869 | |
| 870 | class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]): |
| 871 | """A dictionary of processed keyword arguments to pass to the model, |
| 872 | corresponding to a single item in |
| 873 | [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems]. |
| 874 | """ |
| 875 | |
| 876 | @staticmethod |
| 877 | def dummy(nbytes: int = 1): |
| 878 | """Convenience class for testing.""" |
| 879 | mm_elem = MultiModalFieldElem( |
| 880 | data=torch.empty(nbytes, dtype=torch.uint8), |
| 881 | field=MultiModalSharedField(batch_size=1), |
| 882 | ) |
| 883 | return MultiModalKwargsItem({"dummy": mm_elem}) |
| 884 | |
| 885 | def get_data(self) -> dict[str, NestedTensors]: |
| 886 | return {key: elem.data for key, elem in self.items()} |
| 887 | |
| 888 | |
| 889 | _I = TypeVar( |
| 890 | "_I", |
| 891 | MultiModalKwargsItem, |
| 892 | MultiModalKwargsItem | None, |
| 893 | default=MultiModalKwargsItem, |
| 894 | ) |
| 895 | |
| 896 | |
| 897 | class MultiModalKwargsItems(UserDict[str, Sequence[_I]]): |
| 898 | """A dictionary of processed multi-modal inputs by modality. |
| 899 | |
| 900 | For example, given a processor that processes |
| 901 | images into `pixel_values` and `image_grid_thw`, |
| 902 | and audios into `input_audio_features`, |
| 903 | a prompt with 2 images and 1 audio will be processed |
| 904 | into a `MultiModalKwargsItems` with the following structure: |
| 905 | |
| 906 | ```python |
| 907 | MultiModalKwargsItems( |
| 908 | { |
| 909 | "image": [ |
| 910 | # For the first image |
| 911 | MultiModalKwargsItem({"pixel_values": ..., "image_grid_thw": ...}), |
| 912 | # For the second imgae |
| 913 | MultiModalKwargsItem({"pixel_values": ..., "image_grid_thw": ...}), |
| 914 | ], |
| 915 | "audio": [ |
| 916 | # For the first audio |
| 917 | MultiModalKwargsItem({"input_audio_features": ...}), |
| 918 | ], |
| 919 | } |
| 920 | ) |
| 921 | ``` |
| 922 | |
| 923 | Unlike HF processing which returns all items |
| 924 | in a single dictionary with batched keyword arguments, |
| 925 | we split up the items because some of them may already be cached. |
| 926 | Also, items from multiple requests may be batched together to improve throughput, |
| 927 | using the logic defined by the |
| 928 | [`BaseMultiModalField`][vllm.multimodal.inputs.BaseMultiModalField] |
| 929 | for each keyword argument. |
| 930 | """ |
| 931 | |
| 932 | @staticmethod |
| 933 | def from_hf_inputs( |
| 934 | hf_inputs: "BatchFeature", |
| 935 | config_by_key: Mapping[str, MultiModalFieldConfig], |
| 936 | ): |
| 937 | # NOTE: This skips fields in `hf_inputs` that are not in `config_by_key` |
| 938 | # We assume that those fields are not used in vLLM |
| 939 | elems_by_key = dict[str, Sequence[MultiModalFieldElem]]() |
| 940 | keys_by_modality = defaultdict[str, set[str]](set) |
| 941 | for key, config in config_by_key.items(): |
| 942 | batch = hf_inputs.get(key) |
| 943 | if batch is not None: |
| 944 | elems = config.build_elems(key, batch) |
| 945 | if len(elems) > 0: |
| 946 | elems_by_key[key] = elems |
| 947 | keys_by_modality[config.modality].add(key) |
| 948 | |
| 949 | items_by_modality = dict[str, list[MultiModalKwargsItem]]() |
| 950 | for modality, keys in keys_by_modality.items(): |
| 951 | elems_in_modality = {k: elems_by_key[k] for k in keys} |
| 952 | batch_sizes = {k: len(v) for k, v in elems_in_modality.items()} |
| 953 | |
| 954 | if len(set(batch_sizes.values())) > 1: |
| 955 | raise ValueError( |
| 956 | f"Cannot merge different batch sizes for {modality=}! " |
| 957 | f"Found: {batch_sizes=}" |
| 958 | ) |
| 959 | |
| 960 | batch_size = next(iter(batch_sizes.values())) |
| 961 | items_by_modality[modality] = [ |
| 962 | MultiModalKwargsItem({k: v[i] for k, v in elems_in_modality.items()}) |
| 963 | for i in range(batch_size) |
| 964 | ] |
| 965 | |
| 966 | return MultiModalKwargsItems(items_by_modality) |
| 967 | |
| 968 | def __getitem__(self, modality: str) -> Sequence[_I]: |
| 969 | if modality not in self: |
| 970 | raise KeyError( |
| 971 | f"Modality {modality!r} not found. " |
| 972 | f"Available modalities: {set(self.keys())}" |
| 973 | ) |
| 974 | |
| 975 | return super().__getitem__(modality) # type: ignore[return-value] |
| 976 | |
| 977 | def require_data(self) -> "MultiModalKwargsItems[MultiModalKwargsItem]": |
| 978 | for modality, items in self.items(): |
| 979 | for i, item in enumerate(items): |
| 980 | if item is None: |
| 981 | raise RuntimeError(f"Found empty mm_items[{modality}][{i}]") |
| 982 | |
| 983 | return self # type: ignore[return-value] |
| 984 | |
| 985 | def get_data( |
| 986 | self, |
| 987 | *, |
| 988 | device: torch.types.Device = None, |
| 989 | pin_memory: bool = False, |
| 990 | ) -> BatchedTensorInputs: |
| 991 | """Construct a dictionary of keyword arguments to pass to the model.""" |
| 992 | from .utils import group_and_batch_mm_items |
| 993 | |
| 994 | items_by_modality = self.require_data() |
| 995 | batches_by_modality = { |
| 996 | modality: [ |
| 997 | data |
| 998 | for _, data in group_and_batch_mm_items( |
| 999 | items, |
| 1000 | device=device, |
| 1001 | pin_memory=pin_memory, |
| 1002 | ) |
| 1003 | ] |
| 1004 | for modality, items in items_by_modality.items() |
| 1005 | if len(items) > 0 |
| 1006 | } |
| 1007 | |
| 1008 | out_data: BatchedTensorInputs = {} |
| 1009 | for _, batches in batches_by_modality.items(): |
| 1010 | if len(batches) != 1: |
| 1011 | num_batches_by_modality = { |
| 1012 | modality: len(batches) |
| 1013 | for modality, batches in batches_by_modality.items() |
| 1014 | } |
| 1015 | |
| 1016 | raise RuntimeError( |
| 1017 | f"Some modalities cannot be merged into a single batch " |
| 1018 | f"({num_batches_by_modality=})" |
| 1019 | ) |
| 1020 | |
| 1021 | out_data.update(batches[0]) |
| 1022 | |
| 1023 | return out_data |
| 1024 | |
| 1025 | |
| 1026 | MultiModalKwargsOptionalItems: TypeAlias = ( |
| 1027 | MultiModalKwargsItems[MultiModalKwargsItem] |
| 1028 | | MultiModalKwargsItems[MultiModalKwargsItem | None] |
| 1029 | ) |