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
+    # The TP all-gather needs a common first dimension. Allocate that final
-    # for tensor_model_parallel_all_gather to work
+    # shape directly instead of materializing a padding fragment and catting it.
-    current_len = image_embeds_local.shape[0]
+    image_embeds_local_padded = _pad_mrope_vision_embeddings_for_tp_gather(
-    if current_len < max_len_per_rank:
+        image_embeds_local, 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
 
     # Do all_gather to collect embeddings from all ranks
     gathered_embeds = get_parallel().attn_tp_group.all_gather(

test/registered/unit/multimodal/test_mrope_encoder_utils.py

@@ -0,0 +1,92 @@
+"""CPU coverage for mRoPE DP vision-encoder helpers."""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.multimodal.mm_utils import (
+    _pad_mrope_vision_embeddings_for_tp_gather,
+    run_dp_sharded_mrope_vision_model,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class _RecordingGather:
+    def __init__(self):
+        self.input = None
+
+    def all_gather(self, input_, dim):
+        self.input = input_
+        rank_zero_embeddings = torch.full_like(input_, 99)
+        return torch.cat([rank_zero_embeddings, input_], dim=dim)
+
+
+class _Rope2dVisionTower:
+    merge_kernel_size = (1, 1)
+    config = SimpleNamespace(hidden_size=1)
+
+    def __call__(self, pixel_values, grid_hw, max_seqlen):
+        return pixel_values.reshape(-1, 1, 1)
+
+
+class TestMropeVisionEncoderPadding(CustomTestCase):
+    def test_padding_preserves_2d_embedding_prefix(self):
+        embeddings = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+
+        padded = _pad_mrope_vision_embeddings_for_tp_gather(embeddings, 5)
+
+        self.assertEqual(padded.shape, (5, 4))
+        self.assertTrue(torch.equal(padded[:3], embeddings))
+
+    def test_padding_preserves_3d_embedding_prefix(self):
+        embeddings = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)
+
+        padded = _pad_mrope_vision_embeddings_for_tp_gather(embeddings, 5)
+
+        self.assertEqual(padded.shape, (5, 3, 4))
+        self.assertTrue(torch.equal(padded[:2], embeddings))
+
+    def test_empty_rank_gets_a_gatherable_shape(self):
+        embeddings = torch.empty((0, 3, 4), dtype=torch.bfloat16)
+
+        padded = _pad_mrope_vision_embeddings_for_tp_gather(embeddings, 5)
+
+        self.assertEqual(padded.shape, (5, 3, 4))
+        self.assertEqual(padded.dtype, torch.bfloat16)
+
+    def test_already_full_embedding_is_not_copied(self):
+        embeddings = torch.randn(5, 4)
+
+        padded = _pad_mrope_vision_embeddings_for_tp_gather(embeddings, 5)
+
+        self.assertIs(padded, embeddings)
+
+    def test_dp_encoder_reconstructs_an_underfilled_rope2d_rank(self):
+        gather = _RecordingGather()
+        pixel_values = torch.tensor([[1], [2], [3], [4], [7]], dtype=torch.float32)
+
+        with get_parallel().override(
+            attn_tp_size=2,
+            attn_tp_rank=1,
+            attn_tp_group=gather,
+        ):
+            embeddings = run_dp_sharded_mrope_vision_model(
+                _Rope2dVisionTower(),
+                pixel_values,
+                [[2, 2], [1, 1]],
+                rope_type="rope_2d",
+            )
+
+        self.assertEqual(gather.input.shape, (4, 1, 1))
+        self.assertTrue(torch.equal(gather.input[:1], torch.tensor([[[7.0]]])))
+        self.assertTrue(torch.equal(embeddings[:4], torch.full((4, 1, 1), 99.0)))
+        self.assertTrue(torch.equal(embeddings[4:], torch.tensor([[[7.0]]])))
+
+
+if __name__ == "__main__":
+    unittest.main()