| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Implementation of Shared Memory Cache (see docs/configuration/optimization.md).""" |
| 4 | |
| 5 | from collections.abc import Sequence |
| 6 | from multiprocessing.synchronize import Lock as LockType |
| 7 | from typing import TYPE_CHECKING, TypeAlias |
| 8 | |
| 9 | from typing_extensions import override |
| 10 | |
| 11 | import vllm.envs as envs |
| 12 | from vllm.config import VllmConfig |
| 13 | from vllm.distributed.device_communicators.shm_object_storage import ( |
| 14 | MsgpackSerde, |
| 15 | SingleWriterShmObjectStorage, |
| 16 | SingleWriterShmRingBuffer, |
| 17 | ) |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.utils.cache import CacheInfo |
| 20 | from vllm.utils.mem_constants import GiB_bytes, MiB_bytes |
| 21 | |
| 22 | from ..inputs import ( |
| 23 | MultiModalBatchedField, |
| 24 | MultiModalFeatureSpec, |
| 25 | MultiModalFieldElem, |
| 26 | MultiModalKwargsItem, |
| 27 | ) |
| 28 | from .base import ( |
| 29 | BaseMultiModalProcessorCache, |
| 30 | BaseMultiModalReceiverCache, |
| 31 | MultiModalProcessorCacheInItem, |
| 32 | MultiModalProcessorCacheOutItem, |
| 33 | ) |
| 34 | |
| 35 | if TYPE_CHECKING: |
| 36 | from ..processing.processor import ResolvedPromptUpdate |
| 37 | |
| 38 | logger = init_logger(__name__) |
| 39 | |
| 40 | ShmHandle: TypeAlias = tuple[int, int, list[int]] |
| 41 | |
| 42 | |
| 43 | def _get_shm_handle(mm_item: MultiModalKwargsItem) -> ShmHandle | None: |
| 44 | if "address" not in mm_item: |
| 45 | return None |
| 46 | if "monotonic_id" not in mm_item or "signature" not in mm_item: |
| 47 | raise ValueError("Missing SHM handle signature for cache key.") |
| 48 | |
| 49 | address = mm_item["address"].data |
| 50 | monotonic_id = mm_item["monotonic_id"].data |
| 51 | signature = mm_item["signature"].data |
| 52 | if ( |
| 53 | type(address) is not int |
| 54 | or not 0 <= address < 2**64 |
| 55 | or type(monotonic_id) is not int |
| 56 | or not 0 <= monotonic_id < 2**32 |
| 57 | or not isinstance(signature, list) |
| 58 | ): |
| 59 | raise ValueError("Invalid SHM handle signature for cache key.") |
| 60 | |
| 61 | return address, monotonic_id, signature |
| 62 | |
| 63 | |
| 64 | class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache): |
| 65 | """The cache which is used on P0 when SHM caching is enabled. |
| 66 | |
| 67 | How to update each item: |
| 68 | |
| 69 | - If the item is already in the cache, clear the input to avoid |
| 70 | unnecessary IPC. |
| 71 | |
| 72 | - If the item is not in the cache, store the data in shared memory. |
| 73 | """ |
| 74 | |
| 75 | def __init__(self, vllm_config: VllmConfig) -> None: |
| 76 | super().__init__() |
| 77 | |
| 78 | self.world_size = vllm_config.parallel_config.world_size |
| 79 | mm_config = vllm_config.model_config.get_multimodal_config() |
| 80 | |
| 81 | ring_buffer = SingleWriterShmRingBuffer( |
| 82 | data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes), |
| 83 | name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME, |
| 84 | create=True, # sender is the writer |
| 85 | ) |
| 86 | self._shm_cache = SingleWriterShmObjectStorage( |
| 87 | max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes, |
| 88 | n_readers=self.world_size, |
| 89 | ring_buffer=ring_buffer, |
| 90 | serde_class=MsgpackSerde, |
| 91 | ) |
| 92 | # cache prompt_updates for P0 only |
| 93 | self._p0_cache: dict[str, Sequence[ResolvedPromptUpdate]] = {} |
| 94 | |
| 95 | self._hits = 0 |
| 96 | self._total = 0 |
| 97 | self._last_info = CacheInfo(hits=0, total=0) |
| 98 | |
| 99 | def _stat(self, *, delta: bool = False) -> CacheInfo: |
| 100 | info = CacheInfo(hits=self._hits, total=self._total) |
| 101 | |
| 102 | if delta: |
| 103 | info_delta = info - self._last_info |
| 104 | self._last_info = info |
| 105 | info = info_delta |
| 106 | |
| 107 | return info |
| 108 | |
| 109 | @override |
| 110 | def is_cached_item(self, mm_hash: str) -> bool: |
| 111 | return self._shm_cache.is_cached(mm_hash) |
| 112 | |
| 113 | @override |
| 114 | def get_and_update_item( |
| 115 | self, |
| 116 | mm_item: MultiModalProcessorCacheInItem, |
| 117 | mm_hash: str, |
| 118 | ) -> MultiModalProcessorCacheOutItem: |
| 119 | if self._shm_cache.is_cached(mm_hash): |
| 120 | self._hits += 1 |
| 121 | self._total += 1 |
| 122 | |
| 123 | address, monotonic_id = self._shm_cache.get_cached(mm_hash) |
| 124 | signature = self._shm_cache.get_signature(mm_hash) |
| 125 | prompt_updates = self._p0_cache[mm_hash] |
| 126 | return ( |
| 127 | self.address_as_item(address, monotonic_id, signature), |
| 128 | prompt_updates, |
| 129 | ) |
| 130 | |
| 131 | assert mm_item is not None, f"Expected a cached item for {mm_hash=}" |
| 132 | item, prompt_updates = mm_item |
| 133 | |
| 134 | self._total += 1 |
| 135 | |
| 136 | try: |
| 137 | address, monotonic_id = self._shm_cache.put(mm_hash, item) |
| 138 | signature = self._shm_cache.get_signature(mm_hash) |
| 139 | # Try to remove dangling items if p0 cache is too large. |
| 140 | if len(self._p0_cache) >= 2 * len(self._shm_cache.key_index): |
| 141 | self.remove_dangling_items() |
| 142 | |
| 143 | self._p0_cache[mm_hash] = prompt_updates |
| 144 | return ( |
| 145 | self.address_as_item(address, monotonic_id, signature), |
| 146 | prompt_updates, |
| 147 | ) |
| 148 | except ValueError as e: |
| 149 | # `put` raises ValueError either for an oversize item or for a |
| 150 | # duplicate key (concurrent insert); the latter is benign so we |
| 151 | # only warn on the oversize case. Subsequent UUID-only requests |
| 152 | # for an oversize item will fail with a cache miss. |
| 153 | if "already exists" not in str(e): |
| 154 | logger.warning_once( |
| 155 | "mm_input %s too large to cache; " |
| 156 | "raise --mm-shm-cache-max-object-size-mb. (%s)", |
| 157 | mm_hash, |
| 158 | str(e), |
| 159 | ) |
| 160 | return mm_item |
| 161 | except MemoryError as e: |
| 162 | # Cache full and protected items prevent eviction. |
| 163 | logger.debug( |
| 164 | "mm_input %s not cached; shm cache full, " |
| 165 | "consider raising --mm-processor-cache-gb. (%s)", |
| 166 | mm_hash, |
| 167 | str(e), |
| 168 | ) |
| 169 | return mm_item |
| 170 | |
| 171 | @override |
| 172 | def touch_sender_cache_item(self, mm_hash: str) -> None: |
| 173 | """Touch the item in shared memory cache to prevent eviction.""" |
| 174 | self._shm_cache.touch(mm_hash) |
| 175 | |
| 176 | @override |
| 177 | def release_sender_touches(self) -> None: |
| 178 | self._shm_cache.release_touches() |
| 179 | |
| 180 | @override |
| 181 | def validate_input_item( |
| 182 | self, |
| 183 | mm_item: MultiModalKwargsItem, |
| 184 | mm_hash: str, |
| 185 | ) -> None: |
| 186 | if (handle := _get_shm_handle(mm_item)) is not None: |
| 187 | address, monotonic_id, signature = handle |
| 188 | self._shm_cache.verify_signature( |
| 189 | mm_hash, |
| 190 | address, |
| 191 | monotonic_id, |
| 192 | signature, |
| 193 | ) |
| 194 | |
| 195 | @override |
| 196 | def clear_cache(self) -> None: |
| 197 | self._shm_cache.clear() |
| 198 | self._p0_cache.clear() |
| 199 | |
| 200 | self._hits = 0 |
| 201 | self._total = 0 |
| 202 | self._last_info = CacheInfo(hits=0, total=0) |
| 203 | |
| 204 | @override |
| 205 | def make_stats(self, *, delta: bool = False) -> CacheInfo: |
| 206 | return self._stat(delta=delta) |
| 207 | |
| 208 | @override |
| 209 | def close(self) -> None: |
| 210 | self._shm_cache.close() |
| 211 | |
| 212 | def remove_dangling_items(self) -> None: |
| 213 | """Remove items that are no longer in the shared memory cache.""" |
| 214 | cached_hashes = self._shm_cache.key_index.keys() |
| 215 | dangling_hashes = set(self._p0_cache.keys()) - cached_hashes |
| 216 | for mm_hash in dangling_hashes: |
| 217 | del self._p0_cache[mm_hash] |
| 218 | |
| 219 | def address_as_item( |
| 220 | self, |
| 221 | address: int, |
| 222 | monotonic_id: int, |
| 223 | signature: list[int], |
| 224 | ) -> MultiModalKwargsItem: |
| 225 | addr_elem = MultiModalFieldElem( |
| 226 | data=address, |
| 227 | field=MultiModalBatchedField(), |
| 228 | ) |
| 229 | id_elem = MultiModalFieldElem( |
| 230 | data=monotonic_id, |
| 231 | field=MultiModalBatchedField(), |
| 232 | ) |
| 233 | |
| 234 | signature_elem = MultiModalFieldElem( |
| 235 | data=signature, |
| 236 | field=MultiModalBatchedField(), |
| 237 | ) |
| 238 | |
| 239 | return MultiModalKwargsItem( |
| 240 | { |
| 241 | "address": addr_elem, |
| 242 | "monotonic_id": id_elem, |
| 243 | "signature": signature_elem, |
| 244 | } |
| 245 | ) |
| 246 | |
| 247 | |
| 248 | class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache): |
| 249 | """The cache which is used on P1 Worker Process when SHM caching is enabled. |
| 250 | |
| 251 | How to update each item: |
| 252 | |
| 253 | - If the item has an address, replace the input with the cached item. |
| 254 | - If not, return the input. |
| 255 | """ |
| 256 | |
| 257 | def __init__( |
| 258 | self, |
| 259 | vllm_config: VllmConfig, |
| 260 | shared_worker_lock: LockType, |
| 261 | ) -> None: |
| 262 | super().__init__() |
| 263 | |
| 264 | self.world_size = vllm_config.parallel_config.world_size |
| 265 | mm_config = vllm_config.model_config.get_multimodal_config() |
| 266 | |
| 267 | ring_buffer = SingleWriterShmRingBuffer( |
| 268 | data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes), |
| 269 | name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME, |
| 270 | create=False, # Server is a reader |
| 271 | ) |
| 272 | self._shm_cache = SingleWriterShmObjectStorage( |
| 273 | max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes, |
| 274 | n_readers=self.world_size, |
| 275 | ring_buffer=ring_buffer, |
| 276 | serde_class=MsgpackSerde, |
| 277 | reader_lock=shared_worker_lock, |
| 278 | ) |
| 279 | |
| 280 | @override |
| 281 | def get_and_update_features( |
| 282 | self, mm_features: list["MultiModalFeatureSpec"] |
| 283 | ) -> list["MultiModalFeatureSpec"]: |
| 284 | # strip_covered_mm_data preserves address items, so None here represents |
| 285 | # a stripped uncached payload and has no SHM reference to acknowledge. |
| 286 | features_with_data = [ |
| 287 | feature for feature in mm_features if feature.data is not None |
| 288 | ] |
| 289 | super().get_and_update_features(features_with_data) |
| 290 | return mm_features |
| 291 | |
| 292 | @override |
| 293 | def get_and_update_item( |
| 294 | self, |
| 295 | mm_item: MultiModalKwargsItem | None, |
| 296 | mm_hash: str, |
| 297 | ) -> MultiModalKwargsItem: |
| 298 | assert mm_item is not None, f"Expected an address item for {mm_hash=}" |
| 299 | if (handle := _get_shm_handle(mm_item)) is not None: |
| 300 | address, monotonic_id, signature = handle |
| 301 | return self._shm_cache.get(address, monotonic_id, signature, mm_hash) |
| 302 | |
| 303 | return mm_item |
| 304 | |
| 305 | @override |
| 306 | def touch_receiver_cache_item( |
| 307 | self, |
| 308 | mm_hash: str, |
| 309 | mm_item: MultiModalKwargsItem | None = None, |
| 310 | ) -> None: |
| 311 | """Validate the item's handle in shared memory cache.""" |
| 312 | assert mm_item is not None |
| 313 | if (handle := _get_shm_handle(mm_item)) is not None: |
| 314 | address, monotonic_id, signature = handle |
| 315 | self._shm_cache.touch( |
| 316 | mm_hash, |
| 317 | address=address, |
| 318 | monotonic_id=monotonic_id, |
| 319 | signature=signature, |
| 320 | ) |
| 321 | |
| 322 | @override |
| 323 | def clear_cache(self) -> None: |
| 324 | self._shm_cache.clear() |