diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 1c165c08ef..0a2fa9a885 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -258,28 +258,18 @@ class PaddleOCRVLMultiModalProcessor( hf_kwargs, ) - def _postprocess_hf_mm_data( - self, - hf_data: Mapping[str, object], - hf_kwargs: Mapping[str, object], - processed_data: BatchFeature, - ) -> BatchFeature: - if not hf_data: - return processed_data - - num_patches_per_image = processed_data["image_grid_thw"].prod(-1) - processed_data["pixel_values"] = processed_data["pixel_values"].split( - num_patches_per_image.tolist() - ) - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, hf_processor_mm_kwargs: Mapping[str, object], ) -> Mapping[str, MultiModalFieldConfig]: + image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) + image_grid_sizes = image_grid_thw.prod(-1) + return dict( - pixel_values=MultiModalFieldConfig.batched("image"), + pixel_values=MultiModalFieldConfig.flat_from_sizes( + "image", image_grid_sizes + ), image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), ) @@ -347,12 +337,12 @@ class Projector(nn.Module): ) -> torch.Tensor: m1, m2 = self.merge_kernel_size if isinstance(image_features, (list, tuple)): - processed_features = list() - for image_feature, image_grid in zip(image_features, image_grid_thw): - image_feature = self.pre_norm(image_feature) - t, h, w = image_grid - - image_feature = rearrange( + # Norm and MLP are per-token, so run them once over all images; + # only the 2x2 merge depends on each image's grid. + sizes = [image_feature.shape[0] for image_feature in image_features] + normed_features = self.pre_norm(torch.cat(image_features)).split(sizes) + merged_features = [ + rearrange( image_feature, "(t h p1 w p2) d -> (t h w) (p1 p2 d)", t=t, @@ -361,12 +351,13 @@ class Projector(nn.Module): w=w // m2, p2=m2, ) - hidden_states = self.linear_1(image_feature) - hidden_states = self.act(hidden_states) - hidden_states = self.linear_2(hidden_states) - processed_features.append(hidden_states) + for image_feature, (t, h, w) in zip(normed_features, image_grid_thw) + ] + hidden_states = self.linear_1(torch.cat(merged_features)) + hidden_states = self.act(hidden_states) + hidden_states = self.linear_2(hidden_states) - return processed_features + return hidden_states.split([f.shape[0] for f in merged_features]) dims = image_features.shape[:-1] dim = image_features.shape[-1] @@ -383,7 +374,7 @@ class PaddleOCRImagePixelInputs(TensorSchema): type: Literal["pixel_values"] pixel_values: Annotated[ torch.Tensor, - TensorShape("bn", "p", 3, "patch_size", "patch_size", dynamic_dims={"p"}), + TensorShape("np", 3, "patch_size", "patch_size"), ] image_grid_thw: Annotated[ torch.Tensor, @@ -790,10 +781,10 @@ class SiglipEncoder(nn.Module): if cu_seqlens is None: raise ValueError("cu_seqlens cannot be None for SiglipEncoder.") if not isinstance(cu_seqlens, torch.Tensor): - cu_seqlens = async_tensor_h2d(cu_seqlens, dtype=torch.int32, device=device) - else: - cu_seqlens = cu_seqlens.to(device=device, non_blocking=True) + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32) + # Computed before the H2D copy: the attention wrappers call `.item()` + # on it, which would sync on every layer for a device tensor. max_seqlen = None if self.attn_backend in { AttentionBackendEnum.FLASH_ATTN, @@ -802,6 +793,11 @@ class SiglipEncoder(nn.Module): }: max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() + if cu_seqlens.is_cpu: + cu_seqlens = async_tensor_h2d(cu_seqlens, dtype=torch.int32, device=device) + else: + cu_seqlens = cu_seqlens.to(device=device, non_blocking=True) + hidden_states = inputs_embeds for encoder_layer in self.layers: hidden_states = encoder_layer( @@ -1104,55 +1100,43 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support ) def encode_image( - self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor + self, + pixel_values: torch.Tensor, + image_grid_thw: list[tuple[int, int, int]], ) -> torch.Tensor: + """Encode the flat patches of all images in one vision tower call.""" pixel_values = pixel_values.type(self.visual.dtype) - siglip_position_ids = list() - image_grid_hws = list() - cu_seqlens = [0] + seq_lens = [math.prod(grid) for grid in image_grid_thw] - thw_tuple = tuple(image_grid_thw.tolist()) - numel = np.prod(thw_tuple) - image_grid_hws.append(thw_tuple) - image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:]) - siglip_position_ids.append(image_position_ids) - cu_seqlens.append(cu_seqlens[-1] + numel) - - # Both are built on the host; concat straight into a pinned buffer - # so the H2D copy stays non-blocking. + # Built on the host; concat straight into a pinned buffer so the + # H2D copy stays non-blocking. siglip_position_ids = torch.concat( - siglip_position_ids, + [ + torch.arange(seq_len) % (h * w) + for seq_len, (_, h, w) in zip(seq_lens, image_grid_thw) + ], dim=0, - out=torch.empty( - sum(t.numel() for t in siglip_position_ids), - dtype=torch.int64, - pin_memory=PIN_MEMORY, - ), + out=torch.empty(sum(seq_lens), dtype=torch.int64, pin_memory=PIN_MEMORY), ).to(pixel_values.device, non_blocking=True) - cu_seqlens = async_tensor_h2d( - cu_seqlens, dtype=torch.int32, device=pixel_values.device - ) + # One segment per image so that images never attend to each other. + cu_seqlens = torch.from_numpy(np.cumsum([0, *seq_lens], dtype=np.int32)) - vision_outputs = self.visual( + return self.visual( pixel_values=pixel_values, - image_grid_thw=image_grid_hws, + image_grid_thw=image_grid_thw, position_ids=siglip_position_ids, interpolate_pos_encoding=True, cu_seqlens=cu_seqlens, ) - return vision_outputs def _process_image_input( self, image_input: PaddleOCRImagePixelInputs ) -> MultiModalEmbeddings: - pixel_values = image_input.pixel_values image_grid_thw = image_input.image_grid_thw - vision_outputs = tuple( - self.encode_image(pixel, grid).squeeze(0) - for pixel, grid in zip(pixel_values, image_grid_thw) - ) - image_embeds = self.mlp_AR(vision_outputs, image_grid_thw) - return image_embeds + grids = [tuple(grid) for grid in image_grid_thw.tolist()] + vision_outputs = self.encode_image(image_input.pixel_values, grids).squeeze(0) + vision_outputs = vision_outputs.split([math.prod(grid) for grid in grids]) + return self.mlp_AR(vision_outputs, image_grid_thw) def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings: image_input = self._parse_and_validate_image_input(**kwargs)