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, 17:15 UTC
[Misc] Avoid direct access of global `mm_registry` in `compute_encoder_budget` (#15621)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Mar 28, 2025 at 01:52 UTC
13ac9cab21e3cd12acd0c94376bc2f6da3dca5cd
4 files changed
+19
-7
vllm/v1/core/encoder_cache_manager.py
+12
-4
@@ -3,7 +3,7 @@
3
from typing import TYPE_CHECKING
4
5
from vllm.logger import init_logger
6
-from vllm.multimodal import MULTIMODAL_REGISTRY
6
+from vllm.multimodal import MultiModalRegistry
7
from vllm.v1.request import Request
8
9
if TYPE_CHECKING:
@@ -67,6 +67,7 @@ class EncoderCacheManager:
67
def compute_encoder_budget(
68
model_config: "ModelConfig",
69
scheduler_config: "SchedulerConfig",
70
+ mm_registry: MultiModalRegistry,
71
) -> tuple[int, int]:
72
"""Compute the encoder cache budget based on the model and scheduler
73
configurations.
@@ -74,6 +75,7 @@ def compute_encoder_budget(
75
Args:
76
model_config: Model configuration.
77
scheduler_config: Scheduler configuration.
78
+ mm_registry: Provides information about the token cost.
79
80
Returns:
81
- Compute budget for encoder execution, in unit of number of tokens
@@ -89,7 +91,11 @@ def compute_encoder_budget(
91
(
92
encoder_compute_budget,
93
encoder_cache_size,
92
- ) = _compute_encoder_budget_multimodal(model_config, scheduler_config)
94
+ ) = _compute_encoder_budget_multimodal(
95
+ model_config,
96
+ scheduler_config,
97
+ mm_registry,
98
+ )
99
100
return encoder_compute_budget, encoder_cache_size
101
@@ -97,6 +103,7 @@ def compute_encoder_budget(
103
def _compute_encoder_budget_multimodal(
104
model_config: "ModelConfig",
105
scheduler_config: "SchedulerConfig",
106
+ mm_registry: MultiModalRegistry,
107
) -> tuple[int, int]:
108
"""Compute the encoder cache budget based on the model and scheduler
109
configurations for a multimodal model.
@@ -104,6 +111,7 @@ def _compute_encoder_budget_multimodal(
111
Args:
112
model_config: Model configuration.
113
scheduler_config: Scheduler configuration.
114
+ mm_registry: Provides information about the token cost.
115
116
Returns:
117
- Compute budget for encoder execution, in unit of number of tokens
@@ -112,8 +120,8 @@ def _compute_encoder_budget_multimodal(
120
in the input sequence.
121
"""
122
115
- max_tokens_by_modality_dict = MULTIMODAL_REGISTRY.get_max_tokens_per_item_by_nonzero_modality( # noqa: E501
116
- model_config)
123
+ max_tokens_by_modality_dict = mm_registry \
124
+ .get_max_tokens_per_item_by_nonzero_modality(model_config)
125
126
if not max_tokens_by_modality_dict:
127
logger.warning(
vllm/v1/core/sched/scheduler.py
+3
@@ -10,6 +10,7 @@ from typing import Optional, Union
10
from vllm.config import (CacheConfig, LoRAConfig, ModelConfig, SchedulerConfig,
11
SpeculativeConfig)
12
from vllm.logger import init_logger
13
+from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
14
from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
15
compute_encoder_budget)
16
from vllm.v1.core.kv_cache_manager import KVCacheManager
@@ -38,6 +39,7 @@ class Scheduler(SchedulerInterface):
39
speculative_config: Optional[SpeculativeConfig],
40
log_stats: bool,
41
structured_output_manager: StructuredOutputManager,
42
+ mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
43
) -> None:
44
self.scheduler_config = scheduler_config
45
self.cache_config = cache_config
@@ -93,6 +95,7 @@ class Scheduler(SchedulerInterface):
95
encoder_compute_budget, encoder_cache_size = compute_encoder_budget(
96
model_config=model_config,
97
scheduler_config=scheduler_config,
98
+ mm_registry=mm_registry,
99
)
100
101
# NOTE(woosuk): Here, "encoder" includes the vision encoder (and
vllm/v1/worker/gpu_model_runner.py
+3
-3
@@ -137,6 +137,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
137
encoder_compute_budget, encoder_cache_size = compute_encoder_budget(
138
model_config=model_config,
139
scheduler_config=scheduler_config,
140
+ mm_registry=self.mm_registry,
141
)
142
self.max_num_encoder_input_tokens = encoder_compute_budget
143
self.encoder_cache_size = encoder_cache_size
@@ -1439,9 +1440,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1440
# NOTE: Currently model is profiled with a single non-text
1441
# modality with the max possible input tokens even when
1442
# it supports multiple.
1442
- max_tokens_by_modality_dict = (
1443
- MULTIMODAL_REGISTRY.
1444
- get_max_tokens_per_item_by_nonzero_modality(self.model_config))
1443
+ max_tokens_by_modality_dict = self.mm_registry \
1444
+ .get_max_tokens_per_item_by_nonzero_modality(self.model_config)
1445
dummy_data_modality, max_tokens_per_mm_item = max(
1446
max_tokens_by_modality_dict.items(), key=lambda item: item[1])
1447
vllm/v1/worker/tpu_model_runner.py
+1
@@ -109,6 +109,7 @@ class TPUModelRunner:
109
encoder_compute_budget, encoder_cache_size = compute_encoder_budget(
110
model_config=model_config,
111
scheduler_config=scheduler_config,
112
+ mm_registry=self.mm_registry,
113
)
114
self.max_num_encoder_input_tokens = encoder_compute_budget
115
self.encoder_cache_size = encoder_cache_size