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

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

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):