@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 39b643dc1a

Fork 0

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 (G) (#26117)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Oct 3, 2025 at 13:38 UTC 39b643dc1ac68bb595a9e0a3ea0834e9adc00e92

5 files changed +56 -108

vllm/model_executor/models/gemma3_mm.py

+14 -21

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

36 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

37 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)

42 logger = init_logger(__name__)

@@ -289,7 +289,7 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):

289 processor=hf_processor)

290 for size in image_sizes

291 ]

292 - processed_outputs["num_crops"] = torch.tensor(num_crops)

292 + processed_outputs["num_patches"] = torch.tensor(num_crops) + 1

293

294 return processed_outputs

295

@@ -298,12 +298,12 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):

298 hf_inputs: BatchFeature,

299 hf_processor_mm_kwargs: Mapping[str, object],

300 ) -> Mapping[str, MultiModalFieldConfig]:

301 - num_crops = hf_inputs.get("num_crops", torch.empty(0))

301 + num_patches = hf_inputs.get("num_patches", torch.empty(0))

302

303 return dict(

304 pixel_values=MultiModalFieldConfig.flat_from_sizes(

305 - "image", num_crops + 1),

306 - num_crops=MultiModalFieldConfig.batched("image"),

305 + "image", num_patches),

306 + num_patches=MultiModalFieldConfig.batched("image"),

307 )

308

309 def _get_prompt_updates(

@@ -460,6 +460,8 @@ class Gemma3MultiModalProjector(nn.Module):

460 dummy_inputs=Gemma3DummyInputsBuilder)

461 class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

462 SupportsLoRA):

463 + merge_by_field_config = True

464 +

465 packed_modules_mapping = {

466 "qkv_proj": [

467 "q_proj",

@@ -526,29 +528,20 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

528 def _parse_and_validate_image_input(

529 self, **kwargs: object) -> Optional[Gemma3ImageInputs]:

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

529 - num_crops = kwargs.pop("num_crops", None)

531 + num_patches = kwargs.pop("num_patches", None)

532 image_embeds = kwargs.pop("image_embeds", None)

533 assert image_embeds is None, "Gemma3 does not support image_embeds."

534 if pixel_values is None:

535 return None

536

535 - if not isinstance(pixel_values, (torch.Tensor, list)):

536 - raise ValueError("Incorrect type of pixel values. "

537 - f"Got type: {type(pixel_values)}")

538 -

539 - if not isinstance(num_crops, (torch.Tensor, list)):

540 - raise ValueError("Incorrect type of num_crops. "

541 - f"Got type: {type(num_crops)}")

542 -

537 image_size = self.config.vision_config.image_size

538

545 - return Gemma3ImagePixelInputs(

546 - pixel_values=flatten_bn(pixel_values, concat=True),

547 - num_patches=flatten_bn(num_crops, concat=True) + 1,

548 - resolve_bindings={

549 - "h": image_size,

550 - "w": image_size

551 - })

539 + return Gemma3ImagePixelInputs(pixel_values=pixel_values,

540 + num_patches=num_patches,

541 + resolve_bindings={

542 + "h": image_size,

543 + "w": image_size

544 + })

545

546 def _image_pixels_to_features(

547 self,

vllm/model_executor/models/gemma3n_mm.py

+31 -38

@@ -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 Any, Literal, Optional, TypedDict, Union, cast

4 +from typing import Annotated, Any, Literal, Optional, Union, cast

6 import numpy as np

7 import torch

@@ -41,6 +41,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,

41 # yapf: enable

42 from vllm.multimodal.profiling import BaseDummyInputsBuilder

43 from vllm.sequence import IntermediateTensors

44 +from vllm.utils.tensor_schema import TensorSchema, TensorShape

46 from .interfaces import (MultiModalEmbeddings, SupportsMultiModal,

47 SupportsTranscription)

@@ -54,17 +55,28 @@ TOKENS_PER_IMAGE = 256

55 TOKENS_PER_AUDIO = 188

57 -class Gemma3nImagePixelInputs(TypedDict):

58 - pixel_values: torch.Tensor

59 - """Shape: `(batch_size * num_images, num_channels, height, width)`"""

58 +class Gemma3nImagePixelInputs(TensorSchema):

59 + """

60 + Dimensions:

61 + - bn: Batch size * number of images

62 + - c: Number of channels (3)

63 + - h: Height of each patch

64 + - w: Width of each patch

65 + """

66 + type: Literal["pixel_values"] = "pixel_values"

67 + pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]

62 -class Gemma3nAudioInputs(TypedDict):

63 - input_features: Union[torch.Tensor, list[torch.Tensor]]

64 - input_features_padded: torch.Tensor

65 - """Shape: `(batch_size * num_audio, seq_length, num_features)`"""

66 - input_features_mask: torch.Tensor

67 - """Shape: `(batch_size * num_audio, seq_length)`"""

70 +class Gemma3nAudioInputs(TensorSchema):

71 + """

72 + Dimensions:

73 + - bn: Batch size * number of audios

74 + - s: seq_length

75 + - f: num_features

76 + """

77 + type: Literal["audio"] = "audio"

78 + input_features_padded: Annotated[torch.Tensor, TensorShape("bn", "s", "f")]

79 + input_features_mask: Annotated[torch.Tensor, TensorShape("bn", "s")]

82 Gemma3nImageInputs = Gemma3nImagePixelInputs

@@ -212,9 +224,9 @@ class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]

224

225 return dict(

226 pixel_values=MultiModalFieldConfig.batched("image"),

215 - input_features=MultiModalFieldConfig.batched("audio"),

227 input_features_padded=MultiModalFieldConfig.batched("audio"),

217 - input_features_mask=MultiModalFieldConfig.batched("audio"))

228 + input_features_mask=MultiModalFieldConfig.batched("audio"),

229 + )

230

231 def _get_prompt_updates(

232 self,

@@ -422,6 +434,7 @@ class Gemma3nMultimodalEmbedder(nn.Module):

434 dummy_inputs=Gemma3nDummyInputsBuilder)

435 class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal,

436 SupportsTranscription):

437 + merge_by_field_config = True

438 supported_languages = ISO639_1_SUPPORTED_LANGS

439

440 packed_modules_mapping = {

@@ -482,14 +495,6 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal,

495 device=self.language_model.model.embed_tokens.weight.device,

496 dtype=self.language_model.model.embed_tokens.weight.dtype)

497

485 - @property

486 - def dtype(self):

487 - return next(self.parameters()).dtype

488 -

489 - def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:

490 - # TODO check if there are any

491 - return data

492 -

498 def _parse_and_validate_image_input(

499 self, **kwargs: object) -> Optional[Gemma3nImageInputs]:

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

@@ -499,34 +504,22 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal,

504 if pixel_values is None:

505 return None

506

502 - if not isinstance(pixel_values, (torch.Tensor, list)):

503 - raise ValueError("Incorrect type of pixel values. "

504 - f"Got type: {type(pixel_values)}")

505 -

506 - pixel_values = flatten_bn(pixel_values, concat=True)

507 - pixel_values = pixel_values.contiguous()

508 -

509 - return Gemma3nImagePixelInputs(

510 - pixel_values=self._validate_pixel_values(pixel_values), )

507 + return Gemma3nImagePixelInputs(pixel_values=pixel_values)

508

509 def _parse_and_validate_audio_input(

510 self, **kwargs: object) -> Optional[Gemma3nAudioInputs]:

514 - input_features = kwargs.pop("input_features", None)

515 - if input_features is None:

511 +

512 + input_features_padded = kwargs.pop("input_features_padded", None)

513 + if input_features_padded is None:

514 return None

515

516 input_features_mask = kwargs.pop("input_features_mask", None)

517 if input_features_mask is None:

518 return None

519

522 - input_features_padded = kwargs.pop("input_features_padded", None)

523 - if input_features_padded is None:

524 - return None

525 -

520 return Gemma3nAudioInputs(

527 - input_features=input_features,

528 - input_features_mask=input_features_mask,

521 input_features_padded=input_features_padded,

522 + input_features_mask=input_features_mask,

523 )

524

525 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

@@ -539,7 +532,7 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal,

532 ) and "image" not in mm_input_by_modality:

533 mm_input_by_modality[

534 "image"] = self._parse_and_validate_image_input(**kwargs)

542 - if input_key == "input_features" \

535 + if input_key == "input_features_padded" \

536 and "audio" not in mm_input_by_modality:

537 mm_input_by_modality[

538 "audio"] = self._parse_and_validate_audio_input(**kwargs)

vllm/model_executor/models/glm4_1v.py

+2 -36

@@ -1319,6 +1319,8 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]):

1319 )

1320 class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1321 SupportsLoRA, SupportsPP):

1322 + merge_by_field_config = True

1323 +

1324 packed_modules_mapping = {

1325 "qkv_proj": [

1326 "q_proj",

@@ -1381,22 +1383,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1383 self.make_empty_intermediate_tensors = (

1384 self.language_model.make_empty_intermediate_tensors)

1385

1384 - def _validate_and_reshape_mm_tensor(self, mm_input: object,

1385 - name: str) -> torch.Tensor:

1386 - if not isinstance(mm_input, (torch.Tensor, list)):

1387 - raise ValueError(

1388 - f"Incorrect type of {name}. Got type: {type(mm_input)}")

1389 - if isinstance(mm_input, torch.Tensor):

1390 - if mm_input.ndim == 2:

1391 - return mm_input

1392 - if mm_input.ndim != 3:

1393 - raise ValueError(f"{name} should be 2D or batched 3D tensor. "

1394 - f"Got ndim: {mm_input.ndim} "

1395 - f"(shape={mm_input.shape})")

1396 - return mm_input.reshape(-1, mm_input.shape[-1])

1397 - else:

1398 - return torch.concat(mm_input)

1399 -

1386 def _parse_and_validate_image_input(

1387 self, **kwargs: object) -> Optional[Glm4vImageInputs]:

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

@@ -1407,11 +1393,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1393 return None

1394

1395 if pixel_values is not None:

1410 - pixel_values = self._validate_and_reshape_mm_tensor(

1411 - pixel_values, "image pixel values")

1412 - image_grid_thw = self._validate_and_reshape_mm_tensor(

1413 - image_grid_thw, "image grid_thw")

1414 -

1396 return Glm4vImagePixelInputs(

1397 type="pixel_values",

1398 pixel_values=pixel_values,

@@ -1419,11 +1400,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1400 )

1401

1402 if image_embeds is not None:

1422 - image_embeds = self._validate_and_reshape_mm_tensor(

1423 - image_embeds, "image embeds")

1424 - image_grid_thw = self._validate_and_reshape_mm_tensor(

1425 - image_grid_thw, "image grid_thw")

1426 -

1403 return Glm4vImageEmbeddingInputs(

1404 type="image_embeds",

1405 image_embeds=image_embeds,

@@ -1440,11 +1416,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1416 return None

1417

1418 if pixel_values_videos is not None:

1443 - pixel_values_videos = self._validate_and_reshape_mm_tensor(

1444 - pixel_values_videos, "video pixel values")

1445 - video_grid_thw = self._validate_and_reshape_mm_tensor(

1446 - video_grid_thw, "video grid_thw")

1447 -

1419 return Glm4vVideoPixelInputs(

1420 type="pixel_values_videos",

1421 pixel_values_videos=pixel_values_videos,

@@ -1452,11 +1423,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1423 )

1424

1425 if video_embeds is not None:

1455 - video_embeds = self._validate_and_reshape_mm_tensor(

1456 - video_embeds, "video embeds")

1457 - video_grid_thw = self._validate_and_reshape_mm_tensor(

1458 - video_grid_thw, "video grid_thw")

1459 -

1426 return Glm4vVideoEmbeddingInputs(

1427 type="video_embeds",

1428 video_embeds=video_embeds,

vllm/model_executor/models/glm4v.py

+6 -9

@@ -43,7 +43,6 @@ 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

48 class GLMVImagePixelInputs(TensorSchema):

@@ -529,8 +528,9 @@ class GLM4VMultiModalProcessor(BaseMultiModalProcessor[GLM4VProcessingInfo]):

528 @MULTIMODAL_REGISTRY.register_processor(GLM4VMultiModalProcessor,

529 info=GLM4VProcessingInfo,

530 dummy_inputs=GLM4VDummyInputsBuilder)

532 -class GLM4VForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP,

533 - SupportsMultiModal):

531 +class GLM4VForCausalLM(ChatGLMBaseModel, SupportsMultiModal, SupportsLoRA,

532 + SupportsPP):

533 + merge_by_field_config = True

534

535 packed_modules_mapping = {

536 "query_key_value": ["query_key_value"],

@@ -574,14 +574,9 @@ class GLM4VForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP,

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

575

576 if pixel_values is not None:

577 - if not isinstance(pixel_values, (torch.Tensor, list)):

578 - raise ValueError("Incorrect type of pixel values. "

579 - f"Got type: {type(pixel_values)}")

580 -

577 expected_h = expected_w = self.config.vision_config["image_size"]

578 return GLMVImagePixelInputs(type="pixel_values",

583 - data=flatten_bn(pixel_values,

584 - concat=True),

579 + data=pixel_values,

580 resolve_bindings={

581 "h": expected_h,

582 "w": expected_w

@@ -598,6 +593,8 @@ class GLM4VForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP,

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

594 return self.transformer

595

596 + get_input_embeddings = SupportsMultiModal.get_input_embeddings

597 +

598 def get_multimodal_embeddings(self,

599 **kwargs: object) -> MultiModalEmbeddings:

600 image_input = self._parse_and_validate_image_input(**kwargs)

vllm/model_executor/models/granite_speech.py

+3 -4

@@ -168,10 +168,8 @@ class GraniteSpeechMultiModalProcessor(

168 # Calculate the number of audio tokens per entry in the batch;

169 # This is used to split the batch back out after padding.

170 audio_token_index = self.info.get_hf_config().audio_token_index

171 - processed_outputs["audio_embed_sizes"] = [

172 - torch.sum(indices == audio_token_index).item()

173 - for indices in processed_outputs["input_ids"]

174 - ]

171 + processed_outputs["audio_embed_sizes"] = (

172 + processed_outputs["input_ids"] == audio_token_index).sum(-1)

173

174 return processed_outputs

175

@@ -527,6 +525,7 @@ class GraniteSpeechForConditionalGeneration(

525 SupportsPP,

526 SupportsLoRA,

527 ):

528 + merge_by_field_config = True

529

530 packed_modules_mapping = {

531 "qkv_proj": [