"vscode:/vscode.git/clone" did not exist on "e2090bf3af96843c899d6f5c85d9c12b03b5cabb"
inputs.py 26.2 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
from abc import ABC, abstractmethod
4
from collections import UserDict, defaultdict
5
6
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
7
from functools import partial
8
from itertools import accumulate
9
10
from typing import (TYPE_CHECKING, Any, Literal, Optional, TypedDict, TypeVar,
                    Union, cast, final)
11
12

import numpy as np
13
from typing_extensions import NotRequired, TypeAlias
14

15
from vllm.jsontree import JSONTree, json_map_leaves
16
from vllm.utils import LazyLoader, full_groupby, is_list_of
17

18
if TYPE_CHECKING:
19
20
21
22
23
    import torch
    import torch.types
    from PIL.Image import Image
    from transformers.feature_extraction_utils import BatchFeature

24
    from .hasher import MultiModalHashDict
25
26
else:
    torch = LazyLoader("torch", globals(), "torch")
27

28
29
_T = TypeVar("_T")

30
HfImageItem: TypeAlias = Union["Image", np.ndarray, "torch.Tensor"]
31
"""
32
33
A {class}`transformers.image_utils.ImageInput` representing a single image
item, which can be passed to a HuggingFace `ImageProcessor`.
34
35
"""

36
37
HfVideoItem: TypeAlias = Union[list["Image"], np.ndarray, "torch.Tensor",
                               list[np.ndarray], list["torch.Tensor"]]
38
"""
39
40
A {class}`transformers.image_utils.VideoInput` representing a single video
item, which can be passed to a HuggingFace `VideoProcessor`.
41
42
"""

43
HfAudioItem: TypeAlias = Union[list[float], np.ndarray, "torch.Tensor"]
44
"""
45
Represents a single audio
46
item, which can be passed to a HuggingFace `AudioProcessor`.
47
48
"""

49
ImageItem: TypeAlias = Union[HfImageItem, "torch.Tensor"]
50
"""
51
52
A {class}`transformers.image_utils.ImageInput` representing a single image
item, which can be passed to a HuggingFace `ImageProcessor`.
53
54
55
56
57
58

Alternatively, a 3-D tensor or batch of 2-D tensors,
which are treated as image embeddings;
these are directly passed to the model without HF processing.
"""

59
VideoItem: TypeAlias = Union[HfVideoItem, "torch.Tensor"]
60
"""
61
62
A {class}`transformers.image_utils.VideoInput` representing a single video
item, which can be passed to a HuggingFace `VideoProcessor`.
63
64
65
66
67
68
69

Alternatively, a 3-D tensor or batch of 2-D tensors,
which are treated as video embeddings;
these are directly passed to the model without HF processing.
"""

AudioItem: TypeAlias = Union[HfAudioItem, tuple[np.ndarray, float],
70
                             "torch.Tensor"]
71
72
"""
Represents a single audio
73
item, which can be passed to a HuggingFace `AudioProcessor`.
74
75
76
77
78
79
80
81
82
83
84

Alternatively, a tuple `(audio, sampling_rate)`, where the sampling rate
is different from that expected by the model;
these are resampled to the model's sampling rate before being processed by HF.

Alternatively, a 3-D tensor or batch of 2-D tensors,
which are treated as audio embeddings;
these are directly passed to the model without HF processing.
"""

ModalityData: TypeAlias = Union[_T, list[_T]]
85
86
87
88
"""
Either a single data item, or a list of data items.

The number of data items allowed per modality is restricted by
89
`--limit-mm-per-prompt`.
90
91
92
93
94
95
96
"""


@final
class MultiModalDataBuiltins(TypedDict, total=False):
    """Type annotations for modality types predefined by vLLM."""

97
    image: ModalityData[ImageItem]
98
99
    """The input image(s)."""

100
    video: ModalityData[VideoItem]
101
102
    """The input video(s)."""

103
    audio: ModalityData[AudioItem]
104
105
106
    """The input audio(s)."""


107
MultiModalDataDict: TypeAlias = Mapping[str, ModalityData[Any]]
108
109
"""
A dictionary containing an entry for each modality type to input.
110

111
The built-in modalities are defined by {class}`MultiModalDataBuiltins`.
112
113
114
"""


115
116
@dataclass(frozen=True)
class PlaceholderRange:
117
118
119
    """
    Placeholder location information for multi-modal data.

120
121
    Example:

122
    Prompt: `AAAA BBBB What is in these images?`
123

124
    Images A and B will have:
125

126
127
128
129
    ```
    A: PlaceholderRange(offset=0, length=4)
    B: PlaceholderRange(offset=5, length=4)
    ```
130
131
132
133
134
135
136
137
    """

    offset: int
    """The start index of the placeholder in the prompt."""

    length: int
    """The length of the placeholder."""

138
    is_embed: Optional["torch.Tensor"] = None
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
    """
    A boolean mask of shape `(length,)` indicating which positions
    between `offset` and `offset + length` to assign embeddings to.
    """

    def get_num_embeds(self) -> int:
        if self.is_embed is None:
            return self.length

        return int(self.is_embed.sum().item())

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, self.__class__):
            return False
        if not (self.offset, self.length) == (other.offset, other.length):
            return False

        if self.is_embed is None:
            return other.is_embed is None
        if other.is_embed is None:
            return self.is_embed is None

        return nested_tensors_equal(self.is_embed, other.is_embed)

163

164
165
NestedTensors: TypeAlias = Union[list["NestedTensors"], list["torch.Tensor"],
                                 "torch.Tensor", tuple["torch.Tensor", ...]]
166
167
168
169
"""
Uses a list instead of a tensor if the dimensions of each element do not match.
"""

170
171

def nested_tensors_equal(a: NestedTensors, b: NestedTensors) -> bool:
172
    """Equality check between {data}`NestedTensors` objects."""
173
    if isinstance(a, torch.Tensor):
174
        return isinstance(b, torch.Tensor) and torch.equal(a, b)
175
    elif isinstance(b, torch.Tensor):
176
        return isinstance(a, torch.Tensor) and torch.equal(b, a)
177
178
179
180
181
182
183
184
185
186
187
188
189

    if isinstance(a, list):
        return (isinstance(b, list)
                and all(nested_tensors_equal(a_, b_) for a_, b_ in zip(a, b)))
    if isinstance(b, list):
        return (isinstance(a, list)
                and all(nested_tensors_equal(b_, a_) for b_, a_ in zip(b, a)))

    # Both a and b are scalars
    return a == b


BatchedTensorInputs: TypeAlias = Mapping[str, NestedTensors]
190
191
"""
A dictionary containing nested tensors which have been batched via
192
{meth}`MultiModalKwargs.batch`.
193
194
195
"""


196
@dataclass(frozen=True)
197
class MultiModalFieldElem:
198
199
    """
    Represents a keyword argument corresponding to a multi-modal item
200
    in {class}`MultiModalKwargs`.
201
202
203
204
205
206
207
208
209
210
    """

    modality: str
    """
    The modality of the corresponding multi-modal item.
    Each multi-modal item can consist of multiple keyword arguments.
    """

    key: str
    """
211
    The key of this field in {class}`MultiModalKwargs`,
212
213
214
    i.e. the name of the keyword argument to be passed to the model.
    """

215
    data: NestedTensors
216
    """
217
    The tensor data of this field in {class}`MultiModalKwargs`,
218
219
220
221
222
223
224
225
    i.e. the value of the keyword argument to be passed to the model.
    """

    field: "BaseMultiModalField"
    """
    Defines how to combine the tensor data of this field with others
    in order to batch multi-modal items together for model inference.
    """
226
227
228
229
230

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, self.__class__):
            return False

231
232
233
        return ((self.modality, self.key) == (other.modality, other.key)
                and nested_tensors_equal(self.data, other.data)
                and type(self.field) == type(other.field))  # noqa: E721
234
235
236
237


@dataclass(frozen=True)
class BaseMultiModalField(ABC):
238
239
    """
    Defines how to interpret tensor data belonging to a keyword argument in
240
    {class}`MultiModalKwargs` for multiple multi-modal items, and vice versa.
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
    """

    def _field_factory(self, *, modality: str, key: str):
        f = partial(
            MultiModalFieldElem,
            modality=modality,
            key=key,
            field=self,
        )

        # Allow passing data as positional argument
        def factory(data: NestedTensors) -> MultiModalFieldElem:
            return f(data=data)

        return factory
256
257

    @abstractmethod
258
259
260
261
262
263
264
    def build_elems(
        self,
        modality: str,
        key: str,
        data: NestedTensors,
    ) -> Sequence[MultiModalFieldElem]:
        """
265
        Construct {class}`MultiModalFieldElem` instances to represent
266
        the provided data.
267

268
        This is the inverse of {meth}`reduce_data`.
269
        """
270
271
        raise NotImplementedError

272
273
274
    @abstractmethod
    def _reduce_data(self, batch: list[NestedTensors]) -> NestedTensors:
        raise NotImplementedError
275

276
277
    def reduce_data(self, elems: list[MultiModalFieldElem]) -> NestedTensors:
        """
278
        Merge the data from multiple instances of {class}`MultiModalFieldElem`.
279

280
        This is the inverse of {meth}`build_elems`.
281
282
283
284
        """
        field_types = [type(item.field) for item in elems]
        if len(set(field_types)) > 1:
            raise ValueError(f"Cannot merge different {field_types=}")
285

286
        return self._reduce_data([item.data for item in elems])
287
288
289
290
291


@dataclass(frozen=True)
class MultiModalBatchedField(BaseMultiModalField):
    """
292
293
    Info:
        [MultiModalFieldConfig.batched][]
294
295
    """

296
297
298
299
300
301
302
303
    def build_elems(
        self,
        modality: str,
        key: str,
        data: NestedTensors,
    ) -> Sequence[MultiModalFieldElem]:
        field_factory = self._field_factory(modality=modality, key=key)
        return [field_factory(item) for item in data]
304
305
306

    def _reduce_data(self, batch: list[NestedTensors]) -> NestedTensors:
        if len(batch) > 0 and is_list_of(batch, torch.Tensor, check="all"):
307
308
309
310
311
            if len(batch) == 1:
                # An optimization when `batch` contains only one tensor:
                # - produce exactly same result as `torch.stack(batch)`
                # - will achieve zero-copy if the tensor is contiguous
                return batch[0].unsqueeze(0).contiguous()
312
            first_shape = batch[0].shape
313
            if all(elem.shape == first_shape for elem in batch):
314
315
316
317
318
319
320
321
                return torch.stack(batch)

        return batch


@dataclass(frozen=True)
class MultiModalFlatField(BaseMultiModalField):
    """
322
323
324
    Info:
        [MultiModalFieldConfig.flat][]
        [MultiModalFieldConfig.flat_from_sizes][]
325
    """
326
327
    slices: Union[Sequence[slice], Sequence[Sequence[slice]]]
    dim: int = 0
328

329
    def build_elems(
330
        self,
331
332
333
334
335
        modality: str,
        key: str,
        data: NestedTensors,
    ) -> Sequence[MultiModalFieldElem]:
        field_factory = self._field_factory(modality=modality, key=key)
336
337
338
339
        if not is_list_of(self.slices, slice, check="all"):
            assert isinstance(data, torch.Tensor), \
                "torch.Tensor is required for multiple slices"
        return [field_factory(data[cast(slice, s)]) for s in self.slices]
340
341
342

    def _reduce_data(self, batch: list[NestedTensors]) -> NestedTensors:
        if len(batch) > 0 and is_list_of(batch, torch.Tensor, check="all"):
343
344
345
346
347
            if len(batch) == 1:
                # An optimization when `batch` contains only one tensor:
                # - produce exactly same result as `torch.concat(batch)`
                # - will achieve zero-copy if the tensor is contiguous
                return batch[0].contiguous()
348

349
350
351
352
353
354
355
356
357
            def _expect_same_shape(tensor: torch.Tensor):
                return tensor.shape[:self.dim] + tensor.shape[self.dim + 1:]

            first_shape = _expect_same_shape(batch[0])

            if all(_expect_same_shape(elem) == first_shape for elem in batch):
                return torch.concat(batch, dim=self.dim)

        assert self.dim == 0, "dim == 0 is required for nested list"
358
        return [e for elem in batch for e in elem]
359
360


361
362
363
@dataclass(frozen=True)
class MultiModalSharedField(BaseMultiModalField):
    """
364
365
    Info:
        [MultiModalFieldConfig.shared][]
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
    """
    batch_size: int

    def build_elems(
        self,
        modality: str,
        key: str,
        data: NestedTensors,
    ) -> Sequence[MultiModalFieldElem]:
        field_factory = self._field_factory(modality=modality, key=key)
        return [field_factory(data)] * self.batch_size

    def _reduce_data(self, batch: list[NestedTensors]) -> NestedTensors:
        return batch[0]


382
383
384
385
class MultiModalFieldConfig:

    @staticmethod
    def batched(modality: str):
386
387
388
389
390
391
392
393
394
395
        """
        Defines a field where an element in the batch is obtained by
        indexing into the first dimension of the underlying data.

        Args:
            modality: The modality of the multi-modal item that uses this
                keyword argument.

        Example:

396
397
398
399
400
401
402
403
404
405
406
        ```
        Input:
            Data: [[AAAA]
                [BBBB]
                [CCCC]]

        Output:
            Element 1: [AAAA]
            Element 2: [BBBB]
            Element 3: [CCCC]
        ```
407
        """
408
        return MultiModalFieldConfig(
409
            field=MultiModalBatchedField(),
410
411
412
413
            modality=modality,
        )

    @staticmethod
414
415
416
    def flat(modality: str,
             slices: Union[Sequence[slice], Sequence[Sequence[slice]]],
             dim: int = 0):
417
418
419
420
421
422
423
        """
        Defines a field where an element in the batch is obtained by
        slicing along the first dimension of the underlying data.

        Args:
            modality: The modality of the multi-modal item that uses this
                keyword argument.
424
            slices: For each multi-modal item, a slice (dim=0) or a tuple of
425
                slices (dim>0) that is used to extract the data corresponding
426
427
                to it.
            dim: The dimension to extract data, default to 0.
428
429
430

        Example:

431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
        ```
        Given:
            slices: [slice(0, 3), slice(3, 7), slice(7, 9)]

        Input:
            Data: [AAABBBBCC]

        Output:
            Element 1: [AAA]
            Element 2: [BBBB]
            Element 3: [CC]
        ```

        ```
        Given:
            slices: [
                (slice(None), slice(0, 3)),
                (slice(None), slice(3, 7)),
                (slice(None), slice(7, 9))]
            dim: 1

        Input:
            Data: [[A],[A],[A],[B],[B],[B],[B],[C],[C]]

        Output:
            Element 1: [[A],[A],[A]]
            Element 2: [[B],[B],[B],[B]]
            Element 3: [[C],[C]]
        ```
460
        """
461
        return MultiModalFieldConfig(
462
            field=MultiModalFlatField(slices=slices, dim=dim),
463
464
465
            modality=modality,
        )

466
    @staticmethod
467
    def flat_from_sizes(modality: str,
468
                        size_per_item: "torch.Tensor",
469
                        dim: int = 0):
470
471
472
473
474
475
476
477
478
        """
        Defines a field where an element in the batch is obtained by
        slicing along the first dimension of the underlying data.

        Args:
            modality: The modality of the multi-modal item that uses this
                keyword argument.
            slices: For each multi-modal item, the size of the slice that
                is used to extract the data corresponding to it.
479
            dim: The dimension to slice, default to 0.
480
481
482

        Example:

483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
        ```
        Given:
            size_per_item: [3, 4, 2]

        Input:
            Data: [AAABBBBCC]

        Output:
            Element 1: [AAA]
            Element 2: [BBBB]
            Element 3: [CC]
        ```

        ```
        Given:
            slices: [3, 4, 2]
            dim: 1

        Input:
            Data: [[A],[A],[A],[B],[B],[B],[B],[C],[C]]

        Output:
            Element 1: [[A],[A],[A]]
            Element 2: [[B],[B],[B],[B]]
            Element 3: [[C],[C]]
        ```

510
511
        Info:
            [MultiModalFieldConfig.flat][]
512
513
        """

514
515
516
517
        if size_per_item.ndim != 1:
            raise ValueError("size_per_item should be a 1-D tensor, "
                             f"but found shape: {size_per_item.shape}")

518
        slice_idxs = [0, *accumulate(size_per_item)]
519
520
521
        slices = [(slice(None, None, None), ) * dim +
                  (slice(slice_idxs[i], slice_idxs[i + 1]), )
                  for i in range(len(size_per_item))]
522

523
        return MultiModalFieldConfig.flat(modality, slices, dim=dim)
524

525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
    @staticmethod
    def shared(modality: str, batch_size: int):
        """
        Defines a field where an element in the batch is obtained by
        taking the entirety of the underlying data.

        This means that the data is the same for each element in the batch.

        Args:
            modality: The modality of the multi-modal item that uses this
                keyword argument.
            batch_size: The number of multi-modal items which share this data.

        Example:

540
541
542
        ```
        Given:
            batch_size: 4
543

544
545
        Input:
            Data: [XYZ]
546

547
548
549
550
551
552
        Output:
            Element 1: [XYZ]
            Element 2: [XYZ]
            Element 3: [XYZ]
            Element 4: [XYZ]
        ```
553
554
555
556
557
558
559
        """
        return MultiModalFieldConfig(
            field=MultiModalSharedField(batch_size),
            modality=modality,
        )

    def __init__(self, field: BaseMultiModalField, modality: str) -> None:
560
561
        super().__init__()

562
        self.field = field
563
        self.modality = modality
564

565
    def build_elems(
566
567
568
        self,
        key: str,
        batch: NestedTensors,
569
    ) -> Sequence[MultiModalFieldElem]:
570
        return self.field.build_elems(self.modality, key, batch)
571
572


573
574
class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]):
    """
575
576
    A collection of {class}`MultiModalFieldElem`
    corresponding to a data item in {class}`MultiModalDataItems`.
577
    """
578

579
580
    @staticmethod
    def from_elems(elems: Sequence[MultiModalFieldElem]):
581
        return MultiModalKwargsItem({elem.key: elem for elem in elems})
582

583
584
    @property
    def modality(self) -> str:
585
        modalities = {elem.modality for elem in self.data.values()}
586
587
        assert len(modalities) == 1, f"Found different modalities={modalities}"
        return next(iter(modalities))
588
589


590
591
592
593
594
# NOTE: UserDict is for V0 compatibility.
# V1 should access individual items via `get_item`.
class MultiModalKwargs(UserDict[str, NestedTensors]):
    """
    A dictionary that represents the keyword arguments to
595
    {meth}`~torch.nn.Module.forward`.
596

597
598
599
    The metadata `items` enables us to obtain the keyword arguments
    corresponding to each data item in {class}`MultiModalDataItems`, via
    {meth}`get_item` and {meth}`get_items`.
600
601
    """

602
603
    @staticmethod
    def from_hf_inputs(
604
        hf_inputs: "BatchFeature",
605
606
607
608
        config_by_key: Mapping[str, MultiModalFieldConfig],
    ):
        # NOTE: This skips fields in `hf_inputs` that are not in `config_by_key`
        # We assume that those fields are not used in vLLM
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
        elems_by_key = dict[str, Sequence[MultiModalFieldElem]]()
        keys_by_modality = defaultdict[str, set[str]](set)
        for key, config in config_by_key.items():
            batch = hf_inputs.get(key)
            if batch is not None:
                elems = config.build_elems(key, batch)
                if len(elems) > 0:
                    elems_by_key[key] = elems
                    keys_by_modality[config.modality].add(key)

        items = list[MultiModalKwargsItem]()
        for modality, keys in keys_by_modality.items():
            elems_in_modality = {k: elems_by_key[k] for k in keys}
            batch_sizes = {k: len(v) for k, v in elems_in_modality.items()}

            if len(set(batch_sizes.values())) > 1:
                raise ValueError(
                    f"Cannot merge different batch sizes for {modality=}! "
                    f"Found: {batch_sizes=}")

            batch_size = next(iter(batch_sizes.values()))
            for item_idx in range(batch_size):
                elems = [v[item_idx] for v in elems_in_modality.values()]
                items.append(MultiModalKwargsItem.from_elems(elems))

        return MultiModalKwargs.from_items(items)
635
636

    @staticmethod
637
    def from_items(items: Sequence[MultiModalKwargsItem]):
638
        """Construct a new {class}`MultiModalKwargs` from multiple items."""
639
640
641
642
643
        elems_by_key = defaultdict[str, list[MultiModalFieldElem]](list)
        for item in items:
            for key, elem in item.items():
                elems_by_key[key].append(elem)

644
        data = {
645
            key: elems[0].field.reduce_data(elems)
646
            for key, elems in elems_by_key.items() if len(elems) > 0
647
648
        }

649
        return MultiModalKwargs(data, items=items)
650
651
652
653
654

    def __init__(
        self,
        data: Mapping[str, NestedTensors],
        *,
655
        items: Optional[Sequence[MultiModalKwargsItem]] = None,
656
657
658
    ) -> None:
        super().__init__(data)

659
660
        items_by_modality = full_groupby(items or [], key=lambda x: x.modality)
        self._items_by_modality = dict(items_by_modality)
661

662
663
664
    @property
    def modalities(self):
        return self._items_by_modality.keys()
665

666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
    @staticmethod
    def _try_stack(nested_tensors: NestedTensors) -> NestedTensors:
        """
        Stack the inner dimensions that have the same shape in
        a nested list of tensors.

        Thus, a dimension represented by a list means that the inner
        dimensions are different for each element along that dimension.
        """
        if isinstance(nested_tensors, torch.Tensor):
            return nested_tensors

        # TODO: Remove these once all models have been migrated
        if isinstance(nested_tensors, np.ndarray):
            return torch.from_numpy(nested_tensors)
        if isinstance(nested_tensors, (int, float)):
            return torch.tensor(nested_tensors)

        stacked = [MultiModalKwargs._try_stack(t) for t in nested_tensors]
        if not is_list_of(stacked, torch.Tensor, check="all"):
            # Only tensors (not lists) can be stacked.
            return stacked

689
        tensors_ = cast(list[torch.Tensor], stacked)
690
691
692
693
694
695
        if len(tensors_) == 1:
            # An optimization when `tensors_` contains only one tensor:
            # - produce exactly same result as `torch.stack(tensors_)`
            # - will achieve zero-copy if the tensor is contiguous
            return tensors_[0].unsqueeze(0).contiguous()

696
697
698
699
700
701
702
        if any(t.shape != tensors_[0].shape for t in tensors_):
            # The tensors have incompatible shapes and can't be stacked.
            return tensors_

        return torch.stack(tensors_)

    @staticmethod
703
    def batch(inputs_list: list["MultiModalKwargs"]) -> BatchedTensorInputs:
704
705
706
707
708
709
710
711
712
713
714
715
716
717
        """
        Batch multiple inputs together into a dictionary.

        The resulting dictionary has the same keys as the inputs.
        If the corresponding value from each input is a tensor and they all
        share the same shape, the output value is a single batched tensor;
        otherwise, the output value is a list containing the original value
        from each input.
        """
        if len(inputs_list) == 0:
            return {}

        # We need to consider the case where each item in the batch
        # contains different modalities (i.e. different keys).
718
        item_lists = defaultdict[str, list[NestedTensors]](list)
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743

        for inputs in inputs_list:
            for k, v in inputs.items():
                item_lists[k].append(v)

        return {
            k: MultiModalKwargs._try_stack(item_list)
            for k, item_list in item_lists.items()
        }

    @staticmethod
    def as_kwargs(
        batched_inputs: BatchedTensorInputs,
        *,
        device: torch.types.Device,
    ) -> BatchedTensorInputs:
        json_inputs = cast(JSONTree[torch.Tensor], batched_inputs)

        json_mapped = json_map_leaves(
            lambda x: x.to(device, non_blocking=True),
            json_inputs,
        )

        return cast(BatchedTensorInputs, json_mapped)

744
745
746
747
748
749
750
    def __delitem__(self, key: str) -> None:
        super().__delitem__(key)

        for items in self._items_by_modality.values():
            for item in items:
                item.pop(key, None)

751
752
753
    def __eq__(self, other: object) -> bool:
        if not isinstance(other, self.__class__):
            return False
754
        if self._items_by_modality != other._items_by_modality:
755
756
757
758
759
760
            return False

        ks = self.keys()
        return (ks == other.keys()
                and all(nested_tensors_equal(self[k], other[k]) for k in ks))

761
762
763
764
765
    def _validate_modality(self, method_name: str, modality: str) -> None:
        if not self._items_by_modality:
            raise RuntimeError(
                f"`{method_name}` is not supported when "
                "MultiModalKwargs is not initialized with `items`")
766

767
768
        if modality not in self._items_by_modality:
            available_modalities = set(self._items_by_modality.keys())
769
770
771
            raise KeyError(f"Modality {modality!r} not found. "
                           f"Available modalities: {available_modalities}")

772
773
774
775
    def get_item_count(self, modality: str) -> int:
        """Get the number of items belonging to a modality."""
        self._validate_modality("get_item_count", modality)
        return len(self._items_by_modality[modality])
776

777
778
779
780
781
782
783
    def get_item(self, modality: str, item_index: int) -> MultiModalKwargsItem:
        """
        Get the keyword arguments corresponding to an item identified by
        its modality and index.
        """
        self._validate_modality("get_item", modality)
        return self._items_by_modality[modality][item_index]
784

785
    def get_items(self, modality: str) -> Sequence[MultiModalKwargsItem]:
786
        """
787
788
        Get the keyword arguments corresponding to each item belonging to
        a modality.
789
        """
790
791
        self._validate_modality("get_items", modality)
        return self._items_by_modality[modality]
792

793

794
MultiModalPlaceholderDict: TypeAlias = Mapping[str, Sequence[PlaceholderRange]]
795
"""
796
A dictionary containing placeholder ranges for each modality.
797
798
799
"""


800
class MultiModalInputs(TypedDict):
801
    """
802
    Represents the outputs of
803
    {class}`vllm.multimodal.processing.BaseMultiModalProcessor`,
804
805
806
807
808
809
810
    ready to be passed to vLLM internals.
    """

    type: Literal["multimodal"]
    """The type of inputs."""

    prompt: str
811
    """The processed prompt text."""
812

813
    prompt_token_ids: list[int]
814
815
    """The processed token IDs which includes placeholder tokens."""

816
    token_type_ids: NotRequired[list[int]]
817
818
    """The token type IDs of the prompt."""

819
820
821
    mm_kwargs: MultiModalKwargs
    """Keyword arguments to be directly passed to the model after batching."""

822
    mm_hashes: Optional["MultiModalHashDict"]
823
824
    """The hashes of the multi-modal data."""

825
    mm_placeholders: "MultiModalPlaceholderDict"
826
827
    """
    For each modality, information about the placeholder tokens in
828
    `prompt_token_ids`.
829
    """
830

831
832
833
834
835
    cache_salt: NotRequired[str]
    """
    Optional cache salt to be used for prefix caching.
    """

836
837
838

class MultiModalEncDecInputs(MultiModalInputs):
    """
839
    Represents the outputs of {class}`vllm.multimodal.EncDecMultiModalProcessor`
840
841
842
843
844
845
846
847
848
849
850
    ready to be passed to vLLM internals.
    """

    encoder_prompt: str
    """The processed encoder prompt text."""

    encoder_prompt_token_ids: list[int]
    """The processed token IDs of the encoder prompt."""

    encoder_token_type_ids: NotRequired[list[int]]
    """The token type IDs of the encoder prompt."""