mm_input_cache.py 5.79 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
from typing import Any, Optional
4
5

from vllm.config import ModelConfig
6
from vllm.envs import VLLM_MM_INPUT_CACHE_GIB
7
from vllm.logger import init_logger
8
9
from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalDataDict,
                             MultiModalKwargs, MultiModalRegistry)
10
from vllm.multimodal.processing import ProcessingCache
11
12
13

logger = init_logger(__name__)

14
15
16
# The idea of multimodal preprocessing caching is based on having a client and
# a server, where the client executes in the frontend process (=P0) and the
# server in the core process (=P1).
17
#
18
19
20
21
22
23
24
25
# -- Client:
#  - Apply legacy input_mapper (if one exists) to generate MultiModalKwargs.
#  - Perform caching of the generated MultiModalKwargs.
#  - This client can be deprecated once all mutimodal models migrate to use
#    merged preprocessor with built-in caching functionality.
#
# -- Server:
#  - Perform caching of the received MultiModalKwargs.
26
27
28
29
#
# The caching for both client and server is mirrored/similar, and this allows us
# to avoid the serialization of "mm_inputs" (like pixel values) between
# client (=P0) and server (=P1) processes.
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_GIB.
34

35

36
37
38
# TODO(ywang96): Deprecate this class once all multimodal models migrate to use
# merged preprocessor with built-in caching functionality.
class MMInputCacheClient:
39
40
41
42
43
44

    def __init__(
        self,
        model_config: ModelConfig,
        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
    ):
45
        self.model_config = model_config
46
47
48
49
50
        self.mm_registry = mm_registry
        self.multi_modal_input_mapper = mm_registry.create_input_mapper(
            model_config)
        self.mm_registry.init_mm_limits_per_prompt(model_config)

51
        # Init cache
52
        self.use_cache = not model_config.disable_mm_preprocessor_cache
53
54
        self.mm_cache = ProcessingCache.get_lru_cache(VLLM_MM_INPUT_CACHE_GIB,
                                                      MultiModalKwargs)
55
56
57

        # DEBUG: Set to None to disable
        self.mm_debug_cache_hit_ratio_steps = None
58
59
        self.mm_debug_cache_hits = 0
        self.mm_debug_cache_total = 0
60

61
    def cache_hit_ratio(self, steps):
62
63
64
        total = self.mm_debug_cache_total

        if total > 0 and total % steps == 0:
65
            logger.debug("MMInputMapper: cache_hit_ratio = %.2f ",
66
                         self.mm_debug_cache_hits / total)
67

68
69
    # NOTE: process_inputs only supports image inputs since all multimodal
    # models with other modalities have migrated to use merged preprocessor.
70
71
72
    def process_inputs(
        self,
        mm_data: MultiModalDataDict,
73
74
75
        mm_hashes: Optional[list[str]],
        mm_processor_kwargs: Optional[dict[str, Any]],
        precomputed_mm_inputs: Optional[list[MultiModalKwargs]],
76
    ) -> list[Optional[MultiModalKwargs]]:
77
78
79
80
81
82
83
84
        if precomputed_mm_inputs is None:
            image_inputs = mm_data["image"]
            if not isinstance(image_inputs, list):
                image_inputs = [image_inputs]
            num_inputs = len(image_inputs)
        else:
            num_inputs = len(precomputed_mm_inputs)

85
86
        # Sanity
        if self.use_cache:
87
            assert mm_hashes is not None
88
            assert num_inputs == len(mm_hashes)
89
90
91
92

        # Process each image input separately, so that later we can schedule
        # them in a fine-grained manner.
        # Apply caching (if enabled) and reuse precomputed inputs (if provided)
93
        ret_inputs: list[Optional[MultiModalKwargs]] = []
94
95
96
97
98
        for input_id in range(num_inputs):
            if self.mm_debug_cache_hit_ratio_steps is not None:
                self.cache_hit_ratio(self.mm_debug_cache_hit_ratio_steps)

            mm_input = None
99
            if self.use_cache:
100
                assert mm_hashes is not None
101
102
103
                mm_hash = mm_hashes[input_id]
                mm_input = self.mm_cache.get(mm_hash)

104
            self.mm_debug_cache_total += 1
105
106
107
108
109
            if mm_input is None:
                if precomputed_mm_inputs is not None:
                    # Reuse precomputed input (for merged preprocessor)
                    mm_input = precomputed_mm_inputs[input_id]
                else:
110
                    # Apply legacy input_mapper
111
112
113
114
115
                    mm_input = self.multi_modal_input_mapper(
                        {"image": [image_inputs[input_id]]},
                        mm_processor_kwargs=mm_processor_kwargs,
                    )

116
                if self.use_cache:
117
                    # Add to cache
118
                    assert mm_hash is not None
119
                    self.mm_cache[mm_hash] = mm_input
120
            else:
121
                self.mm_debug_cache_hits += 1
122
123
124
125
                mm_input = None  # Avoids sending mm_input to Server

            ret_inputs.append(mm_input)

126
        return ret_inputs
127
128


129
class MMInputCacheServer:
130

131
    def __init__(self, model_config):
132
        self.use_cache = not model_config.disable_mm_preprocessor_cache
133
134
        self.mm_cache = ProcessingCache.get_lru_cache(VLLM_MM_INPUT_CACHE_GIB,
                                                      MultiModalKwargs)
135

136
    def get_and_update(
137
        self,
138
139
        mm_inputs: list[Optional[MultiModalKwargs]],
        mm_hashes: list[str],
140
    ) -> list[Optional[MultiModalKwargs]]:
141
142
        assert len(mm_inputs) == len(mm_hashes)

143
144
145
        if not self.use_cache:
            return mm_inputs

146
147
        full_mm_inputs = []
        for mm_input, mm_hash in zip(mm_inputs, mm_hashes):
148
            assert mm_hash is not None
149
150
151
152
            if mm_input is None:
                mm_input = self.mm_cache.get(mm_hash)
                assert mm_input is not None
            else:
153
                self.mm_cache[mm_hash] = mm_input
154
155
156
157

            full_mm_inputs.append(mm_input)

        return full_mm_inputs