| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import operator |
| 4 | import sys |
| 5 | from abc import ABC, abstractmethod |
| 6 | from collections.abc import Mapping, Sequence |
| 7 | from typing import TYPE_CHECKING, Generic, TypeAlias, TypeVar |
| 8 | |
| 9 | import torch |
| 10 | from typing_extensions import override |
| 11 | |
| 12 | from vllm.config import ModelConfig |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.utils.cache import CacheInfo, LRUCache |
| 15 | from vllm.utils.jsontree import json_count_leaves, json_map_leaves, json_reduce_leaves |
| 16 | from vllm.utils.mem_constants import GiB_bytes |
| 17 | from vllm.utils.mem_utils import format_gib |
| 18 | |
| 19 | from ..inputs import ( |
| 20 | MultiModalFeatureSpec, |
| 21 | MultiModalFieldElem, |
| 22 | MultiModalKwargsItem, |
| 23 | MultiModalKwargsItems, |
| 24 | NestedTensors, |
| 25 | ) |
| 26 | |
| 27 | if TYPE_CHECKING: |
| 28 | from ..processing.processor import ResolvedPromptUpdate |
| 29 | |
| 30 | logger = init_logger(__name__) |
| 31 | |
| 32 | |
| 33 | class MultiModalCacheMissError(RuntimeError): |
| 34 | """Raised by the P1 receiver cache when items are requested with no data and |
| 35 | are not cached. |
| 36 | |
| 37 | P0 (frontend) keeps a metadata-only shadow of P1 (engine) and sends |
| 38 | ``data=None`` on a shadow hit. The two caches are updated in different orders |
| 39 | across processes, so they can drift -- leaving P0 referencing items P1 has |
| 40 | evicted. Raising (instead of asserting) lets the engine return a retryable |
| 41 | response and have P0 drop the stale entries |
| 42 | (``BaseMultiModalProcessorCache.invalidate``) so the client resends the data. |
| 43 | Carries every drifted ``mm_hash`` in the request so P0 can drop them all in one |
| 44 | pass -- one retry then recovers the whole request, not one item per retry. |
| 45 | """ |
| 46 | |
| 47 | def __init__(self, mm_hashes: list[str]) -> None: |
| 48 | super().__init__( |
| 49 | f"Multi-modal items {mm_hashes} are not in the receiver (P1) cache and " |
| 50 | "no data was provided to recompute them (P0/P1 cache drift); the request " |
| 51 | "should be retried with the multi-modal data attached." |
| 52 | ) |
| 53 | self.mm_hashes = mm_hashes |
| 54 | |
| 55 | |
| 56 | class MultiModalProcessorCacheItem: |
| 57 | """The data to store inside `MultiModalProcessorOnlyCache`. |
| 58 | |
| 59 | Args: |
| 60 | item: The processed tensor data corresponding to a multi-modal item. |
| 61 | prompt_updates: The prompt updates corresponding to `item`. |
| 62 | |
| 63 | """ |
| 64 | |
| 65 | def __init__( |
| 66 | self, |
| 67 | item: MultiModalKwargsItem, |
| 68 | prompt_updates: Sequence["ResolvedPromptUpdate"], |
| 69 | ) -> None: |
| 70 | super().__init__() |
| 71 | |
| 72 | self.item = item |
| 73 | self.prompt_updates = prompt_updates |
| 74 | |
| 75 | |
| 76 | class MultiModalProcessorCacheItemMetadata: |
| 77 | """The metadata to store inside `MultiModalProcessorSenderCache`. |
| 78 | |
| 79 | Args: |
| 80 | item: The processed tensor data corresponding to a multi-modal item. |
| 81 | Since P1 already stores the tensor data, we only store its size |
| 82 | metadata in P0 to reduce memory usage. The size metadata is still |
| 83 | needed to keep the same cache eviction policy as P0. |
| 84 | prompt_updates: The prompt updates corresponding to `item`. |
| 85 | This needs to stay on P0 because for some models, they are |
| 86 | dependent on the processed tensor data (cached on P1). |
| 87 | |
| 88 | """ |
| 89 | |
| 90 | def __init__( |
| 91 | self, |
| 92 | item: MultiModalKwargsItem, |
| 93 | prompt_updates: Sequence["ResolvedPromptUpdate"], |
| 94 | ) -> None: |
| 95 | super().__init__() |
| 96 | |
| 97 | self.item_size = MultiModalCache.get_item_size(item) |
| 98 | self.prompt_updates = prompt_updates |
| 99 | |
| 100 | |
| 101 | MultiModalCacheValue: TypeAlias = ( |
| 102 | MultiModalProcessorCacheItem |
| 103 | | MultiModalProcessorCacheItemMetadata |
| 104 | | MultiModalKwargsItems |
| 105 | | MultiModalKwargsItem |
| 106 | | Mapping[str, NestedTensors] |
| 107 | ) |
| 108 | |
| 109 | _V = TypeVar("_V", bound=MultiModalCacheValue) |
| 110 | |
| 111 | |
| 112 | class MultiModalCache: |
| 113 | @classmethod |
| 114 | def get_leaf_size(cls, leaf: object) -> int: |
| 115 | if isinstance(leaf, MultiModalProcessorCacheItem): |
| 116 | return cls.get_leaf_size(leaf.item) |
| 117 | if isinstance(leaf, MultiModalProcessorCacheItemMetadata): |
| 118 | return leaf.item_size |
| 119 | |
| 120 | # These are not subclasses of dict |
| 121 | if isinstance( |
| 122 | leaf, |
| 123 | (MultiModalKwargsItems, MultiModalKwargsItem, MultiModalFieldElem), |
| 124 | ): |
| 125 | return cls.get_item_size(leaf.data) # type: ignore |
| 126 | |
| 127 | # sys.getsizeof doesn't work for tensors |
| 128 | if isinstance(leaf, torch.Tensor): |
| 129 | return leaf.nbytes |
| 130 | |
| 131 | return sys.getsizeof(leaf) |
| 132 | |
| 133 | @classmethod |
| 134 | def get_item_size( |
| 135 | cls, |
| 136 | value: MultiModalCacheValue, |
| 137 | *, |
| 138 | debug: bool = False, |
| 139 | ) -> int: |
| 140 | size = json_reduce_leaves( |
| 141 | operator.add, json_map_leaves(cls.get_leaf_size, value) |
| 142 | ) |
| 143 | |
| 144 | if debug: |
| 145 | leaf_count = json_count_leaves(value) |
| 146 | logger.debug( |
| 147 | "Calculated size of %s to be %s GiB (%d leaves)", |
| 148 | type(value), |
| 149 | format_gib(size), |
| 150 | leaf_count, |
| 151 | ) |
| 152 | |
| 153 | return size |
| 154 | |
| 155 | @classmethod |
| 156 | def get_item_complexity(cls, value: MultiModalCacheValue) -> int: |
| 157 | """Get the number of leaf elements in a multi-modal cache value. |
| 158 | |
| 159 | This provides a measure of structural complexity that can be useful |
| 160 | for debugging cache performance and understanding data patterns. |
| 161 | |
| 162 | Args: |
| 163 | value: The multi-modal cache value to analyze. |
| 164 | |
| 165 | Returns: |
| 166 | The number of leaf elements in the nested structure. |
| 167 | |
| 168 | """ |
| 169 | return json_count_leaves(value) |
| 170 | |
| 171 | @classmethod |
| 172 | def get_lru_cache( |
| 173 | cls, |
| 174 | capacity_gb: float, |
| 175 | value_type: type[_V], |
| 176 | *, |
| 177 | debug: bool = False, |
| 178 | ) -> LRUCache[str, _V]: |
| 179 | return LRUCache( |
| 180 | GiB_bytes * capacity_gb, |
| 181 | getsizeof=lambda x: cls.get_item_size(x, debug=debug), |
| 182 | ) |
| 183 | |
| 184 | |
| 185 | _I = TypeVar("_I", contravariant=True) |
| 186 | _O = TypeVar("_O", covariant=True) |
| 187 | |
| 188 | |
| 189 | class BaseMultiModalCache(ABC, Generic[_I, _O]): |
| 190 | """Abstract base class to read/write multi-modal items from cache. |
| 191 | |
| 192 | The idea of multi-modal caching is based on having a client and server |
| 193 | where the client executes in the frontend process (=P0) and |
| 194 | the server in the core process (=P1). The data flow is as follows: |
| 195 | |
| 196 | ``` |
| 197 | is_cached() x N get_and_update() |
| 198 | P0: From API -----------------> -----------------> To P1 |
| 199 | |
| 200 | get_and_update() |
| 201 | P1: From P0 -----------------> To model |
| 202 | ``` |
| 203 | |
| 204 | `is_cached()` can be called any number of times in P0. However, |
| 205 | `get_and_update()` must be called in P0 and P1 one after another |
| 206 | so that their cache eviction order remains the same. |
| 207 | |
| 208 | This ensures that the keys in P0 and P1 caches are mirrored, |
| 209 | allowing us to determine whether a key is cached in P1 by looking |
| 210 | up the P0 cache, without having to communicate with P1. |
| 211 | """ |
| 212 | |
| 213 | @abstractmethod |
| 214 | def get_and_update_item( |
| 215 | self, |
| 216 | mm_item: _I, |
| 217 | mm_hash: str, |
| 218 | ) -> _O: |
| 219 | """Possibly update a multi-modal item based on whether it is |
| 220 | in the underlying cache. |
| 221 | |
| 222 | This update is done out-of-place and updates the cache eviction order. |
| 223 | |
| 224 | Args: |
| 225 | mm_item: The multi-modal item to update. |
| 226 | mm_hash: The hash of `mm_item`. |
| 227 | |
| 228 | Returns: |
| 229 | The update multi-modal item. |
| 230 | |
| 231 | """ |
| 232 | raise NotImplementedError |
| 233 | |
| 234 | def get_and_update( |
| 235 | self, |
| 236 | mm_items: Sequence[_I], |
| 237 | mm_hashes: list[str], |
| 238 | ) -> list[_O]: |
| 239 | """Possibly update a sequence of multi-modal items based on whether they |
| 240 | are in the underlying cache. |
| 241 | |
| 242 | This update is done out-of-place and updates the cache eviction order. |
| 243 | |
| 244 | Args: |
| 245 | mm_items: The multi-modal items to update. |
| 246 | mm_hashes: The hash of each item in `mm_items`. |
| 247 | |
| 248 | Returns: |
| 249 | A new list of updated multi-modal items. |
| 250 | |
| 251 | """ |
| 252 | assert len(mm_items) == len(mm_hashes) |
| 253 | |
| 254 | return [ |
| 255 | self.get_and_update_item(mm_item, mm_hash) |
| 256 | for mm_item, mm_hash in zip(mm_items, mm_hashes) |
| 257 | ] |
| 258 | |
| 259 | def cache_if_fits( |
| 260 | self, |
| 261 | cache: LRUCache[str, _V], |
| 262 | key: str, |
| 263 | value: _V, |
| 264 | ) -> bool: |
| 265 | """Insert `value` if it fits in `cache`. |
| 266 | |
| 267 | `cachetools.Cache` raises `ValueError("value too large")` when a |
| 268 | single item exceeds `maxsize`. An item bigger than the whole |
| 269 | processor cache can never be a hit, so skip the insert and serve it |
| 270 | uncached instead of aborting engine startup. |
| 271 | |
| 272 | LRU P0/P1 caches call this so they stay mirrored. SHM subclasses do |
| 273 | not use it; they already skip oversize items in `put()`. |
| 274 | |
| 275 | Args: |
| 276 | cache: The LRU cache to update. |
| 277 | key: Cache key (typically the multi-modal item hash). |
| 278 | value: Value to insert. |
| 279 | |
| 280 | Returns: |
| 281 | `True` if the item was cached, otherwise `False`. |
| 282 | |
| 283 | """ |
| 284 | if cache.put_if_fits(key, value): |
| 285 | return True |
| 286 | logger.warning_once( |
| 287 | "Skipping multi-modal processor cache insert for an item of " |
| 288 | "%s GiB because it exceeds --mm-processor-cache-gb=%s. " |
| 289 | "The item will be processed uncached; increase " |
| 290 | "--mm-processor-cache-gb to cache items of this size.", |
| 291 | format_gib(int(cache.getsizeof(value))), |
| 292 | format_gib(int(cache.maxsize)), |
| 293 | ) |
| 294 | return False |
| 295 | |
| 296 | @abstractmethod |
| 297 | def clear_cache(self) -> None: |
| 298 | """Clear the underlying cache.""" |
| 299 | raise NotImplementedError |
| 300 | |
| 301 | |
| 302 | MultiModalProcessorCacheInItem: TypeAlias = ( |
| 303 | tuple[MultiModalKwargsItem, Sequence["ResolvedPromptUpdate"]] | None |
| 304 | ) |
| 305 | |
| 306 | |
| 307 | MultiModalProcessorCacheOutItem: TypeAlias = tuple[ |
| 308 | MultiModalKwargsItem | None, Sequence["ResolvedPromptUpdate"] |
| 309 | ] |
| 310 | |
| 311 | |
| 312 | class BaseMultiModalProcessorCache( |
| 313 | BaseMultiModalCache[MultiModalProcessorCacheInItem, MultiModalProcessorCacheOutItem] |
| 314 | ): |
| 315 | """The required interface for caches on P0.""" |
| 316 | |
| 317 | @abstractmethod |
| 318 | def is_cached_item(self, mm_hash: str) -> bool: |
| 319 | """Check whether a multi-modal item is |
| 320 | in the underlying cache. |
| 321 | |
| 322 | This **DOES NOT** update the cache eviction order. |
| 323 | |
| 324 | Args: |
| 325 | mm_hash: The hash of the item to check. |
| 326 | |
| 327 | Returns: |
| 328 | `True` if the item is cached, otherwise `False`. |
| 329 | |
| 330 | """ |
| 331 | raise NotImplementedError |
| 332 | |
| 333 | def is_cached(self, mm_hashes: list[str]) -> list[bool]: |
| 334 | """Check whether a sequence of multi-modal items are |
| 335 | in the underlying cache. |
| 336 | |
| 337 | This **DOES NOT** update the cache eviction order. |
| 338 | |
| 339 | Args: |
| 340 | mm_hashes: The hash of each item to check. |
| 341 | |
| 342 | Returns: |
| 343 | For each item, `True` if the item is cached, otherwise `False`. |
| 344 | |
| 345 | """ |
| 346 | return [self.is_cached_item(mm_hash) for mm_hash in mm_hashes] |
| 347 | |
| 348 | def invalidate(self, mm_hash: str) -> None: |
| 349 | """Drop ``mm_hash`` from this P0 shadow cache to recover from P0/P1 drift. |
| 350 | |
| 351 | No-op by default; shadow caches that can drift from P1 override this. |
| 352 | """ |
| 353 | pass |
| 354 | |
| 355 | def close(self) -> None: |
| 356 | """Close the underlying cache, if needed.""" |
| 357 | pass |
| 358 | |
| 359 | def validate_input_item( |
| 360 | self, |
| 361 | mm_item: MultiModalKwargsItem, |
| 362 | mm_hash: str, |
| 363 | ) -> None: |
| 364 | """Validate externally supplied cache metadata before engine handoff.""" |
| 365 | return None |
| 366 | |
| 367 | @abstractmethod |
| 368 | def touch_sender_cache_item(self, mm_hash: str) -> None: |
| 369 | """Update the cache eviction order for a multi-modal item. |
| 370 | |
| 371 | This is used to touch the item in the cache without changing |
| 372 | its value. |
| 373 | |
| 374 | Args: |
| 375 | mm_hash: The hash of the multi-modal item. |
| 376 | |
| 377 | """ |
| 378 | raise NotImplementedError |
| 379 | |
| 380 | def release_sender_touches(self) -> None: |
| 381 | """Release the items that were touched but not updated afterwards.""" |
| 382 | return None |
| 383 | |
| 384 | @abstractmethod |
| 385 | def make_stats(self, *, delta: bool = False) -> CacheInfo: |
| 386 | """Get (and reset) the multi-modal cache stats. |
| 387 | |
| 388 | Returns: |
| 389 | The current multi-modal caching stats. |
| 390 | |
| 391 | """ |
| 392 | raise NotImplementedError |
| 393 | |
| 394 | |
| 395 | class MultiModalProcessorOnlyCache(BaseMultiModalProcessorCache): |
| 396 | """The cache which is used on P0 when IPC caching is disabled. |
| 397 | |
| 398 | How to update each item: |
| 399 | |
| 400 | - If the item is in the cache, replace the input with the cached item. |
| 401 | - If the item is not in the cache, store that item (which includes |
| 402 | tensor data and metadata) into the cache, and return the input. |
| 403 | """ |
| 404 | |
| 405 | def __init__(self, model_config: ModelConfig) -> None: |
| 406 | super().__init__() |
| 407 | |
| 408 | mm_config = model_config.get_multimodal_config() |
| 409 | |
| 410 | self._cache = MultiModalCache.get_lru_cache( |
| 411 | mm_config.mm_processor_cache_gb, |
| 412 | MultiModalProcessorCacheItem, |
| 413 | ) |
| 414 | |
| 415 | @override |
| 416 | def is_cached_item(self, mm_hash: str) -> bool: |
| 417 | return mm_hash in self._cache |
| 418 | |
| 419 | @override |
| 420 | def get_and_update_item( |
| 421 | self, |
| 422 | mm_item: MultiModalProcessorCacheInItem, |
| 423 | mm_hash: str, |
| 424 | ) -> MultiModalProcessorCacheOutItem: |
| 425 | if (cached_item := self._cache.get(mm_hash)) is not None: |
| 426 | return cached_item.item, cached_item.prompt_updates |
| 427 | |
| 428 | assert mm_item is not None, f"Expected a cached item for {mm_hash=}" |
| 429 | |
| 430 | self.cache_if_fits(self._cache, mm_hash, MultiModalProcessorCacheItem(*mm_item)) |
| 431 | return mm_item |
| 432 | |
| 433 | @override |
| 434 | def touch_sender_cache_item(self, mm_hash: str) -> None: |
| 435 | self._cache.touch(mm_hash) |
| 436 | |
| 437 | @override |
| 438 | def clear_cache(self) -> None: |
| 439 | self._cache.clear() |
| 440 | |
| 441 | @override |
| 442 | def make_stats(self, *, delta: bool = False) -> CacheInfo: |
| 443 | return self._cache.stat(delta=delta) |
| 444 | |
| 445 | |
| 446 | class BaseMultiModalReceiverCache( |
| 447 | BaseMultiModalCache[MultiModalKwargsItem | None, MultiModalKwargsItem] |
| 448 | ): |
| 449 | """The required interface for caches on P1.""" |
| 450 | |
| 451 | def get_and_update_features( |
| 452 | self, |
| 453 | mm_features: list["MultiModalFeatureSpec"], |
| 454 | ) -> list["MultiModalFeatureSpec"]: |
| 455 | """Update multimodal features with cached encoder outputs. |
| 456 | Touch all identifier at first before update to avoid |
| 457 | item in updated list evict during update. |
| 458 | |
| 459 | Uses mm_hash for cache key to share across LoRAs (falls back to |
| 460 | identifier for backward compatibility). |
| 461 | """ |
| 462 | for feature in mm_features: |
| 463 | cache_key = feature.mm_hash or feature.identifier |
| 464 | self.touch_receiver_cache_item(cache_key, feature.data) |
| 465 | |
| 466 | missing_mm_hashes: list[str] = [] |
| 467 | for feature in mm_features: |
| 468 | cache_key = feature.mm_hash or feature.identifier |
| 469 | try: |
| 470 | feature.data = self.get_and_update_item(feature.data, cache_key) |
| 471 | except MultiModalCacheMissError as e: |
| 472 | # Collect every drifted hash in this request before raising, so the |
| 473 | # engine can have P0 invalidate them all at once -- otherwise a |
| 474 | # request with k drifted items needs k client retries (each resend |
| 475 | # only un-shadows the one reported hash). |
| 476 | missing_mm_hashes.extend(e.mm_hashes) |
| 477 | if missing_mm_hashes: |
| 478 | raise MultiModalCacheMissError(missing_mm_hashes) |
| 479 | return mm_features |
| 480 | |
| 481 | @abstractmethod |
| 482 | def touch_receiver_cache_item( |
| 483 | self, |
| 484 | mm_hash: str, |
| 485 | mm_item: MultiModalKwargsItem | None = None, |
| 486 | ) -> None: |
| 487 | """Update the cache eviction order for a multi-modal item. |
| 488 | |
| 489 | This is used to touch the item in the cache without changing |
| 490 | its value. |
| 491 | |
| 492 | Args: |
| 493 | mm_hash: The hash of the multi-modal item. |
| 494 | mm_item: The multi-modal item itself. This is optional and |
| 495 | may not be needed by some cache implementations. |
| 496 | |
| 497 | """ |
| 498 | raise NotImplementedError |