| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections import OrderedDict |
| 5 | from collections.abc import Mapping |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | from vllm.config import VllmConfig |
| 9 | from vllm.config.ec_manager_config import EncoderCacheManagerMetadata |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.v1.request import Request |
| 12 | |
| 13 | if TYPE_CHECKING: |
| 14 | from vllm.config import SchedulerConfig |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | class EncoderCacheManager: |
| 20 | """Manages caching of encoder outputs for multimodal models in vLLM V1. |
| 21 | |
| 22 | The EncoderCacheManager handles the lifecycle of multimodal encoder outputs |
| 23 | (such as vision embeddings from images) during request processing. It |
| 24 | provides memory-aware caching to avoid recomputing encoder outputs when the |
| 25 | same multimodal inputs appear in different stages of request processing. |
| 26 | |
| 27 | This manager is particularly important for: |
| 28 | - Vision-language models (e.g., LLaVA) where image encoder outputs are |
| 29 | cached |
| 30 | - Any multimodal model where encoder computation is expensive and |
| 31 | cacheable |
| 32 | |
| 33 | The cache operates at the granularity of individual multimodal input items |
| 34 | within requests, allowing for fine-grained memory management and enabling |
| 35 | chunked processing of multimodal inputs. |
| 36 | |
| 37 | Cache is enabled to share embeddings of same multimodal data |
| 38 | item (identified by their hash value) between different requests, |
| 39 | and eviction takes place at allocation time when there's no free |
| 40 | space for new embeddings. |
| 41 | Oldest cached embeddings with no request referenced will be first evicted. |
| 42 | |
| 43 | NOTE: The EncoderCacheManager operates on the level of multimodal embeddings |
| 44 | instead of encoder tokens (i.e. all tokens that represent the multimodal data |
| 45 | in the input sequence). This means all break/text tokens in-between multimodal |
| 46 | embeddings are not considered with respect to the cache size and the number |
| 47 | of free slots. |
| 48 | |
| 49 | Args: |
| 50 | cache_size: Limit the size of the cache, measured by the number of |
| 51 | encoder embeddings from the input sequence. |
| 52 | |
| 53 | Attributes: |
| 54 | cache_size: Total cache capacity in encoder embeddings. |
| 55 | num_free_slots: Current available cache capacity in encoder embeddings. |
| 56 | num_freeable_slots: Capacity that can be immediately reclaimed by |
| 57 | evicting entries with zero references (in encoder embeddings). |
| 58 | cached: Mapping from mm_hash to a set of request IDs that currently |
| 59 | reference the cached entry. If the set is empty, the entry exists |
| 60 | but is not referenced by any request and is eligible for |
| 61 | reclamation. |
| 62 | num_embeds_by_hash: Mapping from mm_hash to the number of encoder |
| 63 | embeddings stored for that cache entry. |
| 64 | freeable: List of tuples (mm_hash, num_encoder_embeds) representing entries |
| 65 | whose no current running request is needed and that can be freed to |
| 66 | make space when needed. |
| 67 | freed: List of mm_hash strings that were actually evicted since the |
| 68 | last call to get_freed_mm_hashes(). This list is cleared on return. |
| 69 | |
| 70 | """ |
| 71 | |
| 72 | @classmethod |
| 73 | def create_manager( |
| 74 | cls, *, cache_size: int, vllm_config: "VllmConfig" |
| 75 | ) -> "EncoderCacheManager": |
| 76 | return cls(cache_size=cache_size) |
| 77 | |
| 78 | def __init__(self, cache_size: int): |
| 79 | self.cache_size = cache_size |
| 80 | self.num_free_slots = cache_size |
| 81 | self.num_freeable_slots = cache_size |
| 82 | |
| 83 | # mm_hash of mm_data => ids of requests that reference the mm_data |
| 84 | self.cached: dict[str, set[str]] = {} |
| 85 | # mm_hash of mm_data => number of cached encoder embeddings |
| 86 | self.num_embeds_by_hash: dict[str, int] = {} |
| 87 | # request_id => set of input_ids cached for that request |
| 88 | self.request_cached_ids: dict[str, set[int]] = {} |
| 89 | |
| 90 | # mm_hash of mm_data => num_encoder_embeds of the mm_data |
| 91 | self.freeable: OrderedDict[str, int] = OrderedDict() |
| 92 | self.freed: list[str] = [] |
| 93 | |
| 94 | def reset(self) -> None: |
| 95 | """Reset the encoder cache to its initial state. |
| 96 | |
| 97 | This clears all cached encoder outputs and resets capacity tracking. |
| 98 | Called when model weights are updated to invalidate stale embeddings. |
| 99 | """ |
| 100 | self.cached.clear() |
| 101 | self.num_embeds_by_hash.clear() |
| 102 | self.request_cached_ids.clear() |
| 103 | self.freeable.clear() |
| 104 | self.freed.clear() |
| 105 | self.num_free_slots = self.cache_size |
| 106 | self.num_freeable_slots = self.cache_size |
| 107 | |
| 108 | def check_and_update_cache(self, request: Request, input_id: int) -> bool: |
| 109 | """Check if encoder output for a specific multimodal input is cached. |
| 110 | |
| 111 | If the encoder output is cached, update `cached` to add the request id |
| 112 | to the set of request ids that reference the cached encoder output. |
| 113 | If the encoder output was previously not referenced by any request, |
| 114 | update `freeable` and `num_freeable_slots` accordingly. |
| 115 | |
| 116 | Args: |
| 117 | request: The request containing the multimodal input |
| 118 | input_id: Index of the multimodal input within the request |
| 119 | |
| 120 | Returns: |
| 121 | True if the encoder output for this input is already cached |
| 122 | |
| 123 | """ |
| 124 | mm_hash = request.mm_features[input_id].identifier |
| 125 | # Not cached at all |
| 126 | if mm_hash not in self.cached: |
| 127 | return False |
| 128 | |
| 129 | # Cached but currently not referenced by any request |
| 130 | if not self.cached[mm_hash]: |
| 131 | num_encoder_embeds = self.freeable.pop(mm_hash) |
| 132 | self.num_freeable_slots -= num_encoder_embeds |
| 133 | |
| 134 | self.cached[mm_hash].add(request.request_id) |
| 135 | self.request_cached_ids.setdefault(request.request_id, set()).add(input_id) |
| 136 | return True |
| 137 | |
| 138 | def get_cached_num_encoder_embeds( |
| 139 | self, request: Request, input_id: int |
| 140 | ) -> int | None: |
| 141 | """Return the encoder-embedding count stored for a cached input.""" |
| 142 | mm_hash = request.mm_features[input_id].identifier |
| 143 | return self.num_embeds_by_hash.get(mm_hash) |
| 144 | |
| 145 | def can_allocate( |
| 146 | self, |
| 147 | request: Request, |
| 148 | input_id: int, |
| 149 | encoder_compute_budget: int, |
| 150 | num_embeds_to_schedule: int, |
| 151 | ) -> bool: |
| 152 | """Check if there's sufficient cache space for a multimodal input. |
| 153 | If there is, return True and update EncoderCacheManager state. |
| 154 | |
| 155 | If there is not enough free space in `num_free_slots` but there is |
| 156 | enough reclaimable space in `num_freeable_slots`, entries will be |
| 157 | evicted from `freeable` (their mm_hash appended to `freed`) until |
| 158 | enough space is available, and then this method returns True. |
| 159 | Older entries are evicted first. |
| 160 | |
| 161 | Returns False only if the requested number of tokens exceeds both |
| 162 | the free and reclaimable capacities combined. |
| 163 | |
| 164 | Args: |
| 165 | request: The request containing the multimodal input. |
| 166 | input_id: Index of the multimodal input within the request. |
| 167 | encoder_compute_budget: Number of encoder embeddings allowed to be |
| 168 | computed when this method is invoked. |
| 169 | num_embeds_to_schedule: Number of encoder embeddings already scheduled to be |
| 170 | allocated with cache space when this method is invoked. |
| 171 | |
| 172 | Returns: |
| 173 | True if there's enough capacity to hold the encoder output for this |
| 174 | input (possibly after reclaiming `freeable` entries); otherwise |
| 175 | False. |
| 176 | |
| 177 | Note: This method does not allocate physical memory for the encoder |
| 178 | output but only the state of EncoderCacheManager. |
| 179 | |
| 180 | """ |
| 181 | num_embeds = request.get_num_encoder_embeds(input_id) |
| 182 | |
| 183 | # Not enough compute budget |
| 184 | if num_embeds > encoder_compute_budget: |
| 185 | return False |
| 186 | |
| 187 | num_embeds += num_embeds_to_schedule |
| 188 | |
| 189 | # Enough free slots |
| 190 | if num_embeds <= self.num_free_slots: |
| 191 | return True |
| 192 | |
| 193 | # Not enough reclaimable slots |
| 194 | if num_embeds > self.num_freeable_slots: |
| 195 | return False |
| 196 | |
| 197 | # Not enough free slots but enough reclaimable slots |
| 198 | # NOTE: Eviction takes place here, but physical memory is not freed |
| 199 | # until model runner is notified by the scheduler output. |
| 200 | while num_embeds > self.num_free_slots: |
| 201 | mm_hash, num_free_embeds = self.freeable.popitem(last=False) |
| 202 | del self.cached[mm_hash] |
| 203 | del self.num_embeds_by_hash[mm_hash] |
| 204 | self.freed.append(mm_hash) |
| 205 | self.num_free_slots += num_free_embeds |
| 206 | return True |
| 207 | |
| 208 | def allocate(self, request: Request, input_id: int) -> None: |
| 209 | """Allocate cache space for a multimodal input's encoder output. |
| 210 | |
| 211 | This reserves cache space for storing the encoder output of the |
| 212 | specified multimodal input. The actual encoder output storage happens in |
| 213 | the model runner; this method updates the manager's bookkeeping. |
| 214 | |
| 215 | Note: |
| 216 | This method assumes can_allocate() returned True for the same input. |
| 217 | |
| 218 | """ |
| 219 | mm_hash = request.mm_features[input_id].identifier |
| 220 | request_id = request.request_id |
| 221 | if mm_hash not in self.cached: |
| 222 | self.cached[mm_hash] = set() |
| 223 | |
| 224 | num_encoder_embeds = request.get_num_encoder_embeds(input_id) |
| 225 | cached_num_encoder_embeds = self.num_embeds_by_hash.setdefault( |
| 226 | mm_hash, num_encoder_embeds |
| 227 | ) |
| 228 | # _try_schedule_encoder_inputs rejects local-hit and same-step |
| 229 | # identifier/count mismatches before allocation. |
| 230 | assert cached_num_encoder_embeds == num_encoder_embeds |
| 231 | |
| 232 | # NOTE: Encoder cache should always have enough space for encoder inputs |
| 233 | # that are scheduled since eviction takes place at can_allocate(). |
| 234 | assert self.num_free_slots >= num_encoder_embeds |
| 235 | assert self.num_freeable_slots >= num_encoder_embeds |
| 236 | |
| 237 | self.cached[mm_hash].add(request_id) |
| 238 | self.request_cached_ids.setdefault(request_id, set()).add(input_id) |
| 239 | self.num_free_slots -= num_encoder_embeds |
| 240 | self.num_freeable_slots -= num_encoder_embeds |
| 241 | |
| 242 | def get_cached_input_ids(self, request: Request) -> set[int]: |
| 243 | """Get all cached multimodal input IDs for a request.""" |
| 244 | return self.request_cached_ids.get(request.request_id, set()) |
| 245 | |
| 246 | def free_encoder_input(self, request: Request, input_id: int) -> None: |
| 247 | """Free the request's reference to the encoder input (`mm_data`). |
| 248 | |
| 249 | When the reference set for the corresponding `mm_hash` becomes empty, |
| 250 | the entry is appended to `freeable` and `num_freeable_slots` is |
| 251 | increased by the number of encoder embeddings for that input. |
| 252 | |
| 253 | The entry is NOT physically freed until capacity is needed (e.g., by |
| 254 | `can_allocate`). |
| 255 | """ |
| 256 | req_id = request.request_id |
| 257 | mm_hash = request.mm_features[input_id].identifier |
| 258 | # Always clean up request_cached_ids, even if the mm_hash was |
| 259 | # already evicted from cache (e.g. by can_allocate). |
| 260 | if req_id in self.request_cached_ids: |
| 261 | self.request_cached_ids[req_id].discard(input_id) |
| 262 | if not self.request_cached_ids[req_id]: |
| 263 | del self.request_cached_ids[req_id] |
| 264 | # The mm_hash not in cache or the req_id set is empty |
| 265 | if not self.cached.get(mm_hash, None): |
| 266 | return |
| 267 | # `cached` counts referencing requests, not positions, so one request |
| 268 | # that repeats an item (an image carried across conversation turns) has |
| 269 | # a single reference covering every occurrence. Hold it until the last |
| 270 | # occurrence is freed: dropping it at the first makes the entry |
| 271 | # evictable while the request still needs it, and the encoder then |
| 272 | # recomputes an item it already has. |
| 273 | if any( |
| 274 | request.mm_features[other_id].identifier == mm_hash |
| 275 | for other_id in self.request_cached_ids.get(req_id, ()) |
| 276 | ): |
| 277 | return |
| 278 | self.cached[mm_hash].discard(req_id) |
| 279 | if not self.cached[mm_hash]: |
| 280 | num_encoder_embeds = self.num_embeds_by_hash[mm_hash] |
| 281 | self.freeable[mm_hash] = num_encoder_embeds |
| 282 | self.num_freeable_slots += num_encoder_embeds |
| 283 | |
| 284 | def free(self, request: Request) -> None: |
| 285 | """Free all encoder input cache reference held by *request*. |
| 286 | |
| 287 | For each cached input ID, `free_encoder_input` is invoked. |
| 288 | The data stays in memory until eviction is triggered by a future |
| 289 | attempt allocation called by 'can_allocate'. |
| 290 | |
| 291 | Typically called when a request is finished, cancelled, or aborted. |
| 292 | """ |
| 293 | for input_id in list(self.get_cached_input_ids(request)): |
| 294 | self.free_encoder_input(request, input_id) |
| 295 | |
| 296 | def get_freed_mm_hashes(self) -> list[str]: |
| 297 | """Get and clear the list of recently freed encoder cache entries. |
| 298 | |
| 299 | Returns: |
| 300 | List of mm_hash strings that were actually evicted since the last |
| 301 | call to be used by the scheduler to notify workers about which |
| 302 | encoder outputs can be removed from their caches. The internal |
| 303 | list is cleared after this call. |
| 304 | |
| 305 | """ |
| 306 | # An entry evicted early in the scheduling pass can be allocated again |
| 307 | # later in the same pass. Keep its worker-side tensor in that case. |
| 308 | freed = [mm_hash for mm_hash in self.freed if mm_hash not in self.cached] |
| 309 | self.freed = [] |
| 310 | return freed |
| 311 | |
| 312 | def get_manager_metadata(self) -> EncoderCacheManagerMetadata | None: |
| 313 | return None |
| 314 | |
| 315 | |
| 316 | def compute_mm_encoder_budget( |
| 317 | scheduler_config: "SchedulerConfig", |
| 318 | mm_max_toks_per_item: Mapping[str, int], |
| 319 | ) -> tuple[int, int]: |
| 320 | """Compute the encoder cache budget based on the model and scheduler |
| 321 | configurations for a multimodal model. |
| 322 | |
| 323 | Args: |
| 324 | scheduler_config: Scheduler configuration. |
| 325 | mm_max_toks_per_item: The maximum number of tokens per item for each |
| 326 | non-text modality. |
| 327 | |
| 328 | Returns: |
| 329 | - Compute budget for encoder execution, measured in number of tokens |
| 330 | from the input sequence. |
| 331 | - Space budget for encoder cache size, measured in number of tokens |
| 332 | from the input sequence. |
| 333 | |
| 334 | """ |
| 335 | if not mm_max_toks_per_item: |
| 336 | logger.warning( |
| 337 | "All non-text modalities supported by the model have been " |
| 338 | "explicitly disabled via limit_mm_per_prompt. Encoder cache will " |
| 339 | "not be initialized." |
| 340 | ) |
| 341 | return 0, 0 |
| 342 | |
| 343 | max_tokens_per_mm_item = max(mm_max_toks_per_item.values()) |
| 344 | |
| 345 | if ( |
| 346 | scheduler_config.disable_chunked_mm_input |
| 347 | and max_tokens_per_mm_item > scheduler_config.max_num_batched_tokens |
| 348 | ): |
| 349 | raise ValueError( |
| 350 | "Chunked MM input disabled but max_tokens_per_mm_item " |
| 351 | f"({max_tokens_per_mm_item}) is larger than max_num_batched_tokens" |
| 352 | f" ({scheduler_config.max_num_batched_tokens}). Please increase " |
| 353 | "max_num_batched_tokens." |
| 354 | ) |
| 355 | |
| 356 | encoder_compute_budget = max( |
| 357 | scheduler_config.max_num_encoder_input_tokens, max_tokens_per_mm_item |
| 358 | ) |
| 359 | encoder_cache_size = max( |
| 360 | scheduler_config.encoder_cache_size, max_tokens_per_mm_item |
| 361 | ) |
| 362 | |
| 363 | return encoder_compute_budget, encoder_cache_size |
| 364 | |
| 365 | |
| 366 | # NOTE (NickLucche): Temporary implementation for encoder-decoder models that only |
| 367 | # use the manager for scheduling purposes. Encoder-decoder models will eventually |
| 368 | # utilize the cache and this class will fold into EncoderCacheManager, as |
| 369 | # differences with MM models shrink. |
| 370 | class EncoderDecoderCacheManager(EncoderCacheManager): |
| 371 | def __init__(self, cache_size: int): |
| 372 | self.cache_size = cache_size |
| 373 | self.num_free_slots = cache_size |
| 374 | self.allocated: list[str] = [] |
| 375 | self.to_free: list[str] = [] |
| 376 | |
| 377 | def reset(self) -> None: |
| 378 | """Reset the encoder cache to its initial state.""" |
| 379 | self.num_free_slots = self.cache_size |
| 380 | self.allocated.clear() |
| 381 | self.to_free.clear() |
| 382 | |
| 383 | def check_and_update_cache(self, request: Request, input_id: int) -> bool: |
| 384 | return False |
| 385 | |
| 386 | def can_allocate( |
| 387 | self, |
| 388 | request: Request, |
| 389 | input_id: int, |
| 390 | encoder_compute_budget: int, |
| 391 | num_embeds_to_schedule: int, |
| 392 | ) -> bool: |
| 393 | num_encoder_embeds = request.get_num_encoder_embeds(input_id) |
| 394 | # Not enough compute budget |
| 395 | if num_encoder_embeds > encoder_compute_budget: |
| 396 | return False |
| 397 | |
| 398 | num_encoder_embeds += num_embeds_to_schedule |
| 399 | # Enough free slots |
| 400 | return num_encoder_embeds <= self.num_free_slots |
| 401 | |
| 402 | def allocate(self, request: Request, input_id: int) -> None: |
| 403 | num_encoder_embeds = request.get_num_encoder_embeds(input_id) |
| 404 | self.num_free_slots -= num_encoder_embeds |
| 405 | |
| 406 | mm_hash = request.mm_features[input_id].identifier |
| 407 | self.allocated.append(mm_hash) |
| 408 | |
| 409 | def free(self, request: Request) -> None: |
| 410 | for input_id in range(len(request.mm_features)): |
| 411 | self.free_encoder_input(request, input_id) |
| 412 | |
| 413 | def get_cached_input_ids(self, request: Request) -> set[int]: |
| 414 | return set(range(len(request.mm_features))) |
| 415 | |
| 416 | def get_freed_mm_hashes(self) -> list[str]: |
| 417 | # As encoder cache is not used for enc-dec models, we can free the entries here |
| 418 | # The actual free happens in the runner, *before* the model is executed. |
| 419 | # Therefore, `freeable` acts as a buffer to free the entries only after the |
| 420 | # model is executed, mimicking the state transition of `EncoderCacheManager`. |
| 421 | to_free = self.to_free |
| 422 | self.to_free = self.allocated |
| 423 | self.allocated = [] |
| 424 | return to_free |
| 425 | |
| 426 | def free_encoder_input(self, request: Request, input_id: int) -> None: |
| 427 | num_encoder_embeds = request.get_num_encoder_embeds(input_id) |
| 428 | self.num_free_slots += num_encoder_embeds |