@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e484e02857
Fork 0

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,