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
Remove index_put from MM embeddings merging (#22105)
Co-authored-by: Chenxi Yang <cxyang@meta.com>
Chenxi Yang committed
Aug 3, 2025 at 22:15 UTC
e5949e5ae013692ba09cc52472cf441675f5a270
1 file changed
+24
-18
vllm/model_executor/models/utils.py
+24
-18
@@ -393,7 +393,7 @@ def merge_multimodal_embeddings_from_map(
393
inputs_embeds: torch.Tensor, multimodal_embeddings: NestedTensors,
394
placeholder_map: MultiModalPlaceholderMap.IndexMap) -> torch.Tensor:
395
"""
396
- Merge ``multimodal_embeddings`` into ``inputs_embeds`` using the provided
396
+ Merge ``multimodal_embeddings`` into ``inputs_embeds`` using the provided
397
placeholder map .
398
399
Note:
@@ -418,17 +418,23 @@ def _merge_multimodal_embeddings(
418
Note:
419
This updates ``inputs_embeds`` in place.
420
"""
421
- num_expected_tokens = is_multimodal.sum().item()
422
- assert isinstance(num_expected_tokens, int)
423
-
421
flattened = _flatten_embeddings(multimodal_embeddings)
425
- if flattened.shape[0] != num_expected_tokens:
426
- expr = _embedding_count_expression(multimodal_embeddings)
427
- raise ValueError(
428
- f"Attempted to assign {expr} = {flattened.shape[0]} "
429
- f"multimodal tokens to {num_expected_tokens} placeholders")
422
+ try:
423
+ # This is equivalent to: inputs_embeds[is_multimodal] = flattened.
424
+ inputs_embeds.masked_scatter_(is_multimodal.unsqueeze(-1), flattened)
425
+ except RuntimeError as e:
426
+ num_expected_tokens = is_multimodal.sum().item()
427
+ assert isinstance(num_expected_tokens, int)
428
+
429
+ if flattened.shape[0] != num_expected_tokens:
430
+ expr = _embedding_count_expression(multimodal_embeddings)
431
+ raise ValueError(
432
+ f"Attempted to assign {expr} = {flattened.shape[0]} "
433
+ f"multimodal tokens to {num_expected_tokens} placeholders"
434
+ ) from e
435
+ else:
436
+ raise ValueError("Error during masked scatter operation") from e
437
431
- inputs_embeds[is_multimodal] = flattened
438
return inputs_embeds
439
440
@@ -478,11 +484,11 @@ def merge_multimodal_embeddings(
484
Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
485
positions in ``inputs_embeds`` corresponding to placeholder tokens in
486
``input_ids``.
481
-
482
- ``placeholder_token_id`` can be a list of token ids (e.g, token ids
483
- of img_start, img_break, and img_end tokens) when needed: This means
484
- the order of these tokens in the ``input_ids`` MUST MATCH the order of
485
- their embeddings in ``multimodal_embeddings`` since we need to
487
+
488
+ ``placeholder_token_id`` can be a list of token ids (e.g, token ids
489
+ of img_start, img_break, and img_end tokens) when needed: This means
490
+ the order of these tokens in the ``input_ids`` MUST MATCH the order of
491
+ their embeddings in ``multimodal_embeddings`` since we need to
492
slice-merge instead of individually scattering.
493
494
For example, if input_ids is "TTTTTSIIIBIIIBIIIETTT", where
@@ -491,9 +497,9 @@ def merge_multimodal_embeddings(
497
- I is image embedding token
498
- B is image break token
499
- E is image end token.
494
-
495
- Then the image embeddings (that correspond to I's) from vision encoder
496
- must be padded with embeddings of S, B, and E in the same order of
500
+
501
+ Then the image embeddings (that correspond to I's) from vision encoder
502
+ must be padded with embeddings of S, B, and E in the same order of
503
input_ids for a correct embedding merge.
504
505
Note: