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 10, 04:00 UTC
[2/N] Move cache factories to MM registry (#32382)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Jan 15, 2026 at 17:02 UTC
cbbae38f9368b6c35d9b9295bf4ceee1e6452750
10 files changed
+124
-154
tests/multimodal/test_cache.py
+4
-6
@@ -19,8 +19,6 @@ from vllm.multimodal.cache import (
19
MultiModalReceiverCache,
20
ShmObjectStoreReceiverCache,
21
ShmObjectStoreSenderCache,
22
- engine_receiver_cache_from_config,
23
- processor_cache_from_config,
22
)
23
from vllm.multimodal.hasher import MultiModalHasher
24
from vllm.multimodal.inputs import (
@@ -132,10 +130,10 @@ def _compare_caches(
130
n_iter: int = 100,
131
seed: int = 0,
132
):
135
- cache_0_p0 = processor_cache_from_config(config_0, MULTIMODAL_REGISTRY)
136
- cache_0_p1 = engine_receiver_cache_from_config(config_0, MULTIMODAL_REGISTRY)
137
- cache_1_p0 = processor_cache_from_config(config_1, MULTIMODAL_REGISTRY)
138
- cache_1_p1 = engine_receiver_cache_from_config(config_1, MULTIMODAL_REGISTRY)
133
+ cache_0_p0 = MULTIMODAL_REGISTRY.processor_cache_from_config(config_0)
134
+ cache_0_p1 = MULTIMODAL_REGISTRY.engine_receiver_cache_from_config(config_0)
135
+ cache_1_p0 = MULTIMODAL_REGISTRY.processor_cache_from_config(config_1)
136
+ cache_1_p1 = MULTIMODAL_REGISTRY.engine_receiver_cache_from_config(config_1)
137
138
cache_size_gb = max(
139
config_0.model_config.multimodal_config.mm_processor_cache_gb,
tests/v1/engine/test_process_multi_modal_uuids.py
+3
-4
@@ -6,9 +6,8 @@ import pytest
6
from vllm.assets.image import ImageAsset
7
from vllm.assets.video import VideoAsset
8
from vllm.config import CacheConfig, DeviceConfig, ModelConfig, VllmConfig
9
-from vllm.multimodal import MultiModalUUIDDict
9
+from vllm.multimodal import MultiModalRegistry, MultiModalUUIDDict
10
from vllm.sampling_params import SamplingParams
11
-from vllm.v1.engine import input_processor as input_processor_mod
11
from vllm.v1.engine.input_processor import InputProcessor
12
13
cherry_pil_image = ImageAsset("cherry_blossom").pil_image
@@ -36,9 +35,9 @@ def _mock_input_processor(
35
raising=True,
36
)
37
monkeypatch.setattr(
39
- input_processor_mod,
38
+ MultiModalRegistry,
39
"processor_cache_from_config",
41
- lambda vllm_config, mm_registry: None,
40
+ lambda self, vllm_config: None,
41
raising=True,
42
)
43
vllm/lora/model_manager.py
+10
-10
@@ -135,9 +135,15 @@ class LoRAModelManager:
135
llm_punica_wrapper
136
)
137
138
- def _maybe_init_mm(self, vllm_config: VllmConfig, max_num_batched_tokens) -> None:
139
- self.supports_tower_connector_lora = False
138
+ def _maybe_init_mm(
139
+ self,
140
+ vllm_config: VllmConfig,
141
+ max_num_batched_tokens: int,
142
+ ) -> None:
143
model_config: ModelConfig = vllm_config.model_config
144
+ mm_registry = MULTIMODAL_REGISTRY
145
+
146
+ self.supports_tower_connector_lora = False
147
self.mm_mapping: MultiModelKeys = self.model.get_mm_mapping()
148
149
# Only one language model can be included in the model.
@@ -154,9 +160,7 @@ class LoRAModelManager:
160
self.punica_wrapper_mapping[lm_prefix] = llm_punica_wrapper
161
162
if self.lora_config.enable_tower_connector_lora:
157
- self.mm_processor_info = MULTIMODAL_REGISTRY.create_processor(
158
- model_config
159
- ).info
163
+ self.mm_processor_info = mm_registry.create_processor(model_config).info
164
self.supports_tower_connector_lora = self.supports_mm and hasattr(
165
self.model, "get_num_mm_encoder_tokens"
166
)
@@ -169,11 +173,7 @@ class LoRAModelManager:
173
"GitHub if you encounter them."
174
)
175
172
- mm_budget = MultiModalBudget(
173
- model_config,
174
- vllm_config.scheduler_config,
175
- MULTIMODAL_REGISTRY,
176
- )
176
+ mm_budget = MultiModalBudget(vllm_config, mm_registry)
177
limit_per_prompt: int = max(
178
self.mm_processor_info.get_allowed_mm_limits().values()
179
)
vllm/multimodal/cache.py
-109
@@ -35,7 +35,6 @@ if TYPE_CHECKING:
35
from vllm.config import ModelConfig, VllmConfig
36
37
from .processing.processor import ResolvedPromptUpdate
38
- from .registry import MultiModalRegistry
38
39
logger = init_logger(__name__)
40
@@ -561,67 +560,6 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
560
return mm_item
561
562
564
-def _enable_processor_cache(
565
- model_config: "ModelConfig",
566
- mm_registry: "MultiModalRegistry",
567
-) -> bool:
568
- if not mm_registry.supports_multimodal_inputs(model_config):
569
- return False
570
-
571
- mm_config = model_config.get_multimodal_config()
572
- return mm_config.mm_processor_cache_gb > 0
573
-
574
-
575
-def _enable_ipc_cache(vllm_config: "VllmConfig") -> bool:
576
- parallel_config = vllm_config.parallel_config
577
- supports_ipc_cache = (
578
- parallel_config._api_process_count == 1
579
- and parallel_config.data_parallel_size == 1
580
- ) or parallel_config.data_parallel_external_lb
581
-
582
- return supports_ipc_cache
583
-
584
-
585
-def _enable_mm_input_shm_cache(vllm_config: "VllmConfig") -> bool:
586
- """Whether the shared memory based cache should be enabled."""
587
-
588
- if not _enable_ipc_cache(vllm_config):
589
- return False
590
-
591
- mm_config = vllm_config.model_config.get_multimodal_config()
592
-
593
- return mm_config.mm_processor_cache_type == "shm"
594
-
595
-
596
-def processor_cache_from_config(
597
- vllm_config: "VllmConfig",
598
- mm_registry: "MultiModalRegistry",
599
-) -> BaseMultiModalProcessorCache | None:
600
- """Return a `BaseMultiModalProcessorCache`, if enabled."""
601
- model_config = vllm_config.model_config
602
-
603
- if not _enable_processor_cache(model_config, mm_registry):
604
- return None
605
-
606
- if not _enable_ipc_cache(vllm_config):
607
- return MultiModalProcessorOnlyCache(model_config)
608
-
609
- if not _enable_mm_input_shm_cache(vllm_config):
610
- return MultiModalProcessorSenderCache(model_config)
611
- return ShmObjectStoreSenderCache(vllm_config)
612
-
613
-
614
-def processor_only_cache_from_config(
615
- model_config: "ModelConfig",
616
- mm_registry: "MultiModalRegistry",
617
-):
618
- """Return a `MultiModalProcessorOnlyCache`, if enabled."""
619
- if not _enable_processor_cache(model_config, mm_registry):
620
- return None
621
-
622
- return MultiModalProcessorOnlyCache(model_config)
623
-
624
-
563
class BaseMultiModalReceiverCache(
564
BaseMultiModalCache[MultiModalKwargsItem | None, MultiModalKwargsItem]
565
):
@@ -780,50 +718,3 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):
718
@override
719
def clear_cache(self) -> None:
720
self._shm_cache.clear()
783
-
784
-
785
-def engine_receiver_cache_from_config(
786
- vllm_config: "VllmConfig",
787
- mm_registry: "MultiModalRegistry",
788
-) -> BaseMultiModalReceiverCache | None:
789
- """
790
- This is used in the engine process.
791
- Return a `BaseMultiModalReceiverCache` only when IPC caching is enabled and
792
- mm_processor_cache_type=="lru".
793
- """
794
- model_config = vllm_config.model_config
795
-
796
- if not _enable_processor_cache(model_config, mm_registry):
797
- return None
798
-
799
- if not _enable_ipc_cache(vllm_config):
800
- return None
801
-
802
- if not _enable_mm_input_shm_cache(vllm_config):
803
- return MultiModalReceiverCache(model_config)
804
-
805
- return None
806
-
807
-
808
-def worker_receiver_cache_from_config(
809
- vllm_config: "VllmConfig",
810
- mm_registry: "MultiModalRegistry",
811
- shared_worker_lock: LockType,
812
-) -> BaseMultiModalReceiverCache | None:
813
- """
814
- This is used in the worker process.
815
- Return a `BaseMultiModalReceiverCache` only when IPC caching is enabled and
816
- mm_processor_cache_type=="shm".
817
- """
818
- model_config = vllm_config.model_config
819
-
820
- if not _enable_processor_cache(model_config, mm_registry):
821
- return None
822
-
823
- if not _enable_ipc_cache(vllm_config):
824
- return None
825
-
826
- if not _enable_mm_input_shm_cache(vllm_config):
827
- return None
828
-
829
- return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock)
vllm/multimodal/registry.py
+93
-3
@@ -2,14 +2,23 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Mapping
4
from dataclasses import dataclass
5
-from typing import TYPE_CHECKING, Generic, Protocol, TypeVar, cast
5
+from multiprocessing.synchronize import Lock as LockType
6
+from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeVar, cast
7
8
from vllm.config.multimodal import BaseDummyOptions
9
from vllm.config.observability import ObservabilityConfig
10
from vllm.logger import init_logger
11
from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
12
12
-from .cache import BaseMultiModalProcessorCache
13
+from .cache import (
14
+ BaseMultiModalProcessorCache,
15
+ BaseMultiModalReceiverCache,
16
+ MultiModalProcessorOnlyCache,
17
+ MultiModalProcessorSenderCache,
18
+ MultiModalReceiverCache,
19
+ ShmObjectStoreReceiverCache,
20
+ ShmObjectStoreSenderCache,
21
+)
22
from .inputs import MultiModalInputs
23
from .processing import (
24
BaseDummyInputsBuilder,
@@ -19,7 +28,7 @@ from .processing import (
28
)
29
30
if TYPE_CHECKING:
22
- from vllm.config import ModelConfig, ObservabilityConfig
31
+ from vllm.config import ModelConfig, ObservabilityConfig, VllmConfig
32
from vllm.model_executor.models.interfaces import SupportsMultiModal
33
34
logger = init_logger(__name__)
@@ -355,3 +364,84 @@ class MultiModalRegistry:
364
365
first_modality = next(iter(max_tokens))
366
return max_tokens[first_modality]
367
+
368
+ def _get_cache_type(
369
+ self,
370
+ vllm_config: "VllmConfig",
371
+ ) -> Literal[None, "processor_only", "lru", "shm"]:
372
+ model_config = vllm_config.model_config
373
+ if not self.supports_multimodal_inputs(model_config):
374
+ return None
375
+
376
+ # Check if the cache is disabled.
377
+ mm_config = model_config.get_multimodal_config()
378
+ if mm_config.mm_processor_cache_gb <= 0:
379
+ return None
380
+
381
+ # Check if IPC caching is supported.
382
+ parallel_config = vllm_config.parallel_config
383
+ is_ipc_supported = parallel_config._api_process_count == 1 and (
384
+ parallel_config.data_parallel_size == 1
385
+ or parallel_config.data_parallel_external_lb
386
+ )
387
+
388
+ if not is_ipc_supported:
389
+ return "processor_only"
390
+
391
+ mm_config = model_config.get_multimodal_config()
392
+ return mm_config.mm_processor_cache_type
393
+
394
+ def processor_cache_from_config(
395
+ self,
396
+ vllm_config: "VllmConfig",
397
+ ) -> BaseMultiModalProcessorCache | None:
398
+ """Return a `BaseMultiModalProcessorCache`, if enabled."""
399
+ cache_type = self._get_cache_type(vllm_config)
400
+ if cache_type is None:
401
+ return None
402
+ elif cache_type == "processor_only":
403
+ return MultiModalProcessorOnlyCache(vllm_config.model_config)
404
+ elif cache_type == "lru":
405
+ return MultiModalProcessorSenderCache(vllm_config.model_config)
406
+ elif cache_type == "shm":
407
+ return ShmObjectStoreSenderCache(vllm_config)
408
+ else:
409
+ raise ValueError(f"Unknown cache type: {cache_type!r}")
410
+
411
+ def processor_only_cache_from_config(
412
+ self,
413
+ vllm_config: "VllmConfig",
414
+ ) -> MultiModalProcessorOnlyCache | None:
415
+ """Return a `MultiModalProcessorOnlyCache`, if enabled."""
416
+ cache_type = self._get_cache_type(vllm_config)
417
+ if cache_type is None:
418
+ return None
419
+
420
+ return MultiModalProcessorOnlyCache(vllm_config.model_config)
421
+
422
+ def engine_receiver_cache_from_config(
423
+ self,
424
+ vllm_config: "VllmConfig",
425
+ ) -> BaseMultiModalReceiverCache | None:
426
+ """Return a `BaseMultiModalReceiverCache` for the engine process."""
427
+ cache_type = self._get_cache_type(vllm_config)
428
+ if cache_type in (None, "processor_only", "shm"):
429
+ return None
430
+ elif cache_type == "lru":
431
+ return MultiModalReceiverCache(vllm_config.model_config)
432
+ else:
433
+ raise ValueError(f"Unknown cache type: {cache_type!r}")
434
+
435
+ def worker_receiver_cache_from_config(
436
+ self,
437
+ vllm_config: "VllmConfig",
438
+ shared_worker_lock: LockType,
439
+ ) -> BaseMultiModalReceiverCache | None:
440
+ """Return a `BaseMultiModalReceiverCache` for the worker process."""
441
+ cache_type = self._get_cache_type(vllm_config)
442
+ if cache_type in (None, "processor_only", "lru"):
443
+ return None
444
+ elif cache_type == "shm":
445
+ return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock)
446
+ else:
447
+ raise ValueError(f"Unknown cache type: {cache_type!r}")
vllm/v1/engine/core.py
+2
-3
@@ -23,7 +23,6 @@ from vllm.logger import init_logger
23
from vllm.logging_utils.dump_input import dump_engine_exception
24
from vllm.lora.request import LoRARequest
25
from vllm.multimodal import MULTIMODAL_REGISTRY
26
-from vllm.multimodal.cache import engine_receiver_cache_from_config
26
from vllm.tasks import POOLING_TASKS, SupportedTask
27
from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
28
from vllm.utils.gc_utils import (
@@ -149,8 +148,8 @@ class EngineCore:
148
self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore
149
150
self.mm_registry = mm_registry = MULTIMODAL_REGISTRY
152
- self.mm_receiver_cache = engine_receiver_cache_from_config(
153
- vllm_config, mm_registry
151
+ self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config(
152
+ vllm_config
153
)
154
155
# If a KV connector is initialized for scheduler, we want to collect
vllm/v1/engine/input_processor.py
+1
-2
@@ -14,7 +14,6 @@ from vllm.inputs.preprocess import InputPreprocessor
14
from vllm.logger import init_logger
15
from vllm.lora.request import LoRARequest
16
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
17
-from vllm.multimodal.cache import processor_cache_from_config
17
from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict
18
from vllm.multimodal.parse import MultiModalDataParser
19
from vllm.multimodal.processing.context import set_request_id
@@ -58,7 +57,7 @@ class InputProcessor:
57
self.generation_config_fields = self.model_config.try_get_generation_config()
58
59
self.mm_registry = mm_registry
61
- self.mm_processor_cache = processor_cache_from_config(vllm_config, mm_registry)
60
+ self.mm_processor_cache = mm_registry.processor_cache_from_config(vllm_config)
61
62
self.input_preprocessor = InputPreprocessor(
63
self.model_config,
vllm/v1/worker/gpu_model_runner.py
+1
-5
@@ -623,11 +623,7 @@ class GPUModelRunner(
623
self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)
624
625
self.mm_budget = (
626
- MultiModalBudget(
627
- self.model_config,
628
- self.scheduler_config,
629
- self.mm_registry,
630
- )
626
+ MultiModalBudget(self.vllm_config, self.mm_registry)
627
if self.supports_mm_inputs
628
else None
629
)
vllm/v1/worker/utils.py
+5
-7
@@ -8,11 +8,10 @@ import torch
8
from typing_extensions import deprecated
9
10
from vllm.attention.layer import Attention
11
-from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
11
+from vllm.config import CacheConfig, VllmConfig
12
from vllm.logger import init_logger
13
from vllm.model_executor.models.interfaces import MultiModalEmbeddings
14
from vllm.model_executor.models.utils import extract_layer_index
15
-from vllm.multimodal.cache import processor_only_cache_from_config
15
from vllm.multimodal.registry import MultiModalRegistry
16
from vllm.platforms import current_platform
17
from vllm.utils.mem_utils import MemorySnapshot, format_gib
@@ -28,16 +27,15 @@ class MultiModalBudget:
27
28
def __init__(
29
self,
31
- model_config: ModelConfig,
32
- scheduler_config: SchedulerConfig,
30
+ vllm_config: VllmConfig,
31
mm_registry: MultiModalRegistry,
32
) -> None:
33
super().__init__()
34
37
- self.model_config = model_config
38
- self.scheduler_config = scheduler_config
35
+ self.model_config = model_config = vllm_config.model_config
36
+ self.scheduler_config = scheduler_config = vllm_config.scheduler_config
37
self.mm_registry = mm_registry
40
- self.cache = cache = processor_only_cache_from_config(model_config, mm_registry)
38
+ self.cache = cache = mm_registry.processor_only_cache_from_config(vllm_config)
39
40
self.max_model_len = model_config.max_model_len
41
self.max_num_reqs = scheduler_config.max_num_seqs
vllm/v1/worker/worker_base.py
+5
-5
@@ -12,7 +12,6 @@ from vllm.config import VllmConfig, set_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.lora.request import LoRARequest
14
from vllm.multimodal import MULTIMODAL_REGISTRY
15
-from vllm.multimodal.cache import worker_receiver_cache_from_config
15
from vllm.utils.import_utils import resolve_obj_by_qualname
16
from vllm.utils.system_utils import update_environment_variables
17
from vllm.v1.kv_cache_interface import KVCacheSpec
@@ -303,10 +302,11 @@ class WorkerWrapperBase:
302
303
self.mm_receiver_cache = None
304
else:
306
- self.mm_receiver_cache = worker_receiver_cache_from_config(
307
- vllm_config,
308
- MULTIMODAL_REGISTRY,
309
- shared_worker_lock,
305
+ self.mm_receiver_cache = (
306
+ MULTIMODAL_REGISTRY.worker_receiver_cache_from_config(
307
+ vllm_config,
308
+ shared_worker_lock,
309
+ )
310
)
311
312
with set_current_vllm_config(self.vllm_config):