A high-throughput and memory-efficient inference and serving engine for LLMs
[BugFix] Fix prompt_embeds for multimodal models (#45383)
Signed-off-by: ruinan ma <r7ma3088@gmail.com>
Michael Ma committed
Jun 14, 2026 at 01:44 UTC
c621af16908f05270e033afd4237509902b7ba4d
2 files changed
+53
-7
vllm/config/vllm.py
+19
@@ -966,6 +966,15 @@ class VllmConfig:
966
"Async scheduling is not compatible with "
967
"disable_padded_drafter_batch=True."
968
)
969
+ if (
970
+ self.model_config is not None
971
+ and self.model_config.enable_prompt_embeds
972
+ and self.model_config.is_multimodal_model
973
+ ):
974
+ raise ValueError(
975
+ "Async scheduling is not yet supported with prompt embeds "
976
+ "for multimodal models."
977
+ )
978
if not executor_supports_async_sched:
979
raise ValueError(
980
f"`{executor_backend}` does not support async scheduling yet."
@@ -1009,6 +1018,16 @@ class VllmConfig:
1018
executor_backend,
1019
)
1020
self.scheduler_config.async_scheduling = False
1021
+ elif (
1022
+ self.model_config is not None
1023
+ and self.model_config.enable_prompt_embeds
1024
+ and self.model_config.is_multimodal_model
1025
+ ):
1026
+ logger.warning_once(
1027
+ "Async scheduling is not yet supported with prompt embeds "
1028
+ "for multimodal models and will be disabled."
1029
+ )
1030
+ self.scheduler_config.async_scheduling = False
1031
else:
1032
self.scheduler_config.async_scheduling = True
1033
vllm/v1/worker/gpu_model_runner.py
+34
-7
@@ -3452,14 +3452,41 @@ class GPUModelRunner(
3452
# NOTE(woosuk): To unify token ids and soft tokens (vision
3453
# embeddings), we always use embeddings (rather than token ids)
3454
# as input to the multimodal model, even when the input is text.
3455
- inputs_embeds_scheduled = self.model.embed_input_ids(
3456
- self.input_ids.gpu[:num_scheduled_tokens],
3457
- multimodal_embeddings=mm_embeds,
3458
- is_multimodal=is_mm_embed,
3459
- )
3455
+ if self.enable_prompt_embeds and self.input_batch.req_prompt_embeds:
3456
+ # Some positions carry precomputed prompt_embeds: they are
3457
+ # already in self.inputs_embeds and marked is_token_ids=False.
3458
+ # Embed only the token-id positions (zeroing the placeholder ids
3459
+ # at prompt_embeds positions so the embedding gather cannot read
3460
+ # out-of-range ids), and write them back without clobbering the
3461
+ # prompt_embeds positions.
3462
+ is_token_ids = self.is_token_ids.gpu[:num_scheduled_tokens]
3463
+ safe_input_ids = torch.where(
3464
+ is_token_ids,
3465
+ self.input_ids.gpu[:num_scheduled_tokens],
3466
+ 0,
3467
+ )
3468
+ inputs_embeds_scheduled = self.model.embed_input_ids(
3469
+ safe_input_ids,
3470
+ multimodal_embeddings=mm_embeds,
3471
+ is_multimodal=is_mm_embed,
3472
+ )
3473
+ target = self.inputs_embeds.gpu[:num_scheduled_tokens]
3474
+ self.inputs_embeds.gpu[:num_scheduled_tokens] = torch.where(
3475
+ is_token_ids.unsqueeze(-1),
3476
+ inputs_embeds_scheduled,
3477
+ target,
3478
+ )
3479
+ else:
3480
+ inputs_embeds_scheduled = self.model.embed_input_ids(
3481
+ self.input_ids.gpu[:num_scheduled_tokens],
3482
+ multimodal_embeddings=mm_embeds,
3483
+ is_multimodal=is_mm_embed,
3484
+ )
3485
3461
- # TODO(woosuk): Avoid the copy. Optimize.
3462
- self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
3486
+ # TODO(woosuk): Avoid the copy. Optimize.
3487
+ self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(
3488
+ inputs_embeds_scheduled
3489
+ )
3490
3491
input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens)
3492
model_kwargs = {