| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | from vllm.multimodal.inputs import MultiModalFeatureSpec |
| 6 | |
| 7 | |
| 8 | class EncoderCache: |
| 9 | def __init__(self): |
| 10 | # req_id -> MM features |
| 11 | self.mm_features: dict[str, list[MultiModalFeatureSpec]] = {} |
| 12 | # MM hash -> encoder outputs |
| 13 | self.encoder_outputs: dict[str, torch.Tensor] = {} |
| 14 | |
| 15 | def __len__(self) -> int: |
| 16 | return len(self.encoder_outputs) |
| 17 | |
| 18 | def add_request( |
| 19 | self, req_id: str, mm_features: list[MultiModalFeatureSpec] |
| 20 | ) -> None: |
| 21 | self.mm_features[req_id] = mm_features |
| 22 | |
| 23 | def remove_request(self, req_id: str) -> None: |
| 24 | self.mm_features.pop(req_id, None) |
| 25 | |
| 26 | def reset_mm_cache(self) -> None: |
| 27 | """Clear the multi-modal cache that was used during profiling, |
| 28 | but no longer needed during inference. |
| 29 | """ |
| 30 | # NOTE: v2 encoder cache profiling skips the multi-modal cache |
| 31 | pass |
| 32 | |
| 33 | def reset_encoder_cache(self) -> None: |
| 34 | """Clear the GPU-side encoder cache storing vision embeddings. |
| 35 | |
| 36 | This should be called when model weights are updated to ensure |
| 37 | stale embeddings computed with old weights are not reused. |
| 38 | """ |
| 39 | self.encoder_outputs.clear() |
| 40 | |
| 41 | def free_encoder_cache(self, mm_hash: str) -> None: |
| 42 | self.encoder_outputs.pop(mm_hash, None) |