执行摘要

• 一句话：单图 DP 视觉嵌入改 broadcast，耗时降低 7.6x
• 推荐动作：值得精读：该 PR 展示了如何用更轻量的 collective（broadcast）替换对称 all-gather，并借助手动验证脚本证明 bitwise 等价。关注点：长度计算的统一、attn_tp_group.broadcast 的抽象契约、以及单图分支的边界条件。建议后续补充跨后端（AMD/NPU/CPU）的 broadcast 兼容性测试。

功能与动机

PR body 指出：单图像请求时对称 all-gather 移动 tp_size x max_len 行，其中 (tp_size-1)/tp_size 是填充，每个 rank 还要分配 tp_size 倍大的缓冲区。作者希望改成从 owner rank broadcast，实现「一次拷贝而不是 tp_size 次」，并在 8x B300 上给出 7.6x 的实测加速。

实现拆解

1. 重构长度计算：在 python/sglang/srt/multimodal/mm_utils.py 的 run_dp_sharded_mrope_vision_model 中，将原先基于像素 patch 数的 grouped_pixel_values_len 改为基于输出 token 数的 output_tokens_per_image 和 grouped_output_lengths，使 max_len_per_rank 直接表示各 rank 输出 embedding 的实际行数，为 broadcast 路径提供准确的边界。
2. 新增单图快速路径：当 len(grid_thw_list) == 1 时，owner rank（由 image_to_tp_rank[0] 确定）直接使用 image_embeds_local.contiguous() 作为输出，其余 rank 分配 (n_tok, *shape[1:]) 空 buffer，然后调用 attn_tp_group.broadcast(out_embeddings, src=owner_local) 填充。该路径绕过了 _pad_mrope_vision_embeddings_for_tp_gather 和 all_gather，且保证多图分支完全不受影响。
3. 顺带修复 CPU 同步问题：packed_2d_rope 分支的 local_grid_thw 设备改为 None（CPU），避免 CUDA 上的 .tolist() 调用触发 host 同步，与广播路径的性能目标一致。
4. 测试配套：修改 test/registered/unit/models/test_kimi_k25.py，新增 test_dp_helper_broadcasts_a_single_image_from_its_owner_rank 验证单图走 broadcast 且 src 为 owner，并断言不会调用 all-gather；同时更新空 rank 场景测试以覆盖 broadcast 填充逻辑。新增 test/manual/vlm/verify_single_image_gather.py 手动验证脚本，用于在真实 NCCL 环境对比 bitwise 一致性和耗时。
5. 提交演进：4 个提交逐步完成——引入 broadcast、精简注释、覆盖 owner/receiver 双向测试、将 reproducer 移出包体放入 test/manual。

关键文件：

• python/sglang/srt/multimodal/mm_utils.py（模块 多模态；类别 source；类型 core-logic；符号 run_dp_sharded_mrope_vision_model）: 核心实现文件：新增单图 broadcast 快速路径，并重构长度计算，是性能提升的关键。
• test/registered/unit/models/test_kimi_k25.py（模块 模型测试；类别 test；类型 test-coverage；符号 test_dp_helper_broadcasts_a_single_image_from_its_owner_rank, test_dp_helper_uses_config_hidden_size_for_empty_moonvit3d_rank, _GatherGroup）: 新增单图 broadcast 路径的单元测试，验证 owner 源和禁止走 all-gather，保证回归安全。
• test/manual/vlm/verify_single_image_gather.py（模块 验证脚本；类别 test；类型 test-coverage；符号 main, path_a, path_b, timeit）: 手动验证脚本：在真实 NCCL 环境对比 all-gather 与 broadcast 的 bitwise 一致性和耗时，是性能结论的证据来源。

关键符号：run_dp_sharded_mrope_vision_model, test_dp_helper_broadcasts_a_single_image_from_its_owner_rank, test_dp_helper_uses_config_hidden_size_for_empty_moonvit3d_rank, main

关键源码片段

python/sglang/srt/multimodal/mm_utils.py

核心实现文件：新增单图 broadcast 快速路径，并重构长度计算，是性能提升的关键。

# 单图快速路径：单张图像只有一个 rank 持有 embedding
# 直接由 owner 广播，避免 pad-to-max all-gather 的冗余拷贝
# 条件：grid_thw_list 长度恰为 1（单个图像）
if len(grid_thw_list) == 1:
    owner_local = image_to_tp_rank[0] # owner 是负载均衡分配给该图的 rank
    n_tok = output_tokens_per_image[0] # 该图输出的 token 数（已除 embed_dim_reduction_factor）
    if tp_rank_local == owner_local:
        # owner rank 直接使用自己的输出，contiguous 保证 broadcast 源连续
        out_embeddings = image_embeds_local.contiguous()
    else:
        # 其他 rank 分配同样形状的空 buffer，由 broadcast 填充
        out_embeddings = torch.empty(
            (n_tok, *image_embeds_local.shape[1:]),
            dtype=input_dtype,
            device=input_device,
        )
    # 从 owner 广播，所有 rank 得到 bit-identical 结果
    get_parallel().attn_tp_group.broadcast(out_embeddings, src=owner_local)
    return out_embeddings

test/registered/unit/models/test_kimi_k25.py

新增单图 broadcast 路径的单元测试，验证 owner 源和禁止走 all-gather，保证回归安全。

def test_dp_helper_broadcasts_a_single_image_from_its_owner_rank():
    broadcast_src = []
​
    class _GatherGroup:
        # 单图场景绝不允许走到 all-gather，一旦触发即失败
        def all_gather(self, tensor, dim):
            raise AssertionError("a single image must not reach the all-gather")
​
        # 记录 broadcast 源，并直接把输入 tensor 当作返回值（形状断言在外部）
        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] # rank 0 是 owner

test/manual/vlm/verify_single_image_gather.py

手动验证脚本：在真实 NCCL 环境对比 all-gather 与 broadcast 的 bitwise 一致性和耗时，是性能结论的证据来源。

def path_a(): # 旧路径：pad-to-max all-gather + 重建 owner 行
    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(): # 新路径：从 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
​
# 逐位相等性校验：A==truth, A==B, B==truth
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)

评论区精华

PR 无 review 评论，但 Issue 评论中有两次 /rerun-test 操作。首次重跑 3 个测试在 ubuntu-latest 上失败，第二次扩展到 5 个测试后全部通过（含 test_mrope_encoder_utils.py、test_kimi_k25.py、test_kimi_vl.py、test_server_args.py、test_encoder_dp.py），说明测试不稳定或存在环境依赖，最终修复通过。

• CI 重跑验证单图 broadcast 路径 (testing): 最终所有指定测试通过，包括 test_mrope_encoder_utils、test_kimi_k25、test_kimi_vl、test_server_args、test_encoder_dp。

风险与影响

• 风险：
  1. broadcast API 依赖：快速路径假设 attn_tp_group 一定实现 broadcast 方法，若某些后端（如自定义并行组）只实现了 all_gather，会在单图请求时崩溃。需确认所有 TP 组实现均支持。
  2. 单图判定条件：len(grid_thw_list) == 1 是进入快速路径的唯一条件，若未来出现单图但跨 rank 分片（非 DP 分配）的情况，image_to_tp_rank[0] 可能不是唯一 owner，导致广播源错误。
  3. 空图防护：当 pixel_values_local 为空且 tp_rank_local == owner_local 时，image_embeds_local 是空张量，contiguous() 后广播空 buffer，其他 rank 填充后形状依赖 image_embeds_local.shape[1:]，可能存在形状不一致。
  4. 性能回归风险：单图路径现在不再走 all-gather，若 broadcast 实现效率低于预期或驱动有 bug，可能影响单图场景的端到端时延。
     - 影响：影响范围集中在多模态 DP 编码器的单图请求路径，覆盖 Kimi K25 等使用 rope_2d_packed 的模型。对多图请求和 tp_size == 1 路径无影响。实测单图 embedding 同步耗时从 0.69ms 降至 0.09ms，减少约 45MB 的跨 rank 数据移动，对高并发小图场景有明显收益。团队后续需要关注广播 API 在不同硬件（AMD、NPU）上的可用性与性能。
     - 风险标记：核心路径变更, 依赖 broadcast API 可用性, 单图分支边界条件

关联脉络

• 暂无明显关联 PR