A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
[VLM] Limit multimodal input cache by memory (#14805)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Mar 15, 2025 at 17:52 UTC
3556a414341033aad1bbb84674ec16b235324b25
13 files changed
+159
-55
.pre-commit-config.yaml
+1
-1
@@ -53,7 +53,7 @@ repos:
53
entry: tools/mypy.sh 0 "local"
54
language: python
55
types: [python]
56
- additional_dependencies: &mypy_deps [mypy==1.11.1, types-setuptools, types-PyYAML, types-requests]
56
+ additional_dependencies: &mypy_deps [mypy==1.11.1, types-cachetools, types-setuptools, types-PyYAML, types-requests]
57
stages: [pre-commit] # Don't run in CI
58
- id: mypy-3.9 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
59
name: Run mypy for Python 3.9
requirements/common.txt
+1
@@ -1,3 +1,4 @@
1
+cachetools
2
psutil
3
sentencepiece # Required for LLaMA tokenizer.
4
numpy < 2.0.0
requirements/docs.txt
+1
@@ -9,6 +9,7 @@ msgspec
9
cloudpickle
10
11
# packages to install to build the documentation
12
+cachetools
13
pydantic >= 2.8
14
-f https://download.pytorch.org/whl/cpu
15
torch
tests/models/multimodal/processing/test_common.py
+1
-1
@@ -48,7 +48,7 @@ def _test_processing_correctness(
48
tokenizer=cached_tokenizer_from_config(model_config),
49
)
50
# Ensure that it can fit all of the data
51
- cache = ProcessingCache(capacity=1 << 30)
51
+ cache = ProcessingCache(capacity_gb=2048)
52
53
processing_info = factories.info(ctx)
54
supported_mm_limits = processing_info.get_supported_mm_limits()
vllm/envs.py
+5
-6
@@ -56,7 +56,7 @@ if TYPE_CHECKING:
56
VLLM_IMAGE_FETCH_TIMEOUT: int = 5
57
VLLM_VIDEO_FETCH_TIMEOUT: int = 30
58
VLLM_AUDIO_FETCH_TIMEOUT: int = 10
59
- VLLM_MM_INPUT_CACHE_SIZE: int = 256
59
+ VLLM_MM_INPUT_CACHE_GIB: int = 8
60
VLLM_TARGET_DEVICE: str = "cuda"
61
MAX_JOBS: Optional[str] = None
62
NVCC_THREADS: Optional[str] = None
@@ -432,11 +432,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
432
"VLLM_AUDIO_FETCH_TIMEOUT":
433
lambda: int(os.getenv("VLLM_AUDIO_FETCH_TIMEOUT", "10")),
434
435
- # Cache size for multimodal feature/input cache for multimodal models
436
- # in unit of number of multimodal data items (e.g. image, video, audio).
437
- # Default is 256 multimodal data items.
438
- "VLLM_MM_INPUT_CACHE_SIZE":
439
- lambda: int(os.getenv("VLLM_MM_INPUT_CACHE_SIZE", "256")),
435
+ # Cache size (in GiB) for multimodal input cache
436
+ # Default is 8GiB
437
+ "VLLM_MM_INPUT_CACHE_GIB":
438
+ lambda: int(os.getenv("VLLM_MM_INPUT_CACHE_GIB", "8")),
439
440
# Path to the XLA persistent cache directory.
441
# Only used for XLA devices such as TPUs.
vllm/jsontree.py
new
+79
@@ -0,0 +1,79 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Helper functions to work with nested JSON structures."""
3
+from collections.abc import Iterable
4
+from functools import reduce
5
+from typing import Callable, TypeVar, Union, overload
6
+
7
+_T = TypeVar("_T")
8
+_U = TypeVar("_U")
9
+
10
+JSONTree = Union[dict[str, "JSONTree[_T]"], list["JSONTree[_T]"],
11
+ tuple["JSONTree[_T]", ...], _T]
12
+"""A nested JSON structure where the leaves need not be JSON-serializable."""
13
+
14
+
15
+def json_iter_leaves(value: JSONTree[_T]) -> Iterable[_T]:
16
+ """Iterate through each leaf in a nested JSON structure."""
17
+ if isinstance(value, dict):
18
+ for v in value.values():
19
+ yield from json_iter_leaves(v)
20
+ elif isinstance(value, (list, tuple)):
21
+ for v in value:
22
+ yield from json_iter_leaves(v)
23
+ else:
24
+ yield value
25
+
26
+
27
+def json_map_leaves(
28
+ func: Callable[[_T], _U],
29
+ value: JSONTree[_T],
30
+) -> JSONTree[_U]:
31
+ """Apply a function to each leaf in a nested JSON structure."""
32
+ if isinstance(value, dict):
33
+ return {k: json_map_leaves(func, v) for k, v in value.items()}
34
+ elif isinstance(value, list):
35
+ return [json_map_leaves(func, v) for v in value]
36
+ elif isinstance(value, tuple):
37
+ return tuple(json_map_leaves(func, v) for v in value)
38
+ else:
39
+ return func(value)
40
+
41
+
42
+@overload
43
+def json_reduce_leaves(
44
+ func: Callable[[_T, _T], _T],
45
+ value: JSONTree[_T],
46
+ /,
47
+) -> _T:
48
+ ...
49
+
50
+
51
+@overload
52
+def json_reduce_leaves(
53
+ func: Callable[[_U, _T], _U],
54
+ value: JSONTree[_T],
55
+ initial: _U,
56
+ /,
57
+) -> _U:
58
+ ...
59
+
60
+
61
+def json_reduce_leaves(
62
+ func: Callable[..., Union[_T, _U]],
63
+ value: JSONTree[_T],
64
+ initial: _U = ..., # type: ignore[assignment]
65
+ /,
66
+) -> Union[_T, _U]:
67
+ """
68
+ Apply a function of two arguments cumulatively to each leaf in a
69
+ nested JSON structure, from left to right, so as to reduce the
70
+ sequence to a single value.
71
+ """
72
+ if initial is ...:
73
+ return reduce(func, json_iter_leaves(value)) # type: ignore[arg-type]
74
+
75
+ return reduce(
76
+ func, # type: ignore[arg-type]
77
+ json_iter_leaves(value),
78
+ initial,
79
+ )
vllm/model_executor/models/llava.py
+2
-1
@@ -18,6 +18,7 @@ from transformers.models.pixtral import PixtralProcessor
18
19
from vllm.config import VllmConfig
20
from vllm.inputs import InputProcessingContext
21
+from vllm.jsontree import JSONTree, json_map_leaves
22
from vllm.model_executor.layers.activation import get_act_fn
23
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
24
RowParallelLinear)
@@ -35,7 +36,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
36
PromptReplacement, PromptUpdate)
37
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
38
from vllm.sequence import IntermediateTensors
38
-from vllm.utils import JSONTree, flatten_2d_lists, json_map_leaves
39
+from vllm.utils import flatten_2d_lists
40
41
from .clip import CLIPVisionModel
42
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
vllm/model_executor/models/molmo.py
+2
-1
@@ -24,6 +24,7 @@ from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
24
get_tensor_model_parallel_world_size,
25
split_tensor_along_last_dim,
26
tensor_model_parallel_all_gather)
27
+from vllm.jsontree import JSONTree, json_map_leaves
28
from vllm.model_executor import SamplingMetadata
29
from vllm.model_executor.layers.activation import (MulAndSilu, QuickGELU,
30
SiluAndMul)
@@ -50,7 +51,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
51
PromptInsertion, PromptUpdate)
52
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
53
from vllm.sequence import IntermediateTensors
53
-from vllm.utils import JSONTree, flatten_2d_lists, json_map_leaves
54
+from vllm.utils import flatten_2d_lists
55
56
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
57
SupportsMultiModal, SupportsPP, SupportsQuant)
vllm/multimodal/inputs.py
+2
-1
@@ -16,7 +16,8 @@ from PIL.Image import Image
16
from transformers import BatchFeature
17
from typing_extensions import NotRequired, TypeAlias
18
19
-from vllm.utils import JSONTree, full_groupby, is_list_of, json_map_leaves
19
+from vllm.jsontree import JSONTree, json_map_leaves
20
+from vllm.utils import full_groupby, is_list_of
21
22
if TYPE_CHECKING:
23
from .hasher import MultiModalHashDict
vllm/multimodal/processing.py
+43
-8
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-
2
import re
3
+import sys
4
from abc import ABC, abstractmethod
5
from collections import defaultdict
6
from collections.abc import (Callable, Generator, ItemsView, Iterable, Mapping,
@@ -11,14 +11,17 @@ from functools import lru_cache
11
from typing import (TYPE_CHECKING, Generic, NamedTuple, Optional, Protocol,
12
TypeVar, Union, cast)
13
14
+import torch
15
+from cachetools import LRUCache
16
from transformers import BatchFeature, PretrainedConfig, ProcessorMixin
17
from typing_extensions import assert_never
18
19
from vllm.inputs import InputProcessingContext
20
+from vllm.jsontree import json_map_leaves, json_reduce_leaves
21
from vllm.logger import init_logger
22
from vllm.transformers_utils.tokenizer import (AnyTokenizer, decode_tokens,
23
encode_tokens)
21
-from vllm.utils import LRUCache, flatten_2d_lists, full_groupby
24
+from vllm.utils import GiB_bytes, flatten_2d_lists, full_groupby
25
26
from .hasher import MultiModalHasher
27
from .inputs import (MultiModalDataDict, MultiModalEncDecInputs,
@@ -812,25 +815,50 @@ def find_mm_placeholders(
815
return dict(full_groupby_modality(it))
816
817
818
+_V = TypeVar("_V", bound="Union[MultiModalKwargs, MultiModalKwargsItem]")
819
+
820
+
821
class ProcessingCache:
822
817
- def __init__(self, capacity: int) -> None:
823
+ @staticmethod
824
+ def get_lru_cache(
825
+ capacity_gb: int,
826
+ value_type: type[_V],
827
+ ) -> LRUCache[str, _V]:
828
+
829
+ def get_size(leaf: object) -> int:
830
+ if isinstance(leaf, torch.Tensor):
831
+ return leaf.nbytes # sys.getsizeof doesn't work for tensors
832
+
833
+ return sys.getsizeof(leaf)
834
+
835
+ return LRUCache[str, _V](
836
+ GiB_bytes * capacity_gb,
837
+ getsizeof=lambda x: json_reduce_leaves(
838
+ lambda a, b: a + b,
839
+ json_map_leaves(get_size, x),
840
+ ),
841
+ )
842
+
843
+ def __init__(self, capacity_gb: int) -> None:
844
super().__init__()
845
846
# DEBUG: Set to None to disable
847
self.debug_cache_hit_ratio_steps: Optional[int] = None
848
+ self.debug_cache_hits = 0
849
+ self.debug_cache_total = 0
850
823
- self._cache = LRUCache[str, MultiModalKwargsItem](capacity)
851
+ self._cache = self.get_lru_cache(capacity_gb, MultiModalKwargsItem)
852
853
def _maybe_log_cache_stats(self) -> None:
854
steps = self.debug_cache_hit_ratio_steps
855
if not steps:
856
return
857
830
- cache_stats = self._cache.stat()
831
- if cache_stats.total % steps == 0:
858
+ total = self.debug_cache_total
859
+ if total > 0 and total % steps == 0:
860
logger.debug("ProcessingCache: hit_ratio = %.2f",
833
- cache_stats.hit_ratio)
861
+ self.debug_cache_hits / total)
862
863
def get(
864
self,
@@ -853,6 +881,13 @@ class ProcessingCache:
881
cache_key = MultiModalHasher.hash_kwargs(model_id=model_id,
882
**{modality: input_item},
883
**input_kwargs)
884
+
885
+ if self.debug_cache_hit_ratio_steps:
886
+ if cache_key in self._cache:
887
+ self.debug_cache_hits += 1
888
+
889
+ self.debug_cache_total += 1
890
+
891
return self._cache.get(cache_key)
892
893
def put(
@@ -870,7 +905,7 @@ class ProcessingCache:
905
cache_key = MultiModalHasher.hash_kwargs(model_id=model_id,
906
**{modality: input_item},
907
**input_kwargs)
873
- self._cache.put(cache_key, output_kwargs)
908
+ self._cache[cache_key] = output_kwargs
909
910
911
class BaseProcessingInfo:
vllm/multimodal/registry.py
+2
-2
@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Generic, Optional, Protocol, TypeVar
8
9
import torch.nn as nn
10
11
-from vllm.envs import VLLM_MM_INPUT_CACHE_SIZE
11
+from vllm.envs import VLLM_MM_INPUT_CACHE_GIB
12
from vllm.inputs import InputProcessingContext
13
from vllm.logger import init_logger
14
from vllm.transformers_utils.tokenizer import (AnyTokenizer,
@@ -119,7 +119,7 @@ class MultiModalRegistry:
119
120
self._limits_by_model = _MultiModalLimits()
121
122
- self._processing_cache = ProcessingCache(VLLM_MM_INPUT_CACHE_SIZE)
122
+ self._processing_cache = ProcessingCache(VLLM_MM_INPUT_CACHE_GIB)
123
124
def register_plugin(self, plugin: MultiModalPlugin) -> None:
125
"""
vllm/utils.py
-16
@@ -845,22 +845,6 @@ def is_list_of(
845
assert_never(check)
846
847
848
-JSONTree = Union[dict[str, "JSONTree[T]"], list["JSONTree[T]"],
849
- tuple["JSONTree[T]", ...], T]
850
-"""A nested JSON structure where the leaves need not be JSON-serializable."""
851
-
852
-
853
-def json_map_leaves(func: Callable[[T], U], value: JSONTree[T]) -> JSONTree[U]:
854
- if isinstance(value, dict):
855
- return {k: json_map_leaves(func, v) for k, v in value.items()}
856
- elif isinstance(value, list):
857
- return [json_map_leaves(func, v) for v in value]
858
- elif isinstance(value, tuple):
859
- return tuple(json_map_leaves(func, v) for v in value)
860
- else:
861
- return func(value)
862
-
863
-
848
def flatten_2d_lists(lists: list[list[T]]) -> list[T]:
849
"""Flatten a list of lists to a single list."""
850
return [item for sublist in lists for item in sublist]
vllm/v1/engine/mm_input_cache.py
+20
-18
@@ -3,11 +3,11 @@
3
from typing import Any, Optional
4
5
from vllm.config import ModelConfig
6
-from vllm.envs import VLLM_MM_INPUT_CACHE_SIZE
6
+from vllm.envs import VLLM_MM_INPUT_CACHE_GIB
7
from vllm.logger import init_logger
8
from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalDataDict,
9
MultiModalKwargs, MultiModalRegistry)
10
-from vllm.utils import LRUCache
10
+from vllm.multimodal.processing import ProcessingCache
11
12
logger = init_logger(__name__)
13
@@ -30,7 +30,7 @@ logger = init_logger(__name__)
30
31
# Both Client and Server must use the same cache size
32
# (to perform mirrored caching). This cache size is set by the environment
33
-# variable VLLM_MM_INPUT_CACHE_SIZE.
33
+# variable VLLM_MM_INPUT_CACHE_GIB.
34
35
36
# TODO(ywang96): Deprecate this class once all multimodal models migrate to use
@@ -50,18 +50,20 @@ class MMInputCacheClient:
50
51
# Init cache
52
self.use_cache = not model_config.disable_mm_preprocessor_cache
53
- self.mm_cache = LRUCache[str,
54
- MultiModalKwargs](VLLM_MM_INPUT_CACHE_SIZE)
53
+ self.mm_cache = ProcessingCache.get_lru_cache(VLLM_MM_INPUT_CACHE_GIB,
54
+ MultiModalKwargs)
55
56
# DEBUG: Set to None to disable
57
self.mm_debug_cache_hit_ratio_steps = None
58
- self.mm_cache_hits = 0
59
- self.mm_cache_total = 0
58
+ self.mm_debug_cache_hits = 0
59
+ self.mm_debug_cache_total = 0
60
61
def cache_hit_ratio(self, steps):
62
- if self.mm_cache_total > 0 and self.mm_cache_total % steps == 0:
62
+ total = self.mm_debug_cache_total
63
+
64
+ if total > 0 and total % steps == 0:
65
logger.debug("MMInputMapper: cache_hit_ratio = %.2f ",
64
- self.mm_cache_hits / self.mm_cache_total)
66
+ self.mm_debug_cache_hits / total)
67
68
# NOTE: process_inputs only supports image inputs since all multimodal
69
# models with other modalities have migrated to use merged preprocessor.
@@ -71,7 +73,7 @@ class MMInputCacheClient:
73
mm_hashes: Optional[list[str]],
74
mm_processor_kwargs: Optional[dict[str, Any]],
75
precomputed_mm_inputs: Optional[list[MultiModalKwargs]],
74
- ) -> list[MultiModalKwargs]:
76
+ ) -> list[Optional[MultiModalKwargs]]:
77
if precomputed_mm_inputs is None:
78
image_inputs = mm_data["image"]
79
if not isinstance(image_inputs, list):
@@ -88,7 +90,7 @@ class MMInputCacheClient:
90
# Process each image input separately, so that later we can schedule
91
# them in a fine-grained manner.
92
# Apply caching (if enabled) and reuse precomputed inputs (if provided)
91
- ret_inputs: list[MultiModalKwargs] = []
93
+ ret_inputs: list[Optional[MultiModalKwargs]] = []
94
for input_id in range(num_inputs):
95
if self.mm_debug_cache_hit_ratio_steps is not None:
96
self.cache_hit_ratio(self.mm_debug_cache_hit_ratio_steps)
@@ -99,7 +101,7 @@ class MMInputCacheClient:
101
mm_hash = mm_hashes[input_id]
102
mm_input = self.mm_cache.get(mm_hash)
103
102
- self.mm_cache_total += 1
104
+ self.mm_debug_cache_total += 1
105
if mm_input is None:
106
if precomputed_mm_inputs is not None:
107
# Reuse precomputed input (for merged preprocessor)
@@ -114,9 +116,9 @@ class MMInputCacheClient:
116
if self.use_cache:
117
# Add to cache
118
assert mm_hash is not None
117
- self.mm_cache.put(mm_hash, mm_input)
119
+ self.mm_cache[mm_hash] = mm_input
120
else:
119
- self.mm_cache_hits += 1
121
+ self.mm_debug_cache_hits += 1
122
mm_input = None # Avoids sending mm_input to Server
123
124
ret_inputs.append(mm_input)
@@ -128,14 +130,14 @@ class MMInputCacheServer:
130
131
def __init__(self, model_config):
132
self.use_cache = not model_config.disable_mm_preprocessor_cache
131
- self.mm_cache = LRUCache[str,
132
- MultiModalKwargs](VLLM_MM_INPUT_CACHE_SIZE)
133
+ self.mm_cache = ProcessingCache.get_lru_cache(VLLM_MM_INPUT_CACHE_GIB,
134
+ MultiModalKwargs)
135
136
def get_and_update(
137
self,
138
mm_inputs: list[Optional[MultiModalKwargs]],
139
mm_hashes: list[str],
138
- ) -> list[MultiModalKwargs]:
140
+ ) -> list[Optional[MultiModalKwargs]]:
141
assert len(mm_inputs) == len(mm_hashes)
142
143
if not self.use_cache:
@@ -148,7 +150,7 @@ class MMInputCacheServer:
150
mm_input = self.mm_cache.get(mm_hash)
151
assert mm_input is not None
152
else:
151
- self.mm_cache.put(mm_hash, mm_input)
153
+ self.mm_cache[mm_hash] = mm_input
154
155
full_mm_inputs.append(mm_input)
156