@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 59a85c366e

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 (H-L) (#26230)

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

Cyrus Leung committed Oct 5, 2025 at 11:54 UTC 59a85c366ef3666d22b57f952979a3f74ee50f61

6 files changed +29 -161

examples/offline_inference/vision_language_multi_image.py

+1 -1

@@ -548,7 +548,7 @@ def load_keye_vl1_5(question: str, image_urls: list[str]) -> ModelRequestData:

548 engine_args = EngineArgs(

549 model=model_name,

550 trust_remote_code=True,

551 - max_model_len=8192,

551 + max_model_len=32768,

552 max_num_seqs=5,

553 limit_mm_per_prompt={"image": len(image_urls)},

554 )

vllm/model_executor/models/idefics3.py

+9 -23

@@ -53,7 +53,7 @@ from .idefics2_vision_model import (

53 # yapf: enable

54 from .interfaces import MultiModalEmbeddings, SupportsLoRA, SupportsMultiModal

55 from .llama import LlamaModel

56 -from .utils import AutoWeightsLoader, flatten_bn, maybe_prefix

56 +from .utils import AutoWeightsLoader, maybe_prefix

59 class Idefics3ImagePixelInputs(TensorSchema):

@@ -67,7 +67,7 @@ class Idefics3ImagePixelInputs(TensorSchema):

67 """

68 type: Literal["pixel_values"]

69 pixel_values: Annotated[torch.Tensor, TensorShape("bnp", 3, "h", "w")]

70 - pixel_attention_mask: torch.Tensor

70 + pixel_attention_mask: Annotated[torch.Tensor, TensorShape("bnp", "h", "w")]

71 num_patches: Annotated[torch.Tensor, TensorShape("bn")]

@@ -569,6 +569,8 @@ class Idefics3Model(nn.Module):

569 dummy_inputs=Idefics3DummyInputsBuilder)

570 class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,

571 SupportsLoRA):

572 + merge_by_field_config = True

573 +

574 packed_modules_mapping = {

575 "qkv_proj": [

576 "q_proj",

@@ -621,37 +623,21 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,

623 return None

624

625 if image_embeds is not None:

624 - if not isinstance(image_embeds, (torch.Tensor, list)):

625 - raise ValueError("Incorrect type of image embeddings. "

626 - f"Got type: {type(image_embeds)}")

627 -

626 return Idefics3ImageEmbeddingInputs(

627 type="image_embeds",

630 - data=flatten_bn(image_embeds, concat=True),

628 + data=image_embeds,

629 )

630

631 if pixel_values is not None:

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

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

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

637 -

632 pixel_attention_mask = kwargs.pop("pixel_attention_mask")

639 - if not isinstance(pixel_attention_mask, (torch.Tensor, list)):

640 - raise ValueError("Incorrect type of pixel_attention_mask. "

641 - f"Got type: {type(pixel_attention_mask)}")

642 -

633 num_patches = kwargs.pop("num_patches")

644 - if not isinstance(num_patches, (torch.Tensor, list)):

645 - raise ValueError("Incorrect type of num_patches. "

646 - f"Got type: {type(num_patches)}")

647 -

634 expected_h = expected_w = self.config.vision_config.image_size

635 +

636 return Idefics3ImagePixelInputs(

637 type="pixel_values",

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

652 - pixel_attention_mask=flatten_bn(pixel_attention_mask,

653 - concat=True),

654 - num_patches=flatten_bn(num_patches, concat=True),

638 + pixel_values=pixel_values,

639 + pixel_attention_mask=pixel_attention_mask,

640 + num_patches=num_patches,

641 resolve_bindings={

642 "h": expected_h,

643 "w": expected_w

vllm/model_executor/models/keye.py

+7 -52

@@ -30,7 +30,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

30 from vllm.model_executor.model_loader.weight_utils import (

31 default_weight_loader, maybe_remap_kv_scale_name)

32 from vllm.model_executor.models.module_mapping import MultiModelKeys

33 -from vllm.multimodal import MULTIMODAL_REGISTRY, NestedTensors

33 +from vllm.multimodal import MULTIMODAL_REGISTRY

34 from vllm.multimodal.inputs import (ImageItem, ModalityData,

35 MultiModalDataDict, MultiModalFieldConfig,

36 MultiModalKwargsItems, VideoItem)

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

42 PromptUpdate)

43 from vllm.multimodal.profiling import BaseDummyInputsBuilder

44 from vllm.sequence import IntermediateTensors

45 -from vllm.utils import is_list_of

45 from vllm.utils.tensor_schema import TensorSchema, TensorShape

47 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

@@ -100,8 +99,7 @@ def smart_resize(

99 class KeyeImagePixelInputs(TensorSchema):

100 """

101 Dimensions:

103 - - b: Batch size

104 - - np: Number of patches

102 + - bnp: Batch size * Number of patches

103 - c: Number of channels

104 - ps: Patch size

105 - ni: Number of images

@@ -110,7 +108,7 @@ class KeyeImagePixelInputs(TensorSchema):

108 type: Literal["pixel_values"]

109 pixel_values: Annotated[

110 torch.Tensor,

113 - TensorShape("b", "np", 3, "ps", "ps", dynamic_dims={"np"})]

111 + TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"})]

112 image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]

113

114

@@ -134,8 +132,7 @@ KeyeImageInputs = Union[KeyeImagePixelInputs, KeyeImageEmbeddingInputs]

132 class KeyeVideoPixelInputs(TensorSchema):

133 """

134 Dimensions:

137 - - b: Batch size

138 - - np: Number of patches

135 + - bnp: Batch size * Number of patches

136 - c: Number of channels

137 - ps: Patch size

138 - ni: Number of images

@@ -144,7 +141,7 @@ class KeyeVideoPixelInputs(TensorSchema):

141 type: Literal["pixel_values_videos"]

142 pixel_values_videos: Annotated[

143 torch.Tensor,

147 - TensorShape("b", "np", 3, "ps", "ps", dynamic_dims={"np"})]

144 + TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"})]

145 video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]

146

147

@@ -1258,6 +1255,8 @@ class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]):

1255

1256

1257 class BaseKeyeModule(nn.Module):

1258 + merge_by_field_config = True

1259 +

1260 packed_modules_mapping = {

1261 "qkv_proj": [

1262 "q_proj",

@@ -1524,28 +1523,6 @@ class KeyeForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

1523 prefix: str = "") -> nn.Module:

1524 return Projector(text_config, vision_config, quant_config, prefix)

1525

1527 - def _validate_and_reshape_mm_tensor(

1528 - self, mm_input: NestedTensors,

1529 - name: str) -> Union[torch.Tensor, list[torch.Tensor]]:

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

1531 - raise ValueError(f"Incorrect type of {name}. "

1532 - f"Got type: {type(mm_input)}")

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

1534 - if mm_input.ndim == 2:

1535 - return mm_input

1536 - if mm_input.ndim == 5:

1537 - return mm_input

1538 - if mm_input.ndim != 3:

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

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

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

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

1543 - elif is_list_of(mm_input, torch.Tensor):

1544 - if all(p.dim() == 4 for p in mm_input) or all(p.dim() == 2

1545 - for p in mm_input):

1546 - return mm_input

1547 - return torch.concat(mm_input)

1548 -

1526 def _parse_and_validate_image_input(

1527 self, **kwargs: object) -> Optional[KeyeImageInputs]:

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

@@ -1556,11 +1533,6 @@ class KeyeForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

1533 return None

1534

1535 if pixel_values is not None:

1559 - pixel_values = self._validate_and_reshape_mm_tensor(

1560 - pixel_values, "image pixel values")

1561 - image_grid_thw = self._validate_and_reshape_mm_tensor(

1562 - image_grid_thw, "image grid_thw")

1563 -

1536 return KeyeImagePixelInputs(

1537 type="pixel_values",

1538 pixel_values=pixel_values,

@@ -1568,11 +1540,6 @@ class KeyeForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

1540 )

1541

1542 if image_embeds is not None:

1571 - image_embeds = self._validate_and_reshape_mm_tensor(

1572 - image_embeds, "image embeds")

1573 - image_grid_thw = self._validate_and_reshape_mm_tensor(

1574 - image_grid_thw, "image grid_thw")

1575 -

1543 return KeyeImageEmbeddingInputs(

1544 type="image_embeds",

1545 image_embeds=image_embeds,

@@ -1589,13 +1556,6 @@ class KeyeForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

1556 return None

1557

1558 if pixel_values_videos is not None:

1592 - pixel_values_videos = self._validate_and_reshape_mm_tensor(

1593 - pixel_values_videos,

1594 - "video pixel values",

1595 - )

1596 - video_grid_thw = self._validate_and_reshape_mm_tensor(

1597 - video_grid_thw, "video grid_thw")

1598 -

1559 return KeyeVideoPixelInputs(

1560 type="pixel_values_videos",

1561 pixel_values_videos=pixel_values_videos,

@@ -1603,11 +1563,6 @@ class KeyeForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

1563 )

1564

1565 if video_embeds is not None:

1606 - video_embeds = self._validate_and_reshape_mm_tensor(

1607 - video_embeds, "video embeds")

1608 - video_grid_thw = self._validate_and_reshape_mm_tensor(

1609 - video_grid_thw, "video grid_thw")

1610 -

1566 return KeyeVideoEmbeddingInputs(

1567 type="video_embeds",

1568 video_embeds=video_embeds,

vllm/model_executor/models/keye_vl1_5.py

+5 -51

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

18 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

19 RowParallelLinear)

20 from vllm.model_executor.layers.quantization import QuantizationConfig

21 -from vllm.multimodal import MULTIMODAL_REGISTRY, NestedTensors

21 +from vllm.multimodal import MULTIMODAL_REGISTRY

22 from vllm.multimodal.inputs import (ImageItem, ModalityData,

23 MultiModalFieldConfig,

24 MultiModalKwargsItems, VideoItem)

@@ -100,8 +100,7 @@ def get_num_patches(grid_thw: torch.Tensor,

100 class KeyeVL1_5ImagePixelInputs(TensorSchema):

101 """

102 Dimensions:

103 - - b: Batch size

104 - - np: Number of patches

103 + - bnp: Batch size * Number of patches

104 - c: Number of channels

105 - ps: Patch size

106 - ni: Number of images

@@ -111,7 +110,7 @@ class KeyeVL1_5ImagePixelInputs(TensorSchema):

110

111 pixel_values: Annotated[

112 torch.Tensor,

114 - TensorShape("np", 3, "ps", "ps", dynamic_dims={"np"})]

113 + TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"})]

114

115 image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]

116

@@ -137,8 +136,7 @@ KeyeVL1_5ImageInputs = Union[KeyeVL1_5ImagePixelInputs,

136 class KeyeVL1_5VideoPixelInputs(TensorSchema):

137 """

138 Dimensions:

140 - - b: Batch size

141 - - np: Number of patches

139 + - bnp: Batch size * Number of patches

140 - c: Number of channels

141 - ps: Patch size

142 - ni: Number of images

@@ -147,7 +145,7 @@ class KeyeVL1_5VideoPixelInputs(TensorSchema):

145 type: Literal["pixel_values_videos"]

146 pixel_values_videos: Annotated[

147 torch.Tensor,

150 - TensorShape("np", 3, "ps", "ps", dynamic_dims={"np"})]

148 + TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"})]

149 video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]

150

151 num_frames: torch.Tensor

@@ -483,24 +481,6 @@ class KeyeVL1_5ForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

481 self.merge_size = config.vision_config.spatial_merge_size

482 super().__init__(vllm_config=vllm_config, prefix=prefix)

483

486 - def _validate_and_reshape_mm_tensor(self, mm_input: NestedTensors,

487 - expected_dim: int, name: str):

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

489 - raise ValueError(f"Incorrect type of {name}. "

490 - f"Got type: {type(mm_input)}")

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

492 - if mm_input.ndim == expected_dim:

493 - return mm_input

494 - elif mm_input.ndim == expected_dim + 1:

495 - return mm_input.reshape(-1, *mm_input.shape[2:])

496 - else:

497 - raise ValueError(

498 - f"{name} should be {expected_dim}D or "

499 - f"batched {expected_dim}D tensor."

500 - f"Got ndim: {mm_input.ndim} (shape={mm_input.shape})")

501 - else:

502 - return torch.concat(mm_input)

503 -

484 def _parse_and_validate_image_input(

485 self, **kwargs: object) -> Optional[KeyeVL1_5ImageInputs]:

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

@@ -511,11 +491,6 @@ class KeyeVL1_5ForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

491 return None

492

493 if pixel_values is not None:

514 - pixel_values = self._validate_and_reshape_mm_tensor(

515 - pixel_values, expected_dim=4, name="image pixel values")

516 - image_grid_thw = self._validate_and_reshape_mm_tensor(

517 - image_grid_thw, expected_dim=2, name="image grid_thw")

518 -

494 return KeyeVL1_5ImagePixelInputs(

495 type="pixel_values",

496 pixel_values=pixel_values,

@@ -523,11 +498,6 @@ class KeyeVL1_5ForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

498 )

499

500 if image_embeds is not None:

526 - image_embeds = self._validate_and_reshape_mm_tensor(

527 - image_embeds, expected_dim=2, name="image embeds")

528 - image_grid_thw = self._validate_and_reshape_mm_tensor(

529 - image_grid_thw, expected_dim=2, name="image grid_thw")

530 -

501 return KeyeVL1_5ImageEmbeddingInputs(

502 type="image_embeds",

503 image_embeds=image_embeds,

@@ -545,17 +515,6 @@ class KeyeVL1_5ForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

515 return None

516

517 if pixel_values_videos is not None:

548 - pixel_values_videos = self._validate_and_reshape_mm_tensor(

549 - pixel_values_videos,

550 - expected_dim=4,

551 - name="video pixel values",

552 - )

553 - video_grid_thw = self._validate_and_reshape_mm_tensor(

554 - video_grid_thw, expected_dim=2, name="video grid_thw")

555 -

556 - num_frames = self._validate_and_reshape_mm_tensor(

557 - num_frames, expected_dim=1, name="video num frames")

558 -

518 return KeyeVL1_5VideoPixelInputs(

519 type="pixel_values_videos",

520 pixel_values_videos=pixel_values_videos,

@@ -563,11 +522,6 @@ class KeyeVL1_5ForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

522 num_frames=num_frames)

523

524 if video_embeds is not None:

566 - video_embeds = self._validate_and_reshape_mm_tensor(

567 - video_embeds, expected_dim=2, name="video embeds")

568 - video_grid_thw = self._validate_and_reshape_mm_tensor(

569 - video_grid_thw, expected_dim=2, name="video grid_thw")

570 -

525 return KeyeVL1_5VideoEmbeddingInputs(type="video_embeds",

526 video_embeds=video_embeds,

527 video_grid_thw=video_grid_thw,

vllm/model_executor/models/kimi_vl.py

+1 -32

@@ -283,6 +283,7 @@ class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]):

283 dummy_inputs=KimiVLDummyInputsBuilder)

284 class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

285 SupportsPP):

286 + merge_by_field_config = True

287

288 supports_encoder_tp_data = True

289

@@ -342,23 +343,6 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

343 config.vocab_size, logit_scale)

344 self.media_placeholder: int = self.config.media_placeholder_token_id

345

345 - # ref: qwen2_vl.py

346 - def _validate_and_reshape_mm_tensor(self, mm_input: object,

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

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

349 - raise ValueError(f"Incorrect type of {name}. "

350 - f"Got type: {type(mm_input)}")

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

352 - if mm_input.ndim == 2:

353 - return mm_input

354 - if mm_input.ndim != 3:

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

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

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

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

359 - else:

360 - return torch.concat(mm_input)

361 -

346 def _parse_and_validate_image_input(

347 self, **kwargs: object) -> Optional[KimiVLImageInputs]:

348 # image input type must be pixel values now

@@ -368,21 +352,6 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal,

352 if pixel_values is None:

353 return None

354

371 - image_grid_hws = self._validate_and_reshape_mm_tensor(

372 - image_grid_hws, "image grid hws")

373 - # pixel_values may have complex shapes

374 - num_channels = 3

375 - patch_size = self.config.vision_config.patch_size

376 - if isinstance(pixel_values, list):

377 - pixel_values = torch.cat([

378 - x.reshape(-1, num_channels, patch_size, patch_size)

379 - for x in pixel_values

380 - ])

381 - else:

382 - pixel_values = pixel_values.reshape(-1, num_channels, patch_size,

383 - patch_size)

384 - pixel_values = pixel_values.to(self.vision_tower.dtype)

385 -

355 return KimiVLImagePixelInputs(

356 type="pixel_values",

357 pixel_values=pixel_values,

vllm/utils/tensor_schema.py

+6 -2

@@ -164,7 +164,9 @@ class TensorSchema:

164

165 if len(actual_shape) != len(expected_shape):

166 raise ValueError(f"{field_name} has rank {len(actual_shape)} "

167 - f"but expected {len(expected_shape)}")

167 + f"but expected {len(expected_shape)}. "

168 + f"Expected shape: {expected_shape}, "

169 + f"but got {actual_shape}")

170

171 for i, dim in enumerate(expected_shape):

172 if dim in dynamic_dims:

@@ -172,7 +174,9 @@ class TensorSchema:

174 elif isinstance(dim, int):

175 if actual_shape[i] != dim:

176 raise ValueError(f"{field_name} dim[{i}] expected "

175 - f"{dim}, got {actual_shape[i]}")

177 + f"{dim}, got {actual_shape[i]}. "

178 + f"Expected shape: {expected_shape}, "

179 + f"but got {actual_shape}")

180 elif isinstance(dim, str):

181 if dim in shape_env:

182 if actual_shape[i] != shape_env[dim]: