DeepInfer 源码拆解（二十八）：userbuffer，四类通信缓冲与对称内存

NCCL 的抽象：集合通信 + 内部缓冲

torch.distributed 与 NCCL 的组合，面向的是“语义清晰的集合操作”：all_reduce、all_gather、reduce_scatter、all_to_all。你给它一组张量，它内部决定切块、分阶段、走哪条路径（Ring、Tree、NVLS），并自己管理收发缓冲。这套抽象极大降低了写分布式代码的门槛，但它有几个隐含假设：

• 通信是独立于计算的一类操作。 NCCL 内核需要 SM 才能跑，但它不知道你的 GEMM 正在用多少 SM，也无法主动让自己的内核“只占 16 个 SM”去和计算重叠；
• 缓冲区是库的内部细节。 你看不到、也改不了收发缓冲的布局，因此无法让“对端写入的位置”与“本地 GEMM 的输入指针”在物理上重合；
• 一次调用一次同步。 多个张量的 all_gather 需要多次调用，或者自己拼成一个大张量再调用。

DeepInfer 面对的场景把这三点都推到了极限。

DeepInfer 的通信形态

回到 DeepInfer 实际的工作负载，它的通信集中在四个地方：

这四类需求的共同点是：它们的性能上限由“SM 与带宽如何共享、收发缓冲的物理位置、是否能零拷贝”决定，而不是由“集合操作的语义”决定。

第一类（MoE）要的是“通信只占一部分 SM、把缓冲切成小 buckets 与计算流水”。第二类（CP）要的是“多个小张量一次会话完成，输出直接是注册缓冲里的视图”。第三类（Engram）要的是“远端存储注册成可被 RDMA 直接读的区域”。第四类（MegaMoE）甚至没有独立的通信内核——dispatch、GEMM、combine 是同一个 kernel，rank 之间通过对称内存指针读写。

NCCL 的通用集合通信抽象覆盖不了这四种形态，这就是 userbuffer 出现的原因。它借用了 DeepEP 与 DeepGEMM 提供的底层原语（EPBuffer、BucketBuffer、EngramBuffer、SymmBuffer），把“尺寸推导、分配时机、生命周期”这三件跨框架通用的事情收拢到一处。

（图源：DeepEP 官方仓库 figures/normal.png。DeepInfer 的 EPBuffer 正是这幅图里两类内核的统一入口——同一个缓冲对象，可以按调用参数切换高吞吐与低延迟形态。）

分配发生在什么时候

分配时机在这里至关重要。deepseek.py 里相关代码的先后顺序是：

1
# Allocate userbuffer.
2
torch.cuda.empty_cache()
3
self.allocate_userbuffer()

torch.cuda.empty_cache() 先执行，把 CUDA caching allocator 里所有空闲块还给驱动，再开始分配。这不是仪式感——EPBuffer 与 EngramBuffer 内部的收发区是普通 CUDA 分配，而 SymmBuffer 依赖 DeepGEMM 的对称内存机制；如果先分配再 empty_cache()，这些区域不会被回收（它们是活跃分配），但其他临时张量的空洞会占住地址空间，导致对称内存对齐或总显存估算出现偏差。

在 allocate_userbuffer() 内部，顺序也是固定的（下面是精简后的关键结构）：

1
def allocate_userbuffer(self) -> None:
2
    max_ag_bytes_per_token = self.hidden_size * 2  # hidden_state in sample
3

4
    if envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact" and self.model.yoco_boundary_layer_idx is not None:
5
        # yoco compact 开 mirrorcp 需要传 HC residual（mask 模式不做边界 gather，无此需求）
6
        max_ag_bytes_per_token = max(max_ag_bytes_per_token, getattr(self.cfg, "hc_mult", 1) * self.hidden_size * 2)
7

8
    decoder_layers = list(self.model.layers)
9
    if self.num_mtp_moe_layers > 0:
10
        decoder_layers += self.MTP_module.moe_layer_list()
11
    for decoder_layer in decoder_layers:
12
        if isinstance(decoder_layer.mlp, DeepseekMoE):
13
            decoder_layer.mlp.allocate_userbuffer()
14
        if isinstance(decoder_layer.self_attn, V4AttentionBase):
15
            max_ag_bytes_per_token = max(max_ag_bytes_per_token, decoder_layer.self_attn.get_max_ag_bytes_per_token())
16

17
    userbuffer.allocate_cp(
18
        PG.context_parallel_group(),
19
        max_num_tokens=self.cfg.max_num_prefilling_tokens,
20
        max_ag_bytes_per_token=max_ag_bytes_per_token,
21
    )
22

23
    # Engram
24
    if self.model.engram_runtime is not None:
25
        self.model.engram_runtime.initialize()

三件事按顺序发生：

1. MoE 层逐层分配 EP 或 SymmBuffer。 DeepseekMoE.allocate_userbuffer 根据 use_mega_moe 二选一：MegaMoE 走 allocate_sym_buffer，否则走 allocate_ep。MTP 层的 MoE 也被拉进来——投机解码的草稿头如果在同一进程里跑，它的 MoE 前向同样需要 dispatcher，这一行避免了一个很隐蔽的问题：MTP 层的 allocate_userbuffer 被漏掉，前向时才报“EP buffer 未初始化”。
2. 注意力层聚合 CP 的 all-gather 字节数。 每层调用 get_max_ag_bytes_per_token()，取最大值，并且以一个“输出侧 hidden 必须 gather”的基线 hidden_size * 2（BF16 的两个字节）起步。为什么是最大值而不是总和？因为 CP 的 all-gather 是逐层发生的，同一时刻只有一层的张量在缓冲里，所以缓冲按“单层最大需求 × token 数 × CP 大小”估算即可。
3. CP 缓冲一次性分配，Engram 随后初始化。 Engram 的 initialize() 内部会调用到 userbuffer.allocate_engram：GPU 表建立零拷贝视图后交给 EngramBuffer，这里的 allocate_engram 正是完成这一步的地方。

这个基线值本身也值得一说。CP 的最后一步是注意力输出之后的 hidden_states all-gather（下文的 deepseek.py 第 242 行），它要 gather 的是每 token hidden_size × 2 字节的 BF16 隐状态；而各注意力层内部 gather 的 K/V 分片可能更大也可能更小，取决于 qk_head_dim 与 o_groups 的配置。基线保证了即使所有层内部的聚合字节都不大，输出侧的 gather 也不会因为缓冲不够而越界；各层的 get_max_ag_bytes_per_token() 再在这个基线上取最大。

（本站绘制，依据 userbuffer.py 与 deepseek.py 的调用顺序。图中“模型层调用点”对应 DeepseekMoE.allocate_userbuffer() 与 V4AttentionBase.get_max_ag_bytes_per_token()；下方“生命周期”一栏的第六步 close() 在 worker/base.py 的 stop() 里被调用，不在模型内部。）

这张图把四类缓冲并列摆在一起，容易让人误以为它们是“一个东西的四种模式”。实际上它们的性质差别很大，下一节逐个拆。

显存预算：C 侧算出来的一个数

EPBuffer 的构造并不在 Python 里逐字段累加，而是把参数交给 C++：

1
# Calculate buffer size (already 2 MB-aligned from hint functions / calculate_ep_buffer_size)
2
if num_bytes is None:
3
    num_bytes = _C.calculate_ep_buffer_size(
4
        self.nccl_comm_handle.get(),
5
        num_max_tokens_per_rank, hidden, num_topk, use_fp8_dispatch,
6
        allow_hybrid_mode, allow_multiple_reduction)

calculate_ep_buffer_size 在 DeepEP 的 C++ 侧实现，返回值已经按 2 MB 对齐（get_buffer_size_hint 复用同一个函数，所以复用断言比较的也是同一口径）。它的输入里没有“层数”“序列长度”“并发请求数”，只有每 rank 的最大 token 数、隐维、topk、是否 FP8，再加上模式开关。这说明一件事：EPBuffer 是 scratch，不是按“模型总吞吐”预留的容量池。 它只需要容纳一次 dispatch/combine 的最大瞬时数据量，具体多少层、多少请求都共用这一份。

num_max_tokens_per_rank 从哪来？mlp.py 里：

1
self.max_num_tokens = self.max_num_decoding_tokens if self.engine_type == EngineType.COMPLETION else self.max_num_prefilling_tokens

而 model_loader.py 里这两个值的推导是：

1
model_cfg.gpt.max_num_decoding_tokens = model_cfg.gpt.max_num_seqs * (1 + envs.DEEPINFER_USE_NEXTN_DECODE)
2
model_cfg.gpt.max_num_prefilling_tokens = model_cfg.gpt.max_num_batched_tokens // PG.context_parallel_group().size()

把它们连起来看，含义很清晰：

• Decode 引擎（COMPLETION）按 max_num_seqs × (1 + 投机深度) 分配。 多出来的 1 + N 因子是给投机解码草稿 token 留的位置——每个请求一步可能带 1 个目标 token 加 N 个草稿 token 进入 MoE dispatch，如果缓冲只按 max_num_seqs 算，投机路径会把 buffer 写穿。这一条与投机解码一篇直接对应。
• Prefill 引擎按 max_num_batched_tokens // cp_size 分配。 因为 prefill 的 token 会先被 CP 切成 cp_size 份，每个 rank 只看自己那一段，所以单 rank 的 MoE token 上界要除以 CP 大小。这个除法如果漏掉，CP 开启时会多分配 cp_size 倍显存。
• 一个进程只跑一种引擎类型。 DeepInfer 的 Prefill/Decode 是分离部署的，所以 engine_type 在进程启动时就固定，缓冲尺寸也只按一种形态算。这与“同一个进程里 prefill 和 decode 混跑”的框架不同，后者必须按两者的最大值分配。

use_fp8_dispatch 则直接来自 MLP 的 fp8_mlp 配置：

1
self.use_fp8_act = cfg.get("fp8_mlp", False)

开启后，dispatch 发送的是 FP8 张量加 scale factor：接收端拿到的是 (FP8 tensor, scale factors) 这个二元组，数据字节数大约减半，但缓冲里要额外留 scale 的空间。这就是为什么 C 侧函数把它作为独立开关，而不是简单按 hidden 乘一个系数。

SM 预算：通信让出多少计算资源

EPBuffer 的 SM 数不是显式参数，而是 get_theoretical_num_sms(num_experts, num_topk) 算出来的。DeepInfer 又在这个结果上做了一次“保留计算 SM”的裁剪：

1
device_num_sms = torch.cuda.get_device_properties().multi_processor_count
2
compute_num_sms = min(32, device_num_sms // 4)
3
envs.DEEPINFER_USERBUFFER_COMM_SMS = min(num_sms, device_num_sms - compute_num_sms)

以 H800 的 132 个 SM 为例：compute_num_sms = min(32, 33) = 32，所以通信最多用 100 个 SM。为什么是“32 或四分之一，取小”？这是两个约束的交集：

• 四分之一：给计算留足至少 25% 的 SM，保证 decode 的 GEMM 不会因为通信抢占而排队；
• 32：在 SM 数较少的卡（如 80 SM、96 SM 的卡）上，四分之一也会缩小到二十来个 SM，这时用 32 作上限只是给“大卡”封顶，保证计算侧不会被压得过低。它是一个保守的常数，而不是根据 GEMM 形状推出来的理论值。

这个结果写进环境变量 DEEPINFER_USERBUFFER_COMM_SMS，后续每次 dispatch/combine 都显式把它作为 num_sms 传入：

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

那行 TODO 值得注意：dispatch 与 combine 的带宽模型并不相同，但当前 DeepEP 只接受一个 num_sms，所以 DeepInfer 用一个保守值同时服务两者。 dispatch 是“读本地 token、写远端缓冲”，combine 是“读远端结果、累加回本地”，两者对 SM 与 QP 的需求不同。统一值意味着其中一个阶段一定不是最优——这是工程上的已知折衷，而不是设计目标。

一个容易忽略的细节：get_theoretical_num_sms() 在 DeepEP 内部有缓存，get_theoretical_num_qps() 则是一个纯函数式的推导。更重要的是，SM 数只在 DEEPINFER_USERBUFFER_COMM_SMS == 0（即用户没有手动指定）且第一次分配时被推导一次，之后所有 MoE 层、包括后续可能的第二次 allocate_ep 调用，都复用这个环境变量。这符合 DeepEP 的约束：一个 EPBuffer 里的通信内核在构造时就绑定了 SM 与 QP 数，运行时每个 dispatch 调用的 num_sms 只是在“预分配的网络上下文”里选择用多少个，而不能超过构造时准备的 QP 上限。

理论 SM 数是怎么推出来的

get_theoretical_num_sms 的推导不算短，但逻辑是清晰的。它要回答：给定专家数、topk、带宽，通信内核最少需要多少 SM 才能不成为瓶颈？

第一步是估算“一个 token 的 topk 选择平均会落到几个不同的 rank 上”。对均衡 gate，公式是：

其中 $E$ 是专家总数，$k$ 是 topk，$G$ 是分组数（在 EP 语境里就是 rank 数，每个 rank 负责 $E/G$ 个专家）。这个式子的含义是：先算“一个 token 的 $k$ 个选择全部避开某个固定组”的概率，再用 1 减去它，得到“至少命中该组一个专家”的概率；乘以组数，就是命中的不同组数的期望。对 DeepSeek 的 group-limited gate，这个公式只是近似，DeepEP 的注释也明确写了“请勿用于 V3.0 的 group-limited gate”，但作为一阶估算足够。

第二步是把流量拆成四类，全部按“每 token、每 hidden 字节”归一化：

• sm_read：dispatch 阶段从 HBM 读 token 的量；
• sm_write：dispatch 阶段写 send buffer 与发出 NVLink 的量；
• nvlink_traffic：走 scale-up 域（NVLink）的字节量；
• rdma_traffic：走 scale-out 域（RDMA）的字节量。

对单节点的纯 NVLink EP（num_scaleout_ranks == 1），代码里的累加是：

1
# Write send buffer
2
if self.num_rdma_ranks > 1:
3
    sm_write += 1 / num_expected_topk
4

5
# Issue NVLink
6
sm_write += self.num_nvlink_ranks / self.num_ranks
7

8
# NVLink and RDMA traffic
9
nvlink_traffic += self.num_nvlink_ranks / self.num_ranks * (1 - 1 / self.num_nvlink_ranks)  # Except local bypass
10
rdma_traffic += (self.num_ranks - self.num_nvlink_ranks) / self.num_ranks

1 / num_expected_topk 这一项对应“每个 token 平均被复制到几个 rank”，即写放大；(1 - 1/R_nvlink) 对应“除本地 rank 外，其余 NVLink rank 都要收一份”，即扇出。对跨节点 EP，还会有 num_expected_scaleout_topk 与转发路径的额外项。

第三步是找瓶颈域：比较 rdma_traffic / rdma_gbs 与 nvlink_traffic / nvlink_gbs，取较大的那个作为 bounded_traffic 与 bounded_gbs。这相当于“哪条链路的归一化时间更长，就以它为瓶颈”。

第四步是算 SM：

1
num_sms = max(
2
    # Dispatch
3
    bounded_gbs / bounded_traffic * sm_read / sm_read_gbs,
4
    bounded_gbs / bounded_traffic * sm_write / sm_write_gbs,
5
    # Combine
6
    bounded_gbs / bounded_traffic * (sm_write - nvlink_traffic) / sm_read_gbs,
7
    bounded_gbs / bounded_traffic * (sm_read + nvlink_traffic) / sm_write_gbs,
8
)
9
num_sms = align(max(4, math.ceil(num_sms * 1.3)), 4)
10
num_sms = num_sms if self.prefer_overlap_with_compute else max(num_sms, 64)
11
num_sms = min(num_sms, num_device_sms // 2 * 2)

四个比值分别对应 dispatch 的读/写与 combine 的读/写；乘以 1.3 是工程裕量，向上取整到 4 的倍数是为了让 SM 分配能均分到偶数个 channel；max(num_sms, 64) 只在“不需要与计算重叠”时生效，因为独占通信时多给 SM 能更快跑完；最后把上限压到设备 SM 数的一半，避免通信把整卡吃满。

DeepInfer 调用时没有传 prefer_overlap_with_compute，所以走 EPBuffer 的默认值 True，不会触发 64 的下限；随后又被 device_num_sms - compute_num_sms 二次裁剪。两层裁剪叠加后，DeepInfer 实际使用的是“理论值、设备 SM 减保留计算 SM”两者的最小值（理论值本身已经不超过设备 SM 的一半）。 这就是为什么在 132 SM 的卡上，通信通常只拿到 64–100 个 SM，而不是理论公式给出的上限。

（图源：DeepEP 官方仓库 figures/low-latency.png。DeepInfer 的 EPBuffer 在 decode 阶段更接近这幅图：低延迟模式把收发槽位固定下来，用 RDMA 直接写对端缓冲；而 prefill 的大批量 token 更适合 Normal 模式的转发路径。）

QP 预算：一条 channel 一条队列

RDMA 的 QP（Queue Pair）数决定能同时有多少个未完成的 RDMA 操作。EPBuffer 在构造时先定一个上限：

1
if num_allocated_qps == 0:
2
    # Hybrid mode will consume more QPs
3
    # The extra QP is for notify warps
4
    if self.allow_hybrid_mode:
5
        num_allocated_qps = 65 if check_fast_rdma_atomic_support() else 129
6
    else:
7
        num_allocated_qps = 17

为什么是 65 和 129？混合模式下，每个 channel 需要一条数据 QP，另外还要一条 notify QP 用来发完成通知；如果网卡支持快速 RDMA atomic，可以用一条 QP 同时承担数据与原子操作，于是 64 个 channel + 1 条 notify = 65；不支持时要把数据与 atomic 拆开，变成 128 + 1 = 129。非混合模式只有 16 个 channel，所以是 16 + 1 = 17。

实际使用的 QP 数又由 SM 数推导：

1
# For direct mode, we encourage less QPs to reduce DB ringing overhead
2
num_qps = min(num_sms, 8 + 1)
3

4
# For hybrid mode, we encourage every channel (and notify) to have an independent QP
5
if self.allow_hybrid_mode:
6
    num_qps = num_sms * 16 + 1
7

8
return min(num_qps, self.num_allocated_qps)

混合模式下每个 SM 展开 16 条 channel，所以 num_sms * 16 + 1；再被 num_allocated_qps 截断。这个设计把“SM 数”和“QP 数”绑在一起：多给 SM 不只是多给计算单元，还会多开 QP，从而增加网卡队列压力。 反过来说，如果网卡 QP 上限较低，SM 数也不能无限加，num_allocated_qps 就是那道天花板。dispatch/combine 调用时若传入的 num_qps 超过 num_allocated_qps，DeepEP 会直接 assert num_qps <= self.num_allocated_qps 失败——所以这不是软限制，而是硬约束。

一个容易忽略的互斥：batch-invariant 与 EPLB

allow_multiple_reduction 的取值不是常量，而是一个条件表达式：

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

这三个符号分别是什么意思？

• DEEPINFER_BATCH_INVARIANT：要求同一请求在不同 batch 组成下产生逐位相同的结果。batch-invariant 那一篇讲过，DeepInfer 为此固定了归约顺序、禁用了会改变累加顺序的优化。
• DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL > 0：EPLB 在线重排专家，每隔若干步就会改变“哪个 rank 持有哪个专家”。
• allow_multiple_reduction：DeepEP 的 combine 是否允许分多次归约。允许时可以利用更多并行度，但归约顺序会随 rank 拓扑变化；关闭时只在 combine 的 epilogue 里做一次归约，精度最好，代价是传输量可能变大。

三者同时成立时，如果允许多次归约，EPLB 重排后的归约顺序可能与重排前不同，同一请求的输出就不再逐位一致。所以 DeepInfer 把 allow_multiple_reduction 关掉：宁可 combine 慢一点，也要保证 batch-invariant 的承诺。这是一条典型的“正确性优先于吞吐”的约束，也解释了为什么 userbuffer.py 里会出现一个看似与缓冲无关的布尔表达式。

BucketBuffer：一块连续对称存储 + 会话语义

CPManager 本身很薄，真正的实现在 DeepEP 的 BucketBuffer 里：

1
class CPManager:
2
    """Manage CP communication storage, sessions, and all-gather operations.
3

4
    Wait callbacks must be consumed within the session. Gathered tensors alias
5
    the buffer storage, which is reused by the next session.
6
    """
7

8
    def __init__(self, group: dist.ProcessGroup, num_bytes: int) -> None:
9
        self._buffer = BucketBuffer(group, num_bytes, explicitly_destroy=True)
10
        self._rank = group.rank()
11
        self._size = group.size()
12
        self._session: BucketSession | None = None

那段 docstring 是整段代码里信息密度最高的地方，它说了两件必须同时成立的约束：

1. 等待回调必须在会话内消费。 all_gather 返回的是一个 handle，handle.wait() 之后才拿到结果；但会话退出时缓冲的“借用”关系会被撤销。如果会话外还握着 handle 再调用 wait()，拿到的可能是下一次会话已经写过的数据。
2. 聚合结果是缓冲的别名。 wait() 返回的张量没有复制，直接是把注册存储里的那一段 view 出来。这意味着使用方必须在会话内把结果读完，或者自己复制出去；否则下一层 CP 通信会把这块区域覆盖。

BucketBuffer 的设计目标由此清晰：它是一块按“会话”复用的共享存储，而不是按张量分配的内存池。 传统做法是每次 all-gather 都 torch.empty 一块接收区，NCCL 从发送缓冲拷进去；BucketBuffer 则要求发送方与接收方都指向同一块注册存储的不同切片，通信内核直接写这块存储，省掉一次复制。

尺寸怎么算

allocate_cp 的逻辑用十几行代码就能读完：

1
def allocate_cp(
2
    self,
3
    cp_group: dist.ProcessGroup,
4
    max_num_tokens: int,
5
    max_ag_bytes_per_token: int,
6
    max_num_ags_per_session: int = 8,
7
) -> CPManager | None:
8
    if cp_group.size() <= 1:
9
        assert self.cp is None, "cp buffer was already allocated for a multi-rank CP group"
10
        return None
11
    max_ag_bytes_per_session = max_num_tokens * max_ag_bytes_per_token * cp_group.size()
12
    max_ag_bytes_per_session += max_num_ags_per_session * deep_ep.get_num_rdma_alignment()
13
    alignment = deep_ep.get_num_allocation_alignment()
14
    num_bytes = ((max_ag_bytes_per_session + alignment - 1) // alignment) * alignment
15
    if self.cp is not None:
16
        assert self.cp.num_bytes >= num_bytes
17
        return self.cp
18

19
    self.cp = CPManager(cp_group, num_bytes)
20
    return self.cp

先看尺寸公式：

各符号的含义：

• $T_{\max}$：单 rank 的最大 token 数，传的是 max_num_prefilling_tokens（注意不是 CP 切分前的总 token 数）；
• $B_{\text{token}}$：每个 token 参与 CP all-gather 的字节数，由各层注意力 get_max_ag_bytes_per_token() 与输出侧基线 hidden_size * 2 的最大值给出；
• $N_{\text{cp}}$：CP 组大小，因为聚合后在每个 rank 上都要放完整的 $T_{\max} \times N_{\text{cp}}$ 个 token；
• $N_{\text{ags}}$：一个会话里最多有多少次分配的额外对齐，默认 8；
• $A_{\text{rdma}}$：一次 RDMA 对齐单位（get_num_rdma_alignment()）。

为什么要乘以 CP 大小？ 因为 all_gather 的输出是“所有 rank 的本地分片拼接”，每个 rank 最终持有 $T_{\max} \times N_{\text{cp}}$ 个 token 的数据。如果只按本地 token 数分配，输出区根本放不下。

为什么要加对齐字节？ BucketSession._allocate 每次分配都会把字节数向上对齐到 get_num_rdma_alignment()；一个会话里最多 8 次分配，就会有最多 8 次对齐浪费。预留 8 个对齐单位，是为了保证会话内 8 次分配不会因为对齐向上取整而超出总大小。这个细节在代码里没有注释，只能从 _allocate 的实现反推：

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

注意 offset + allocated_bytes > self.buffer.num_bytes 这个判断用的是对齐后的 allocated_bytes，所以真正越界的是对齐后的大小。allocate_cp 里预留的 N_ags × alignment 恰好覆盖了这段对齐增量。

最后一次向上取整到 get_num_allocation_alignment()：这是 DeepEP 对缓冲存储起始地址与总大小的要求，因为 RDMA 注册与 BucketBuffer 的对称切分都按这个粒度为前提。

max_ag_bytes_per_token 从哪来

这个值由两部分取最大：输出侧基线，与各注意力层的估计。输出侧基线在 allocate_userbuffer 里是 self.hidden_size * 2，对应 CP 最后一步把完整隐状态 gather 回来的长度；注意力层的估计则来自 V4AttentionBase.get_max_ag_bytes_per_token()：

1
def get_max_ag_bytes_per_token(self) -> int:
2
    return self.qk_head_dim * torch.bfloat16.itemsize

qk_head_dim 是 MLA 压缩后的 KV 隐维（DeepSeek 系模型里是 576，即 512 的 latent 加 64 维 RoPE），乘 BF16 的 2 字节得到每 token 的 KV 分片字节数。这一层要 gather 的是“投影后、RoPE 前”的 KV——注意 _prepare_swa_kv 里 fused_norm_rope_hadamard_cast_store_kvcache 的 out_args 直接指向 userbuffer.cp.get_inplace_tensor(...) 返回的注册切片，也就是说投影、归一化、RoPE、量化存储这一整条融合 kernel 的“输出”并不是在写 KV Cache，而是在把 RoPE 前的结果写进 CP 的注册存储，等 all-gather 之后再统一落盘。

get_max_ag_bytes_per_token() 在整个 V4AttentionBase 里只返回这一个常数，说明 CP 的 all-gather 在注意力层里只针对主 KV 分片这一种数据；不同层取最大值得到的其实是同一个数。但缓冲的尺寸仍然按“最大”而不是“总和”算——因为同一时刻只有一层在会话里。

这也解释了一个容易被误解的行为：如果 cp_group.size() <= 1，allocate_cp 返回 None，整个 CP 缓冲不分配，cp_session() 直接空转。 这行断言 assert self.cp is None 不是多余的——它防止“先按多卡跑了一次、再改成单卡”的进程里残留旧值；由于 userbuffer 是单例，跨模型实例复用时这类残留是真实存在的。

会话如何把“未注册张量”变成“注册切片”

BucketSession 最巧妙的部分，是它在会话期间临时替换 Buffer 的集合通信方法：

1
def __enter__(self) -> 'BucketSession':
2
    # NOTES: Users must ensure themselves all the allocated bytes in the buffer can be
3
    # safely reused before entering this session.
4
    assert not any(name in self.buffer.__dict__ for name in self.collective_methods), \
5
        'Nested sessions are not supported'
6
    self.num_allocated_bytes = 0
7
    self.context_idx = None
8
    for name in self.collective_methods:
9
        method = partial(self._collective_func, name)
10
        setattr(self.buffer, name, method)
11
    return self

进入会话时，all_gather/reduce_scatter/all_reduce 三个方法被替换成 _collective_func 的偏函数；退出时再删除这些实例属性，恢复成类上的原始方法。于是：

• 会话外调用 buffer.all_gather(srcs) 走原生路径，要求 srcs 必须属于这块缓冲（已经是注册切片）；
• 会话内调用同样的方法，则先检查 srcs 是否已注册；如果没注册，就从会话游标处分配一块注册切片、torch._foreach_copy_ 把数据拷进去，再调用原生方法。

这就是 _collective_func 里那段 is_src_registered 判断的意义：

1
is_src_registered = storage_begin <= src.data_ptr() < storage_end
2
...
3
if not is_src_registered and not accept_unregistered_src:
4
    copy_tasks.append((src_buffer, src))
5
...
6
if copy_tasks:
7
    torch._foreach_copy_(*zip(*copy_tasks))

torch._foreach_copy_ 是批量拷贝，把多个小张量一次提交，减少 Python 侧循环与 kernel 启动开销。accept_unregistered_src 的取值是 name == 'all_gather' and num_rdma_ranks == 1，即只在纯 NVLink、且只做 all-gather 时为真——此时 DeepEP 可以直接从任意 CUDA 张量读，不需要先拷进注册区；其余情况（RDMA、reduce-scatter、all-reduce）都必须先落到注册存储里。

_wrap_handle 则处理“结果落在哪里”的问题：如果调用方给了 dsts，而 dsts 不在缓冲里，会话会把聚合结果先写进临时注册切片，等 wait() 后再批量拷回 dsts：

1
def hook_after_wait():
2
    dst_buffers = hook()
3
    copy_pairs = [(dst, dst_buffer) for dst_buffer, dst in zip(dst_buffers, dsts, strict=True)
4
                  if dst.data_ptr() != dst_buffer.data_ptr()]
5
    if copy_pairs:
6
        torch._foreach_copy_([dst.view(-1) for dst, _ in copy_pairs],
7
                             [dst_buffer.view(dst.dtype) for dst_buffer, dst in copy_pairs])
8
    return result

dst.data_ptr() != dst_buffer.data_ptr() 这一项是为了避免自拷贝：如果调用方本来就是用注册缓冲做输出，就不要再拷一遍。

把这条链路连起来，一次会话内的 cp.all_gather(x) 完整流程是：

1. 检查 x 是否已注册；未注册则从游标分配一块，并拷贝进去；
2. 调用 buffer.all_gather(src_buffers)，结果同样落在游标分配的注册区里；
3. 返回 handle，wait() 时把聚合结果的 rank 分片按列拼成 $N_{\text{cp}} \times T$ 再 flatten(0, 1)；
4. 若调用方给了外部 dsts，在 hook_after_wait 里拷回。

这与并行拓扑一篇里“gathered_hidden_states = userbuffer.cp.all_gather(hidden_states)()”的调用一一对应：外面看到的是一个同步调用，里面是三段式：拷贝、通信、回填。

DeepInfer 在哪些地方开 CP 会话

CP 会话在 DeepInfer 里有五处调用点，基本覆盖了 prefill 阶段所有“需要序列维度完整视野”的通信：

其中 update_main_cache 和 _prepare_swa_kv 展示了两种不同的“写入方式”：

1
def update_main_cache(self, main_k, kv_cache, main_positions, slot_mapping_store_swa, mirror_cp_reverse_index) -> None:
2
    use_cp = mirror_cp_reverse_index is not None
3
    with userbuffer.cp_session(enabled=use_cp):
4
        if use_cp:
5
            main_k = self.kv_norm(
6
                main_k,
7
                out=userbuffer.cp.get_inplace_tensor(main_k.shape, main_k.dtype),
8
            )
9
            wait_main_k = userbuffer.cp.all_gather(main_k.unsqueeze(1))
10
            main_k = gather(wait_main_k(), mirror_cp_reverse_index)
11

12
    fused_norm_rope_hadamard_cast_store_kvcache(
13
        main_k,
14
        norm_args=None if use_cp else NormArgs(self.kv_norm.weight, self.kv_norm.variance_epsilon),
15
        ...

这里 get_inplace_tensor 是 CPManager 提供的“预分配”接口：它在会话里按 (N_cp, *shape) 分配一整块聚合区，返回本 rank 对应的那一片。于是一个很值钱的优化出现了：归一化 kernel 的输出直接写进注册区，省掉一次“算到临时张量、再拷进缓冲”的往返；随后 all_gather 看到源已注册，也就不再拷贝。注意 norm_args=None if use_cp else ...：CP 开启时归一化已经在 kv_norm(..., out=...) 里做过了，融合 kernel 不再重复做归一化。

而 _prepare_swa_kv 走的是另一条路：直接让 fused_norm_rope_hadamard_cast_store_kvcache 的 out_args 指向 get_inplace_tensor() 的返回切片。这条路径上的 CastArgs 带 cast_back=True，out_stage 保持默认的 transform：内核按 e4m3、每通道元素数 swa_cache_num_per_channels 做一遍量化-反量化的往返，输出的是“经过 RoPE、Hadamard 变换、并按缓存精度舍入过”的 BF16 张量，写进 CP 注册区。之所以要先舍入再聚合，是为了让聚合后的 kv_swa 与直接写进 SWA 缓存的那份在数值上同一口径——prefill 之后 store_kv_cache(kv_swa, kv_cache.swa_cache, ...) 用的就是这个聚合结果，后续层读缓存时看到的也就是它，两边如果精度不一致，同一段上下文会在不同层得到不同的 K。

这里有一个容易看反的地方：CP 开启时 out_args 是 CP 的注册切片，kv_cache 并不直接参与写入，cast_back 的意义不在“写缓存”，而在“让聚合前后的数值口径对齐”。两条路径的共同点是：通信的源是注册缓冲，不是普通张量；差别只在于谁负责往里写、写的是哪一个阶段的结果。

会话本身的开销也要算清楚。cp_session() 是一个上下文管理器，进入时替换三个实例属性、退出时删除；每次进出的 Python 开销很小，但它约束了“会话内所有集合调用必须用同一组”（_check_group），也禁止嵌套。DeepInfer 的调用点都是“一次会话、一到两次集合调用”，所以这个约束不会带来实际麻烦。

（本站绘制，依据 BucketSession 的游标分配与 hook_after_wait 回填逻辑。图中左侧是本 rank 的源张量，中间是会话游标分配的注册区域——每次分配按 RDMA 对齐向上取整，右侧是聚合结果在 flatten(0, 1) 之后的形状。会话结束时游标归零，下一层可以复用同一块区域。）

决策一：为什么只允许一次

assert self.engram is None, "allocate_engram can only be called once"。

与 EP 缓冲“多层共享”不同，Engram 缓冲只分配一次，并且不能被复用。原因在于 EngramBuffer 的状态不只是存储，还包括 set_config 写入的表结构：层数、每层条目数、每 token 取几条、dtype、scale 打包方式。这些一旦写入，fetch(indices) 的索引张量形状就固定了——它的索引是 [num_layers, num_tokens, num_entries_per_token]；如果复用旧缓冲去装另一个结构的表，索引会静默错位。

DeepInfer 的 Engram 是所有层共用一张表族（num_entries_per_layer 是列表，长度等于层数），所以一个缓冲就够；这个断言防止的是“模型里意外出现两个 Engram 运行时”这种情形。

决策二：两个尺寸，两种存储

get_storage_size_hint 返回的是 (num_gpu_bytes, num_rdma_storage_size)，它本身只是把参数转交 C 侧的 calculate_storage_size：

1
return _C.EngramBuffer.calculate_storage_size(
2
    comm.get_nccl_comm_handle(group).get(),
3
    num_entries_per_layer, hidden, dtype.itemsize,
4
    num_max_tokens, num_entries_per_token, num_sf_packs)

docstring 把两个数的含义写得很清楚：

• num_gpu_bytes：GPU 侧的“fetch 接收区 + 本 rank 的那一份 FP8 scale 因子表”。接收区要能同时容纳一次 fetch 返回的所有层的结果，所以它随 max_tokens × entries_per_token × hidden × itemsize × num_layers 增长；
• num_rdma_storage_size：每个 rank 注册的存储大小，也就是这个 rank 负责的那份表分片的字节数。

后者默认放在 CPU 内存里：

1
use_cpu_rdma_storage = not envs.DEEPINFER_ENGRAM_FORCE_GPU_STORAGE

这个默认值是整个 Engram 设计的关键。表的总量按“层数 × 条目数 × 隐维 × dtype”算，FP8 下也已经远超单卡显存；放在 CPU 里，通过 RDMA 直接读，意味着不需要把表复制到 GPU 再读。代价是网络延迟：一次 fetch 要走 PCIe + 网卡，延迟远高于 HBM 访问；但 Engram 的访问模式是一次性、大粒度的 gather，延迟可以被批量掩盖。

EngramBuffer 构造时还有一个处理注册大小的细节：

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

混合模式下，NCCL 的 window 需要把“本节点上所有 rank 的注册存储”都映射进一个地址空间；按 16 个本地 rank 预留，再加上 4 GB 的余量，是为了容纳 GPU 接收区、各 rank 的存储、以及 NCCL 自身的簿记。如果这个总量超过设备总显存，就调大 NCCL_WIN_STRIDE，让 window 的虚拟地址跨度足以容纳它——这是一个“先按最坏情形设地址空间、再允许真实分配超过显存”的技巧，因为注册存储本身多半在 CPU，真正的 GPU 占用只有 num_gpu_bytes。如果本地 rank 数确实不是 16，EP_NUM_MAX_LOCAL_RANKS 可以覆盖这个假设。

决策三：QP 数与队列深度按 NIC 读能力推

get_theoretical_config 返回三个值，背后是一个很直接的估算：

1
nic_read_mpps = 90 if check_fast_rdma_atomic_support() else 45
2
return _C.EngramBuffer.get_theoretical_config(
3
    comm.get_nccl_comm_handle(group).get(),
4
    num_layers, max_num_tokens, num_entries_per_token,
5
    nic_read_mpps, num_max_read_requests_per_qp)

nic_read_mpps 是网卡每秒能完成的读请求数（单位是百万次，Mpps），90 与 45 对应新一代与老一代网卡的能力；num_max_read_requests_per_qp 默认 16，即每个 QP 上同时挂多少个未完成的读。C 侧的估算逻辑可以写成：

其中 $L$ 是层数，$T$ 是最大 token 数，$k_{\text{per token}}$ 是每 token 取几条条目（num_entries_per_token），qp_depth 是每个 QP 上并发的读请求上限。

这个推导的含义是：QP 数不是拍脑袋定的，而是“要让所有层、所有 token 的读请求能在单步内发出”所需的最小并发度。 如果 QP 太少，读请求会排队，fetch 会退化成串行；如果 QP 太多，网卡的连接表与内存开销会变大。get_theoretical_config 上还挂了 functools.lru_cache，说明这个推导按参数只需要算一次。

restrict_rd_atomic 是第三个返回值：它被传给构造器，再以环境变量的形式传给 NCCL：

1
with EnvOverride([
2
    ('NCCL_GIN_TYPE', '3', True),
3
    # TODO: Configure RD atomics through NCCL devComm config when supported.
4
    ('NCCL_GIN_GDAKI_MAX_QP_RD_ATOMIC', '8', restrict_rd_atomic),
5
]):
6
    self.nccl_comm_handle = comm.get_nccl_comm_handle(group, force_new_comm=True)

NCCL_GIN_TYPE=3 指定使用 GDAKI 后端（GPU 直接访问网卡）以加速初始化，并在构造期间临时覆盖环境变量，构造结束后恢复。这是为了让 Engram 的初始化走“快速路径”，而不受进程其他部分环境配置的影响；force_new_comm=True 意味着它不复用 EP 的 NCCL 通信器，而是单独建一个——两种流量的 QP 配额、超时、后端参数本来就不同。

决策四：RDMA service level 的流量隔离

allocate_engram 里有一行注释：

1
# EP 通信走 nvlink，此处 Engram sl 保持默认 sl3（sl2 留给 kv cache 流量隔离）

InfiniBand/RoCE 的 service level（SL）决定流量进入哪个虚拟通道，不同 SL 在网卡上有不同的缓冲与优先级。DeepInfer 把 SL2 留给 KV Cache 的跨实例传输（P/D 分离、HiCache 之类的场景），把 Engram 的读请求留在默认的 SL3。这是流量工程，而不是性能调优：如果把 Engram 的大量小读请求与 KV Cache 的大块搬运放在同一个 SL，队列会互相干扰，P/D 分离的延迟会出现毛刺。

决策五：set_config 不是可选的

分配完 EngramBuffer 之后立刻调用 set_config：

1
self.engram.set_config(
2
    num_entries_per_layer,
3
    hidden_size,
4
    max_num_tokens,
5
    num_entries_per_token,
6
    dtype,
7
    num_sf_packs,
8
)

EngramBuffer 的构造参数只告诉 C++“多大内存、多少 QP、多久超时”，而 set_config 告诉它“这些内存怎么解释”：每层条目数、条目宽度、一次最多取多少、元素字节数、scale 因子怎么打包。没有这一步，fetch(indices) 无法把索引张量翻译成 RDMA 目标地址。

write(storages, sfs) 则把表内容写进注册存储。在 DeepInfer 的装载流程里，这一步发生在 initialize() 内部：先把 HF 权重里的表整理成 CPU/GPU 上的分片，再调用 write 写进注册存储，最后释放原始张量。CP 组对象广播那一篇已经提到，这类跨 rank 的初始化需要在所有 rank 同时进行。这也解释了为什么 Engram 的加载在 userbuffer 之后——它需要一个已经能接收 RDMA 写的内存区域，而这个区域是 allocate_engram 才建出来的。

fetch 的返回则是一个 per-layer 的等待回调列表：索引一次发出、所有层的 RDMA 读同时进行，调用方逐层 wait() 再取结果。这与 CP 的 all_gather 返回 handle 的设计是同一种思路——先把通信发出去，把等待推迟到真正要用数据的那一刻，从而让通信与前向的其他计算重叠。

用在哪里

mlp.py 里只有一处调用，在 MoE 的 combine 之后、共享专家前向之前：

1
# Combine
2
combine_pre_evt = self.ep_buffer.capture()
3
with limit_num_sms(envs.DEEPINFER_USERBUFFER_COMM_SMS, enabled=PLATFORM.is_gpu):
4
    shared_out = self.shared_forward(hidden_states)
5
event_overlap = self.combine(
6
    moe_out_hidden,
7
    shared_out,
8
    handle,
9
    combine_pre_evt,
10
)
11
moe_combined_hidden, _combined_topk_weights = event_overlap.current_stream_wait()

看一下时序：combine 先 capture() 一个事件，然后在共享专家前向期间把计算用的 SM 压小，最后 combine() 把通信与共享专家前向一起发出去，用事件在流上等待两者都完成。这是一个典型的“通信与计算重叠”切片：combine 在跑的时候，共享专家的 GEMM 也在跑，但后者被限制在 $N_{ ext{compute}}$ 个 SM 内，保证两者不会因为争抢 SM 而互相拖慢。

如果只看 limit_num_sms 的名字，很容易以为它“给通信让出 SM”；更准确的说法是“给计算限配额”。通信侧的 num_sms 是请求值，计算侧的 SM 上限是保护值；capture() 加 current_stream_wait() 则保证两个任务在同一个流上的完成顺序正确——只有两者都结束，才把合并后的 hidden 交给下一层。

enabled=PLATFORM.is_gpu 说明在非 GPU 平台（例如 Ascend）上这个分区不启用：那里的通信与计算资源划分由平台自己的调度器负责，Python 侧改 SM 预算没有意义。

一个限制：GPU 上的全局口径

num_sms 是 DeepGEMM 与 TileLang 的进程级全局设置，不是 per-stream、per-thread 的。这意味着：

• 在 with 块内，任何并发执行的 DeepGEMM kernel 都会看到缩小后的 SM 数；
• 块内有 Python 异常时，finally 保证恢复。但若同样的写法出现在 CUDA Graph 捕获过程中，图里就会永久写死缩小后的 SM 配置（捕获只记录当时的状态，finally 恢复的只是捕获之后的新状态）——这也是这类块只应出现在捕获范围之外的原因；
• 多个 worker 线程共享这个设置，所以要保证同一时刻只有一个地方在改它。DeepInfer 的做法是只在 MLP 前向的单一位置调用，天然避免了竞争。

读到这里可以回头看 allocate_ep 里的 SM 推导：它给出的是一个“单步内通信能用多少 SM”的静态估计，而 limit_num_sms 是在执行期把这个估计施加到计算侧。前者决定缓冲区尺寸与 QP 数，后者决定这两笔预算在时间轴上的实际占用。

为什么要“视图”

MegaMoE 把 gate、dispatch、grouped GEMM、combine 融进一个内核，rank 之间通过对称内存指针交换数据（并行拓扑一篇讲到过 deep_gemm.fp8_fp4_mega_moe）。DeepGEMM 的 SymmBuffer 需要一组固定的参数：专家数、每 rank 最大 token、topk、隐维、中间维、共享专家、MMA 类型、激活类型。不同模型层这些参数通常相同，所以按规格缓存是自然的。

但 DeepInfer 里有两类层会用到它：主干 MoE 层与 MTP 层。它们的专家数或中间维可能相同，也可能不同——MTP 的草稿头往往是一个更小的 MoE。如果每层都独立分配，显存会随层数线性翻倍；SymmBuffer 的“base 视图”机制恰好解决了这一点：物理存储只分配一次，其他规格只是在这块存储上按不同形状重新切。

use_mega_moe 的判定也决定了它与 allocate_ep 是二选一：

1
self.use_mega_moe = (
2
    (PLATFORM.is_gpu or self.mma_type == "fp8xfp4")
3
    and self.mma_type in ("fp8xfp4", "fp8xfp8", "bf16xbf16")
4
    and PG.is_within_nvl_domain(PG.expert_parallel_group())
5
)

也就是说，MegaMoE 要求 EP 组落在同一个 NVLink 域内（is_within_nvl_domain）。一旦 EP 跨节点，对称内存指针就不再指向同一地址空间里的对端，MegaMoE 无法工作，只能退回 DeepEP 的 EPBuffer。这个判定解释了生产拓扑里的一个隐含约束：Wide-EP 跨节点部署时，MoE 前向走的是 EPBuffer 路径，userbuffer.sym_buffers 大概率是空的。

“首个必须最大”为什么是正确性约束

base = next(iter(self._sym_buffers.values()), None) 取的是字典里第一个插入的缓冲，而不是“当前最大”的缓冲。这意味着：

• 如果第一次分配的是 num_experts 已包含冗余专家、hidden 与 intermediate 都是主干规格的层，后续 MTP 层的规格更小，视图形状都能被包含在 base 的存储里；
• 如果第一次分配的是 MTP 的小规格，后续主干层视图需要更大的存储，就会越界——而 SymmBuffer 的越界不一定抛异常，可能是读到相邻 rank 的数据，表现为“结果偶发错误”。

因此 docstring 里那句“第一个规格必须最大”不是建议，而是分配顺序的硬约束。它的实现依赖调用顺序，而 allocate_userbuffer 的遍历顺序恰好保证了这一点：decoder_layers 先列 self.model.layers 的全部主干层，之后才 += MTP_module.moe_layer_list()。 如果哪天有人调整了遍历顺序，比如先处理 MTP 再处理主干，这条约束就会被悄悄破坏，且不会有断言拦住。

这是一类典型的“文档约束 + 顺序依赖”的隐患，也提示了一个通用原则：当缓存键与生命周期绑定到“首个实例”时，顺序就是接口的一部分。

为什么要先做 SM locality 初始化

base is None and PLATFORM.is_gpu 分支里的这段代码看起来与缓冲无关：

1
locality_init = getattr(getattr(deep_gemm, "locality_domain", None), "get_balanced_sm_locality_domains", None)
2
if locality_init is not None:
3
    locality_init()

注释已经把原因写清楚了：新版 DeepGEMM 会在首次使用时探测并缓存 SM locality，这个探测会同步。 如果它发生在 guarded forward 或 CUDA Graph 捕获期间，就会导致捕获失败——图里不能出现同步点。所以 userbuffer 在分配第一个 SymmBuffer 之前主动触发一次，把同步行为挪到“一切都还没开始”的时刻。

这是一个很值得记住的工程模式：任何“首次调用会同步/分配/初始化”的库函数，都必须在 CUDA Graph 捕获之前预热。 CUDA Graph 捕获那一篇提到过类似的预热逻辑，这里是从缓冲分配侧又补了一处。

getattr(...) 的双层防御式写法则说明 DeepInfer 要兼容不同版本的 DeepGEMM：老版本没有 locality_domain，就用 getattr 返回 None 跳过；新版本有则调用。这种“可选能力探测”在多后端框架里很常见（讲 TileLang 双平台时也说过同样的思路）。

_asdict() 与“fail fast”

deep_gemm.SymmBuffer(**spec._asdict(), base=base) 这一行把 NamedTuple 直接展开成关键字参数。这样做的好处是规格与构造参数一一对应，不会漏传或多传——_SymBufferSpec 的 docstring 也明确写了“fields are exactly the deep_gemm.SymmBuffer kwargs”。缺点是一旦 DeepGEMM 改了构造签名，这里会直接 TypeError，而不是静默行为变化。这是刻意的“fail fast”设计：通信缓冲的规格错配如果静默发生，后果是显存越界或数据错位，比启动失败更难排查。

SymmBuffer 构造时还会在 EP 组内做一次 barrier（它和 EPBuffer、EngramBuffer 一样，都在构造末尾 torch.cuda.synchronize() 后 group.barrier()）。所以首次分配 SymmBuffer 也是一个跨 rank 的同步点——它必须在所有 rank 都会走到的代码路径里调用，不能出现在“某个 rank 恰好分支不同”的地方。

与 allocate_ep 的二选一

DeepseekMoE.allocate_userbuffer 里：

1
def allocate_userbuffer(self) -> None:
2
    if self.use_mega_moe:
3
        self.sym_buffer = userbuffer.allocate_sym_buffer(...)
4
    else:
5
        self.ep_buffer = userbuffer.allocate_ep(...)

一个 MLP 层只持有其中一种。但 userbuffer 里 ep 与 _sym_buffers 可以并存：如果同一个模型里只有部分层满足 MegaMoE 条件，另一部分层走 EPBuffer，那么两者都会被分配，显存预算要按两者之和算。实践中 use_mega_moe 依赖的是 EP 组是否在同一 NVLink 域，这是拓扑级属性，对同一组的所有层都成立，所以正常部署里并存的概率不高——但代码并没有对此做断言，不能默认二者互斥。

close() 的顺序也反映了这一点：先销毁 ep/cp/engram，再逆序销毁 sym 视图。先销毁 base 再销毁视图会出错，因为视图引用了 base 的存储；逆序（后分配的视图先销毁）恰好保证 base 最后销毁。vars(buffer).clear() 在销毁后显式清空对象字典，是为了让 Python 层立刻丢掉对底层句柄的引用，不再依赖 GC 时机。

阶段一：构造之前

Worker.init_model 调用 DeepseekForCausalLM.__init__，在其末尾分配 userbuffer，而这一时刻前面还有几件事已经完成：

• HF config.json 读入后，finalize_model_config 补全 max_num_decoding_tokens、max_num_prefilling_tokens 与专家冗余布局；
• 权重装载完成，专家按 EP rank 物化（见权重装载一篇），所以每层都知道自己的 num_physical_experts 与 num_physical_topk；
• EPLB 初始化完成（见 EPLB 一篇），冗余专家数已确定。

也就是说，缓冲尺寸依赖的不是“配置里写了什么”，而是“物化之后每个 rank 实际持有多少专家”。 如果权重装载阶段因为 EP 冗余把每个 rank 的物理专家数从“总专家数 / EP”变成了更大的数，allocate_sym_buffer 和 allocate_ep 里的 num_experts 就必须是这个物化后的数，否则 SymmBuffer 的视图会少切一列、EPBuffer 的尺寸估算也会偏小。这一依赖关系解释了为什么 allocate_userbuffer 必须放在模型构造的最后，而不能放在 __init__ 开头。

阶段二：分配

torch.cuda.empty_cache() 之后按 “MoE 层 → CP → Engram” 的顺序分配。这个顺序不是随意的：

• MoE 在前：EP/SymmBuffer 的峰值显存最大，先分配给它们，避免被其他分配切碎地址空间；
• CP 在中：它依赖各注意力层的 get_max_ag_bytes_per_token()，需要遍历所有层；
• Engram 在后：Engram 表需要 set_config + write，而且它的存储默认在 CPU，分配的是 RDMA 注册，与 GPU 地址空间的关系最小。

如果顺序反过来（先 Engram），Engram 的 NCCL window 会先占掉一段很大的虚拟地址空间，后续 EP 的对称内存对齐可能被挤压。

阶段三：捕获与重放

CUDA Graph 捕获时（见 CUDA Graph 一篇），userbuffer 里的缓冲已经存在且尺寸固定。这是图能被捕获的前提：图里每一次通信内核的指针、大小、offset 都在捕获时被写死。运行时如果重新分配缓冲，图重放会写到旧地址。所以 userbuffer 的分配时机与 CUDA Graph 是硬耦合的——先有缓冲，后有图；缓冲变了，图就作废。

阶段四：关闭

worker/base.py 的 Worker.stop 里，关闭顺序写得很讲究：

1
def stop(self) -> None:
2
    # `os._exit(0)` below skips destructors, which would leave the driver worker's UB
3
    # stuck at its destructor barrier, so close userbuffer explicitly between barriers.
4
    dist.barrier()
5
    userbuffer.close()
6
    dist.barrier()
7

8
    # NOTE necessary to avoid destroy process group hangs.
9
    # The reason is that when using NCCL inside CUDA graphs, NCCL allocates some extra space (a kind of dependency).
10
    # When destroying the process group, NCCL checks whether all dependencies have been released.
11
    # SEE https://github.com/pytorch/pytorch/issues/115388
12
    ...

注释里那句“os._exit(0) 跳过了析构函数，所以 driver worker 的 UB 会卡在它的析构 barrier 上”值得展开：

• 为什么是 barrier 夹住 close？ userbuffer.close() 会调用每个 DeepEP/DeepGEMM 缓冲的 destroy()，而 destroy() 内部涉及跨 rank 的协调（例如 NCCL 的 window 注销需要所有 rank 参与）。如果某个 rank 先销毁、另一个 rank 还在跑，就会挂起。两侧的 dist.barrier() 保证所有 rank 同时进入、同时退出。
• 为什么不能依赖析构函数？ 进程退出时（os._exit(0)）Python 不会跑 __del__，C 侧的析构器也不会按预期顺序执行；DeepEP 缓冲的 C++ 对象析构时可能等一个跨 rank 的 barrier，若某个 rank 已经退出，就会永久阻塞。显式 close 把“同步销毁”挪到一个所有 rank 都还活着的时刻。
• 为什么先关 userbuffer，再销毁 CUDA Graph？ 注释紧接着说明了原因：CUDA Graph 里用了 NCCL 时，NCCL 会分配“图依赖”的额外空间，destroy_process_group 会检查这些依赖是否已释放（PyTorch issue 115388）。userbuffer 的销毁本身不需要图存在，放在销毁图之前是安全的，也更早释放跨 rank 资源。

这条顺序可以概括为：barrier → close → barrier → 销毁图 → 清理 shm 广播 → 销毁进程组。

阶段五：cleanup_shm_broadcast

PG.cleanup_shm_broadcast() 会主动关闭并 unlink MessageQueue 的共享内存：

1
def cleanup_shm_broadcast(self) -> None:
2
    """Explicitly close and unlink shared memory used by MessageQueue,
3
    so that the resource_tracker won't report leaked shared_memory warnings."""
4
    if self.mq is not None:
5
        if self.mq.buffer is not None:
6
            self.mq.buffer.shared_memory.close()
7
            if self.mq.buffer.is_creator:
8
                self.mq.buffer.shared_memory.unlink()
9
        self.mq = None

这与 userbuffer 不是同一个对象，但属于同一类问题：跨 rank 的共享资源必须在进程退出前显式析构，否则要么泄漏，要么挂起。 CP 组对象广播一篇拆过这个 MessageQueue，这里只是补上它的收尾。

vLLM

vLLM 的 EP 通信同样基于 DeepEP，但 buffer 的组织方式不同：

• 按 Worker 持有，按模型层共享。 DeepEP 的封装在 Worker 初始化时创建 buffer，所有 MoE 层复用；
• SM 数可以通过配置显式给出。 配置里通常有 num_sms 或 num_max_dispatch_tokens_per_rank 之类的字段，不够时直接报错，而不是像 userbuffer 那样做“理论值 → 保留计算 SM → 设备一半”的三层裁剪；
• 没有 CP 的 BucketBuffer：CP 在 vLLM 里主要由 NCCL 进程组承担，没有会话式零拷贝缓冲。

差异的根源在于拓扑：vLLM 支持 TP，CP 的通信与 TP 的通信在同一个 NCCL 进程组体系里，没必要单独为 CP 造一层缓冲。

SGLang

SGLang 的 DeepEPConfig 更接近 DeepInfer 的思路：它按带宽模型算 SM 数，也按 max_num_tokens 与 hidden 算 buffer 大小。差别在三点：

1. SGLang 把 DeepEP 配置做成独立类，而 DeepInfer 把它塞进 userbuffer 的一个方法，尺寸推导与模型结构强耦合（依赖物化后的物理专家数）；
2. SGLang 没有 Engram 缓冲，也没有 SymmBuffer 的“视图缓存”机制；它的对称内存用法更多是 DeepGEMM 的普通调用，而不是按规格缓存的视图管理；
3. SGLang 的弹性 EP 需要支持动态扩容，通信缓冲要能在“世界变大”后重建；DeepInfer 的 userbuffer 是“一次分配、不复用更大规格”，更适合固定规模的集群。

TensorRT-LLM

TensorRT-LLM 的 BufferManager 与 userbuffer 在职责上最接近：它统一管理 KV 池、workspace 与 CUDA Graph 缓冲，按预算分配（见 BufferManager 一篇）。区别在于：

• TensorRT-LLM 的通信走 NCCL/Mpi，BufferManager 主要管显存，不管通信缓冲；
• DeepInfer 把通信缓冲也纳入预算，因为它的通信不走通用 NCCL，而是 DeepEP/DeepGEMM 的专用原语，缓冲的尺寸直接决定通信性能。

llama.cpp

llama.cpp 不涉及多卡 EP 缓冲，但在 CP 的思路上有一个可对照的点：它的 server 用 slot 复用一个 KV 缓存，等价于“会话式复用”的低配版。DeepInfer 的 BucketSession 比它多了一层“通信期间的临时分配游标”，本质相同：所有临时资源走同一块区域，退出即失效。