[Perf] Broadcast single-image DP vision embedding instead of pad-to-max all-gather (#33307)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/srt/multimodal/mm_utils.py

@@ -640,8 +640,8 @@ def run_dp_sharded_mrope_vision_model(
     # image_to_tp_rank = [0, 2, 1, 3]
     # gpu_sample_counts = [1, 3]
     # grouped_pixel_values_len = [1000, 350]
-    image_to_tp_rank, gpu_sample_counts, grouped_pixel_values_len = (
-        get_dp_encoder_lb_assignment(patches_per_image, tp_size)
+    image_to_tp_rank, gpu_sample_counts, _ = get_dp_encoder_lb_assignment(
+        patches_per_image, tp_size
     )
 
     # cu_gpu_sample_counts = [0, 1, 4]
@@ -680,18 +680,35 @@ def run_dp_sharded_mrope_vision_model(
             vision_model.spatial_merge_size * vision_model.spatial_merge_size
         )
 
+    output_tokens_per_image = [
+        math.prod(grid) // embed_dim_reduction_factor for grid in grid_thw_list
+    ]
+    grouped_output_lengths = []
+    assignment_offset = 0
+    for sample_count in gpu_sample_counts:
+        rank_images = image_to_tp_rank[
+            assignment_offset : assignment_offset + sample_count
+        ]
+        grouped_output_lengths.append(
+            sum(output_tokens_per_image[i] for i in rank_images)
+        )
+        assignment_offset += sample_count
+
     # Find the max length across all ranks
     # The output embedding of every DP rank has to be
     # padded to this length for tensor_model_parallel_all_gather
     # to work
-    max_len_per_rank = max(grouped_pixel_values_len) // embed_dim_reduction_factor
+    max_len_per_rank = max(grouped_output_lengths)
     local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]
 
     # Run the vision model on the local pixel_values_local
     if packed_2d_rope:
         if pixel_values_local is not None and pixel_values_local.shape[0] > 0:
+            # Packed MoonViT reads grid_thw as CPU shape metadata. Placing it
+            # on CUDA would make each .tolist() call synchronize with the host.
             local_grid_thw = torch.tensor(
-                local_grid_thw_list, device=pixel_values_local.device
+                local_grid_thw_list,
+                device=(pixel_values_local.device if rope_type == "rope_2d" else None),
             )
             if rope_type == "rope_2d":
                 image_embeds_local = vision_model(
@@ -729,6 +746,22 @@ def run_dp_sharded_mrope_vision_model(
                 dtype=input_dtype,
             )
 
+    # Single-image fast path. Bit-identical to the all-gather below, which for
+    # one image just pads the owner's rows and slices them back out.
+    if len(grid_thw_list) == 1:
+        owner_local = image_to_tp_rank[0]
+        n_tok = output_tokens_per_image[0]
+        if tp_rank_local == owner_local:
+            out_embeddings = image_embeds_local.contiguous()
+        else:
+            out_embeddings = torch.empty(
+                (n_tok, *image_embeds_local.shape[1:]),
+                dtype=input_dtype,
+                device=input_device,
+            )
+        get_parallel().attn_tp_group.broadcast(out_embeddings, src=owner_local)
+        return out_embeddings
+
     # 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(
@@ -744,15 +777,9 @@ def run_dp_sharded_mrope_vision_model(
     rank_embeddings = list[torch.Tensor]()
     for rank in range(tp_size):
         start_idx = rank * max_len_per_rank
-        end_idx = start_idx + (
-            grouped_pixel_values_len[rank] // embed_dim_reduction_factor
-        )
+        end_idx = start_idx + grouped_output_lengths[rank]
         rank_embeddings.append(gathered_embeds[start_idx:end_idx])
 
-    patches_per_output_image = [
-        (patch_size // embed_dim_reduction_factor) for patch_size in patches_per_image
-    ]
-
     # Reconstruct embeddings in the original order
     original_order_embeddings = [None] * len(grid_thw_list)
     current_idx = 0
@@ -768,7 +795,7 @@ def run_dp_sharded_mrope_vision_model(
             # Split rank embeddings back to individual images
             embed_start = 0
             for img_idx in rank_images:
-                img_patches = patches_per_output_image[img_idx]
+                img_patches = output_tokens_per_image[img_idx]
                 original_order_embeddings[img_idx] = rank_embed[
                     embed_start : embed_start + img_patches
                 ]

test/manual/vlm/verify_single_image_gather.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+"""Verify the single-image broadcast fast path vs the pad-to-max all_gather:
+bitwise equivalence + timing, on real NCCL over 8 ranks. Mirrors what
+run_dp_sharded_mrope_vision_model does for a single image (one owner rank
+holds the embedding, the rest are empty). torchrun --nproc_per_node=8."""
+
+import time
+
+import torch
+import torch.distributed as dist
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank, world = dist.get_rank(), dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+    owner = 3  # arbitrary non-zero owner, as LB would pick
+    n_tok, hidden = 5476, 4096  # ~2048^2 image, typical tower output (~44MB bf16)
+
+    # deterministic ground-truth owner embedding, known to every rank
+    gen = torch.Generator(device=dev).manual_seed(12345)
+    owner_truth = torch.randn(
+        n_tok, hidden, dtype=torch.bfloat16, device=dev, generator=gen
+    )
+    emb = (
+        owner_truth.clone()
+        if rank == owner
+        else torch.empty(0, hidden, dtype=torch.bfloat16, device=dev)
+    )
+    max_len = n_tok  # single image: max over ranks == owner's length
+
+    def path_a():  # current: pad-to-max all_gather + reconstruct owner rows
+        padded = torch.empty(max_len, hidden, dtype=torch.bfloat16, device=dev)
+        if emb.shape[0] > 0:
+            padded[: emb.shape[0]].copy_(emb)
+        gathered = [
+            torch.empty(max_len, hidden, dtype=torch.bfloat16, device=dev)
+            for _ in range(world)
+        ]
+        dist.all_gather(gathered, padded)
+        return gathered[owner][:n_tok]
+
+    def path_b():  # fast path: broadcast from owner
+        buf = (
+            emb.contiguous()
+            if rank == owner
+            else torch.empty(n_tok, hidden, dtype=torch.bfloat16, device=dev)
+        )
+        dist.broadcast(buf, src=owner)
+        return buf
+
+    out_a, out_b = path_a(), path_b()
+    eq_truth = torch.equal(out_a, owner_truth)
+    eq_ab = torch.equal(out_a, out_b)
+    eq_b_truth = torch.equal(out_b, owner_truth)
+
+    def timeit(fn, n=100):
+        for _ in range(15):
+            fn()
+        torch.cuda.synchronize()
+        dist.barrier()
+        t0 = time.perf_counter()
+        for _ in range(n):
+            fn()
+        torch.cuda.synchronize()
+        dist.barrier()
+        return (time.perf_counter() - t0) / n * 1000
+
+    ta, tb = timeit(path_a), timeit(path_b)
+    # gather correctness flags from all ranks
+    flags = torch.tensor(
+        [eq_ab and eq_truth and eq_b_truth], device=dev, dtype=torch.int32
+    )
+    dist.all_reduce(flags, op=dist.ReduceOp.MIN)
+    if rank == 0:
+        print(
+            f"world={world} owner={owner} shape=[{n_tok},{hidden}] "
+            f"(~{n_tok*hidden*2/1e6:.0f}MB) | all_ranks_bitwise_ok={bool(flags.item())} "
+            f"(A==truth={eq_truth} A==B={eq_ab}) | all_gather {ta:.3f}ms "
+            f"broadcast {tb:.3f}ms speedup {ta/tb:.2f}x",
+            flush=True,
+        )
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

test/registered/unit/models/test_kimi_k25.py

@@ -126,10 +126,21 @@ def test_dp_helper_can_lazily_load_kimi_features_on_tp1():
 
 
 def test_dp_helper_uses_config_hidden_size_for_empty_moonvit3d_rank():
+    # Single image, so this empty rank takes the broadcast fast path: the
+    # buffer it allocates is shaped from config.hidden_size, then filled by
+    # the owner rank.
+    owner_embedding = torch.arange(8, dtype=torch.float32).reshape(1, 4, 2)
+    broadcast_src = []
+
     class _GatherGroup:
         def all_gather(self, tensor, dim):
             return torch.cat([torch.ones_like(tensor), tensor], dim=dim)
 
+        def broadcast(self, tensor, src):
+            broadcast_src.append(src)
+            tensor.copy_(owner_embedding)
+            return tensor
+
     tower = _MoonViT3dTower()
     parallel = SimpleNamespace(
         attn_tp_size=2,
@@ -146,9 +157,42 @@ def test_dp_helper_uses_config_hidden_size_for_empty_moonvit3d_rank():
         )
 
     assert output.shape == (1, 4, 2)
+    assert torch.equal(output, owner_embedding)
+    assert broadcast_src == [0]
     assert tower.grid_thws is None
 
 
+def test_dp_helper_broadcasts_a_single_image_from_its_owner_rank():
+    broadcast_src = []
+
+    class _GatherGroup:
+        def all_gather(self, tensor, dim):
+            raise AssertionError("a single image must not reach the all-gather")
+
+        def broadcast(self, tensor, src):
+            broadcast_src.append(src)
+            return tensor
+
+    tower = _MoonViT3dTower()
+    pixel_values = torch.randn(4, 2)
+    parallel = SimpleNamespace(
+        attn_tp_size=2,
+        attn_tp_rank=0,
+        attn_tp_group=_GatherGroup(),
+    )
+
+    with patch("sglang.srt.multimodal.mm_utils.get_parallel", return_value=parallel):
+        output = run_dp_sharded_mrope_vision_model(
+            tower,
+            pixel_values,
+            [[1, 2, 2]],
+            rope_type="rope_2d_packed",
+        )
+
+    assert torch.equal(output, pixel_values.reshape(1, 4, 2))
+    assert broadcast_src == [0]
+
+
 def test_dp_helper_lazily_loads_only_its_local_image_shard():
     class _GatherGroup:
         def all_gather(self, tensor, dim):