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, 14:45 UTC
[Bugfix] Avoid transferring cached multi-modal items from P0 to P1 (#16273)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Apr 9, 2025 at 15:51 UTC
e484e028575e670137f8267a56247a1eb04fb884
5 files changed
+65
-22
vllm/v1/engine/__init__.py
+2
-1
@@ -2,6 +2,7 @@
2
3
import enum
4
import time
5
+from collections.abc import Sequence
6
from typing import Any, Optional, Union
7
8
import msgspec
@@ -52,7 +53,7 @@ class EngineCoreRequest(
53
# Detokenizer, but set to None when it is added to EngineCoreClient.
54
prompt: Optional[str]
55
prompt_token_ids: list[int]
55
- mm_inputs: Optional[list[MultiModalKwargs]]
56
+ mm_inputs: Optional[Sequence[Optional[MultiModalKwargs]]]
57
mm_hashes: Optional[list[str]]
58
mm_placeholders: Optional[list[PlaceholderRange]]
59
sampling_params: SamplingParams
vllm/v1/engine/core.py
+3
-3
@@ -31,7 +31,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
31
from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler
32
from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
33
EngineCoreRequestType, UtilityOutput)
34
-from vllm.v1.engine.mm_input_cache import MMInputCacheServer
34
+from vllm.v1.engine.mm_input_cache import MirroredProcessingCache
35
from vllm.v1.executor.abstract import Executor
36
from vllm.v1.kv_cache_interface import KVCacheConfig
37
from vllm.v1.outputs import ModelRunnerOutput
@@ -105,7 +105,7 @@ class EngineCore:
105
)
106
107
# Setup MM Input Mapper.
108
- self.mm_input_cache_server = MMInputCacheServer(
108
+ self.mm_input_cache_server = MirroredProcessingCache(
109
vllm_config.model_config)
110
111
# Setup batch queue for pipeline parallelism.
@@ -173,7 +173,7 @@ class EngineCore:
173
# anything that has a hash must have a HIT cache entry here
174
# as well.
175
assert request.mm_inputs is not None
176
- request.mm_inputs = self.mm_input_cache_server.get_and_update(
176
+ request.mm_inputs = self.mm_input_cache_server.get_and_update_p1(
177
request.mm_inputs, request.mm_hashes)
178
179
req = Request.from_engine_core_request(request)
vllm/v1/engine/mm_input_cache.py
+34
-7
@@ -1,8 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+from collections.abc import Sequence
3
+from typing import Optional
4
5
from vllm.envs import VLLM_MM_INPUT_CACHE_GIB
6
from vllm.multimodal import MultiModalKwargs
7
from vllm.multimodal.processing import ProcessingCache
8
+from vllm.utils import is_list_of
9
10
# The idea of multimodal preprocessing caching is based on having a client and
11
# a server, where the client executes in the frontend process (=P0) and the
@@ -11,9 +14,11 @@ from vllm.multimodal.processing import ProcessingCache
14
# -- Client:
15
# - BaseMultiModalProcessor to process MultiModalData into MultiModalKwargs
16
# with built-in caching functionality, with mm_hash as its identifier.
17
+# - MirroredProcessingCache to keep track of the cached entries and
18
+# determine whether to send the MultiModalKwargs to P1.
19
#
20
# -- Server:
16
-# - MMInputCacheServer to perform caching of the received MultiModalKwargs.
21
+# - MirroredProcessingCache to store the MultiModalKwargs from P0.
22
#
23
# The caching for both client and server is mirrored, and this allows us
24
# to avoid the serialization of "mm_inputs" (like pixel values) between
@@ -25,26 +30,48 @@ from vllm.multimodal.processing import ProcessingCache
30
# variable VLLM_MM_INPUT_CACHE_GIB.
31
32
28
-class MMInputCacheServer:
33
+class MirroredProcessingCache:
34
35
def __init__(self, model_config):
36
self.use_cache = not model_config.disable_mm_preprocessor_cache
37
self.mm_cache = ProcessingCache.get_lru_cache(VLLM_MM_INPUT_CACHE_GIB,
38
MultiModalKwargs)
39
35
- def get_and_update(
40
+ def get_and_update_p0(
41
self,
37
- mm_inputs: list[MultiModalKwargs],
42
+ mm_inputs: Sequence[MultiModalKwargs],
43
mm_hashes: list[str],
39
- ) -> list[MultiModalKwargs]:
44
+ ) -> Sequence[Optional[MultiModalKwargs]]:
45
assert len(mm_inputs) == len(mm_hashes)
46
47
if not self.use_cache:
48
+ assert is_list_of(mm_inputs, MultiModalKwargs)
49
return mm_inputs
50
45
- full_mm_inputs = []
51
+ full_mm_inputs = list[Optional[MultiModalKwargs]]()
52
+ for mm_input, mm_hash in zip(mm_inputs, mm_hashes):
53
+ if mm_hash in self.mm_cache:
54
+ mm_input = None
55
+ else:
56
+ self.mm_cache[mm_hash] = mm_input
57
+
58
+ full_mm_inputs.append(mm_input)
59
+
60
+ return full_mm_inputs
61
+
62
+ def get_and_update_p1(
63
+ self,
64
+ mm_inputs: Sequence[Optional[MultiModalKwargs]],
65
+ mm_hashes: list[str],
66
+ ) -> Sequence[MultiModalKwargs]:
67
+ assert len(mm_inputs) == len(mm_hashes)
68
+
69
+ if not self.use_cache:
70
+ assert is_list_of(mm_inputs, MultiModalKwargs)
71
+ return mm_inputs
72
+
73
+ full_mm_inputs = list[MultiModalKwargs]()
74
for mm_input, mm_hash in zip(mm_inputs, mm_hashes):
47
- assert mm_hash is not None
75
if mm_input is None:
76
mm_input = self.mm_cache[mm_hash]
77
else:
vllm/v1/engine/processor.py
+17
-6
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import time
4
-from collections.abc import Mapping
4
+from collections.abc import Mapping, Sequence
5
from typing import Literal, Optional, Union
6
7
from vllm.config import VllmConfig
@@ -19,6 +19,7 @@ from vllm.prompt_adapter.request import PromptAdapterRequest
19
from vllm.sampling_params import SamplingParams
20
from vllm.transformers_utils.tokenizer_group import BaseTokenizerGroup
21
from vllm.v1.engine import EngineCoreRequest
22
+from vllm.v1.engine.mm_input_cache import MirroredProcessingCache
23
from vllm.v1.structured_output.backend_guidance import (
24
validate_guidance_grammar)
25
from vllm.v1.structured_output.utils import (
@@ -47,6 +48,8 @@ class Processor:
48
self.tokenizer,
49
mm_registry)
50
51
+ self.mm_input_cache_client = MirroredProcessingCache(self.model_config)
52
+
53
# Multi-modal hasher (for images)
54
self.use_hash = (
55
not self.model_config.disable_mm_preprocessor_cache) or \
@@ -231,7 +234,7 @@ class Processor:
234
self.tokenizer.get_lora_tokenizer(lora_request))
235
236
# Multimodal related.
234
- sorted_mm_inputs: Optional[list[MultiModalKwargs]] = None
237
+ sorted_mm_inputs: Optional[Sequence[Optional[MultiModalKwargs]]] = None
238
sorted_mm_positions: Optional[list[PlaceholderRange]] = None
239
sorted_mm_hashes: Optional[list[str]] = None
240
if decoder_inputs["type"] == "multimodal":
@@ -256,20 +259,28 @@ class Processor:
259
# are multiple modalities.
260
unique_modalities = set(sorted_item_modalities)
261
if len(unique_modalities) > 1:
259
- sorted_mm_inputs = []
262
+ orig_sorted_mm_inputs = []
263
used_indices = {modality: 0 for modality in unique_modalities}
264
+
265
for modality in sorted_item_modalities:
266
items = decoder_mm_inputs.get_items(modality)
267
item = items[used_indices[modality]]
264
- sorted_mm_inputs.append(MultiModalKwargs.from_items([item
265
- ]))
268
+
269
+ orig_sorted_mm_inputs.append(
270
+ MultiModalKwargs.from_items([item]))
271
used_indices[modality] += 1
272
else:
268
- sorted_mm_inputs = [
273
+ orig_sorted_mm_inputs = [
274
MultiModalKwargs.from_items([item]) for item in
275
decoder_mm_inputs.get_items(sorted_item_modalities[0])
276
]
277
278
+ if sorted_mm_hashes is not None:
279
+ sorted_mm_inputs = self.mm_input_cache_client.get_and_update_p0(
280
+ orig_sorted_mm_inputs, sorted_mm_hashes)
281
+ else:
282
+ sorted_mm_inputs = orig_sorted_mm_inputs
283
+
284
return EngineCoreRequest(
285
request_id=request_id,
286
prompt=decoder_inputs.get("prompt"),
vllm/v1/request.py
+9
-5
@@ -3,17 +3,16 @@
3
import enum
4
from typing import TYPE_CHECKING, Optional, Union
5
6
+from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
7
from vllm.sampling_params import SamplingParams
8
+from vllm.utils import is_list_of
9
from vllm.v1.engine import (EngineCoreEvent, EngineCoreEventType,
10
EngineCoreRequest, FinishReason)
11
from vllm.v1.structured_output.request import StructuredOutputRequest
12
from vllm.v1.utils import ConstantList
13
14
if TYPE_CHECKING:
13
-
15
from vllm.lora.request import LoRARequest
15
- from vllm.multimodal import MultiModalKwargs
16
- from vllm.multimodal.inputs import PlaceholderRange
16
17
18
class Request:
@@ -23,9 +22,9 @@ class Request:
22
request_id: str,
23
prompt: Optional[str],
24
prompt_token_ids: list[int],
26
- multi_modal_inputs: Optional[list["MultiModalKwargs"]],
25
+ multi_modal_inputs: Optional[list[MultiModalKwargs]],
26
multi_modal_hashes: Optional[list[str]],
28
- multi_modal_placeholders: Optional[list["PlaceholderRange"]],
27
+ multi_modal_placeholders: Optional[list[PlaceholderRange]],
28
sampling_params: SamplingParams,
29
eos_token_id: Optional[int],
30
arrival_time: float,
@@ -75,6 +74,11 @@ class Request:
74
75
@classmethod
76
def from_engine_core_request(cls, request: EngineCoreRequest) -> "Request":
77
+ if request.mm_inputs is not None:
78
+ assert isinstance(request.mm_inputs, list)
79
+ assert is_list_of(request.mm_inputs, MultiModalKwargs), (
80
+ "mm_inputs was not updated in EngineCore.add_request")
81
+
82
return cls(
83
request_id=request.request_id,
84
prompt=request.prompt,