perf: avoid temporary VLM encoder gather padding (#31301)

python/sglang/srt/multimodal/mm_utils.py

@@ -486,6 +486,30 @@ def get_dp_encoder_lb_assignment(
     return (shuffle_indices, gpu_sample_counts, gpu_loads)
 
 
+def _pad_mrope_vision_embeddings_for_tp_gather(
+    image_embeds_local: torch.Tensor, max_len_per_rank: int
+) -> torch.Tensor:
+    """Pad the DP encoder output for a fixed-shape TP all-gather.
+
+    Allocating the padding fragment and then concatenating it creates two
+    temporary buffers on every underfilled rank. Allocate the final
+    fixed-shape input directly and copy just the valid embeddings instead.
+    """
+
+    current_len = image_embeds_local.shape[0]
+    if current_len >= max_len_per_rank:
+        return image_embeds_local
+
+    padded = torch.empty(
+        (max_len_per_rank, *image_embeds_local.shape[1:]),
+        dtype=image_embeds_local.dtype,
+        device=image_embeds_local.device,
+    )
+    if current_len > 0:
+        padded[:current_len].copy_(image_embeds_local)
+    return padded
+
+
 # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/vision.py
 def run_dp_sharded_vision_model(
     image_input: torch.Tensor, vision_model: torch.nn.Module
@@ -705,30 +729,11 @@ def run_dp_sharded_mrope_vision_model(
                 dtype=input_dtype,
             )
 
-    # Pad the output based on max_len_per_rank
-    # for tensor_model_parallel_all_gather to work
-    current_len = image_embeds_local.shape[0]
-    if current_len < max_len_per_rank:
-        padding_size = max_len_per_rank - current_len
-        if packed_2d_rope:
-            padding = torch.empty(
-                (
-                    padding_size,
-                    image_embeds_local.shape[1],
-                    image_embeds_local.shape[2],
-                ),
-                dtype=image_embeds_local.dtype,
-                device=image_embeds_local.device,
-            )
-        else:
-            padding = torch.empty(
-                (padding_size, image_embeds_local.shape[1]),
-                dtype=image_embeds_local.dtype,
-                device=image_embeds_local.device,
-            )
-        image_embeds_local_padded = torch.cat([image_embeds_local, padding], dim=0)
-    else:
-        image_embeds_local_padded = image_embeds_local
+    # The TP all-gather needs a common first dimension. Allocate that final
+    # shape directly instead of materializing a padding fragment and catting it.
+    image_embeds_local_padded = _pad_mrope_vision_embeddings_for_tp_gather(
+        image_embeds_local, max_len_per_rank
+    )
 
     # Do all_gather to collect embeddings from all ranks
     gathered_embeds = get_parallel().attn_tp_group.all_gather(