DeepInfer 源码拆解（四十）：集合通信的账本——五张进程组、DeepEP 的三类数据面与 NCCL handle 的复用规则

从“谁和谁通信”开始

前面几篇已经分别拆过 DeepInfer 的进程组构造、MoE 层的前向、上下文并行的 KV 收集与 Engram 的远程查表。把它们放在一起看，会浮现出一个更基础、也更少被单独讨论的问题：这些通信到底走的是哪一条通道，每条通道又和哪张进程组绑定？

这个问题值得单独成篇，原因有三个。

第一，DeepInfer 的并行拓扑（DP Attention + Wide-EP，不做张量并行）决定了它不会像 vLLM 或 TensorRT-LLM 那样在每层 attention 和 dense MLP 后面挂 all-reduce。它主要的通信发生在 MoE 的 dispatch/combine、CP 的 all-gather 以及 Engram 的远端查表上，三者用的是三套完全不同的数据面：DeepEP 的 EP 全互联缓冲、DeepEP 的通用集合缓冲、DeepEP 的 RDMA 存储缓冲。读代码时如果只看 userbuffer.allocate_* 这几个函数名，很容易误以为它们只是同一套缓冲的不同尺寸版本。

第二，PyTorch 进程组和 NCCL communicator 之间多了一层由 DeepEP 维护的 handle 缓存。torch.distributed 建一张组，DeepEP 拿这张组去取 NCCL 通信子；取的时候是复用 PyTorch 内部的 comm 指针，还是自己再建一个全新的 communicator，由一组环境变量和缓存键决定。对同一张组，谁先来谁就决定后面所有缓冲用的是同一个 communicator，这直接影响显存与通信资源的归属。

第三，这些通信并不是“能用 NCCL 就用 NCCL”。DeepInfer 的绝大多数走量通信，实际上绕开了 NCCL 的集合通信 kernel，换成了 DeepEP 自己的 dispatch/combine 内核与 DeepGEMM 的对称内存内核。

本文按“组 → 通道 → 缓冲 → 内核”的层次展开：

1. 五张核心进程组各自承担哪一类通信，哪些走 GPU 后端、哪些额外开 gloo；
2. DeepEP 的数据面分三类，分别是什么语义、由谁持有；
3. 一次 MoE 前向里 dispatch 与 combine 的两次搬运分别落在这三类数据面的哪一类上；
4. NCCL handle 的复用与新建规则，以及它为什么会影响显存；
5. 用一张账本把一次 decode step 的通信次数与字节数对齐。

先把集合的语义对齐

在进入 DeepInfer 的代码之前，值得把四个最常用的集合原语按“每个 rank 拿到的结果”重新定义一遍。很多实现细节的取舍，其实都来自语义差异，而不是性能偏好。

设有 $k$ 个 rank，每个 rank 持有一个长度为 $N$ 的数组 $\mathbf{in}_r$：

• AllReduce：每个 rank 都拿到同一个归约结果， $$\mathbf{out}[i] = \sum_{r=0}^{k-1} \mathbf{in}_r[i], \quad i = 0,\dots,N-1.$$
• ReduceScatter：先做同样的归约，但把结果按 rank 切成 $k$ 块，rank $r$ 只保留第 $r$ 块。每块长度为 $N/k$，总输出为 $N$。
• AllGather：不做归约，只把各 rank 的输入按 rank 序号首尾相接，输出长度为 $k \times N$；rank $r$ 的那一段固定落在输出第 $r$ 个槽位。
• AllToAll：每个 rank 的输入被切成 $k$ 段，第 $j$ 段发给 rank $j$；同时从每个 rank 接收一段，输出按来源 rank 排序。

这四个原语有两条恒等式，几乎决定了所有集合通信算法的形态：

第一条说明：AllReduce 可以拆成“先归约并分散、再全收集”两段，每段的通信量都是 $N(k-1)/k$ 量级，总通信量接近 $2N$，而不是把 $N$ 份数据广播 $k$ 次的 $kN$。这就是 NCCL 的环算法与双树算法都围绕 ReduceScatter/AllGather 组织的原因。第二条说明：AllReduce 也可以拆成“先归约到根、再广播”，适合小消息或树形拓扑。

AllReduce 的语义是“所有 rank 的输出完全相同”。在张量并行里，它对应每层 RowParallelLinear 之后把各卡的部分和加回去；DeepInfer 不做张量并行，所以这条路径在它的前向里不存在。

ReduceScatter 把归约结果按 rank 切开。每个 rank 的输出只有 $N/k$，比 AllReduce 的 $N$ 少，但要求下游能接受“每个 rank 只持有一段”的布局。DeepInfer 的 BucketBuffer 在 reduce_scatter 分支里按这个语义推导 SM 数，并在 hybrid 路径下按“multimem 归约 + RDMA 归约”两段计算读写量。

AllGather 是 CP 最常用的原语。它不做归约，只做拼接；输出顺序严格按 rank 序号，而不是按 token 在原始序列中的位置。DeepInfer 的 CPManager.all_gather 正是利用这一点：把 [cp_size, ...] 的结果按第 0 维展平，得到“按 CP rank 分组”的 token 块。至于这些块如何映射回全局 token 顺序，取决于 CP 的 zigzag/mirror 布局，而不是 all-gather 本身。

AllToAll 则是 MoE 的语义基础：每个 rank 的 token 要去往不同的专家 rank，且每个 token 的 top-k 路由都不一样，发送量按 rank 不规则。它既不是“所有人拿同一个结果”，也不是“简单拼接”，而是“每个 rank 给每个 rank 发不同的数据”。

AllToAll 的图示很直观：输入被等分成 $k$ 份，第 $j$ 份发往 rank $j$；输出也按来源 rank 等分成 $k$ 份。这正是 DeepEP 的 dispatch 与 combine 要解决的问题，也是它们不能用 NCCL all_to_all 直接替代的原因：后者要求每个 rank 的发送缓冲被等分成 $k$ 份，而 MoE 的每份长度由路由结果决定。

以上四张图的原始版本见 NCCL Collective Operations 文档。NCCL 的文档同时给出了这几条恒等式，后文讨论 DeepEP 的算法选择时会反复用到。

五张进程组：NCCL 负责走量，gloo 负责控制

先回顾一下 DeepInfer 源码拆解：并行拓扑 里已经出现的构造代码。ParallelProcessGroup.create_process_groups 用四次 einops rearrange 把全局 rank 切成五类组：

1
ranks = torch.arange(max(enp_size, dp_size * cp_size))
2
backend = PLATFORM.distributed_backend          # 对 CUDA 是 "nccl"，对 Ascend 是 "hccl"
3

4
# DP 组：固定 (replica, cp)，沿 dp 维建组
5
dp_ranks = rearrange(ranks, "(replica dp cp) -> (replica cp) dp", dp=dp_size, cp=cp_size)
6
self.dp = _split_group(dp_ranks.tolist(), backend=backend, pg_options=options)
7
self.cpu_dp_group = _split_group(dp_ranks.tolist(), backend="gloo")
8

9
# CP 组：固定 dp，沿 cp 维建组
10
cp_ranks = rearrange(ranks, "(dp cp) -> dp cp", cp=cp_size)
11
self.cp = _split_group(cp_ranks.tolist(), backend=backend, pg_options=options)
12
self.cpu_cp_group = _split_group(cp_ranks.tolist(), backend="gloo")
13

14
# EP 组：连续的 ep_size 个 rank
15
ep_ranks = rearrange(ranks, "(edp ep) -> edp ep", ep=ep_size)
16
self.ep = _split_group(ep_ranks.tolist(), backend=backend, pg_options=options)
17

18
# Engram（ENP）组：沿 enp 维建组
19
enp_ranks = rearrange(ranks, "(replica enp) -> replica enp", enp=enp_size)
20
self.enp = _split_group(enp_ranks.tolist(), backend=backend, pg_options=options)
21

22
# EP-in-DP 组：与 EP 相同的连续切法
23
self.ep_in_dp = _split_group(ep_in_dp_ranks.tolist(), backend=backend, pg_options=options)
24
self.cpu_ep_in_dp = _split_group(ep_in_dp_ranks.tolist(), backend="gloo")

拆开看，GPU 后端的组一共有 world、local、dp、cp、ep、enp、ep_in_dp 七个对象，但承担实际数据量的是 ep、cp 与 enp 三个：

1
@functools.lru_cache
2
def _split_group_cached(split_ranks, backend, pg_options=None):
3
    my_group = None
4
    for ranks in split_ranks:
5
        group = dist.new_group(list(ranks), backend=backend, pg_options=pg_options)
6
        if dist.get_rank() in ranks:
7
            my_group = group
8
    return my_group

1
def init_device_groups(self) -> None:
2
    seen: set[int] = set()
3
    for value in self.__dict__.values():
4
        try:
5
            if not isinstance(value, ProcessGroup) or dist.get_backend(value) not in ("nccl", "hccl"):
6
                continue
7
        except Exception:
8
            continue
9
        if id(value) not in seen:
10
            seen.add(id(value))
11
            dist.all_reduce(torch.zeros(1).cuda(), group=value)

1
单 rank 组也不能跳过：engram deep_ep 经 `backend._comm_ptr()` 复用这里建出的 comm，空指针会直接段错误。

DeepEP 的三类数据面

DeepEP 当前的版本里，deep_ep 包对外暴露的缓冲类不止一个。从 deep_ep/__init__.py 的导出列表可以直接看到三类：

1
from .buffers.ep import EPBuffer, EPHandle
2
from .buffers.engram import EngramBuffer
3
from .buffers.bucket import BucketBuffer, BucketSession

三者都继承自同一个 BufferBase，共享“显式销毁”的生命周期，但语义完全不同：

• EPBuffer：专家并行的全互联缓冲。入口是 dispatch 与 combine 两个方法，参数是 token、topk_idx、topk_weights，返回的是被路由后的 token 张量与一个 EPHandle。它解决的是“每个 token 按路由去往不同 rank 的专家”这种不规则通信，而不是固定切分下的规约。
• BucketBuffer：通用的集合通信缓冲，方法有 all_reduce、reduce_scatter、all_gather 与 barrier。它解决的是“在已注册的对称存储上做定长、批量、可融合的通信”，典型调用者是 CP。
• EngramBuffer：远端查表的存储缓冲。它有 num_gpu_bytes（接收区）与 num_rdma_storage_bytes（注册给 RDMA 的存储）两个尺寸参数，入口是查表请求而不是搬 token。

一张表可以概括这三类数据面在 DeepInfer 里的对应关系：

第四行不是 DeepEP 的类，而是 DeepGEMM 的对称内存缓冲。它和 EPBuffer 绑定同一张组，但语义完全不同：EPBuffer 是两个独立 kernel 与中间的显存暂存，SymmBuffer 是一块所有 rank 都可直接读写的对称地址空间，供融合后的 MegaMoE 内核在计算中途交换数据。这一层在 userbuffer，四类通信缓冲与对称内存 一篇里已经展开，本文只把它放进账本里作对照。

1
def dispatch(self, x, topk_idx=None, topk_weights=None, ...,
2
             num_experts=None, num_max_tokens_per_rank=None,
3
             expert_alignment=None, num_sms=0, num_qps=0,
4
             previous_event=None, async_with_compute_stream=False,
5
             allocate_on_comm_stream=False, handle=None,
6
             do_expand=False, defer_epilogue=False):

1
event_overlap = self.ep_buffer.dispatch(
2
    x=hidden_states,
3
    num_sms=envs.DEEPINFER_USERBUFFER_COMM_SMS,
4
    topk_idx=topk_idx,
5
    topk_weights=topk_weights,
6
    cumulative_local_expert_recv_stats=self._eplb_layer().expert_count,
7
    num_experts=self.num_physical_experts,
8
    num_max_tokens_per_rank=(self.max_num_prefilling_tokens
9
                             if input_metadata.is_prompt
10
                             else self.max_num_decoding_tokens),
11
    expert_alignment=deep_gemm.get_mk_alignment_for_contiguous_layout(),
12
    previous_event=dispatch_pre_evt,
13
    async_with_compute_stream=True,
14
    allocate_on_comm_stream=dispatch_pre_evt is not None,
15
    do_expand=True,
16
    use_tma_aligned_col_major_sf=True,
17
    defer_epilogue=True,
18
)

1
event_overlap = self.ep_buffer.combine(
2
    x=hidden_states,
3
    handle=handle,
4
    num_sms=envs.DEEPINFER_USERBUFFER_COMM_SMS,
5
    bias=shared_hidden_states,
6
    previous_event=combine_pre_evt,
7
    async_with_compute_stream=True,
8
    allocate_on_comm_stream=combine_pre_evt is not None,
9
    defer_epilogue=True,
10
)

1
num_sms=envs.DEEPINFER_USERBUFFER_COMM_SMS,  # TODO: Use all SMs when DeepEP support different #SM for dispatch and combine.

SM 预算是怎么算出来的

get_theoretical_num_sms 是前面多次提到的“带宽模型”，值得单独展开，因为它是唯一一处把“通信要占多少 SM”写成闭式公式的地方。

它先解决一个概率问题：在均衡门控下，一个 token 的 $K$ 个 top-k 专家里，期望有多少个落在“其它”组？DeepEP 给出的答案是

其中 $E$ 是专家总数，$G$ 是组的个数（可以是 rank 数，也可以是 scale-out rank 数），$K$ 是 top-k。直觉是：从 $E$ 个专家里选 $K$ 个，全部避开某个组（组内有 $E/G$ 个专家）的概率是 $\binom{E-E/G}{K}/\binom{E}{K}$，所以至少命中一次的概率就是括号里的部分；再乘上组数，就是“至少命中一个其它组、且把每个其它组按期望计一次”的近似。

这个公式把“路由分布”压缩成一个标量。以 $E=256$、$K=6$、$G=64$ 为例，每个组有 4 个专家，$K$ 个选择全部落在同一组的概率约为 0.909，期望有 0.091 的组被命中，乘上 64 就是约 5.8 个 top-k 落在“其它组”上。换成 $G=8$、每组 32 个专家，概率降到 0.445，期望命中数升到 4.44。代码里的注释特意写了“这是针对均衡门控的估计；DeepSeek-V3 系列的分组限制门控请不要用这个函数”，因为分组门控会强制至少命中一个分组，上述概率计算不再成立。

拿到期望值后，函数按流量守恒估 SM 数。以跨机 EP（num_scaleout_ranks > 1）为例：

• 读 token：每个 token 平均需要读 $1 / \mathbb{E}[\text{topk}]$ 份 hidden，所以 sm_read += 1 / num_expected_topk。
• 写发送缓冲：如果跨节点，先写一份发送缓冲，sm_write += 1 / num_expected_topk。
• scale-out 流量：其中一份需要跨到节点外，另一份留在本地，rdma_traffic 与本地写各占一定比例。
• 转发 warp：跨节点的那部分需要在本地再经 NVLink 转发到具体专家所在 rank，所以 nvlink_traffic 增加。

然后取 RDMA 与 NVLink 两条链路中“时间更长”的那条作为瓶颈（bounded_traffic / bounded_gbs），并以它为基准换算 SM：

这里的 $s_{\text{read}}$ 与 $s_{\text{write}}$ 是单位数据经过 SM 的读写次数，$B$ 是单个 SM 的 HBM 读写带宽。直观上，$C$ 是“每单位流量可以给每个 SM 多少时间”，乘上读写次数再除以 SM 带宽，就是“一个 SM 在链路上能压住多少流量”。

最后再做三件事：取 max(4, ceil(num_sms * 1.3)) 留 30% 余量并对齐到 4；如果 prefer_overlap_with_compute 为真，则允许用更少的 SM；最后把上限压到不超过设备 SM 数的最大偶数（num_device_sms // 2 * 2）。前两个调整是工程余量，最后一个调整保证通信不会申请超过整卡的 SM 数，并保持偶数便于对齐。

这个模型有三条明确的边界：不支持 do_expand，不支持 allow_multiple_reduction，且只针对均衡门控。DeepInfer 调用时开启了 do_expand=True，所以实际使用的 SM 数由 DEEPINFER_USERBUFFER_COMM_SMS 环境变量或 allocate_ep 里的“扣掉计算 SM”逻辑决定，而不是直接使用模型结果。模型的价值在于给出一个量级正确的起点，避免用户凭感觉填一个数。

QP（Queue Pair）数量的推导紧随其后：

1
# For direct mode, we encourage less QPs to reduce DB ringing overhead
2
num_qps = min(num_sms, 8 + 1)
3
# For hybrid mode, we encourage every channel (and notify) to have an independent QP
4
if self.allow_hybrid_mode:
5
    num_qps = num_sms * 16 + 1
6
return min(num_qps, self.num_allocated_qps)

在纯 RDMA 的 direct 模式下，QP 数被压到 8，因为多个 QP 会对同一块门铃寄存器反复写，产生“DB ringing”这种不必要的开销。在 hybrid 模式下，每个通道都需要独立的 QP 以并行读写，所以是 SM 数的 16 倍再加 1 个用于通知。allocate_ep 里默认分配的 QP 数遵循同样思路：hybrid 模式且 NIC 支持快速原子操作时 65 个，否则 129 个；非 hybrid 模式 17 个。这些数字不是魔法，而是“每通道 N 个 + 通知 1 个”的固定配置。

1
self.runtime = _C.BucketBuffer(
2
    [item.rank() for item in self.groups],
3
    [item.size() for item in self.groups],
4
    [item.get() for item in self.nccl_comm_handles],
5
    num_bytes, sl_idx, num_gpu_timeout_secs, nvl_gbs, rdma_gbs,
6
    explicitly_destroy)
7
self.contexts = self.runtime.contexts
8
self.storage  = self.runtime.storage

1
if num_rdma_ranks == 1 and num_nvl_ranks > 1:          # 纯 NVLink
2
    dst_gbs = get_nvlink_gbs() / num_nvl_ranks
3
    num_sms = max(dst_gbs / get_sm_read_gbs(), dst_gbs / get_sm_write_gbs())
4
    return align(max(6, math.ceil(num_sms)), 2)
5

6
if num_nvl_ranks > 1 and num_rdma_ranks > 1:          # Hybrid
7
    dst_gbs = min(get_nvlink_gbs() / num_nvl_ranks,
8
                  get_rdma_gbs() * num_rdma_ranks / (num_rdma_ranks - 1))
9
    # Multimem 与本地归约各需要一次读与一次写
10
    num_sms = max(2 * dst_gbs / get_sm_read_gbs(),
11
                  2 * dst_gbs / get_sm_write_gbs())
12
    return align(max(8, math.ceil(num_sms * 1.2)), 2)

1
num_sms: the number of SMs to use. Must be ``0`` or ``None`` (copy-engine driven).

1
if op == 'all_gather':
2
    return 0

1
def all_gather(self, x):
2
    assert self._session is not None
3
    if not x.is_contiguous():
4
        contiguous_x = self.get_inplace_tensor(x.shape, x.dtype)
5
        contiguous_x.copy_(x)
6
        x = contiguous_x
7
    handle = self._buffer.all_gather(x)
8
    gathered_shape = (self._size, *x.shape)
9
    return lambda: handle.wait().view(gathered_shape).flatten(0, 1)

1
class BucketSession:
2
    def __init__(self, buffer):
3
        self.buffer = buffer
4
        self.alignment = _C.get_num_rdma_alignment()
5
        self.num_allocated_bytes = 0
6
        self.context_idx = None
7
        self.collective_methods = ('all_gather', 'reduce_scatter', 'all_reduce')
8

9
    def __enter__(self):
10
        assert not any(name in self.buffer.__dict__ for name in self.collective_methods), \
11
            'Nested sessions are not supported'
12
        self.num_allocated_bytes = 0
13
        self.context_idx = None
14
        for name in self.collective_methods:
15
            method = partial(self._collective_func, name)
16
            setattr(self.buffer, name, method)
17
        return self
18

19
    def __exit__(self, exc_type, exc_val, exc_tb):
20
        for name in self.collective_methods:
21
            delattr(self.buffer, name)

1
def _allocate(self, num_bytes, dtype=torch.uint8):
2
    allocated_bytes = align(num_bytes, self.alignment)
3
    offset = self.num_allocated_bytes
4
    if offset + allocated_bytes > self.buffer.num_bytes:
5
        raise RuntimeError(f'BucketBuffer session out of storage, ...')
6
    self.num_allocated_bytes += allocated_bytes
7
    return self.buffer.storage.narrow(0, offset, num_bytes).view(dtype)

1
@contextmanager
2
def session(self):
3
    assert self._session is None
4
    with self._buffer.session() as session:
5
        self._session = session
6
        try:
7
            yield session
8
        finally:
9
            self._session = None

1
num_registered_bytes = num_gpu_bytes + num_rdma_storage_bytes * num_max_local_ranks + (1 << 32)
2
num_total_gpu_bytes = torch.cuda.get_device_properties('cuda').total_memory
3
if num_registered_bytes > num_total_gpu_bytes:
4
    os.environ['NCCL_WIN_STRIDE'] = str(align(num_registered_bytes, 1 << 32))

1
use_cpu_rdma_storage = not envs.DEEPINFER_ENGRAM_FORCE_GPU_STORAGE

NCCL handle：进程组之上的一层缓存

五张进程组与三类数据面之间，还有一层经常被忽略的胶水：DeepEP 在 comm/handle.py 里维护的 NCCLCommHandle 缓存。

1
def get_nccl_comm_handle(group, force_new_comm=False):
2
    global _storage
3
    if not force_new_comm and group in _storage:
4
        return _storage[group]
5

6
    backend = group._get_backend(torch.device('cuda'))
7
    if (not force_new_comm and hasattr(backend, '_comm_ptr')
8
            and int(os.getenv('EP_REUSE_NCCL_COMM', '1'))):
9
        _storage[group] = NCCLCommHandle(backend._comm_ptr(), False)
10
        return _storage[group]
11

12
    # 老版本 PyTorch 或强制新建：自己 all_gather 一个 unique id 出来
13
    nccl_unique_ids = [None] * group.size()
14
    dist.all_gather_object(nccl_unique_ids, _C.get_local_nccl_unique_id(), group)
15
    ...

这几行分支构成了三条路径：

1. 缓存命中。同一个 ProcessGroup 对象第二次来取 handle，直接返回第一次建好的那个。注意缓存键是组对象本身，所以 PG.ep 与 PG.ep_in_dp 即使 rank 集合完全一样，只要是两个不同的 Python 对象，就会各自拿到一个 handle。
2. 复用 PyTorch 的 comm 指针。如果当前 PyTorch 版本支持 backend._comm_ptr()，并且没有设置 EP_REUSE_NCCL_COMM=0，DeepEP 就直接借用 PyTorch 为这张组建好的 NCCL communicator。这条路径不新分配通信器，也不会多占一份 NCCL 内部缓冲。
3. 强制新建。老版本 PyTorch 没有 _comm_ptr，或者调用者传了 force_new_comm=True 时，DeepEP 会自己 all_gather_object 一个 unique id，再调用 _C 侧的 ncclCommInitRank 建一个全新的 communicator，由 DeepEP 自己管理销毁。

为什么“复用指针”比“新建”好？在 MoE 场景下，EP 组往往就是 world_size 全部 rank，如果 PG.ep 与 PG.enp 都各自新建 communicator，同一批 rank 之间就会存在多个 NCCL communicator。每个 NCCL communicator 都会占用一定的显存用于内部缓冲与连接元数据，而且按 rank 数、通道数与协议开销增长；同一批 rank 上多建几个 communicator，多出来的显存就是纯开销。DeepInfer 的做法是让 DeepEP 尽量借用 PyTorch 已有的 communicator，只在必要时才新建。

get_nccl_comm_handle 的复用还影响销毁顺序。PyTorch 的 process group 要在所有 worker 的通信完成后才销毁；如果 DeepEP 新建了独立 communicator，它需要在 destroy_all_managed_nccl_comm 里主动清空缓存。DeepInfer 的 WorkerBase.stop 里就有相应的处理：

1
dist.barrier()
2
userbuffer.close()
3
dist.barrier()
4
self.model_runner.destroy_cuda_graphs()
5
PG.cleanup_shm_broadcast()
6
dist.destroy_process_group()

两个 barrier 夹着 userbuffer.close()。前一个 barrier 保证所有 rank 都完成了最后一轮通信；close() 里先销毁 EP、CP、Engram 缓冲，再逆序销毁 SymmBuffer 视图；后一个 barrier 保证所有 rank 的缓冲都释放完毕，再统一销毁 process group。如果顺序反了——先销毁 process group 再销毁缓冲——NCCL communicator 可能在缓冲析构时已经被释放，DeepEP 的析构路径会访问空指针。这是这类“库中库”结构反复出现的坑，worker 握手 与 userbuffer 四类缓冲 两篇里都提过。

通信与计算的交叠：事件、流与所有权

DeepEP 的 dispatch 与 combine 都返回 EventOverlap，而不是直接返回张量。这个返回类型是整套异步设计的核心。

EventOverlap 同时承担三件事：

1. 记录通信内核在通信流上的完成事件；
2. 提供一个 current_stream_wait()，把当前计算流挂到该事件上；
3. 允许注册一个“等待之后才执行”的回调，用来做延迟的归约尾部（deferred epilogue）。

在 MoE 前向里，调用顺序是：

1
event_overlap = self.dispatch(hidden_states, topk_idx, topk_weights, input_metadata, None)
2
moe_in_hidden, _, dispatched_topk_weights, handle = event_overlap.current_stream_wait()
3

4
moe_out_hidden = self.routed_experts_forward(
5
    moe_in_hidden, dispatched_topk_weights,
6
    handle.psum_num_recv_tokens_per_expert)
7

8
combine_pre_evt = self.ep_buffer.capture()
9
with limit_num_sms(envs.DEEPINFER_USERBUFFER_COMM_SMS, enabled=PLATFORM.is_gpu):
10
    shared_out = self.shared_forward(hidden_states)
11

12
event_overlap = self.combine(moe_out_hidden, shared_out, handle, combine_pre_evt)
13
moe_combined_hidden, _ = event_overlap.current_stream_wait()

这条链路里，dispatch 返回后并没有立刻阻塞；current_stream_wait() 才把计算流与通信流对齐。对齐之后，专家 GEMM 在计算流上跑。与此同时，共享专家的 GEMM 与 combine 的通信被安排成可重叠的两段：capture() 记录共享专家开始前的事件，combine(previous_event=combine_pre_evt) 让通信内核等待这个事件；limit_num_sms 又让共享专家只使用一部分 SM。最终，combine 的数据搬运与共享专家的计算在时间上交错，等到 current_stream_wait() 时，两者都已经接近完成。

previous_event 的参数在 dispatch 里也有。它的语义是“通信内核发射前，先等待这个事件”。如果上一次计算留下了未完成的依赖，传入它就能把顺序串起来，而不需要调用 torch.cuda.synchronize()。最后一个参数 allocate_on_comm_stream=True 则把通信过程中新分配的张量（例如接收缓冲）归到通信流上；否则 PyTorch 的 caching allocator 会在计算流上做一次记录，可能引入跨流同步。

defer_epilogue=True 的机制稍复杂。combine 的归约尾部（把各 rank 的部分和按 top-k 权重加起来）通常紧跟在数据搬运之后；如果立即执行，通信流要等归约完成才能发完成事件，计算流又得等这个事件。DeepInfer 把归约推迟到 current_stream_wait() 触发时执行，这样通信流可以先完成“搬数据”这一段，计算流可以在等待点附近做归约，重叠窗口更大。代价是归约所需的数据必须在缓冲里多存活一段时间，缓冲的生命周期要覆盖到等待点之后。

CUDA Graph 是这条链路的另一个约束。DeepInfer 在 warmup 阶段捕获计算图与通信图，捕获之后所有内核的指针、参数地址、事件依赖都被固化。任何在捕获之后重新分配或移动通信缓冲的操作，都会让图里的指针失效。因此 userbuffer 的所有分配都必须发生在图捕获之前；SymmBuffer 的 locality_init() 也必须在捕获之前完成。DeepInfer 的启动顺序保证在此之前存在一个明确的分配窗口：worker 的 __init__ 先完成 init_dist，随后模型装载阶段才构造 MoE 层并调用 allocate_userbuffer，显存 profile、KV 初始化与 warmup 排在其后。这个顺序不是风格问题，而是图捕获的硬性前提：allocate_userbuffer 在模型构造期被调用（下文用“构造期”指代这个窗口，而不是仓库里某个未见到的顶层函数名），图捕获则在所有分配完成之后才开始。

初始化顺序：什么时候可以安全地建缓冲

把启动链按时间排开，可以看到缓冲分配卡在两个约束之间：

1
Worker.__init__
2
  ├─ init_model()        设置设备、NUMA、CPU 亲和、随机种子
3
  ├─ init_dist()         gloo 默认组 → 建五张子组 → init_device_groups 哑巴 all_reduce
4
  └─ start_slave_worker_loop()（非 driver）
5
load_model()
6
  ├─ 模型构造期 → MoE 层 allocate_userbuffer → EP/CP/Engram/SymmBuffer 分配
7
  ├─ profile_num_available_blocks()  测峰值显存，算 KV 块数
8
  ├─ init_cache_engine()
9
  └─ warm_up_model()     构图、捕获 CUDA Graph

三个时间点各自解决一个问题：

• init_dist 之前不能建缓冲。 DeepEP 的缓冲构造需要 ProcessGroup 对象，而 ProcessGroup 需要 dist.init_process_group 先完成。更重要的是，EPBuffer、BucketBuffer、EngramBuffer 的构造都会调用 group.barrier() 或 torch.cuda.synchronize()，如果组还没完成初始化，barrier 会直接卡死。
• load_model 之前不能建缓冲。 DeepInfer 的缓冲尺寸由模型参数决定：hidden 维度、专家数、top-k、最大 token 数都来自模型配置。模型层构造时才知道这些值，所以缓冲只能在模型层里分配。
• warm_up_model 之前必须建完缓冲。 CUDA Graph 捕获会把通信内核的地址写进图里；缓冲一旦移动，图就要重新捕获。DeepInfer 选择在模型装载阶段一次性分配，之后不再扩缩容。

EPBuffer 构造函数的末尾有一个容易被忽略的同步：

1
torch.cuda.synchronize()
2
group.barrier()
3
torch.cuda.synchronize()

两个 torch.cuda.synchronize() 夹着一个 group.barrier()。第一个 synchronize 保证本 rank 在默认流上已经完成的内存操作（包括构造过程中为对称存储做的准备）被设备真正执行完，而不是还挂在流的队列里；barrier 保证所有 rank 都到达同一点；第二个 synchronize 保证 barrier 返回后本 rank 的默认流状态干净，后续第一次 dispatch 在同一流上发射时不会与构造阶段的残余工作冲突。

构造函数注释里的那句 “Call a barrier to ensure initialization visibility for all peers” 说明了 barrier 的作用：对等 rank 在本端缓冲构造完成之后才开始使用它。没有这个 barrier，一个较快的 rank 可能已经发出第一次 dispatch，而较慢的 rank 还在构造，接收区尚未准备好。前后的 synchronize 则是为了把“本端已经准备好的状态”从 CPU 侧确认下来，再交给集合操作；少了它们，barrier 只能保证 CPU 到达同一个调用点，不能保证此前发出的 CUDA 工作都已完成。

这个模式在 EngramBuffer 里以另一种形式出现：它需要先 create_shared_handle，再用 dist.all_gather_object 收集所有 rank 的共享句柄，然后才能构造 C++ 运行时。共享句柄的收集是阻塞的集合调用，所以 Engram 缓冲的构造本身就是一个同步点。

确定性与精度：两处容易被忽略的约束

MoE 的 combine 是“按 top-k 权重做归约”，归约顺序会影响浮点结果。DeepEP 给了两个开关：

• allow_multiple_reduction：默认允许在 combine 尾部做多次归约。文档写得很清楚：关闭时“只在 combine 的 epilogue 做一次归约，精度最好，但可能增加数据传输量”；开启时可以把归约拆到多段，减少传输，但累加顺序会变。
• deterministic：构造 EPBuffer 时传入，要求路由算法使用确定性排序。combine 开头会调用 check_torch_deterministic()，在 PyTorch 的确定性模式被破坏时直接报错，而不是静默地给出不确定结果。

DeepInfer 在 allocate_ep 里把这两个开关与 batch-invariant 绑定：

1
allow_multiple_reduction=not (
2
    envs.DEEPINFER_BATCH_INVARIANT and envs.DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL > 0
3
)

也就是说，只有在“开启 batch-invariant 且 EPLB 会定期重排专家”时，才关闭多次归约。原因是 batch-invariant 要求同一个请求在不同 batch 组成下得到完全相同的结果；EPLB 重排会改变专家所在的 rank，如果归约顺序又随 batch 变化，同一个 token 的浮点累加顺序就会变，输出无法逐位一致。关闭多次归约，把累加固定成一次，是以更多传输量换确定性。相关背景在 Batch-invariant 推理完全拆解 一篇里有跨框架的对照。

另一个与精度相关的细节是 expert_alignment 的补零。dispatch 会把每个专家收到的 token 数补齐到对齐边界，补出来的行是零。这些零行会参与 grouped GEMM（乘出来的结果仍是零），但 combine 不会把它们送回源 token，因为 handle 里的 num_unaligned_recv_tokens_per_expert 记录了真实行数。补零只影响计算侧的行数对齐，不影响最终归约的结果。这也是为什么 handle 必须同时保存“对齐后的前缀和”与“未对齐的真实数量”：两者缺一，要么 GEMM 行数不对，要么 combine 多送零行。

一次 decode step 的通信账本

把前面的组件串起来，可以按一次 decode step 的时间推进，给每一次通信点名、编号、记字节数。以 V4.1-Flash 的生产配置为例：40 层、每层都有 MoE，DP Attention 复制注意力，Wide-EP 切分专家，CP 关闭，MTP 开启且投机深度为 $k$。

表中 $B$ 是本 DP rank 上的 token 数，$H$ 是 hidden 维度，$k$ 是每 token 选中的专家数（包含冗余专家带来的物理 top-k）。这里的字节数只是量级估计，具体数值取决于模型配置。

这张账本能解释几个常见误解：

其一，注意力与 dense 部分没有集合通信。 DeepInfer 不做张量并行，注意力层的输出不经过 all-reduce；共享专家与门控也是复制的，不需要跨 rank 同步。这意味着在 decode 阶段，通信被集中在每层 MoE 的两次 dispatch/combine 上。若某层不是 MoE 层（例如部分混合模型的 dense 层），该层通信为零。

其二，dispatch 的字节数与 top-k 成正比，而不是与 hidden 成正比。 一个 token 如果被路由到 $k$ 个专家，且这些专家分布在其它 rank 上，就要搬 $k$ 份 hidden state。DeepInfer 用 do_expand=True 之后，每个 (token, 专家) 组合是一行，所以 dispatch 侧的字节数约为“token 数 × 物理 top-k × $H$ × dtype”。combine 则按专家返回的结果行数计算，通常在专家对齐后略多于 dispatch。

其三，MTP 让通信次数翻倍。 投机解码的验证批需要把草稿 token 与主 token 一起送进 MoE，一次 decode step 实际执行两次（或多次）MoE 前向。DeepInfer 的 combine 在变长验证时会同时处理多个接受长度的行，输出与 KV 回滚的具体机制在此前的 投机解码 一篇里展开。对通信侧来说，这意味着每层要额外付一次 dispatch/combine 的钱。

其四，CP 一旦打开，账本会变。 CP 的 all-gather 每层至少一次，并额外有首尾各一次。DeepInfer 的 CP 主要用于 prefill，此时 all-gather 的数据量按 token 数乘 head 数乘 head 维度乘 CP size 计算，量级与 MoE 通信可比甚至更大。这也是 CPManager.all_gather 选择 copy engine 而不是 SM 内核的直接原因：prefill 阶段计算侧 SM 已经很满。

其五，MegaMoE 把第 2、3、4 三项并成一个内核。 在 use_mega_moe 分支下，MoE 前向不再调用 EPBuffer.dispatch/combine，而是走 deep_gemm.fp8_fp4_mega_moe。这个 kernel 用 SymmBuffer 在 rank 间交换数据，通信与计算在同一内核内交错。从账本的角度看，是“把 2、3、4 三项归并成一项”，而不是“把通信去掉”。通信字节数不变，变化的是这些字节走过的路径：从 DeepEP 的通信 kernel 变成了 DeepGEMM 的融合 kernel。

MegaMoE 的条件与 SM 预算的联动

use_mega_moe 不是运行时按需切换的，而是在 MoE 层初始化时就按 EP 组的规模确定。源码中的判断条件写在 Moe 类构造路径上，含义是“EP 组整体落在同一个物理域内”。DeepEP 提供了两个查询函数：

1
def get_physical_domain_size(group_or_buffer) -> Tuple[int, int]:
2
    """Return physical RDMA and NVLink domain sizes ..."""
3
    context = getattr(group_or_buffer, 'context', None)
4
    if context is not None:
5
        return context.get_physical_domain_size()
6
    return _C.get_physical_domain_size(get_nccl_comm_handle(group_or_buffer).get())

它返回的是“本机 RDMA 域大小”与“本机 NVLink 域大小”两个数。ParallelProcessGroup.is_within_nvl_domain 直接用这两个数判断：

1
@staticmethod
2
def is_within_nvl_domain(group) -> bool:
3
    if PLATFORM.is_npu:
4
        return True
5
    return group.size() == 1 or group.size() <= deep_ep.get_physical_domain_size(group)[1]

group.size() <= NVLink 域大小 成立时，EP 组不跨节点，MegaMoE 可以启用。否则走 EPBuffer 的标准两段式。

SymmBuffer 的“首个规格必须最大”约束也与这一路径相关。在 allocate_sym_buffer 里，第一次分配时创建实际的对称存储，之后相同或更小的规格都只是它的视图：

1
base = next(iter(self._sym_buffers.values()), None)  # get the base buffer
2
if base is None and PLATFORM.is_gpu:
3
    locality_init = getattr(getattr(deep_gemm, "locality_domain", None),
4
                            "get_balanced_sm_locality_domains", None)
5
    if locality_init is not None:
6
        locality_init()
7
buffer = deep_gemm.SymmBuffer(**spec._asdict(), base=base)

locality_init() 在第一次构造前调用，是为了让 DeepGEMM 提前探测并缓存 SM 局部性信息。这一步本身会触发同步，所以被放在受保护的 forward 和 CUDA Graph 捕获之前——如果放在图捕获期间执行，会破坏图中的异步语义。

为什么“首个必须最大”不是性能提示而是正确性前提？把 base 的逻辑写开就很清楚：deep_gemm.SymmBuffer(**spec._asdict(), base=base) 在 base 非空时不是新建一块存储，而是在 base 的底层张量上按新规格创建一组视图。新规格所需的字节数如果超过 base 的尺寸，narrow 或 view 会直接越界。DeepInfer 的调用顺序保证了这一点：MegaMoE 的 allocate_userbuffer 用“最大 token 数、最大 physical top-k、最大专家数”的规格来分配第一个 SymmBuffer；之后若有别的层或功能需要更小的规格（例如只含共享专家、或激活不同），才会命中 base 视图。如果调用方在构造早期用一个小规格先分配，后续大规格就会静默地越界写入另一块显存——这是 CUDA 程序里最难排查的崩溃之一。

这也解释了 _sym_buffers 为什么用完整的 _SymBufferSpec 作为键，而不是只按 token 数或专家数作键。内核需要的布局参数（hidden、intermediate_hidden、num_shared_experts、mma_type、activation）任意一项不同，视图的切法就不同；用完整规格作键，才能保证“同规格共享、异规格新建视图”的语义精确。

SM 预算的联动也值得说明。MegaMoE 路径下，通信不再是独立 kernel，而是与 GEMM 交错，DEEPINFER_USERBUFFER_COMM_SMS 不再用于 dispatch/combine。但 limit_num_sms 仍然在共享专家计算时生效：

1
with limit_num_sms(envs.DEEPINFER_USERBUFFER_COMM_SMS, enabled=PLATFORM.is_gpu):
2
    shared_out = self.shared_forward(hidden_states)

limit_num_sms 的实现是把 DeepGEMM 与 TileKernels 的 SM 数临时减掉一个值，再用 finally 恢复：

1
prev_sms = deep_gemm.get_num_sms()
2
tile_kernels.set_num_sms(prev_sms - num_sms)
3
deep_gemm.set_num_sms(prev_sms - num_sms)
4
try:
5
    yield
6
finally:
7
    tile_kernels.set_num_sms(prev_sms)
8
    deep_gemm.set_num_sms(prev_sms)

也就是说，共享专家的 GEMM 在运行时用的是“总 SM 数减去通信预留数”的结果。这样在 combine 的内核已经发射、需要在通信流上排队的同时，共享专家的 GEMM 不会把全部 SM 抢光，从而给 combine 留下足够的 SM。这个“让路”的粒度是 kernel 级别，而不是流级别。

与其他框架的对照

把 DeepInfer 的做法放到其他推理框架里对照，会更容易看出哪些是通用选择、哪些是它自己的取舍。

vLLM 在 TP 路径上的集合通信完全交给 NCCL：每个 RowParallelLinear 之后都有一次张量并行的 all-reduce，TP 组的建立与销毁由 PyTorch 负责，vLLM 本身不维护额外的通信缓冲。MoE 的专家全互联则通过 DeepEP 的 EPBuffer 完成，与 DeepInfer 在 MoE 侧的思路一致。不同点在于 vLLM 同时承担 TP 与 EP 两套通信，DeepInfer 只有 EP，并且把 CP、Engram 也收进同一个 DeepEP 体系里。TP 的完整实现可以对照站内 vLLM 源码完全拆解（十一）：模型并行与 Worker。

SGLang 与 DeepInfer 都采用 DP Attention + DeepEP dispatch/combine 的组合，差异在于 SGLang 的 CP 通信通常走 NCCL 的 all_gather 与 all_gather_into_tensor，而不是 DeepEP 的 BucketBuffer。DeepInfer 把 CP 放在用户态缓冲里，换取“按会话复用一次存储”的能力——一次 CP 会话里可能有多次 all-gather，会话结束后缓冲整体复用给下一个会话。这在 userbuffer 四类缓冲 一篇的 cp_session 一节里有专门讨论；SGLang 侧的对应实现在 SGLang 源码完全拆解（十七）：DP Attention 与专家并行。

TensorRT-LLM 的 Communicator 与 DeviceMesh 是另一种组织方式：它显式维护 rank 坐标系，能在运行时按需求构造任意组合的通信域，并支持自定义集合通信原语（如跨 host 的 all-reduce、基于 MNNVL 的对称内存）。DeepInfer 没有这么通用，但它把每种通信和具体的数据面对应写死，换来的是更少的抽象层和更容易推理的显存账。相关实现在 TensorRT-LLM 源码完全拆解：并行拓扑与通信 一文中有完整展开。

llama.cpp 的 RPC 后端与多卡 MetaDevice 走的是另一条路线：用 RPC 把一张虚拟设备分片到多机，AllReduce 视后端不同可以走 NCCL 或 RPC 自带实现。它没有 DeepEP 这种专门的 MoE 通信数据面，适合规模较小的部署。这一点在 llama.cpp 源码完全拆解：Meta device 与张量并行 与 RPC 后端 两篇里有讨论。

共同的结论是：大 MoE 推理的通信瓶颈在专家全互联，不在注意力规约。 四个框架都在 MoE 侧引入 DeepEP 或类似的自研数据面，而在注意力侧要么不做 TP（DeepInfer、SGLang 的 DP Attention），要么把 TP 规约保留为 NCCL 的常规用法。

读这部分代码容易踩的坑

1. PG.ep 与 PG.ep_in_dp 通常是同一个 ProcessGroup 对象，但不要依赖这一点。 它们按同一种 ep_size 连续切分，_split_group_cached 的 lru_cache 会让它们命中同一份缓存，所以属性值相等。但缓存键是 split_ranks 元组；如果以后某次改动让两者的切分方式不同，它们就会变成两个独立的组，各自拿到独立的 NCCL communicator。DeepInfer 的 UserBufferManager.allocate_ep 传入的是 self.ep_group，实际来源是 ParallelProcessGroup.ep；写新代码时按“名字”取组，不要跨属性混用。
2. get_physical_domain_size 与 get_logical_domain_size 不是一回事。 前者按硬件拓扑（NVLink 域、RDMA 域）返回，后者按配置与 hybrid 模式返回“scaleout/scaleup domain”。is_within_nvl_domain 用的是前者，BucketBuffer 的 SM 推导用的是后者拆出来的物理域。
3. all_gather 的 num_sms 必须是 0。 如果对 BucketBuffer 的 all_gather 传了一个非零的 SM 数，会直接触发断言失败，因为 all-gather 是 copy engine 驱动，不占 SM。
4. allocate_ep 的幂等条件基于字节数而非完全相等。 同一次加载里若两次请求的规格不同，第二次会断言 self.ep.num_bytes >= hint 并复用旧缓冲。如果第二次规格更大，断言会失败——这通常意味着调用方在“最大规格”之前就分配了。allocate_sym_buffer 有同样的问题，且更隐蔽：较大的规格会复用 base 缓冲的视图，视图越界。
5. 销毁顺序：先缓冲，后 process group。 任何绕过 userbuffer.close() 直接 dist.destroy_process_group() 的路径，都会让 NCCL communicator 在缓冲析构时已被释放，产生难以排查的崩溃。

小结

DeepInfer 的通信层可以概括成三句话：

• 组只解决“谁和谁”。五张进程组分别绑定了 DP 存活对齐、CP 上下文并行、EP 专家并行、ENP 远端查表与 EP-in-DP 控制面。gloo 副本只在需要 CPU 侧协调的地方（ZOS 的存活对齐、CP 的对象广播）额外开一份。
• 通道由 DeepEP 与 DeepGEMM 提供，而不是 NCCL 的集合函数。EP 全互联用 EPBuffer，通用集合用 BucketBuffer，远端查表用 EngramBuffer，MegaMoE 用 SymmBuffer。四者共享组、共享部分 NCCL handle，但内核是各自的。
• 账本靠“每次前向都记账”维持。dispatch 与 combine 按 token 与路由计字节；all-gather 按 token × head × CP size 计；Engram 按被查表的行数计。MegaMoE 只是把 MoE 的三项归并成一项，通信的字节量并不会消失。

理解了这一层，再回头看 userbuffer 四类缓冲、MoE 的 MLP 与 EPLB 三篇里散落的细节，它们其实是同一张账本在不同视角下的记录：谁占了多少缓冲、谁在什么时机发起通信、谁在哪个进程里完成回收。

参考资料

1. DeepInfer GitHub 仓库
2. DeepEP GitHub 仓库
3. NCCL Collective Operations 文档
4. PyTorch Distributed 文档
5. DeepInfer parallel_state.py
6. DeepInfer userbuffer.py
7. DeepEP buffers/ep.py
8. DeepEP buffers/bucket.py