执行摘要

• 一句话：NCCL对称内存扩展至AllGather与ReduceScatter
• 推荐动作：值得精读：展现了将硬件加速（NCCL symmetric memory/NVLS）从单一操作扩展到多个collectives的工程模式。关注 _get_symm_scratch 的持久化设计、per-communicator隔离的修复方式，以及review中关于uniform size标准化的设计讨论。如有EP场景，建议跟进修复reduce_scatterv的标准化问题。

功能与动机

Enable NCCL symmetric memory and NVLS (NVLink SHARP) for AllGather and ReduceScatter collectives (used by MoE dispatch/combine in Expert Parallel deployments), extending the existing symmetric-memory support that was limited to AllReduce. Also fixes per-communicator window registration so multiple NCCL communicators can share the symmetric memory pool without double-registering segments.

实现拆解

1. 修改 CudaCommunicator 路由逻辑：重写 all_gather、reduce_scatter、reduce_scatterv 方法，当启用对称内存且操作为dim=0 uniform时，转发到新增的 _all_gather_symm_mem 和 _reduce_scatter_symm_mem 私有方法；否则回退到PyNCCL或基类实现。新增 _get_symm_scratch 管理持久化的对称内存scratch buffer，每个rank分配并注册一次。
2. 新增决策函数：在 all_reduce_utils.py 中添加 should_nccl_symm_mem_ag_rs，用于快速判断是否启用AG/RS对称内存路径（仅检查环境变量和batch invariant），与已有的 should_nccl_symm_mem_allreduce 的tensor-size权衡逻辑不同。
3. 修复多communicator窗口注册：在 pynccl_allocator.py 中将全局集合 _registered_base_addrs 从 set[address] 改为 dict[bytes, set]，以NCCL comm unique id隔离不同通信器的已注册地址，防止重复注册导致NVLS失败。
4. 整合并扩展测试：删除单文件 test_nccl_symm_mem_allreduce.py，新建 test_nccl_symm_mem.py 包含allreduce/allgather/reduce_scatter三类测试，每个测试在独立进程中spawn多rank worker验证正确性。
5. 更新CI配置：将 .buildkite/test_areas/distributed.yaml 中的测试命令指向新文件。

关键文件：

• vllm/distributed/device_communicators/cuda_communicator.py（模块 通信层；类别 source；类型 core-logic；符号 all_gather, _get_symm_scratch, _reduce_scatter_symm_mem, _all_gather_symm_mem）: 核心修改文件，实现所有新方法并修改路由逻辑
• tests/distributed/test_nccl_symm_mem.py（模块 测试；类别 test；类型 test-coverage；符号 nccl_symm_mem_allreduce_worker, test_nccl_symm_mem_allreduce, nccl_symm_mem_allgather_worker, test_nccl_symm_mem_allgather）: 新增整合测试文件，覆盖所有三种操作
• tests/distributed/test_nccl_symm_mem_allreduce.py（模块 测试；类别 test；类型 deletion；符号 nccl_symm_mem_allreduce_worker, test_nccl_symm_mem_allreduce）: 删除旧测试文件，内容已合并到新文件
• vllm/distributed/device_communicators/all_reduce_utils.py（模块 通信层；类别 source；类型 core-logic；符号 should_nccl_symm_mem_ag_rs）: 新增 should_nccl_symm_mem_ag_rs 决策函数
• vllm/distributed/device_communicators/pynccl_allocator.py（模块 通信层；类别 source；类型 core-logic）: 修复多comm窗口注册重复问题
• .buildkite/test_areas/distributed.yaml（模块 CI配置；类别 config；类型 configuration）: CI 配置更新，指向新测试文件

关键符号：all_gather, _get_symm_scratch, _reduce_scatter_symm_mem, _all_gather_symm_mem, _all_gather_batched_symm_mem, should_nccl_symm_mem_ag_rs, nccl_symm_mem_allreduce_worker, nccl_symm_mem_allgather_worker, nccl_symm_mem_reduce_scatter_worker, test_nccl_symm_mem_allreduce, test_nccl_symm_mem_allgather, test_nccl_symm_mem_reduce_scatter

关键源码片段

vllm/distributed/device_communicators/cuda_communicator.py

核心修改文件，实现所有新方法并修改路由逻辑

# vllm/distributed/device_communicators/cuda_communicator.py
​
class CudaCommunicator(DeviceCommunicatorBase):
    # ... (init, all_reduce unchanged) ...
​
    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
        # Route uniform dim-0 all-gathers through NVLS symmetric memory when
        # enabled (mirrors reduce_scatter); otherwise fall back to the
        # base-class ring all-gather. Sequence parallelism's gather-before-GEMM
        # uses dim=0 with tp-aligned (uniform) shards.
        if dim < 0:
            dim += input_.dim()
        if dim == 0 and should_nccl_symm_mem_ag_rs():
            # NVLS requires contiguous tensors
            return self._all_gather_symm_mem(input_.contiguous())
        return super().all_gather(input_, dim)
​
    def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
        world_size = self.world_size
        pynccl_comm = self.pynccl_comm
        assert pynccl_comm is not None
        if dim < 0:
            dim += input_.dim()
        # movedim to make dim-0 the scatter dimension
        input_tensor = input_.movedim(0, dim).contiguous()
        assert input_tensor.shape[0] % world_size == 0
        chunk_size = input_tensor.shape[0] // world_size
        output_shape = (chunk_size,) + input_tensor.shape[1:]
​
        if should_nccl_symm_mem_ag_rs():
            output = self._reduce_scatter_symm_mem(input_tensor)
        else:
            output = torch.empty(output_shape, dtype=input_tensor.dtype,
                                 device=input_tensor.device)
            pynccl_comm.reduce_scatter(output, input_tensor)
​
        return output.movedim(0, dim).contiguous()
​
    def reduce_scatterv(self, input_: torch.Tensor, dim: int = -1,
                        sizes: list[int] | None = None):
        # ... dimension normalization ...
        # Symmetric memory is only used when all ranks have uniform sizes.
        # ncclCommWindowRegister is collective: asymmetric pool allocations
        # from variable per-rank sizes cause deadlocks.
        use_symm_mem = sizes is None and should_nccl_symm_mem_ag_rs()
        if use_symm_mem:
            output = self._reduce_scatter_symm_mem(input_tensor)
        else:
            output = torch.empty(output_shape, dtype=input_tensor.dtype,
                                 device=input_tensor.device)
            if sizes is not None and sizes.count(sizes[0]) != len(sizes):
                pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
            else:
                pynccl_comm.reduce_scatter(output, input_tensor)
        return output.movedim(0, dim).contiguous()

tests/distributed/test_nccl_symm_mem.py

新增整合测试文件，覆盖所有三种操作

# tests/distributed/test_nccl_symm_mem.py
​
def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int):
    monkeypatch = pytest.MonkeyPatch()
    with monkeypatch.context() as m:
        m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
        dtype = torch.bfloat16
        device = torch.device(f"cuda:{local_rank}")
        torch.accelerator.set_device_index(device)
        torch.set_default_device(device)
        torch.set_default_dtype(dtype)
        update_environment_variables({
            "RANK": str(local_rank),
            "LOCAL_RANK": str(local_rank),
            "WORLD_SIZE": str(world_size),
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": "12346",
        })
​
        init_distributed_environment()
        with ensure_current_vllm_config():
            initialize_model_parallel(tensor_model_parallel_size=world_size)
​
        # 获取 CudaCommunicator 实例
        cuda_communicator = typing.cast(
            CudaCommunicator, get_tp_group().device_communicator)
        if get_nccl_mem_pool() is None:
            pytest.skip("NCCL allocator compilation failed")
        if not is_symmetric_memory_enabled():
            pytest.skip("NCCL symmetric memory is disabled.")
​
        # 每个 rank 准备本地输入，然后通过 CudaCommunicator 的 all_gather 进行集合通信
        per_rank_size = test_size_elements
        input_tensor = torch.rand(per_rank_size, dtype=dtype, device=device)
        gathered = cuda_communicator.all_gather(input_tensor, dim=0)
​
        # 验证：使用 PyTorch 原生 all_gather 作为参考
        ref_list = [torch.empty_like(input_tensor) for _ in range(world_size)]
        dist.all_gather(ref_list, input_tensor, group=get_tp_group().device_group)
        ref = torch.cat(ref_list, dim=0)
        torch.testing.assert_close(gathered, ref, atol=2.5, rtol=0.1)
​
​
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
@pytest.mark.parametrize("world_size", [2])
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="CUDA only")
def test_nccl_symm_mem_allgather(monkeypatch, world_size):
    if world_size > torch.accelerator.device_count():
        pytest.skip("Not enough GPUs")
    # 启用 SymmMemCommunicator 和相关环境变量
    monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
    monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
    monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
    mp.spawn(nccl_symm_mem_allgather_worker, args=(world_size,), nprocs=world_size)
    cleanup_dist_env_and_memory()

vllm/distributed/device_communicators/all_reduce_utils.py

新增 should_nccl_symm_mem_ag_rs 决策函数

# vllm/distributed/device_communicators/all_reduce_utils.py
​
def should_nccl_symm_mem_ag_rs() -> bool:
    """Check whether NCCL symmetric memory should be used for
    AllGather / ReduceScatter collectives."""
    from vllm.distributed.device_communicators.pynccl_allocator import (
        is_symmetric_memory_enabled,
    )
    # batch invariant mode disables custom collectives
    if envs.VLLM_BATCH_INVARIANT:
        return False
    # 仅依赖全局是否启用对称内存，不进行 tensor size 权衡
    return is_symmetric_memory_enabled()

评论区精华

claude[bot] 在 reduce_scatterv 的diff上指出：该函数在 sizes 非None时直接跳过NVLS路径，而 all_gatherv 会将uniform sizes标准化为None从而走NVLS路径。在EP combine流中（all2all.py:137），sizes 总是提供且steady-state下uniform，导致 combine 的 reduce_scatter 无法利用NVLS，抵消了本次PR的部分收益。建议在 reduce_scatterv 顶部添加相同的 if sizes is not None and all(s == sizes[0] for s in sizes): sizes = None 标准化逻辑。该评论未在合并前得到明确回应，当前代码仍未包含该标准化，可能需后续修复。

• reduce_scatterv 中 uniform sizes 未标准化导致 NVLS 路径绕过 (design): PR 已合并但代码中未包含该标准化，可能需后续 PR 修复。

风险与影响

• 风险：
  • 回归风险：修改了核心通信路径 all_gather、reduce_scatter、reduce_scatterv，条件判断错误可能导致结果错误或死锁。
  • 兼容性：所有变化在 VLLM_USE_NCCL_SYMM_MEM=1 后，默认关闭，不影响现有用户。
  • 窗口注册修复：将全局集合改为per-comm字典，可能影响之前多communicator场景下的memory pool共享行为，但更正确。
  • 未解决问题：reduce_scatterv 的 uniform size 标准化缺失可能在实际EP场景中绕过NVLS，降低部分收益。
• 影响：
  • 用户：默认无影响；启用后，使用Expert Parallelism且NVLS兼容硬件的用户可显著提升MoE dispatch/combine性能。
  • 系统：依赖NVLink SHARP（NCCL_NVLS_ENABLE=1）和CUDMEM（NCCL_CUMEM_ENABLE=1）支持。
  • 团队：需要维护两个决策函数（AR vs AG/RS），并确保后续优化不破坏优先级顺序。
  • 风险标记：核心通信路径变更, 默认关闭降低风险, 讨论中提及未解决问题

关联脉络

• PR #38549 [Original] NCCL symmetric memory AllReduce (squashed base): 本 PR 是其 squashed rebase，将对称内存从 AllReduce 扩展到 AllGather/ReduceScatter，并包含了其修复。