From 2eed35d738dbc7b0a5704ec6a5f928db3ed6144b Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 20 Jul 2026 12:54:29 +0800 Subject: [PATCH] perf: avoid temporary VLM encoder gather padding (#31301) --- python/sglang/srt/multimodal/mm_utils.py | 53 ++++++----- .../multimodal/test_mrope_encoder_utils.py | 92 +++++++++++++++++++ 2 files changed, 121 insertions(+), 24 deletions(-) create mode 100644 test/registered/unit/multimodal/test_mrope_encoder_utils.py diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index e50f50081..f57a97b41 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/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( diff --git a/test/registered/unit/multimodal/test_mrope_encoder_utils.py b/test/registered/unit/multimodal/test_mrope_encoder_utils.py new file mode 100644 index 000000000..a725c8853 --- /dev/null +++ b/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()