TensorRT-LLM 源码完全拆解（十一）：并行拓扑与通信——Mapping、DeviceMesh 与 Communicator 如何组织一次多卡前向

第一层校验：主并行的乘法关系

构造函数里最先做的是一组硬校验。最关键的一条是：

1
if tp_size * pp_size * cp_size != world_size:
2
    raise ValueError(
3
        "world_size must equal to tp_size * pp_size * cp_size, "
4
        f"but got {world_size} != {tp_size} * {pp_size} * {cp_size}.")

这条约束把 TP、PP、CP 定义为正交的三个维度，三者相乘恰好覆盖全部 rank，不允许重叠也不允许留空。原因在于后面所有组划分都基于”全局 rank 号做一个混合进制分解”：

这条公式直接对应 MpiTopology 里的三个 property：

1
@property
2
def tp_rank(self) -> int:
3
    return self.rank % (self.tp_size * self.cp_size) // self.cp_size
4

5
@property
6
def pp_rank(self) -> int:
7
    return self.rank // (self.tp_size * self.cp_size)
8

9
@property
10
def cp_rank(self) -> int:
11
    return self.rank % self.cp_size

用大白话讲：PP 是最高位、TP 是中间位、CP 是最低位。全局 rank 先按 PP 分段，段内再按 TP 分段，每段内连续的 cp_size 个 rank 是同一个 CP 组。

顺序不是随便定的，它决定了通信组在物理拓扑上是否连续。CP 放在最低位意味着一个 CP 组是连续编号的 rank，能落在同一台机器内部的连续 GPU 上；TP 组则是步长为 cp_size 的交错 rank 集合。

Mapping 的 docstring 给了几个具体例子，值得逐个对照：

1
A node with 8 GPUs, tp_size = 4, cp_size = 1, pp_size = 2
2

3
2 tp groups:
4
- [0, 1, 2, 3]
5
- [4, 5, 6, 7]
6

7
4 pp groups:
8
- [0, 4]
9
- [1, 5]
10
- [2, 6]
11
- [3, 7]

tp_size=4, cp_size=1, pp_size=2，world_size=8。按公式，tp_rank = rank % 4，pp_rank = rank // 4。rank 0–3 是 pp_rank 0，rank 4–7 是 pp_rank 1。TP 组是同一 pp_rank 内所有 rank，所以是 [0,1,2,3] 和 [4,5,6,7]；PP 组是”相同 tp_rank、不同 pp_rank”，即 [0,4]、[1,5]……

再看 CP 非 1 的情况：

1
A node with 8 GPUs, tp_size = 4, cp_size = 2, pp_size = 1
2

3
4 cp groups:
4
- [0, 1]
5
- [2, 3]
6
- [4, 5]
7
- [6, 7]
8

9
2 tp groups:
10
- [0, 2, 4, 6]
11
- [1, 3, 5, 7]

这里 CP 组是连续对 [0,1]，TP 组变成步长 2 的交错集 [0,2,4,6]。这正是”CP 放最低位”的直接后果。把 CP 放在最低位，是为了让上下文并行组内的通信走机内 NVLink——CP 的通信量（每个 token 的所有 KV）远大于 TP 的通信量（每个 token 的 hidden state），更需要高带宽域。

第二层：MoE 并行的三重复合关系

MoE 层的并行比稠密层复杂得多。Mapping 里与 MoE 相关的字段有四个，彼此之间还有一个必须成立的等式：

1
moe_world_size = tp_size if cp_type == CpType.ULYSSES else tp_size * cp_size
2

3
moe_tp_ep_size = moe_tp_size * moe_ep_size
4
self.moe_tp_cluster_ep_size = moe_tp_ep_size * moe_cluster_size
5
if not (dwdp_size > 1) and self.moe_tp_cluster_ep_size != moe_world_size:
6
    raise ValueError(
7
        "moe_tp_size * moe_ep_size * moe_cluster_size must equal to moe_world_size, "
8
        f"but got {self.moe_tp_cluster_ep_size} != {moe_world_size}")

即

为什么必须相等？因为专家的总张量副本数是被并行度整除的关系决定的。一个 MoE 层共有 $E$ 个专家，如果 moe_ep_size = 8，那么每个 rank 只装 $E/8$ 个专家。同时 moe_tp_size 又把每个专家自己的权重切开，moe_cluster_size 则是把若干相邻的 EP 组捆成一个”cluster”（用于让一个 token 在 cluster 内部就能找到它需要的所有专家，从而避免跨 cluster 的 all-to-all）。三个维度各自乘起来，必须等于”所有参与 MoE 计算的 rank 数”——否则专家的账目对不上。

注意 moe_world_size 的取值依赖 cp_type：ULYSSES 模式下 CP 与 TP 是”同一批物理卡上的两种视角”，所以 MoE 的世界大小等于 tp_size；而 HELIX 模式下 CP 的这批 rank 在 FFN 层会被重新用作 TP（见后面的 repurpose_helix_cp_to_tp），所以 MoE 世界大小扩到 tp_size * cp_size。

第三层：注意力层的独立切法

attn_tp_size 和 attn_cp_size 是 TensorRT-LLM 里比较特殊的一组参数。它们的约束是：

1
attn_tp_cp_size = attn_tp_size * attn_cp_size
2
if attn_tp_cp_size != tp_size * cp_size:
3
    raise ValueError(
4
        "tp_size * cp_size must equal to attn_tp_size * attn_cp_size, ...")
5

6
if attn_cp_size != 1 and cp_type == CpType.ULYSSES:
7
    raise ValueError(
8
        f"attn_cp_size must be 1 for now for ulysses, but got {attn_tp_size}, {attn_cp_size}.")

含义是：注意力层可以把 TP 和 CP 的”合并维度”重新分配。总容量 $tp\_size \times cp\_size$ 不变，但其中多少用于 TP、多少用于 CP 可以由注意力层自己决定。这在 HELIX 模式下很有用——当 cp_size=8, tp_size=1 时，注意力层可以把 attn_cp_size 设得小一些、attn_tp_size 设得大一些，让注意力也吃到张量并行的好处。

ULYSSES 模式不允许 attn_cp_size != 1，因为 Ulysses 的 CP 是”按 head 维度切分再 all-to-all”的经典实现，其通信组与 TP 组是同一个物理集合的不同视图，把 CP 再拆开就失去了意义。

组划分：MpiTopology._init_parallel_groups

MappingBase.__init__ 只负责校验和保存标量；真正的组划分发生在 MpiTopology._init_parallel_groups（mapping.py）。它是一段非常值得逐行读的代码，因为它用四条 range 表达式把前面那张拓扑表全部实现了：

1
def _init_parallel_groups(self):
2
    # init pp group
3
    for i in range(self.tp_size * self.cp_size):
4
        ranks = range(i, self.world_size, self.tp_size * self.cp_size)
5
        self.pp_groups.append(list(ranks))
6

7
    # init cp group (consecutive ranks within each tp slice).
8
    for i in range(self.pp_size):
9
        for j in range(self.tp_size):
10
            ranks = range(
11
                i * self.tp_size * self.cp_size + j * self.cp_size,
12
                i * self.tp_size * self.cp_size + (j + 1) * self.cp_size)
13
            self.cp_groups.append(list(ranks))
14

15
    # init tp group (interleaved ranks with stride of cp_size).
16
    for i in range(self.pp_size):
17
        for j in range(self.cp_size):
18
            ranks = range(i * self.tp_size * self.cp_size + j,
19
                          (i + 1) * self.tp_size * self.cp_size + j,
20
                          self.cp_size)
21
            self.tp_groups.append(list(ranks))

逐个理解：

PP 组：步长是 tp_size * cp_size，起点从 0 到 tp_size * cp_size - 1。这就是”固定 (tp_rank, cp_rank)，扫过所有 pp_rank”。共有 tp_size * cp_size 个 PP 组。

CP 组：在每个 (pp_rank, tp_rank) 组合下，取一段长度为 cp_size 的连续区间。共有 pp_size * tp_size 个 CP 组。

TP 组：在每个 (pp_rank, cp_rank) 组合下，从 i * tp_size * cp_size + j 开始，以 cp_size 为步长走到该 PP 段的末尾。共有 pp_size * cp_size 个 TP 组。

三段代码和前面 tp_rank/pp_rank/cp_rank 的公式完全自洽。以 8 卡、tp=4, cp=2, pp=1 为例验证一下 CP 组：

• i = 0（pp_rank 0），j 从 0 到 3：
  • j=0: range(0*8 + 0*2, 0*8 + 1*2) = [0, 1]
  • j=1: range(0 + 2, 0 + 4) = [2, 3]
  • j=2: [4, 5]，j=3: [6, 7]

得到 4 个 CP 组 [0,1] [2,3] [4,5] [6,7]，与 docstring 一致。

MoE 组的划分：为什么索引表达式这么长

MoE 的那三段循环索引表达式看起来吓人，但只要记住 “MoE 世界是三层嵌套：cluster 在外、EP 在中、TP 在内” 就能读懂：

1
# init moe tp group
2
for i in range(self.pp_size):
3
    for j in range(self.moe_cluster_size * self.moe_ep_size):
4
        ranks = range(i * self.moe_tp_cluster_ep_size + j,
5
                      (i + 1) * self.moe_tp_cluster_ep_size,
6
                      self.moe_cluster_size * self.moe_ep_size)
7
        self.moe_tp_groups.append(list(ranks))

MoE TP 组：固定 (pp_rank, moe_cluster_rank, moe_ep_rank)，扫过 moe_tp_size 个 rank，步长是 moe_cluster_size * moe_ep_size。也就是”同一份专家权重切片在不同 TP 分片上的那些卡”。

1
# init moe ep group
2
for i in range(self.pp_size):
3
    for j in range(self.moe_tp_size):
4
        for k in range(self.moe_cluster_size):
5
            ranks = range(
6
                i * self.moe_tp_cluster_ep_size +
7
                j * self.moe_cluster_size * self.moe_ep_size +
8
                k * self.moe_ep_size,
9
                i * self.moe_tp_cluster_ep_size +
10
                j * self.moe_cluster_size * self.moe_ep_size +
11
                (k + 1) * self.moe_ep_size)
12
            self.moe_ep_groups.append(list(ranks))

MoE EP 组：固定 (pp_rank, moe_tp_rank, moe_cluster_rank)，取连续的一段 moe_ep_size 个 rank。这是 all-to-all 通信真正走的组——每个 rank 把自己的 token 按路由结果发给同 EP 组内的其他 rank，再把收到的 token 算完发回去。

图源：TensorRT-LLM 官方文档（docs/source/media/moe-parallelism.png）。专家维度（EP）决定”哪些专家放在哪个 rank 上”，张量维度（TP）决定”每个专家自己的权重再切成几份”。图里几组配置的对比正是后面那条 moe_tp_size × moe_ep_size × moe_cluster_size = moe_world_size 等式的图解。

对照 docstring 的例子验证 moe_tp_size = 2, moe_ep_size = 4（8 卡）：

• MoE TP 组：j 从 0 到 moe_cluster_size(1) * moe_ep_size(4) - 1 = 3，步长 4，长度 8。j=0 得 [0, 4]，j=1 得 [1, 5]，j=2 得 [2, 6]，j=3 得 [3, 7]。4 个组，与 docstring 一致。
• MoE EP 组：j 只有 0（moe_tp_size=2 → j 取 0,1；但 moe_cluster_size=1 时 j 循环的是 moe_tp_size），k=0，得 range(0, 4) = [0,1,2,3]；j=1 得 range(4, 8) = [4,5,6,7]。2 个组，与 docstring 一致。

属性的取值索引也是精心设计的

有了组列表之后，每个 rank 要能”取出自己所属的那个组”。MpiTopology 靠的是与组生成顺序严格对应的索引公式：

1
@property
2
def cp_group(self) -> List[int]:
3
    return self.cp_groups[self.pp_rank * self.tp_size + self.tp_rank]
4

5
@property
6
def tp_group(self) -> List[int]:
7
    return self.tp_groups[self.pp_rank * self.cp_size + self.cp_rank]
8

9
@property
10
def pp_group(self) -> List[int]:
11
    return self.pp_groups[self.tp_rank * self.cp_size + self.cp_rank]

这三个索引公式和上面三个生成循环的循环变量顺序是一一对应的：cp_groups 是 (i=pp_rank, j=tp_rank) 双层循环，所以索引是 pp_rank * tp_size + tp_rank；tp_groups 是 (i=pp_rank, j=cp_rank)，所以是 pp_rank * cp_size + cp_rank；pp_groups 只有一层且以”固定的 (tp_rank, cp_rank)”为准，所以是 tp_rank * cp_size + cp_rank。

MoE 的三个 group 索引同理，只是层级更深：

1
@property
2
def moe_ep_group(self) -> List[int]:
3
    return self.moe_ep_groups[self.pp_rank * self.moe_tp_size *
4
                              self.moe_cluster_size +
5
                              self.moe_tp_rank * self.moe_cluster_size +
6
                              self.moe_cluster_rank]

这类”生成顺序 = 索引公式”的写法在分布式代码里很常见，也很容易写错。读这段代码时最实用的方法是：先在纸上写出组列表的遍历顺序树，再去核对 property 里的乘法权重。

repurpose_helix_cp_to_tp：HELIX 模式的”身份切换”

HELIX 上下文并行有一条独特的设计：CP 只对注意力层有意义，到了 FFN/MoE 层，这批 rank 被重新当成 TP 用。代码里体现为一个方法：

1
def repurpose_helix_cp_to_tp(self):
2
    # In helix parallelism, CP is relevant only for the attention layer.
3
    # These ranks are repurposed to TP for FFN layers.
4
    assert self.has_cp_helix()
5
    return Mapping(
6
        world_size=self.world_size,
7
        rank=self.rank,
8
        gpus_per_node=self.gpus_per_node,
9
        cp_size=1,
10
        cp_config={},
11
        tp_size=self.tp_size * self.cp_size,
12
        pp_size=self.pp_size,
13
        ...
14
    )

它返回一个新的 Mapping 对象：cp_size 归零，tp_size 变成原来的 tp_size * cp_size。也就是说，同一张卡在同一个模型里持有两个 Mapping——注意力层用前者，FFN 层用后者。这不是”改配置”，而是同一批硬件上两种不同的逻辑视图，这也是 HELIX 能把 CP 的通信开销摊薄的原因：CP 期间被切开的卡，在 FFN 期间全部拿去做 TP，算力不浪费。

张量并行的权重切分：从 Mapping 到 Linear

拓扑表建好了，具体到每一层怎么切权重？tensorrt_llm/_torch/modules/linear.py 里的 TensorParallelMode 只有两个取值，却定义了整套切分语义：

1
class TensorParallelMode(str, enum.Enum):
2
    COLUMN = 'column'
3
    ROW = 'row'
4

5
    @classmethod
6
    def split_dim(cls, mode):
7
        return 1 if mode == cls.ROW else 0
8

9
    # Helper to shard the corresponding per-channel activation scales
10
    # which are sharded along the dimension orthogonal to the weights
11
    @classmethod
12
    def flip(cls, mode):
13
        return cls.ROW if mode == cls.COLUMN else cls.COLUMN

split_dim 给出的是权重矩阵的切分维度：COLUMN 切 dim 0（输出维），ROW 切 dim 1（输入维）。flip 则是给量化 per-channel scale 用的——scale 的切分维度与权重正交，所以要从 COLUMN 翻成 ROW、从 ROW 翻成 COLUMN。这个细节在很多手写 TP 实现里会被忽略，结果是 scale 和权重错位，数值悄悄出错。

Linear.__init__ 里保存的 self.tp_mode、self.tp_rank、self.tp_size 全部来自 Mapping，切分本身在权重加载阶段完成（load_shard）。真正与通信相关的是 forward：

1
def forward(self, input, *, all_reduce_params=None, lora_params=None,
2
            layer_idx=None):
3
    if self.tp_mode == TensorParallelMode.ROW:
4
        ...
5
        bias = None if (self.tp_rank > 0) else self.bias
6
        if self.reduce_output:
7
            if use_fused_gemm_allreduce:
8
                output = self.apply_linear_allreduce(input, self.bias, layer_idx)
9
            else:
10
                ...
11
                use_nccl_symmetric_memory_window = (...)
12
                output = self.all_reduce(output, all_reduce_params=all_reduce_params)
13
        else:
14
            output = self.apply_linear(input, bias, lora_params, layer_idx)
15
    elif self.tp_mode == TensorParallelMode.COLUMN:
16
        output = self.apply_linear(input, self.bias, lora_params, layer_idx)
17
        if self.gather_output:
18
            from ..distributed import allgather
19
            output = allgather(output, self.mapping, sizes=self.gather_output_sizes)
20
    else:
21
        output = self.apply_linear(input, self.bias, lora_params, layer_idx)
22
    return output

三条分支对应三件事，值得逐一说明：

1. ROW + reduce_output：算完之后 all-reduce。这是主干路径，也是后面所有”融合”优化的落点。注意 bias = None if self.tp_rank > 0 else self.bias——bias 只在 rank 0 上加，因为 all-reduce 会把各 rank 的部分和相加，如果每张卡都加一遍 bias，最终结果会多出 tp_size 份。这是个很容易写错的细节。
2. COLUMN + gather_output：算完之后 all-gather。但 Transformer 里几乎不用它——列并行的输出通常直接喂给下一层，不需要拼回完整张量；只有需要完整输出的地方（比如 LM head 之后再切）才用。
3. 无 tp_mode：权重不切，纯本地计算。

还有一行容易被忽略的逻辑：

1
self.gather_output_sizes = None
2
if self.gather_output and self.tp_size > 1 \
3
        and self.tp_mode == TensorParallelMode.COLUMN:
4
    if self._tp_sharding_is_auto:
5
        sizes = [end - start for start, end in ...]
6
        ...
7
        self.gather_output_sizes = sizes

allgather 带 sizes 参数是为了支持不均匀切分（uneven TP）——当输出维度不能被 tp_size 整除时，各 rank 分到的宽度不同，all-gather 必须显式知道每段的长度，否则按等长拼接会错位。

MPIDist：把 MPI 当通信后端

MPIDist 直接持有 MPI 通信子：

1
class MPIDist(Distributed):
2
    tp_comm: MPI.Comm
3

4
    def __init__(self, mapping: Mapping):
5
        super().__init__(mapping)
6
        self._cp_comm = None
7
        self._tp_comm = None
8
        self._pp_comm = None
9

10
    def _validate_world_size(self):
11
        """Validate world size before creating sub-communicators to prevent segfaults."""
12
        if ENABLE_MULTI_DEVICE:
13
            actual_world_size = mpi_world_size()
14
            max_rank_needed = self.mapping.world_size
15
            if max_rank_needed > actual_world_size:
16
                raise RuntimeError(
17
                    f"Mapping requires world_size={max_rank_needed} "
18
                    f"(tp_size={self.mapping.tp_size} * pp_size={self.mapping.pp_size} * cp_size={self.mapping.cp_size}), "
19
                    f"but MPI world size is only {actual_world_size}. ")

_validate_world_size 的存在理由写在注释里——在建子通信子之前先校验，否则会 segfault。这是分布式编程里最经典的踩坑点：MPI_Comm_split 之类的调用如果 rank 数对不上，行为是未定义的，很多 MPI 实现直接崩掉而不是返回错误。把一个”配置错误”翻译成一次段错误，排查成本极高，所以这里宁可显式抛异常。

通信组按需延迟创建，三个 _*_comm 初始都是 None：

1
@property
2
def tp_comm(self):
3
    ...
4
@property
5
def cp_comm(self):
6
    ...
7
@property
8
def pp_comm(self):
9
    ...

只有真正用到某个维度的通信时才建对应的子通信子。这带来的好处是：如果一个作业里 CP 从不使用（cp_size=1），就永远不会创建 CP 通信子，也就不会有它的建组开销。

TorchDist：Ray 模式下的初始化顺序

Ray 模式下没有 MPI，通信后端是 PyTorch 原生的 torch.distributed。TorchDist.__init__ 藏着一段顺序关键的逻辑：

1
def __init__(self, mapping: Mapping):
2
    super().__init__(mapping)
3
    assert dist.is_initialized(
4
    ), "torch.distributed should be initialized before TorchDist"
5

6
    self.cluster_info = None
7

8
    from tensorrt_llm._utils import set_torch_comm
9
    set_torch_comm(self)  # Set as global instance
10
    mapping.build_mesh()
11

12
    self.setup_local_comm()
13
    self.default_store = torch.distributed.distributed_c10d._get_default_store()
14

15
    init_pg(torch.distributed.group.WORLD, self.local_comm,
16
            torch_pybind11_abi())

四个步骤，每一步都有理由：

1. set_torch_comm(self)：把自己注册成全局单例。C++ 侧的 custom op（比如 torch.ops.trtllm.allreduce_pg）在没有 MPI 时需要通过这个全局引用回拿通信对象。
2. mapping.build_mesh()：建 DeviceMesh。这一步必须在建 local comm 之前，因为 DeviceMeshTopology 的 tp_group_pg 等属性依赖网格。
3. setup_local_comm()：按节点划分本地通信组。
4. init_pg(...)：把 WORLD 组和 local 组的句柄传给 C++ 侧的进程组管理器（init_pg），让 custom op 能直接拿到 ProcessGroup。

setup_local_comm 的实现值得细看：

1
def setup_local_comm(self):
2
    self._get_cluster_info()
3

4
    # node IP -> list of ranks
5
    ip_to_ranks = {}
6
    for rank, (node_ip, _) in enumerate(self.cluster_info):
7
        ip_to_ranks.setdefault(node_ip, []).append(int(rank))
8

9
    self.local_comm = None
10
    for ranks in ip_to_ranks.values():
11
        # All global ranks from the default process group to participate in the call,
12
        # even if some ranks are not part of the new process group being created
13
        pg = dist.new_group(ranks=ranks, backend='cuda:nccl,cpu:gloo')
14
        if int(self.rank) in ranks:
15
            self.local_comm = pg

注释里那句话说的是一个非常具体的坑：dist.new_group 是一个集合操作，所有 rank 都必须调用它，哪怕自己不属于这个新组。这是 PyTorch 分布式的一个硬性要求——建组的内部握手需要在默认组上完成。所以代码写成”所有人都循环、只有属于该组的 rank 把结果存下来”。如果只让组内 rank 调用，其余 rank 会卡死在对默认组的集合操作里。

_get_cluster_info 用三次 all_gather_object 收集每个 rank 的 (node_ip, gpu_index, rank)，再拼成”按 rank 排序的节点信息表”。它只在 Ray 初始化后才可用（ray.util.get_node_ip_address()），否则直接抛 RuntimeError。这里体现的差异是：MPI 模式下进程拓扑由 MPI 运行时管理，节点分组可以靠 local_mpi_size() 直接问；Ray 模式下每个 actor 是独立进程，必须自己靠通信把拓扑拼出来。

流水线通信：PPCommNCCL 与 PPCommTorch

PP 的通信模式和 TP/CP 很不相同。TP/CP 用的是集合通信（all-reduce、all-gather），而 PP 用的是点对点 send/recv：第 $i$ 阶段的输出 tensor 直接发给第 $i+1$ 阶段。

communicator.py 为此单独做了一组：

1
class PPCommNCCL:
2
    def __init__(self, global_mapping: Mapping):
3
        ...
4
        self.nccl_comm = torch.classes.trtllm.NcclCommunicatorOp(...)
5

6
    def send(self, tensor: torch.Tensor, dest: Optional[int] = None):
7
        ...
8
    def recv(self, tensor: torch.Tensor, src: Optional[int] = None):
9
        ...
10

11
class PPCommTorch:
12
    def __init__(self, global_mapping: Mapping):
13
        ...
14
    def _global_to_local_rank(self, global_rank: int):
15
        ...

图源：TensorRT-LLM 官方文档（docs/source/media/Pipeline_Parallel_Vis.svg）。图中的横轴是时间、纵轴是 PP 阶段，每个色块是一次 microbatch 的前向；阶段之间的箭头对应下面要讲的 PPCommNCCL.send/recv。

PPCommNCCL 直接建了一个”世界通信子”（world communicator），而不是为每个 PP 组各建一个。为什么可以这么简化？因为 PP 的发送目标是确定的：从 pp_rank 到 pp_rank + 1，对应的全局 rank 可以通过 next_pp_rank 算出来。用世界通信子 + 点对点 send/recv 就足够了，不需要建 per-group 的通信子。

init_pp_comm 里有一大段关于通信子复用的注释，讲的是一个真实的生产问题：

1
def init_pp_comm(mapping):
2
    """Initialize PPComm once at startup"""
3
    global _pp_comm
4
    if mpi_disabled():
5
        _pp_comm = PPCommTorch(mapping)
6
    elif isinstance(_pp_comm, PPCommNCCL) and \
7
            _pp_comm.mapping.world_size == mapping.world_size:
8
        # Reuse the existing world NCCL communicator across LLM instances that
9
        # share the same worker processes (e.g. a reused MpiPoolSession). The
10
        # underlying comm depends only on (world_size, rank) -- it is a world
11
        # communicator, independent of the pp/tp/ep layout -- so only the
12
        # routing mapping needs refreshing. Recreating it would drop the old
13
        # comm and trigger a collective ncclCommDestroy at an unsynchronized
14
        # point during the next model build, which can deadlock on reused
15
        # workers. ...
16
        _pp_comm.mapping = mapping
17
    else:
18
        if _pp_comm is not None:
19
            logger.warning(
20
                "init_pp_comm: replacing existing PP comm (world_size "
21
                f"{_pp_comm.mapping.world_size} -> {mapping.world_size}) on a "
22
                "live process; this can deadlock on reused MPI workers.")
23
        _pp_comm = PPCommNCCL(mapping)
24
    init_helix_cp_comm(mapping)

三个分支的逻辑：

• 复用分支：world_size 相同时只换 mapping 引用，不重建通信子。理由是底层 NCCL communicator 只取决于 (world_size, rank)，与 pp/tp/ep 的具体布局无关。而重建会触发 ncclCommDestroy，这是一个集合操作，在 worker 复用场景下会在未同步的时间点执行，导致死锁。
• 替换分支：world_size 变了，不得不重建，但至少打一条 warning 把风险暴露出来，而不是静默挂住。
• 新建分支：首次调用，正常建。

这类注释在 TensorRT-LLM 源码里不少，它们记录的是”踩过的坑”——读源码时这类段落的价值往往高于主逻辑，因为它解释了为什么不能写成更简洁的形式。

MULTIMEM 的硬件约束

SymmetricMemoryAllReduce（symm_mem_allreduce.py）的类 docstring 直接把支持范围写了出来：

1
class SymmetricMemoryAllReduce(nn.Module):
2
    """
3
    AllReduce implementation using PyTorch's symmetric memory operations.
4
    This leverages MULTIMEM hardware instructions for faster allreduce operations.
5

6
    Supported configurations (world_size):
7
    - SM 9.0: 4, 6, 8 GPUs
8
    - SM 10.0: 6, 8 GPUs
9
    """
10

11
    # World sizes that support MULTIMEM instructions
12
    _WORLD_SIZES_MULTIMEM = {
13
        "9.0": [4, 6, 8],
14
        "10.0": [6, 8],
15
    }

以及一组按 SM 版本和 world size 区分的大小上限：

1
    MiB = 1024 * 1024
2
    # Maximum buffer sizes for symmetric memory (bytes)
3
    _MAX_SIZES = {
4
        "9.0": {
5
            2: 64 * MiB,   # 64 MB
6
            4: 32 * MiB,   # 32 MB
7
            6: 64 * MiB,
8
            8: 64 * MiB,
9
        },
10
        "10.0": {
11
            2: 8 * MiB,    # 8 MB
12
            4: 32 * MiB,
13
            6: 128 * MiB,
14
            8: 128 * MiB,
15
        },
16
    }

为什么 world_size 只能是 4/6/8？因为 MULTIMEM 是走 NVLink 多播（multicast）硬件的——一次写操作同时落到多个 GPU 的显存里。多播组的规模受限于交换机拓扑（NVSwitch 的端口分组），4/6/8 正好对应几种常见的域划分。超出这个范围，要么没有硬件路径，要么收益不足以覆盖开销，所以直接声明不支持。

_MAX_SIZES 的存在则是因为多播缓冲区是预分配的对称内存：world size 越大，每张卡要预留的窗口越多，总占用按平方级增长。SM 9.0（Hopper）和 SM 10.0（Blackwell）的数值不同，反映的是两代 NVLink 交换能力的差异——Blackwell 在 world_size 8 时能给到 128 MiB，而 Hopper 只有 64 MiB。

workspace 大小怎么算

CustomAllReduceHelper.max_workspace_size_auto（allreduce_helper.py）的 docstring 把 workspace 的用途和容量算法讲得很清楚：

1
@staticmethod
2
def max_workspace_size_auto(tp_size: int, support_deterministic=True) -> int:
3
    """Calculate workspace size for allreduce fusion kernel.
4

5
    The workspace is used for lamport buffers in the fusion kernel.
6
    Required size calculation:
7
    - Each GPU needs 3 sub-buffers (for triple buffering)
8
    - Each sub-buffer stores: max_num_tokens * hidden_size * dtype_size (bf16=2)
9
    - The lamport allocation multiplies by tp_size, so:
10
      lamport_size = 3 * size * tp_size (per GPU)
11

12
    Example: Llama 8B (hidden=4096), max_tokens=8192, bf16, TP=4
13
    - Data per sub-buffer: 8192 * 4096 * 2 = 64 MiB
14
    - Total lamport: 3 * 64MB * 4 = 768 MiB per GPU
15
    - Required 'size' parameter: 64 MiB (gets multiplied by tp_size in allocation)
16
    ...
17
    """

这段注释值得停下来算一遍。Lamport buffer 是自定义 all-reduce kernel 里用于 rank 间同步的机制——每张卡往自己的 buffer 里写、再读别人的 buffer，不需要 NCCL 的 barrier。三个 sub-buffer 是三缓冲（triple buffering），让写下一批和读上一批能重叠。

按注释里的例子：max_num_tokens = 8192、hidden_size = 4096、bf16，单个 sub-buffer 是 $8192 \times 4096 \times 2 = 64$ MiB。乘 3（三缓冲）再乘 tp_size（每张卡要持有所有 rank 的 buffer 才做得了本地归约），一张卡要 768 MiB。这不是小开销——用显存换通信延迟就是自定义 all-reduce 的本质取舍。

实际分配时（allocate_allreduce_fusion_workspace）：

1
@staticmethod
2
def allocate_allreduce_fusion_workspace(
3
    mapping: Mapping, size: int
4
) -> Tuple[List[IpcMemory], "torch.tensor"]:
5
    is_p2p_supported = can_access_peer(mapping)
6
    ipc_buffers_size = size * mapping.tp_size
7
    ipc_buffers = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported)
8
    ipc_barriers = IpcMemory(mapping, 256 * mapping.tp_size, is_p2p_supported)
9
    lamport_buffers_size = size * mapping.tp_size
10
    lamport_buffers = IpcMemory(mapping, 3 * lamport_buffers_size, is_p2p_supported)
11
    if is_p2p_supported:
12
        lamport_initialize(lamport_buffers.local_ptr, 3 * lamport_buffers_size)
13
    # flag_buffer[0], atomic flag read counter
14
    # flag_buffer[1], non-lamport flag
15
    # flag_buffer[2], lamport flag
16
    flag_buffer = torch.tensor([0, 0, 0], dtype=torch.int, device="cuda")
17
    # layout_buffer[0], clear size for next lamport kernel
18
    # layout_buffer[1], triple buffer offset for lamport kernel
19
    layout_buffer = torch.tensor([0, lamport_buffers_size], dtype=torch.int64, device="cuda")
20
    ...

五个组成部分，每个都有明确用途：

1. ipc_buffers：数据缓冲区，大小 size * tp_size。
2. ipc_barriers：同步屏障，每 rank 256 字节。
3. lamport_buffers：三缓冲的 Lamport 区，大小 3 * size * tp_size——与注释里的公式严格对应。
4. flag_buffer：三个 int，注释直接给出了各自的语义（原子读计数、非 lamport 标志、lamport 标志）。
5. layout_buffer：两个 int64，记录”下一轮 lamport kernel 要清多少字节”和”三缓冲当前偏移”。

can_access_peer(mapping) 决定这些 buffer 是否能用 P2P 直访（NVLink/P2P 可用时，GPU 之间可以直接读写对方的显存；不可用时只能走主机内存中转）。lamport_initialize 只在 P2P 可用时调用——Lamport 协议的初始化（把 buffer 填成特定的哨兵值）本身要通过 P2P 完成。

默认大小：

1
        # Default: 64 MiB - supports most common model configurations
2
        # Increase via env var if you see CUDA illegal memory access errors with large models
3
        default_size = 67_108_864  # Exactly 64 MiB
4
        return default_size

64 MiB 恰好对应注释例子里的”hidden=4096 × max_tokens=8192 × bf16”。注释里那句”遇到 CUDA illegal memory access 就调大”是经验之谈——workspace 不够时报的不是”空间不足”，而是内核越界访存，因为 kernel 会按 layout_buffer 里的偏移去写。这个排查线索很实用。

非确定性的开关

1
def force_all_reduce_deterministic():
2
    return (
3
        os.getenv("FORCE_DETERMINISTIC", "0") == "1"
4
        or os.getenv("FORCE_ALL_REDUCE_DETERMINISTIC", "0") == "1"
5
    )

以及它对应的处理：

1
if force_all_reduce_deterministic() and support_deterministic:
2
    workspace_size = os.getenv("FORCE_ALLREDUCE_KERNEL_WORKSPACE_SIZE", "1000000000")
3
    return int(workspace_size)

为什么自定义 all-reduce 会非确定？因为它的归约顺序取决于各 rank 的到达时间——Lamport 协议允许”谁先到谁先写”，浮点加法不满足结合律，顺序不同结果就不同。开启确定性模式后，用更大的 workspace（默认 1 GB）换取固定顺序的归约。这条开关在需要复现实验或对拍数值时会用到。

融合算子：把 all-reduce、残差、归一化、量化压进一次访存

标准 Transformer 的残差块在 TP 下长这样：

1
h = h + attn_out          # 残差相加
2
h = rms_norm(h)           # 归一化

而 attn_out 是行并行层的输出，需要先 all-reduce。串起来就是三件事：all-reduce → 加残差 → RMSNorm。如果逐个算子做，数据要进出显存三趟：all-reduce 写一遍、加残差读写一遍、归一化读写一遍。以 hidden=8192、bf16、batch 1024 token 计，每次就是 $1024 \times 8192 \times 2 = 16$ MiB 的读写，乘以三趟和层数就是可观的开销。

TensorRT-LLM 的做法是把它们融成一个 kernel。融合模式的枚举定义在 functional.py：

1
class AllReduceFusionOp(IntEnum):
2
    NONE = 0
3
    RESIDUAL_RMS_NORM = 1
4
    LAST_PROCESS_FOR_UB = 2
5
    RESIDUAL_RMS_PREPOST_NORM = 3
6
    RESIDUAL_RMS_NORM_QUANT_FP8 = 4
7
    RESIDUAL_RMS_NORM_QUANT_NVFP4 = 5
8
    RESIDUAL_RMS_NORM_OUT_QUANT_FP8 = 6
9
    RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7
10
    MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM = 8
11
    RMS_NORM = 9

按用途分四组理解：

1. NONE：纯 all-reduce。
2. 归一化类：RMS_NORM（只归一化，不做残差）、RESIDUAL_RMS_NORM（加残差再归一化）、RESIDUAL_RMS_PREPOST_NORM（pre/post 双归一化，某些模型结构需要）。
3. 量化类：带 QUANT_FP8 / QUANT_NVFP4 后缀表示归一化之后直接量化输出。OUT_QUANT_* 与不带 OUT 的区别在于输出的是”归一化结果 + 量化结果”还是只有量化结果——AllReduce.forward 的 docstring 把这组返回值的顺序写得很明确：

1
NONE: [hidden_states]
2
RMS_NORM: [hidden_states]
3
RESIDUAL_RMS_NORM: [hidden_states, residual]
4
RESIDUAL_RMS_NORM_QUANT_FP8: [norm_quant, residual]
5
RESIDUAL_RMS_NORM_OUT_QUANT_FP8: [norm, norm_quant, residual]
6
RESIDUAL_RMS_NORM_QUANT_NVFP4: [norm_quant_fp4, scale_factor, residual]
7
RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4: [norm, norm_quant_fp4, scale_factor, residual]

这个”返回几个张量、顺序是什么”的表格是接这条路径时最容易出错的地方。特别是 NVFP4 那条：多出一个 scale_factor——因为 NVFP4 是块量化格式，每 16 个元素共享一个 FP8 缩放因子，缩放因子必须跟着数据一起产出。

4. MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM：MoE 专用，把”专家输出归约 + 共享专家相加 + all-reduce + 残差 + 归一化”五件事合成一个 kernel。这是整个融合体系里最激进的一条，下面单独讲。

融合的载体是 AllReduceParams：

1
class AllReduceParams:
2

3
    def __init__(
4
        self,
5
        strategy: AllReduceStrategy = AllReduceStrategy.AUTO,
6
        fusion_op: AllReduceFusionOp = AllReduceFusionOp.NONE,
7
        bias: Optional[Tensor] = None,
8
        residual: Optional[Tensor] = None,
9
        norm_weight: Optional[Tensor] = None,
10
        scale: Optional[Tensor] = None,
11
        norm_pre_residual_weight: Optional[Tensor] = None,
12
        eps: float = 1e-06,
13
        enable_allreduce: bool = True,
14
        trigger_completion_at_end: bool = True,
15
    ):
16
        ...
17
        assert fusion_op in (AllReduceFusionOp.NONE.value,
18
                             AllReduceFusionOp.RMS_NORM.value) or (residual
19
                                                                   is not None)

最后那个 assert 是一条契约：只要融合模式涉及残差，就必须传 residual。原因很直观——残差加法的输入不给，kernel 就算不出来。把它写成 assert 而不是运行期检查，是因为这个错误只可能来自模型定义代码写错，属于开发期就该暴露的问题。

还有三个小工具方法：

1
    def has_affine(self):
2
        return 1 if self.norm_weight is not None else 0
3

4
    def has_bias(self):
5
        return 1 if self.bias is not None else 0
6

7
    def has_scale(self):
8
        return 1 if self.scale is not None else 0

它们把”参数是否为空”编码成 kernel 的模板参数或标志位。这样同一个 kernel 二进制可以覆盖”有/无 bias”、“有/无 scale”两种变体，不需要各编译一份。

AllReduce.forward：策略是分层回退的

forward 的执行路径是一条逐层回退链，值得按顺序读：

1
def forward(self, input, *, all_reduce_params=None):
2
    if self.mapping.tp_size == 1 or (all_reduce_params is not None
3
                                     and all_reduce_params.enable_allreduce == False):
4
        return input
5

6
    input = input.contiguous()  # Underlying op requires contiguous input
7

8
    allreduce_strategy = self.strategy
9

10
    if all_reduce_params is None:
11
        all_reduce_params = AllReduceParams()
12

13
    # Try Symmetric Memory AllReduce first if available
14
    if self.symm_mem_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE:
15
        symm_mem_output = self.symm_mem_allreduce(input)
16
        if symm_mem_output is not None:
17
            return symm_mem_output
18
    elif self.symm_mem_allreduce and all_reduce_params.fusion_op != AllReduceFusionOp.NONE:
19
        logger.debug_once(
20
            f"Skipping SymmetricMemoryAllReduce for fused operation (fusion_op={all_reduce_params.fusion_op}), using regular allreduce",
21
            key=(self.mapping.tp_rank, all_reduce_params.fusion_op, "debug_fusion_skip"),
22
        )
23

24
    # Try MNNVL AllReduce if symm_mem didn't handle it
25
    if self.mnnvl_allreduce:
26
        mnnvl_output = self.mnnvl_allreduce(input, all_reduce_params=all_reduce_params)
27
        if mnnvl_output is not None:
28
            return mnnvl_output
29

30
    # Fall back to regular AllReduce
31
    if allreduce_strategy in (AllReduceStrategy.MNNVL, AllReduceStrategy.SYMM_MEM):
32
        allreduce_strategy = AllReduceStrategy.AUTO
33
    ...

四条关键信息：

第一，短路条件有两个。tp_size == 1 是显然的；enable_allreduce == False 则是给上层用的开关——比如 HELIX CP 的注意力输出投影会传 AllReduceParams(enable_allreduce=False)，因为在那种布局下它要改用 reduce-scatter，而不是 all-reduce。这比”用两个不同的类”要轻。

第二，contiguous() 不是可有可无的。底层 custom op 按连续内存做向量化访问，非连续输入会直接算错或者报错。注释写明了是为了”满足底层 op 的要求”。

第三，SYMM_MEM 只覆盖 NONE，遇到融合就跳过。原因是 MULTIMEM 是硬件多播指令，只做”把数据同时写到多个 GPU 并求和”这一件事，没法在里面插一段 RMSNorm。代码用 debug_once 记录跳过原因，而不是每次调用都打日志——这种”只打一次”的日志模式在热路径代码里很常见。

第四，策略名会被”改写”：

1
if allreduce_strategy in (AllReduceStrategy.MNNVL, AllReduceStrategy.SYMM_MEM):
2
    allreduce_strategy = AllReduceStrategy.AUTO

注释解释得很直接：底层的 allreduceOp 里没有 MNNVL/SYMM_MEM 的分支，所以走到普通路径时必须先把策略名换掉，否则会被判为非法策略。这是一类很典型的”两层策略命名空间不重合一”的问题——上层策略是给”选择逻辑”用的，下层策略是给”kernel 分发”用的。

最后是分发：

1
    if allreduce_strategy == AllReduceStrategy.AUTO and not disable_allreduce_autotune and not self._disable_mpi:
2
        if _ALLREDUCE_AUTOTUNER_TUNING_MODE or not _ALLREDUCE_NATIVE_AUTOTUNER_ENABLED:
3
            output = torch.ops.trtllm.tunable_allreduce(...)
4
        else:
5
            output = torch.ops.trtllm.autotuned_allreduce(...)
6
    else:
7
        output = self.all_reduce_op(...)

AUTO 策略下会在开跑时自动调优：tunable_allreduce 会在首个 batch 上试跑几种 kernel 变体，选最快的那个，然后把结果缓存下来供后续批次直接查表（autotuned_allredcüce 走的是查表路径）。这解释了为什么 TensorRT-LLM 的第一次推理往往比后续慢一点——首轮包含 autotune 的开销。

环境变量 TLLM_DISABLE_ALLREDUCE_AUTOTUNE=1 可以关掉这个自动调优（注释里写着是因为曾经观察到性能回退，留了后门）。

GEMM 与 all-reduce 的零拷贝融合

前面提到的 use_fused_gemm_allreduce/NCCL 窗口两条路径，开关条件写在 Linear._refresh_fused_gemm_allreduce 里：

1
enabled = all([
2
    self.reduce_output,
3
    not mpi_disabled(),
4
    self.dtype in (torch.float16, torch.bfloat16),
5
    self.in_features % 128 == 0,
6
    self.out_features % 64 == 0,
7
    self.tp_mode == TensorParallelMode.ROW,
8
    self.tp_size > 1,
9
    quant_valid,
10
    get_sm_version() >= 100,
11
    self._enable_gemm_allreduce_fusion,
12
    os.environ.get("TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "0") == "1",
13
])
14
self.use_fused_gemm_allreduce = enabled and ipc_nvls_supported()

十一个条件逐个看，就是一份”这条路径为什么难做”的清单：

• reduce_output + ROW + tp_size > 1：必须是行并行且需要归约的层。
• dtype in (fp16, bf16)：低精度才有融合的意义（FP32 通信量太大，融合收益不明显）。
• in_features % 128 == 0、out_features % 64 == 0：kernel 的 tile 边界要求。
• quant_valid：要求量化配置里激活也是 NVFP4/FP8 量化的（quant_config_has_nvfp4_activation_quantization）。因为融合路径最终要把 GEMM 的输出直接写进通信缓冲，如果后续还要做量化，必须一并考虑。
• get_sm_version() >= 100：Blackwell 及以上。这条路径依赖新硬件的特性。
• ipc_nvls_supported()：需要有 NVLS（NVLink SHARP）支持，因为要跨卡直接写多播缓冲。
• 一个环境变量开关 TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED 默认关闭。

默认关闭这一点很值得注意：这类”极致优化”的路径在成熟产品里通常先以 opt-in 形式发布，经过充分验证再默认打开。读源码时看到 os.environ.get(...) == "1" 这类默认关闭的开关，基本可以判断这是一条实验性/高风险的路径。

对应的 forward 分支也有注释解释为什么某些情况必须绕开：

1
if self.tp_mode == TensorParallelMode.ROW:
2
    use_locality_domain = self.partition_plan.enabled
3
    # The fused GEMM-allreduce and NCCL-window GEMM paths consume the
4
    # module's full weight tensors. locality domain releases those tensors after
5
    # creating localized shards, so ROW-TP must first run the shard-aware GEMM
6
    # and then all-reduce its partial output.
7
    use_fused_gemm_allreduce = (not use_locality_domain
8
                                and self.use_fused_gemm_allreduce
9
                                and lora_params is None)

两个互斥条件：locality domain（把权重按局部性重新分片的技术，会释放完整权重张量）和 LoRA（需要额外注入低秩分支）都会破坏融合路径的前提，所以直接短路。

关于 gather_output 的 sizes：

1
self.gather_output_sizes = None
2
if self.gather_output and self.tp_size > 1 \
3
        and self.tp_mode == TensorParallelMode.COLUMN:
4
    if self._tp_sharding_is_auto:
5
        sizes = [
6
            end - start for start, end in (
7
                ...
8
            )
9
        ]

当启用不均匀 TP（uneven TP）时，输出维不能被 tp_size 整除，每张卡的宽度不同。此时 all-gather 必须显式传 sizes，否则按等长切片拼接会错位。同一逻辑在 allgather 的 docstring 里也有：

1
If 'sizes' is 'None', the input tensors in the different ranks must have the same shape.
2
Otherwise, 'sizes[i]' must be 'input.shape[dim]' at rank i, and the input tensors in
3
the different ranks can only differ in shape at dimension `dim`.

MoEAllReduce：五算子融合