| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Helpers for render-time multimodal feature extraction and generate input.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from collections.abc import Callable, Collection, Sequence |
| 8 | from typing import cast |
| 9 | |
| 10 | from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import ( |
| 11 | decode_mm_kwargs_item, |
| 12 | encode_mm_kwargs_item, |
| 13 | ) |
| 14 | from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( |
| 15 | MultiModalFeatures, |
| 16 | PlaceholderRangeInfo, |
| 17 | ) |
| 18 | from vllm.inputs import ( |
| 19 | EngineInput, |
| 20 | MultiModalHashes, |
| 21 | MultiModalInput, |
| 22 | MultiModalPlaceholders, |
| 23 | ) |
| 24 | from vllm.multimodal.inputs import ( |
| 25 | MultiModalKwargsItem, |
| 26 | MultiModalKwargsOptionalItems, |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | def merge_mm_kwargs_items( |
| 31 | kwargs_item: MultiModalKwargsItem | None, |
| 32 | metadata_item: MultiModalKwargsItem | None, |
| 33 | ) -> MultiModalKwargsItem | None: |
| 34 | """Merge full multimodal data with its metadata-only counterpart.""" |
| 35 | if kwargs_item is None: |
| 36 | return metadata_item |
| 37 | if metadata_item is None: |
| 38 | return kwargs_item |
| 39 | return MultiModalKwargsItem({**kwargs_item, **metadata_item}) |
| 40 | |
| 41 | |
| 42 | def _encode_metadata_items( |
| 43 | items: Sequence[MultiModalKwargsItem | None], |
| 44 | *, |
| 45 | declared: Collection[str], |
| 46 | ) -> list[str | None]: |
| 47 | """Serialize placeholder-metadata and ``keep_on_cpu`` fields per item.""" |
| 48 | metadata_items: list[str | None] = [] |
| 49 | for item in items: |
| 50 | if item is None: |
| 51 | metadata_items.append(None) |
| 52 | continue |
| 53 | |
| 54 | metadata_item = MultiModalKwargsItem( |
| 55 | { |
| 56 | key: elem |
| 57 | for key, elem in item.items() |
| 58 | if elem.field.keep_on_cpu or key in declared |
| 59 | } |
| 60 | ) |
| 61 | metadata_items.append( |
| 62 | encode_mm_kwargs_item(metadata_item) if metadata_item else None |
| 63 | ) |
| 64 | return metadata_items |
| 65 | |
| 66 | |
| 67 | def _encode_mm_kwargs_with_metadata( |
| 68 | raw_mm_kwargs: MultiModalKwargsOptionalItems, |
| 69 | *, |
| 70 | metadata_fields_for: Callable[[str], Collection[str]] | None = None, |
| 71 | ) -> tuple[dict[str, list[str | None]], dict[str, list[str | None]] | None]: |
| 72 | """Serialize full kwargs and their metadata-only subsets per modality.""" |
| 73 | kwargs_data: dict[str, list[str | None]] = {} |
| 74 | metadata_by_modality: dict[str, list[str | None]] = {} |
| 75 | |
| 76 | for modality, items in raw_mm_kwargs.items(): |
| 77 | kwargs_data[modality] = [ |
| 78 | encode_mm_kwargs_item(item) if item is not None else None for item in items |
| 79 | ] |
| 80 | |
| 81 | declared = ( |
| 82 | set(metadata_fields_for(modality)) |
| 83 | if metadata_fields_for is not None |
| 84 | else set() |
| 85 | ) |
| 86 | metadata_items = _encode_metadata_items(items, declared=declared) |
| 87 | if any(item is not None for item in metadata_items): |
| 88 | metadata_by_modality[modality] = metadata_items |
| 89 | |
| 90 | mm_metadata = metadata_by_modality or None |
| 91 | return kwargs_data, mm_metadata |
| 92 | |
| 93 | |
| 94 | def mm_kwargs_from_features( |
| 95 | features: MultiModalFeatures, |
| 96 | ) -> dict[str, list[MultiModalKwargsItem | None]]: |
| 97 | """Deserialize ``features`` into per-modality kwargs for ``mm_input``.""" |
| 98 | mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {} |
| 99 | kwargs_data = features.kwargs_data or {} |
| 100 | mm_metadata = features.mm_metadata or {} |
| 101 | for modality, hashes in features.mm_hashes.items(): |
| 102 | n = len(hashes) |
| 103 | kwargs_items = [ |
| 104 | decode_mm_kwargs_item(item) if item is not None else None |
| 105 | for item in kwargs_data.get(modality, [None] * n) |
| 106 | ] |
| 107 | metadata_items = [ |
| 108 | decode_mm_kwargs_item(item) if item is not None else None |
| 109 | for item in mm_metadata.get(modality, [None] * n) |
| 110 | ] |
| 111 | mm_kwargs[modality] = [ |
| 112 | merge_mm_kwargs_items(kwargs_item, metadata_item) |
| 113 | for kwargs_item, metadata_item in zip( |
| 114 | kwargs_items, metadata_items, strict=True |
| 115 | ) |
| 116 | ] |
| 117 | return mm_kwargs |
| 118 | |
| 119 | |
| 120 | def placeholder_ranges_from_engine_input( |
| 121 | engine_input: EngineInput, |
| 122 | ) -> dict[str, list[PlaceholderRangeInfo]] | None: |
| 123 | """Return per-modality placeholder ranges, or ``None`` for text-only prompts.""" |
| 124 | if engine_input.get("type") != "multimodal": |
| 125 | return None |
| 126 | |
| 127 | raw_placeholders: MultiModalPlaceholders = cast(MultiModalInput, engine_input)[ |
| 128 | "mm_placeholders" |
| 129 | ] |
| 130 | return { |
| 131 | modality: [ |
| 132 | PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges |
| 133 | ] |
| 134 | for modality, ranges in raw_placeholders.items() |
| 135 | } |
| 136 | |
| 137 | |
| 138 | def extract_mm_features( |
| 139 | engine_input: EngineInput, |
| 140 | *, |
| 141 | metadata_fields_for: Callable[[str], Collection[str]] | None = None, |
| 142 | ) -> MultiModalFeatures | None: |
| 143 | """Extract multimodal features from a rendered engine prompt. |
| 144 | |
| 145 | Returns ``None`` for text-only prompts. ``mm_metadata`` keeps the |
| 146 | intersection of processed kwargs that prefill needs after EC transfer: |
| 147 | fields declared as embedding metadata, plus fields marked |
| 148 | ``keep_on_cpu`` (for example M-RoPE grid dims). |
| 149 | """ |
| 150 | if engine_input.get("type") != "multimodal": |
| 151 | return None |
| 152 | |
| 153 | mm_engine_input = cast(MultiModalInput, engine_input) |
| 154 | mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"] |
| 155 | mm_placeholders = placeholder_ranges_from_engine_input(engine_input) |
| 156 | assert mm_placeholders is not None |
| 157 | |
| 158 | kwargs_data: dict[str, list[str | None]] | None = None |
| 159 | mm_metadata: dict[str, list[str | None]] | None = None |
| 160 | if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"): |
| 161 | kwargs_data, mm_metadata = _encode_mm_kwargs_with_metadata( |
| 162 | raw_mm_kwargs, |
| 163 | metadata_fields_for=metadata_fields_for, |
| 164 | ) |
| 165 | |
| 166 | return MultiModalFeatures( |
| 167 | mm_hashes=mm_hashes, |
| 168 | mm_placeholders=mm_placeholders, |
| 169 | kwargs_data=kwargs_data, |
| 170 | mm_metadata=mm_metadata, |
| 171 | ) |