27d7638b94 · tduncan/Nvidia-Optimization-Stack-VLLM

docs/contributing/model/multimodal.md

+5 -28

@@ -66,35 +66,12 @@ Further update the model as follows:

66 !!! important

67 The returned `multimodal_embeddings` must be either a **3D [torch.Tensor][]** of shape `(num_items, feature_size, hidden_size)`, or a **list / tuple of 2D [torch.Tensor][]'s** of shape `(feature_size, hidden_size)`, so that `multimodal_embeddings[i]` retrieves the embeddings generated from the `i`-th multimodal data item (e.g, image) of the request.

68

69 -- Implement [get_input_embeddings][vllm.model_executor.models.interfaces.SupportsMultiModal.get_input_embeddings] to merge `multimodal_embeddings` with text embeddings from the `input_ids`. If input processing for the model is implemented correctly (see sections below), then you can leverage the utility function we provide to easily merge the embeddings.

70 -

71 - ??? code

72 -

73 - ```python

74 - from .utils import merge_multimodal_embeddings

75 -

76 - class YourModelForImage2Seq(nn.Module):

77 - ...

78 -

79 - def get_input_embeddings(

80 - self,

81 - input_ids: torch.Tensor,

82 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

83 - ) -> torch.Tensor:

84 -

85 - # `get_input_embeddings` should already be implemented for the language

86 - # model as one of the requirements of basic vLLM model implementation.

87 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

88 -

89 - if multimodal_embeddings is not None:

90 - inputs_embeds = merge_multimodal_embeddings(

91 - input_ids=input_ids,

92 - inputs_embeds=inputs_embeds,

93 - multimodal_embeddings=multimodal_embeddings,

94 - placeholder_token_id=self.config.image_token_index)

69 +!!! note

70 + By default, vLLM merges the multimodal embeddings into text embeddings depending on the information of their locations defined in

71 + [PlaceholderRange][vllm.multimodal.inputs.PlaceholderRange] from input processing.

72 + This logic can be found at [get_input_embeddings][vllm.model_executor.models.interfaces.SupportsMultiModal.get_input_embeddings].

73

96 - return inputs_embeds

97 - ```

74 + You may override this method if additional logic is required for your model when merging embeddings.

75

76 - Implement [get_language_model][vllm.model_executor.models.interfaces.SupportsMultiModal.get_language_model] getter to provide stable access to the underlying language model.

77

vllm/config/model.py

+6 -1

@@ -509,9 +509,14 @@ class ModelConfig:

509 else: # task == "auto"

510 pass

511 else:

512 + debug_info = {

513 + "architectures": architectures,

514 + "is_generative_model": is_generative_model,

515 + "is_pooling_model": is_pooling_model,

516 + }

517 raise AssertionError("The model should be a generative or "

518 "pooling model when task is set to "

514 - f"{self.task!r}.")

519 + f"{self.task!r}. Found: {debug_info}")

520

521 self.runner = runner

522 self.convert = convert

vllm/model_executor/models/aria.py

+6 -19

@@ -38,8 +38,7 @@ from .idefics2_vision_model import (

38 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant

39 from .llama import LlamaDecoderLayer, LlamaMLP, LlamaModel

40 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

41 - is_pp_missing_parameter, maybe_prefix,

42 - merge_multimodal_embeddings)

41 + is_pp_missing_parameter, maybe_prefix)

42

43

44 class AriaImagePixelInputs(TensorSchema):

@@ -605,19 +604,6 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

604 multimodal_embeddings = self._process_image_input(image_input)

605 return multimodal_embeddings

606

608 - def get_input_embeddings(

609 - self,

610 - input_ids: torch.Tensor,

611 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

612 - ) -> torch.Tensor:

613 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

614 - if multimodal_embeddings is not None \

615 - and len(multimodal_embeddings) != 0:

616 - inputs_embeds = merge_multimodal_embeddings(

617 - input_ids, inputs_embeds, multimodal_embeddings,

618 - self.config.image_token_index)

619 - return inputs_embeds

620 -

607 def forward(

608 self,

609 input_ids: torch.Tensor,

@@ -628,10 +614,11 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

614 ) -> Union[torch.Tensor, IntermediateTensors]:

615 if inputs_embeds is None:

616 multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

631 - # always pass the input via `inputs_embeds`

632 - # to make sure the computation graph is consistent

633 - inputs_embeds = self.get_input_embeddings(input_ids,

634 - multimodal_embeddings)

617 + inputs_embeds = self.get_input_embeddings(

618 + input_ids,

619 + multimodal_embeddings,

620 + is_multimodal=input_ids == self.config.image_token_index,

621 + )

622 input_ids = None

623

624 hidden_states = self.language_model(

vllm/model_executor/models/aya_vision.py

+6 -21

@@ -33,8 +33,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

33 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

34 from .siglip import SiglipVisionModel

35 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

36 - init_vllm_registered_model, maybe_prefix,

37 - merge_multimodal_embeddings)

36 + init_vllm_registered_model, maybe_prefix)

37

38

39 class AyaVisionImagePixelInputs(TensorSchema):

@@ -417,23 +416,6 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,

416

417 return self._process_image_input(image_input, **kwargs)

418

420 - def get_input_embeddings(

421 - self,

422 - input_ids: torch.Tensor,

423 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

424 - ) -> torch.Tensor:

425 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

426 - if multimodal_embeddings is not None \

427 - and len(multimodal_embeddings) != 0:

428 - inputs_embeds = merge_multimodal_embeddings(

429 - input_ids=input_ids,

430 - inputs_embeds=inputs_embeds,

431 - multimodal_embeddings=multimodal_embeddings,

432 - placeholder_token_id=self.config.image_token_index,

433 - )

434 -

435 - return inputs_embeds

436 -

419 def forward(

420 self,

421 input_ids: torch.Tensor,

@@ -449,8 +431,11 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,

431 # condition is for v0 compatibility.

432 elif inputs_embeds is None:

433 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

452 - inputs_embeds = self.get_input_embeddings(input_ids,

453 - vision_embeddings)

434 + inputs_embeds = self.get_input_embeddings(

435 + input_ids,

436 + vision_embeddings,

437 + is_multimodal=input_ids == self.config.image_token_index,

438 + )

439 input_ids = None

440

441 hidden_states = self.language_model.model(

vllm/model_executor/models/bert.py

+12

@@ -348,6 +348,9 @@ class BertModel(nn.Module, SupportsQuant):

348 self.encoder = BertEncoder(vllm_config=vllm_config,

349 prefix=f"{prefix}.encoder")

350

351 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

352 + return self.embeddings(input_ids)

353 +

354 def forward(

355 self,

356 input_ids: torch.Tensor,

@@ -457,6 +460,9 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):

460 prefix=maybe_prefix(prefix, "model"))

461 self.pooler = self._build_pooler(pooler_config)

462

463 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

464 + return self.model.get_input_embeddings(input_ids)

465 +

466 def forward(

467 self,

468 input_ids: torch.Tensor,

@@ -588,6 +594,9 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,

594 ),

595 })

596

597 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

598 + return self.bert.get_input_embeddings(input_ids)

599 +

600 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

601 loader = AutoWeightsLoader(self)

602 loaded_params = loader.load_weights(weights)

@@ -637,6 +646,9 @@ class BertForTokenClassification(nn.Module):

646 Pooler.for_encode(pooler_config),

647 })

648

649 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

650 + return self.bert.get_input_embeddings(input_ids)

651 +

652 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

653 loader = AutoWeightsLoader(self)

654 loaded_params = loader.load_weights(weights)

vllm/model_executor/models/bert_with_rope.py

+6

@@ -426,6 +426,9 @@ class BertWithRope(nn.Module, SupportsQuant):

426 prefix=f"{prefix}.encoder")

427 self.pooler = BertPooler(self.config) if add_pooling_layer else None

428

429 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

430 + return self.embeddings(input_ids)

431 +

432 def forward(

433 self,

434 input_ids: torch.Tensor,

@@ -673,6 +676,9 @@ class GteNewForSequenceClassification(nn.Module, SupportsCrossEncoding):

676 loaded_params = loader.load_weights(weights)

677 return loaded_params

678

679 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

680 + return self.new.get_input_embeddings(input_ids)

681 +

682 def forward(

683 self,

684 input_ids: Optional[torch.Tensor],

vllm/model_executor/models/blip2.py

+6 -16

@@ -27,7 +27,7 @@ from .blip import BlipVisionModel

27 from .interfaces import (MultiModalEmbeddings, SupportsMultiModal, SupportsPP,

28 SupportsQuant)

29 from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,

30 - maybe_prefix, merge_multimodal_embeddings)

30 + maybe_prefix)

31

32 # We use this internally as placeholders since there is no image token

33 # defined on the HuggingFace repo

@@ -631,19 +631,6 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

631 vision_embeddings = self._process_image_input(image_input)

632 return vision_embeddings

633

634 - def get_input_embeddings(

635 - self,

636 - input_ids: torch.Tensor,

637 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

638 - ) -> torch.Tensor:

639 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

640 - if multimodal_embeddings is not None \

641 - and len(multimodal_embeddings) != 0:

642 - inputs_embeds = merge_multimodal_embeddings(

643 - input_ids, inputs_embeds, multimodal_embeddings,

644 - _IMAGE_TOKEN_ID)

645 - return inputs_embeds

646 -

634 def forward(

635 self,

636 input_ids: torch.Tensor,

@@ -689,8 +676,11 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

676 # condition is for v0 compatibility.

677 elif inputs_embeds is None:

678 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

692 - inputs_embeds = self.get_input_embeddings(input_ids,

693 - vision_embeddings)

679 + inputs_embeds = self.get_input_embeddings(

680 + input_ids,

681 + vision_embeddings,

682 + is_multimodal=input_ids == _IMAGE_TOKEN_ID,

683 + )

684 input_ids = None

685

686 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/chameleon.py

+7 -17

@@ -44,7 +44,7 @@ from .interfaces import (MultiModalEmbeddings, SupportsMultiModal, SupportsPP,

44 SupportsQuant)

45 from .utils import (flatten_bn, is_pp_missing_parameter,

46 make_empty_intermediate_tensors_factory, make_layers,

47 - maybe_prefix, merge_multimodal_embeddings)

47 + maybe_prefix)

48

49 logger = init_logger(__name__)

50

@@ -1002,20 +1002,6 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,

1002 vision_embeddings = self.model.get_input_embeddings(image_tokens)

1003 return vision_embeddings

1004

1005 - def get_input_embeddings(

1006 - self,

1007 - input_ids: torch.Tensor,

1008 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1009 - ) -> torch.Tensor:

1010 -

1011 - inputs_embeds = self.model.get_input_embeddings(input_ids)

1012 - if multimodal_embeddings is not None \

1013 - and len(multimodal_embeddings) != 0:

1014 - inputs_embeds = merge_multimodal_embeddings(

1015 - input_ids, inputs_embeds, multimodal_embeddings,

1016 - self.model.vocabulary_mapping.image_token_id)

1017 - return inputs_embeds

1018 -

1005 def forward(

1006 self,

1007 input_ids: torch.Tensor,

@@ -1032,8 +1018,12 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,

1018 # condition is for v0 compatibility.

1019 elif inputs_embeds is None:

1020 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

1035 - inputs_embeds = self.get_input_embeddings(input_ids,

1036 - vision_embeddings)

1021 + image_token_id = self.model.vocabulary_mapping.image_token_id

1022 + inputs_embeds = self.get_input_embeddings(

1023 + input_ids,

1024 + vision_embeddings,

1025 + is_multimodal=input_ids == image_token_id,

1026 + )

1027 input_ids = None

1028

1029 hidden_states = self.model(input_ids,

vllm/model_executor/models/chatglm.py

+3

@@ -433,6 +433,9 @@ class ChatGLMBaseModel(nn.Module):

433 self.make_empty_intermediate_tensors = (

434 self.transformer.make_empty_intermediate_tensors)

435

436 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

437 + return self.transformer.get_input_embeddings(input_ids)

438 +

439 def compute_logits(

440 self,

441 hidden_states: torch.Tensor,

vllm/model_executor/models/cohere2_vision.py

+6 -21

@@ -37,8 +37,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

37 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

38 from .siglip import SiglipVisionModel

39 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

40 - init_vllm_registered_model, maybe_prefix,

41 - merge_multimodal_embeddings)

40 + init_vllm_registered_model, maybe_prefix)

41

42

43 class Cohere2VisionImagePixelInputs(TensorSchema):

@@ -430,23 +429,6 @@ class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal,

429

430 return self._process_image_input(image_input, **kwargs)

431

433 - def get_input_embeddings(

434 - self,

435 - input_ids: torch.Tensor,

436 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

437 - ) -> torch.Tensor:

438 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

439 - if multimodal_embeddings is not None \

440 - and len(multimodal_embeddings) != 0:

441 - inputs_embeds = merge_multimodal_embeddings(

442 - input_ids=input_ids,

443 - inputs_embeds=inputs_embeds,

444 - multimodal_embeddings=multimodal_embeddings,

445 - placeholder_token_id=self.config.image_token_id,

446 - )

447 -

448 - return inputs_embeds

449 -

432 def forward(

433 self,

434 input_ids: torch.Tensor,

@@ -462,8 +444,11 @@ class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal,

444 # condition is for v0 compatibility.

445 elif inputs_embeds is None:

446 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

465 - inputs_embeds = self.get_input_embeddings(input_ids,

466 - vision_embeddings)

447 + inputs_embeds = self.get_input_embeddings(

448 + input_ids,

449 + vision_embeddings,

450 + is_multimodal=input_ids == self.config.image_token_id,

451 + )

452 input_ids = None

453

454 hidden_states = self.language_model.model(

vllm/model_executor/models/deepseek_eagle.py

+6

@@ -66,6 +66,9 @@ class DeepseekV2Model(nn.Module):

66 self.norm = RMSNorm(self.config.hidden_size,

67 eps=self.config.rms_norm_eps)

68

69 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

70 + return self.embed_tokens(input_ids)

71 +

72 def forward(

73 self,

74 input_ids: torch.Tensor,

@@ -205,6 +208,9 @@ class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):

208 self.logits_processor = LogitsProcessor(self.config.vocab_size,

209 scale=logit_scale)

210

211 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

212 + return self.model.get_input_embeddings(input_ids)

213 +

214 def forward(

215 self,

216 input_ids: torch.Tensor,

vllm/model_executor/models/deepseek_mtp.py

+6

@@ -101,6 +101,9 @@ class DeepSeekMultiTokenPredictor(nn.Module):

101 )

102 self.logits_processor = LogitsProcessor(config.vocab_size)

103

104 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

105 + return self.embed_tokens(input_ids)

106 +

107 def forward(

108 self,

109 input_ids: torch.Tensor,

@@ -142,6 +145,9 @@ class DeepSeekMTP(nn.Module, SupportsPP):

145 prefix=maybe_prefix(

146 prefix, "model"))

147

148 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

149 + return self.model.get_input_embeddings(input_ids)

150 +

151 def forward(

152 self,

153 input_ids: torch.Tensor,

vllm/model_executor/models/deepseek_vl2.py

+7 -18

@@ -41,8 +41,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

41

42 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

43 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

44 - init_vllm_registered_model, maybe_prefix,

45 - merge_multimodal_embeddings)

44 + init_vllm_registered_model, maybe_prefix)

45

46 # The image token id may be various

47 _IMAGE_TOKEN = "<image>"

@@ -346,7 +345,7 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

345

346 model_config = vllm_config.model_config

347 tokenizer = cached_tokenizer_from_config(model_config)

349 - self.image_token_id = tokenizer.vocab[_IMAGE_TOKEN]

348 + self.image_token_id: int = tokenizer.vocab[_IMAGE_TOKEN]

349

350 self.vision = self._init_vision_module(self.vision_config,

351 quant_config,

@@ -605,19 +604,6 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

604 vision_embeddings = self._process_image_input(image_input)

605 return vision_embeddings

606

608 - def get_input_embeddings(

609 - self,

610 - input_ids: torch.Tensor,

611 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

612 - ) -> torch.Tensor:

613 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

614 - if multimodal_embeddings is not None \

615 - and len(multimodal_embeddings) != 0:

616 - inputs_embeds = merge_multimodal_embeddings(

617 - input_ids, inputs_embeds, multimodal_embeddings,

618 - self.image_token_id)

619 - return inputs_embeds

620 -

607 def forward(self,

608 input_ids: torch.Tensor,

609 positions: torch.Tensor,

@@ -632,8 +618,11 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

618 # condition is for v0 compatibility

619 elif inputs_embeds is None:

620 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

635 - inputs_embeds = self.get_input_embeddings(input_ids,

636 - vision_embeddings)

621 + inputs_embeds = self.get_input_embeddings(

622 + input_ids,

623 + vision_embeddings,

624 + is_multimodal=input_ids == self.image_token_id,

625 + )

626 input_ids = None

627

628 hidden_states = self.language_model(input_ids,

vllm/model_executor/models/dots_ocr.py

+12 -32

@@ -34,8 +34,7 @@ from vllm.model_executor.models.qwen2_vl import (Qwen2VLDummyInputsBuilder,

34 Qwen2VLProcessingInfo)

35 from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,

36 init_vllm_registered_model,

37 - maybe_prefix,

38 - merge_multimodal_embeddings)

37 + maybe_prefix)

38 from vllm.model_executor.models.vision import get_vit_attn_backend

39 from vllm.multimodal import MULTIMODAL_REGISTRY

40 from vllm.multimodal.inputs import MultiModalDataDict

@@ -796,33 +795,17 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

795 def get_language_model(self) -> torch.nn.Module:

796 return self.language_model

797

799 - def get_multimodal_embeddings(

800 - self, **kwargs: object) -> Optional[MultiModalEmbeddings]:

798 + def get_multimodal_embeddings(self,

799 + **kwargs: object) -> MultiModalEmbeddings:

800 image_input = self._parse_and_validate_image_input(**kwargs)

801 if image_input is None:

802 return []

803 vision_embeddings = self._process_image_input(image_input)

804 return vision_embeddings

805

807 - def get_input_embeddings(

808 - self,

809 - input_ids: torch.Tensor,

810 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

811 - ) -> torch.Tensor:

812 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

813 - if multimodal_embeddings is not None:

814 - inputs_embeds = merge_multimodal_embeddings(

815 - input_ids,

816 - inputs_embeds,

817 - multimodal_embeddings,

818 - self.config.image_token_id,

819 - )

820 -

821 - return inputs_embeds

822 -

806 def forward(

807 self,

825 - input_ids: Optional[torch.Tensor],

808 + input_ids: torch.Tensor,

809 positions: torch.Tensor,

810 intermediate_tensors: Optional[IntermediateTensors] = None,

811 inputs_embeds: Optional[torch.Tensor] = None,

@@ -830,17 +813,14 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

813 ) -> Union[torch.Tensor, IntermediateTensors]:

814 if intermediate_tensors is not None:

815 inputs_embeds = None

833 - elif inputs_embeds is None and kwargs.get("pixel_values") is not None:

834 - image_input = self._parse_and_validate_image_input(**kwargs)

835 - if image_input is None:

836 - inputs_embeds = None

837 - else:

838 - assert input_ids is not None

839 - inputs_embeds = self.get_multimodal_embeddings(

840 - input_ids,

841 - image_input=image_input,

842 - )

843 - input_ids = None

816 + elif inputs_embeds is None:

817 + vision_embeddings = self.get_multimodal_embeddings(**kwargs)

818 + inputs_embeds = self.get_input_embeddings(

819 + input_ids,

820 + vision_embeddings,

821 + is_multimodal=input_ids == self.config.image_token_id,

822 + )

823 + input_ids = None

824

825 hidden_states = self.language_model(

826 input_ids=input_ids,

vllm/model_executor/models/ernie45_vl.py

+18 -13

@@ -60,8 +60,7 @@ from vllm.sequence import IntermediateTensors

60 from .ernie45_vl_moe import Ernie4_5_VLMoeForCausalLM

61 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

62 SupportsMultiModal, SupportsPP)

63 -from .utils import (AutoWeightsLoader, WeightsMapper, maybe_prefix,

64 - merge_multimodal_embeddings)

63 +from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix

64 from .vision import get_vit_attn_backend

65

66 logger = init_logger(__name__)

@@ -1467,18 +1466,24 @@ class Ernie4_5_VLMoeForConditionalGeneration(nn.Module, SupportsMultiModal,

1466 self,

1467 input_ids: torch.Tensor,

1468 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1469 + *,

1470 + is_multimodal: Optional[torch.Tensor] = None,

1471 + handle_oov_mm_token: bool = False,

1472 ) -> torch.Tensor:

1471 -

1472 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1473 -

1474 - if multimodal_embeddings is None:

1475 - return inputs_embeds

1476 -

1477 - self._set_visual_token_mask(input_ids)

1478 - inputs_embeds = merge_multimodal_embeddings(input_ids, inputs_embeds,

1479 - multimodal_embeddings,

1480 - [self.config.im_patch_id])

1481 - return inputs_embeds

1473 + if multimodal_embeddings is not None and len(

1474 + multimodal_embeddings) > 0:

1475 + self._set_visual_token_mask(input_ids)

1476 +

1477 + # This is to satisfy the type checker for each overload

1478 + if multimodal_embeddings is None or is_multimodal is None:

1479 + return super().get_input_embeddings(input_ids)

1480 +

1481 + return super().get_input_embeddings(

1482 + input_ids,

1483 + multimodal_embeddings=multimodal_embeddings,

1484 + is_multimodal=is_multimodal,

1485 + handle_oov_mm_token=handle_oov_mm_token,

1486 + )

1487

1488 def forward(

1489 self,

vllm/model_executor/models/ernie_mtp.py

+6

@@ -116,6 +116,9 @@ class ErnieMultiTokenPredictor(nn.Module):

116 )

117 self.logits_processor = LogitsProcessor(config.vocab_size)

118

119 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

120 + return self.embed_tokens(input_ids)

121 +

122 def forward(

123 self,

124 input_ids: torch.Tensor,

@@ -160,6 +163,9 @@ class ErnieMTP(nn.Module, SupportsPP):

163 if self.config.tie_word_embeddings:

164 self.lm_head.weight = self.model.embed_tokens.weight

165

166 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

167 + return self.model.get_input_embeddings(input_ids)

168 +

169 def forward(

170 self,

171 input_ids: torch.Tensor,

vllm/model_executor/models/fuyu.py

+6 -20

@@ -42,8 +42,7 @@ from vllm.sequence import IntermediateTensors

42 from vllm.utils.tensor_schema import TensorSchema, TensorShape

43

44 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

45 -from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn, maybe_prefix,

46 - merge_multimodal_embeddings)

45 +from .utils import AutoWeightsLoader, WeightsMapper, flatten_bn, maybe_prefix

46

47 # Cannot find the following 2 numbers from hf config.

48 _IMAGE_TOKEN_ID = 71011

@@ -342,22 +341,6 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

341

342 return self._process_image_input(image_input)

343

345 - def get_input_embeddings(

346 - self,

347 - input_ids: torch.Tensor,

348 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

349 - ) -> torch.Tensor:

350 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

351 - if multimodal_embeddings is not None \

352 - and len(multimodal_embeddings) != 0:

353 - inputs_embeds = merge_multimodal_embeddings(

354 - input_ids,

355 - inputs_embeds,

356 - multimodal_embeddings,

357 - _IMAGE_TOKEN_ID,

358 - )

359 - return inputs_embeds

360 -

344 def forward(

345 self,

346 input_ids: torch.Tensor,

@@ -373,8 +356,11 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

356 # condition is for v0 compatibility.

357 elif inputs_embeds is None:

358 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

376 - inputs_embeds = self.get_input_embeddings(input_ids,

377 - vision_embeddings)

359 + inputs_embeds = self.get_input_embeddings(

360 + input_ids,

361 + vision_embeddings,

362 + is_multimodal=input_ids == _IMAGE_TOKEN_ID,

363 + )

364 input_ids = None

365

366 hidden_states = self.language_model(

vllm/model_executor/models/gemma3_mm.py

+6 -20

@@ -37,8 +37,7 @@ from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

37 SupportsMultiModal, SupportsPP)

38 from .siglip import SiglipVisionModel

39 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

40 - init_vllm_registered_model, maybe_prefix,

41 - merge_multimodal_embeddings)

40 + init_vllm_registered_model, maybe_prefix)

41

42 logger = init_logger(__name__)

43

@@ -588,22 +587,6 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

587

588 return self._process_image_input(image_input)

589

591 - def get_input_embeddings(

592 - self,

593 - input_ids: torch.Tensor,

594 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

595 - ) -> torch.Tensor:

596 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

597 - if multimodal_embeddings is not None \

598 - and len(multimodal_embeddings) != 0:

599 - inputs_embeds = merge_multimodal_embeddings(

600 - input_ids,

601 - inputs_embeds,

602 - multimodal_embeddings,

603 - self.config.image_token_index,

604 - )

605 - return inputs_embeds

606 -

590 def forward(self,

591 input_ids: torch.Tensor,

592 positions: torch.Tensor,

@@ -618,8 +601,11 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

601 elif inputs_embeds is None:

602 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

603

621 - inputs_embeds = self.get_input_embeddings(input_ids,

622 - vision_embeddings)

604 + inputs_embeds = self.get_input_embeddings(

605 + input_ids,

606 + vision_embeddings,

607 + is_multimodal=input_ids == self.config.image_token_index,

608 + )

609 if (vision_embeddings is not None) and len(vision_embeddings) != 0:

610 kwargs = self.prepare_attn_masks(

611 input_ids,

vllm/model_executor/models/gemma3n_mm.py

+13 -10

@@ -632,8 +632,10 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal,

632 self,

633 input_ids: torch.Tensor,

634 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

635 + *,

636 + is_multimodal: Optional[torch.Tensor] = None,

637 + handle_oov_mm_token: bool = False,

638 ) -> torch.Tensor:

636 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

639 # NOTE (NickLucche) Each pass needs tokens to compute PLE so we cache

640 # them here, as the model forward has only access to the input_embeds.

641 if input_ids is not None:

@@ -645,15 +647,16 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal,

647 self.per_layer_embeddings[:per_layer_inputs.shape[0]].copy_(

648 per_layer_inputs)

649

648 - if multimodal_embeddings is not None \

649 - and len(multimodal_embeddings) != 0:

650 - inputs_embeds = merge_multimodal_embeddings(

651 - input_ids,

652 - inputs_embeds,

653 - multimodal_embeddings,

654 - # NOTE: this order of processing mm items is important

655 - [self.config.image_token_id, self.config.audio_token_id])

656 - return inputs_embeds

650 + # This is to satisfy the type checker for each overload

651 + if multimodal_embeddings is None or is_multimodal is None:

652 + return super().get_input_embeddings(input_ids)

653 +

654 + return super().get_input_embeddings(

655 + input_ids,

656 + multimodal_embeddings=multimodal_embeddings,

657 + is_multimodal=is_multimodal,

658 + handle_oov_mm_token=handle_oov_mm_token,

659 + )

660

661 def forward(self,

662 input_ids: torch.Tensor,

vllm/model_executor/models/glm4_1v.py

-17

@@ -1552,23 +1552,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1552 multimodal_embeddings += video_embeddings

1553 return multimodal_embeddings

1554

1555 - def get_input_embeddings(

1556 - self,

1557 - input_ids: torch.Tensor,

1558 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1559 - ) -> torch.Tensor:

1560 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1561 - if (multimodal_embeddings is not None

1562 - and len(multimodal_embeddings) != 0

1563 - and all(embed.numel() > 0 for embed in multimodal_embeddings)):

1564 - inputs_embeds = merge_multimodal_embeddings(

1565 - input_ids,

1566 - inputs_embeds,

1567 - multimodal_embeddings,

1568 - [self.config.image_token_id, self.config.video_token_id],

1569 - )

1570 - return inputs_embeds

1571 -

1555 def get_input_embeddings_v0(

1556 self,

1557 input_ids: torch.Tensor,

vllm/model_executor/models/glm4_moe_mtp.py

+6

@@ -132,6 +132,9 @@ class Glm4MoeMultiTokenPredictor(nn.Module):

132 )

133 self.logits_processor = LogitsProcessor(config.vocab_size)

134

135 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

136 + return self.embed_tokens(input_ids)

137 +

138 def forward(

139 self,

140 input_ids: torch.Tensor,

@@ -173,6 +176,9 @@ class Glm4MoeMTP(nn.Module, SupportsPP):

176 prefix=maybe_prefix(

177 prefix, "model"))

178

179 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

180 + return self.model.get_input_embeddings(input_ids)

181 +

182 def forward(

183 self,

184 input_ids: torch.Tensor,

vllm/model_executor/models/glm4v.py

+10 -25

@@ -43,7 +43,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

43 from .chatglm import ChatGLMBaseModel, ChatGLMModel

44 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

45 SupportsMultiModal, SupportsPP)

46 -from .utils import flatten_bn, merge_multimodal_embeddings

46 +from .utils import flatten_bn, isin_list

47

48

49 class GLMVImagePixelInputs(TensorSchema):

@@ -607,28 +607,6 @@ class GLM4VForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP,

607 vision_embeddings = self._process_image_input(image_input)

608 return vision_embeddings

609

610 - def get_input_embeddings(

611 - self,

612 - input_ids: torch.Tensor,

613 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

614 - ) -> torch.Tensor:

615 - inputs_embeds = self.transformer.get_input_embeddings(input_ids)

616 -

617 - if multimodal_embeddings is not None \

618 - and len(multimodal_embeddings) != 0:

619 - inputs_embeds = merge_multimodal_embeddings(

620 - input_ids=input_ids,

621 - inputs_embeds=inputs_embeds,

622 - multimodal_embeddings=multimodal_embeddings,

623 - placeholder_token_id=[

624 - self.config.boi_token_id,

625 - self.config.pad_token_id,

626 - self.config.eoi_token_id,

627 - ],

628 - )

629 -

630 - return inputs_embeds

631 -

610 def forward(

611 self,

612 input_ids: torch.Tensor,

@@ -644,8 +622,15 @@ class GLM4VForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP,

622 # condition is for v0 compatibility.

623 elif inputs_embeds is None:

624 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

647 - inputs_embeds = self.get_input_embeddings(input_ids,

648 - vision_embeddings)

625 + inputs_embeds = self.get_input_embeddings(

626 + input_ids,

627 + vision_embeddings,

628 + is_multimodal=isin_list(input_ids, [

629 + self.config.boi_token_id,

630 + self.config.pad_token_id,

631 + self.config.eoi_token_id,

632 + ]),

633 + )

634 input_ids = None

635

636 hidden_states = self.transformer(input_ids, positions,

vllm/model_executor/models/granite_speech.py

+21 -13

@@ -52,8 +52,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

52 from .blip2 import Blip2QFormerModel

53 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

54 SupportsMultiModal, SupportsPP)

55 -from .utils import (AutoWeightsLoader, embed_multimodal,

56 - init_vllm_registered_model, maybe_prefix)

55 +from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix

56

57

58 ### Audio Input

@@ -720,6 +719,9 @@ class GraniteSpeechForConditionalGeneration(

719 # Split variable length features into a tuple

720 return torch.split(masked_embeds, audio_input["audio_embed_sizes"])

721

722 + def get_language_model(self) -> torch.nn.Module:

723 + return self.language_model

724 +

725 def get_multimodal_embeddings(

726 self,

727 **kwargs: object,

@@ -728,7 +730,7 @@ class GraniteSpeechForConditionalGeneration(

730 audio_input = self._parse_and_validate_audio_input(**kwargs)

731 if audio_input is None:

732 return []

731 - return None

733 +

734 audio_features = self._process_audio_input(audio_input)

735 return audio_features

736

@@ -736,19 +738,21 @@ class GraniteSpeechForConditionalGeneration(

738 self,

739 input_ids: torch.Tensor,

740 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

741 + *,

742 + is_multimodal: Optional[torch.Tensor] = None,

743 + # Multi-modal token ID may exceed vocab size

744 + handle_oov_mm_token: bool = True,

745 ) -> torch.Tensor:

740 - """Compute the merged LLM / audio embeddings."""

741 - if multimodal_embeddings is None \

742 - or len(multimodal_embeddings) == 0:

743 - return self.language_model.get_input_embeddings(input_ids)

746 + # This is to satisfy the type checker for each overload

747 + if multimodal_embeddings is None or is_multimodal is None:

748 + return super().get_input_embeddings(input_ids)

749

745 - inputs_embeds = embed_multimodal(

750 + return super().get_input_embeddings(

751 input_ids,

747 - self.config.audio_token_index,

748 - self.language_model.model.get_input_embeddings,

749 - multimodal_embeddings,

752 + multimodal_embeddings=multimodal_embeddings,

753 + is_multimodal=is_multimodal,

754 + handle_oov_mm_token=handle_oov_mm_token,

755 )

751 - return inputs_embeds

756

757 def forward(

758 self,

@@ -765,7 +769,11 @@ class GraniteSpeechForConditionalGeneration(

769 # condition is for v0 compatibility.

770 elif inputs_embeds is None:

771 audio_embeds = self.get_multimodal_embeddings(**kwargs)

768 - inputs_embeds = self.get_input_embeddings(input_ids, audio_embeds)

772 + inputs_embeds = self.get_input_embeddings(

773 + input_ids,

774 + audio_embeds,

775 + is_multimodal=input_ids == self.config.audio_token_index,

776 + )

777 input_ids = None

778

779 model_output = self.language_model(input_ids, positions,

vllm/model_executor/models/hunyuan_v1.py

+3

@@ -989,6 +989,9 @@ class HunYuanV1Base(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):

989 moe.n_redundant_experts = self.num_redundant_experts

990 moe.experts.update_expert_map()

991

992 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

993 + return self.model.get_input_embeddings(input_ids)

994 +

995 def forward(

996 self,

997 input_ids: torch.Tensor,

vllm/model_executor/models/hyperclovax_vision.py

+10 -25

@@ -45,8 +45,8 @@ from vllm.sequence import IntermediateTensors

45 from .clip import CLIPVisionModel

46 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

47 from .siglip import SiglipVisionModel

48 -from .utils import (AutoWeightsLoader, init_vllm_registered_model,

49 - maybe_prefix, merge_multimodal_embeddings)

48 +from .utils import (AutoWeightsLoader, init_vllm_registered_model, isin_list,

49 + maybe_prefix)

50 from .vision import get_vision_encoder_info

51

52 EOT = "<|endofturn|>"

@@ -691,7 +691,7 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

691 def get_multimodal_embeddings(

692 self,

693 **kwargs: Unpack[HCXVisionMultimodalInputs],

694 - ) -> Optional[MultiModalEmbeddings]:

694 + ) -> MultiModalEmbeddings:

695

696 multimodal_embeddings = list()

697 if kwargs.get("pixel_values_images") is not None:

@@ -736,26 +736,6 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

736 multimodal_embeddings.append(_multimodal_embeddings_videos)

737 return multimodal_embeddings

738

739 - def get_input_embeddings(

740 - self,

741 - input_ids: torch.Tensor,

742 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

743 - ) -> torch.Tensor:

744 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

745 - if multimodal_embeddings is not None \

746 - and len(multimodal_embeddings) != 0:

747 - inputs_embeds = merge_multimodal_embeddings(

748 - input_ids,

749 - inputs_embeds,

750 - multimodal_embeddings,

751 - placeholder_token_id=[

752 - self.config.image_token_id,

753 - self.config.video_token_id,

754 - ],

755 - )

756 -

757 - return inputs_embeds

758 -

739 def forward(

740 self,

741 input_ids: torch.Tensor,

@@ -771,8 +751,13 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

751 # condition is for v0 compatibility.

752 elif inputs_embeds is None:

753 multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

774 - inputs_embeds = self.get_input_embeddings(input_ids,

775 - multimodal_embeddings)

754 + inputs_embeds = self.get_input_embeddings(

755 + input_ids,

756 + multimodal_embeddings,

757 + is_multimodal=isin_list(

758 + input_ids,

759 + [self.config.image_token_id, self.config.video_token_id]),

760 + )

761 input_ids = None

762 hidden_states = self.language_model.model(input_ids,

763 positions,

vllm/model_executor/models/idefics3.py

+7 -24

@@ -52,8 +52,7 @@ from .idefics2_vision_model import (

52 # yapf: enable

53 from .interfaces import MultiModalEmbeddings, SupportsLoRA, SupportsMultiModal

54 from .llama import LlamaModel

55 -from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,

56 - merge_multimodal_embeddings)

55 +from .utils import AutoWeightsLoader, flatten_bn, maybe_prefix

56

57

58 class Idefics3ImagePixelInputs(TensorSchema):

@@ -539,10 +538,7 @@ class Idefics3Model(nn.Module):

538

539 return image_hidden_states

540

542 - def get_input_embeddings(

543 - self,

544 - input_ids: torch.Tensor,

545 - ) -> torch.Tensor:

541 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

542 return self.text_model.get_input_embeddings(input_ids)

543

544 def forward(

@@ -695,22 +691,6 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,

691

692 return self._process_image_input(image_input)

693

698 - def get_input_embeddings(

699 - self,

700 - input_ids: torch.Tensor,

701 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

702 - ) -> torch.Tensor:

703 - inputs_embeds = self.model.get_input_embeddings(input_ids)

704 - if multimodal_embeddings is not None \

705 - and len(multimodal_embeddings) != 0:

706 - inputs_embeds = merge_multimodal_embeddings(

707 - input_ids,

708 - inputs_embeds,

709 - multimodal_embeddings,

710 - self.config.image_token_id,

711 - )

712 - return inputs_embeds

713 -

694 def forward(

695 self,

696 input_ids: torch.Tensor,

@@ -726,8 +706,11 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,

706 # condition is for v0 compatibility.

707 elif inputs_embeds is None:

708 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

729 - inputs_embeds = self.get_input_embeddings(input_ids,

730 - vision_embeddings)

709 + inputs_embeds = self.get_input_embeddings(

710 + input_ids,

711 + vision_embeddings,

712 + is_multimodal=input_ids == self.config.image_token_id,

713 + )

714 input_ids = None

715

716 hidden_states = self.model.text_model(input_ids,

vllm/model_executor/models/interfaces.py

+75 -8

@@ -2,8 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Iterable, Mapping, MutableSequence

5 -from typing import (TYPE_CHECKING, ClassVar, Literal, Optional, Protocol,

6 - Union, overload, runtime_checkable)

5 +from typing import (TYPE_CHECKING, Callable, ClassVar, Literal, Optional,

6 + Protocol, Union, overload, runtime_checkable)

7

8 import numpy as np

9 import torch

@@ -20,7 +20,7 @@ from vllm.model_executor.layers.quantization.base_config import (

20 QuantizationConfig)

21 from vllm.utils import supports_kw

22

23 -from .interfaces_base import is_pooling_model

23 +from .interfaces_base import VllmModel, is_pooling_model

24

25 if TYPE_CHECKING:

26 from vllm.config import VllmConfig

@@ -90,7 +90,7 @@ class SupportsMultiModal(Protocol):

90 """

91 ...

92

93 - def get_language_model(self) -> torch.nn.Module:

93 + def get_language_model(self) -> VllmModel:

94 """

95 Returns the underlying language model used for text generation.

96

@@ -102,17 +102,84 @@ class SupportsMultiModal(Protocol):

102 """

103 ...

104

105 + @overload

106 + def get_input_embeddings(self, input_ids: Tensor) -> Tensor:

107 + ...

108 +

109 + @overload

110 + def get_input_embeddings(

111 + self,

112 + input_ids: Tensor,

113 + multimodal_embeddings: MultiModalEmbeddings,

114 + *,

115 + is_multimodal: torch.Tensor,

116 + handle_oov_mm_token: bool = False,

117 + ) -> Tensor:

118 + ...

119 +

120 + def _get_text_embeddings(

121 + self,

122 + input_ids: Tensor,

123 + get_input_embeddings: Callable[[Tensor], Tensor],

124 + *,

125 + is_multimodal: Optional[Tensor],

126 + handle_oov_mm_token: bool,

127 + ) -> Tensor:

128 + if handle_oov_mm_token and is_multimodal is not None:

129 + is_text = ~is_multimodal

130 + text_embeds = get_input_embeddings(input_ids[is_text])

131 +

132 + return torch.empty(

133 + (input_ids.shape[0], text_embeds.shape[1]),

134 + dtype=text_embeds.dtype,

135 + device=text_embeds.device,

136 + ).masked_scatter_(is_text.unsqueeze_(-1), text_embeds)

137 +

138 + return get_input_embeddings(input_ids)

139 +

140 def get_input_embeddings(

141 self,

142 input_ids: Tensor,

143 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

144 + *,

145 + is_multimodal: Optional[Tensor] = None,

146 + handle_oov_mm_token: bool = False,

147 ) -> Tensor:

148 """

111 - Returns the input embeddings merged from the text embeddings from

112 - input_ids and the multimodal embeddings generated from multimodal

113 - kwargs.

149 + Apply token embeddings to `input_ids`.

150 +

151 + If `multimodal_embeddings` is passed, scatter them into

152 + `input_ids` according to the mask `is_multimodal`.

153 +

154 + In case the multi-modal token IDs exceed the vocabulary size of

155 + the language model, you can set `handle_oov_mm_token=False`

156 + to avoid calling the language model's `get_input_embeddings` method

157 + on those tokens. Note however that doing so increases memory usage

158 + as an additional buffer is needed to hold the input embeddings.

159 """

115 - ...

160 + from .utils import _merge_multimodal_embeddings

161 +

162 + inputs_embeds = self._get_text_embeddings(

163 + input_ids,

164 + self.get_language_model().get_input_embeddings,

165 + is_multimodal=is_multimodal,

166 + handle_oov_mm_token=handle_oov_mm_token,

167 + )

168 +

169 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:

170 + return inputs_embeds

171 +

172 + if is_multimodal is None:

173 + raise ValueError(

174 + "`get_input_embeddings` now requires `is_multimodal` arg, "

175 + "please update your model runner according to "

176 + "https://github.com/vllm-project/vllm/pull/16229.")

177 +

178 + return _merge_multimodal_embeddings(

179 + inputs_embeds=inputs_embeds,

180 + multimodal_embeddings=multimodal_embeddings,

181 + is_multimodal=is_multimodal,

182 + )

183

184

185 @runtime_checkable

vllm/model_executor/models/interfaces_base.py

+23 -1

@@ -41,6 +41,13 @@ class VllmModel(Protocol[T_co]):

41 ) -> None:

42 ...

43

44 + def get_input_embeddings(

45 + self,

46 + input_ids: torch.Tensor,

47 + ) -> torch.Tensor:

48 + """Apply token embeddings to `input_ids`."""

49 + ...

50 +

51 def forward(

52 self,

53 input_ids: torch.Tensor,

@@ -54,6 +61,19 @@ def _check_vllm_model_init(model: Union[type[object], object]) -> bool:

61 return supports_kw(model_init, "vllm_config")

62

63

64 +def _check_vllm_model_get_input_embeddings(

65 + model: Union[type[object], object]) -> bool:

66 + model_get_input_embeddings = getattr(model, "get_input_embeddings", None)

67 + if not callable(model_get_input_embeddings):

68 + logger.warning(

69 + "The model (%s) is missing the `get_input_embeddings` method.",

70 + model,

71 + )

72 + return False

73 +

74 + return True

75 +

76 +

77 def _check_vllm_model_forward(model: Union[type[object], object]) -> bool:

78 model_forward = getattr(model, "forward", None)

79 if not callable(model_forward):

@@ -88,7 +108,9 @@ def is_vllm_model(model: object) -> TypeIs[VllmModel]:

108 def is_vllm_model(

109 model: Union[type[object], object],

110 ) -> Union[TypeIs[type[VllmModel]], TypeIs[VllmModel]]:

91 - return _check_vllm_model_init(model) and _check_vllm_model_forward(model)

111 + return (_check_vllm_model_init(model)

112 + and _check_vllm_model_get_input_embeddings(model)

113 + and _check_vllm_model_forward(model))

114

115

116 @runtime_checkable

vllm/model_executor/models/interns1.py

+27 -20

@@ -40,8 +40,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

40 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

41 SupportsMultiModal, SupportsPP)

42 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

43 - init_vllm_registered_model, maybe_prefix,

44 - merge_multimodal_embeddings)

43 + init_vllm_registered_model, isin_list, maybe_prefix)

44

45

46 class InternS1MultiModalProjector(nn.Module):

@@ -767,24 +766,24 @@ class InternS1ForConditionalGeneration(nn.Module, SupportsMultiModal,

766 self,

767 input_ids: torch.Tensor,

768 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

769 + *,

770 + is_multimodal: Optional[torch.Tensor] = None,

771 + handle_oov_mm_token: bool = False,

772 ) -> torch.Tensor:

771 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

772 - if multimodal_embeddings is not None \

773 - and len(multimodal_embeddings) != 0:

774 - context_token_ids = [

775 - token_id for token_id in (self.img_context_token_id,

776 - self.video_context_token_id)

777 - if token_id is not None

778 - ]

779 - assert len(context_token_ids) >= 1

773 + if multimodal_embeddings is not None and len(

774 + multimodal_embeddings) > 0:

775 self._set_visual_token_mask(input_ids)

781 - inputs_embeds = merge_multimodal_embeddings(

782 - input_ids,

783 - inputs_embeds,

784 - multimodal_embeddings,

785 - context_token_ids,

786 - )

787 - return inputs_embeds

776 +

777 + # This is to satisfy the type checker for each overload

778 + if multimodal_embeddings is None or is_multimodal is None:

779 + return super().get_input_embeddings(input_ids)

780 +

781 + return super().get_input_embeddings(

782 + input_ids,

783 + multimodal_embeddings=multimodal_embeddings,

784 + is_multimodal=is_multimodal,

785 + handle_oov_mm_token=handle_oov_mm_token,

786 + )

787

788 def forward(

789 self,

@@ -802,9 +801,17 @@ class InternS1ForConditionalGeneration(nn.Module, SupportsMultiModal,

801 # NOTE: In v1, inputs_embeds is always generated at model runner, this

802 # condition is for v0 compatibility.

803 elif inputs_embeds is None:

804 + context_token_ids = [

805 + token_id for token_id in (self.img_context_token_id,

806 + self.video_context_token_id)

807 + if token_id is not None

808 + ]

809 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

806 - inputs_embeds = self.get_input_embeddings(input_ids,

807 - vision_embeddings)

810 + inputs_embeds = self.get_input_embeddings(

811 + input_ids,

812 + vision_embeddings,

813 + is_multimodal=isin_list(input_ids, context_token_ids),

814 + )

815 input_ids = None

816

817 forward_kwargs = {

vllm/model_executor/models/internvl.py

+27 -19

@@ -43,7 +43,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

43 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

44 SupportsMultiModal, SupportsPP)

45 from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,

46 - maybe_prefix, merge_multimodal_embeddings)

46 + isin_list, maybe_prefix)

47

48 IMG_START = '<img>'

49 IMG_END = '</img>'

@@ -1339,24 +1339,24 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,

1339 self,

1340 input_ids: torch.Tensor,

1341 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1342 + *,

1343 + is_multimodal: Optional[torch.Tensor] = None,

1344 + handle_oov_mm_token: bool = False,

1345 ) -> torch.Tensor:

1343 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1344 - if multimodal_embeddings is not None \

1345 - and len(multimodal_embeddings) != 0:

1346 - context_token_ids = [

1347 - token_id for token_id in (self.img_context_token_id,

1348 - self.video_context_token_id)

1349 - if token_id is not None

1350 - ]

1351 - assert len(context_token_ids) >= 1

1346 + if multimodal_embeddings is not None and len(

1347 + multimodal_embeddings) > 0:

1348 self._set_visual_token_mask(input_ids)

1353 - inputs_embeds = merge_multimodal_embeddings(

1354 - input_ids,

1355 - inputs_embeds,

1356 - multimodal_embeddings,

1357 - context_token_ids,

1358 - )

1359 - return inputs_embeds

1349 +

1350 + # This is to satisfy the type checker for each overload

1351 + if multimodal_embeddings is None or is_multimodal is None:

1352 + return super().get_input_embeddings(input_ids)

1353 +

1354 + return super().get_input_embeddings(

1355 + input_ids,

1356 + multimodal_embeddings=multimodal_embeddings,

1357 + is_multimodal=is_multimodal,

1358 + handle_oov_mm_token=handle_oov_mm_token,

1359 + )

1360

1361 def forward(

1362 self,

@@ -1374,9 +1374,17 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,

1374 # NOTE: In v1, inputs_embeds is always generated at model runner, this

1375 # condition is for v0 compatibility.

1376 elif inputs_embeds is None:

1377 + context_token_ids = [

1378 + token_id for token_id in (self.img_context_token_id,

1379 + self.video_context_token_id)

1380 + if token_id is not None

1381 + ]

1382 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

1378 - inputs_embeds = self.get_input_embeddings(input_ids,

1379 - vision_embeddings)

1383 + inputs_embeds = self.get_input_embeddings(

1384 + input_ids,

1385 + vision_embeddings,

1386 + is_multimodal=isin_list(input_ids, context_token_ids),

1387 + )

1388 input_ids = None

1389

1390 forward_kwargs = {

vllm/model_executor/models/keye.py

-18

@@ -1450,24 +1450,6 @@ class BaseKeyeModule(nn.Module):

1450 multimodal_embeddings += video_embeddings

1451 return multimodal_embeddings

1452

1453 - def get_input_embeddings(

1454 - self,

1455 - input_ids: torch.Tensor,

1456 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1457 - ) -> torch.Tensor:

1458 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1459 - if multimodal_embeddings is not None:

1460 - inputs_embeds = merge_multimodal_embeddings(

1461 - input_ids,

1462 - inputs_embeds,

1463 - multimodal_embeddings,

1464 - [

1465 - self.config.image_token_id,

1466 - self.config.video_token_id,

1467 - ],

1468 - )

1469 - return inputs_embeds

1470 -

1453 def get_input_embeddings_v0(

1454 self,

1455 input_ids: torch.Tensor,

vllm/model_executor/models/kimi_vl.py

+3 -26

@@ -66,7 +66,6 @@ from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model

66 from vllm.model_executor.models.interfaces import (SupportsMultiModal,

67 SupportsPP)

68 from vllm.model_executor.models.moonvit import MoonVitPretrainedModel

69 -from vllm.model_executor.models.utils import merge_multimodal_embeddings

69 from vllm.multimodal import MULTIMODAL_REGISTRY

70 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

71 MultiModalKwargsItems, NestedTensors)

@@ -424,26 +423,6 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

423 vision_embeddings = self._process_image_input(image_input)

424 return vision_embeddings

425

427 - def get_input_embeddings(

428 - self,

429 - input_ids: torch.Tensor,

430 - multimodal_embeddings: Optional[NestedTensors] = None,

431 - ) -> torch.Tensor:

432 -

433 - # `get_input_embeddings` should already be implemented for the language

434 - # model as one of the requirements of basic vLLM model implementation.

435 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

436 -

437 - if multimodal_embeddings is not None and len(

438 - multimodal_embeddings) != 0:

439 - inputs_embeds = merge_multimodal_embeddings(

440 - input_ids=input_ids,

441 - inputs_embeds=inputs_embeds,

442 - multimodal_embeddings=multimodal_embeddings,

443 - placeholder_token_id=self.config.media_placeholder_token_id)

444 -

445 - return inputs_embeds

446 -

426 def forward(

427 self,

428 input_ids: torch.Tensor,

@@ -462,14 +441,12 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

441 if image_input is None:

442 inputs_embeds = None

443 else:

465 - inputs_embeds = self.get_input_embeddings(input_ids)

444 image_embeds = self._process_image_input(image_input)

467 - inputs_embeds = merge_multimodal_embeddings(

445 + inputs_embeds = self.get_input_embeddings(

446 input_ids,

469 - inputs_embeds,

447 image_embeds,

471 - placeholder_token_id=self.config.

472 - media_placeholder_token_id,

448 + is_multimodal=input_ids ==

449 + self.config.media_placeholder_token_id,

450 )

451 input_ids = None

452

vllm/model_executor/models/lfm2.py

+3

@@ -522,6 +522,9 @@ class Lfm2ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

522 self.make_empty_intermediate_tensors = (

523 self.model.make_empty_intermediate_tensors)

524

525 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

526 + return self.model.get_input_embeddings(input_ids)

527 +

528 def forward(

529 self,

530 input_ids: torch.Tensor,

vllm/model_executor/models/llama4_eagle.py

+8 -23

@@ -37,9 +37,9 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader

37 from vllm.model_executor.models.llama4 import (Llama4DecoderLayer,

38 Llama4ForCausalLM)

39 from vllm.model_executor.models.utils import extract_layer_index

40 -from vllm.multimodal.inputs import NestedTensors

40

42 -from .utils import AutoWeightsLoader, maybe_prefix, merge_multimodal_embeddings

41 +from .interfaces import SupportsMultiModal

42 +from .utils import AutoWeightsLoader, maybe_prefix

43

44 logger = init_logger(__name__)

45

@@ -79,10 +79,7 @@ class LlamaModel(nn.Module):

79 self.norm = RMSNorm(self.config.hidden_size,

80 eps=self.config.rms_norm_eps)

81

82 - def get_input_embeddings(

83 - self,

84 - input_ids: torch.Tensor,

85 - ) -> torch.Tensor:

82 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

83 return self.embed_tokens(input_ids)

84

85 def forward(

@@ -194,6 +191,11 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):

191 self.logits_processor = LogitsProcessor(self.config.vocab_size,

192 scale=logit_scale)

193

194 + def get_language_model(self) -> torch.nn.Module:

195 + return self.model

196 +

197 + get_input_embeddings = SupportsMultiModal.get_input_embeddings # type: ignore

198 +

199 def forward(

200 self,

201 input_ids: torch.Tensor,

@@ -220,20 +222,3 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):

222 skip_prefixes=(["lm_head."]),

223 )

224 loader.load_weights(map(transform, weights))

223 -

224 - def get_input_embeddings(

225 - self,

226 - input_ids: torch.Tensor,

227 - multimodal_embeddings: Optional[NestedTensors] = None,

228 - ) -> torch.Tensor:

229 - inputs_embeds = self.model.get_input_embeddings(input_ids)

230 -

231 - if multimodal_embeddings is not None:

232 - inputs_embeds = merge_multimodal_embeddings(

233 - input_ids,

234 - inputs_embeds,

235 - multimodal_embeddings,

236 - self.config.image_token_index,

237 - )

238 -

239 - return inputs_embeds

vllm/model_executor/models/llama_eagle.py

+6

@@ -73,6 +73,9 @@ class LlamaModel(nn.Module):

73 self.config.hidden_size,

74 bias=False)

75

76 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

77 + return self.embed_tokens(input_ids)

78 +

79 def forward(

80 self,

81 input_ids: torch.Tensor,

@@ -149,6 +152,9 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):

152 self.logits_processor = LogitsProcessor(self.config.vocab_size,

153 scale=logit_scale)

154

155 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

156 + return self.model.get_input_embeddings(input_ids)

157 +

158 def forward(

159 self,

160 input_ids: torch.Tensor,

vllm/model_executor/models/llama_eagle3.py

+4 -13

@@ -18,7 +18,6 @@ from vllm.model_executor.layers.quantization.base_config import (

18 from vllm.model_executor.layers.vocab_parallel_embedding import (

19 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

20 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

21 -from vllm.model_executor.models.interfaces import MultiModalEmbeddings

21 from vllm.model_executor.models.llama import (LlamaDecoderLayer,

22 LlamaForCausalLM)

23

@@ -144,10 +143,7 @@ class LlamaModel(nn.Module):

143 eps=self.config.rms_norm_eps,

144 )

145

147 - def get_input_embeddings(

148 - self,

149 - input_ids: torch.Tensor,

150 - ) -> torch.Tensor:

146 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

147 return self.embed_tokens(input_ids)

148

149 def forward(

@@ -239,6 +235,9 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

235 requires_grad=False,

236 )

237

238 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

239 + return self.model.get_input_embeddings(input_ids)

240 +

241 def forward(

242 self,

243 input_ids: torch.Tensor,

@@ -302,11 +301,3 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

301 skip_substrs=skip_substrs,

302 )

303 loader.load_weights(model_weights.items())

305 -

306 - def get_input_embeddings(

307 - self,

308 - input_ids: torch.Tensor,

309 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

310 - ) -> torch.Tensor:

311 - inputs_embeds = self.model.get_input_embeddings(input_ids)

312 - return inputs_embeds

vllm/model_executor/models/llava.py

+6 -20

@@ -41,8 +41,7 @@ from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

41 from .pixtral import PixtralHFEncoderInfo, PixtralHFVisionModel

42 from .siglip import SiglipVisionModel

43 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

44 - init_vllm_registered_model, maybe_prefix,

45 - merge_multimodal_embeddings)

44 + init_vllm_registered_model, maybe_prefix)

45 from .vision import get_vision_encoder_info

46

47

@@ -676,22 +675,6 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

675

676 return self._process_image_input(image_input)

677

679 - def get_input_embeddings(

680 - self,

681 - input_ids: torch.Tensor,

682 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

683 - ) -> torch.Tensor:

684 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

685 - if multimodal_embeddings is not None \

686 - and len(multimodal_embeddings) != 0:

687 - inputs_embeds = merge_multimodal_embeddings(

688 - input_ids,

689 - inputs_embeds,

690 - multimodal_embeddings,

691 - self.config.image_token_index,

692 - )

693 - return inputs_embeds

694 -

678 def forward(

679 self,

680 input_ids: torch.Tensor,

@@ -744,8 +727,11 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

727 # condition is for v0 compatibility.

728 elif inputs_embeds is None:

729 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

747 - inputs_embeds = self.get_input_embeddings(input_ids,

748 - vision_embeddings)

730 + inputs_embeds = self.get_input_embeddings(

731 + input_ids,

732 + vision_embeddings,

733 + is_multimodal=input_ids == self.config.image_token_index,

734 + )

735 input_ids = None

736

737 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/llava_next.py

+18 -13

@@ -25,8 +25,8 @@ from .llava import (BaseLlavaMultiModalProcessor, BaseLlavaProcessingInfo,

25 LlavaDummyInputsBuilder, LlavaLikeConfig,

26 LlavaMultiModalProjector, init_vision_tower_for_llava)

27 from .siglip import SiglipVisionModel

28 -from .utils import (AutoWeightsLoader, WeightsMapper, embed_multimodal,

29 - flatten_bn, init_vllm_registered_model, maybe_prefix)

28 +from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

29 + init_vllm_registered_model, maybe_prefix)

30

31

32 class LlavaNextImagePixelInputs(TensorSchema):

@@ -474,19 +474,21 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,

474 self,

475 input_ids: torch.Tensor,

476 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

477 + *,

478 + is_multimodal: Optional[torch.Tensor] = None,

479 + # Multi-modal token ID may exceed vocab size

480 + handle_oov_mm_token: bool = True,

481 ) -> torch.Tensor:

482 + # This is to satisfy the type checker for each overload

483 + if multimodal_embeddings is None or is_multimodal is None:

484 + return super().get_input_embeddings(input_ids)

485

479 - if multimodal_embeddings is None \

480 - or len(multimodal_embeddings) == 0:

481 - return self.language_model.get_input_embeddings(input_ids)

482 -

483 - inputs_embeds = embed_multimodal(

486 + return super().get_input_embeddings(

487 input_ids,

485 - self.config.image_token_index,

486 - self.language_model.model.get_input_embeddings,

487 - multimodal_embeddings,

488 + multimodal_embeddings=multimodal_embeddings,

489 + is_multimodal=is_multimodal,

490 + handle_oov_mm_token=handle_oov_mm_token,

491 )

489 - return inputs_embeds

492

493 def forward(

494 self,

@@ -549,8 +551,11 @@ model_executor.models.llava_next.LlavaNextProcessingInfo.get_num_image_tokens].

551 # condition is for v0 compatibility.

552 elif inputs_embeds is None:

553 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

552 - inputs_embeds = self.get_input_embeddings(input_ids,

553 - vision_embeddings)

554 + inputs_embeds = self.get_input_embeddings(

555 + input_ids,

556 + vision_embeddings,

557 + is_multimodal=input_ids == self.config.image_token_index,

558 + )

559 input_ids = None

560

561 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/llava_next_video.py

+6 -17

@@ -30,8 +30,7 @@ from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

30 from .llava import init_vision_tower_for_llava

31 from .siglip import SiglipVisionModel

32 from .utils import (AutoWeightsLoader, WeightsMapper,

33 - init_vllm_registered_model, maybe_prefix,

34 - merge_multimodal_embeddings)

33 + init_vllm_registered_model, maybe_prefix)

34 from .vision import get_vision_encoder_info

35

36

@@ -415,19 +414,6 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal,

414 vision_embeddings = self._process_video_pixels(video_input)

415 return vision_embeddings

416

418 - def get_input_embeddings(

419 - self,

420 - input_ids: torch.Tensor,

421 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

422 - ) -> torch.Tensor:

423 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

424 - if multimodal_embeddings is not None \

425 - and len(multimodal_embeddings) != 0:

426 - inputs_embeds = merge_multimodal_embeddings(

427 - input_ids, inputs_embeds, multimodal_embeddings,

428 - self.config.video_token_index)

429 - return inputs_embeds

430 -

417 def forward(

418 self,

419 input_ids: torch.Tensor,

@@ -449,8 +435,11 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal,

435 # condition is for v0 compatibility.

436 elif inputs_embeds is None:

437 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

452 - inputs_embeds = self.get_input_embeddings(input_ids,

453 - vision_embeddings)

438 + inputs_embeds = self.get_input_embeddings(

439 + input_ids,

440 + vision_embeddings,

441 + is_multimodal=input_ids == self.config.video_token_index,

442 + )

443 input_ids = None

444

445 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/llava_onevision.py

-13

@@ -850,19 +850,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

850

851 return multimodal_embeddings

852

853 - def get_input_embeddings(

854 - self,

855 - input_ids: torch.Tensor,

856 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

857 - ) -> torch.Tensor:

858 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

859 - if multimodal_embeddings is not None \

860 - and len(multimodal_embeddings) != 0:

861 - inputs_embeds = merge_multimodal_embeddings(

862 - input_ids, inputs_embeds, multimodal_embeddings,

863 - [self.config.image_token_index, self.config.video_token_index])

864 - return inputs_embeds

865 -

853 def get_input_embeddings_v0(

854 self,

855 input_ids: torch.Tensor,

vllm/model_executor/models/midashenglm.py

+6 -19

@@ -54,8 +54,7 @@ from vllm.sequence import IntermediateTensors

54 from vllm.transformers_utils.configs.midashenglm import DashengConfig

55

56 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

57 -from .utils import (AutoWeightsLoader, init_vllm_registered_model,

58 - maybe_prefix, merge_multimodal_embeddings)

57 +from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix

58

59 _Tuple2 = Union[int, tuple[int, int], Sequence[int]]

60

@@ -744,21 +743,6 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

743 return []

744 return self._process_audio_input(audio_input)

745

747 - def get_input_embeddings(

748 - self,

749 - input_ids: torch.Tensor,

750 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

751 - ) -> torch.Tensor:

752 - inputs_embeds = self.decoder.get_input_embeddings(input_ids)

753 - if multimodal_embeddings and len(multimodal_embeddings) > 0:

754 - inputs_embeds = merge_multimodal_embeddings(

755 - input_ids,

756 - inputs_embeds,

757 - multimodal_embeddings,

758 - self.config.audio_token_id,

759 - )

760 - return inputs_embeds

761 -

746 def forward(

747 self,

748 input_ids: torch.Tensor,

@@ -771,8 +755,11 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

755 inputs_embeds = None

756 elif inputs_embeds is None:

757 multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

774 - inputs_embeds = self.get_input_embeddings(input_ids,

775 - multimodal_embeddings)

758 + inputs_embeds = self.get_input_embeddings(

759 + input_ids,

760 + multimodal_embeddings,

761 + is_multimodal=input_ids == self.config.audio_token_id,

762 + )

763 input_ids = None

764

765 return self.decoder.model(input_ids,

vllm/model_executor/models/mimo_mtp.py

+6

@@ -117,6 +117,9 @@ class MiMoMultiTokenPredictor(nn.Module):

117

118 self.logits_processor = LogitsProcessor(config.vocab_size)

119

120 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

121 + return self.embed_tokens(input_ids)

122 +

123 def forward(

124 self,

125 input_ids: torch.Tensor,

@@ -158,6 +161,9 @@ class MiMoMTP(nn.Module):

161 self.config.hidden_size,

162 prefix=maybe_prefix(prefix, "lm_head"))

163

164 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

165 + return self.model.get_input_embeddings(input_ids)

166 +

167 def forward(

168 self,

169 input_ids: torch.Tensor,

vllm/model_executor/models/minicpmv.py

+6 -21

@@ -71,8 +71,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

71 from .idefics2_vision_model import Idefics2VisionTransformer

72 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

73 SupportsMultiModal, SupportsPP)

74 -from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,

75 - merge_multimodal_embeddings)

74 +from .utils import AutoWeightsLoader, flatten_bn, isin_list, maybe_prefix

75

76 # For profile run

77 _MAX_FRAMES_PER_VIDEO = 16

@@ -1144,23 +1143,6 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1143

1144 return self._process_multimodal_inputs(modalities)

1145

1147 - def get_input_embeddings(

1148 - self,

1149 - input_ids: torch.Tensor,

1150 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1151 - ) -> torch.Tensor:

1152 - inputs_embeds = self.llm.get_input_embeddings(input_ids)

1153 - if multimodal_embeddings is not None \

1154 - and len(multimodal_embeddings) != 0:

1155 - assert len(self.mm_token_ids) > 0

1156 - inputs_embeds = merge_multimodal_embeddings(

1157 - input_ids,

1158 - inputs_embeds,

1159 - multimodal_embeddings,

1160 - list(self.mm_token_ids),

1161 - )

1162 - return inputs_embeds

1163 -

1146 def forward(

1147 self,

1148 input_ids: torch.Tensor,

@@ -1178,8 +1160,11 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1160 elif inputs_embeds is None:

1161 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

1162

1181 - inputs_embeds = self.get_input_embeddings(input_ids,

1182 - vision_embeddings)

1163 + inputs_embeds = self.get_input_embeddings(

1164 + input_ids,

1165 + vision_embeddings,

1166 + is_multimodal=isin_list(input_ids, list(self.mm_token_ids)),

1167 + )

1168 input_ids = None

1169

1170 hidden_states = self.llm.model(

vllm/model_executor/models/minimax_text_01.py

+2 -8

@@ -592,10 +592,7 @@ class MiniMaxText01Model(nn.Module):

592 dtype=torch.long)

593 minimax_cache_tensors[:, slots_tensor, ...] = 0

594

595 - def get_input_embeddings(

596 - self,

597 - input_ids: torch.Tensor,

598 - ) -> torch.Tensor:

595 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

596 return self.embed_tokens(input_ids)

597

598 def forward(self,

@@ -687,10 +684,7 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):

684 return self.model.minimax_cache.get_seqlen_agnostic_capture_inputs(

685 batch_size)

686

690 - def get_input_embeddings(

691 - self,

692 - input_ids: torch.Tensor,

693 - ) -> torch.Tensor:

687 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

688 return self.model.get_input_embeddings(input_ids)

689

690 def forward(self,

vllm/model_executor/models/minimax_vl_01.py

+6 -19

@@ -28,7 +28,7 @@ from .llava_next import LlavaNextProcessingInfo

28 from .pixtral import PixtralHFVisionModel

29 from .siglip import SiglipVisionModel

30 from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,

31 - maybe_prefix, merge_multimodal_embeddings)

31 + maybe_prefix)

32

33

34 class MiniMaxVL01ImagePixelInputs(TensorSchema):

@@ -218,22 +218,6 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal,

218 self.make_empty_intermediate_tensors = (

219 self.language_model.make_empty_intermediate_tensors)

220

221 - def get_input_embeddings(

222 - self,

223 - input_ids: torch.Tensor,

224 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

225 - ) -> torch.Tensor:

226 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

227 - if multimodal_embeddings is not None \

228 - and len(multimodal_embeddings) != 0:

229 - inputs_embeds = merge_multimodal_embeddings(

230 - input_ids,

231 - inputs_embeds,

232 - multimodal_embeddings,

233 - self.config.image_token_index,

234 - )

235 - return inputs_embeds

236 -

221 def get_language_model(self) -> torch.nn.Module:

222 return self.language_model

223

@@ -403,8 +387,11 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal,

387 inputs_embeds = None

388 elif inputs_embeds is None:

389 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

406 - inputs_embeds = self.get_input_embeddings(input_ids,

407 - vision_embeddings)

390 + inputs_embeds = self.get_input_embeddings(

391 + input_ids,

392 + vision_embeddings,

393 + is_multimodal=input_ids == self.config.image_token_index,

394 + )

395 input_ids = None

396

397 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/mistral3.py

+6 -20

@@ -38,8 +38,7 @@ from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

38 SupportsMultiModal, SupportsPP)

39 from .pixtral import PixtralHFEncoderInfo, PixtralHFVisionModel

40 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

41 - init_vllm_registered_model, maybe_prefix,

42 - merge_multimodal_embeddings)

41 + init_vllm_registered_model, maybe_prefix)

42 from .vision import get_vision_encoder_info

43

44

@@ -524,22 +523,6 @@ class Mistral3ForConditionalGeneration(nn.Module, SupportsLoRA,

523

524 return vision_embeddings

525

527 - def get_input_embeddings(

528 - self,

529 - input_ids: torch.Tensor,

530 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

531 - ) -> torch.Tensor:

532 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

533 - if multimodal_embeddings is not None \

534 - and len(multimodal_embeddings) != 0:

535 - inputs_embeds = merge_multimodal_embeddings(

536 - input_ids,

537 - inputs_embeds,

538 - multimodal_embeddings,

539 - self.config.image_token_index,

540 - )

541 - return inputs_embeds

542 -

526 def forward(

527 self,

528 input_ids: torch.Tensor,

@@ -592,8 +575,11 @@ class Mistral3ForConditionalGeneration(nn.Module, SupportsLoRA,

575 # condition is for v0 compatibility.

576 elif inputs_embeds is None:

577 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

595 - inputs_embeds = self.get_input_embeddings(input_ids,

596 - vision_embeddings)

578 + inputs_embeds = self.get_input_embeddings(

579 + input_ids,

580 + vision_embeddings,

581 + is_multimodal=input_ids == self.config.image_token_index,

582 + )

583 input_ids = None

584

585 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/mllama4.py

+7 -23

@@ -42,7 +42,7 @@ from vllm.model_executor.model_loader.utils import initialize_model

42 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

43 from vllm.multimodal import MULTIMODAL_REGISTRY

44 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

45 - MultiModalKwargsItems, NestedTensors)

45 + MultiModalKwargsItems)

46 from vllm.multimodal.parse import (ImageProcessorItems, ImageSize,

47 MultiModalDataItems)

48 from vllm.multimodal.processing import (BaseMultiModalProcessor,

@@ -56,8 +56,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

56

57 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

58 from .llama4 import Llama4ForCausalLM

59 -from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,

60 - merge_multimodal_embeddings)

59 +from .utils import AutoWeightsLoader, flatten_bn, maybe_prefix

60 from .vision import run_dp_sharded_vision_model

61

62

@@ -813,24 +812,6 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,

812

813 return self._process_image_input(image_input)

814

816 - def get_input_embeddings(

817 - self,

818 - input_ids: torch.Tensor,

819 - multimodal_embeddings: Optional[NestedTensors] = None,

820 - ) -> torch.Tensor:

821 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

822 -

823 - if multimodal_embeddings is not None and len(

824 - multimodal_embeddings) != 0:

825 - inputs_embeds = merge_multimodal_embeddings(

826 - input_ids,

827 - inputs_embeds,

828 - multimodal_embeddings,

829 - self.config.image_token_index,

830 - )

831 -

832 - return inputs_embeds

833 -

815 def forward(

816 self,

817 input_ids: torch.Tensor,

@@ -846,8 +827,11 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,

827 # this condition is for v0 compatibility.

828 elif inputs_embeds is None:

829 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

849 - inputs_embeds = self.get_input_embeddings(input_ids,

850 - vision_embeddings)

830 + inputs_embeds = self.get_input_embeddings(

831 + input_ids,

832 + vision_embeddings,

833 + is_multimodal=input_ids == self.config.image_token_index,

834 + )

835 input_ids = None

836

837 return self.language_model(input_ids, positions, intermediate_tensors,

vllm/model_executor/models/modernbert.py

+9

@@ -43,6 +43,9 @@ class ModernBertEmbeddings(nn.Module):

43 eps=config.layer_norm_eps,

44 bias=config.norm_bias)

45

46 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

47 + return self.tok_embeddings(input_ids)

48 +

49 def forward(

50 self,

51 input_ids: torch.Tensor,

@@ -220,6 +223,9 @@ class ModernBertModel(nn.Module):

223 eps=config.norm_eps,

224 bias=config.norm_bias)

225

226 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

227 + return self.embeddings.get_input_embeddings(input_ids)

228 +

229 def load_weights(self, weights: Iterable[tuple[str,

230 torch.Tensor]]) -> set[str]:

231 weights = self.hf_to_vllm_mapper.apply(weights)

@@ -333,6 +339,9 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):

339 ),

340 })

341

342 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

343 + return self.model.get_input_embeddings(input_ids)

344 +

345 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

346

347 self_weights = []

vllm/model_executor/models/molmo.py

+7 -25

@@ -58,7 +58,7 @@ from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

58 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

59 is_pp_missing_parameter,

60 make_empty_intermediate_tensors_factory, make_layers,

61 - maybe_prefix, merge_multimodal_embeddings)

61 + maybe_prefix)

62

63 # TODO: hard-coded for now. Consider making it configurable.

64 VIT_LAYERS = [-2, -9]

@@ -819,10 +819,7 @@ class MolmoModel(nn.Module, SupportsQuant):

819 make_empty_intermediate_tensors_factory(

820 ["hidden_states", "residual"], config.hidden_size))

821

822 - def get_input_embeddings(

823 - self,

824 - input_ids: torch.Tensor,

825 - ) -> torch.Tensor:

822 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

823 return self.embed_tokens(input_ids)

824

825 def forward(

@@ -1481,24 +1478,6 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1478

1479 return self._process_image_input(image_input)

1480

1484 - def get_input_embeddings(

1485 - self,

1486 - input_ids: torch.Tensor,

1487 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1488 - ) -> torch.Tensor:

1489 - inputs_embeds = self.model.get_input_embeddings(input_ids)

1490 - if multimodal_embeddings is not None \

1491 - and len(multimodal_embeddings) != 0:

1492 - assert self.img_patch_id is not None

1493 -

1494 - inputs_embeds = merge_multimodal_embeddings(

1495 - input_ids,

1496 - inputs_embeds,

1497 - multimodal_embeddings,

1498 - self.img_patch_id,

1499 - )

1500 - return inputs_embeds

1501 -

1481 def forward(

1482 self,

1483 input_ids: torch.LongTensor,

@@ -1515,8 +1494,11 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1494 # condition is for v0 compatibility.

1495 elif inputs_embeds is None:

1496 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

1518 - inputs_embeds = self.get_input_embeddings(input_ids,

1519 - vision_embeddings)

1497 + inputs_embeds = self.get_input_embeddings(

1498 + input_ids,

1499 + vision_embeddings,

1500 + is_multimodal=input_ids == self.img_patch_id,

1501 + )

1502 input_ids = None

1503

1504 hidden_states = self.model(input_ids,

vllm/model_executor/models/nano_nemotron_vl.py

+13 -30

@@ -35,8 +35,7 @@ from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM

35 from vllm.model_executor.models.radio import RadioModel

36 from vllm.model_executor.models.utils import (flatten_bn,

37 init_vllm_registered_model,

38 - maybe_prefix,

39 - merge_multimodal_embeddings)

38 + isin_list, maybe_prefix)

39 from vllm.multimodal import MULTIMODAL_REGISTRY

40 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

41 MultiModalKwargs, MultiModalKwargsItems,

@@ -1096,8 +1095,8 @@ class NemotronH_Nano_VL_V2(nn.Module, HasInnerState, IsHybrid,

1095

1096 return modalities

1097

1099 - def get_multimodal_embeddings(

1100 - self, **kwargs: object) -> Optional[MultiModalEmbeddings]:

1098 + def get_multimodal_embeddings(self,

1099 + **kwargs: object) -> MultiModalEmbeddings:

1100 # Validate the multimodal input keyword arguments

1101 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)

1102 if modalities is None:

@@ -1121,30 +1120,6 @@ class NemotronH_Nano_VL_V2(nn.Module, HasInnerState, IsHybrid,

1120

1121 return multimodal_embeddings

1122

1124 - def get_input_embeddings(

1125 - self,

1126 - input_ids: torch.Tensor,

1127 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1128 - ) -> torch.Tensor:

1129 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1130 -

1131 - if (multimodal_embeddings is not None

1132 - and len(multimodal_embeddings) != 0):

1133 - context_token_ids = [

1134 - token_id for token_id in (self.img_context_token_id,

1135 - self.video_context_token_id)

1136 - if token_id is not None

1137 - ]

1138 - assert len(context_token_ids) >= 1

1139 - inputs_embeds = merge_multimodal_embeddings(

1140 - input_ids,

1141 - inputs_embeds,

1142 - multimodal_embeddings,

1143 - context_token_ids,

1144 - )

1145 -

1146 - return inputs_embeds

1147 -

1123 def get_language_model(self) -> torch.nn.Module:

1124 return self.language_model

1125

@@ -1163,9 +1138,17 @@ class NemotronH_Nano_VL_V2(nn.Module, HasInnerState, IsHybrid,

1138 # NOTE: In v1, inputs_embeds is always generated at model runner, this

1139 # condition is for v0 compatibility.

1140 elif inputs_embeds is None:

1141 + context_token_ids = [

1142 + token_id for token_id in (self.img_context_token_id,

1143 + self.video_context_token_id)

1144 + if token_id is not None

1145 + ]

1146 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

1167 - inputs_embeds = self.get_input_embeddings(input_ids,

1168 - vision_embeddings)

1147 + inputs_embeds = self.get_input_embeddings(

1148 + input_ids,

1149 + vision_embeddings,

1150 + is_multimodal=isin_list(input_ids, context_token_ids),

1151 + )

1152 input_ids = None

1153

1154 hidden_states = self.language_model(

vllm/model_executor/models/nemotron_vl.py

+22 -15

@@ -38,7 +38,7 @@ from vllm.transformers_utils.tokenizer import AnyTokenizer

38 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

39 SupportsMultiModal, SupportsPP)

40 from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,

41 - maybe_prefix, merge_multimodal_embeddings)

41 + maybe_prefix)

42

43 IMG_START = '<img>'

44 IMG_END = '</img>'

@@ -576,20 +576,24 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,

576 self,

577 input_ids: torch.Tensor,

578 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

579 + *,

580 + is_multimodal: Optional[torch.Tensor] = None,

581 + handle_oov_mm_token: bool = False,

582 ) -> torch.Tensor:

580 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

581 - if multimodal_embeddings is not None \

582 - and len(multimodal_embeddings) != 0:

583 - context_token_ids = [self.img_context_token_id]

584 - assert len(context_token_ids) >= 1

583 + if multimodal_embeddings is not None and len(

584 + multimodal_embeddings) > 0:

585 self._set_visual_token_mask(input_ids)

586 - inputs_embeds = merge_multimodal_embeddings(

587 - input_ids,

588 - inputs_embeds,

589 - multimodal_embeddings,

590 - context_token_ids,

591 - )

592 - return inputs_embeds

586 +

587 + # This is to satisfy the type checker for each overload

588 + if multimodal_embeddings is None or is_multimodal is None:

589 + return super().get_input_embeddings(input_ids)

590 +

591 + return super().get_input_embeddings(

592 + input_ids,

593 + multimodal_embeddings=multimodal_embeddings,

594 + is_multimodal=is_multimodal,

595 + handle_oov_mm_token=handle_oov_mm_token,

596 + )

597

598 def forward(

599 self,

@@ -608,8 +612,11 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,

612 # condition is for v0 compatibility.

613 elif inputs_embeds is None:

614 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

611 - inputs_embeds = self.get_input_embeddings(input_ids,

612 - vision_embeddings)

615 + inputs_embeds = self.get_input_embeddings(

616 + input_ids,

617 + vision_embeddings,

618 + is_multimodal=input_ids == self.img_context_token_id,

619 + )

620 input_ids = None

621

622 forward_kwargs = {

vllm/model_executor/models/olmo2.py

+6

@@ -295,6 +295,9 @@ class Olmo2Model(nn.Module):

295 make_empty_intermediate_tensors_factory(["hidden_states"],

296 self.config.hidden_size))

297

298 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

299 + return self.embed_tokens(input_ids)

300 +

301 def forward(

302 self,

303 input_ids: torch.Tensor,

@@ -408,6 +411,9 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

411 self.make_empty_intermediate_tensors = (

412 self.model.make_empty_intermediate_tensors)

413

414 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

415 + return self.model.get_input_embeddings(input_ids)

416 +

417 def forward(

418 self,

419 input_ids: torch.Tensor,

vllm/model_executor/models/ovis.py

+5 -16

@@ -48,7 +48,6 @@ from vllm.transformers_utils.processors.ovis import OvisProcessor

48 from vllm.utils.tensor_schema import TensorSchema, TensorShape

49

50 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

51 -from .utils import merge_multimodal_embeddings

51

52 # Cannot find the following number from hf config.

53 IMAGE_TOKEN = "<image>"

@@ -501,19 +500,6 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):

500

501 return image_features

502

504 - def get_input_embeddings(

505 - self,

506 - input_ids: torch.Tensor,

507 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

508 - ) -> torch.Tensor:

509 - inputs_embeds = self.llm.get_input_embeddings(input_ids)

510 - if multimodal_embeddings is not None \

511 - and len(multimodal_embeddings) != 0:

512 - inputs_embeds = merge_multimodal_embeddings(

513 - input_ids, inputs_embeds, multimodal_embeddings,

514 - self.image_pad_token_id)

515 - return inputs_embeds

516 -

503 def forward(

504 self,

505 input_ids: torch.Tensor,

@@ -529,8 +515,11 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):

515 # condition is for v0 compatibility.

516 elif inputs_embeds is None:

517 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

532 - inputs_embeds = self.get_input_embeddings(input_ids,

533 - vision_embeddings)

518 + inputs_embeds = self.get_input_embeddings(

519 + input_ids,

520 + vision_embeddings,

521 + is_multimodal=input_ids == self.image_pad_token_id,

522 + )

523 input_ids = None

524

525 # up until here we have an inputs_embeds 100% numerical identity

vllm/model_executor/models/ovis2_5.py

+5 -13

@@ -585,17 +585,6 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

585

586 return multimodal_embeddings

587

588 - def get_input_embeddings(

589 - self,

590 - input_ids: torch.Tensor,

591 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

592 - ) -> torch.Tensor:

593 - inputs_embeds = self.llm.get_input_embeddings(input_ids)

594 - if multimodal_embeddings is not None:

595 - tmp = torch.concat(multimodal_embeddings, dim=0)

596 - inputs_embeds[input_ids == self.image_pad_token_id] = tmp

597 - return inputs_embeds

598 -

588 def forward(

589 self,

590 input_ids: torch.Tensor,

@@ -612,8 +601,11 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

601 elif inputs_embeds is None:

602 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

603

615 - inputs_embeds = self.get_input_embeddings(input_ids,

616 - vision_embeddings)

604 + inputs_embeds = self.get_input_embeddings(

605 + input_ids,

606 + vision_embeddings,

607 + is_multimodal=input_ids == self.image_pad_token_id,

608 + )

609 input_ids = None

610

611 # up until here we have a inputs_embeds 100% numerical identity

vllm/model_executor/models/paligemma.py

+6 -17

@@ -26,8 +26,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

26 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

27 from .siglip import SiglipVisionModel

28 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

29 - init_vllm_registered_model, maybe_prefix,

30 - merge_multimodal_embeddings)

29 + init_vllm_registered_model, maybe_prefix)

30 from .vision import get_vision_encoder_info

31

32 logger = init_logger(__name__)

@@ -362,19 +361,6 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,

361 vision_embeddings = vision_embeddings * (self.config.hidden_size**-0.5)

362 return vision_embeddings

363

365 - def get_input_embeddings(

366 - self,

367 - input_ids: torch.Tensor,

368 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

369 - ) -> torch.Tensor:

370 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

371 - if multimodal_embeddings is not None \

372 - and len(multimodal_embeddings) != 0:

373 - inputs_embeds = merge_multimodal_embeddings(

374 - input_ids, inputs_embeds, multimodal_embeddings,

375 - self.config.image_token_index)

376 - return inputs_embeds

377 -

364 def forward(self,

365 input_ids: torch.Tensor,

366 positions: torch.Tensor,

@@ -388,8 +374,11 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,

374 # condition is for v0 compatibility.

375 elif inputs_embeds is None:

376 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

391 - inputs_embeds = self.get_input_embeddings(input_ids,

392 - vision_embeddings)

377 + inputs_embeds = self.get_input_embeddings(

378 + input_ids,

379 + vision_embeddings,

380 + is_multimodal=input_ids == self.config.image_token_index,

381 + )

382 input_ids = None

383

384 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/phi3v.py

+32 -12

@@ -51,9 +51,9 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

51 from .clip import CLIPVisionModel

52 from .interfaces import (MultiModalEmbeddings, SupportsMultiModal, SupportsPP,

53 SupportsQuant)

54 -from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

55 - init_vllm_registered_model, maybe_prefix,

56 - merge_multimodal_embeddings)

54 +from .utils import (AutoWeightsLoader, WeightsMapper,

55 + _merge_multimodal_embeddings, flatten_bn,

56 + init_vllm_registered_model, maybe_prefix)

57

58 logger = init_logger(__name__)

59

@@ -643,14 +643,31 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

643 self,

644 input_ids: torch.Tensor,

645 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

646 + *,

647 + is_multimodal: Optional[torch.Tensor] = None,

648 + handle_oov_mm_token: bool = False,

649 ) -> torch.Tensor:

647 - inputs_embeds = self.embed_tokens(input_ids)

648 - if multimodal_embeddings is not None \

649 - and len(multimodal_embeddings) != 0:

650 - inputs_embeds = merge_multimodal_embeddings(

651 - input_ids, inputs_embeds, multimodal_embeddings,

652 - self.image_token_id)

653 - return inputs_embeds

650 + inputs_embeds = self._get_text_embeddings(

651 + input_ids,

652 + self.embed_tokens,

653 + is_multimodal=is_multimodal,

654 + handle_oov_mm_token=handle_oov_mm_token,

655 + )

656 +

657 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:

658 + return inputs_embeds

659 +

660 + if is_multimodal is None:

661 + raise ValueError(

662 + "`get_input_embeddings` now requires `is_multimodal` arg, "

663 + "please update your model runner according to "

664 + "https://github.com/vllm-project/vllm/pull/16229.")

665 +

666 + return _merge_multimodal_embeddings(

667 + inputs_embeds=inputs_embeds,

668 + multimodal_embeddings=multimodal_embeddings,

669 + is_multimodal=is_multimodal,

670 + )

671

672 def forward(self,

673 input_ids: torch.Tensor,

@@ -666,8 +683,11 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

683 # condition is for v0 compatibility

684 elif inputs_embeds is None:

685 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

669 - inputs_embeds = self.get_input_embeddings(input_ids,

670 - vision_embeddings)

686 + inputs_embeds = self.get_input_embeddings(

687 + input_ids,

688 + vision_embeddings,

689 + is_multimodal=self.image_token_id,

690 + )

691 input_ids = None

692

693 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/phi4_multimodal.py

+3 -15

@@ -1342,12 +1342,12 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1342 image_attention_mask)

1343 return image_embeds

1344

1345 - def get_multimodal_embeddings(

1346 - self, **kwargs: object) -> Optional[MultiModalEmbeddings]:

1345 + def get_multimodal_embeddings(self,

1346 + **kwargs: object) -> MultiModalEmbeddings:

1347

1348 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)

1349 if not modalities:

1350 - return None

1350 + return []

1351

1352 # The result multimodal_embeddings is tuple of tensors, with each

1353 # tensor corresponding to a multimodal data item (image or video).

@@ -1371,18 +1371,6 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1371

1372 return multimodal_embeddings

1373

1374 - def get_input_embeddings(

1375 - self,

1376 - input_ids: torch.Tensor,

1377 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1378 - ) -> torch.Tensor:

1379 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1380 - if multimodal_embeddings is not None:

1381 - inputs_embeds = merge_multimodal_embeddings(

1382 - input_ids, inputs_embeds, multimodal_embeddings,

1383 - [_IMAGE_PLACEHOLDER_TOKEN_ID, _AUDIO_PLACEHOLDER_TOKEN_ID])

1384 - return inputs_embeds

1385 -

1374 def get_input_embeddings_v0(

1375 self,

1376 input_ids: torch.Tensor,

vllm/model_executor/models/phi4mm.py

-14

@@ -1151,7 +1151,6 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1151 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)

1152 if not modalities:

1153 return []

1154 - return None

1154

1155 # The result multimodal_embeddings is tuple of tensors, with each

1156 # tensor corresponding to a multimodal data item (image or video).

@@ -1175,19 +1174,6 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1174

1175 return multimodal_embeddings

1176

1178 - def get_input_embeddings(

1179 - self,

1180 - input_ids: torch.Tensor,

1181 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1182 - ) -> torch.Tensor:

1183 - inputs_embeds = self.model.embed_tokens(input_ids)

1184 - if multimodal_embeddings is not None and len(

1185 - multimodal_embeddings) != 0:

1186 - inputs_embeds = merge_multimodal_embeddings(

1187 - input_ids, inputs_embeds, multimodal_embeddings,

1188 - [_IMAGE_PLACEHOLDER_TOKEN_ID, _AUDIO_PLACEHOLDER_TOKEN_ID])

1189 - return inputs_embeds

1190 -

1177 def get_input_embeddings_v0(

1178 self,

1179 input_ids: torch.Tensor,

vllm/model_executor/models/pixtral.py

+6 -20

@@ -50,8 +50,7 @@ from vllm.transformers_utils.tokenizer import (MistralTokenizer,

50 from vllm.utils.tensor_schema import TensorSchema, TensorShape

51

52 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

53 -from .utils import (flatten_bn, init_vllm_registered_model, maybe_prefix,

54 - merge_multimodal_embeddings)

53 +from .utils import flatten_bn, init_vllm_registered_model, maybe_prefix

54 from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs

55

56 try:

@@ -433,22 +432,6 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

432

433 return self._process_image_input(image_input)

434

436 - def get_input_embeddings(

437 - self,

438 - input_ids: torch.Tensor,

439 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

440 - ) -> torch.Tensor:

441 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

442 - if multimodal_embeddings is not None \

443 - and len(multimodal_embeddings) != 0:

444 - inputs_embeds = merge_multimodal_embeddings(

445 - input_ids,

446 - inputs_embeds,

447 - multimodal_embeddings,

448 - self.vision_args.image_token_id,

449 - )

450 - return inputs_embeds

451 -

435 def forward(

436 self,

437 input_ids: torch.Tensor,

@@ -465,8 +448,11 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

448 # condition is for v0 compatibility.

449 elif inputs_embeds is None:

450 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

468 - inputs_embeds = self.get_input_embeddings(input_ids,

469 - vision_embeddings)

451 + inputs_embeds = self.get_input_embeddings(

452 + input_ids,

453 + vision_embeddings,

454 + is_multimodal=input_ids == self.vision_args.image_token_id,

455 + )

456 input_ids = None

457

458 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/qwen2_5_omni_thinker.py

+15 -13

@@ -865,24 +865,26 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

865 multimodal_embeddings += audio_embeddings

866 return multimodal_embeddings

867

868 + # TODO (ywang96): support overlapping modality embeddings so that

869 + # `use_audio_in_video` will work on V1.

870 def get_input_embeddings(

871 self,

872 input_ids: torch.Tensor,

873 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

874 + *,

875 + is_multimodal: Optional[torch.Tensor] = None,

876 + handle_oov_mm_token: bool = False,

877 ) -> torch.Tensor:

873 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

874 - if multimodal_embeddings is not None \

875 - and len(multimodal_embeddings) != 0:

876 -

877 - # TODO (ywang96): support overlapping modality embeddings so that

878 - # `use_audio_in_video` will work on V1.

879 - inputs_embeds = merge_multimodal_embeddings(

880 - input_ids, inputs_embeds, multimodal_embeddings, [

881 - self.config.image_token_index,

882 - self.config.video_token_index,

883 - self.config.audio_token_index

884 - ])

885 - return inputs_embeds

878 + # This is to satisfy the type checker for each overload

879 + if multimodal_embeddings is None or is_multimodal is None:

880 + return super().get_input_embeddings(input_ids)

881 +

882 + return super().get_input_embeddings(

883 + input_ids,

884 + multimodal_embeddings=multimodal_embeddings,

885 + is_multimodal=is_multimodal,

886 + handle_oov_mm_token=handle_oov_mm_token,

887 + )

888

889 def get_multimodal_embeddings_v0(

890 self, **kwargs: object) -> Optional[NestedTensors]:

vllm/model_executor/models/qwen2_5_vl.py

-13

@@ -1365,19 +1365,6 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1365 multimodal_embeddings += video_embeddings

1366 return multimodal_embeddings

1367

1368 - def get_input_embeddings(

1369 - self,

1370 - input_ids: torch.Tensor,

1371 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1372 - ) -> torch.Tensor:

1373 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1374 - if multimodal_embeddings is not None \

1375 - and len(multimodal_embeddings) != 0:

1376 - inputs_embeds = merge_multimodal_embeddings(

1377 - input_ids, inputs_embeds, multimodal_embeddings,

1378 - [self.config.image_token_id, self.config.video_token_id])

1379 - return inputs_embeds

1380 -

1368 def get_input_embeddings_v0(

1369 self,

1370 input_ids: torch.Tensor,

vllm/model_executor/models/qwen2_audio.py

+6 -17

@@ -49,8 +49,7 @@ from vllm.sequence import IntermediateTensors

49 from vllm.utils.tensor_schema import TensorSchema, TensorShape

50

51 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

52 -from .utils import (AutoWeightsLoader, init_vllm_registered_model,

53 - maybe_prefix, merge_multimodal_embeddings)

52 +from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix

53

54

55 # # === Audio Inputs === #

@@ -438,19 +437,6 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal,

437 masked_audio_features = self._process_audio_input(audio_input)

438 return masked_audio_features

439

441 - def get_input_embeddings(

442 - self,

443 - input_ids: torch.Tensor,

444 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

445 - ) -> torch.Tensor:

446 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

447 - if multimodal_embeddings is not None \

448 - and len(multimodal_embeddings) != 0:

449 - inputs_embeds = merge_multimodal_embeddings(

450 - input_ids, inputs_embeds, multimodal_embeddings,

451 - self.config.audio_token_index)

452 - return inputs_embeds

453 -

440 def forward(

441 self,

442 input_ids: torch.Tensor,

@@ -467,8 +453,11 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal,

453 # condition is for v0 compatibility.

454 elif inputs_embeds is None:

455 multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

470 - inputs_embeds = self.get_input_embeddings(input_ids,

471 - multimodal_embeddings)

456 + inputs_embeds = self.get_input_embeddings(

457 + input_ids,

458 + multimodal_embeddings,

459 + is_multimodal=input_ids == self.config.audio_token_index,

460 + )

461 input_ids = None

462

463 hidden_states = self.language_model.model(input_ids,

vllm/model_executor/models/qwen2_vl.py

-13

@@ -1459,19 +1459,6 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1459

1460 return multimodal_embeddings

1461

1462 - def get_input_embeddings(

1463 - self,

1464 - input_ids: torch.Tensor,

1465 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1466 - ) -> torch.Tensor:

1467 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1468 - if multimodal_embeddings is not None \

1469 - and len(multimodal_embeddings) != 0:

1470 - inputs_embeds = merge_multimodal_embeddings(

1471 - input_ids, inputs_embeds, multimodal_embeddings,

1472 - [self.config.image_token_id, self.config.video_token_id])

1473 - return inputs_embeds

1474 -

1462 def get_input_embeddings_v0(

1463 self,

1464 input_ids: torch.Tensor,

vllm/model_executor/models/qwen3_vl.py

+59 -30

@@ -79,7 +79,8 @@ from .qwen2_5_vl import (Qwen2_5_VisionAttention,

79 from .qwen2_vl import Qwen2VLProcessingInfo

80 from .qwen3 import Qwen3ForCausalLM, Qwen3Model

81 from .utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,

82 - maybe_prefix, merge_multimodal_embeddings)

82 + _merge_multimodal_embeddings, maybe_prefix,

83 + merge_multimodal_embeddings)

84 from .vision import get_vit_attn_backend, run_dp_sharded_mrope_vision_model

85

86 logger = init_logger(__name__)

@@ -1324,17 +1325,22 @@ class Qwen3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1325 return multimodal_embeddings

1326

1327 def _compute_deepstack_embeds(

1327 - self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor,

1328 - multimodal_embeddings: MultiModalEmbeddings) -> torch.Tensor:

1329 - visual_lens = [

1330 - x.shape[0] if isinstance(x, torch.Tensor) else len(x)

1331 - for x in multimodal_embeddings

1332 - ]

1328 + self,

1329 + inputs_embeds: torch.Tensor,

1330 + multimodal_embeddings: MultiModalEmbeddings,

1331 + is_multimodal: torch.Tensor,

1332 + ) -> tuple[torch.Tensor, MultiModalEmbeddings]:

1333 + visual_lens = [len(x) for x in multimodal_embeddings]

1334 multimodal_embeddings_cat = torch.cat(multimodal_embeddings, dim=0)

1335

1335 - multimodal_embeddings_main, multimodal_embeddings_multiscale = torch.split( # noqa:E501

1336 - multimodal_embeddings_cat, [self.visual_dim, self.multiscale_dim],

1337 - dim=-1)

1336 + (

1337 + multimodal_embeddings_main,

1338 + multimodal_embeddings_multiscale,

1339 + ) = torch.split(

1340 + multimodal_embeddings_cat,

1341 + [self.visual_dim, self.multiscale_dim],

1342 + dim=-1,

1343 + )

1344

1345 multimodal_embeddings = torch.split(multimodal_embeddings_main,

1346 visual_lens,

@@ -1346,39 +1352,62 @@ class Qwen3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1352 inputs_embeds.size(0),

1353 self.deepstack_num_level * inputs_embeds.size(1))

1354

1349 - deepstack_input_embeds = merge_multimodal_embeddings(

1350 - input_ids,

1351 - deepstack_input_embeds,

1352 - multimodal_embeddings_multiscale,

1353 - placeholder_token_id=[

1354 - self.config.image_token_id, self.config.video_token_id

1355 - ],

1355 + deepstack_input_embeds = _merge_multimodal_embeddings(

1356 + inputs_embeds=deepstack_input_embeds,

1357 + multimodal_embeddings=multimodal_embeddings_multiscale,

1358 + is_multimodal=is_multimodal,

1359 )

1360 deepstack_input_embeds = deepstack_input_embeds.view(

1361 inputs_embeds.shape[0], self.deepstack_num_level, self.visual_dim)

1362 deepstack_input_embeds = deepstack_input_embeds.permute(1, 0, 2)

1363 +

1364 return deepstack_input_embeds, multimodal_embeddings

1365

1366 def get_input_embeddings(

1367 self,

1368 input_ids: torch.Tensor,

1369 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1370 + *,

1371 + is_multimodal: Optional[torch.Tensor] = None,

1372 + handle_oov_mm_token: bool = False,

1373 ) -> torch.Tensor:

1367 - deepstack_input_embeds = None

1368 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1369 - if multimodal_embeddings is not None:

1370 - if self.use_deepstack:

1371 - deepstack_input_embeds, multimodal_embeddings = self._compute_deepstack_embeds( # noqa:E501

1372 - input_ids, inputs_embeds, multimodal_embeddings)

1373 - inputs_embeds = merge_multimodal_embeddings(

1374 - input_ids, inputs_embeds, multimodal_embeddings,

1375 - [self.config.image_token_id, self.config.video_token_id])

1374 + inputs_embeds = self._get_text_embeddings(

1375 + input_ids,

1376 + self.language_model.get_input_embeddings,

1377 + is_multimodal=is_multimodal,

1378 + handle_oov_mm_token=handle_oov_mm_token,

1379 + )

1380 +

1381 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:

1382 + return inputs_embeds

1383 +

1384 + if is_multimodal is None:

1385 + raise ValueError(

1386 + "`get_input_embeddings` now requires `is_multimodal` arg, "

1387 + "please update your model runner according to "

1388 + "https://github.com/vllm-project/vllm/pull/16229.")

1389

1390 if self.use_deepstack:

1378 - if deepstack_input_embeds is None:

1379 - deepstack_input_embeds = torch.zeros_like(

1380 - inputs_embeds).unsqueeze(0).repeat(

1381 - self.deepstack_num_level, 1, 1).contiguous()

1391 + (

1392 + deepstack_input_embeds,

1393 + multimodal_embeddings,

1394 + ) = self._compute_deepstack_embeds(

1395 + inputs_embeds=inputs_embeds,

1396 + multimodal_embeddings=multimodal_embeddings,

1397 + is_multimodal=is_multimodal,

1398 + )

1399 + else:

1400 + deepstack_input_embeds = None

1401 +

1402 + inputs_embeds = _merge_multimodal_embeddings(

1403 + inputs_embeds=inputs_embeds,

1404 + multimodal_embeddings=multimodal_embeddings,

1405 + is_multimodal=is_multimodal,

1406 + )

1407 +

1408 + if deepstack_input_embeds is not None:

1409 + deepstack_input_embeds = torch.zeros_like(inputs_embeds).unsqueeze(

1410 + 0).repeat(self.deepstack_num_level, 1, 1).contiguous()

1411 self._set_deepstack_input_embeds(deepstack_input_embeds)

1412

1413 return inputs_embeds

vllm/model_executor/models/qwen_vl.py

+7 -18

@@ -45,7 +45,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

45 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

46 SupportsMultiModal, SupportsPP)

47 from .qwen import QWenBaseModel, QWenModel

48 -from .utils import flatten_bn, merge_multimodal_embeddings

48 +from .utils import flatten_bn

49

50

51 class QwenImagePixelInputs(TensorSchema):

@@ -756,21 +756,6 @@ class QwenVLForConditionalGeneration(QWenBaseModel, SupportsPP, SupportsLoRA,

756 vision_embeddings = self._process_image_input(image_input)

757 return vision_embeddings

758

759 - def get_input_embeddings(

760 - self,

761 - input_ids: torch.Tensor,

762 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

763 - ) -> torch.Tensor:

764 - inputs_embeds = self.transformer.get_input_embeddings(input_ids)

765 -

766 - if multimodal_embeddings is not None \

767 - and len(multimodal_embeddings) != 0:

768 - inputs_embeds = merge_multimodal_embeddings(

769 - input_ids, inputs_embeds, multimodal_embeddings,

770 - self.transformer.visual.image_pad_id)

771 -

772 - return inputs_embeds

773 -

759 def forward(

760 self,

761 input_ids: torch.Tensor,

@@ -786,8 +771,12 @@ class QwenVLForConditionalGeneration(QWenBaseModel, SupportsPP, SupportsLoRA,

771 # condition is for v0 compatibility.

772 elif inputs_embeds is None:

773 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

789 - inputs_embeds = self.get_input_embeddings(input_ids,

790 - vision_embeddings)

774 + inputs_embeds = self.get_input_embeddings(

775 + input_ids,

776 + vision_embeddings,

777 + is_multimodal=input_ids ==

778 + self.transformer.visual.image_pad_id,

779 + )

780 input_ids = None

781

782 hidden_states = self.transformer(input_ids, positions,

vllm/model_executor/models/roberta.py

+3

@@ -218,6 +218,9 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):

218 loader = AutoWeightsLoader(self)

219 return loader.load_weights(weights, mapper=self.jina_to_vllm_mapper)

220

221 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

222 + return self.roberta.get_input_embeddings(input_ids)

223 +

224 def forward(

225 self,

226 input_ids: Optional[torch.Tensor],

vllm/model_executor/models/skyworkr1v.py

+22 -14

@@ -38,7 +38,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

38

39 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

40 from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,

41 - maybe_prefix, merge_multimodal_embeddings)

41 + maybe_prefix)

42

43 IMG_START = '<img>'

44 IMG_END = '</img>'

@@ -842,19 +842,24 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

842 self,

843 input_ids: torch.Tensor,

844 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

845 + *,

846 + is_multimodal: Optional[torch.Tensor] = None,

847 + handle_oov_mm_token: bool = False,

848 ) -> torch.Tensor:

846 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

847 - if multimodal_embeddings is not None \

848 - and len(multimodal_embeddings) != 0:

849 - assert self.img_context_token_id is not None

849 + if multimodal_embeddings is not None and len(

850 + multimodal_embeddings) > 0:

851 self._set_visual_token_mask(input_ids)

851 - inputs_embeds = merge_multimodal_embeddings(

852 - input_ids,

853 - inputs_embeds,

854 - multimodal_embeddings,

855 - self.img_context_token_id,

856 - )

857 - return inputs_embeds

852 +

853 + # This is to satisfy the type checker for each overload

854 + if multimodal_embeddings is None or is_multimodal is None:

855 + return super().get_input_embeddings(input_ids)

856 +

857 + return super().get_input_embeddings(

858 + input_ids,

859 + multimodal_embeddings=multimodal_embeddings,

860 + is_multimodal=is_multimodal,

861 + handle_oov_mm_token=handle_oov_mm_token,

862 + )

863

864 def forward(

865 self,

@@ -873,8 +878,11 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

878 # condition is for v0 compatibility.

879 elif inputs_embeds is None:

880 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

876 - inputs_embeds = self.get_input_embeddings(input_ids,

877 - vision_embeddings)

881 + inputs_embeds = self.get_input_embeddings(

882 + input_ids,

883 + vision_embeddings,

884 + is_multimodal=input_ids == self.img_context_token_id,

885 + )

886 input_ids = None

887

888 forward_kwargs = {

vllm/model_executor/models/solar.py

+3

@@ -483,6 +483,9 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

483 self.make_empty_intermediate_tensors = (

484 self.model.make_empty_intermediate_tensors)

485

486 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

487 + return self.model.get_input_embeddings(input_ids)

488 +

489 def forward(

490 self,

491 input_ids: torch.Tensor,

vllm/model_executor/models/step3_text.py

+3

@@ -395,6 +395,9 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):

395 self.make_empty_intermediate_tensors = (

396 self.model.make_empty_intermediate_tensors)

397

398 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

399 + return self.model.get_input_embeddings(input_ids)

400 +

401 def forward(self,

402 input_ids: torch.Tensor,

403 positions: torch.Tensor,

vllm/model_executor/models/step3_vl.py

+26 -26

@@ -25,7 +25,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

25 from vllm.model_executor.layers.quantization import QuantizationConfig

26 from vllm.multimodal import MULTIMODAL_REGISTRY

27 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

28 - MultiModalKwargsItems, NestedTensors)

28 + MultiModalKwargsItems)

29 from vllm.multimodal.parse import ImageSize, MultiModalDataItems

30 from vllm.multimodal.processing import (BaseMultiModalProcessor,

31 BaseProcessingInfo, PromptReplacement,

@@ -37,8 +37,7 @@ from vllm.transformers_utils.tokenizer import AnyTokenizer

37

38 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

39 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

40 - init_vllm_registered_model, maybe_prefix,

41 - merge_multimodal_embeddings)

40 + init_vllm_registered_model, maybe_prefix)

41 from .vision import run_dp_sharded_vision_model

42

43

@@ -996,10 +995,13 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

995 1 else cur_feature[0])

996 return merged_image_features

997

999 - def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:

998 + def get_language_model(self) -> torch.nn.Module:

999 + return self.language_model

1000 +

1001 + def get_multimodal_embeddings(self, **kwargs) -> MultiModalEmbeddings:

1002 image_input = self._parse_and_validate_image_input(**kwargs)

1003 if image_input is None:

1002 - return None

1004 + return []

1005 vision_embeddings = self._process_image_input(image_input)

1006 return vision_embeddings

1007

@@ -1007,24 +1009,21 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1009 self,

1010 input_ids: torch.Tensor,

1011 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1012 + *,

1013 + is_multimodal: Optional[torch.Tensor] = None,

1014 + # Multi-modal token ID may exceed vocab size

1015 + handle_oov_mm_token: bool = True,

1016 ) -> torch.Tensor:

1011 - if multimodal_embeddings is None:

1012 - inputs_embeds = self.language_model.model.get_input_embeddings(

1013 - input_ids)

1014 - else:

1015 - is_text = input_ids != self.config.image_token_id

1016 - text_ids = input_ids[is_text]

1017 - text_embeds = self.language_model.model.get_input_embeddings(

1018 - text_ids)

1019 - inputs_embeds = torch.empty(input_ids.shape[0],

1020 - text_embeds.shape[-1],

1021 - dtype=text_embeds.dtype,

1022 - device=text_embeds.device)

1023 - inputs_embeds[is_text] = text_embeds

1024 - inputs_embeds = merge_multimodal_embeddings(

1025 - input_ids, inputs_embeds, multimodal_embeddings,

1026 - self.config.image_token_id)

1027 - return inputs_embeds

1017 + # This is to satisfy the type checker for each overload

1018 + if multimodal_embeddings is None or is_multimodal is None:

1019 + return super().get_input_embeddings(input_ids)

1020 +

1021 + return super().get_input_embeddings(

1022 + input_ids,

1023 + multimodal_embeddings=multimodal_embeddings,

1024 + is_multimodal=is_multimodal,

1025 + handle_oov_mm_token=handle_oov_mm_token,

1026 + )

1027

1028 def forward(

1029 self,

@@ -1038,10 +1037,11 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1037 inputs_embeds = None

1038 elif inputs_embeds is None:

1039 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

1041 - # always pass the input via `inputs_embeds`

1042 - # to make sure the computation graph is consistent

1043 - inputs_embeds = self.get_input_embeddings(input_ids,

1044 - vision_embeddings)

1040 + inputs_embeds = self.get_input_embeddings(

1041 + input_ids,

1042 + vision_embeddings,

1043 + is_multimodal=input_ids == self.config.image_token_id,

1044 + )

1045 input_ids = None

1046

1047 hidden_states = self.language_model(input_ids,

vllm/model_executor/models/tarsier.py

+6 -19

@@ -40,7 +40,7 @@ from .clip import CLIPVisionModel

40 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

41 from .siglip import SiglipVisionModel

42 from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,

43 - maybe_prefix, merge_multimodal_embeddings)

43 + maybe_prefix)

44 from .vision import VisionEncoderInfo, get_vision_encoder_info

45

46

@@ -589,22 +589,6 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal,

589 return []

590 return self._process_image_input(image_input)

591

592 - def get_input_embeddings(

593 - self,

594 - input_ids: torch.Tensor,

595 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

596 - ) -> torch.Tensor:

597 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

598 - if multimodal_embeddings is not None \

599 - and len(multimodal_embeddings) != 0:

600 - inputs_embeds = merge_multimodal_embeddings(

601 - input_ids,

602 - inputs_embeds,

603 - multimodal_embeddings,

604 - self.config.image_token_index,

605 - )

606 - return inputs_embeds

607 -

592 def forward(

593 self,

594 input_ids: torch.Tensor,

@@ -617,8 +601,11 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal,

601 inputs_embeds = None

602 elif inputs_embeds is None:

603 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

620 - inputs_embeds = self.get_input_embeddings(input_ids,

621 - vision_embeddings)

604 + inputs_embeds = self.get_input_embeddings(

605 + input_ids,

606 + vision_embeddings,

607 + is_multimodal=input_ids == self.config.image_token_index,

608 + )

609 input_ids = None

610 hidden_states = self.language_model.model(

611 input_ids=input_ids,

vllm/model_executor/models/terratorch.py

+3

@@ -233,6 +233,9 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):

233 self,

234 input_ids: torch.Tensor,

235 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

236 + *,

237 + is_multimodal: Optional[torch.Tensor] = None,

238 + handle_oov_mm_token: bool = False,

239 ) -> torch.Tensor:

240 # We do not really use any input tokens and therefore no embeddings

241 # to be calculated. However, due to the mandatory token ids in

vllm/model_executor/models/transformers.py

+50 -14

@@ -52,8 +52,8 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder

52 from vllm.sequence import IntermediateTensors

53 from vllm.utils import is_list_of

54

55 -from .interfaces import (SupportsLoRA, SupportsMultiModal, SupportsPP,

56 - SupportsQuant)

55 +from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

56 + SupportsMultiModal, SupportsPP, SupportsQuant)

57 from .utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,

58 flatten_bn, make_empty_intermediate_tensors_factory,

59 maybe_prefix)

@@ -797,6 +797,9 @@ class TransformersForCausalLM(TransformersBase):

797 else:

798 self.lm_head = PPMissingLayer()

799

800 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

801 + return self.model.get_input_embeddings()(input_ids)

802 +

803 def compute_logits(

804 self,

805 hidden_states: torch.Tensor,

@@ -873,13 +876,19 @@ class TransformersForMultimodalLM(TransformersForCausalLM, SupportsMultiModal):

876 multimodal_embeds = self.get_multimodal_embeddings(**kwargs)

877 if multimodal_embeds is not None:

878 inputs_embeds = self.get_input_embeddings(

876 - input_ids, multimodal_embeds)

879 + input_ids,

880 + multimodal_embeds,

881 + is_multimodal=input_ids == self.config.image_token_id,

882 + )

883 input_ids = None

884

885 model_output = super().forward(input_ids, positions,

886 intermediate_tensors, inputs_embeds)

887 return model_output

888

889 + def get_language_model(self) -> torch.nn.Module:

890 + return self.model

891 +

892 def get_multimodal_embeddings(self, **kwargs):

893 pixel_values = kwargs.pop("pixel_values", None)

894 pixel_values = pixel_values if pixel_values is not None else kwargs.pop(

@@ -934,15 +943,42 @@ class TransformersForMultimodalLM(TransformersForCausalLM, SupportsMultiModal):

943 def get_input_embeddings(

944 self,

945 input_ids: torch.Tensor,

937 - multimodal_embeddings=None,

946 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

947 + *,

948 + is_multimodal: Optional[torch.Tensor] = None,

949 + handle_oov_mm_token: bool = False,

950 ) -> torch.Tensor:

939 - inputs_embeds = self.model.get_input_embeddings()(input_ids)

940 - if (multimodal_embeddings is not None

941 - and len(multimodal_embeddings) != 0):

942 - mask = (input_ids == self.config.image_token_id)

943 - mask = mask.unsqueeze(-1).expand_as(inputs_embeds)

944 - multimodal_embeddings = torch.cat(multimodal_embeddings)

945 -

946 - inputs_embeds = inputs_embeds.masked_scatter(

947 - mask, multimodal_embeddings)

948 - return inputs_embeds

951 + """

952 + Apply token embeddings to `input_ids`.

953 +

954 + If `multimodal_embeddings` is passed, scatter them into

955 + `input_ids` according to the mask `is_multimodal`.

956 +

957 + In case the multi-modal token IDs exceed the vocabulary size of

958 + the language model, you can set `handle_oov_mm_token=False`

959 + to avoid calling the language model's `get_input_embeddings` method

960 + on those tokens.

961 + """

962 + from .utils import _merge_multimodal_embeddings

963 +

964 + inputs_embeds = self._get_text_embeddings(

965 + input_ids,

966 + self.model.get_input_embeddings(),

967 + is_multimodal=is_multimodal,

968 + handle_oov_mm_token=handle_oov_mm_token,

969 + )

970 +

971 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:

972 + return inputs_embeds

973 +

974 + if is_multimodal is None:

975 + raise ValueError(

976 + "`get_input_embeddings` now requires `is_multimodal` arg, "

977 + "please update your model runner according to "

978 + "https://github.com/vllm-project/vllm/pull/16229.")

979 +

980 + return _merge_multimodal_embeddings(

981 + inputs_embeds=inputs_embeds,

982 + multimodal_embeddings=multimodal_embeddings,

983 + is_multimodal=is_multimodal,

984 + )

vllm/model_executor/models/ultravox.py

+20 -16

@@ -33,8 +33,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

33 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

34 SupportsMultiModal, SupportsPP)

35 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

36 - init_vllm_registered_model, maybe_prefix,

37 - merge_multimodal_embeddings)

36 + init_vllm_registered_model, maybe_prefix)

37

38 _AUDIO_PLACEHOLDER_OVERRIDE = "<|audio|>"

39 _MAX_ENCODER_BATCH_SIZE = 16

@@ -555,19 +554,21 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

554 self,

555 input_ids: torch.Tensor,

556 multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

557 + *,

558 + is_multimodal: Optional[torch.Tensor] = None,

559 + # Multi-modal token ID may exceed vocab size

560 + handle_oov_mm_token: bool = True,

561 ) -> torch.Tensor:

559 - # The audio token index is not included in the embedding table

560 - # We need to remove it before embedding lookup

561 - safe_input_ids = input_ids.clone()

562 - safe_input_ids[safe_input_ids == self.config.audio_token_index] = 0

563 - inputs_embeds = self.language_model.get_input_embeddings(

564 - safe_input_ids)

565 - if multimodal_embeddings is not None and len(

566 - multimodal_embeddings) > 0:

567 - inputs_embeds = merge_multimodal_embeddings(

568 - input_ids, inputs_embeds, multimodal_embeddings,

569 - self.config.audio_token_index)

570 - return inputs_embeds

562 + # This is to satisfy the type checker for each overload

563 + if multimodal_embeddings is None or is_multimodal is None:

564 + return super().get_input_embeddings(input_ids)

565 +

566 + return super().get_input_embeddings(

567 + input_ids,

568 + multimodal_embeddings=multimodal_embeddings,

569 + is_multimodal=is_multimodal,

570 + handle_oov_mm_token=handle_oov_mm_token,

571 + )

572

573 def forward(self,

574 input_ids: torch.Tensor,

@@ -601,8 +602,11 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

602 elif inputs_embeds is None:

603 multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

604

604 - inputs_embeds = self.get_input_embeddings(input_ids,

605 - multimodal_embeddings)

605 + inputs_embeds = self.get_input_embeddings(

606 + input_ids,

607 + multimodal_embeddings,

608 + is_multimodal=input_ids == self.config.audio_token_index,

609 + )

610 input_ids = None

611

612 language_model = self.language_model

vllm/model_executor/models/utils.py

+37 -57

@@ -4,7 +4,7 @@

4 import itertools

5 from collections.abc import Iterable, Mapping

6 from dataclasses import dataclass, field

7 -from typing import Any, Callable, Literal, Optional, Protocol, Union, overload

7 +from typing import Any, Literal, Optional, Protocol, Union, overload

8

9 import torch

10 import torch.nn as nn

@@ -391,8 +391,8 @@ def _embedding_count_expression(embeddings: NestedTensors) -> str:

391

392 def _merge_multimodal_embeddings(

393 inputs_embeds: torch.Tensor,

394 - is_multimodal: torch.Tensor,

394 multimodal_embeddings: NestedTensors,

395 + is_multimodal: torch.Tensor,

396 ) -> torch.Tensor:

397 """

398 Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the

@@ -402,61 +402,35 @@ def _merge_multimodal_embeddings(

402 Note:

403 This updates ``inputs_embeds`` in place.

404 """

405 - flattened = _flatten_embeddings(multimodal_embeddings)

405 + if len(multimodal_embeddings) == 0:

406 + return inputs_embeds

407 +

408 + mm_embeds_flat = _flatten_embeddings(multimodal_embeddings)

409 + input_dtype = inputs_embeds.dtype

410 +

411 try:

407 - # This is equivalent to: inputs_embeds[is_multimodal] = flattened.

412 + # For debugging

413 + # inputs_embeds[is_multimodal] = mm_embeds_flat.to(dtype=input_dtype)

414 +

415 + # NOTE: This can avoid D2H sync (#22105), but fails to

416 + # raise an error if is_multimodal.sum() < len(mm_embeds_flat)

417 inputs_embeds.masked_scatter_(is_multimodal.unsqueeze(-1),

409 - flattened.to(dtype=inputs_embeds.dtype))

418 + mm_embeds_flat.to(dtype=input_dtype))

419 except RuntimeError as e:

420 + num_actual_tokens = len(mm_embeds_flat)

421 num_expected_tokens = is_multimodal.sum().item()

412 - assert isinstance(num_expected_tokens, int)

422

414 - if flattened.shape[0] != num_expected_tokens:

423 + if num_actual_tokens != num_expected_tokens:

424 expr = _embedding_count_expression(multimodal_embeddings)

425 +

426 raise ValueError(

417 - f"Attempted to assign {expr} = {flattened.shape[0]} "

427 + f"Attempted to assign {expr} = {num_actual_tokens} "

428 f"multimodal tokens to {num_expected_tokens} placeholders"

429 ) from e

420 - else:

421 - raise ValueError("Error during masked scatter operation") from e

422 -

423 - return inputs_embeds

424 -

425 -

426 -def embed_multimodal(

427 - input_ids: torch.Tensor,

428 - multimodal_token_id: int,

429 - get_text_embeds: Callable[[torch.Tensor], torch.Tensor],

430 - multimodal_embeds: NestedTensors,

431 -) -> torch.Tensor:

432 - """

433 - Embed token IDs and multimodal inputs and combine their embeddings.

434 -

435 - ``multimodal_token_id`` is used to determine whether a token ID should

436 - be embedded using ``get_text_embeds`` or ``get_multimodal_embeds``.

430

438 - Compared to ``merge_multimodal_embeddings`, this avoids running

439 - ``get_text_embeds`` on ``input_ids[input_ids == multimodal_token_id]``

440 - which causes issues when the placeholder token ID exceeds the

441 - vocabulary size of the language model.

442 - """

443 - is_multimodal = input_ids == multimodal_token_id

444 - is_text = ~is_multimodal

445 -

446 - text_embeds = get_text_embeds(input_ids[is_text])

447 - merged_embeds = torch.empty(

448 - (input_ids.shape[0], text_embeds.shape[1]),

449 - dtype=text_embeds.dtype,

450 - device=text_embeds.device,

451 - )

431 + raise ValueError("Error during masked scatter operation") from e

432

453 - merged_embeds[is_text] = text_embeds

454 -

455 - return _merge_multimodal_embeddings(

456 - merged_embeds,

457 - is_multimodal,

458 - multimodal_embeds,

459 - )

433 + return inputs_embeds

434

435

436 def merge_multimodal_embeddings(

@@ -491,23 +465,29 @@ def merge_multimodal_embeddings(

465 This updates ``inputs_embeds`` in place.

466 """

467 if isinstance(placeholder_token_id, list):

494 - placeholder_token_id = torch.tensor(

495 - placeholder_token_id,

496 - pin_memory=is_pin_memory_available()).to(device=input_ids.device,

497 - non_blocking=True)

498 - return _merge_multimodal_embeddings(

499 - inputs_embeds,

500 - torch.isin(input_ids, placeholder_token_id),

501 - multimodal_embeddings,

502 - )

468 + is_multimodal = isin_list(input_ids, placeholder_token_id)

469 + else:

470 + is_multimodal = (input_ids == placeholder_token_id)

471

472 return _merge_multimodal_embeddings(

473 inputs_embeds,

506 - (input_ids == placeholder_token_id),

507 - multimodal_embeddings,

474 + multimodal_embeddings=multimodal_embeddings,

475 + is_multimodal=is_multimodal,

476 )

477

478

479 +def isin_list(

480 + elements: torch.Tensor,

481 + test_elements_list: list[int],

482 +) -> torch.Tensor:

483 + test_elements = torch.tensor(

484 + test_elements_list,

485 + pin_memory=is_pin_memory_available(),

486 + ).to(device=elements.device, non_blocking=True)

487 +

488 + return torch.isin(elements, test_elements)

489 +

490 +

491 class LayerFn(Protocol):

492

493 def __call__(self, prefix: str) -> torch.nn.Module:

vllm/model_executor/models/voxtral.py

+9 -20

@@ -45,10 +45,8 @@ from vllm.sequence import IntermediateTensors

45 from vllm.transformers_utils.tokenizer import (MistralTokenizer,

46 cached_tokenizer_from_config)

47

48 -from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

49 - SupportsMultiModal, SupportsTranscription)

50 -from .utils import (flatten_bn, init_vllm_registered_model, maybe_prefix,

51 - merge_multimodal_embeddings)

48 +from .interfaces import SupportsLoRA, SupportsMultiModal, SupportsTranscription

49 +from .utils import flatten_bn, init_vllm_registered_model, maybe_prefix

50

51 logger = init_logger(__name__)

52

@@ -376,9 +374,14 @@ class VoxtralForConditionalGeneration(nn.Module, SupportsMultiModal,

374 # NOTE: In v1, inputs_embeds is always generated at model runner, this

375 # condition is for v0 compatibility.

376 elif inputs_embeds is None:

377 + audio_encoder = self.tokenizer.instruct.audio_encoder

378 + audio_tok_id = audio_encoder.audio_token

379 audio_embeddings = self.get_multimodal_embeddings(**kwargs)

380 - inputs_embeds = self.get_input_embeddings(input_ids,

381 - audio_embeddings)

380 + inputs_embeds = self.get_input_embeddings(

381 + input_ids,

382 + audio_embeddings,

383 + is_multimodal=input_ids == audio_tok_id,

384 + )

385 input_ids = None

386

387 hidden_states = self.language_model.model(input_ids,

@@ -421,20 +424,6 @@ class VoxtralForConditionalGeneration(nn.Module, SupportsMultiModal,

424

425 return audio_embeddings

426

424 - def get_input_embeddings(

425 - self,

426 - input_ids: torch.Tensor,

427 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

428 - ) -> torch.Tensor:

429 - audio_encoder = self.tokenizer.instruct.audio_encoder

430 - audio_tok_id = audio_encoder.audio_token

431 -

432 - inputs_embeds = self.language_model.get_input_embeddings(input_ids)

433 - if multimodal_embeddings is not None:

434 - inputs_embeds = merge_multimodal_embeddings(

435 - input_ids, inputs_embeds, multimodal_embeddings, audio_tok_id)

436 - return inputs_embeds

437 -

427 def _parse_and_validate_audio_arrays(

428 self, **kwargs: object) -> Union[list[torch.Tensor], None]:

429 audio_arrays = kwargs.pop("audio_arrays", None)

vllm/model_executor/models/whisper.py

+5 -5

@@ -579,10 +579,7 @@ class WhisperDecoder(nn.Module):

579 hidden_states = self.layer_norm(hidden_states)

580 return hidden_states

581

582 - def get_input_embeddings(

583 - self,

584 - input_ids: torch.Tensor,

585 - ) -> torch.Tensor:

582 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

583 return self.embed_tokens(input_ids)

584

585

@@ -916,7 +913,10 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,

913 def get_input_embeddings(

914 self,

915 input_ids: torch.Tensor,

919 - multimodal_embeddings: Optional[NestedTensors] = None,

916 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

917 + *,

918 + is_multimodal: Optional[torch.Tensor] = None,

919 + handle_oov_mm_token: bool = False,

920 ) -> torch.Tensor:

921 # This method just returns the decoder sequence embeddings since

922 # Whisper does not have encoder text tokens.

vllm/v1/spec_decode/eagle.py

+27 -19

@@ -18,6 +18,7 @@ from vllm.logger import init_logger

18 from vllm.model_executor.model_loader import get_model

19 from vllm.model_executor.models import supports_multimodal

20 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

21 +from vllm.multimodal import MULTIMODAL_REGISTRY

22 from vllm.platforms import current_platform

23 from vllm.utils import is_pin_memory_available

24 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata

@@ -64,8 +65,10 @@ class EagleProposer:

65 # hidden size (e.g., Llama 3.3 70B).

66 self.hidden_size = self.draft_model_config.get_hidden_size()

67

67 - self.is_multimodal_model = vllm_config.model_config \

68 - .is_multimodal_model

68 + # Multi-modal data support

69 + self.mm_registry = MULTIMODAL_REGISTRY

70 + self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(

71 + vllm_config.model_config)

72

73 self.attn_metadata_builder: Optional[AttentionMetadataBuilder] = None

74

@@ -175,7 +178,8 @@ class EagleProposer:

178 last_token_indices: Optional[torch.Tensor],

179 common_attn_metadata: CommonAttentionMetadata,

180 sampling_metadata: SamplingMetadata,

178 - mm_embeds: Optional[list[torch.Tensor]] = None,

181 + mm_embed_inputs: Optional[tuple[list[torch.Tensor],

182 + torch.Tensor]] = None,

183 ) -> torch.Tensor:

184 num_tokens = target_token_ids.shape[0]

185 batch_size = next_token_ids.shape[0]

@@ -219,18 +223,21 @@ class EagleProposer:

223 # copy inputs to buffer for cudagraph

224 self._set_positions(num_tokens, target_positions)

225 self.hidden_states[:num_tokens] = target_hidden_states

222 - if self.is_multimodal_model:

223 - input_ids = self.input_ids[:num_tokens]

224 - inputs_embeds = self.model.get_input_embeddings(

225 - input_ids,

226 - multimodal_embeddings=mm_embeds or None,

226 +

227 + if self.supports_mm_inputs:

228 + mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

229 +

230 + self.inputs_embeds[:num_tokens] = self.model.get_input_embeddings(

231 + self.input_ids[:num_tokens],

232 + multimodal_embeddings=mm_embeds,

233 + is_multimodal=is_mm_embed,

234 )

228 - self.inputs_embeds[:num_tokens] = inputs_embeds

229 - inputs_embeds = self.inputs_embeds[:num_input_tokens]

235 +

236 input_ids = None

237 + inputs_embeds = self.inputs_embeds[:num_input_tokens]

238 else:

232 - inputs_embeds = None

239 input_ids = self.input_ids[:num_input_tokens]

240 + inputs_embeds = None

241

242 with set_forward_context(per_layer_attn_metadata,

243 self.vllm_config,

@@ -372,14 +379,15 @@ class EagleProposer:

379 self.input_ids[:batch_size] = input_ids

380 self._set_positions(batch_size, clamped_positions)

381 self.hidden_states[:batch_size] = hidden_states

375 - if self.is_multimodal_model:

376 - inputs_embeds = self.model.get_input_embeddings(input_ids)

377 - self.inputs_embeds[:batch_size] = inputs_embeds

378 - inputs_embeds = self.inputs_embeds[:input_batch_size]

382 + if self.supports_mm_inputs:

383 + self.inputs_embeds[:batch_size] = \

384 + self.model.get_input_embeddings(input_ids)

385 +

386 input_ids = None

387 + inputs_embeds = self.inputs_embeds[:input_batch_size]

388 else:

381 - inputs_embeds = None

389 input_ids = self.input_ids[:input_batch_size]

390 + inputs_embeds = None

391

392 # Run the model.

393 with set_forward_context(per_layer_attn_metadata,

@@ -849,7 +857,7 @@ class EagleProposer:

857

858 self.attn_layer_names = list(draft_attn_layer_names)

859

852 - if self.is_multimodal_model:

860 + if self.supports_mm_inputs:

861 # Even if the target model is multimodal, we can also use

862 # text-only draft models

863 try:

@@ -861,7 +869,7 @@ class EagleProposer:

869 logger.warning(

870 "Draft model does not support multimodal inputs, "

871 "falling back to text-only mode")

864 - self.is_multimodal_model = False

872 + self.supports_mm_inputs = False

873

874 if supports_multimodal(target_model):

875 # handle multimodality

@@ -933,7 +941,7 @@ class EagleProposer:

941 ) -> None:

942 with set_forward_context(None, self.vllm_config,

943 num_tokens=num_tokens):

936 - if self.is_multimodal_model:

944 + if self.supports_mm_inputs:

945 input_ids = None

946 inputs_embeds = self.inputs_embeds[:num_tokens]

947 else:

vllm/v1/worker/gpu_model_runner.py

+39 -13

@@ -368,6 +368,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

368 self.num_accepted_tokens = self._make_buffer(self.max_num_reqs,

369 dtype=torch.int64)

370

371 + # Only relevant for multimodal models

372 + if self.supports_mm_inputs:

373 + self.is_mm_embed = self._make_buffer(self.max_num_tokens,

374 + dtype=torch.bool)

375 +

376 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)

377 if self.uses_mrope:

378 # NOTE: `mrope_positions` is implemented with one additional dummy

@@ -1627,9 +1632,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1632 self,

1633 scheduler_output: "SchedulerOutput",

1634 shift_computed_tokens: int = 0,

1630 - ) -> list[torch.Tensor]:

1635 + ) -> tuple[list[torch.Tensor], torch.Tensor]:

1636 + total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

1637 +

1638 + mm_embeds = list[torch.Tensor]()

1639 + is_mm_embed = self.is_mm_embed.cpu

1640 + is_mm_embed[:total_num_scheduled_tokens] = False

1641 +

1642 + req_start_idx = 0

1643 should_sync_mrope_positions = False

1632 - mm_embeds: list[torch.Tensor] = []

1644 +

1645 for req_id in self.input_batch.req_ids:

1646 mm_embeds_req: list[torch.Tensor] = []

1647

@@ -1638,6 +1650,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1650 req_state = self.requests[req_id]

1651 num_computed_tokens = \

1652 req_state.num_computed_tokens + shift_computed_tokens

1653 +

1654 for mm_feature in req_state.mm_features:

1655 pos_info = mm_feature.mm_position

1656 start_pos = pos_info.offset

@@ -1670,6 +1683,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1683 if (is_embed := pos_info.is_embed) is not None:

1684 is_embed = is_embed[start_idx:end_idx]

1685

1686 + req_start_pos = req_start_idx + start_pos - num_computed_tokens

1687 + is_mm_embed[req_start_pos+start_idx:req_start_pos + end_idx] \

1688 + = True if is_embed is None else is_embed

1689 +

1690 mm_embeds_item = gather_mm_placeholders(

1691 encoder_output[start_idx:end_idx],

1692 is_embed=is_embed,

@@ -1677,6 +1694,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1694 mm_embeds_req.append(mm_embeds_item)

1695

1696 if self.is_multimodal_pruning_enabled and self.uses_mrope:

1697 + assert req_state.mrope_positions is not None

1698 should_sync_mrope_positions = True

1699 mm_embeds_req, new_mrope_positions, new_delta = (

1700 self.model.recompute_mrope_positions(

@@ -1685,18 +1703,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1703 mrope_positions=req_state.mrope_positions,

1704 num_computed_tokens=req_state.num_computed_tokens,

1705 ))

1688 - assert req_state.mrope_positions is not None

1706 req_state.mrope_positions.copy_(new_mrope_positions)

1707 req_state.mrope_position_delta = new_delta

1708

1709 mm_embeds.extend(mm_embeds_req)

1710 + req_start_idx += num_scheduled_tokens

1711 +

1712 + is_mm_embed = self.is_mm_embed.copy_to_gpu(total_num_scheduled_tokens)

1713

1714 if should_sync_mrope_positions:

1715 self._calc_mrope_positions(scheduler_output)

1696 - self.mrope_positions.copy_to_gpu(

1697 - scheduler_output.total_num_scheduled_tokens)

1716 + self.mrope_positions.copy_to_gpu(total_num_scheduled_tokens)

1717

1699 - return mm_embeds

1718 + return mm_embeds, is_mm_embed

1719

1720 def _extract_encoder_inputs(

1721 self,

@@ -1990,14 +2009,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2009 and not self.model_config.is_encoder_decoder):

2010 # Run the multimodal encoder if any.

2011 self._execute_mm_encoder(scheduler_output)

1993 - mm_embeds = self._gather_mm_embeddings(scheduler_output)

2012 + mm_embeds, is_mm_embed = self._gather_mm_embeddings(

2013 + scheduler_output)

2014

2015 # NOTE(woosuk): To unify token ids and soft tokens (vision

2016 # embeddings), we always use embeddings (rather than token ids)

2017 # as input to the multimodal model, even when the input is text.

2018 inputs_embeds_scheduled = self.model.get_input_embeddings(

1999 - input_ids=self.input_ids.gpu[:num_scheduled_tokens],

2000 - multimodal_embeddings=mm_embeds or None,

2019 + self.input_ids.gpu[:num_scheduled_tokens],

2020 + multimodal_embeddings=mm_embeds,

2021 + is_multimodal=is_mm_embed,

2022 )

2023

2024 # TODO(woosuk): Avoid the copy. Optimize.

@@ -2586,10 +2607,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2607 [h[token_indices] for h in aux_hidden_states], dim=-1)

2608 else:

2609 target_hidden_states = hidden_states[token_indices]

2589 - mm_embeds = None

2610 +

2611 if self.supports_mm_inputs:

2591 - mm_embeds = self._gather_mm_embeddings(scheduler_output,

2592 - shift_computed_tokens=1)

2612 + mm_embed_inputs = self._gather_mm_embeddings(

2613 + scheduler_output,

2614 + shift_computed_tokens=1,

2615 + )

2616 + else:

2617 + mm_embed_inputs = None

2618

2619 draft_token_ids = self.drafter.propose(

2620 target_token_ids=target_token_ids,

@@ -2599,8 +2624,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2624 last_token_indices=token_indices_to_sample,

2625 sampling_metadata=sampling_metadata,

2626 common_attn_metadata=common_attn_metadata,

2602 - mm_embeds=mm_embeds,

2627 + mm_embed_inputs=mm_embed_inputs,

2628 )

2629 +

2630 return draft_token_ids

2631

2632 def update_config(self, overrides: dict[str, Any]) -> None:

vllm/v1/worker/tpu_model_runner.py

+66 -13

@@ -263,6 +263,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

263 pin_memory=self.pin_memory)

264 self.seq_lens_np = self.seq_lens_cpu.numpy()

265

266 + # Only relevant for multimodal models

267 + if self.supports_mm_inputs:

268 + self.is_mm_embed_cpu = torch.zeros(self.max_num_tokens,

269 + dtype=torch.bool,

270 + device="cpu",

271 + pin_memory=self.pin_memory)

272 +

273 # Range tensor with values [0 .. self.max_num_tokens - 1].

274 # Used to initialize positions / context_lens / seq_lens

275 # Keep in int64 to avoid overflow with long context

@@ -879,13 +886,22 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

886 def _gather_mm_embeddings(

887 self,

888 scheduler_output: "SchedulerOutput",

882 - ) -> list[torch.Tensor]:

883 - mm_embeds: list[torch.Tensor] = []

889 + ) -> tuple[list[torch.Tensor], torch.Tensor]:

890 + total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

891 + padded_total_num_scheduled_tokens = _get_padded_token_len(

892 + self.num_tokens_paddings, total_num_scheduled_tokens)

893 +

894 + is_mm_embed = self.is_mm_embed_cpu

895 + is_mm_embed[:padded_total_num_scheduled_tokens] = False

896 + mm_embeds = list[torch.Tensor]()

897 + req_start_idx = 0

898 +

899 for req_id in self.input_batch.req_ids:

900 num_scheduled_tokens = scheduler_output.num_scheduled_tokens[

901 req_id]

902 req_state = self.requests[req_id]

903 num_computed_tokens = req_state.num_computed_tokens

904 +

905 # TODO unroll loop and assume/enforce --disable_chunked_mm_input

906 # NOTE (NickLucche) here we diverge from logic in other runners, as

907 # we assume to only have whole mm items to process. Hence we avoid

@@ -906,26 +922,53 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

922 # The encoder output is already processed and stored

923 # in the decoder's KV cache.

924 continue

925 +

926 + start_idx = max(num_computed_tokens - start_pos, 0)

927 + end_idx = min(

928 + num_computed_tokens - start_pos + num_scheduled_tokens,

929 + num_encoder_tokens,

930 + )

931 + assert start_idx < end_idx

932 +

933 mm_hash = mm_feature.identifier

934 encoder_output = self.encoder_cache.get(mm_hash, None)

935 assert encoder_output is not None,\

936 f"Encoder cache miss for {mm_hash}."

937 +

938 assert pos_info.is_embed is None, "Expected all positions to"\

939 " be contiguous and embeddings."

915 - encoder_output = self.encoder_cache[mm_hash]

940 +

941 + req_start_pos = req_start_idx + start_pos - num_computed_tokens

942 + is_mm_embed[req_start_pos+start_idx:req_start_pos + end_idx] \

943 + = True

944 +

945 + # Only whole mm items are processed

946 mm_embeds.append(encoder_output)

917 - return mm_embeds

947

919 - def _get_model_inputs(self, input_ids: torch.Tensor,

920 - mm_embeds: list[torch.Tensor]):

948 + req_start_idx += num_scheduled_tokens

949 +

950 + is_mm_embed = is_mm_embed[:padded_total_num_scheduled_tokens] \

951 + .to(self.device)

952 +

953 + return mm_embeds, is_mm_embed

954 +

955 + def _get_model_inputs(

956 + self,

957 + input_ids: torch.Tensor,

958 + mm_embed_inputs: Optional[tuple[list[torch.Tensor], torch.Tensor]],

959 + ):

960 if self.supports_mm_inputs:

961 + mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

962 +

963 # NOTE(woosuk): To unify token ids and soft tokens (vision

964 # embeddings), we always use embeddings (rather than token ids)

965 # as input to the multimodal model, even when the input is text.

966 inputs_embeds = self.model.get_input_embeddings(

926 - input_ids=input_ids,

967 + input_ids,

968 multimodal_embeddings=mm_embeds,

969 + is_multimodal=is_mm_embed,

970 )

971 +

972 return None, inputs_embeds

973 else:

974 # For text-only models, we use token ids as input.

@@ -953,9 +996,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

996 if self.supports_mm_inputs:

997 # Run the multimodal encoder if any.

998 self._execute_mm_encoder(scheduler_output)

956 - mm_embeds = self._gather_mm_embeddings(scheduler_output)

999 + mm_embed_inputs = self._gather_mm_embeddings(scheduler_output)

1000 else:

958 - mm_embeds = []

1001 + mm_embed_inputs = None

1002 +

1003 torch_xla.sync(wait=False)

1004 # Prepare inputs, the requests might be split into multiple

1005 # executions, combine the result of each execution.

@@ -972,7 +1016,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1016 attn_metadata, logits_indices, padded_num_reqs, num_reqs,\

1017 end_index = self._prepare_inputs(scheduler_output, start_index)

1018 input_ids, inputs_embeds = self._get_model_inputs(

975 - self.input_ids, mm_embeds)

1019 + self.input_ids, mm_embed_inputs)

1020 torch_xla.sync(wait=False)

1021 # Run the decoder

1022 with set_forward_context(

@@ -1325,9 +1369,15 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1369 hf_config.image_token_index

1370

1371 placeholders_ids = placeholders_ids.to(self.device)

1372 +

1373 + mm_mask = torch.tensor([False] * num_tokens)

1374 + mm_mask[:items_size] = True

1375 + mm_mask = mm_mask.to(self.device)

1376 # Assign outputs or the graph will be cut short.

1329 - a, b = self._get_model_inputs(placeholders_ids,

1330 - [mm_embeds])

1377 + a, b = self._get_model_inputs(

1378 + placeholders_ids,

1379 + mm_embed_inputs=([mm_embeds], mm_mask),

1380 + )

1381 assert a is None

1382 torch_xla.sync(wait=False)

1383

@@ -1338,7 +1388,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1388 dtype=torch.int32,

1389 device="cpu")

1390 placeholders_ids = placeholders_ids.to(self.device)

1341 - a, b = self._get_model_inputs(placeholders_ids, [])

1391 + a, b = self._get_model_inputs(

1392 + placeholders_ids,

1393 + mm_embed_inputs=None,

1394 + )

1395 assert a is None

1396 torch_xla.sync(wait=False)

1397

[Bugfix] Merge MM embeddings by index instead of token IDs (#16229)