SGLang 源码完全拆解（十七）：DP Attention 与专家并行——注意力数据并行、DeepEP 通信与 EPLB 负载均衡

TP 下 KV Cache 的复制问题

设 TP 大小为 8。注意力层的权重是切分的：Q/K/V 的投影矩阵按输出维切成 8 份，每张卡算自己那一份 head 的输出。算完之后要做一次 all-reduce，因为输出投影 W_O 需要看到全部 head 的结果。

问题的关键在于：每张卡要自己算自己的那些 head，就必须自己读自己那些 head 的 K 和 V。而 K 和 V 是历史所有 token 的，来自上一层。于是每张卡都必须把完整的 KV Cache 存在本地——否则它算不出自己那份 head 的注意力。

一台 8 卡机器跑一个 70B 模型、128K 上下文：按 GQA 8 个 KV 头、头维 128、80 层、FP16 算，单个请求的 KV 是

如果每张卡都存一份，8 卡一共吃掉 344 GB 显存——而真正需要的数据只有 43 GB。换句话说，TP 把 KV Cache 放大了 TP 倍。这在短上下文、小模型时不明显，在长上下文场景下是致命的：它直接把可并发的请求数除以 8。

EP 下专家权重切不均匀的问题

MoE 层反过来：如果按 TP 切专家权重，每个专家都要被切开。DeepSeek-V3 每个专家的 gate_proj 是 [2048, 7168]，up_proj 同尺寸，down_proj 是 [7168, 2048]。按 TP=16 切，中间维 2048 被切成 128 行——GEMM 的 N 维只有 128，在 Hopper Tensor Core 上远低于最小有效 tile，kernel 要么退化要么直接不支持。

更合理的做法是不切专家，而是把 256 个专家分给各个 rank：EP=16 时每 rank 拿 16 个专家，专家内部的 GEMM 保持完整尺寸。代价是：token 必须被送到”持有它要用的专家”的那个 rank 上去。这就是 all-to-all。

所以一个自然的结论：attention 适合 DP（避免 KV 复制），MoE 适合 EP（避免专家切碎）。而这两者在同一层里必须拼接，就产生了一个”进出各一次全局通信”的结构。

图片来源：SGLang 大规模 EP 设计博客 的整体架构图

图里值得注意的是两件事：一是 attention 部分被画成了若干个并排的副本（DP），二是 MoE 部分被画成一个横跨所有 GPU 的专家池（EP）。中间的箭头就是本篇要拆的 gather/scatter 与 A2A。

名称上的一个坑：dp_size 与 attn_dp_size

SGLang 的并行度参数在过去一年里被重命名过一次，读源码时必须分清：

• --dp-size 是真正的数据并行副本数，每个副本有独立的模型权重副本、独立调度器。这是”多副本服务”的语义。
• --attn-dp-size 是TP 组内的注意力数据并行宽度：在同一个 TP 组内部，把 token 分给若干个 attention-DP 副本，attention 是数据并行的，FFN 仍然是张量并行的。这是”混合”的语义。
• --enable-dp-attention 是已废弃的写法。在固定提交 d88d5c3 上，它已经被解析成 attn_dp_size：server_args.py 里明确写着解析逻辑是 resolved["enable_dp_attention"] = resolving_view(self).attn_dp_size > 1，并且注释标注 TODO: remove after 2026-12-31。CLI 上仍保留这个 flag，但它实际报的是”是否配置了 attention-DP 宽度”。

参数文档里对 attn_dp_size 的定义是：attention 数据并行尺寸，即 TP 组内部的数据并行注意力组数，每组一个调度器，而 FFN 保持张量并行。这句话是整个设计的浓缩版，后面所有机制都在兑现它。

于是 tp_size 被拆成两个正交的乘积：

其中 attn_cp_size 是注意力的上下文并行（把序列维切开做长上下文 prefill）。attn_tp_size 由 derive_attn_tp_size 导出：

1
def derive_attn_tp_size(*, tp_size: int, attn_cp_size: int, attn_dp_size: int) -> int:
2
    return tp_size // attn_dp_size // attn_cp_size

MoE 侧是另一套乘积：

两套分解同时成立、共用同一批物理 rank，只是分组方式不同。initialize_model_parallel 的文档注释给了一个具体例子：8 张卡，attn_cp_size=2、attn_tp_size=4，同时 moe_dp_size=2、moe_ep_size=4，那么：

• attention context-parallel 组：[g0,g4] [g1,g5] [g2,g6] [g3,g7]
• attention tensor-parallel 组：[g0,g1,g2,g3] [g4,g5,g6,g7]
• MoE expert-parallel 组：[g0,g1,g2,g3] [g4,g5,g6,g7]
• MoE data-parallel 组：[g0,g4] [g1,g5] [g2,g6] [g3,g7]

同一个 rank 同时属于四个不同名字的组。这是理解后面所有 gather/scatter 的坐标系——通信到底在哪个组上做，取决于这一维是谁切的。

全局缓冲区的尺寸与语义

_DpGatheredBufferWrapper 是一个状态门面，把”全局 gather 缓冲区”的元数据（hidden size、dtype、device）挂在 flags.dp 上，把每次前向的尺寸四元组保留为类属性：

1
_global_dp_buffer_len: int = 0
2
_local_dp_buffer_len: int = 0
3
_dp_max_padding: bool = False
4
_global_num_tokens: Optional[List[int]] = None
5
_global_num_tokens_gpu: Optional[torch.Tensor] = None

源码注释解释了为什么尺寸要放类属性而不是 contextvar：这些值会在被 torch.compile 追踪的模型代码里读到，而属性来源的 int 会被 dynamo 当作自动动态（automatic dynamic），contextvar 则无法追踪，dict 槽位还会触发 recompile 上限——每个不同的 size 都重编译一次。这是一个典型的性能工程细节：为了不让编译器为每种 token 数重新编译一遍图，把这些标量放在”编译器看得懂”的位置上。

global_num_tokens 是一个长度为 DP 副本数的列表：第 i 个元素是第 i 个 DP 副本这一轮有多少 token。global_num_tokens_gpu 是它在 GPU 上的张量版本，用来在 kernel 里算偏移。

本 rank 的切片位置怎么算

get_dp_local_info 回答一个问题：在拼接后的全局 buffer 里，我这一段从哪一行开始、有多长？

1
def get_dp_local_info(forward_batch: ForwardBatch) -> Tuple[torch.Tensor, torch.Tensor]:
2
    if forward_batch.dp_local_start_pos is None:
3
        dp_rank = dp_slot_in(forward_batch.global_num_tokens_gpu)
4
        cumtokens = torch.cumsum(forward_batch.global_num_tokens_gpu, dim=0)
5
        if dp_rank == 0:
6
            local_start_pos = torch.zeros_like(cumtokens[0])
7
        else:
8
            local_start_pos = cumtokens[dp_rank - 1]
9
        local_num_tokens = forward_batch.global_num_tokens_gpu[dp_rank]
10
        forward_batch.dp_local_start_pos = local_start_pos
11
        forward_batch.dp_local_num_tokens = local_num_tokens
12
    return forward_batch.dp_local_start_pos, forward_batch.dp_local_num_tokens

先 cumsum 得到前缀和，rank 0 从 0 开始，rank i 从”前面所有 rank 的 token 数之和”开始。这就是把各副本的 token 顺序拼成一个全局序列的下标计算。注意这是在 GPU 上算的（张量运算），结果缓存在 ForwardBatch 上避免重复计算。

两种 gather：all_reduce 与 all_gather

SGLang 为 gather 写了两条完全不同的实现路径，选哪条取决于拓扑。

路径一：_dp_gather_via_all_reduce。 思路是”反向 all-reduce”：每个 rank 先把全局 buffer 清零，然后把自己那一段写进对应位置，再做一次 all-reduce，各段的贡献加起来正好拼出完整 buffer。

1
    global_tokens.fill_(0)
2
    ...
3
    writes = (is_partial or get_parallel().attn_tp_rank == 0) and (
4
        cp_shard_counts is not None or get_parallel().attn_cp_rank == 0
5
    )
6
    if local_tokens.shape[0] > 0 and writes:
7
        ...
8
        memcpy(global_tokens, local_tokens, 0, local_start_pos, local_num_tokens, False)

这里有个关键设计：只有 attn_tp_rank == 0 才写。因为 attention-TP 组内的各 rank 持有的是同一批 token 的切片（sequence parallel 或 head 切分），只有一个 rank 该负责写入自己的 DP 段，否则同一份数据会被写多次、all-reduce 之后数值翻倍。all-reduce 完成后，所有 rank 都拿到了完整的全局 buffer——这正是 EP 需要的（每个 rank 都要看到全部 token 才能做正确的路由）。

注释还提到一个细节：输入 ID 是 int32 类型，所以本地情形下要用 inplace_all_reduce（custom all-reduce 的路径），

1
        NUM_GPUS_PER_NODE = 8
2
        if (
3
            not local_tokens.dtype.is_floating_point
4
            and get_parallel().tp_size <= NUM_GPUS_PER_NODE
5
        ):
6
            from sglang.srt.distributed.parallel_state import inplace_all_reduce
7

8
            inplace_all_reduce(
9
                global_tokens, group_name=get_parallel().tp_group.unique_name
10
            )

整数类型的 all-reduce 走 NCCL 的原地路径更快，且只在单机 8 卡以内有效——跨机时自定义 all-reduce（通常基于 NVLink + P2P 实现）不再适用。

路径二：_dp_gather_via_all_gather。 如果 attn_tp_size == 1，说明每个 DP 副本就是一个 rank，直接 all-gather 即可：

1
    if get_parallel().attn_tp_size == 1:
2
        if use_world:
3
            all_gather_single(global_tokens, local_tokens, group=torch.distributed.group.WORLD)
4
        else:
5
            get_parallel().tp_group.all_gather_into_tensor(global_tokens, local_tokens)
6
        return

如果 attn_tp_size > 1，则要先在 attention-TP 组内把切片拼回本 DP 副本的完整 token，再跨 DP 做 gather，代码用 reduce_scatter_tensor 与前缀拼接处理：

1
    if not is_partial:
2
        if get_parallel().attn_tp_rank != 0:
3
            local_tokens.fill_(0)
4
    scattered_local_tokens = local_tokens.tensor_split(get_parallel().attn_tp_size)[
5
        get_parallel().attn_tp_rank
6
    ]
7
    get_parallel().attn_tp_group.reduce_scatter_tensor(
8
        scattered_local_tokens, local_tokens
9
    )

这里出现了一个重要概念：is_partial。它表示”这批 hidden state 本身还是 attention-TP 意义上的部分和”，即各 attn_tp rank 各持一部分贡献，需要先 reduce 才能用。这对应着一层里 attention 输出还没做 all-reduce 的情形（融合了 reduce 的路径，后面会讲）。

DpPaddingMode：补齐到最长还是补齐到总和

既然要拼成一个全局 buffer，就有一个不可避免的问题：各 DP 副本本轮 token 数不一样。副本 A 有 3000 个 token，副本 B 只有 200 个，拼起来没问题（用 cumsum 定位即可），但集合通信往往要求各 rank 的输入形状一致。SGLang 因此定义了两个 padding 模式：

1
class DpPaddingMode(IntEnum):
2
    # Padding tokens to max length and then gather tokens using `all_gather_into_tensor`
3
    MAX_LEN = auto()
4
    # Padding tokens to sum length and then gather tokens using `all_reduce`
5
    SUM_LEN = auto()

• MAX_LEN：每个副本都补齐到 max(global_num_tokens)，然后用 all_gather_into_tensor（要求等长）。
• SUM_LEN：每个副本补齐到”全体之和”，然后用 all_reduce（all-reduce 对不等长的分段求和）。

选择逻辑在 get_dp_padding_mode，这是一段非常值得读的启发式：

1
        if moe_a2a_backend.is_pplx():
2
            return DpPaddingMode.MAX_LEN
3

4
        if moe_a2a_backend.is_deepep_v2() and envs.SGLANG_DEEPEP_V2_FORCE_MAX_LEN.get():
5
            return DpPaddingMode.MAX_LEN
6

7
        if is_extend_in_batch and gather_width > 1:
8
            if get_flags().dp.max_len_with_idle and min(global_num_tokens) == 0:
9
                return DpPaddingMode.MAX_LEN
10
            return DpPaddingMode.SUM_LEN
11

12
        # we choose the mode that minimizes the communication cost
13
        max_len = max(global_num_tokens)
14
        sum_len = sum(global_num_tokens)
15
        if sum_len * 2 >= max_len * gather_width:
16
            return cls.MAX_LEN
17
        else:
18
            return cls.SUM_LEN

逐条解释这四个分支：

1. pplx 后端强制 MAX_LEN。源码注释给了原因：pplx-kernels 的 a2a 是一个对称集合通信，每个 EP rank 必须 dispatch 相同数量的 token，否则设备侧握手死锁——空闲的 DP rank（0 token）永远不发信号，对端就永远等。这是一个真实的死锁，不是性能问题。
2. DeepEP v2 强制 MAX_LEN（可通过环境变量开关）。同样是按阶段选择的通信拓扑对形状对称性的要求。
3. prefill 阶段（is_extend_in_batch）且 gather 跨多个 DP rank 时用 SUM_LEN。注释说得很直接：避免 token 分布不均带来的 padding 开销。prefill 阶段各副本的 token 数差异可以极大（有人在做 4K 的 prompt，有人在解码），补到最长意味着大量无效行白算。例外是”混合 SSM 模型”（max_len_with_idle）且某个副本恰好为空时需要 MAX_LEN，源码注释说这是因为这类模型要用”fabricated-row conversion”来实体化空闲 rank。
4. 默认走代价最小化：sum_len * 2 >= max_len * gather_width 时选 MAX_LEN。这个不等式的含义是——如果按 MAX_LEN 补，通信总量是 max_len × width；如果按 SUM_LEN，通信总量是 sum_len。当 sum_len 超过 max_len × width 的一半时，两者相近，而 MAX_LEN 还能顺带启用对称内存优化（use_symmetric_memory，后面提），因此倾向 MAX_LEN。

在 CUDA Graph 场景下还有默认模式：

1
    @classmethod
2
    def get_default_mode_in_cuda_graph(cls) -> DpPaddingMode:
3
        if _USE_ROCM700A_WA:
4
            return cls.SUM_LEN
5
        else:
6
            return cls.MAX_LEN

CUDA Graph 要求形状固定，所以必须选一个确定的 padding 模式；MAX_LEN 在等长情形下更容易做出稳定的图。

scatter：把 MoE 的输出还回本副本

相反方向的操作在 layers/layer_boundary/ops.py。_dp_scatter_step 是 dispatch 表里的一项，真正的实现在 dp_scatter：

1
def dp_reduce_scatterv(
2
    local_hidden_states: torch.Tensor,
3
    hidden_states: torch.Tensor,
4
    forward_batch: ForwardBatch,
5
) -> None:
6
    get_parallel().tp_group.reduce_scatterv(
7
        hidden_states,
8
        output=local_hidden_states,
9
        sizes=get_dp_global_num_tokens(),
10
    )

reduce_scatterv 是”不等长版本的 reduce-scatter”：按 sizes 指定的切分点，把输入切给各 rank，同时做求和。这正好是 gather 的逆操作——gather 时用 all-reduce 拼，scatter 时就用 reduce-scatterv 拆，且 sizes 就是 global_num_tokens。一致性由此保证：同一份 global_num_tokens 既决定拼的偏移，也决定拆的长度。

还有一个值得看的群体选择函数：

1
def _dp_scatter_group() -> GroupCoordinator:
2
    parallel = get_parallel()
3
    if parallel.tp_size == parallel.attn_dp_size:
4
        return parallel.tp_group
5
    return parallel.attn_tp_group

当 tp_size == attn_dp_size 时，attn_tp_size 必然为 1，每一个 DP 副本就是一个 rank（没有 attention-TP 切片），这时用整个 TP 组做 scatter；否则用 attention-TP 组（因为切片要先在组内合并）。通信组的选取完全由”这一维是否真的被切过”决定——这是 layer_boundary 那套抽象的核心思想。

每个副本为什么必须有调度器

DP Attention 的一个硬性后果是：每个 attention-DP 副本都需要自己的调度器进程。因为每个副本的 batch 是独立的——副本 A 的请求队列和副本 B 完全不同，KV Cache 也是各自独立的池。所以 SGLang 在 attn_dp_size > 1 时的进程拓扑是：

这个数字在参数文档里被明确写出：num_dp_ranks 的定义就是”这个部署用多少个数据并行 rank 来服务，每个一个调度器：dp_size * attn_dp_size”。

多个调度器之间需要同步，是因为 MoE 的 EP 通信跨越所有 rank：如果各调度器各自决定”我这轮跑多少 token”，A2A 的 buffer 对不上，通信就会挂住。所以每轮前向之前，各调度器要交换一轮元数据。

11 个字段的元数据

MLPSyncBatchInfo 把要同步的东西打包成一个固定长度的张量：

1
    def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
2
        return torch.tensor(
3
            [
4
                self.num_tokens,
5
                self.num_tokens_for_logprob,
6
                int(self.can_run_decode_cuda_graph),
7
                int(self.is_extend_in_batch),
8
                int(self.local_can_run_tbo),
9
                self.local_forward_mode,
10
                int(self.can_run_prefill_cuda_graph),
11
                self.prefill_cuda_graph_max_prefix_len,
12
                int(self.can_run_draft_cuda_graph),
13
                int(self.is_health_check),
14
                self.attention_pairs,
15
            ],
16
            device=device,
17
            dtype=dtype,
18
        )

这 11 个字段不是随便选的，它们覆盖了”跨副本必须一致”的每一件事：

• num_tokens / num_tokens_for_logprob：本副本的 token 数，用来算全局偏移与 padding 目标。
• is_extend_in_batch：这一轮全局是不是 prefill 批次。这个信息决定 DeepEP 用 normal 还是 low-latency（后面讲），也决定 padding 模式。
• can_run_decode_cuda_graph / can_run_prefill_cuda_graph / can_run_draft_cuda_graph：本副本是否能用 CUDA Graph。这是一个必须全局一致的决策——如果用图的 rank 和不用图的 rank 混在一起，通信形状就可能对不上。
• local_can_run_tbo：本副本能否参与 two-batch overlap。
• local_forward_mode：本副本的 forward mode（decode / extend / idle / …），用来推导”全局 forward mode”。
• prefill_cuda_graph_max_prefix_len：prefill 图的形状分桶依据。
• is_health_check：这一批是否全是一次健康检查探针。健康检查要特别对待，因为它出现在系统本就空闲的时刻，把它算进负载不均衡统计会得到假象。
• attention_pairs：这一轮本副本处理的因果 (query, key) 对数量。这是个比 token 数更能反映真实计算量的量——decode 时每个请求 1 个 query 但要看全部历史 KV，token 数完全反映不出负载。

一个 D2H 而不是十一个

同步的实现是 all-gather 全部副本的这 11 个 int64：

1
        global_info_tensor = fallback_tensor.repeat(
2
            self.num_dp_ranks, self.tp_size * self.cp_size, 1
3
        )

形状是 [num_dp_ranks, tp_size * cp_size, 11]。tp_size * cp_size 这一维存在是因为 attention-TP 组内的各 rank 也需要参与（它们构成 DP 副本的切片）。每个槽位先填 fallback 值，再把活跃 rank 的真实值写进去：

1
        if device == "cpu":
2
            tp_active_ranks = get_parallel().tp_group.active_ranks_cpu
3
        else:
4
            tp_active_ranks = get_parallel().tp_group.active_ranks
5
        if tp_active_ranks.shape[0] < num_ranks_in_tp_info:
6
            tp_active_ranks = torch.ones(
7
                num_ranks_in_tp_info,
8
                dtype=tp_active_ranks.dtype,
9
                device=tp_active_ranks.device,
10
            )
11
        tp_info[tp_active_ranks[:num_ranks_in_tp_info] == 0] = fallback_tensor

active_ranks 是弹性 EP 场景下的”当前有哪些 rank 活着”。不活跃的槽位会被强制回填成 fallback，避免读到垃圾数据。

最后一步是把结果搬到 CPU。这里有一个被显式注释的性能点：

1
        # One D2H for every field: each `.item()` / `.tolist()` on a device
2
        # tensor is its own stream sync. Copy the whole tensor, not the
3
        # `[:, 0, :]` slice -- that slice is non-contiguous once
4
        # attn_tp * attn_cp > 1, adding a gather kernel inside the wait.
5
        tp0_info_cpu = global_info_tensor.cpu()[:, 0, :]

注释说了两件事：第一，每个 .item() 都是一次独立的流同步，所以要把整块搬回来一次性 .tolist()；第二，不能直接拷 [:, 0, :] 这个切片，因为当 attn_tp * attn_cp > 1 时它在内存上不连续，.cpu() 会额外插一个 gather kernel 到等待路径里——直接把整块拷回来再索引反而更快。这类”看着多做了、实际更快”的取舍，是推理系统代码里很常见的东西。

只取 [:, 0, :]（tp rank 0 那一列）是因为同一 DP 副本的各 attn_tp rank 持有相同元数据，取一个代表即可。

从”本副本的 vote”到”全局的决策”

元数据回来之后，各个字段用不同的归约：

1
        self.global_num_tokens = tp0_info_cpu[:, 0].tolist()
2
        self.global_num_tokens_for_logprob = tp0_info_cpu[:, 1].tolist()
3
        self.can_run_decode_cuda_graph = bool(tp0_info_cpu[:, 2].min())
4
        self.is_extend_in_batch = bool(tp0_info_cpu[:, 3].max())
5
        self.can_run_prefill_cuda_graph = bool(tp0_info_cpu[:, 6].min())
6
        self.prefill_cuda_graph_max_prefix_len = int(tp0_info_cpu[:, 7].max())
7
        self.can_run_draft_cuda_graph = bool(tp0_info_cpu[:, 8].min())
8
        self.any_health_check = bool(tp0_info_cpu[:, 9].max())
9
        self.global_attention_pairs = tp0_info_cpu[:, 10].tolist()

• token 数和 attention_pairs 要”全部保留”（list），因为后面要用 cumsum 定位每个人的段。
• CUDA Graph 资格用 min()：只要有一个副本不能用图，全体都不能用。这叫”一票否决”，理由是形状必须全局对齐。
• is_extend_in_batch 用 max()：只要有任何一个副本有 prefill，就算这一轮是 prefill 轮。这是为了让 DeepEP 选择 normal dispatch——所有副本必须用同一种 dispatch 模式。
• prefill_cuda_graph_max_prefix_len 用 max()：取最长的前缀长度，图的分桶要覆盖最坏情况。
• 健康检查用 max()：任何一个副本在跑探针，全轮都要特殊标记。

注意投票函数的单侧语义——_local_decode_cuda_graph_vote 里空闲和 prebuilt 批次都被认为是”可以用图”（返回 True）：

1
def _local_decode_cuda_graph_vote(
2
    *,
3
    local_batch: Optional[ScheduleBatch],
4
    disable_cuda_graph: bool,
5
) -> bool:
6
    """This rank's vote for the decode graph (min-reduced across dp ranks)."""
7
    if disable_cuda_graph:
8
        return False
9
    return (
10
        local_batch is None
11
        or local_batch.forward_mode.is_decode_or_idle()
12
        or local_batch.forward_mode.is_prebuilt()
13
    )

这很重要：空闲的副本不应该否决别人的图，它自己没有形状要求。

空闲副本也必须参与

一个容易被忽略的事实是：即使某个 DP 副本本轮一个请求都没有，它也必须参加这次 gather。原因就是前面 pplx 那段注释说的对称性——集合通信要求所有 rank 到场。为此 SGLang 会为它构造一个 IDLE 批次：

1
    if skip_all_gather:
2
        # Skip idle batch when attn-dp=1 (and always under DWDP: ranks run independently)
3
        need_idle_batch = not dwdp and num_dp_ranks > 1
4
    else:
5
        need_idle_batch = max(mlp_sync_info.global_num_tokens) > 0

如果一轮里所有副本都是 0 token（整个集群完全空闲），那就没必要造 idle 批次了——max(...) > 0 这个条件就是在剪掉这种情形。反过来只要有任何一个副本有活，其他空的副本都要造一个 idle 批次陪着走完这轮通信。

fallback 张量里的 forward mode 就是 IDLE：

1
    def _get_fallback_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
2
        return torch.tensor(
3
            [
4
                0,  # num_tokens
5
                0,  # num_tokens_for_logprob
6
                1,  # can_run_decode_cuda_graph
7
                0,  # is_extend_in_batch
8
                1,  # local_can_run_tbo
9
                ForwardMode.IDLE.value,  # local_forward_mode
10
                0,  # can_run_prefill_cuda_graph
11
                0,  # prefill_cuda_graph_max_prefix_len
12
                1,  # can_run_draft_cuda_graph
13
                0,  # is_health_check
14
                0,  # attention_pairs
15
            ],
16
            device=device,
17
            dtype=dtype,
18
        )

注意空闲 rank 的 can_run_decode_cuda_graph 和 can_run_draft_cuda_graph 都是 1——同样的道理，它没有形状要求，不应该拖累别人。

should_skip_scheduler_all_gather：什么时候可以省掉这次通信

1
def should_skip_scheduler_all_gather(num_dp_ranks: int) -> bool:
2
    """Return whether scheduler metadata is already local and rank-invariant.
3

4
    With one attention-DP rank there is no cross-DP state to reconcile.  The
5
    TP schedulers consume the same broadcast request stream, so gathering the
6
    identical batch mode, graph eligibility, and token counts only adds a
7
    device collective plus host synchronization.  Preserve the environment
8
    override for deployments that explicitly guarantee this invariant beyond
9
    DP1.
10
    """
11

12
    return num_dp_ranks == 1 or envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get()

当 num_dp_ranks == 1 时，没有跨 DP 的状态要协调，这次 all-gather 纯属浪费（一次设备集合通信加一次主机同步）。环境变量 SGLANG_SCHEDULER_SKIP_ALL_GATHER 是给那些”明确保证不变量”的部署的逃生门。

还有一个细节：sync_wait_seconds 被单独记账并用 DPSyncWaitCarry 跨轮传递：

1
    sync_wait_seconds = mlp_sync_info.sync_wait_seconds
2
    if sync_wait_carry is not None:
3
        sync_wait_seconds += sync_wait_carry.seconds
4
        sync_wait_carry.seconds = 0.0

如果某轮没产出 batch（drop_sync_wait_carry），这次等待的时间会被记到下一轮，避免统计上丢掉。因为”等待最慢的副本”正是 DP Attention 最典型的开销，把它单独测量出来才能判断批处理是否失衡。

全局 forward mode 与 TBO 的耦合

MLPSyncBatchInfo 里还有一个 global_forward_mode，它由 TboDPAttentionPreparer.compute_output 产出：

1
    @staticmethod
2
    def _compute_global_forward_mode(forward_modes):
3
        forward_modes_excluding_idle_and_prebuilt = [
4
            x
5
            for x in forward_modes
6
            if x != ForwardMode.IDLE.value and x != ForwardMode.PREBUILT.value
7
        ]
8

9
        if not forward_modes_excluding_idle_and_prebuilt:
10
            return ForwardMode.IDLE, False
11

12
        forward_mode_agree = TboDPAttentionPreparer._is_all_same(
13
            forward_modes_excluding_idle_and_prebuilt
14
        )
15

16
        global_forward_mode = (
17
            ForwardMode(forward_modes_excluding_idle_and_prebuilt[0])
18
            if forward_mode_agree
19
            else None
20
        )
21
        return global_forward_mode, forward_mode_agree

逻辑是先剔掉 IDLE 与 PREBUILT，再看剩下的是不是全部相同。全部相同才给出全局 mode，否则给 None。为什么这么严格？因为 two-batch overlap 需要把 batch 按 token 维切成两个 micro-batch，切分点的计算依赖 forward mode：

1
def compute_split_seq_index(
2
    forward_mode: ForwardMode,
3
    num_tokens: int,
4
    extend_lens: Optional[Sequence[int]],
5
    token_num_per_seq: Optional[int],
6
) -> Optional[int]:
7
    if forward_mode == ForwardMode.EXTEND or forward_mode == ForwardMode.MIXED:
8
        assert extend_lens is not None
9
        return _split_extend_seqs(extend_lens)
10
    elif forward_mode.is_target_verify() or forward_mode.is_decode():
11
        assert token_num_per_seq is not None
12
        return (num_tokens // token_num_per_seq) // 2
13
    elif forward_mode.is_idle() or forward_mode.is_prebuilt():
14
        assert num_tokens == 0
15
        return 0
16
    else:
17
        raise NotImplementedError()

EXTEND 按序列边界切（_split_extend_seqs），decode 按 token 数对半切（(num_tokens // token_num_per_seq) // 2）。如果各副本的 forward mode 不一致（有人在 prefill 有人在 decode），就没有一个统一的切法，TBO 只能关掉。这就是 DP Attention 与 TBO 的耦合点：TBO 的可用性由 DP 副本的 mode 一致性决定。

Token 轴：把”谁切了什么”显式化

layout.py 定义了三个 token 轴：

1
class TokenAxis(Enum):
2
    """A parallel dimension across which ranks hold different tokens."""
3

4
    ATTN_DP = auto()
5
    ATTN_CP = auto()
6
    # Each attention-TP rank holds a slice of its group's tokens.
7
    ATTN_TP = auto()

配套的 Layout 只记录一件事：哪些轴被真正切过（size 为 1 的轴被 sharded_over() 过滤掉）：

1
    @classmethod
2
    def sharded_over(
3
        cls, *axes: TokenAxis, axis_sizes: Mapping[TokenAxis, int]
4
    ) -> "Layout":
5
        return cls(frozenset(axis for axis in axes if axis_sizes[axis] > 1))

这个设计解决了一个很实际的问题：模型文件不该知道部署拓扑。一个 DeepSeek 层的代码不需要写”如果 attn_dp_size > 1 就 gather”这种分支，它只要声明”我这个阶段的输入/输出在哪些轴上是切分的”，边界层根据 Layout 的差集自动决定要不要通信、在哪个组上通信。move_rows 就是这个决策的执行者：

1
def move_rows(
2
    hidden_states: torch.Tensor,
3
    rows: Layout,
4
    to: Layout,
5
    forward_batch: ForwardBatch,
6
) -> torch.Tensor:
7
    """A complete value from the rows it is on to ``to``: gathered over the
8
    token axes ``to`` does not shard (attention TP, then attention DP), or cut to
9
    this rank's share of those it does (attention DP, then attention TP)."""
10
    gathered, cut = rows.sharded - to.sharded, to.sharded - rows.sharded
11
    if (gathered and cut) or TokenAxis.ATTN_CP in gathered | cut:
12
        raise NotImplementedError(f"{rows=} {to=}")
13
    if TokenAxis.ATTN_TP in gathered:
14
        hidden_states = attn_tp_gather(hidden_states)
15
    if TokenAxis.ATTN_DP in gathered:
16
        hidden_states = dp_gather(hidden_states, forward_batch)
17
    if TokenAxis.ATTN_DP in cut:
18
        hidden_states = to_dp_local(_dp_scatter_step, forward_batch, hidden_states)
19
    if TokenAxis.ATTN_TP in cut:
20
        parallel = get_parallel()
21
        hidden_states = hidden_states.tensor_split(parallel.attn_tp_size)[
22
            parallel.attn_tp_rank
23
        ]
24
    return hidden_states

它的输入是”当前在哪些轴上被切了”和”目标要求哪些轴被切”，输出是执行必要通信后的张量。注意几个约束：不允许同时既有 gather 又有 cut（那说明布局设计有问题，直接 NotImplementedError 而不是默默算错），ATTN_CP 不允许在这里移动（上下文并行的重排有自己的路径）。

边界判定的融合

boundary.py 里有一段逻辑决定”这次 gather 要不要顺手把该做的 reduce 和 norm 一起做掉”。核心判定是：

1
    places_cp_shards = TokenAxis.ATTN_CP in gathered
2
    if (
3
        owes_attention_tp
4
        and not read.reads_before_dp_gather
5
        and is_plain_add
6
        and read.is_plain_norm
7
    ):
8
        return EntryStep(
9
            partial(
10
                _dp_gather_sum_read,
11
                gathers_residual=gathers_residual,
12
                places_cp_shards=places_cp_shards,
13
                read=read,
14
            )
15
        )
16
    return EntryStep(
17
        partial(
18
            _reduce_update_read_dp_gather,
19
            gathers_residual=gathers_residual,
20
            reduces_attention_tp=owes is not None,
21
            group=owes,
22
            places_cp_shards=places_cp_shards,
23
            read=read,
24
        )
25
    )

三个条件同时满足才有”最强融合”：

• owes_attention_tp：attention-TP 组欠一次 all-reduce（输出投影还没做完）；
• not read.reads_before_dp_gather：这次读取在 gather 之后做，而不是之前；
• is_plain_add 且 read.is_plain_norm：残差只是普通加法、norm 只是普通 RMSNorm。

reads_before_dp_gather 是一个很有意思的标记。源码里的解释是：某些残差读取（比如 attention-residual bank 的聚合读取）必须在 DP gather 之前完成，因为它读的是本 rank 自己持有的那些行（bank 就存在本 rank 上），gather 之后行序变了就找不到了。同理，“gated” 类残差（门控残差）也不能被普通 add+norm 的融合 kernel 接管，所以它们的 is_plain_add 是 False。

于是 _reduce_update_read_dp_gather 的实现就是顺理成章的”一次做完三件事”：

1
def _reduce_update_read_dp_gather(
2
    hidden_states: torch.Tensor,
3
    residual: torch.Tensor,
4
    forward_batch: ForwardBatch,
5
    norm: torch.nn.Module,
6
    *,
7
    cache=None,
8
    gathers_residual: bool,
9
    reduces_attention_tp: bool,
10
    places_cp_shards: bool = False,
11
    group: SumGroup = SumGroup.ATTN_TP,
12
    read: ResidualReadout = NORM_READOUT,
13
    update: ResidualUpdate = PLAIN_ADD,
14
):
15
    """Attention DP: complete the attention-TP sum if it is owed, write the
16
    output into the residual and read the FFN input locally, then gather."""
17
    if gathers_residual:
18
        residual = update.gather_residual_attn_tp(residual)
19
    if hidden_states.shape[0] != 0:
20
        if reduces_attention_tp:
21
            hidden_states = (
22
                attention_tensor_model_parallel_all_reduce(hidden_states)
23
                if group is SumGroup.ATTN_TP
24
                else _sum_group(group).all_reduce(hidden_states)
25
            )
26
        with use_symmetric_memory(
27
            get_parallel().tp_group,
28
            disabled=not is_allocation_symmetric(),
29
        ):
30
            hidden_states, residual = read.update_and_read(
31
                update, hidden_states, residual, norm
32
            )
33
    else:
34
        hidden_states, residual = read.update_and_read(
35
            update, hidden_states, residual, norm
36
        )
37
    cp_shard_counts = _cp_shard_token_rows(forward_batch) if places_cp_shards else None
38
    return (
39
        dp_gather(hidden_states, forward_batch, cp_shard_counts),
40
        residual,
41
    )

顺序是：先 all-reduce 完成 attention-TP 的求和 → 再 update_and_read（把结果写进残差、读出 norm 后的 FFN 输入）→ 最后 DP gather。这三步原本是三次独立的算子（一次集合通信、一次残差加法、一次 RMSNorm，再加一次集合通信），现在合成一趟。

use_symmetric_memory 这个上下文管理器也值得注意：它是对称显存分配，让参与通信的 rank 在相同的地址上拥有相同大小的 buffer，从而让通信库可以做更高效的 zero-copy 传输。is_allocation_symmetric() 决定能不能启用——如果某个 rank 的分配不对称（通常是动态形状），就退回普通分配。

hidden_states.shape[0] != 0 这个分支是给空闲 rank 准备的：0 token 的 rank 不做 all-reduce（没意义），但仍要完成 update/read 与 gather（保持对称）。

后端注册表

layers/moe/token_dispatcher/__init__.py 导出了十来个 dispatcher：DeepEPDispatcher、DeepEPv2Dispatcher、MooncakeEPDispatcher、NixlEPDispatcher、MoriEPDispatcher、PplxDispatcher、FlashinferDispatcher、AscendTPDispatcher、StandardDispatcher。对应的 moe_a2a_backend 取值在参数文档里有完整枚举：none / deepep / mooncake / nixl / mori / ascend_fuseep / flashinfer / megamoe / deepep_v2 / ascend_tp / pplx / flashinfer_megamoe。

StandardDispatcher 是”没有 A2A 后端”时的路径，它不搬 token，而是用 all-gatherv 把各 DP 副本的 hidden state 拼起来、用 reduce_scatterv 把结果拆回去——换汤不换药，本质上还是 DP gather/scatter 那一套。所以 DP Attention 在没有 EP 的情况下也能工作，只是专家分片和 token 搬运都由 TP 组承担。

normal 与 low-latency：一对互补的模式

DeepEP 提供两种 dispatch 模式，它们的取舍是整个 EP 设计的枢纽。先看模式枚举：

1
class DeepEPMode(Enum):
2
    NORMAL = "normal"
3
    LOW_LATENCY = "low_latency"
4
    AUTO = "auto"
5

6
    def enable_normal(self) -> bool:
7
        return self in [DeepEPMode.NORMAL, DeepEPMode.AUTO]
8

9
    def enable_low_latency(self) -> bool:
10
        return self in [DeepEPMode.LOW_LATENCY, DeepEPMode.AUTO]
11

12
    def resolve(self, is_extend_in_batch: bool) -> DeepEPMode:
13
        if self != DeepEPMode.AUTO:
14
            return self
15
        if is_extend_in_batch:
16
            return DeepEPMode.NORMAL
17
        else:
18
            return DeepEPMode.LOW_LATENCY

AUTO 的规则一句话：prefill 用 normal，decode 用 low-latency。这不是偏好，而是两种模式的物理特性决定的：

• normal dispatch 为长输入优化。prefill 时一个 batch 可能有几万个 token，通信量巨大，追求的是吞吐。它的输出形状是符号化的（symbolic shape）——每个专家收到多少 token 取决于路由结果，事先不知道。
• low-latency dispatch 为逐 token 生成优化。decode 时每个请求只有 1 个 token，通信量极小但延迟极其关键。它要求预分配固定内存：每个专家最多能收多少 token 是启动时定死的，超出就报运行时错误。

这两条特性各自带来一个后果，而后果决定了能不能配 DP Attention 和 CUDA Graph：

这张表来自 LMSYS 的大规模 EP 博客。读法：

• Normal 支持 DP Attention 但不支持 CUDA Graph。因为符号化形状没法捕获成静态图。
• Low-Latency 支持 CUDA Graph 但只擅长长输出（decode）。
• Auto 两者都好，但不支持 DP Attention——在非 PD 分离的部署下。博客的解释是：auto 模式无法在同一个通信组里同时支持 normal（给 prefill）和 low-latency（给 decode），而这个限制恰好与 DP Attention 冲突。

这解释了一个看起来奇怪的组合最优解：PD 分离 + Auto。把 prefill 和 decode 拆到不同实例之后，prefill 实例（跑 EP + normal dispatch）和 decode 实例（跑 EP + low-latency dispatch）可以各自选适合的模式，两者都在 DP Attention 下工作，资源利用率和性能都更好。这也是为什么站内讲 P/D 分离时总把它和 EP 放在一起——见 SGLang 源码完全拆解（十二）：Disaggregation。

这里的结论对工程很重要：“能不能同时用 DP Attention + CUDA Graph + 混合负载”是一个三角约束，PD 分离是把它解开的那把钥匙。

normal dispatch 的调用链

_DeepEPDispatcherImplNormal._dispatch_core 是真实通信发生的地方：

1
        buffer = self._get_buffer()
2
        (
3
            num_tokens_per_rank,
4
            num_tokens_per_rdma_rank,
5
            num_tokens_per_expert,
6
            is_token_in_rank,
7
            previous_event,
8
        ) = buffer.get_dispatch_layout(
9
            topk_ids,
10
            self.num_experts,
11
            previous_event=previous_event,
12
            async_finish=self.async_finish,
13
            allocate_on_comm_stream=previous_event is not None,
14
        )
15
        ...
16
        (
17
            recv_x,
18
            recv_topk_ids,
19
            recv_topk_weights,
20
            num_recv_tokens_per_expert,
21
            self.handle,
22
            event,
23
        ) = buffer.dispatch(
24
            x,
25
            topk_idx=topk_ids,
26
            topk_weights=topk_weights,
27
            num_tokens_per_rank=num_tokens_per_rank,
28
            num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
29
            is_token_in_rank=is_token_in_rank,
30
            num_tokens_per_expert=num_tokens_per_expert,
31
            previous_event=previous_event,
32
            async_finish=self.async_finish,
33
            allocate_on_comm_stream=(previous_event is not None) and self.async_finish,
34
            expert_alignment=128 if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM else 1,
35
            config=DeepEPConfig.get_instance().normal_dispatch_config,
36
            **npu_quantization_opts,
37
        )

几个关键点：

• get_dispatch_layout 先算出四个统计量：每 rank 收多少 token、每 RDMA rank 收多少（RDMA rank 是”节点”层面的计数，因为跨节点要先到对端节点再分发）、每专家收多少、以及每个 token 属于哪些 rank（is_token_in_rank 是布尔矩阵）。这套 layout 是后续通信的”地址簿”。
• expert_alignment=128（当启用 JIT DeepGEMM 时）：把每个专家的 token 数对齐到 128 的倍数，让后续的 grouped GEMM 能用整齐的 tile，并且能配合 masked 布局。
• self.handle 被存成成员变量，源码注释坦白说这是个妥协：

1
        # FIXME: `handle` should be transmitted with tokens from dispatch to combine.
2
        # However, doing this would incur an unknown synchronization error, but keeping
3
        # `handle` as a member variable works.

handle 里记录的是”combine 阶段要把结果送回哪里”的信息。理论上它应该跟着 token 一起传（因为它是每批数据的一部分），但因为一个尚未定位的同步错误，SGLang 选择把它挂在 dispatcher 对象上。这类”注释里留着 FIXME 但实际必须这么写”的地方，是读源码时最值得记下来的——它说明这个约束来自底层通信库的实现细节，不是设计者的偏好。

低延迟路径对应的 dispatch_a 则先算一个 expected_m：

1
    def dispatch_a(
2
        self,
3
        hidden_states: torch.Tensor,
4
        topk_output: TopKOutput,
5
    ):
6
        buffer = self._get_buffer()
7
        topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
8
        topk_ids = topk_ids.to(torch.int64)
9
        expected_m = (
10
            hidden_states.shape[0] * buffer.group_size * topk_ids.shape[1]
11
            + self.num_experts
12
        ) // self.num_experts
13
        ...

expected_m 是”每个专家平均会收到多少 token”的估计值：(本 rank token 数 × 组大小 × top-k) / 专家数，上取整。为什么要这个估计？因为 low-latency 模式要预分配固定内存，分配多少必须事先知道。这个公式就是”最坏情况下的期望值”，用它作为分配依据。

为什么要拆成 dispatch_a / dispatch_b

两种模式的实现都把 dispatch 拆成两段：dispatch_a 做准备工作（量化、记录事件、算期望值），dispatch_b 真正等待通信完成并组装输出。看 DeepEPDispatcher 的门面：

1
    def dispatch_a(
2
        self,
3
        hidden_states: torch.Tensor,
4
        topk_output: TopKOutput,
5
    ):
6
        self._update_stage(_Stage.INITIAL, _Stage.AFTER_DISPATCH_A)
7
        inner_state = self._get_impl().dispatch_a(
8
            hidden_states=hidden_states,
9
            topk_output=topk_output,
10
        )
11
        self._dispatch_intermediate_state = inner_state
12

13
    def dispatch_b(self):
14
        self._update_stage(_Stage.AFTER_DISPATCH_A, _Stage.AFTER_DISPATCH_B)
15
        inner_state = self._dispatch_intermediate_state
16
        del self._dispatch_intermediate_state
17
        return self._get_impl().dispatch_b(*inner_state)

拆分的目的是插入重叠点。两次调用之间，调度层可以插入别的计算（比如另一个 micro-batch 的 attention），这就是 two-batch overlap 的实现基础。_update_stage 里那些断言（assert self._stage == old_stage）是状态机的守卫：如果一个 dispatcher 被 dispatch_a 之后没有 dispatch_b 又被复用，就会立刻报错而不是产生错误结果。

同时注意 self.handle 那个设计为什么必须如此：因为 dispatch_a 与 dispatch_b 之间状态被暂存在对象上，handle 作为成员变量恰好与这个生命周期吻合。

_dispatch_core 还在 NPU 路径上加了量化选项：

1
    def _get_quantization_kwargs(self, buffer: Buffer) -> dict:
2
        if not _is_npu:
3
            return {}
4
        dispatch_params = inspect.signature(buffer.dispatch).parameters
5
        flag_kwargs = {"use_fp8": self.use_fp8, "use_mxfp4": self.use_mxfp4, "use_mxfp8": self.use_mxfp8}
6
        if all(name in dispatch_params for name in flag_kwargs):
7
            return flag_kwargs
8
        if "quant_mode" in dispatch_params:
9
            if self.use_mxfp4:
10
                quant_mode = "mx_fp4_e2m1"
11
            elif self.use_mxfp8:
12
                quant_mode = "mx_fp8_e4m3"
13
            elif self.use_fp8:
14
                quant_mode = "int8"
15
            else:
16
                quant_mode = "bf16"
17
            return {"quant_mode": quant_mode}

它用 inspect.signature 探测底层 buffer 支持哪套参数名，然后在新旧 API 之间做适配。deepep_dispatcher_output_dtype 这个参数在文档里的取值是 auto / bf16 / fp8 / int8 / nvfp4——通信也可以量化。把 token 用 FP8 或 FP4 发出去，通信量减半甚至降到四分之一，代价是精度损失；在带宽受限的多节点场景下这笔账通常是划算的。

prefll 侧的一个必要步骤：per-token-group 量化

dispatch_a 里还有一段：

1
        topk_ids = topk_ids.to(torch.int64)
2
        if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and self.use_fp8:
3
            # TODO hard code 128 block quant,use fp8 communication
4
            hidden_states = sglang_per_token_group_quant_fp8(
5
                hidden_states,
6
                128,
7
                column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
8
                scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
9
                scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
10
            )

注释里那句 TODO hard code 128 block quant 很直白：量化块大小固定写死成 128，与 expert_alignment=128 是同一个数字，两者是配合关系。如果启用了 DeepGEMM，缩放因子还要按 UE8M0 格式排布并且 TMA 对齐——因为这些缩放因子最终要喂给 DeepGEMM 的 kernel，格式必须匹配。DeepGEMM 的 1×128 细粒度缩放机制见 DeepGEMM 完全拆解。

切分点怎么定

compute_split_seq_index 已经在前面看过：EXTEND 按序列边界切，decode 按 token 数对半切。但还有一个更细的问题——能不能把单条序列切开？_is_two_chunk_split_enabled 处理这个：

1
def _is_two_chunk_split_enabled(extend_lens: Sequence[int]) -> bool:
2
    if extend_lens is None:
3
        return False
4

5
    vanilla_split_seq_index = _split_array_by_balanced_sum(extend_lens)
6
    left_sum = sum(extend_lens[:vanilla_split_seq_index])
7
    overall_sum = sum(extend_lens)
8
    threshold = get_tbo_token_distribution_threshold()
9
    assert threshold <= 0.5, f"{threshold=}"
10
    want_two_chunk = left_sum < overall_sum * threshold or left_sum > overall_sum * (
11
        1 - threshold
12
    )
13
    if not want_two_chunk:
14
        return False
15
    ...

默认策略是按”累积和接近一半”找序列边界切（_split_array_by_balanced_sum）。但如果这样切出来极不平衡——左边占不到 48% 或者超过 52%（tbo_token_distribution_threshold 默认 0.48）——说明某个长序列霸占了大部分 token，按序列边界切会让两个 micro-batch 的负载严重不均。这时改用”把单条序列切成两半”的两块模式。但源码紧接着留了一个保护：

1
    split_seq_index = _split_array_by_cum_less_than_half(extend_lens)
2
    child_a_batch_size = split_seq_index + 1
3
    child_a_num_q_tokens = overall_sum // 2
4
    if child_a_batch_size > child_a_num_q_tokens:
5
        return False
6

7
    return True

注释详细解释了原因：two-chunk 切法让 child_a 跨越 split_seq_index + 1 条序列，但只分到 overall_sum // 2 个 query token。对一个退化 batch（只有一条序列，或者接近空的 DP-sync batch），这个”半个总数”可能只有 0 或极小，于是 child_a 的 batch_size 大于它的 query token 数——这违反了 DSV4 压缩规划器的不变量 batch_size <= num_q_tokens，kernel 会崩溃。所以检测到这种情况就退回按序列边界切，那种切法保证 child_a 每条序列至少贡献 1 个 token，不会出现”有行无 token”。

这就是推理框架里最典型的一类 bug：切分逻辑在正常情况下完全正确，但在极端小的 batch 上产生一个退化输入，然后在某个深层 kernel 的断言上炸掉。

切完之后要重建哪些字段

切成两个 child 之后，很多与序列长度相关的字段需要重新推导。_two_chunk 路径下，derive_fields_related_to_seq_len_for_two_chunk 做的事是：

1
        extend_seq_lens_cpu = batch.extend_seq_lens_cpu
2
        overall_seq_lens_sum = sum(extend_seq_lens_cpu)
3
        half_seq_lens_sum = overall_seq_lens_sum // 2
4
        left_last_seq_token_num = half_seq_lens_sum - sum(
5
            extend_seq_lens_cpu[:tbo_split_seq_index]
6
        )
7
        right_first_seq_token_num = (
8
            extend_seq_lens_cpu[tbo_split_seq_index] - left_last_seq_token_num
9
        )
10

11
        child_a.extend_seq_lens_cpu = copy.deepcopy(child_a.extend_seq_lens_cpu)
12
        child_a.extend_seq_lens_cpu[-1] = left_last_seq_token_num
13
        child_b.extend_seq_lens_cpu = copy.deepcopy(child_b.extend_seq_lens_cpu)
14
        child_b.extend_seq_lens_cpu[0] = right_first_seq_token_num

被切开的那条序列（第 tbo_split_seq_index 条）在两个 child 里各占一段：child_a 拿它的前 left_last_seq_token_num 个 token，child_b 拿剩下的。对应的 seq_lens（含前缀）、extend_prefix_lens 也都要跟着改：

1
        child_a.seq_lens_cpu = copy.deepcopy(child_a.seq_lens_cpu)
2
        child_a.seq_lens_cpu[-1] = (
3
            child_a.extend_seq_lens_cpu[-1] + child_a.extend_prefix_lens_cpu[-1]
4
        )
5
        ...
6
        child_b.extend_prefix_lens_cpu = copy.deepcopy(child_b.extend_prefix_lens_cpu)
7
        child_b.extend_prefix_lens_cpu[0] += left_last_seq_token_num

child_b 的那条被切序列，它的前缀长度要加上 child_a 已经算过的那一段——因为从 child_b 的视角看，child_a 处理过的 token 在因果意义上就是它的前缀。这是块状 prefill 的同一个思想（见 Sarathi-Serve 完全拆解），只是这里切的是 micro-batch 而不是 chunk。

operations 与 yield：TBO 的表达方式

LMSYS 博客里点出了 TBO 实现的两个难点：代码复杂度（直接写会把 micro-batch 管理逻辑复制到各处）和 prefill 阶段的同步问题（normal dispatch 会阻塞 CPU，pipeline 被卡住）。

SGLang 的解法是引入”操作 + 让出点”的抽象。batch_overlap/operations.py 里 Operation 是三种东西的联合：

1
class YieldOperation:
2
    pass
3

4

5
@dataclass
6
class ExecutionOperation:
7
    debug_name: str
8
    fn: Callable
9

10

11
Operation = Union[YieldOperation, ExecutionOperation, Callable]
12
Stage = List[ExecutionOperation]

一个 micro-batch 的整层计算被写成一个扁平的 operation 列表，其中 YieldOperation() 标记”这里可以让另一个 micro-batch 插进来”。以 DeepSeek 的 prefill 策略为例：

1
def _compute_moe_deepseek_blog_prefill(layer):
2
    device_properties = torch.cuda.get_device_properties(device="cuda")
3
    total_num_sms = device_properties.multi_processor_count
4
    deep_gemm_num_sms = None
5
    if not _is_hip:
6
        deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms
7

8
    return OperationsStrategy(
9
        deep_gemm_num_sms=deep_gemm_num_sms,
10
        tbo_delta_stages=0,
11
        operations=[
12
            layer.op_comm_prepare_attn,
13
            layer.self_attn.op_prepare,
14
            layer.self_attn.op_core,
15
            layer.op_comm_prepare_mlp,
16
            layer.mlp.op_gate,
17
            layer.mlp.op_select_experts,
18
            layer.mlp.op_dispatch_a,
19
            operations.YieldOperation(),
20
            layer.mlp.op_dispatch_b,
21
            layer.mlp.op_experts,
22
            layer.mlp.op_combine_a,
23
            operations.YieldOperation(),
24
            layer.mlp.op_shared_experts,
25
            layer.mlp.op_combine_b,
26
            layer.mlp.op_output,
27
            layer.op_comm_postprocess_layer,
28
        ],
29
    )

这份列表就是”顺序执行时的完整流程”，YieldOperation 精确地插在两个地方：dispatch_a 之后（通信已经发起、还没等结果）和 combine_a 之后。这些正是通信可以被打断、让位于另一个 micro-batch 计算的时刻。

deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms 这一行也很关键：通信和计算要真正并行，SM 必须分开。如果 GEMM 把 132 个 SM 全占了，通信 kernel 就没有 SM 可用，重叠是假的。所以 GEMM 只准用”总 SM 数减去通信核占用”的那些。这是 overlap 能不能落地的物理前提，而不是调度技巧。

decode 的策略 tbo_delta_stages=2，意思是 child_b 相对 child_a 延迟 2 个 stage 开始，错开得更精细。

execute_overlapped_operations 是执行器：

1
def execute_overlapped_operations(
2
    inputs_arr: Sequence,
3
    operations_arr: Sequence,
4
    delta_stages: Sequence[int],
5
) -> Sequence:
6
    inputs_a, inputs_b = inputs_arr
7
    operations_a, operations_b = operations_arr
8
    delta_stage_a, delta_stage_b = delta_stages
9
    assert delta_stage_a == 0
10
    delta_stage = delta_stage_b
11
    ...
12
    stages_a = _convert_operations_to_stages(operations_a)
13
    stages_b = _convert_operations_to_stages(operations_b)
14
    executor_a = _StageExecutor("a", stages_a, inputs=inputs_a, child_ctx=child_ctx_a)
15
    executor_b = _StageExecutor("b", stages_b, inputs=inputs_b, child_ctx=child_ctx_b)
16

17
    for _ in range(delta_stage):
18
        executor_a.next()
19

20
    for _ in range(executor_a.num_stages - delta_stage):
21
        executor_a.next()
22
        executor_b.next()
23

24
    for _ in range(delta_stage):
25
        executor_b.next()

模型是”两个游标交错前进”：a 先单独跑 delta_stage 个 stage，然后两者一步步交替，最后 b 独自收尾。stage 划分规则没显式写出来，但从 YieldOperation 作为分界来看，每个 stage 就是”从上一次让出到下一次让出之间的所有操作”——即”发起到等待”为一段。这样交替执行时，a 在等通信的那一段正好 b 在算，反之亦然。

_resolve_tbo_child_contexts 是个容易被忽略的细节：

1
def _resolve_tbo_child_contexts():
2
    """Return (child_ctx_a, child_ctx_b) derived from the active TboAttnBackend,
3
    or (None, None) if the active backend is not a TBO dispatcher."""
4
    from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
5

6
    ctx = get_forward_context()
7
    backend = ctx.attn_backend
8
    if not isinstance(backend, TboAttnBackend):
9
        return None, None
10
    child_a, child_b = backend.children
11
    return (
12
        replace(ctx, attn_backend=child_a),
13
        replace(ctx, attn_backend=child_b),
14
    )

每个 child 需要自己的注意力后端实例——因为两个 child 的 batch 形状和 KV 访问模式不同，注意力后端里缓存的元数据（页表、序列长度等）也不同。TboAttnBackend 就是为此把两个后端包起来、按当前执行的是哪个 child 分派。所以 TBO 不只是切 batch，还要把整条执行上下文复制两份。

TBO 与 DP Attention 的耦合，以及 forward mode 一致性的来源

前面已经看到 TBO 的启用条件包含”全局 forward mode 一致”。把两段代码连起来：

1
        self.global_forward_mode = mlp_sync_info.global_forward_mode
2
        ...
3
    # Decide whether to emit idle batch

global_forward_mode 在 compute_output 里算出，写进 ScheduleBatch，最终影响到是否切分。于是一个 DP 副本在 prefill、另一个在 decode 时，TBO 会自动关闭——因为没有任何一致的切分策略同时适用于两边。

这也是为什么 TBO 只在 decode 或只在 prefill 时最有效：这两种场景下各副本的 mode 天然一致，TBO 可以稳定开启。

不均衡的代价与冗余专家的思路

SGLang 用的度量是 balancedness：

1
def compute_utilization_rate(
2
    gpu_physical_count_of_batch: torch.Tensor,  # (..., num_layer, num_gpu)
3
):
4
    """output: utilization_rate (..., num_layer)"""
5
    gpu_physical_count_of_batch = gpu_physical_count_of_batch.float()
6
    max_gpu_physical_count = einops.reduce(
7
        gpu_physical_count_of_batch,
8
        "... num_layer num_gpu -> ... num_layer",
9
        "max",
10
    )
11
    avg_gpu_physical_count = einops.reduce(
12
        gpu_physical_count_of_batch,
13
        "... num_layer num_gpu -> ... num_layer",
14
        "mean",
15
    )
16
    return (avg_gpu_physical_count + 1e-5) / (max_gpu_physical_count + 1e-5)

先按 GPU 把物理专家的 token 数求和，再取所有 GPU 里的最大值与平均值之比。等于 1 表示完美均衡，越小越不均衡。分母加 1e-5 是防除零。

LMSYS 博客给出的量化：随着节点数（EP 规模）增加，利用率会持续下降，打开 EPLB 后显著回升。这就是这张图要表达的东西：

图片来源：SGLang 大规模 EP 设计博客 的 EPLB 仿真图

解决思路来自 DeepSeek 的 EPLB：允许冗余专家。原本 256 个专家、EP=16 时每 rank 16 个；现在多分配 32 个冗余槽位，组成 288 个物理专家的池子，让高频专家被复制多份摆在不同 GPU 上，中低频专家被打包合并到同一个 GPU 上。

参数上对应的就是：

1
ep_num_redundant_experts: "Allocate this number of redundant experts in expert parallel."  # 默认 0
2
enable_eplb: "Enable EPLB algorithm"                                                       # 默认 False
3
eplb_algorithm: "Chosen EPLB algorithm"                                                    # 默认 "auto"
4
init_expert_location: "Initial location of EP experts."                                    # 默认 "trivial"
5
ep_dispatch_algorithm: "The algorithm to choose ranks for redundant experts in expert parallel."
6
    # 取值 static / dynamic / fake / lp，默认 None

init_expert_location = "trivial" 就是”逻辑专家 i 对应物理专家 i、没有冗余”的平凡布局；ep_num_redundant_experts = 0 时就是这个状态。

博客还指出冗余专家的一个附带好处：并行度不再被限制为 2 的幂。256 个专家时，能整除它的 GPU 数只有 2 的幂次（必须让专家数能被 GPU 数整除）；换成 288（= 256 + 32）之后，12、72 这类规模也变得可行。这直接解释了后面评测里为什么用 EP72。

三步算法：分组、复制、再打包

eplb/eplb_algorithms/deepseek.py 里的 rebalance_experts_hierarchical 是整个算法的入口，它的签名很能说明问题：

1
def rebalance_experts_hierarchical(
2
    weight: torch.Tensor,
3
    num_physical_experts: int,
4
    num_groups: int,
5
    num_nodes: int,
6
    num_gpus: int,
7
):
8
    """
9
    Parameters:
10
        weight: [num_moe_layers, num_logical_experts]
11
        num_physical_experts: number of physical experts after replication
12
        num_groups: number of expert groups
13
        num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
14
        num_gpus: number of GPUs, must be a multiple of `num_nodes`
15
    """

输入是每个逻辑专家的负载权重，输出是物理专家的排布方案。num_nodes 这个参数是算法的核心约束来源：NVLink 比跨节点网络快得多，所以同一个专家组内的通信应当尽量留在节点内。

算法分三步：

第一步：把专家组打包到节点上。

1
    # Step 1: pack groups to nodes
2
    tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
3
    group_pack_index, group_rank_in_pack = balanced_packing(tokens_per_group, num_nodes)
4
    log2mlog = (
5
        (
6
            (group_pack_index * groups_per_node + group_rank_in_pack) * group_size
7
        ).unsqueeze(-1)
8
        + torch.arange(group_size, dtype=torch.int64, device=group_pack_index.device)
9
    ).flatten(-2)
10
    mlog2log = inverse(log2mlog)

先把专家按组聚合负载（tokens_per_group），然后用 balanced_packing 把组均匀分配到各节点上。这一步的目标是让各节点的总负载接近——因为跨节点通信最贵，节点间的失衡代价最大。得到 log2mlog（逻辑→节点内逻辑的映射）和它的逆 mlog2log。

第二步：在节点内构造冗余专家。

1
    # Step 2: construct redundant experts within nodes
2
    # [num_layers * num_nodes, num_logical_experts // num_nodes]
3
    tokens_per_mlog = weight.gather(-1, mlog2log).view(
4
        -1, num_logical_experts // num_nodes
5
    )
6
    phy2mlog, phyrank, mlogcnt = replicate_experts(
7
        tokens_per_mlog, num_physical_experts // num_nodes
8
    )

replicate_experts 就是在每个节点内部做复制：负载高的专家多给几个物理槽位，负载低的合并。它返回三样东西：phy2mlog（物理→节点内逻辑的映射）、phyrank（该物理专家是第几份副本）、mlogcnt（每个节点内逻辑专家被复制了几份）。冗余只发生在节点内——这是刻意的：跨节点复制高频专家会把跨节点通信量翻倍，得不偿失。

第三步：把物理专家打包到 GPU 上。

1
    # Step 3: pack physical_experts to GPUs
2
    # [num_layers * num_nodes, num_physical_experts // num_nodes]
3
    tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
4
    pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes)
5
    phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
6
    pphy2phy = inverse(phy2pphy)

注意 tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog) 这一行：复制出来的副本平摊负载。如果一个专家被复制了 3 份，那每一份名义上承担 1/3 的流量，打包时按平摊后的值算。这不是精确建模（实际路由是”选最不忙的副本”，不是均分），但它是一个合理的近似，让打包结果能把副本分散到不同 GPU。

然后 balanced_packing(tokens_per_phy, num_gpus // num_nodes) 把物理专家均匀分配到节点内的各 GPU 上。

算法骨架的意义：三层嵌套均衡（组→节点、专家→节点内槽位、槽位→GPU）不是随意的，它精确对应了三种通信代价的量级——跨节点最贵、节点内 NVLink 次之、GPU 内最便宜。每一层的均衡目标都是”把最贵的通信路径上的不平衡压掉”。

逻辑与物理的两套映射

结果用 ExpertLocationMetadata 装上：

1
class ExpertLocationMetadata:
2
    physical_to_logical_map: torch.Tensor  # (layers, num_physical_experts)
3
    physical_to_logical_map_cpu: torch.Tensor
4
    logical_to_all_physical_map: torch.Tensor  # (layers, num_logical_experts, X)
5
    logical_to_all_physical_map_cpu: torch.Tensor  # CPU copy for performance
6
    logical_to_all_physical_map_num_valid: torch.Tensor  # (layers, num_logical_experts)
7
    ep_size: int
8
    # (layers, num_logical_experts)
9
    logical_to_rank_dispatch_physical_map: Optional[torch.Tensor]

两套映射是理解 EPLB 的关键：

• physical_to_logical_map：每个物理槽位上放的是哪个逻辑专家。用来更新权重（把专家的参数搬到位）和统计负载（把物理计数还原成逻辑计数）。
• logical_to_all_physical_map：每个逻辑专家被复制到了哪些物理槽位。用来在推理时分派（token 路由到某逻辑专家后，要选一个物理副本）。

logical_to_all_physical_map_num_valid 记录每个逻辑专家实际有几个有效副本（因为 X 是最大副本数的上界，某些专家副本少）。logical_to_rank_dispatch_physical_map 是”某逻辑专家在某 rank 上应该派到哪个物理槽位”的预计算表。

__post_init__ 里的断言检查三者形状一致：

1
    def __post_init__(self):
2
        num_layers_0, num_physical_experts_0 = self.physical_to_logical_map.shape
3
        num_layers_1, num_logical_experts_0, num_physical_experts_1 = (
4
            self.logical_to_all_physical_map.shape
5
        )
6
        num_layers_2, num_logical_experts_1 = (
7
            self.logical_to_all_physical_map_num_valid.shape
8
        )
9
        assert num_layers_0 == num_layers_1 == num_layers_2
10
        assert num_logical_experts_0 == num_logical_experts_1
11
        assert num_physical_experts_0 == num_physical_experts_1

init_trivial 给出无冗余时的布局：

1
    @staticmethod
2
    def init_trivial(model_config: ModelConfig, moe_ep_rank: int):
3
        """Trivial location - logical expert i corresponds to physical expert i"""
4
        ...
5
        physical_to_logical_map = (
6
            torch.arange(0, base_num_physical_experts).repeat(num_layers, 1)
7
            % num_logical_experts
8
        )
9
        physical_to_logical_map = append_trivial_expert_slots(
10
            physical_to_logical_map,
11
            num_physical_experts - base_num_physical_experts,
12
            num_logical_experts,
13
        )

arange(...) % num_logical_experts：物理槽位 i 对应逻辑专家 i % 256。也就是说，当物理专家数超过逻辑专家数时，多出来的槽位是前几个逻辑专家的副本——这是”平凡冗余”的朴素起点，EPLB 会把它算成一个更优的排布。

推理时怎么”选副本”

这是整个 EPLB 里最容易被忽略的一环：冗余专家存在之后，token 路由到逻辑专家 k，应该实际去哪个物理槽位？

expert_location_dispatch.py 里的 ExpertLocationDispatchInfo 有个很关键的注释：

1
    # Whether every rank must pick the same physical expert for a token. True
2
    # without an a2a backend: all EP ranks then run the MoE over the same tokens
3
    # and sum their partial outputs, so a rank-dependent pick counts a replicated
4
    # logical expert several times. With one, each rank dispatches only its own
5
    # tokens and is free to disagree.
6
    rank_invariant: bool = False

以及初始化：

1
        return cls(
2
            ep_dispatch_algorithm=ep_dispatch_algorithm,
3
            rank_invariant=get_exec().moe.moe_a2a_backend == "none",
4
            ...
5
        )