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 9, 02:30 UTC
[Refactor] Move profiling methods to MM budget (#33559)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Feb 2, 2026 at 23:27 UTC
d7e17aaacd5ed1b4b4be6bcfef3a1b7cbc84fc9a
8 files changed
+200
-295
vllm/lora/model_manager.py
+3
-5
@@ -10,7 +10,7 @@ import torch
10
from torch import nn
11
12
from vllm.config import VllmConfig
13
-from vllm.config.lora import LoRAConfig, ModelConfig
13
+from vllm.config.lora import LoRAConfig
14
from vllm.logger import init_logger
15
from vllm.lora.layers import (
16
BaseLayerWithLoRA,
@@ -35,9 +35,9 @@ from vllm.model_executor.models.interfaces import is_pooling_model
35
from vllm.model_executor.models.module_mapping import MultiModelKeys
36
from vllm.model_executor.models.utils import PPMissingLayer
37
from vllm.multimodal import MULTIMODAL_REGISTRY
38
+from vllm.multimodal.budget import MultiModalBudget
39
from vllm.utils.cache import LRUCache
40
from vllm.utils.platform_utils import is_pin_memory_available
40
-from vllm.v1.worker.utils import MultiModalBudget
41
42
logger = init_logger(__name__)
43
@@ -142,7 +142,6 @@ class LoRAModelManager:
142
vllm_config: VllmConfig,
143
max_num_batched_tokens: int,
144
) -> None:
145
- model_config: ModelConfig = vllm_config.model_config
145
mm_registry = MULTIMODAL_REGISTRY
146
147
self.supports_tower_connector_lora = False
@@ -162,7 +161,6 @@ class LoRAModelManager:
161
self.punica_wrapper_mapping[lm_prefix] = llm_punica_wrapper
162
163
if self.lora_config.enable_tower_connector_lora:
165
- 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
)
@@ -176,7 +174,7 @@ class LoRAModelManager:
174
)
175
176
mm_budget = MultiModalBudget(vllm_config, mm_registry)
179
- limit_per_prompt = max(self.mm_processor_info.allowed_mm_limits.values())
177
+ limit_per_prompt = max(mm_budget.mm_max_items_per_prompt.values())
178
num_encoder_tokens = self.model.get_num_mm_encoder_tokens(
179
mm_budget.get_encoder_budget()
180
)
vllm/multimodal/budget.py
new
+154
@@ -0,0 +1,154 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Mapping
4
+
5
+from vllm.config import ModelConfig, VllmConfig
6
+from vllm.multimodal.processing import BaseMultiModalProcessor
7
+from vllm.multimodal.registry import MultiModalRegistry
8
+from vllm.utils.torch_utils import set_default_torch_num_threads
9
+from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
10
+
11
+
12
+def get_mm_max_toks_per_item(
13
+ model_config: ModelConfig,
14
+ mm_registry: MultiModalRegistry,
15
+ processor: BaseMultiModalProcessor,
16
+ mm_counts: Mapping[str, int],
17
+) -> Mapping[str, int]:
18
+ """
19
+ Get the maximum number of tokens per data item from each modality based
20
+ on underlying model configuration.
21
+ """
22
+ max_tokens_per_item = processor.info.get_mm_max_tokens_per_item(
23
+ seq_len=model_config.max_model_len,
24
+ mm_counts=mm_counts,
25
+ )
26
+ if max_tokens_per_item is not None:
27
+ return max_tokens_per_item
28
+
29
+ mm_inputs = mm_registry.get_dummy_mm_inputs(
30
+ model_config,
31
+ mm_counts=mm_counts,
32
+ processor=processor,
33
+ )
34
+
35
+ return {
36
+ modality: sum(item.get_num_embeds for item in placeholders)
37
+ for modality, placeholders in mm_inputs["mm_placeholders"].items()
38
+ }
39
+
40
+
41
+class MultiModalBudget:
42
+ """Helper class to calculate budget information for multi-modal models."""
43
+
44
+ def __init__(
45
+ self,
46
+ vllm_config: VllmConfig,
47
+ mm_registry: MultiModalRegistry,
48
+ ) -> None:
49
+ super().__init__()
50
+
51
+ self.model_config = model_config = vllm_config.model_config
52
+ self.scheduler_config = scheduler_config = vllm_config.scheduler_config
53
+
54
+ self.max_model_len = model_config.max_model_len
55
+ self.max_num_reqs = scheduler_config.max_num_seqs
56
+
57
+ cache = mm_registry.processor_only_cache_from_config(vllm_config)
58
+ processor = mm_registry.create_processor(model_config, cache=cache)
59
+
60
+ self.cache = cache
61
+ self.mm_limits = mm_limits = processor.info.allowed_mm_limits
62
+
63
+ active_modalities = {
64
+ modality for modality, limit in mm_limits.items() if limit > 0
65
+ }
66
+
67
+ with set_default_torch_num_threads(): # Avoid hang during startup
68
+ all_mm_max_toks_per_item = get_mm_max_toks_per_item(
69
+ model_config,
70
+ mm_registry,
71
+ processor,
72
+ mm_counts=dict.fromkeys(active_modalities, 1),
73
+ )
74
+
75
+ mm_max_toks_per_item = {
76
+ modality: all_mm_max_toks_per_item[modality]
77
+ for modality in active_modalities
78
+ }
79
+
80
+ encoder_compute_budget, encoder_cache_size = compute_mm_encoder_budget(
81
+ scheduler_config,
82
+ mm_max_toks_per_item,
83
+ )
84
+
85
+ self.encoder_compute_budget = encoder_compute_budget
86
+ self.encoder_cache_size = encoder_cache_size
87
+
88
+ mm_max_items_per_prompt = dict[str, int]()
89
+ mm_max_items_per_batch = dict[str, int]()
90
+
91
+ for modality, max_toks_per_item in mm_max_toks_per_item.items():
92
+ (
93
+ mm_max_items_per_prompt[modality],
94
+ mm_max_items_per_batch[modality],
95
+ ) = self._get_max_items(modality, max_toks_per_item)
96
+
97
+ self.mm_max_toks_per_item = mm_max_toks_per_item
98
+ self.mm_max_items_per_prompt: Mapping[str, int] = mm_max_items_per_prompt
99
+ self.mm_max_items_per_batch: Mapping[str, int] = mm_max_items_per_batch
100
+
101
+ def _get_max_items(
102
+ self,
103
+ modality: str,
104
+ max_tokens_per_item: int,
105
+ ) -> tuple[int, int]:
106
+ if max_tokens_per_item == 0:
107
+ return 0, 0
108
+
109
+ # Check how many items of this modality can be supported by
110
+ # the encoder budget.
111
+ if (encoder_budget := self.get_encoder_budget()) == 0:
112
+ return 0, 0
113
+
114
+ max_encoder_items_per_batch = encoder_budget // max_tokens_per_item
115
+
116
+ # Check how many items of this modality can be supported by
117
+ # the decoder budget.
118
+ mm_limit = self.mm_limits[modality]
119
+
120
+ max_items_per_prompt = max(
121
+ 1,
122
+ min(mm_limit, self.max_model_len // max_tokens_per_item),
123
+ )
124
+
125
+ scheduler_config = self.scheduler_config
126
+ max_num_reqs = self.max_num_reqs
127
+
128
+ if not scheduler_config.enable_chunked_prefill:
129
+ max_num_reqs = min(
130
+ max_num_reqs,
131
+ scheduler_config.max_num_batched_tokens // max_tokens_per_item,
132
+ )
133
+
134
+ max_decoder_items_per_batch = max_num_reqs * max_items_per_prompt
135
+
136
+ max_items_per_batch = max(
137
+ 1,
138
+ min(max_encoder_items_per_batch, max_decoder_items_per_batch),
139
+ )
140
+
141
+ return max_items_per_prompt, max_items_per_batch
142
+
143
+ def get_modality_with_max_tokens(self) -> str:
144
+ mm_max_toks_per_item = self.mm_max_toks_per_item
145
+ modality, _ = max(mm_max_toks_per_item.items(), key=lambda x: x[1])
146
+
147
+ return modality
148
+
149
+ def get_encoder_budget(self) -> int:
150
+ return min(self.encoder_compute_budget, self.encoder_cache_size)
151
+
152
+ def reset_cache(self) -> None:
153
+ if self.cache is not None:
154
+ self.cache.clear_cache()
vllm/multimodal/registry.py
+4
-95
@@ -130,15 +130,13 @@ class MultiModalRegistry:
130
if not model_config.is_multimodal_model:
131
return False
132
133
- info = self._create_processing_info(model_config, tokenizer=None)
134
- supported_modalities = info.get_supported_mm_limits()
135
-
133
mm_config = model_config.get_multimodal_config()
134
+ info = self._create_processing_info(model_config, tokenizer=None)
135
136
# Check if all supported modalities have limit == 0
137
if all(
138
mm_config.get_limit_per_prompt(modality) == 0
141
- for modality in supported_modalities
139
+ for modality in info.supported_mm_limits
140
):
141
logger.info_once(
142
"All limits of multimodal modalities supported by the model "
@@ -148,70 +146,6 @@ class MultiModalRegistry:
146
147
return True
148
151
- def get_max_tokens_per_item_by_modality(
152
- self,
153
- model_config: "ModelConfig",
154
- *,
155
- cache: BaseMultiModalProcessorCache | None = None,
156
- profiler_limits: Mapping[str, int] | None = None,
157
- observability_config: ObservabilityConfig | None = None,
158
- ) -> Mapping[str, int]:
159
- """
160
- Get the maximum number of tokens per data item from each modality based
161
- on underlying model configuration.
162
- """
163
- if not model_config.is_multimodal_model:
164
- return {}
165
-
166
- processor = self.create_processor(
167
- model_config, observability_config, cache=cache
168
- )
169
-
170
- if profiler_limits is None:
171
- profiler_limits = processor.info.allowed_mm_limits
172
-
173
- mm_counts = {
174
- modality: 1 for modality, limit in profiler_limits.items() if limit > 0
175
- }
176
-
177
- max_tokens_per_item = processor.info.get_mm_max_tokens_per_item(
178
- seq_len=model_config.max_model_len,
179
- mm_counts=mm_counts,
180
- )
181
- if max_tokens_per_item is not None:
182
- return {
183
- modality: max_tokens
184
- for modality, max_tokens in max_tokens_per_item.items()
185
- if mm_counts.get(modality, 0) > 0
186
- }
187
-
188
- mm_inputs = self.get_dummy_mm_inputs(
189
- model_config,
190
- mm_counts=mm_counts,
191
- processor=processor,
192
- )
193
-
194
- return {
195
- modality: sum(item.get_num_embeds for item in placeholders)
196
- for modality, placeholders in mm_inputs["mm_placeholders"].items()
197
- }
198
-
199
- def get_mm_limits_per_prompt(
200
- self,
201
- model_config: "ModelConfig",
202
- *,
203
- observability_config: ObservabilityConfig | None = None,
204
- ) -> Mapping[str, int]:
205
- """
206
- Get the maximum number of multi-modal input instances for each modality
207
- that are allowed per prompt for a model class.
208
- """
209
- if not model_config.is_multimodal_model:
210
- return {}
211
-
212
- info = self._create_processing_info(model_config, observability_config)
213
- return info.allowed_mm_limits
214
-
149
def register_processor(
150
self,
151
processor: MultiModalProcessorFactory[_I],
@@ -303,10 +237,9 @@ class MultiModalRegistry:
237
def get_dummy_mm_inputs(
238
self,
239
model_config: "ModelConfig",
306
- mm_counts: Mapping[str, int] | None = None,
240
+ mm_counts: Mapping[str, int],
241
*,
242
cache: BaseMultiModalProcessorCache | None = None,
309
- observability_config: ObservabilityConfig | None = None,
243
processor: BaseMultiModalProcessor | None = None,
244
) -> MultiModalInputs:
245
"""
@@ -317,11 +250,7 @@ class MultiModalRegistry:
250
seq_len = model_config.max_model_len
251
252
if processor is None:
320
- processor = self.create_processor(
321
- model_config, observability_config, cache=cache
322
- )
323
- if mm_counts is None:
324
- mm_counts = processor.info.allowed_mm_limits
253
+ processor = self.create_processor(model_config, cache=cache)
254
255
processor_inputs = processor.dummy_inputs.get_dummy_processor_inputs(
256
seq_len=seq_len,
@@ -342,26 +271,6 @@ class MultiModalRegistry:
271
272
return mm_inputs
273
345
- def get_encdec_max_encoder_len(self, model_config: "ModelConfig") -> int:
346
- """
347
- Get the maximum length of the encoder input for encoder-decoder models.
348
- """
349
- if not model_config.is_encoder_decoder:
350
- return 0
351
- max_tokens = self.get_max_tokens_per_item_by_modality(model_config)
352
- if not max_tokens:
353
- # TODO - this function assumes encoder-decoder models are
354
- # multimodal. This will need to change when adding support for more
355
- # than whisper.
356
- return 0
357
- assert len(max_tokens) == 1, (
358
- "Encoder-decoder models are expected "
359
- "to implement the multimodal interface with at most one modality."
360
- )
361
-
362
- first_modality = next(iter(max_tokens))
363
- return max_tokens[first_modality]
364
-
274
def _get_cache_type(
275
self,
276
vllm_config: "VllmConfig",
vllm/v1/core/encoder_cache_manager.py
+5
-50
@@ -6,11 +6,10 @@ from collections.abc import Mapping
6
from typing import TYPE_CHECKING
7
8
from vllm.logger import init_logger
9
-from vllm.multimodal import MultiModalRegistry
9
from vllm.v1.request import Request
10
11
if TYPE_CHECKING:
13
- from vllm.config import ModelConfig, SchedulerConfig
12
+ from vllm.config import SchedulerConfig
13
14
logger = init_logger(__name__)
15
@@ -267,60 +266,16 @@ class EncoderCacheManager:
266
return freed
267
268
270
-def compute_encoder_budget(
271
- model_config: "ModelConfig",
272
- scheduler_config: "SchedulerConfig",
273
- mm_registry: MultiModalRegistry,
274
-) -> tuple[int, int]:
275
- """Compute the encoder cache budget based on the model and scheduler
276
- configurations.
277
-
278
- Returns:
279
- - Compute budget for encoder execution, measured in number of tokens
280
- from the input sequence.
281
- - Space budget for encoder cache size, measured in number of tokens
282
- from the input sequence.
283
- """
284
- if mm_registry.supports_multimodal_inputs(model_config):
285
- max_tokens_by_modality = mm_registry.get_max_tokens_per_item_by_modality(
286
- model_config
287
- )
288
-
289
- return compute_mm_encoder_budget(
290
- scheduler_config,
291
- max_tokens_by_modality,
292
- )
293
-
294
- return compute_text_encoder_budget(scheduler_config)
295
-
296
-
297
-def compute_text_encoder_budget(scheduler_config: "SchedulerConfig") -> tuple[int, int]:
298
- """Compute the encoder cache budget based on the model and scheduler
299
- configurations for a text-only model.
300
-
301
- Args:
302
- scheduler_config: Scheduler configuration.
303
-
304
- Returns:
305
- - Compute budget for encoder execution, in unit of number of tokens
306
- in the input sequence.
307
- - Space budget for encoder cache size, in unit of number of tokens
308
- in the input sequence.
309
- """
310
- # Currently text-only encoder-decoder models are not supported
311
- return 0, 0
312
-
313
-
269
def compute_mm_encoder_budget(
270
scheduler_config: "SchedulerConfig",
316
- max_tokens_by_modality: Mapping[str, int],
271
+ mm_max_toks_per_item: Mapping[str, int],
272
) -> tuple[int, int]:
273
"""Compute the encoder cache budget based on the model and scheduler
274
configurations for a multimodal model.
275
276
Args:
277
scheduler_config: Scheduler configuration.
323
- max_tokens_by_modality: The maximum number of tokens for each
278
+ mm_max_toks_per_item: The maximum number of tokens per item for each
279
non-text modality.
280
281
Returns:
@@ -330,7 +285,7 @@ def compute_mm_encoder_budget(
285
from the input sequence.
286
"""
287
333
- if not max_tokens_by_modality:
288
+ if not mm_max_toks_per_item:
289
logger.warning(
290
"All non-text modalities supported by the model have been "
291
"explicitly disabled via limit_mm_per_prompt. Encoder cache will "
@@ -338,7 +293,7 @@ def compute_mm_encoder_budget(
293
)
294
return 0, 0
295
341
- max_tokens_per_mm_item = max(max_tokens_by_modality.values())
296
+ max_tokens_per_mm_item = max(mm_max_toks_per_item.values())
297
298
if (
299
scheduler_config.disable_chunked_mm_input
vllm/v1/core/sched/scheduler.py
+18
-16
@@ -31,10 +31,10 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
31
RoutedExpertsReader,
32
)
33
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
34
+from vllm.multimodal.budget import MultiModalBudget
35
from vllm.v1.core.encoder_cache_manager import (
36
EncoderCacheManager,
37
EncoderDecoderCacheManager,
37
- compute_encoder_budget,
38
)
39
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
40
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
@@ -174,22 +174,22 @@ class Scheduler(SchedulerInterface):
174
175
# Encoder-related.
176
# Calculate encoder cache size if applicable
177
- # NOTE: For now we use the same budget for both compute and space.
178
- # This can be changed when we make encoder cache for embedding caching
179
- # across requests.
180
- encoder_compute_budget, encoder_cache_size = compute_encoder_budget(
181
- model_config=vllm_config.model_config,
182
- scheduler_config=vllm_config.scheduler_config,
183
- mm_registry=mm_registry,
177
+ self.supports_mm_inputs = mm_registry.supports_multimodal_inputs(
178
+ vllm_config.model_config
179
+ )
180
+ self.mm_budget = mm_budget = (
181
+ MultiModalBudget(vllm_config, mm_registry)
182
+ if self.supports_mm_inputs
183
+ else None
184
)
185
186
- # NOTE(woosuk): Here, "encoder" includes the vision encoder (and
187
- # projector if needed) for MM models as well as encoder-decoder
188
- # transformers.
189
- self.max_num_encoder_input_tokens = encoder_compute_budget
190
- # NOTE: For the models without encoder (e.g., text-only models),
191
- # the encoder cache will not be initialized because cache size is 0
192
- # for these models.
186
+ # NOTE: Text-only encoder-decoder models are implemented as
187
+ # multi-modal models for convenience
188
+ # Example: https://github.com/neuralmagic/bart-plugin
189
+ self.max_num_encoder_input_tokens = (
190
+ mm_budget.encoder_compute_budget if mm_budget else 0
191
+ )
192
+ encoder_cache_size = mm_budget.encoder_cache_size if mm_budget else 0
193
self.encoder_cache_manager = (
194
EncoderDecoderCacheManager(cache_size=encoder_cache_size)
195
if self.is_encoder_decoder
@@ -199,7 +199,9 @@ class Scheduler(SchedulerInterface):
199
# Attn blocks, as for Whisper its input is always padded to the maximum length.
200
# TODO (NickLucche): Generalize to models with variable-length encoder inputs.
201
self._num_encoder_max_input_tokens = (
202
- MULTIMODAL_REGISTRY.get_encdec_max_encoder_len(vllm_config.model_config)
202
+ mm_budget.mm_max_toks_per_item[mm_budget.get_modality_with_max_tokens()]
203
+ if mm_budget
204
+ else 0
205
)
206
207
speculative_config = vllm_config.speculative_config
vllm/v1/engine/input_processor.py
+14
-16
@@ -19,6 +19,7 @@ from vllm.inputs.preprocess import InputPreprocessor
19
from vllm.logger import init_logger
20
from vllm.lora.request import LoRARequest
21
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
22
+from vllm.multimodal.budget import MultiModalBudget
23
from vllm.multimodal.inputs import (
24
MultiModalDataDict,
25
MultiModalFeatureSpec,
@@ -34,7 +35,6 @@ from vllm.tokenizers import TokenizerLike
35
from vllm.tokenizers.mistral import MistralTokenizer
36
from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
37
from vllm.utils.torch_utils import set_default_torch_num_threads
37
-from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
38
from vllm.v1.engine import EngineCoreRequest
39
from vllm.v1.metrics.stats import MultiModalCacheStats
40
from vllm.v1.structured_output.backend_guidance import (
@@ -59,32 +59,30 @@ class InputProcessor:
59
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
60
) -> None:
61
self.vllm_config = vllm_config
62
- self.model_config = vllm_config.model_config
62
+ self.model_config = model_config = vllm_config.model_config
63
self.cache_config = vllm_config.cache_config
64
self.lora_config = vllm_config.lora_config
65
+ self.scheduler_config = vllm_config.scheduler_config
66
self.structured_outputs_config = vllm_config.structured_outputs_config
67
+ self.observability_config = vllm_config.observability_config
68
67
- self.generation_config_fields = self.model_config.try_get_generation_config()
69
+ self.generation_config_fields = model_config.try_get_generation_config()
70
71
self.mm_registry = mm_registry
72
self.mm_processor_cache = mm_registry.processor_cache_from_config(vllm_config)
71
- self.mm_encoder_cache_size = None
73
+
74
+ self.mm_encoder_cache_size: int | None = None
75
if (
73
- self.mm_registry.supports_multimodal_inputs(self.model_config)
74
- and not self.model_config.skip_tokenizer_init
76
+ mm_registry.supports_multimodal_inputs(model_config)
77
+ and not model_config.skip_tokenizer_init
78
):
76
- with set_default_torch_num_threads():
77
- max_tokens_by_modality = (
78
- mm_registry.get_max_tokens_per_item_by_modality(self.model_config)
79
- )
80
-
81
- _, self.mm_encoder_cache_size = compute_mm_encoder_budget(
82
- self.vllm_config.scheduler_config, max_tokens_by_modality
83
- )
79
+ mm_budget = MultiModalBudget(vllm_config, mm_registry)
80
+ self.mm_encoder_cache_size = mm_budget.encoder_cache_size
81
+ mm_budget.reset_cache() # Not used anymore
82
83
self.input_preprocessor = InputPreprocessor(
86
- self.model_config,
87
- vllm_config.observability_config,
84
+ model_config,
85
+ self.observability_config,
86
mm_registry,
87
mm_processor_cache=self.mm_processor_cache,
88
)
vllm/v1/worker/gpu_model_runner.py
+2
-2
@@ -82,6 +82,7 @@ from vllm.model_executor.models.interfaces_base import (
82
is_text_generation_model,
83
)
84
from vllm.multimodal import MULTIMODAL_REGISTRY
85
+from vllm.multimodal.budget import MultiModalBudget
86
from vllm.multimodal.inputs import (
87
BatchedTensorInputs,
88
MultiModalKwargsItem,
@@ -180,7 +181,6 @@ from vllm.v1.worker.workspace import lock_workspace
181
182
from .utils import (
183
AttentionGroup,
183
- MultiModalBudget,
184
add_kv_sharing_layers_to_kv_cache_groups,
185
bind_kv_cache,
186
sanity_check_mm_encoder_outputs,
@@ -5104,7 +5104,7 @@ class GPUModelRunner(
5104
# modality with the max possible input tokens even when
5105
# it supports multiple.
5106
dummy_modality = mm_budget.get_modality_with_max_tokens()
5107
- max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[
5107
+ max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[
5108
dummy_modality
5109
]
5110
vllm/v1/worker/utils.py
-111
@@ -11,125 +11,14 @@ from vllm.logger import init_logger
11
from vllm.model_executor.layers.attention import Attention
12
from vllm.model_executor.models.interfaces import MultiModalEmbeddings
13
from vllm.model_executor.models.utils import extract_layer_index
14
-from vllm.multimodal.registry import MultiModalRegistry
14
from vllm.platforms import current_platform
15
from vllm.utils.mem_utils import MemorySnapshot, format_gib
16
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadataBuilder
18
-from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
17
from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec
18
19
logger = init_logger(__name__)
20
21
24
-class MultiModalBudget:
25
- """Helper class to calculate budget information for multi-modal models."""
26
-
27
- def __init__(
28
- self,
29
- vllm_config: VllmConfig,
30
- mm_registry: MultiModalRegistry,
31
- ) -> None:
32
- super().__init__()
33
-
34
- self.model_config = model_config = vllm_config.model_config
35
- self.scheduler_config = scheduler_config = vllm_config.scheduler_config
36
- self.mm_registry = mm_registry
37
- self.cache = cache = mm_registry.processor_only_cache_from_config(vllm_config)
38
-
39
- self.max_model_len = model_config.max_model_len
40
- self.max_num_reqs = scheduler_config.max_num_seqs
41
-
42
- self.mm_limits = mm_registry.get_mm_limits_per_prompt(model_config)
43
-
44
- max_tokens_by_modality = mm_registry.get_max_tokens_per_item_by_modality(
45
- model_config,
46
- cache=cache,
47
- profiler_limits=self.mm_limits,
48
- )
49
-
50
- encoder_compute_budget, encoder_cache_size = compute_mm_encoder_budget(
51
- scheduler_config,
52
- max_tokens_by_modality,
53
- )
54
-
55
- self.encoder_compute_budget = encoder_compute_budget
56
- self.encoder_cache_size = encoder_cache_size
57
-
58
- max_items_per_prompt_by_modality = dict[str, int]()
59
- max_items_per_batch_by_modality = dict[str, int]()
60
-
61
- for modality, max_tokens in max_tokens_by_modality.items():
62
- (
63
- max_items_per_prompt,
64
- max_items_per_batch,
65
- ) = self.get_max_items(modality, max_tokens)
66
-
67
- max_items_per_prompt_by_modality[modality] = max_items_per_prompt
68
- max_items_per_batch_by_modality[modality] = max_items_per_batch
69
-
70
- self.max_tokens_by_modality = max_tokens_by_modality
71
- self.max_items_per_prompt_by_modality = max_items_per_prompt_by_modality
72
- self.max_items_per_batch_by_modality = max_items_per_batch_by_modality
73
-
74
- def get_modality_with_max_tokens(self) -> str:
75
- max_tokens_by_modality = self.max_tokens_by_modality
76
- modality, _ = max(max_tokens_by_modality.items(), key=lambda x: x[1])
77
-
78
- return modality
79
-
80
- def get_encoder_budget(self) -> int:
81
- return min(self.encoder_compute_budget, self.encoder_cache_size)
82
-
83
- def get_max_items(
84
- self,
85
- modality: str,
86
- max_tokens_per_item: int,
87
- ) -> tuple[int, int]:
88
- if max_tokens_per_item == 0:
89
- return 0, 0
90
-
91
- # Check how many items of this modality can be supported by
92
- # the encoder budget.
93
- encoder_budget = self.get_encoder_budget()
94
-
95
- # TODO: handle encoder-decoder models once we support them.
96
- if encoder_budget == 0:
97
- return 0, 0
98
-
99
- max_encoder_items_per_batch = encoder_budget // max_tokens_per_item
100
-
101
- # Check how many items of this modality can be supported by
102
- # the decoder budget.
103
- mm_limit = self.mm_limits[modality]
104
-
105
- max_items_per_prompt = max(
106
- 1,
107
- min(mm_limit, self.max_model_len // max_tokens_per_item),
108
- )
109
-
110
- scheduler_config = self.scheduler_config
111
- max_num_reqs = self.max_num_reqs
112
-
113
- if not scheduler_config.enable_chunked_prefill:
114
- max_num_reqs = min(
115
- max_num_reqs,
116
- scheduler_config.max_num_batched_tokens // max_tokens_per_item,
117
- )
118
-
119
- max_decoder_items_per_batch = max_num_reqs * max_items_per_prompt
120
-
121
- max_items_per_batch = max(
122
- 1,
123
- min(max_encoder_items_per_batch, max_decoder_items_per_batch),
124
- )
125
-
126
- return max_items_per_prompt, max_items_per_batch
127
-
128
- def reset_cache(self) -> None:
129
- if self.cache is not None:
130
- self.cache.clear_cache()
131
-
132
-
22
@dataclass
23
class AttentionGroup:
24
backend: type[AttentionBackend]