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 8, 16:30 UTC
Add embedding input functionality for disabled modalities [remake] (#32493)
Signed-off-by: Reagan Lee <“reaganjlee@gmail.com”> Signed-off-by: Reagan Lee <reaganjlee@gmail.com> Signed-off-by: Reagan Lee <96998476+reaganjlee@users.noreply.github.com> Co-authored-by: Reagan Lee <“reaganjlee@gmail.com”> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Reagan Lee committed
Feb 8, 2026 at 04:57 UTC
c4df59ad43037a846eed353ce4c17dc264d18f4a
10 files changed
+226
-77
tests/entrypoints/llm/test_mm_embeds_only.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import weakref
5
+
6
+import pytest
7
+
8
+from vllm import LLM, SamplingParams
9
+from vllm.assets.image import ImageAsset
10
+from vllm.distributed import cleanup_dist_env_and_memory
11
+
12
+MODEL = "llava-hf/llava-1.5-7b-hf"
13
+PROMPT = "USER: <image>\nDescribe this image briefly.\nASSISTANT:"
14
+TEXT_ONLY_PROMPT = "USER: What is 2 + 2?\nASSISTANT:"
15
+
16
+
17
+@pytest.fixture(scope="module")
18
+def llm():
19
+ """LLM with enable_mm_embeds=True and all modality limits zeroed out."""
20
+ llm = LLM(
21
+ model=MODEL,
22
+ max_model_len=2048,
23
+ enforce_eager=True,
24
+ gpu_memory_utilization=0.8,
25
+ enable_mm_embeds=True,
26
+ limit_mm_per_prompt={"image": 0},
27
+ )
28
+
29
+ yield weakref.proxy(llm)
30
+
31
+ del llm
32
+
33
+ cleanup_dist_env_and_memory()
34
+
35
+
36
+@pytest.mark.skip_global_cleanup
37
+def test_generate_with_embedding(llm: LLM):
38
+ """Pre-computed embedding produces tokens without hanging."""
39
+ embedding = ImageAsset("stop_sign").image_embeds
40
+ outputs = llm.generate(
41
+ {"prompt": PROMPT, "multi_modal_data": {"image": embedding}},
42
+ sampling_params=SamplingParams(max_tokens=32, temperature=0.0),
43
+ )
44
+ assert len(outputs) == 1
45
+ assert len(outputs[0].outputs[0].text) > 0
46
+
47
+
48
+@pytest.mark.skip_global_cleanup
49
+def test_raw_image_rejected(llm: LLM):
50
+ """Raw image input is still rejected when limit=0."""
51
+ raw_image = ImageAsset("stop_sign").pil_image
52
+ with pytest.raises(ValueError, match=r"At most 0 image\(s\)"):
53
+ llm.generate(
54
+ {"prompt": PROMPT, "multi_modal_data": {"image": raw_image}},
55
+ sampling_params=SamplingParams(max_tokens=16),
56
+ )
57
+
58
+
59
+@pytest.mark.skip_global_cleanup
60
+def test_text_only_prompt(llm: LLM):
61
+ """Text-only prompts still work under this config."""
62
+ outputs = llm.generate(
63
+ TEXT_ONLY_PROMPT,
64
+ sampling_params=SamplingParams(max_tokens=16, temperature=0.0),
65
+ )
66
+ assert len(outputs) == 1
67
+ assert len(outputs[0].outputs[0].text) > 0
tests/multimodal/test_processing.py
+42
-32
@@ -901,7 +901,7 @@ def test_find_mm_placeholders(
901
902
@pytest.mark.parametrize("model_id", ["llava-hf/llava-v1.6-mistral-7b-hf"])
903
@pytest.mark.parametrize(
904
- ("limit", "num_supported", "is_valid"),
904
+ ("num_images", "limit", "is_valid"),
905
[
906
(0, 0, True),
907
(0, 1, True),
@@ -912,7 +912,7 @@ def test_find_mm_placeholders(
912
(2, 2, True),
913
],
914
)
915
-def test_limit_mm_per_prompt_dummy(model_id, limit, num_supported, is_valid):
915
+def test_limit_mm_per_prompt_apply(model_id, num_images, limit, is_valid):
916
limit_mm_per_prompt = {"image": limit}
917
918
model_config = ModelConfig(
@@ -921,33 +921,46 @@ def test_limit_mm_per_prompt_dummy(model_id, limit, num_supported, is_valid):
921
)
922
923
processor = MULTIMODAL_REGISTRY.create_processor(model_config)
924
- processor.info.get_supported_mm_limits = lambda: {"image": num_supported}
924
+
925
+ rng = np.random.RandomState(0)
926
+ image = random_image(rng, min_wh=128, max_wh=256)
927
+ if num_images == 0:
928
+ mm_data = {}
929
+ elif num_images == 1:
930
+ mm_data = {"image": image}
931
+ else:
932
+ mm_data = {"image": [image] * num_images}
933
934
exc_ctx = nullcontext() if is_valid else pytest.raises(ValueError, match="At most")
935
936
with exc_ctx:
929
- MULTIMODAL_REGISTRY.get_dummy_mm_inputs(
930
- model_config,
931
- mm_counts=limit_mm_per_prompt,
932
- processor=processor,
937
+ processor.apply(
938
+ "<image>" * num_images,
939
+ mm_items=processor.info.parse_mm_data(mm_data),
940
+ hf_processor_mm_kwargs={},
941
)
942
943
944
@pytest.mark.parametrize("model_id", ["llava-hf/llava-v1.6-mistral-7b-hf"])
945
@pytest.mark.parametrize(
938
- ("num_images", "limit", "is_valid"),
946
+ ("user_limit", "supported_limit"),
947
[
940
- (0, 0, True),
941
- (0, 1, True),
942
- (1, 0, False),
943
- (1, 1, True),
944
- (1, 2, True),
945
- (2, 1, False),
946
- (2, 2, True),
948
+ (0, 0),
949
+ (0, 1),
950
+ (1, 0), # user wants 1, model supports 0 → capped to 0
951
+ (1, 1),
952
+ (1, 2),
953
+ (2, 1), # user wants 2, model supports 1 → capped to 1
954
+ (2, 2),
955
+ (5, 1), # large user limit, low model support → capped to 1
956
+ (1, 5),
957
+ (10, 0), # large user limit, no model support → capped to 0
958
],
959
)
949
-def test_limit_mm_per_prompt_apply(model_id, num_images, limit, is_valid):
950
- limit_mm_per_prompt = {"image": limit}
960
+def test_budget_caps_prevent_dummy_input_validation_failure(
961
+ model_id, user_limit, supported_limit
962
+):
963
+ limit_mm_per_prompt = {"image": user_limit}
964
965
model_config = ModelConfig(
966
model=model_id,
@@ -955,24 +968,21 @@ def test_limit_mm_per_prompt_apply(model_id, num_images, limit, is_valid):
968
)
969
970
processor = MULTIMODAL_REGISTRY.create_processor(model_config)
971
+ processor.info.get_supported_mm_limits = lambda: {"image": supported_limit}
972
959
- rng = np.random.RandomState(0)
960
- image = random_image(rng, min_wh=128, max_wh=256)
961
- if num_images == 0:
962
- mm_data = {}
963
- elif num_images == 1:
964
- mm_data = {"image": image}
965
- else:
966
- mm_data = {"image": [image] * num_images}
973
+ # This is what budget.py uses to derive mm_counts
974
+ allowed = processor.info.allowed_mm_limits
975
968
- exc_ctx = nullcontext() if is_valid else pytest.raises(ValueError, match="At most")
976
+ assert allowed["image"] <= supported_limit, (
977
+ f"allowed_mm_limits['image']={allowed['image']} exceeds "
978
+ f"supported_limit={supported_limit}"
979
+ )
980
970
- with exc_ctx:
971
- processor.apply(
972
- "<image>" * num_images,
973
- mm_items=processor.info.parse_mm_data(mm_data),
974
- hf_processor_mm_kwargs={},
975
- )
981
+ assert allowed["image"] <= user_limit, (
982
+ f"allowed_mm_limits['image']={allowed['image']} exceeds user_limit={user_limit}"
983
+ )
984
+
985
+ assert allowed["image"] == min(user_limit, supported_limit)
986
987
988
class DummyProcessor:
vllm/config/multimodal.py
+5
@@ -76,6 +76,11 @@ class MultiModalConfig:
76
for the OpenAI-compatible server, this refers to chat messages with content
77
`"type": "*_embeds"`.
78
79
+ When enabled with `--limit-mm-per-prompt` set to 0 for a modality,
80
+ precomputed embeddings skip count validation for that modality,
81
+ saving memory by not loading encoder modules while still enabling
82
+ embeddings as an input. Limits greater than 0 still apply to embeddings.
83
+
84
WARNING: The vLLM engine may crash if incorrect shape of embeddings is passed.
85
Only enable this flag for trusted users!"""
86
media_io_kwargs: dict[str, dict[str, Any]] = Field(default_factory=dict)
vllm/entrypoints/chat_utils.py
+11
-1
@@ -528,7 +528,17 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
528
else:
529
num_items = len(self._items_by_modality[original_modality]) + 1
530
531
- self.mm_processor.info.validate_num_items(input_modality, num_items)
531
+ mm_config = self.model_config.multimodal_config
532
+ if (
533
+ mm_config is not None
534
+ and mm_config.enable_mm_embeds
535
+ and mm_config.get_limit_per_prompt(input_modality) == 0
536
+ and original_modality.endswith("_embeds")
537
+ ):
538
+ # Skip validation: embeddings bypass limit when enable_mm_embeds=True
539
+ pass
540
+ else:
541
+ self.mm_processor.info.validate_num_items(input_modality, num_items)
542
543
# Track original modality for vision_chunk items
544
if use_vision_chunk:
vllm/multimodal/budget.py
+39
-6
@@ -3,11 +3,14 @@
3
from collections.abc import Mapping
4
5
from vllm.config import ModelConfig, VllmConfig
6
+from vllm.logger import init_logger
7
from vllm.multimodal.processing import BaseMultiModalProcessor
8
from vllm.multimodal.registry import MultiModalRegistry
9
from vllm.utils.torch_utils import set_default_torch_num_threads
10
from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
11
12
+logger = init_logger(__name__)
13
+
14
15
def get_mm_max_toks_per_item(
16
model_config: ModelConfig,
@@ -59,11 +62,26 @@ class MultiModalBudget:
62
processor = mm_registry.create_processor(model_config, cache=cache)
63
64
self.cache = cache
65
+ mm_config = model_config.get_multimodal_config()
66
+ enable_mm_embeds = mm_config is not None and mm_config.enable_mm_embeds
67
+
68
+ supported_mm_limits = processor.info.supported_mm_limits
69
self.mm_limits = mm_limits = processor.info.allowed_mm_limits
70
64
- active_modalities = {
65
- modality for modality, limit in mm_limits.items() if limit > 0
71
+ # Modalities that pass through the MM encoder tower
72
+ tower_modalities = {
73
+ modality
74
+ for modality in supported_mm_limits
75
+ if mm_limits.get(modality, 0) > 0
76
}
77
+ # Modalities that bypass the tower (pre-computed embeddings only)
78
+ embed_only_modalities = {
79
+ modality
80
+ for modality in supported_mm_limits
81
+ if enable_mm_embeds and mm_limits.get(modality, 0) == 0
82
+ }
83
+
84
+ active_modalities = tower_modalities | embed_only_modalities
85
86
all_mm_max_toks_per_item = get_mm_max_toks_per_item(
87
model_config,
@@ -72,19 +90,32 @@ class MultiModalBudget:
90
mm_counts=dict.fromkeys(active_modalities, 1),
91
)
92
93
+ if embed_only_modalities:
94
+ logger.info_once(
95
+ "enable_mm_embeds is True; modalities handled as embedding-only: %s",
96
+ tuple(embed_only_modalities),
97
+ )
98
+
99
# Some models (e.g., Qwen3Omni with use_audio_in_video=True) share
100
# placeholders between modalities, so not all active modalities will
101
# have their own entry in the returned dict. We filter to only include
102
# modalities that have independent placeholder tokens.
79
- mm_max_toks_per_item = {
103
+ active_mm_max_toks_per_item = {
104
modality: all_mm_max_toks_per_item[modality]
105
for modality in active_modalities
106
if modality in all_mm_max_toks_per_item
107
}
108
+ tower_mm_max_toks_per_item = {
109
+ modality: active_mm_max_toks_per_item[modality]
110
+ for modality in tower_modalities
111
+ if modality in active_mm_max_toks_per_item
112
+ }
113
114
+ # Encoder budget is computed from all active modalities (including
115
+ # embedding-only ones that need encoder cache space).
116
encoder_compute_budget, encoder_cache_size = compute_mm_encoder_budget(
117
scheduler_config,
87
- mm_max_toks_per_item,
118
+ active_mm_max_toks_per_item,
119
)
120
121
self.encoder_compute_budget = encoder_compute_budget
@@ -93,13 +124,15 @@ class MultiModalBudget:
124
mm_max_items_per_prompt = dict[str, int]()
125
mm_max_items_per_batch = dict[str, int]()
126
96
- for modality, max_toks_per_item in mm_max_toks_per_item.items():
127
+ # Per-prompt/per-batch limits are only relevant for tower modalities
128
+ # (embedding-only modalities don't go through the encoder tower).
129
+ for modality, max_toks_per_item in tower_mm_max_toks_per_item.items():
130
(
131
mm_max_items_per_prompt[modality],
132
mm_max_items_per_batch[modality],
133
) = self._get_max_items(modality, max_toks_per_item)
134
102
- self.mm_max_toks_per_item = mm_max_toks_per_item
135
+ self.mm_max_toks_per_item = tower_mm_max_toks_per_item
136
self.mm_max_items_per_prompt: Mapping[str, int] = mm_max_items_per_prompt
137
self.mm_max_items_per_batch: Mapping[str, int] = mm_max_items_per_batch
138
vllm/multimodal/processing/context.py
+12
-6
@@ -681,16 +681,22 @@ class BaseProcessingInfo:
681
mm_items = self.data_parser.parse_mm_data(mm_data)
682
683
if validate:
684
- mm_config = self.ctx.model_config.get_multimodal_config()
685
- if not mm_config.enable_mm_embeds:
686
- for modality, items in mm_items.items():
687
- if isinstance(items, (EmbeddingItems, DictEmbeddingItems)):
684
+ mm_config = self.ctx.get_mm_config()
685
+
686
+ for modality, items in mm_items.items():
687
+ if isinstance(items, (EmbeddingItems, DictEmbeddingItems)):
688
+ if not mm_config.enable_mm_embeds:
689
raise ValueError(
690
f"You must set `--enable-mm-embeds` to input "
691
f"`{modality}_embeds`"
692
)
692
-
693
- for modality, items in mm_items.items():
693
+ if mm_config.get_limit_per_prompt(modality) == 0:
694
+ logger.debug(
695
+ "Skipping count validation for modality "
696
+ "'%s' (embeddings with limit=0)",
697
+ modality,
698
+ )
699
+ continue
700
self.validate_num_items(modality, len(items))
701
702
return mm_items
vllm/multimodal/processing/dummy_inputs.py
+1
-1
@@ -95,7 +95,7 @@ class BaseDummyInputsBuilder(ABC, Generic[_I]):
95
"""
96
dummy_text = self.get_dummy_text(mm_counts)
97
dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
98
- dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
98
+ dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False)
99
100
tokenization_kwargs = {"truncation": False}
101
vllm/multimodal/processing/processor.py
+1
-1
@@ -1395,7 +1395,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
1395
missing_modality_data.append(data)
1396
mm_missing_data[modality] = missing_modality_data
1397
1398
- mm_missing_items = self.info.parse_mm_data(mm_missing_data)
1398
+ mm_missing_items = self.info.parse_mm_data(mm_missing_data, validate=False)
1399
1400
return mm_is_cached, mm_missing_items
1401
vllm/multimodal/registry.py
+5
@@ -138,6 +138,11 @@ class MultiModalRegistry:
138
mm_config.get_limit_per_prompt(modality) == 0
139
for modality in info.supported_mm_limits
140
):
141
+ # If enable_mm_embeds is True, we still need MM infrastructure
142
+ # to process pre-computed embeddings even though encoder won't run
143
+ if mm_config.enable_mm_embeds:
144
+ return True
145
+
146
logger.info_once(
147
"All limits of multimodal modalities supported by the model "
148
"are set to 0, running in text-only mode."
vllm/v1/worker/gpu_model_runner.py
+43
-30
@@ -1259,6 +1259,9 @@ class GPUModelRunner(
1259
mm_budget = self.mm_budget
1260
assert mm_budget is not None
1261
1262
+ if not mm_budget.mm_max_toks_per_item:
1263
+ return {} # No tower modalities (embed-only mode)
1264
+
1265
dummy_modality = mm_budget.get_modality_with_max_tokens()
1266
return self._get_mm_dummy_batch(dummy_modality, num_seqs)
1267
@@ -5116,40 +5119,50 @@ class GPUModelRunner(
5119
assert mm_budget is not None
5120
5121
if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
5119
- # NOTE: Currently model is profiled with a single non-text
5120
- # modality with the max possible input tokens even when
5121
- # it supports multiple.
5122
- dummy_modality = mm_budget.get_modality_with_max_tokens()
5123
- max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[
5124
- dummy_modality
5125
- ]
5122
+ if not mm_budget.mm_max_toks_per_item:
5123
+ # All modality limits are 0 — embedding-only mode.
5124
+ # Budget is non-zero for embedding storage, but
5125
+ # there's no encoder to profile.
5126
+ logger.info(
5127
+ "Skipping encoder profiling for embedding-only "
5128
+ "mode (all modality limits=0 with "
5129
+ "enable_mm_embeds=True).",
5130
+ )
5131
+ else:
5132
+ # NOTE: Currently model is profiled with a single
5133
+ # non-text modality with the max possible input
5134
+ # tokens even when it supports multiple.
5135
+ dummy_modality = mm_budget.get_modality_with_max_tokens()
5136
+ max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[
5137
+ dummy_modality
5138
+ ]
5139
5127
- logger.info(
5128
- "Encoder cache will be initialized with a budget of "
5129
- "%s tokens, and profiled with %s %s items of the "
5130
- "maximum feature size.",
5131
- encoder_budget,
5132
- max_mm_items_per_batch,
5133
- dummy_modality,
5134
- )
5140
+ logger.info(
5141
+ "Encoder cache will be initialized with a "
5142
+ "budget of %s tokens, and profiled with "
5143
+ "%s %s items of the maximum feature size.",
5144
+ encoder_budget,
5145
+ max_mm_items_per_batch,
5146
+ dummy_modality,
5147
+ )
5148
5136
- # Create dummy batch of multimodal inputs.
5137
- batched_dummy_mm_inputs = self._get_mm_dummy_batch(
5138
- dummy_modality,
5139
- max_mm_items_per_batch,
5140
- )
5149
+ # Create dummy batch of multimodal inputs.
5150
+ batched_dummy_mm_inputs = self._get_mm_dummy_batch(
5151
+ dummy_modality,
5152
+ max_mm_items_per_batch,
5153
+ )
5154
5142
- # Run multimodal encoder.
5143
- dummy_encoder_outputs = self.model.embed_multimodal(
5144
- **batched_dummy_mm_inputs
5145
- )
5155
+ # Run multimodal encoder.
5156
+ dummy_encoder_outputs = self.model.embed_multimodal(
5157
+ **batched_dummy_mm_inputs
5158
+ )
5159
5147
- sanity_check_mm_encoder_outputs(
5148
- dummy_encoder_outputs,
5149
- expected_num_items=max_mm_items_per_batch,
5150
- )
5151
- for i, output in enumerate(dummy_encoder_outputs):
5152
- self.encoder_cache[f"tmp_{i}"] = output
5160
+ sanity_check_mm_encoder_outputs(
5161
+ dummy_encoder_outputs,
5162
+ expected_num_items=max_mm_items_per_batch,
5163
+ )
5164
+ for i, output in enumerate(dummy_encoder_outputs):
5165
+ self.encoder_cache[f"tmp_{i}"] = output
5166
5167
# Add `is_profile` here to pre-allocate communication buffers
5168
hidden_states, last_hidden_states = self._dummy_run(