diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index 4abf3dfb2..4065092df 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/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 ] diff --git a/test/manual/vlm/verify_single_image_gather.py b/test/manual/vlm/verify_single_image_gather.py new file mode 100755 index 000000000..5c298e200 --- /dev/null +++ b/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() diff --git a/test/registered/unit/models/test_kimi_k25.py b/test/registered/unit/models/test_kimi_k25.py index 885fbdf20..7f6c18dec 100644 --- a/test/registered/unit/models/test_kimi_k25.py +++ b/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):