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 14, 05:30 UTC
[Refactor] Simplify code for MM budget (#23310)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Aug 21, 2025 at 16:00 UTC
0c6e40bbaa4707528286a1e7bf17c90c88a1d920
4 files changed
+57
-68
vllm/v1/core/encoder_cache_manager.py
+32
-24
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
+from collections.abc import Mapping
4
from typing import TYPE_CHECKING
5
6
from vllm.logger import init_logger
@@ -188,35 +188,47 @@ def compute_encoder_budget(
188
- Space budget for encoder cache size, in unit of number of tokens
189
in the input sequence.
190
"""
191
+ if mm_registry.supports_multimodal_inputs(model_config):
192
+ max_tokens_by_modality = mm_registry \
193
+ .get_max_tokens_per_item_by_nonzero_modality(model_config)
194
192
- if not mm_registry.supports_multimodal_inputs(model_config):
193
- return 0, 0
195
+ return compute_mm_encoder_budget(
196
+ scheduler_config,
197
+ max_tokens_by_modality,
198
+ )
199
195
- # TODO: handle encoder-decoder models once we support them.
196
- (
197
- encoder_compute_budget,
198
- encoder_cache_size,
199
- ) = _compute_encoder_budget_multimodal(
200
- model_config,
201
- scheduler_config,
202
- mm_registry,
203
- )
200
+ return compute_text_encoder_budget(scheduler_config)
201
205
- return encoder_compute_budget, encoder_cache_size
202
203
+def compute_text_encoder_budget(
204
+ scheduler_config: "SchedulerConfig") -> tuple[int, int]:
205
+ """Compute the encoder cache budget based on the model and scheduler
206
+ configurations for a text-only model.
207
208
-def _compute_encoder_budget_multimodal(
209
- model_config: "ModelConfig",
208
+ Args:
209
+ scheduler_config: Scheduler configuration.
210
+
211
+ Returns:
212
+ - Compute budget for encoder execution, in unit of number of tokens
213
+ in the input sequence.
214
+ - Space budget for encoder cache size, in unit of number of tokens
215
+ in the input sequence.
216
+ """
217
+ # Currently text-only encoder-decoder models are not supported
218
+ return 0, 0
219
+
220
+
221
+def compute_mm_encoder_budget(
222
scheduler_config: "SchedulerConfig",
211
- mm_registry: MultiModalRegistry,
223
+ max_tokens_by_modality: Mapping[str, int],
224
) -> tuple[int, int]:
225
"""Compute the encoder cache budget based on the model and scheduler
226
configurations for a multimodal model.
227
228
Args:
217
- model_config: Model configuration.
229
scheduler_config: Scheduler configuration.
219
- mm_registry: Provides information about the token cost.
230
+ max_tokens_by_modality: The maximum number of tokens for each
231
+ non-text modality.
232
233
Returns:
234
- Compute budget for encoder execution, in unit of number of tokens
@@ -225,18 +237,14 @@ def _compute_encoder_budget_multimodal(
237
in the input sequence.
238
"""
239
228
- max_tokens_by_modality_dict = mm_registry \
229
- .get_max_tokens_per_item_by_nonzero_modality(model_config)
230
-
231
- if not max_tokens_by_modality_dict:
240
+ if not max_tokens_by_modality:
241
logger.warning(
242
"All non-text modalities supported by the model have been "
243
"explicitly disabled via limit_mm_per_prompt. Encoder cache will "
244
"not be initialized.")
245
return 0, 0
246
238
- _, max_tokens_per_mm_item = max(max_tokens_by_modality_dict.items(),
239
- key=lambda item: item[1])
247
+ max_tokens_per_mm_item = max(max_tokens_by_modality.values())
248
249
if (scheduler_config.disable_chunked_mm_input and max_tokens_per_mm_item
250
> scheduler_config.max_num_batched_tokens):
vllm/v1/worker/gpu_model_runner.py
+5
-13
@@ -341,10 +341,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
341
self.model_config,
342
self.scheduler_config,
343
self.mm_registry,
344
- max_model_len=self.max_model_len,
345
- max_num_reqs=self.max_num_reqs,
346
- ) if self.supports_mm_inputs \
347
- else None)
344
+ ) if self.supports_mm_inputs else None)
345
346
self.reorder_batch_threshold: Optional[int] = None
347
@@ -669,7 +666,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
666
mm_budget = self.mm_budget
667
assert mm_budget is not None
668
672
- dummy_modality, _ = mm_budget.get_modality_with_max_tokens()
669
+ dummy_modality = mm_budget.get_modality_with_max_tokens()
670
671
return self._get_mm_dummy_batch(dummy_modality, num_seqs)
672
@@ -2595,14 +2592,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2592
# NOTE: Currently model is profiled with a single non-text
2593
# modality with the max possible input tokens even when
2594
# it supports multiple.
2598
- (
2599
- dummy_modality,
2600
- max_tokens,
2601
- ) = mm_budget.get_modality_with_max_tokens()
2602
- (
2603
- max_mm_items_per_prompt,
2604
- max_mm_items_per_batch,
2605
- ) = mm_budget.get_max_items(dummy_modality, max_tokens)
2595
+ dummy_modality = mm_budget.get_modality_with_max_tokens()
2596
+ max_mm_items_per_batch = mm_budget \
2597
+ .max_items_per_batch_by_modality[dummy_modality]
2598
2599
logger.info(
2600
"Encoder cache will be initialized with a budget of "
vllm/v1/worker/tpu_model_runner.py
+3
-10
@@ -292,8 +292,6 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
292
self.model_config,
293
self.scheduler_config,
294
self.mm_registry,
295
- max_model_len=self.max_model_len,
296
- max_num_reqs=self.max_num_reqs,
295
) if self.supports_mm_inputs else None)
296
297
if not self.use_spmd:
@@ -1545,14 +1543,9 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1543
# NOTE: Currently model is profiled with a single non-text
1544
# modality with the max possible input tokens even when
1545
# it supports multiple.
1548
- (
1549
- dummy_modality,
1550
- max_tokens,
1551
- ) = mm_budget.get_modality_with_max_tokens()
1552
- (
1553
- max_mm_items_per_prompt,
1554
- max_mm_items_per_batch,
1555
- ) = mm_budget.get_max_items(dummy_modality, max_tokens)
1546
+ dummy_modality = mm_budget.get_modality_with_max_tokens()
1547
+ max_mm_items_per_batch = mm_budget \
1548
+ .max_items_per_batch_by_modality[dummy_modality]
1549
1550
logger.info(
1551
"Encoder cache will be initialized with a budget of "
vllm/v1/worker/utils.py
+17
-21
@@ -12,7 +12,7 @@ from vllm.model_executor.models.interfaces import MultiModalEmbeddings
12
from vllm.model_executor.models.utils import extract_layer_index
13
from vllm.multimodal.registry import MultiModalRegistry
14
from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
15
-from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
15
+from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
16
from vllm.v1.kv_cache_interface import KVCacheGroupSpec
17
18
if TYPE_CHECKING:
@@ -27,9 +27,6 @@ class MultiModalBudget:
27
model_config: ModelConfig,
28
scheduler_config: SchedulerConfig,
29
mm_registry: MultiModalRegistry,
30
- *,
31
- max_model_len: int,
32
- max_num_reqs: int,
30
) -> None:
31
super().__init__()
32
@@ -37,25 +34,25 @@ class MultiModalBudget:
34
self.scheduler_config = scheduler_config
35
self.mm_registry = mm_registry
36
40
- encoder_compute_budget, encoder_cache_size = compute_encoder_budget(
41
- model_config=model_config,
42
- scheduler_config=scheduler_config,
43
- mm_registry=mm_registry,
37
+ self.max_model_len = model_config.max_model_len
38
+ self.max_num_reqs = scheduler_config.max_num_seqs
39
+
40
+ self.mm_limits = mm_registry.get_mm_limits_per_prompt(model_config)
41
+
42
+ max_tokens_by_modality = mm_registry \
43
+ .get_max_tokens_per_item_by_nonzero_modality(model_config)
44
+
45
+ encoder_compute_budget, encoder_cache_size = compute_mm_encoder_budget(
46
+ scheduler_config,
47
+ max_tokens_by_modality,
48
)
49
46
- self.max_num_encoder_input_tokens = encoder_compute_budget
50
+ self.encoder_compute_budget = encoder_compute_budget
51
self.encoder_cache_size = encoder_cache_size
48
- self.max_model_len = max_model_len
49
- self.max_num_reqs = max_num_reqs
50
-
51
- self.mm_limits = mm_registry.get_mm_limits_per_prompt(model_config)
52
53
max_items_per_prompt_by_modality = dict[str, int]()
54
max_items_per_batch_by_modality = dict[str, int]()
55
56
- max_tokens_by_modality = mm_registry \
57
- .get_max_tokens_per_item_by_nonzero_modality(model_config)
58
-
56
for modality, max_tokens in max_tokens_by_modality.items():
57
(
58
max_items_per_prompt,
@@ -69,15 +66,14 @@ class MultiModalBudget:
66
self.max_items_per_prompt_by_modality = max_items_per_prompt_by_modality
67
self.max_items_per_batch_by_modality = max_items_per_batch_by_modality
68
72
- def get_modality_with_max_tokens(self) -> tuple[str, int]:
69
+ def get_modality_with_max_tokens(self) -> str:
70
max_tokens_by_modality = self.max_tokens_by_modality
74
- modality, max_tokens = max(max_tokens_by_modality.items(),
75
- key=lambda item: item[1])
71
+ modality, _ = max(max_tokens_by_modality.items(), key=lambda x: x[1])
72
77
- return modality, max_tokens
73
+ return modality
74
75
def get_encoder_budget(self) -> int:
80
- return min(self.max_num_encoder_input_tokens, self.encoder_cache_size)
76
+ return min(self.encoder_compute_budget, self.encoder_cache_size)
77
78
def get_max_items(
79
self,