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
[MM] Allow skipping memory profiling for multimodal models. (#22950)
Signed-off-by: Roger Wang <hey@rogerw.me> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Roger Wang committed
Aug 15, 2025 at 04:41 UTC
49252cf59e70b1e1a8bae21da929f6d51e9acce4
4 files changed
+121
-90
vllm/config/__init__.py
+16
-1
@@ -388,6 +388,10 @@ class ModelConfig:
388
interleave_mm_strings: bool = False
389
"""Enable fully interleaved support for multimodal prompts, while using
390
--chat-template-content-format=string. Defaults to False."""
391
+ skip_mm_profiling: bool = False
392
+ """When enabled, skips multimodal memory profiling and only profiles with
393
+ language backbone model during engine initialization.
394
+ """
395
media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
396
"""Additional args passed to process media inputs, keyed by modalities.
397
For example, to set num_frames for video, set
@@ -837,7 +841,8 @@ class ModelConfig:
841
media_io_kwargs=self.media_io_kwargs,
842
mm_processor_kwargs=self.mm_processor_kwargs,
843
mm_processor_cache_gb=self.mm_processor_cache_gb,
840
- interleave_mm_strings=self.interleave_mm_strings)
844
+ interleave_mm_strings=self.interleave_mm_strings,
845
+ skip_mm_profiling=self.skip_mm_profiling)
846
847
return None
848
@@ -2511,6 +2516,16 @@ class MultiModalConfig:
2516
Enable fully interleaved support for multimodal prompts.
2517
"""
2518
2519
+ skip_mm_profiling: bool = False
2520
+ """
2521
+ When enabled, skips multimodal memory profiling and only profiles with
2522
+ language backbone model during engine initialization.
2523
+
2524
+ This reduces engine startup time but shifts the responsibility to users for
2525
+ estimating the peak memory usage of the activation of multimodal encoder and
2526
+ embedding cache.
2527
+ """
2528
+
2529
def compute_hash(self) -> str:
2530
"""
2531
WARNING: Whenever a new field is added to this config,
vllm/engine/arg_utils.py
+4
@@ -350,6 +350,7 @@ class EngineArgs:
350
MultiModalConfig.mm_processor_kwargs
351
disable_mm_preprocessor_cache: bool = False # DEPRECATED
352
mm_processor_cache_gb: int = MultiModalConfig.mm_processor_cache_gb
353
+ skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
354
# LoRA fields
355
enable_lora: bool = False
356
enable_lora_bias: bool = LoRAConfig.bias_enabled
@@ -716,6 +717,8 @@ class EngineArgs:
717
multimodal_group.add_argument(
718
"--interleave-mm-strings",
719
**multimodal_kwargs["interleave_mm_strings"])
720
+ multimodal_group.add_argument("--skip-mm-profiling",
721
+ **multimodal_kwargs["skip_mm_profiling"])
722
723
# LoRA related configs
724
lora_kwargs = get_kwargs(LoRAConfig)
@@ -918,6 +921,7 @@ class EngineArgs:
921
limit_mm_per_prompt=self.limit_mm_per_prompt,
922
interleave_mm_strings=self.interleave_mm_strings,
923
media_io_kwargs=self.media_io_kwargs,
924
+ skip_mm_profiling=self.skip_mm_profiling,
925
use_async_output_proc=not self.disable_async_output_proc,
926
config_format=self.config_format,
927
mm_processor_kwargs=self.mm_processor_kwargs,
vllm/v1/worker/gpu_model_runner.py
+45
-39
@@ -2479,50 +2479,56 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2479
def profile_run(self) -> None:
2480
# Profile with multimodal encoder & encoder cache.
2481
if self.supports_mm_inputs:
2482
- mm_budget = self.mm_budget
2483
- assert mm_budget is not None
2484
-
2485
- # TODO: handle encoder-decoder models once we support them.
2486
- if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
2487
- # NOTE: Currently model is profiled with a single non-text
2488
- # modality with the max possible input tokens even when
2489
- # it supports multiple.
2490
- (
2491
- dummy_modality,
2492
- max_tokens,
2493
- ) = mm_budget.get_modality_with_max_tokens()
2494
- (
2495
- max_mm_items_per_prompt,
2496
- max_mm_items_per_batch,
2497
- ) = mm_budget.get_max_items(dummy_modality, max_tokens)
2498
-
2482
+ if self.model_config.multimodal_config.skip_mm_profiling:
2483
logger.info(
2500
- "Encoder cache will be initialized with a budget of "
2501
- "%s tokens, and profiled with %s %s items of the maximum "
2502
- "feature size.",
2503
- encoder_budget,
2504
- max_mm_items_per_batch,
2505
- dummy_modality,
2506
- )
2484
+ "Skipping memory profiling for multimodal encoder and "
2485
+ "encoder cache.")
2486
+ else:
2487
+ mm_budget = self.mm_budget
2488
+ assert mm_budget is not None
2489
+
2490
+ # TODO: handle encoder-decoder models once we support them.
2491
+ if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
2492
+ # NOTE: Currently model is profiled with a single non-text
2493
+ # modality with the max possible input tokens even when
2494
+ # it supports multiple.
2495
+ (
2496
+ dummy_modality,
2497
+ max_tokens,
2498
+ ) = mm_budget.get_modality_with_max_tokens()
2499
+ (
2500
+ max_mm_items_per_prompt,
2501
+ max_mm_items_per_batch,
2502
+ ) = mm_budget.get_max_items(dummy_modality, max_tokens)
2503
+
2504
+ logger.info(
2505
+ "Encoder cache will be initialized with a budget of "
2506
+ "%s tokens, and profiled with %s %s items of the "
2507
+ "maximum feature size.",
2508
+ encoder_budget,
2509
+ max_mm_items_per_batch,
2510
+ dummy_modality,
2511
+ )
2512
2508
- # Create dummy batch of multimodal inputs.
2509
- batched_dummy_mm_inputs = self._get_mm_dummy_batch(
2510
- dummy_modality,
2511
- max_mm_items_per_batch,
2512
- )
2513
+ # Create dummy batch of multimodal inputs.
2514
+ batched_dummy_mm_inputs = self._get_mm_dummy_batch(
2515
+ dummy_modality,
2516
+ max_mm_items_per_batch,
2517
+ )
2518
2514
- # Run multimodal encoder.
2515
- dummy_encoder_outputs = self.model.get_multimodal_embeddings(
2516
- **batched_dummy_mm_inputs)
2519
+ # Run multimodal encoder.
2520
+ dummy_encoder_outputs = \
2521
+ self.model.get_multimodal_embeddings(
2522
+ **batched_dummy_mm_inputs)
2523
2518
- sanity_check_mm_encoder_outputs(
2519
- dummy_encoder_outputs,
2520
- expected_num_items=max_mm_items_per_batch,
2521
- )
2524
+ sanity_check_mm_encoder_outputs(
2525
+ dummy_encoder_outputs,
2526
+ expected_num_items=max_mm_items_per_batch,
2527
+ )
2528
2523
- # Cache the dummy encoder outputs.
2524
- self.encoder_cache["tmp"] = dict(
2525
- enumerate(dummy_encoder_outputs))
2529
+ # Cache the dummy encoder outputs.
2530
+ self.encoder_cache["tmp"] = dict(
2531
+ enumerate(dummy_encoder_outputs))
2532
2533
# Add `is_profile` here to pre-allocate communication buffers
2534
hidden_states, last_hidden_states \
vllm/v1/worker/tpu_model_runner.py
+56
-50
@@ -1529,60 +1529,66 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1529
) -> None:
1530
# Profile with multimodal encoder & encoder cache.
1531
if self.supports_mm_inputs:
1532
- mm_budget = self.mm_budget
1533
- assert mm_budget is not None
1534
-
1535
- # TODO: handle encoder-decoder models once we support them.
1536
- if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
1537
- # NOTE: Currently model is profiled with a single non-text
1538
- # modality with the max possible input tokens even when
1539
- # it supports multiple.
1540
- (
1541
- dummy_modality,
1542
- max_tokens,
1543
- ) = mm_budget.get_modality_with_max_tokens()
1544
- (
1545
- max_mm_items_per_prompt,
1546
- max_mm_items_per_batch,
1547
- ) = mm_budget.get_max_items(dummy_modality, max_tokens)
1548
-
1532
+ if self.model_config.multimodal_config.skip_mm_profiling:
1533
logger.info(
1550
- "Encoder cache will be initialized with a budget of "
1551
- "%s tokens, and profiled with %s %s items of the maximum "
1552
- "feature size.",
1553
- encoder_budget,
1554
- max_mm_items_per_batch,
1555
- dummy_modality,
1556
- )
1557
-
1558
- # Create dummy batch of multimodal inputs.
1559
- batched_dummy_mm_inputs = self._get_mm_dummy_batch(
1560
- dummy_modality,
1561
- max_mm_items_per_batch,
1562
- )
1534
+ "Skipping memory profiling for multimodal encoder and "
1535
+ "encoder cache.")
1536
+ else:
1537
+ mm_budget = self.mm_budget
1538
+ assert mm_budget is not None
1539
+
1540
+ # TODO: handle encoder-decoder models once we support them.
1541
+ if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
1542
+ # NOTE: Currently model is profiled with a single non-text
1543
+ # modality with the max possible input tokens even when
1544
+ # it supports multiple.
1545
+ (
1546
+ dummy_modality,
1547
+ max_tokens,
1548
+ ) = mm_budget.get_modality_with_max_tokens()
1549
+ (
1550
+ max_mm_items_per_prompt,
1551
+ max_mm_items_per_batch,
1552
+ ) = mm_budget.get_max_items(dummy_modality, max_tokens)
1553
+
1554
+ logger.info(
1555
+ "Encoder cache will be initialized with a budget of "
1556
+ "%s tokens, and profiled with %s %s items of the "
1557
+ "maximum feature size.",
1558
+ encoder_budget,
1559
+ max_mm_items_per_batch,
1560
+ dummy_modality,
1561
+ )
1562
1564
- # Run multimodal encoder.
1565
- # Isolate encoder graph from post-processing to minimize
1566
- # impact of recompilation until it's fixed.
1567
- start = time.perf_counter()
1568
- xm.mark_step()
1569
- dummy_encoder_outputs = self.model.get_multimodal_embeddings(
1570
- **batched_dummy_mm_inputs)
1571
- xm.mark_step()
1572
- xm.wait_device_ops()
1573
- end = time.perf_counter()
1574
- logger.info(
1575
- "Multimodal Encoder profiling finished in in %.2f [secs].",
1576
- end - start)
1563
+ # Create dummy batch of multimodal inputs.
1564
+ batched_dummy_mm_inputs = self._get_mm_dummy_batch(
1565
+ dummy_modality,
1566
+ max_mm_items_per_batch,
1567
+ )
1568
1578
- sanity_check_mm_encoder_outputs(
1579
- dummy_encoder_outputs,
1580
- expected_num_items=max_mm_items_per_batch,
1581
- )
1569
+ # Run multimodal encoder.
1570
+ # Isolate encoder graph from post-processing to minimize
1571
+ # impact of recompilation until it's fixed.
1572
+ start = time.perf_counter()
1573
+ xm.mark_step()
1574
+ dummy_encoder_outputs = \
1575
+ self.model.get_multimodal_embeddings(
1576
+ **batched_dummy_mm_inputs)
1577
+ xm.mark_step()
1578
+ xm.wait_device_ops()
1579
+ end = time.perf_counter()
1580
+ logger.info(
1581
+ "Multimodal Encoder profiling finished in %.2f [secs].",
1582
+ end - start)
1583
+
1584
+ sanity_check_mm_encoder_outputs(
1585
+ dummy_encoder_outputs,
1586
+ expected_num_items=max_mm_items_per_batch,
1587
+ )
1588
1583
- # Cache the dummy encoder outputs.
1584
- self.encoder_cache["tmp"] = dict(
1585
- enumerate(dummy_encoder_outputs))
1589
+ # Cache the dummy encoder outputs.
1590
+ self.encoder_cache["tmp"] = dict(
1591
+ enumerate(dummy_encoder_outputs))
1592
1593
# Trigger compilation for general shape.
1594
self._dummy_run(num_tokens, self.num_reqs_max_model_len,