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

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 15, 17:15 UTC

[Misc] Clean up `scatter_patch_features` (#15559)

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

Cyrus Leung committed Mar 27, 2025 at 15:45 UTC e6c9053f9ec0b41e9af41def67537a4a3097eeb5

6 files changed +84 -138

vllm/model_executor/models/gemma3_mm.py

+6 -11

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

30 # yapf: enable

31 from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

32 from vllm.sequence import IntermediateTensors

33 -from vllm.utils import flatten_2d_lists

34 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

35 SupportsMultiModal, SupportsPP)

@@ -60,7 +59,7 @@ class Gemma3ImagePixelInputs(TypedDict):

59 A boolean mask indicating which image embeddings correspond

60 to patch tokens.

63 - Shape: `(batch_size, num_images, num_embeds)`

62 + Shape: `(batch_size * num_images, num_embeds)`

63 """

@@ -593,6 +592,7 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

592

593 pixel_values = flatten_bn(pixel_values, concat=True)

594 num_crops = flatten_bn(num_crops, concat=True)

595 + embed_is_patch = flatten_bn(embed_is_patch)

596

597 return Gemma3ImagePixelInputs(

598 type="pixel_values",

@@ -635,14 +635,10 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

635

636 image_features = self._process_image_input(image_input)

637

638 - if kwargs.get("v0_path", False):

639 - return image_features

640 -

641 - return flatten_2d_lists(

642 - scatter_patch_features(*args) for args in zip(

643 - image_features,

644 - image_input["embed_is_patch"],

645 - ))

638 + return scatter_patch_features(

639 + image_features,

640 + image_input["embed_is_patch"],

641 + )

642

643 def get_input_embeddings(

644 self,

@@ -671,7 +667,6 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

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

668 # condition is for v0 compatibility.

669 elif inputs_embeds is None:

674 - kwargs.update({"v0_path": True})

670 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

671

672 inputs_embeds = self.get_input_embeddings(input_ids,

vllm/model_executor/models/internvl.py

+9 -12

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

35 from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

36 from vllm.sequence import IntermediateTensors

37 from vllm.transformers_utils.tokenizer import AnyTokenizer

38 -from vllm.utils import flatten_2d_lists

39 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

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

@@ -66,13 +65,13 @@ class InternVLImagePixelInputs(TypedDict):

65 A boolean mask indicating which image embeddings correspond

66 to patch tokens.

69 - Shape: `(batch_size, num_images, num_embeds)`

68 + Shape: `(batch_size * num_images, num_embeds)`

69 """

72 class InternVLImageEmbeddingInputs(TypedDict):

73 type: Literal["image_embeds"]

75 - data: NestedTensors

74 + data: Union[torch.Tensor, list[torch.Tensor]]

75 """

76 A tensor of shape `(num_images, total_image_feature_size, hidden_size)`

77 or a list of tensors of shape `(total_image_feature_size, hidden_size)`

@@ -867,6 +866,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):

866

867 pixel_values_flat = flatten_bn(pixel_values_flat, concat=True)

868 image_num_patches = flatten_bn(image_num_patches, concat=True)

869 + embed_is_patch = flatten_bn(embed_is_patch)

870

871 return InternVLImagePixelInputs(

872 type="pixel_values",

@@ -881,7 +881,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):

881 def _process_image_input(

882 self,

883 image_input: InternVLImageInputs,

884 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

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

885 if image_input["type"] == "image_embeds":

886 return image_input["data"]

887

@@ -921,15 +921,13 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):

921

922 image_features = self._process_image_input(image_input)

923

924 - if (kwargs.get("v0_path", False)

925 - or image_input["type"] != "pixel_values"):

924 + if image_input["type"] != "pixel_values":

925 return image_features

926

928 - return flatten_2d_lists(

929 - scatter_patch_features(*args) for args in zip(

930 - image_features,

931 - image_input["embed_is_patch"],

932 - ))

927 + return scatter_patch_features(

928 + image_features,

929 + image_input["embed_is_patch"],

930 + )

931

932 def get_input_embeddings(

933 self,

@@ -964,7 +962,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):

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

963 # condition is for v0 compatibility.

964 elif inputs_embeds is None:

967 - kwargs.update({"v0_path": True})

965 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

966 inputs_embeds = self.get_input_embeddings(input_ids,

967 vision_embeddings)

vllm/model_executor/models/llava.py

+10 -12

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

35 PromptReplacement, PromptUpdate)

36 from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

37 from vllm.sequence import IntermediateTensors

38 -from vllm.utils import flatten_2d_lists

39 from .clip import CLIPVisionModel

40 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

@@ -73,7 +72,7 @@ class PixtralHFImagePixelInputs(TypedDict):

72 A boolean mask indicating which image embeddings correspond

73 to patch tokens.

76 - Shape: `(batch_size, num_images, num_embeds)`

75 + Shape: `(batch_size * num_images, num_embeds)`

76 """

@@ -618,6 +617,8 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

617 raise ValueError("Incorrect type of embed_is_patch. "

618 f"Got type: {type(embed_is_patch)}")

619

620 + embed_is_patch = flatten_bn(embed_is_patch)

621 +

622 return PixtralHFImagePixelInputs(

623 type="pixel_values_pixtral",

624 pixel_values=flatten_bn(pixel_values),

@@ -713,18 +714,16 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

714 if image_input is None:

715 return None

716

716 - vision_embeddings = self._process_image_input(image_input)

717 + image_features = self._process_image_input(image_input)

718

718 - if (kwargs.get("v0_path", False)

719 - or image_input["type"] != "pixel_values_pixtral"):

719 + if image_input["type"] != "pixel_values_pixtral":

720 # The path is used for pixtral (V0 only) and llava (V0/V1)

721 - return vision_embeddings

721 + return image_features

722

723 - return flatten_2d_lists(

724 - scatter_patch_features(*args) for args in zip(

725 - vision_embeddings,

726 - image_input["embed_is_patch"],

727 - ))

723 + return scatter_patch_features(

724 + image_features,

725 + image_input["embed_is_patch"],

726 + )

727

728 def get_input_embeddings(

729 self,

@@ -790,7 +789,6 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

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

790 # condition is for v0 compatibility.

791 elif inputs_embeds is None:

793 - kwargs.update({"v0_path": True})

792 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

793 inputs_embeds = self.get_input_embeddings(input_ids,

794 vision_embeddings)

vllm/model_executor/models/molmo.py

+32 -73

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

49 PromptInsertion, PromptUpdate)

50 from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

51 from vllm.sequence import IntermediateTensors

52 -from vllm.utils import flatten_2d_lists

53 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

54 SupportsMultiModal, SupportsPP, SupportsQuant)

@@ -72,17 +71,17 @@ POOLING_SIZE = 2

72 class MolmoImageInputs(TypedDict):

73 images: Union[torch.Tensor, list[torch.Tensor]]

75 - """Shape: `(batch_size, num_crops, num_patch, patch_dim)`"""

74 + """Shape: `(batch_size * num_images, num_crops, num_patch, patch_dim)`"""

76 image_masks: Optional[Union[torch.Tensor, list[torch.Tensor]]]

78 - """Shape: `(batch_size, num_crops, num_patch)`"""

77 + """Shape: `(batch_size * num_images, num_crops, num_patch)`"""

79 feat_is_patch: Union[torch.Tensor, list[torch.Tensor]]

80 """

81 A boolean mask indicating which image features correspond

82 to patch tokens.

85 - Shape: `(batch_size, num_crops, num_patch)`

84 + Shape: `(batch_size * num_images, num_crops, num_patch)`

85 """

87 embed_is_patch: Union[torch.Tensor, list[torch.Tensor]]

@@ -90,7 +89,7 @@ class MolmoImageInputs(TypedDict):

89 A boolean mask indicating which image embeddings correspond

90 to patch tokens.

93 - Shape: `(batch_size, num_embeds)`

92 + Shape: `(batch_size * num_images, num_embeds)`

93 """

95 num_crops: Union[torch.Tensor, list[torch.Tensor]]

@@ -696,9 +695,10 @@ class MolmoVisionBackbone(nn.Module, SupportsQuant):

695 return image_features

696

697 def forward(

699 - self, images: torch.Tensor, image_masks: torch.Tensor

700 - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:

701 -

698 + self,

699 + images: torch.Tensor,

700 + image_masks: torch.Tensor,

701 + ) -> torch.Tensor:

702 # image_features: (batch_size, num_crops(=num_image), num_patch, nximage_emb_dim) # noqa: E501

703 batch_size, num_image = images.shape[:2]

704 images = images.to(device=self.device, dtype=self.dtype)

@@ -1491,6 +1491,8 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1491 f"Got type: {type(img_patch_id)}")

1492 self.img_patch_id = img_patch_id.flatten().unique().item()

1493

1494 + embed_is_patch = flatten_bn(embed_is_patch)

1495 +

1496 return MolmoImageInputs(

1497 images=images,

1498 image_masks=image_masks,

@@ -1502,13 +1504,17 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1504 def _process_image_input(

1505 self,

1506 image_input: MolmoImageInputs,

1505 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

1506 - if isinstance(image_input["images"], list):

1507 + ) -> list[torch.Tensor]:

1508 + images = image_input["images"]

1509 + image_masks = image_input["image_masks"]

1510 + feat_is_patch = image_input["feat_is_patch"]

1511 + num_crops = image_input["num_crops"]

1512 +

1513 + if isinstance(images, list):

1514 # Call the vision backbone on the whole batch at once

1508 - images_flat = flatten_bn(image_input["images"], concat=True)

1509 - image_masks_flat = (None if (image_masks :=

1510 - image_input["image_masks"]) is None

1511 - else flatten_bn(image_masks, concat=True))

1515 + images_flat = flatten_bn(images, concat=True)

1516 + image_masks_flat = (None if image_masks is None else flatten_bn(

1517 + image_masks, concat=True))

1518

1519 image_features_flat = self.vision_backbone(

1520 images=images_flat.unsqueeze(0),

@@ -1517,63 +1523,19 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1523 ).squeeze(0)

1524

1525 # Reconstruct the batch dimension

1520 - image_features = image_features_flat.split(

1521 - image_input["num_crops"].sum(-1).tolist())

1526 + num_crops_per_image = [nc.sum().item() for nc in num_crops]

1527 + image_features = image_features_flat.split(num_crops_per_image)

1528 else:

1529 image_features = self.vision_backbone(

1524 - images=image_input["images"],

1525 - image_masks=image_input["image_masks"],

1530 + images=images,

1531 + image_masks=image_masks,

1532 )

1533

1528 - return image_features

1529 -

1530 - def _get_mm_embeds(

1531 - self,

1532 - features: torch.Tensor, # Shape: (num_crop, num_patch, d)

1533 - feat_is_patch: torch.Tensor, # Shape: (num_crop, num_patch)

1534 - num_crops: torch.Tensor, # Shape: (num_images,)

1535 - embed_is_patch: torch.Tensor, # Shape: (num_embeds,)

1536 - ) -> tuple[torch.Tensor, ...]:

1537 - """

1538 - Scatter the patch features into a contiguous tensor that corresponds

1539 - to the embedding tokens defined by the multimodal processor.

1540 -

1541 - Note:

1542 - The original code only considers patch tokens as feature

1543 - tokens, but our processor considers all image-related tokens

1544 - as feature tokens because the feature tokens need to be

1545 - consecutive in `input_ids`.

1546 -

1547 - Example:

1548 - A simplified example for one item in the batch:

1549 -

1550 - .. code-block::

1551 -

1552 - Embedding tokens (from HF processor):

1553 - [<start> <patch> <patch> <col> <patch> <patch> <col> <end> ]

1554 -

1555 - embed_is_patch (from HF processor):

1556 - [ False True True False True True False False ]

1557 -

1558 - Encoder outputs (from model):

1559 - [ p1 p2 0 p3 p4 0 ]

1560 -

1561 - feat_is_patch (from HF processor):

1562 - [ True True False True True False ]

1563 -

1564 - The resulting embedding tensor is:

1565 - [ nan p1 p2 nan p3 p4 nan nan ]

1566 - """

1567 - num_crops_per_image = num_crops.tolist()

1568 - feats_per_image = features.split(num_crops_per_image)

1569 - f_is_patch_per_image = feat_is_patch.split(num_crops_per_image)

1570 -

1571 - features = torch.cat([

1534 + # Only the features corresponding to patch tokens are relevant

1535 + return [

1536 feats[f_is_patch]

1573 - for feats, f_is_patch in zip(feats_per_image, f_is_patch_per_image)

1574 - ])

1575 -

1576 - return scatter_patch_features(features, embed_is_patch)

1537 + for feats, f_is_patch in zip(image_features, feat_is_patch)

1538 + ]

1539

1540 def get_multimodal_embeddings(

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

@@ -1583,13 +1545,10 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1545

1546 image_features = self._process_image_input(image_input)

1547

1586 - return flatten_2d_lists(

1587 - self._get_mm_embeds(*args) for args in zip(

1588 - image_features,

1589 - image_input["feat_is_patch"],

1590 - image_input["num_crops"],

1591 - image_input["embed_is_patch"],

1592 - ))

1548 + return scatter_patch_features(

1549 + image_features,

1550 + image_input["embed_is_patch"],

1551 + )

1552

1553 def get_input_embeddings(

1554 self,

vllm/model_executor/models/pixtral.py

+7 -11

@@ -42,7 +42,6 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

42 from vllm.sequence import IntermediateTensors

43 from vllm.transformers_utils.tokenizer import (MistralTokenizer,

44 cached_tokenizer_from_config)

45 -from vllm.utils import flatten_2d_lists

46 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

47 from .utils import (flatten_bn, init_vllm_registered_model, maybe_prefix,

@@ -74,7 +73,7 @@ class PixtralImagePixelInputs(TypedDict):

73 A boolean mask indicating which image embeddings correspond

74 to patch tokens.

77 - Shape: `(batch_size, num_images, num_embeds)`

76 + Shape: `(batch_size * num_images, num_embeds)`

77 """

@@ -387,6 +386,8 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

386 raise ValueError("Incorrect type of embed_is_patch. "

387 f"Got type: {type(embed_is_patch)}")

388

389 + embed_is_patch = flatten_bn(embed_is_patch)

390 +

391 return PixtralImagePixelInputs(

392 type="pixel_values",

393 images=flatten_bn(images),

@@ -428,14 +429,10 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

429

430 image_features = self._process_image_input(image_input)

431

431 - if kwargs.get("v0_path", False):

432 - return image_features

433 -

434 - return flatten_2d_lists(

435 - scatter_patch_features(*args) for args in zip(

436 - image_features,

437 - image_input["embed_is_patch"],

438 - ))

432 + return scatter_patch_features(

433 + image_features,

434 + image_input["embed_is_patch"],

435 + )

436

437 def get_input_embeddings(

438 self,

@@ -467,7 +464,6 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

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

465 # condition is for v0 compatibility.

466 elif inputs_embeds is None:

470 - kwargs.update({"v0_path": True})

467 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

468 inputs_embeds = self.get_input_embeddings(input_ids,

469 vision_embeddings)

vllm/model_executor/models/vision.py

+20 -19

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from abc import ABC, abstractmethod

4 +from collections.abc import Sequence

5 from typing import Final, Generic, Optional, Protocol, TypeVar, Union, cast

7 import torch

@@ -154,8 +155,8 @@ def resolve_visual_encoder_outputs(

155

156

157 def scatter_patch_features(

157 - features: torch.Tensor,

158 - embed_is_patch: Union[torch.Tensor, list[torch.Tensor]],

158 + patches: Union[torch.Tensor, Sequence[torch.Tensor]],

159 + embed_is_patch: Union[torch.Tensor, Sequence[torch.Tensor]],

160 ) -> tuple[torch.Tensor, ...]:

161 """

162 Scatter the patch features into a contiguous tensor that corresponds

@@ -165,8 +166,8 @@ def scatter_patch_features(

166 can be filtered out by :func`select_patch_features`.

167

168 Args:

168 - features: The patch features, concatenated across each image.

169 - Shape: `(num_patch, feature_depth)`

169 + patches: The patch features for each image.

170 + Shape: `(num_images, <patch_dims>, feature_depth)`

171 embed_is_patch: A boolean mask indicating which image embeddings

172 correspond to patch tokens for each image.

173 Shape: `(num_images, num_embeds)`

@@ -194,21 +195,21 @@ def scatter_patch_features(

195 The resulting embedding tensor is:

196 [ nan p1 p2 nan p3 p4 nan nan ]

197 """

197 - num_embeds_per_image = [

198 - e_is_patch.numel() for e_is_patch in embed_is_patch

199 - ]

200 - if isinstance(embed_is_patch, torch.Tensor):

201 - embed_is_patch_flat = embed_is_patch.view(-1)

202 - else:

203 - embed_is_patch_flat = torch.cat(embed_is_patch)

204 -

205 - embeds_flat = features.new_full(

206 - (sum(num_embeds_per_image), features.shape[-1]),

207 - fill_value=torch.nan,

208 - )

209 - embeds_flat[embed_is_patch_flat] = features.flatten(0, -2)

210 -

211 - return embeds_flat.split(num_embeds_per_image)

198 + if len(patches) != len(embed_is_patch):

199 + raise ValueError(f"Inconsistent num_images: {len(patches)=} vs. "

200 + f"{len(embed_is_patch)=}")

201 +

202 + def get_embed_one(patches_one: torch.Tensor, e_is_patch: torch.Tensor):

203 + embed_one = patches_one.new_full(

204 + (e_is_patch.shape[0], patches_one.shape[-1]),

205 + fill_value=torch.nan,

206 + )

207 + embed_one[e_is_patch] = patches_one.flatten(0, -2)

208 + return embed_one

209 +

210 + return tuple(

211 + get_embed_one(patches_one, e_is_patch)

212 + for patches_one, e_is_patch in zip(patches, embed_is_patch))

213

214

215 def select_patch_features(