| 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 VllmConfig |
| 6 | from vllm.logger import init_logger |
| 7 | from vllm.multimodal.cache import ( |
| 8 | BaseMultiModalProcessorCache, |
| 9 | processor_only_cache_from_config, |
| 10 | ) |
| 11 | from vllm.multimodal.inputs import MultiModalKwargsItem |
| 12 | from vllm.multimodal.processing import BaseMultiModalProcessor |
| 13 | from vllm.multimodal.registry import MultiModalRegistry |
| 14 | from vllm.utils.torch_utils import set_default_torch_num_threads |
| 15 | from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget |
| 16 | |
| 17 | logger = init_logger(__name__) |
| 18 | |
| 19 | |
| 20 | def get_mm_max_toks_per_item( |
| 21 | processor: BaseMultiModalProcessor, |
| 22 | mm_counts: Mapping[str, int], |
| 23 | cache: BaseMultiModalProcessorCache | None, |
| 24 | ) -> Mapping[str, int]: |
| 25 | """Get the maximum number of tokens per data item from each modality based |
| 26 | on underlying model configuration. |
| 27 | """ |
| 28 | max_tokens_per_item = processor.info.get_mm_max_tokens_per_item( |
| 29 | seq_len=processor.info.ctx.model_config.max_model_len, |
| 30 | mm_counts=mm_counts, |
| 31 | ) |
| 32 | if max_tokens_per_item is not None: |
| 33 | return max_tokens_per_item |
| 34 | |
| 35 | mm_inputs = processor.get_dummy_mm_inputs(mm_counts, cache=cache) |
| 36 | |
| 37 | return { |
| 38 | modality: sum(item.get_num_embeds() for item in placeholders) |
| 39 | for modality, placeholders in mm_inputs["mm_placeholders"].items() |
| 40 | } |
| 41 | |
| 42 | |
| 43 | class MultiModalBudget: |
| 44 | """Helper class to calculate budget information for multi-modal models.""" |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | vllm_config: VllmConfig, |
| 49 | mm_registry: MultiModalRegistry, |
| 50 | *, |
| 51 | enable_cache: bool = True, |
| 52 | ) -> None: |
| 53 | super().__init__() |
| 54 | |
| 55 | self.model_config = model_config = vllm_config.model_config |
| 56 | self.scheduler_config = scheduler_config = vllm_config.scheduler_config |
| 57 | |
| 58 | self.max_model_len = model_config.max_model_len |
| 59 | self.max_num_reqs = scheduler_config.max_num_seqs |
| 60 | |
| 61 | with set_default_torch_num_threads(): # Avoid hang during startup |
| 62 | cache = ( |
| 63 | processor_only_cache_from_config(vllm_config) if enable_cache else None |
| 64 | ) |
| 65 | processor = mm_registry.create_processor(model_config) |
| 66 | |
| 67 | self.cache = cache |
| 68 | self.processor = processor |
| 69 | mm_config = model_config.get_multimodal_config() |
| 70 | enable_mm_embeds = mm_config is not None and mm_config.enable_mm_embeds |
| 71 | |
| 72 | supported_mm_limits = processor.info.supported_mm_limits |
| 73 | self.mm_limits = mm_limits = processor.info.allowed_mm_limits |
| 74 | |
| 75 | # Modalities that pass through the MM encoder tower |
| 76 | tower_modalities = { |
| 77 | modality |
| 78 | for modality in supported_mm_limits |
| 79 | if mm_limits.get(modality, 0) > 0 |
| 80 | } |
| 81 | # Modalities that bypass the tower (pre-computed embeddings only) |
| 82 | embed_only_modalities = { |
| 83 | modality |
| 84 | for modality in supported_mm_limits |
| 85 | if enable_mm_embeds and mm_limits.get(modality, 0) == 0 |
| 86 | } |
| 87 | |
| 88 | active_modalities = tower_modalities | embed_only_modalities |
| 89 | |
| 90 | all_mm_max_toks_per_item = get_mm_max_toks_per_item( |
| 91 | processor, |
| 92 | cache=cache, |
| 93 | mm_counts=dict.fromkeys(active_modalities, 1), |
| 94 | ) |
| 95 | |
| 96 | if embed_only_modalities: |
| 97 | logger.info_once( |
| 98 | "enable_mm_embeds is True; modalities handled as embedding-only: %s", |
| 99 | tuple(embed_only_modalities), |
| 100 | ) |
| 101 | |
| 102 | # Some models (e.g., Qwen3Omni with use_audio_in_video=True) share |
| 103 | # placeholders between modalities, so not all active modalities will |
| 104 | # have their own entry in the returned dict. We filter to only include |
| 105 | # modalities that have independent placeholder tokens. |
| 106 | active_mm_max_toks_per_item = { |
| 107 | modality: all_mm_max_toks_per_item[modality] |
| 108 | for modality in active_modalities |
| 109 | if modality in all_mm_max_toks_per_item |
| 110 | } |
| 111 | tower_mm_max_toks_per_item = { |
| 112 | modality: active_mm_max_toks_per_item[modality] |
| 113 | for modality in tower_modalities |
| 114 | if modality in active_mm_max_toks_per_item |
| 115 | } |
| 116 | |
| 117 | # Encoder budget is computed from all active modalities (including |
| 118 | # embedding-only ones that need encoder cache space). |
| 119 | encoder_compute_budget, encoder_cache_size = compute_mm_encoder_budget( |
| 120 | scheduler_config, |
| 121 | active_mm_max_toks_per_item, |
| 122 | ) |
| 123 | |
| 124 | self.encoder_compute_budget = encoder_compute_budget |
| 125 | self.encoder_cache_size = encoder_cache_size |
| 126 | |
| 127 | mm_max_items_per_prompt = dict[str, int]() |
| 128 | mm_max_items_per_batch = dict[str, int]() |
| 129 | |
| 130 | # Per-prompt/per-batch limits are only relevant for tower modalities |
| 131 | # (embedding-only modalities don't go through the encoder tower). |
| 132 | for modality, max_toks_per_item in tower_mm_max_toks_per_item.items(): |
| 133 | ( |
| 134 | mm_max_items_per_prompt[modality], |
| 135 | mm_max_items_per_batch[modality], |
| 136 | ) = self._get_max_items(modality, max_toks_per_item) |
| 137 | |
| 138 | self.mm_max_toks_per_item = tower_mm_max_toks_per_item |
| 139 | self.mm_max_items_per_prompt: Mapping[str, int] = mm_max_items_per_prompt |
| 140 | self.mm_max_items_per_batch: Mapping[str, int] = mm_max_items_per_batch |
| 141 | |
| 142 | def _get_max_items( |
| 143 | self, |
| 144 | modality: str, |
| 145 | max_tokens_per_item: int, |
| 146 | ) -> tuple[int, int]: |
| 147 | if max_tokens_per_item == 0: |
| 148 | return 0, 0 |
| 149 | |
| 150 | # Check how many items of this modality can be supported by |
| 151 | # the encoder budget. |
| 152 | if (encoder_budget := self.get_encoder_budget()) == 0: |
| 153 | return 0, 0 |
| 154 | |
| 155 | max_encoder_items_per_batch = encoder_budget // max_tokens_per_item |
| 156 | |
| 157 | # Check how many items of this modality can be supported by |
| 158 | # the decoder budget. |
| 159 | mm_limit = self.mm_limits[modality] |
| 160 | |
| 161 | max_items_per_prompt = max( |
| 162 | 1, |
| 163 | min(mm_limit, self.max_model_len // max_tokens_per_item), |
| 164 | ) |
| 165 | |
| 166 | scheduler_config = self.scheduler_config |
| 167 | max_num_reqs = self.max_num_reqs |
| 168 | |
| 169 | if not scheduler_config.enable_chunked_prefill: |
| 170 | max_num_reqs = min( |
| 171 | max_num_reqs, |
| 172 | scheduler_config.max_num_batched_tokens // max_tokens_per_item, |
| 173 | ) |
| 174 | |
| 175 | max_decoder_items_per_batch = max_num_reqs * max_items_per_prompt |
| 176 | |
| 177 | max_items_per_batch = max( |
| 178 | 1, |
| 179 | min(max_encoder_items_per_batch, max_decoder_items_per_batch), |
| 180 | ) |
| 181 | |
| 182 | return max_items_per_prompt, max_items_per_batch |
| 183 | |
| 184 | def get_modality_with_max_tokens(self) -> str: |
| 185 | mm_max_toks_per_item = self.mm_max_toks_per_item |
| 186 | modality, _ = max(mm_max_toks_per_item.items(), key=lambda x: (x[1], x[0])) |
| 187 | |
| 188 | return modality |
| 189 | |
| 190 | def get_encoder_budget(self) -> int: |
| 191 | return min(self.encoder_compute_budget, self.encoder_cache_size) |
| 192 | |
| 193 | def get_dummy_encoder_profile_inputs( |
| 194 | self, |
| 195 | modality: str | None = None, |
| 196 | max_items_per_batch: int | None = None, |
| 197 | ) -> list[tuple[str, MultiModalKwargsItem]]: |
| 198 | if self.get_encoder_budget() <= 0 or not self.mm_max_toks_per_item: |
| 199 | return [] |
| 200 | |
| 201 | if modality is None: |
| 202 | modality = self.get_modality_with_max_tokens() |
| 203 | if max_items_per_batch is None: |
| 204 | max_items_per_batch = self.mm_max_items_per_batch[modality] |
| 205 | |
| 206 | # Don't use `max_items_per_batch` here to avoid redundant computation |
| 207 | dummy_mm_inputs = self.processor.get_dummy_mm_inputs( |
| 208 | mm_counts={modality: 1}, |
| 209 | cache=self.cache, |
| 210 | ) |
| 211 | dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0] |
| 212 | |
| 213 | # We use the cache so that the item is saved to the cache, |
| 214 | # but not read from the cache |
| 215 | assert dummy_mm_item is not None, "Item should not already be cached" |
| 216 | |
| 217 | return [(modality, dummy_mm_item)] * max_items_per_batch |
| 218 | |
| 219 | def reset_cache(self) -> None: |
| 220 | if self.cache is not None: |
| 221 | self.cache.clear_cache() |