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 2, 22:00 UTC
[Bugfix] Auto-raise max_num_batched_tokens for prefix-LM multimodal models (#43051)
Signed-off-by: Ashwin Giridharan <girida@amazon.com> Co-authored-by: abinggo <107740309+abinggo@users.noreply.github.com>
Ashwin Giridharan committed
May 22, 2026 at 21:23 UTC
84e351555aac8f80e4bf03fc77d4f5e89cb36de5
2 files changed
+88
tests/v1/engine/test_engine_args.py
+38
@@ -90,3 +90,41 @@ def test_defaults_with_usage_context():
90
vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
91
assert vllm_config.scheduler_config.max_num_seqs == default_max_num_seqs
92
assert vllm_config.scheduler_config.max_num_batched_tokens == default_server_tokens # noqa: E501
93
+
94
+
95
+def test_mm_prefix_lm_raises_batched_tokens_floor():
96
+ """Verify that prefix-LM multimodal models auto-raise
97
+ max_num_batched_tokens to fit at least one multimodal item.
98
+
99
+ Regression test for https://github.com/vllm-project/vllm/issues/42687
100
+ """
101
+ from unittest.mock import patch
102
+
103
+ # Simulate a prefix-LM multimodal model whose largest modality
104
+ # (video) requires 2496 tokens — more than the 2048 default.
105
+ fake_mm_min = (2496, "video")
106
+
107
+ engine_args = EngineArgs(
108
+ model="facebook/opt-125m",
109
+ max_model_len=2048,
110
+ enforce_eager=True,
111
+ )
112
+
113
+ with (
114
+ patch.object(
115
+ type(engine_args),
116
+ "_get_min_mm_batched_tokens",
117
+ staticmethod(lambda _mc: fake_mm_min),
118
+ ),
119
+ patch(
120
+ "vllm.config.ModelConfig.is_multimodal_model",
121
+ new_callable=lambda: property(lambda self: True),
122
+ ),
123
+ patch(
124
+ "vllm.config.ModelConfig.is_mm_prefix_lm",
125
+ new_callable=lambda: property(lambda self: True),
126
+ ),
127
+ ):
128
+ vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
129
+
130
+ assert vllm_config.scheduler_config.max_num_batched_tokens >= 2496
vllm/engine/arg_utils.py
+50
@@ -2439,6 +2439,39 @@ class EngineArgs:
2439
self.reasoning_config = ReasoningConfig()
2440
self.reasoning_config.reasoning_parser = self.reasoning_parser
2441
2442
+ @staticmethod
2443
+ def _get_min_mm_batched_tokens(
2444
+ model_config: ModelConfig,
2445
+ ) -> tuple[int, str] | None:
2446
+ """Get the minimum max_num_batched_tokens needed for a multimodal
2447
+ prefix-LM model to process at least one item of any supported modality.
2448
+
2449
+ Returns (token_count, modality_name) for the most expensive modality,
2450
+ or None if the value cannot be determined at this stage.
2451
+ """
2452
+ try:
2453
+ from vllm.multimodal import MULTIMODAL_REGISTRY
2454
+
2455
+ # get_processing_info returns the model's multimodal processing
2456
+ # metadata (supported modalities, token limits) without loading
2457
+ # model weights or generating dummy data.
2458
+ info = MULTIMODAL_REGISTRY.get_processing_info(model_config)
2459
+ mm_counts = {modality: 1 for modality in info.supported_mm_limits}
2460
+ # get_mm_max_tokens_per_item returns pre-computed per-item token
2461
+ # ceilings for models that override it (e.g., Gemma4), or None
2462
+ # for models that rely on dummy-input profiling. When None is
2463
+ # returned we bail out — no dummy generation is triggered here.
2464
+ max_tokens = info.get_mm_max_tokens_per_item(
2465
+ seq_len=model_config.max_model_len,
2466
+ mm_counts=mm_counts,
2467
+ )
2468
+ if max_tokens is not None:
2469
+ modality = max(max_tokens, key=max_tokens.__getitem__)
2470
+ return (max_tokens[modality], modality)
2471
+ except Exception as e:
2472
+ logger.warning("Failed to determine min multimodal batched tokens: %s", e)
2473
+ return None
2474
+
2475
def _set_default_max_num_seqs_and_batched_tokens_args(
2476
self,
2477
usage_context: UsageContext | None,
@@ -2489,6 +2522,23 @@ class EngineArgs:
2522
self.max_num_batched_tokens,
2523
)
2524
2525
+ # For multimodal prefix-LM models (e.g., Gemma 4) that disable
2526
+ # chunked MM input, a single multimodal item must fit in one batch.
2527
+ # Raise the floor to accommodate the largest per-item token count.
2528
+ if model_config.is_multimodal_model and model_config.is_mm_prefix_lm:
2529
+ result = self._get_min_mm_batched_tokens(model_config)
2530
+ if result is not None and result[0] > self.max_num_batched_tokens:
2531
+ mm_min, modality = result
2532
+ logger.info(
2533
+ "Raising max_num_batched_tokens from %d to %d to "
2534
+ "accommodate '%s' input for prefix-LM model %s.",
2535
+ self.max_num_batched_tokens,
2536
+ mm_min,
2537
+ modality,
2538
+ model_config.model,
2539
+ )
2540
+ self.max_num_batched_tokens = mm_min
2541
+
2542
# When using default settings,
2543
# Ensure max_num_batched_tokens does not exceed model limit.
2544
# Some models (e.g., Whisper) have embeddings tied to max length.