@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 811ac13d03

Fork 0

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

[Core] Factor out common logic for MM budget calculation (#22228)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Aug 5, 2025 at 14:54 UTC 811ac13d039648f2d78a636ce4366e70449380c8

3 files changed +306 -223

vllm/v1/worker/gpu_model_runner.py

+108 -115

@@ -36,7 +36,8 @@ from vllm.model_executor.models.interfaces import (is_mixture_of_experts,

36 from vllm.model_executor.models.interfaces_base import (

37 VllmModelForPooling, is_pooling_model, is_text_generation_model)

38 from vllm.multimodal import MULTIMODAL_REGISTRY

39 -from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange

39 +from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,

40 + PlaceholderRange)

41 from vllm.multimodal.utils import group_mm_inputs_by_modality

42 from vllm.pooling_params import PoolingParams

43 from vllm.sampling_params import SamplingType

@@ -51,7 +52,6 @@ from vllm.v1.attention.backends.utils import (

52 make_kv_sharing_fast_prefill_attention_metadata,

53 make_local_attention_virtual_batches,

54 reorder_batch_to_split_decodes_and_prefills)

54 -from vllm.v1.core.encoder_cache_manager import compute_encoder_budget

55 from vllm.v1.kv_cache_interface import (AttentionSpec,

56 ChunkedLocalAttentionSpec,

57 FullAttentionSpec, KVCacheConfig,

@@ -73,7 +73,7 @@ from vllm.v1.worker.kv_connector_model_runner_mixin import (

73 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

75 from ..sample.logits_processor import LogitsProcessorManager

76 -from .utils import (bind_kv_cache, gather_mm_placeholders,

76 +from .utils import (MultiModalBudget, bind_kv_cache, gather_mm_placeholders,

77 initialize_kv_cache_for_kv_sharing,

78 sanity_check_mm_encoder_outputs, scatter_mm_placeholders)

@@ -148,14 +148,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

148 self.mm_registry = MULTIMODAL_REGISTRY

149 self.uses_mrope = model_config.uses_mrope

150

151 - encoder_compute_budget, encoder_cache_size = compute_encoder_budget(

152 - model_config=model_config,

153 - scheduler_config=scheduler_config,

154 - mm_registry=self.mm_registry,

155 - )

156 - self.max_num_encoder_input_tokens = encoder_compute_budget

157 - self.encoder_cache_size = encoder_cache_size

158 -

151 # Sampler

152 self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)

153

@@ -330,6 +322,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

322 self.kv_sharing_fast_prefill_logits_indices = torch.zeros(

323 self.max_num_tokens, dtype=torch.int32, device=self.device)

324

325 + self.mm_budget = (MultiModalBudget(

326 + self.model_config,

327 + self.scheduler_config,

328 + self.mm_registry,

329 + max_model_len=self.max_model_len,

330 + max_num_reqs=self.max_num_reqs,

331 + ) if self.is_multimodal_model else None)

332 +

333 self.reorder_batch_threshold: Optional[int] = None

334

335 def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:

@@ -578,37 +578,33 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

578 # Refresh batch metadata with any pending updates.

579 self.input_batch.refresh_metadata()

580

581 - def _init_model_kwargs_for_multimodal_model(

581 + def _extract_mm_kwargs(

582 self,

583 - scheduler_output: Optional["SchedulerOutput"] = None,

584 - num_reqs: int = -1,

585 - ) -> dict[str, Any]:

586 -

587 - model_kwargs: dict[str, Any] = {}

588 - if self.is_multimodal_raw_input_supported:

589 - # This model requires the raw multimodal data in input.

583 + scheduler_output: "SchedulerOutput",

584 + ) -> BatchedTensorInputs:

585 + if self.is_multimodal_raw_input_supported: # noqa: SIM102

586 if scheduler_output:

591 - multi_modal_kwargs_list = []

587 + multi_modal_kwargs_list = list[MultiModalKwargs]()

588 for req in scheduler_output.scheduled_new_reqs:

589 req_mm_inputs = req.mm_inputs

590 if not isinstance(req_mm_inputs, list):

591 req_mm_inputs = list(req_mm_inputs)

592 multi_modal_kwargs_list.extend(req_mm_inputs)

597 - multi_modal_kwargs = MultiModalKwargs.batch(

598 - multi_modal_kwargs_list)

599 - else:

600 - # The only case where SchedulerOutput is None is for

601 - # a dummy run let's get some dummy data.

602 - dummy_data = [

603 - self.mm_registry.get_decoder_dummy_data(

604 - model_config=self.model_config,

605 - seq_len=1).multi_modal_data for i in range(num_reqs)

606 - ]

607 - multi_modal_kwargs = MultiModalKwargs.batch(dummy_data)

593

609 - model_kwargs.update(multi_modal_kwargs)

594 + return MultiModalKwargs.batch(multi_modal_kwargs_list)

595

611 - return model_kwargs

596 + return {}

597 +

598 + def _dummy_mm_kwargs(self, num_seqs: int) -> BatchedTensorInputs:

599 + if self.is_multimodal_raw_input_supported:

600 + mm_budget = self.mm_budget

601 + assert mm_budget is not None

602 +

603 + dummy_modality, _ = mm_budget.get_modality_with_max_tokens()

604 +

605 + return self._get_mm_dummy_batch(dummy_modality, num_seqs)

606 +

607 + return {}

608

609 def _get_cumsum_and_arange(

610 self,

@@ -1517,19 +1513,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1513 # NOTE(woosuk): To unify token ids and soft tokens (vision

1514 # embeddings), we always use embeddings (rather than token ids)

1515 # as input to the multimodal model, even when the input is text.

1520 - input_ids = self.input_ids[:num_scheduled_tokens]

1521 -

1522 - model_kwargs = self._init_model_kwargs_for_multimodal_model(

1523 - scheduler_output=scheduler_output)

1524 - inputs_embeds = self.model.get_input_embeddings(

1525 - input_ids=input_ids,

1516 + inputs_embeds_scheduled = self.model.get_input_embeddings(

1517 + input_ids=self.input_ids[:num_scheduled_tokens],

1518 multimodal_embeddings=mm_embeds or None,

1519 )

1520

1521 # TODO(woosuk): Avoid the copy. Optimize.

1530 - self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds)

1531 - inputs_embeds = self.inputs_embeds[:num_input_tokens]

1522 + self.inputs_embeds[:num_scheduled_tokens].copy_(

1523 + inputs_embeds_scheduled)

1524 +

1525 input_ids = None

1526 + inputs_embeds = self.inputs_embeds[:num_input_tokens]

1527 + model_mm_kwargs = self._extract_mm_kwargs(scheduler_output)

1528 else:

1529 # For text-only models, we use token ids as input.

1530 # While it is possible to use embeddings as input just like the

@@ -1537,7 +1532,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1532 # then the embedding layer is not included in the CUDA graph.

1533 input_ids = self.input_ids[:num_input_tokens]

1534 inputs_embeds = None

1540 - model_kwargs = {}

1535 + model_mm_kwargs = {}

1536 if self.uses_mrope:

1537 positions = self.mrope_positions[:, :num_input_tokens]

1538 else:

@@ -1571,7 +1566,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1566 intermediate_tensors=intermediate_tensors,

1567 inputs_embeds=inputs_embeds,

1568 **MultiModalKwargs.as_kwargs(

1574 - model_kwargs,

1569 + model_mm_kwargs,

1570 device=self.device,

1571 ),

1572 )

@@ -2149,6 +2144,30 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2144 yield

2145 input_ids.fill_(0)

2146

2147 + def _get_mm_dummy_batch(

2148 + self,

2149 + modality: str,

2150 + max_items_per_batch: int,

2151 + ) -> BatchedTensorInputs:

2152 + """Dummy data for profiling and precompiling multimodal models."""

2153 + dummy_decoder_data = self.mm_registry.get_decoder_dummy_data(

2154 + model_config=self.model_config,

2155 + seq_len=self.max_num_tokens,

2156 + mm_counts={modality: 1},

2157 + )

2158 + dummy_mm_data = dummy_decoder_data.multi_modal_data

2159 +

2160 + # Result in the maximum GPU consumption of the model

2161 + dummy_mm_item = dummy_mm_data.get_item(modality=modality, item_index=0)

2162 + dummy_mm_kwargs = MultiModalKwargs.from_items([dummy_mm_item])

2163 +

2164 + batched_dummy_mm_inputs = MultiModalKwargs.batch([dummy_mm_kwargs] *

2165 + max_items_per_batch)

2166 + return MultiModalKwargs.as_kwargs(

2167 + batched_dummy_mm_inputs,

2168 + device=self.device,

2169 + )

2170 +

2171 @torch.inference_mode()

2172 def _dummy_run(

2173 self,

@@ -2213,16 +2232,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2232

2233 with self.maybe_dummy_run_with_lora(self.lora_config,

2234 num_scheduled_tokens):

2216 - model = self.model

2235 if self.is_multimodal_model:

2218 - model_kwargs = self._init_model_kwargs_for_multimodal_model(

2219 - num_reqs=num_reqs)

2236 input_ids = None

2237 inputs_embeds = self.inputs_embeds[:num_tokens]

2238 + model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)

2239 else:

2240 input_ids = self.input_ids[:num_tokens]

2241 inputs_embeds = None

2225 - model_kwargs = {}

2242 + model_mm_kwargs = {}

2243

2244 if self.uses_mrope:

2245 positions = self.mrope_positions[:, :num_tokens]

@@ -2247,13 +2264,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2264 self.vllm_config,

2265 num_tokens=num_tokens,

2266 num_tokens_across_dp=num_tokens_across_dp):

2250 - outputs = model(

2267 + outputs = self.model(

2268 input_ids=input_ids,

2269 positions=positions,

2270 intermediate_tensors=intermediate_tensors,

2271 inputs_embeds=inputs_embeds,

2272 **MultiModalKwargs.as_kwargs(

2256 - model_kwargs,

2273 + model_mm_kwargs,

2274 device=self.device,

2275 ),

2276 )

@@ -2423,75 +2440,51 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2440

2441 def profile_run(self) -> None:

2442 # Profile with multimodal encoder & encoder cache.

2426 - # TODO: handle encoder-decoder models once we support them.

2427 - if (self.is_multimodal_model and self.max_num_encoder_input_tokens > 0

2428 - and self.encoder_cache_size > 0):

2429 -

2430 - # NOTE: Currently model is profiled with a single non-text

2431 - # modality with the max possible input tokens even when

2432 - # it supports multiple.

2433 - max_tokens_by_modality_dict = self.mm_registry \

2434 - .get_max_tokens_per_item_by_nonzero_modality(self.model_config)

2435 - dummy_data_modality, max_tokens_per_mm_item = max(

2436 - max_tokens_by_modality_dict.items(), key=lambda item: item[1])

2437 -

2438 - # Check how many items of this modality can be supported by

2439 - # the encoder budget.

2440 - encoder_budget = min(self.max_num_encoder_input_tokens,

2441 - self.encoder_cache_size)

2442 -

2443 - max_num_mm_items_encoder_budget = encoder_budget // \

2444 - max_tokens_per_mm_item

2445 -

2446 - # Check how many items of this modality can be supported by

2447 - # the decoder budget.

2448 - max_mm_items_per_req = self.mm_registry.get_mm_limits_per_prompt(

2449 - self.model_config)[dummy_data_modality]

2450 -

2451 - # NOTE: We do not consider max_num_batched_tokens on purpose

2452 - # because the multimodal embeddings can be generated in advance

2453 - # and chunked prefilled.

2454 - max_num_mm_items_decoder_budget = self.max_num_reqs * \

2455 - max_mm_items_per_req

2456 -

2457 - max_num_mm_items = max(

2458 - 1,

2459 - min(max_num_mm_items_encoder_budget,

2460 - max_num_mm_items_decoder_budget))

2461 -

2462 - logger.info(

2463 - "Encoder cache will be initialized with a budget of %s tokens,"

2464 - " and profiled with %s %s items of the maximum feature size.",

2465 - encoder_budget, max_num_mm_items, dummy_data_modality)

2466 -

2467 - # Create dummy batch of multimodal inputs.

2468 - dummy_mm_kwargs = self.mm_registry.get_decoder_dummy_data(

2469 - model_config=self.model_config,

2470 - seq_len=max_tokens_per_mm_item,

2471 - mm_counts={

2472 - dummy_data_modality: 1

2473 - },

2474 - ).multi_modal_data

2475 -

2476 - batched_dummy_mm_inputs = MultiModalKwargs.batch(

2477 - [dummy_mm_kwargs] * max_num_mm_items,

2478 - pin_memory=self.pin_memory)

2479 - batched_dummy_mm_inputs = MultiModalKwargs.as_kwargs(

2480 - batched_dummy_mm_inputs,

2481 - device=self.device,

2482 - )

2443 + if self.is_multimodal_model:

2444 + mm_budget = self.mm_budget

2445 + assert mm_budget is not None

2446 +

2447 + # TODO: handle encoder-decoder models once we support them.

2448 + if (encoder_budget := mm_budget.get_encoder_budget()) > 0:

2449 + # NOTE: Currently model is profiled with a single non-text

2450 + # modality with the max possible input tokens even when

2451 + # it supports multiple.

2452 + (

2453 + dummy_modality,

2454 + max_tokens,

2455 + ) = mm_budget.get_modality_with_max_tokens()

2456 + (

2457 + max_mm_items_per_prompt,

2458 + max_mm_items_per_batch,

2459 + ) = mm_budget.get_max_items(dummy_modality, max_tokens)

2460 +

2461 + logger.info(

2462 + "Encoder cache will be initialized with a budget of "

2463 + "%s tokens, and profiled with %s %s items of the maximum "

2464 + "feature size.",

2465 + encoder_budget,

2466 + max_mm_items_per_batch,

2467 + dummy_modality,

2468 + )

2469

2484 - # Run multimodal encoder.

2485 - dummy_encoder_outputs = self.model.get_multimodal_embeddings(

2486 - **batched_dummy_mm_inputs)

2470 + # Create dummy batch of multimodal inputs.

2471 + batched_dummy_mm_inputs = self._get_mm_dummy_batch(

2472 + dummy_modality,

2473 + max_mm_items_per_batch,

2474 + )

2475

2488 - sanity_check_mm_encoder_outputs(

2489 - dummy_encoder_outputs,

2490 - expected_num_items=max_num_mm_items,

2491 - )

2476 + # Run multimodal encoder.

2477 + dummy_encoder_outputs = self.model.get_multimodal_embeddings(

2478 + **batched_dummy_mm_inputs)

2479 +

2480 + sanity_check_mm_encoder_outputs(

2481 + dummy_encoder_outputs,

2482 + expected_num_items=max_mm_items_per_batch,

2483 + )

2484

2493 - # Cache the dummy encoder outputs.

2494 - self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))

2485 + # Cache the dummy encoder outputs.

2486 + self.encoder_cache["tmp"] = dict(

2487 + enumerate(dummy_encoder_outputs))

2488

2489 # Add `is_profile` here to pre-allocate communication buffers

2490 hidden_states, last_hidden_states \

vllm/v1/worker/tpu_model_runner.py

+89 -108

@@ -42,7 +42,6 @@ from vllm.v1.attention.backends.pallas import (TPU_STR_DTYPE_TO_TORCH_DTYPE,

42 PallasAttentionBackend,

43 PallasMetadata,

44 get_page_size_bytes)

45 -from vllm.v1.core.encoder_cache_manager import compute_encoder_budget

45 from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,

46 KVCacheConfig, KVCacheSpec,

47 SlidingWindowSpec)

@@ -55,7 +54,8 @@ from vllm.v1.worker.kv_connector_model_runner_mixin import (

54 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

55 from vllm.v1.worker.tpu_input_batch import CachedRequestState, InputBatch

58 -from .utils import (bind_kv_cache, initialize_kv_cache_for_kv_sharing,

57 +from .utils import (MultiModalBudget, bind_kv_cache,

58 + initialize_kv_cache_for_kv_sharing,

59 sanity_check_mm_encoder_outputs)

61 if TYPE_CHECKING:

@@ -195,14 +195,6 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

195 # TODO: Support M-RoPE (e.g, Qwen2-VL)

196 assert not self.uses_mrope, "TPU does not support M-RoPE yet."

197

198 - encoder_compute_budget, encoder_cache_size = compute_encoder_budget(

199 - model_config=model_config,

200 - scheduler_config=scheduler_config,

201 - mm_registry=self.mm_registry,

202 - )

203 - self.max_num_encoder_input_tokens = encoder_compute_budget

204 - self.encoder_cache_size = encoder_cache_size

205 -

198 self._num_slices_per_kv_cache_update_block = \

199 _get_num_slices_per_kv_cache_update_block(get_page_size_bytes(

200 block_size=self.block_size,

@@ -294,36 +286,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

286 self.structured_decode_arange = torch.arange(

287 0, 32, device="cpu", pin_memory=self.pin_memory)

288

297 - # Get maximum number of mm items per modality (batch size).

298 - self.max_num_mm_items_by_modality = dict()

299 - if (self.is_multimodal_model and self.max_num_encoder_input_tokens > 0

300 - and self.encoder_cache_size > 0):

301 - max_tokens_by_modality_dict = (

302 - MULTIMODAL_REGISTRY.

303 - get_max_tokens_per_item_by_nonzero_modality(self.model_config))

304 - for modality, max_tokens in max_tokens_by_modality_dict.items():

305 - # Check how many items of this modality can be supported by

306 - # the encoder budget.

307 - encoder_budget = min(self.max_num_encoder_input_tokens,

308 - self.encoder_cache_size)

309 -

310 - max_num_mm_items_encoder_budget = cdiv(encoder_budget,

311 - max_tokens)

312 -

313 - # Check how many items of this modality can be supported by

314 - # the decoder budget.

315 - max_mm_items_per_req = self.mm_registry.\

316 - get_mm_limits_per_prompt(self.model_config)[modality]

317 -

318 - # NOTE: We do not consider max_num_batched_tokens on purpose

319 - # because the multimodal embeddings can be generated in advance

320 - # and chunked prefilled.

321 - max_num_mm_items_decoder_budget = self.max_num_reqs * \

322 - max_mm_items_per_req

323 -

324 - max_num_mm_items = min(max_num_mm_items_encoder_budget,

325 - max_num_mm_items_decoder_budget)

326 - self.max_num_mm_items_by_modality[modality] = max_num_mm_items

289 + self.mm_budget = (MultiModalBudget(

290 + self.model_config,

291 + self.scheduler_config,

292 + self.mm_registry,

293 + max_model_len=self.max_model_len,

294 + max_num_reqs=self.max_num_reqs,

295 + ) if self.is_multimodal_model else None)

296

297 if not self.use_spmd:

298 self.sample_from_logits_func = torch.compile(

@@ -1335,23 +1304,33 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1304 xm.mark_step() # Captures metadata updates

1305

1306 def _precompile_mm_encoder(self) -> None:

1307 + if not self.is_multimodal_model:

1308 + return

1309 +

1310 # Pre-compile MM encoder for all supported data modalities.

1311 hf_config = self.vllm_config.model_config.hf_config

1340 - for mode, max_items_by_mode in \

1341 - self.max_num_mm_items_by_modality.items():

1312 +

1313 + mm_budget = self.mm_budget

1314 + assert mm_budget is not None

1315 +

1316 + max_items_per_seq_by_modality = mm_budget.max_items_per_batch_by_modality # noqa: E501

1317 +

1318 + for mode, max_items_per_seq in max_items_per_seq_by_modality.items():

1319 logger.info(

1320 "Compiling Multimodal %s Encoder with different input"

1321 " shapes.", mode)

1322 start = time.perf_counter()

1323 # No padding for MM encoder just yet.

1347 - for num_items in range(1, max_items_by_mode + 1):

1324 + for num_items in range(1, max_items_per_seq + 1):

1325 logger.info(" -- mode: %s items: %d", mode, num_items)

1326 batched_dummy_mm_inputs = self._get_mm_dummy_batch(

1350 - mode, num_items)

1327 + mode,

1328 + num_items,

1329 + )

1330 # Run multimodal encoder.

1331 xm.mark_step()

1353 - mm_embeds = self.model.\

1354 - get_multimodal_embeddings(**batched_dummy_mm_inputs)

1332 + mm_embeds = self.model.get_multimodal_embeddings(

1333 + **batched_dummy_mm_inputs)

1334 xm.mark_step()

1335 num_patches = mm_embeds[0].shape[0]

1336 items_size = num_patches * num_items

@@ -1547,51 +1526,61 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1526 num_tokens: int,

1527 ) -> None:

1528 # Profile with multimodal encoder & encoder cache.

1550 - # TODO: handle encoder-decoder models once we support them.

1551 - if (self.is_multimodal_model and self.max_num_encoder_input_tokens > 0

1552 - and self.encoder_cache_size > 0):

1553 -

1554 - # NOTE: Currently model is profiled with a single non-text

1555 - # modality with the max possible input tokens even when

1556 - # it supports multiple.

1557 - dummy_data_modality, max_num_mm_items = max(

1558 - self.max_num_mm_items_by_modality.items(), key=lambda t: t[1])

1559 -

1560 - encoder_budget = min(self.max_num_encoder_input_tokens,

1561 - self.encoder_cache_size)

1562 -

1563 - logger.info(

1564 - "Encoder cache will be initialized with a budget of %d tokens,"

1565 - " and profiled with %s %s items of the maximum feature size.",

1566 - encoder_budget, max_num_mm_items, dummy_data_modality)

1567 -

1568 - # Create dummy batch of multimodal inputs.

1569 - batched_dummy_mm_inputs = self._get_mm_dummy_batch(

1570 - dummy_data_modality, max_num_mm_items)

1529 + if self.is_multimodal_model:

1530 + mm_budget = self.mm_budget

1531 + assert mm_budget is not None

1532 +

1533 + # TODO: handle encoder-decoder models once we support them.

1534 + if (encoder_budget := mm_budget.get_encoder_budget()) > 0:

1535 + # NOTE: Currently model is profiled with a single non-text

1536 + # modality with the max possible input tokens even when

1537 + # it supports multiple.

1538 + (

1539 + dummy_modality,

1540 + max_tokens,

1541 + ) = mm_budget.get_modality_with_max_tokens()

1542 + (

1543 + max_mm_items_per_prompt,

1544 + max_mm_items_per_batch,

1545 + ) = mm_budget.get_max_items(dummy_modality, max_tokens)

1546 +

1547 + logger.info(

1548 + "Encoder cache will be initialized with a budget of "

1549 + "%s tokens, and profiled with %s %s items of the maximum "

1550 + "feature size.",

1551 + encoder_budget,

1552 + max_mm_items_per_batch,

1553 + dummy_modality,

1554 + )

1555

1572 - # Run multimodal encoder.

1573 - # Isolate encoder graph from post-processing to minimize

1574 - # impact of recompilation until it's fixed.

1575 - start = time.perf_counter()

1576 - xm.mark_step()

1577 - dummy_encoder_outputs = self.model.get_multimodal_embeddings(

1578 - **batched_dummy_mm_inputs)

1579 - xm.mark_step()

1580 - xm.wait_device_ops()

1581 - end = time.perf_counter()

1582 - logger.info(

1583 - "Multimodal Encoder profiling finished in in %.2f [secs].",

1584 - end - start)

1556 + # Create dummy batch of multimodal inputs.

1557 + batched_dummy_mm_inputs = self._get_mm_dummy_batch(

1558 + dummy_modality,

1559 + max_mm_items_per_batch,

1560 + )

1561

1586 - assert len(dummy_encoder_outputs) == max_num_mm_items, (

1587 - "Expected dimension 0 of encoder outputs to match the number "

1588 - f"of multimodal data items: {max_num_mm_items}, got "

1589 - f"{len(dummy_encoder_outputs)=} instead. This is most likely "

1590 - "due to the 'get_multimodal_embeddings' method of the model "

1591 - "not implemented correctly.")

1562 + # Run multimodal encoder.

1563 + # Isolate encoder graph from post-processing to minimize

1564 + # impact of recompilation until it's fixed.

1565 + start = time.perf_counter()

1566 + xm.mark_step()

1567 + dummy_encoder_outputs = self.model.get_multimodal_embeddings(

1568 + **batched_dummy_mm_inputs)

1569 + xm.mark_step()

1570 + xm.wait_device_ops()

1571 + end = time.perf_counter()

1572 + logger.info(

1573 + "Multimodal Encoder profiling finished in in %.2f [secs].",

1574 + end - start)

1575 +

1576 + sanity_check_mm_encoder_outputs(

1577 + dummy_encoder_outputs,

1578 + expected_num_items=max_mm_items_per_batch,

1579 + )

1580

1593 - # Cache the dummy encoder outputs.

1594 - self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))

1581 + # Cache the dummy encoder outputs.

1582 + self.encoder_cache["tmp"] = dict(

1583 + enumerate(dummy_encoder_outputs))

1584

1585 # Trigger compilation for general shape.

1586 self._dummy_run(num_tokens, self.num_reqs_max_model_len,

@@ -1809,33 +1798,25 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1798 self.grammar_bitmask_cpu[:num_reqs].to(logits.device), \

1799 self.structured_decode_arange.to(logits.device)

1800

1812 - def _get_mm_dummy_batch(self, modality: str,

1813 - batch_size: int) -> BatchedTensorInputs:

1814 - # Dummy data for pre-compiling multimodal models.

1815 - dummy_request_data = self.mm_registry.get_decoder_dummy_data(

1801 + def _get_mm_dummy_batch(

1802 + self,

1803 + modality: str,

1804 + max_items_per_batch: int,

1805 + ) -> BatchedTensorInputs:

1806 + """Dummy data for profiling and precompiling multimodal models."""

1807 + dummy_decoder_data = self.mm_registry.get_decoder_dummy_data(

1808 model_config=self.model_config,

1809 seq_len=self.max_num_tokens,

1810 + mm_counts={modality: 1},

1811 )

1819 - dummy_mm_data = dummy_request_data.multi_modal_data

1820 -

1821 - # Dummy data definition in V0 may contain multiple multimodal items

1822 - # (e.g, multiple images) for a single request, therefore here we

1823 - # always replicate first item by max_num_mm_items times since in V1

1824 - # they are scheduled to be processed separately.

1825 - assert isinstance(dummy_mm_data, MultiModalKwargs), (

1826 - "Expected dummy multimodal data to be of type "

1827 - f"MultiModalKwargs, got {type(dummy_mm_data)=} instead. "

1828 - "This is most likely due to the model not having a merged "

1829 - "processor.")

1830 -

1831 - # When models have a merged processor, their dummy data is

1832 - # already batched `MultiModalKwargs`, therefore we take the first

1833 - # `MultiModalKwargsItem` from the desired modality to profile on.

1812 + dummy_mm_data = dummy_decoder_data.multi_modal_data

1813 +

1814 + # Result in the maximum GPU consumption of the model

1815 dummy_mm_item = dummy_mm_data.get_item(modality=modality, item_index=0)

1816 dummy_mm_kwargs = MultiModalKwargs.from_items([dummy_mm_item])

1817

1818 batched_dummy_mm_inputs = MultiModalKwargs.batch([dummy_mm_kwargs] *

1838 - batch_size)

1819 + max_items_per_batch)

1820 return MultiModalKwargs.as_kwargs(

1821 batched_dummy_mm_inputs,

1822 device=self.device,

vllm/v1/worker/utils.py

+109

@@ -5,14 +5,123 @@ from typing import TYPE_CHECKING, Optional

6 import torch

8 +from vllm.config import ModelConfig, SchedulerConfig

9 from vllm.model_executor.models.interfaces import MultiModalEmbeddings

10 from vllm.model_executor.models.utils import extract_layer_index

11 +from vllm.multimodal.registry import MultiModalRegistry

12 +from vllm.v1.core.encoder_cache_manager import compute_encoder_budget

13 from vllm.v1.kv_cache_interface import KVCacheGroupSpec

15 if TYPE_CHECKING:

16 from vllm.attention.layer import Attention

19 +class MultiModalBudget:

20 + """Helper class to calculate budget information for multi-modal models."""

21 +

22 + def __init__(

23 + self,

24 + model_config: ModelConfig,

25 + scheduler_config: SchedulerConfig,

26 + mm_registry: MultiModalRegistry,

27 + *,

28 + max_model_len: int,

29 + max_num_reqs: int,

30 + ) -> None:

31 + super().__init__()

32 +

33 + self.model_config = model_config

34 + self.scheduler_config = scheduler_config

35 + self.mm_registry = mm_registry

36 +

37 + encoder_compute_budget, encoder_cache_size = compute_encoder_budget(

38 + model_config=model_config,

39 + scheduler_config=scheduler_config,

40 + mm_registry=mm_registry,

41 + )

42 +

43 + self.max_num_encoder_input_tokens = encoder_compute_budget

44 + self.encoder_cache_size = encoder_cache_size

45 + self.max_model_len = max_model_len

46 + self.max_num_reqs = max_num_reqs

47 +

48 + self.mm_limits = mm_registry.get_mm_limits_per_prompt(model_config)

49 +

50 + max_items_per_prompt_by_modality = dict[str, int]()

51 + max_items_per_batch_by_modality = dict[str, int]()

52 +

53 + max_tokens_by_modality = mm_registry \

54 + .get_max_tokens_per_item_by_nonzero_modality(model_config)

55 +

56 + for modality, max_tokens in max_tokens_by_modality.items():

57 + (

58 + max_items_per_prompt,

59 + max_items_per_batch,

60 + ) = self.get_max_items(modality, max_tokens)

61 +

62 + max_items_per_prompt_by_modality[modality] = max_items_per_prompt

63 + max_items_per_batch_by_modality[modality] = max_items_per_batch

64 +

65 + self.max_tokens_by_modality = max_tokens_by_modality

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 +

69 + def get_modality_with_max_tokens(self) -> tuple[str, int]:

70 + max_tokens_by_modality = self.max_tokens_by_modality

71 + modality, max_tokens = max(max_tokens_by_modality.items(),

72 + key=lambda item: item[1])

73 +

74 + return modality, max_tokens

75 +

76 + def get_encoder_budget(self) -> int:

77 + return min(self.max_num_encoder_input_tokens, self.encoder_cache_size)

78 +

79 + def get_max_items(

80 + self,

81 + modality: str,

82 + max_tokens_per_item: int,

83 + ) -> tuple[int, int]:

84 + if max_tokens_per_item == 0:

85 + return 0, 0

86 +

87 + # Check how many items of this modality can be supported by

88 + # the encoder budget.

89 + encoder_budget = self.get_encoder_budget()

90 +

91 + # TODO: handle encoder-decoder models once we support them.

92 + if encoder_budget == 0:

93 + return 0, 0

94 +

95 + max_encoder_items_per_batch = encoder_budget // max_tokens_per_item

96 +

97 + # Check how many items of this modality can be supported by

98 + # the decoder budget.

99 + mm_limit = self.mm_limits[modality]

100 +

101 + max_items_per_prompt = max(

102 + 1,

103 + min(mm_limit, self.max_model_len // max_tokens_per_item),

104 + )

105 +

106 + scheduler_config = self.scheduler_config

107 + max_num_reqs = self.max_num_reqs

108 +

109 + if not scheduler_config.enable_chunked_prefill:

110 + max_num_reqs = min(

111 + max_num_reqs,

112 + scheduler_config.max_num_batched_tokens // max_tokens_per_item,

113 + )

114 +

115 + max_decoder_items_per_batch = max_num_reqs * max_items_per_prompt

116 +

117 + max_items_per_batch = max(

118 + 1,

119 + min(max_encoder_items_per_batch, max_decoder_items_per_batch),

120 + )

121 +

122 + return max_items_per_prompt, max_items_per_batch

123 +

124 +

125 def sanity_check_mm_encoder_outputs(

126 mm_embeddings: MultiModalEmbeddings,

127 expected_num_items: int,