[Model] MiniCPM-o 4.5: optimize SigLIP position IDs (#5130)

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
Co-authored-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -936,47 +936,86 @@ class SiglipVisionEmbeddings(nn.Module):
         self.num_positions = self.num_patches
         self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
 
+    def _create_grid_position_ids(
+        self,
+        patch_grid_height: int,
+        patch_grid_width: int,
+        boundaries: torch.Tensor,
+    ) -> torch.Tensor:
+        fractional_coords_h = torch.arange(
+            0,
+            1 - 1e-6,
+            1 / patch_grid_height,
+            device=boundaries.device,
+        )
+        fractional_coords_w = torch.arange(
+            0,
+            1 - 1e-6,
+            1 / patch_grid_width,
+            device=boundaries.device,
+        )
+        bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)
+        bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)
+        return (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten()
+
+    def _create_position_ids(
+        self,
+        patch_attention_mask: torch.BoolTensor,
+        tgt_sizes: torch.IntTensor | None,
+        device: torch.device,
+    ) -> torch.Tensor:
+        batch_size = patch_attention_mask.size(0)
+        flat_patch_attention_mask = patch_attention_mask.reshape(batch_size, -1).to(device="cpu")
+
+        if tgt_sizes is None:
+            # As in the legacy path, infer rectangular grid sizes from a (B, H, W) mask.
+            target_sizes = torch.stack(
+                (
+                    patch_attention_mask[:, :, 0].sum(dim=1),
+                    patch_attention_mask[:, 0, :].sum(dim=1),
+                ),
+                dim=1,
+            )
+        else:
+            target_sizes = tgt_sizes
+
+        # Group construction on the host preserves the legacy floating-point
+        # mapping and is faster than device-side construction for vision grids.
+        target_sizes_list = target_sizes.detach().to(device="cpu").tolist()
+        position_ids = torch.zeros(flat_patch_attention_mask.shape, dtype=torch.long)
+        boundaries = torch.arange(
+            1 / self.num_patches_per_side,
+            1.0,
+            1 / self.num_patches_per_side,
+        )
+        grid_position_ids: dict[tuple[int, int], torch.Tensor] = {}
+
+        for batch_idx, (patch_grid_height, patch_grid_width) in enumerate(target_sizes_list):
+            grid_shape = (int(patch_grid_height), int(patch_grid_width))
+            grid_ids = grid_position_ids.get(grid_shape)
+            if grid_ids is None:
+                grid_ids = self._create_grid_position_ids(*grid_shape, boundaries)
+                grid_position_ids[grid_shape] = grid_ids
+
+            position_ids[batch_idx, flat_patch_attention_mask[batch_idx]] = grid_ids
+
+        return position_ids.to(device=device)
+
     def forward(
         self,
         pixel_values: torch.FloatTensor,
         patch_attention_mask: torch.BoolTensor,
         tgt_sizes: torch.IntTensor | None = None,
     ) -> torch.Tensor:
-        batch_size = pixel_values.size(0)
-
         patch_embeds = self.patch_embedding(pixel_values)
         embeddings = patch_embeds.flatten(2).transpose(1, 2)
 
-        max_im_h, max_im_w = pixel_values.size(2), pixel_values.size(3)
-        max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size
-        boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side)
-        position_ids = torch.full(
-            size=(
-                batch_size,
-                max_nb_patches_h * max_nb_patches_w,
-            ),
-            fill_value=0,
+        position_ids = self._create_position_ids(
+            patch_attention_mask,
+            tgt_sizes,
+            device=self.position_embedding.weight.device,
         )
 
-        for batch_idx, p_attn_mask in enumerate(patch_attention_mask):
-            if tgt_sizes is not None:
-                nb_patches_h = tgt_sizes[batch_idx][0]
-                nb_patches_w = tgt_sizes[batch_idx][1]
-            else:
-                nb_patches_h = p_attn_mask[:, 0].sum()
-                nb_patches_w = p_attn_mask[0].sum()
-
-            fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h)
-            fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w)
-
-            bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)
-            bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)
-
-            pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten()
-            position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids
-
-        position_ids = position_ids.to(self.position_embedding.weight.device)
-
         embeddings = embeddings + self.position_embedding(position_ids)
         return embeddings