A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Model] Use `merge_by_field_config` for MM models (A-C) (#26073)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Oct 2, 2025 at 23:17 UTC
7d6fb905d984e2ff1d4c7d68543e4c61e15dba98
5 files changed
+29
-24
vllm/model_executor/models/aria.py
+13
-6
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Mapping, Sequence
4
-from typing import Annotated, Optional, Union
4
+from typing import Annotated, Literal, Optional, Union
5
6
import torch
7
import torch.nn as nn
@@ -38,8 +38,8 @@ from .idefics2_vision_model import (
38
# yapf: enable
39
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant
40
from .llama import LlamaDecoderLayer, LlamaMLP, LlamaModel
41
-from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
42
- is_pp_missing_parameter, maybe_prefix)
41
+from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
42
+ maybe_prefix)
43
44
45
class AriaImagePixelInputs(TensorSchema):
@@ -52,6 +52,8 @@ class AriaImagePixelInputs(TensorSchema):
52
- w: Width of each image
53
"""
54
55
+ type: Literal["pixel_values"]
56
+
57
pixel_values: Annotated[
58
torch.Tensor,
59
TensorShape("bn", 3, "h", "w"),
@@ -485,6 +487,8 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
487
This model combines a vision tower, a multi-modal projector, and a language
488
model to perform tasks that involve both image and text inputs.
489
"""
490
+ merge_by_field_config = True
491
+
492
hf_to_vllm_mapper = WeightsMapper(
493
orig_to_new_prefix={
494
# mapping for new names in checkpoint saved after transformers v4.52
@@ -551,12 +555,15 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
555
return None
556
557
return AriaImagePixelInputs(
554
- pixel_values=flatten_bn(pixel_values, concat=True),
555
- pixel_mask=flatten_bn(pixel_mask, concat=True),
558
+ type="pixel_values",
559
+ pixel_values=pixel_values,
560
+ pixel_mask=pixel_mask,
561
)
562
563
def _create_patch_attention_mask(
559
- self, pixel_mask: Optional[torch.Tensor]) -> torch.Tensor:
564
+ self,
565
+ pixel_mask: Optional[torch.Tensor],
566
+ ) -> Optional[torch.Tensor]:
567
if pixel_mask is None:
568
return None
569
vllm/model_executor/models/aya_vision.py
+4
-3
@@ -31,7 +31,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
31
32
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
33
from .siglip import SiglipVisionModel
34
-from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
34
+from .utils import (AutoWeightsLoader, WeightsMapper,
35
init_vllm_registered_model, maybe_prefix)
36
37
@@ -295,6 +295,7 @@ def _get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:
295
dummy_inputs=AyaVisionDummyInputsBuilder)
296
class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,
297
SupportsPP):
298
+ merge_by_field_config = True
299
300
hf_to_vllm_mapper = WeightsMapper(
301
orig_to_new_prefix={
@@ -379,8 +380,8 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,
380
381
return AyaVisionImagePixelInputs(
382
type="pixel_values",
382
- pixel_values=flatten_bn(pixel_values, concat=True),
383
- num_patches=flatten_bn(num_patches, concat=True),
383
+ pixel_values=pixel_values,
384
+ num_patches=num_patches,
385
resolve_bindings={
386
"h": self.config.vision_config.image_size,
387
"w": self.config.vision_config.image_size,
vllm/model_executor/models/blip2.py
+4
-9
@@ -26,12 +26,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
26
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)
31
-
32
-# We use this internally as placeholders since there is no image token
33
-# defined on the HuggingFace repo
34
-_IMAGE_TOKEN_ID = 50265
29
+from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
30
31
32
class Blip2ImagePixelInputs(TensorSchema):
@@ -514,6 +509,7 @@ class Blip2MultiModalProcessor(BaseMultiModalProcessor[Blip2ProcessingInfo]):
509
dummy_inputs=Blip2DummyInputsBuilder)
510
class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
511
SupportsQuant):
512
+ merge_by_field_config = True
513
514
@classmethod
515
def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
@@ -570,8 +566,7 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
566
if pixel_values is not None:
567
expected_h = expected_w = self.config.vision_config.image_size
568
return Blip2ImagePixelInputs(type="pixel_values",
573
- data=flatten_bn(pixel_values,
574
- concat=True),
569
+ data=pixel_values,
570
resolve_bindings={
571
"h": expected_h,
572
"w": expected_w
@@ -580,7 +575,7 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
575
if image_embeds is not None:
576
return Blip2ImageEmbeddingInputs(
577
type="image_embeds",
583
- data=flatten_bn(image_embeds, concat=True),
578
+ data=image_embeds,
579
)
580
581
raise AssertionError("This line should be unreachable.")
vllm/model_executor/models/chameleon.py
+4
-3
@@ -42,7 +42,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
42
43
from .interfaces import (MultiModalEmbeddings, SupportsMultiModal, SupportsPP,
44
SupportsQuant)
45
-from .utils import (flatten_bn, is_pp_missing_parameter,
45
+from .utils import (is_pp_missing_parameter,
46
make_empty_intermediate_tensors_factory, make_layers,
47
maybe_prefix)
48
@@ -935,6 +935,8 @@ class ChameleonModel(nn.Module):
935
dummy_inputs=ChameleonDummyInputsBuilder)
936
class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,
937
SupportsPP, SupportsQuant):
938
+ merge_by_field_config = True
939
+
940
packed_modules_mapping = {
941
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
942
"gate_up_proj": ["gate_proj", "up_proj"]
@@ -981,8 +983,7 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,
983
expected_h = expected_w = vq_config.resolution
984
985
return ChameleonImagePixelInputs(type="pixel_values",
984
- data=flatten_bn(pixel_values,
985
- concat=True),
986
+ data=pixel_values,
987
resolve_bindings={
988
"h": expected_h,
989
"w": expected_w
vllm/model_executor/models/cohere2_vision.py
+4
-3
@@ -36,7 +36,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
36
37
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
38
from .siglip import SiglipVisionModel
39
-from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
39
+from .utils import (AutoWeightsLoader, WeightsMapper,
40
init_vllm_registered_model, maybe_prefix)
41
42
@@ -317,6 +317,7 @@ class Cohere2VisionMultiModalProcessor(
317
dummy_inputs=Cohere2VisionDummyInputsBuilder)
318
class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal,
319
SupportsPP):
320
+ merge_by_field_config = True
321
322
hf_to_vllm_mapper = WeightsMapper(
323
orig_to_new_prefix={
@@ -399,8 +400,8 @@ class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal,
400
401
return Cohere2VisionImagePixelInputs(
402
type="pixel_values",
402
- pixel_values=flatten_bn(pixel_values, concat=True),
403
- num_patches=flatten_bn(num_patches, concat=True),
403
+ pixel_values=pixel_values,
404
+ num_patches=num_patches,
405
resolve_bindings={
406
"h": self.config.vision_config.image_size,
407
"w": self.config.vision_config.image_size,