[Perf][Model] Batch PaddleOCR-VL vision encoder and projector (#61458)

Signed-off-by: ssingh231219 <simran.singh@artificialy.com>

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)