LLM推理中的多卡切分

1. 为什么需要并行和 Chunked Prefill

从存储和计算两个方面看:用资源/通信换并行处理

1.1 存储压力

大模型推理时,模型显存中会存储:

  • 模型权重
  • KV Cache
  • 中间激活值

当我们使用比较大的模型进行长下文推理时,单卡显存会存不下,我们需要并行技术可以进行多卡推理,如:

TP:按 head 切分模型权重
PP:按 layer 切分模型权重
EP:按 MoE expert 切分专家权重
CP:按 token 切分,每个 rank 推理不同的 token,降低 kv cache
DP:KV Cache 按请求分摊,MLA 模型不会复制冗余 kv cache
Chunked Prefill:降低单次 Prefill 的临时显存峰值

1.2 计算压力

我们需要使用并行等技术分摊计算量:

TP:多卡分摊矩阵计算
PP:多卡分摊不同层计算
EP:多卡分摊不同 expert 计算
CP:多卡分摊长上下文 attention / KV 读取
DP:多张 GPU 并行处理不同请求,提高整体吞吐
Chunked Prefill:单机推理时,降低单次 Prefill 计算压力,避免长 Prefill 独占 GPU

1.3 简要总结

长上下文推理在推理引擎中面临的问题可以概括成三类:

问题 典型表现 SGLang 中相关技术
模型权重大 单卡放不下权重,单层 GEMM 过大 TP、PP、EP、量化
上下文长 prefill TTFT 高,kv cache 存储压力大,长请求阻塞短请求 Chunked Prefill、CP、DP、prefix cache、调度预算、PD 分离
MoE 负载不均衡 expert 参数巨大,token 路由不均衡,decode 阶段 all-to-all 延迟明显 EP、DeepEP、TBO/SBO(计算和通信重叠)、EPLB

针对上面的问题,目前已有的应对方案有以下几种:

技术 定位 主要收益 主要代价
Chunked Prefill 调度和显存水位控制技术 不开chunk,长输入的峰值显存高导致oom服务崩溃。开chunk后,分块多轮forward降低峰值显存,解决oom问题,同时可以设置mem-fraction-static更大,增加kv cache上下文存储空间。 ttft升高:单个长请求需要多轮 forward;增加调度、kernel launch 开销
TP 最通用的权重切分手段 切分模型权重,降低单 rank 存储和计算压力 小 GEMM 利用率下降,引入额外通信代价(all-reduce)
PP 跨设备/跨节点承载模型的 stage 并行 降低单卡权重显存压力,并通过流水线方式提高多卡利用率 引入额外通信代价(p2p);存在 pipeline bubble;
DP(Attention) attention 并行手段 复制注意力权重,切分请求到不同 rank,主要降低单 rank 存储压力 引入额外通信代价(all-gather)
EP MoE 专家并行 切分专家权重,降低单 rank 存储和计算压力 引入额外通信代价(all-to-all)
CP 长上下文 prefill 的 sequence/context 并行 长上下文激活和 KVCache 可以分摊到多卡,主要降低单 rank 存储压力 引入额外通信代价(all-gather)

2. SGLang 并行推理概览

2.1 SGLang 请求的推理流程

1

  1. 请求排队:当前 Scheduler 中无请求,该 req 放入 waiting_queue 中

  2. 控制 batch 内请求数量:创建 PrefillAdder,从 waiting_queue 中取出该请求,调用 PrefillAdder::add_one_req() prefill 会被截断分块,can_run_list 中只有这个 req,单个请求被打包成 ScheduleBatch

  3. 分配空闲 kv cache:调用 ScheduleBatch::prepare_for_extend() ,实际上调用 alloc_for_extend(),分配 req_pool_indices 以及为需要 extend 的长度申请真实的 KV 显存 out_cache_loc,并把对应关系记录到的 req_to_token_pool

  4. 模型推理 (TpModelWorker & ModelRunner)

    • 进入模型层 forward()。
    • Attention 执行 : 调用不同的 attention backend 根据 metadata 进行 attention 计算
    • MLP/MoE:进行 MLP/MoE 计算
    • lm_head:得到最终的 logits
  5. Sample:根据 req 中的 sample params 进行采样

  6. 最后得到 batch_result,流式返回给 Detokenizer

2.2 SGLang 中的 Chunked Prefill

2
SGLang 参数里有 max_prefill_tokens、prefill_max_requests、max_running_requests、chunked_prefill_size 等,它们共同决定一次 prefill 能放多少 token 和多少请求。

SGLang 的 chunked_prefill_size 控制一个请求一次 prefill forward 最多处理多少新 token。调度器会将完整 req tokens 截断到当前剩余 budget。

chunked prefill size 实际上是 MoE 层看到的 token 数,除以 dp_size 得到每个 DP rank 实际使用的 chunk size

最开始提出是为了解决TPOT过高(因为continous batching将prefill和decode拼在一起算,导致prefill很长而decode空等待出现尖刺)问题以及长文本下激活值占用过大导致 OOM 问题

  • chunked prefill size增大,TTFT 减小,但需要预留更多显存给中间激活值,支持最大上下文减少
  • chunked prefill size减小,TTFT 增大,但可以支持更大上下文

2.3 SGLang 中并行推理示例

1
2
3
4
5
6
7
8
9
python -m sglang.launch_server \
--model-path <model> \
--tp 16 \
--dp 2 \
--cp 4 \
--ep 16 \
--enable-dp-attention \
--enable-prefill-context-parallel \
--moe-a2a-backend deepep # 开 deepep 后 shared_experts 不切分(tp_size=1)

核心公式:tp_size = attn_dp_size × attn_cp_size × attn_tp_size

1
2
3
4
5
6
7
8
9
10
11
12
13
world_size  = tp_size * pp_size# 全局卡数
tp_size # 一个 pipeline stage 内参与模型并行的 rank 数
pp_size # 有多少个 pipeline stage

# attention 并行参数
attn_dp_size = attention_data_parallel_size
attn_cp_size = attention_context_model_parallel_size
attn_tp_size = tensor_model_parallel_size // attn_cp_size // attn_dp_size

# moe 并行参数
moe_ep_size = expert_model_parallel_size
moe_dp_size = moe_data_model_parallel_size
moe_tp_size = tensor_model_parallel_size // moe_ep_size // moe_dp_size

2.3.1 Rank 布局

单个 rank layout (dp, cp, tp)

全局 rank Attention 坐标 MoE 坐标
0 DPA0, CP0, AttnTP0 EP0
1 DPA0, CP0, AttnTP1 EP1
2 DPA0, CP1, AttnTP0 EP2
3 DPA0, CP1, AttnTP1 EP3
4 DPA0, CP2, AttnTP0 EP4
5 DPA0, CP2, AttnTP1 EP5
6 DPA0, CP3, AttnTP0 EP6
7 DPA0, CP3, AttnTP1 EP7
8 DPA1, CP0, AttnTP0 EP8
9 DPA1, CP0, AttnTP1 EP9
10 DPA1, CP1, AttnTP0 EP10
11 DPA1, CP1, AttnTP1 EP11
12 DPA1, CP2, AttnTP0 EP12
13 DPA1, CP2, AttnTP1 EP13
14 DPA1, CP3, AttnTP0 EP14
15 DPA1, CP3, AttnTP1 EP15

2.3.2 参数设置(开启 DP, TP, CP, EP)

模型:NSA 模型(nsa indexer + MLA)

16 卡示例参数:

1
2
3
4
5
6
7
8
9
python -m sglang.launch_server \
--model-path <model> \
--tp 16 \
--dp 2 \
--cp 4 \
--ep 16 \
--enable-dp-attention \
--enable-prefill-context-parallel \
--moe-a2a-backend deepep # 开 deepep 后 shared_experts 不切分(tp_size=1)

3. 并行技术详细介绍

3.1 TP:Tensor Parallelism

解决单卡计算瓶颈和单卡显存瓶颈
当前模型都比较大,单张卡不能放得下,需要将模型权重按 hidden_dim 切分到不同卡上,通过通信来聚合中间结果 hidden_states,同时减少计算量和中间值的存储空间。

主要有两种并行方式:Column Parallel (列并行) 和 Row Parallel (行并行)。通常它们会成对出现(例如在 MLP 中:先 Column 后 Row),以最小化通信开销。

以 TP=2 为例:单卡计算量、中间激活和权重都减半,代价是 row parallel 层后多一次 All-Reduce;输入在 column parallel 层是复制的(不切分),到 row parallel 层才按列切分

3
shared_experts内部用的是 MergedColumnParallelLinear 和 RowParallelLinear,实际上会走 tp 切分。

如果设置 moe-a2a-backend shared_experts tp_size=1 不切分,复制到 rank

ColumnParallelLinear:
数学原理:
输入 XX 是完整的(复制在所有 GPU 上)。每个 GPU 计算 Yi=XAiY_i=XA_i​。输出 YY被 切分为 [Y1,Y2,...,Yp][Y_1,Y_2,...,Y_p],即每个 GPU 得到输出向量的一部分特征。
典型应用:

  • Attention 的 QKV Projection (QKVParallelLinear)。
  • MLP 的 Gate / Up Projection (MergedColumnParallelLinear)。

RowParallelLinear:
数学原理:
为了匹配矩阵乘法规则,输入 XX 也必须按列切分为 [X1​,X2​,...,Xp​][X_1​,X_2​,...,X_p​](这正好是 ColumnParallelLinear 的输出格式)。每个 GPU 计算 Yi=XiAiY_i=X_iA_i​。注意,YiY_i 的形状与最终输出 YY 相同,但它只是部分和。最终输出 Y=∑YiY = \sum Y_i​ ,需要一次 All-Reduce (Sum) 操作
典型应用:

  • Attention 的 Output Projection (o_proj)。
  • MLP 的 Down Projection (down_proj)。

3.2 PP:Pipeline Parallelism

解决单卡显存瓶颈和计算瓶颈
pipeline schedule:

  • stage 0 处理 req 的输入和 embedding 层
  • 中间的 stage i th 从 i-1 th 接受上一轮的 output 并发送:
  • 非最后一个 PP rank forward 结束后返回 hidden states;最后一个 PP rank 才做 logits processor / lm_head。

3.3 DP Attention:MLA 系列模型使用

解决 MLA 模型的显存瓶颈:kv cache 大量冗余
MLA 的 KV Cache 与 head 数量无关(因为 KV 只有一个头,无法切分所以只能复制多份),使用常规的 TP 会重复存多份 KV Cache,造成显存浪费。DP attention 就是为了解决这个问题。

  • DP Attention 每个 rank 复制 attention 权重,每个 attention 只存储自己分到这部分 token 的 KV Cache(用权重存储换 KV Cache 存储)

SGLang 的 DP Attention:attention 层按 DP 切请求/token;MoE/MLP 仍然可以用 TP/EP。

4
其目的是将一个大的张量并行(TP)组重组为几个较小的 TP 组,这些较小的 TP 组又形成一个专门用于注意力层的新的数据并行(DP)组。如 TP=8,拆分成DP=2,TP=4。

  1. 对于模型中除 MLP 层以外的部分(如 Embedding, Self-Attention),每个数据并行单元(DP_Rank)内部的张量并行规模(TP_Size)设置为 1,即每个 DP_Rank 独立计算这些部分。
  2. 对于 MLP 层,所有 DP_Rank 则共同组成一个大的张量并行组(TP_Group),该组的大小等于原始的 tp_size(图中示例 attn_tp_size = 1,所以恰好等于 DP 规模)。

也就是说,原始 tp_size 被重新解释成:tp_size = attn_dp_size × attn_cp_size × attn_tp_size

3.3.1 DP Attention 的 gather/scatter

由于 attention 只处理本 DP rank 的 token,但 MLP/MoE 需要全局 token 布局,所以 SGLang 需要在 attention/MLP 边界做 gather/scatter。

没有 MoE all-to-all 后端时,MoE 不能直接把 token 发给专家,只能先在 DP 组内把各 rank 的 token 收集成全局布局。为了满足 all-gather / all-reduce 的形状要求,需要根据 prefill/decode 和 token 分布,选择把每个 rank 的 batch padding 到 max_len 还是 sum_len,对应使用 all-gather 或 all-reduce;MoE 算完后再 scatter 回各 DP rank。常见情况是 prefill 用 SUM,decode 用 MAX,但 decode 分布很不均时会退化成 SUM。

3.4 EP:Expert Parallelism

现在参数量比较大的 MoE 模型,专家数量可能是几十到几百个。如果每张 GPU 都用TP做切分,在总卡数相同时不会增加每卡显存;主要问题是每个激活专家都要一次 TP all-reduce,不同专家无法并行,导致通信开销大、计算效率低、扩展性差。

EP 一般流程

MoE 层的 Dispatcher

1
2
3
4
5
6
7
8
9
hidden_states
-> gate/router logits
-> top-k expert ids + top-k weights
-> dispatch (all-to-all):按 expert id 把 token 发到专家所在 GPU
-> 本地按专家分组 -> grouped GEMM(专家 FFN)
-> [若专家内部有 TP:TP all-reduce / reduce-scatter + all-gather]
-> combine (反向 all-to-all):把专家输出送回原 token 所在 GPU
-> 按 top-k weights 加权求和
-> output

3.4.1 Standard 模式(无 moe-a2a-backend)

Standard 模式 不做 All-to-All 通信,而是采用 AllGather 全量 token + 本地只算自己负责的专家 + AllReduce 归约的策略。

5

3.4.1.1 特点
  1. 无 token 物理搬移:dispatch 阶段不移动 token 数据,只重映射 expert ID
  2. 每个 rank 处理全部 token:但只在本地专家上计算,其余跳过(id=-1 的 token 贡献为 0)
  3. 靠 AllReduce 聚合:各 rank 的部分结果相加得到最终结果
  • 路由计算 (TopK Router)

    • 计算 Token 路由,输出全局专家 ID (topk_ids) 及其对应权重。
  • 本地分发 (Dispatch)

    • 将全局专家 ID 映射为本地专家 ID(属于当前 Rank 的分配 Local ID,非本 Rank 的标记为 -1)。
  • 专家执行 (MoE Runner)

    • 重排 (Pre-permute): 按专家 ID 对 Token 排序,直接过滤/跳过 ID 为 -1 的 Token。
    • 计算 (Grouped GEMM): 仅对分配到本地的专家执行核心矩阵运算 (W13 + W2)。
    • 还原 (Post-permute): 将计算后的 Token 恢复至原始输入顺序。
  • 状态合并 (Combine)

    • 输出当前 Rank 计算完成的隐藏层状态 (Hidden States)。
  • 全局聚合 (AllReduce)

    • 跨所有 EP/TP Rank 执行通信 (AllReduce),将各节点分散计算的结果相加,得到最终完整输出。

3.4.2 DeepEP 模式(–moe-a2a-backend deepep)

DeepEP 模式使用真正的 All-to-All 通信:将 token 物理发送到目标专家所在的 rank,每个 rank 只接收和处理路由到自己的 token。
传统固定形状 All-to-All:
Rank0 → Rank1: 5 个 token(其中 3 个是 padding)
Rank0 → Rank2: 5 个 token(其中 5 个是 padding,Rank2 根本不需要)
Rank0 → Rank3: 5 个 token(其中 4 个是 padding)

DeepEP 动态 All-to-All:
Rank0 → Rank1: 2 个真实 token
Rank0 → Rank2: 0 个(根本不发)
Rank0 → Rank3: 1 个真实 token

分为两种模式 Normal 和 LowLatency 模式

  • prefill 使用 normal 模式
  • decode 使用 low-latency 模式

这是两者最根本的区别,源于对硬件拓扑的不同利用方式。

Normal 模式(两级接力):它利用了节点内 (Intra-node) NVLink 的高带宽(~160 GB/s) 和节点间 (Inter-node) RDMA 的相对低带宽(~50 GB/s) 的不对称性。其策略是:

第一级 (RDMA):源 GPU 只将一份数据通过 RDMA 发送到目标节点上编号相同的 GPU(例如,0号卡发给0号卡)。这大大减少了昂贵的跨节点数据传输量。
第二级 (NVLink):接收节点上的 GPU 再通过高速 NVLink,将数据“扇出”给节点内所有需要此 token 的 GPU。
这种设计以牺牲一点路径长度为代价,换取了对稀缺 RDMA 带宽的高效利用,非常适合需要传输大量 token 的 Prefill 阶段。
Low-Latency 模式(一级直达):它完全放弃了 NVLink 转发,所有通信都通过纯 RDMA 完成。其策略是:

直接发送:每个 token 的通信请求,都通过 RDMA 直接发送到目标专家所在的 GPU,不经过中间节点转发。
虽然这可能增加 RDMA 的跳数,但路径最短,没有中间环节的等待和开销,因此单次通信延迟极低,完美契合 Decode 阶段每次只处理少量 token 的场景。

3.4.3 通信计算 overlap 机制

目的:为了提高 GPU 利用率

3.4.3.1 TBO(Two-Batch Overlap 双批次重叠)

TBO 工作原理:

  • 将一个 batch 拆成 2 个 sub-batch,在 attention 和 dispatch/combine 之间交错执行,隐藏通信延迟
  • Sub-batch 0 的 MoE compute 与 Sub-batch 1 的 dispatch 通信重叠
  • 需要 2 个独立 dispatcher(各自有独立的 handle 和 buffer 状态)

7

重叠点:
当微批次 0 在进行 MoE 的 dispatch(发送 token)通信时,微批次 1 正在执行其注意力(Attention)的计算。反之亦然。

效果:
计算流(Compute Stream)和通信流(Comm Stream)被交织在一起,原本串行的“计算-通信-计算”链条被打破,通信时间被有效利用起来做有用功。

TBO 并非在所有场景下都能带来收益,它的效果与批次大小和部署环境密切相关。

最佳场景:
大规模专家并行(EP):特别是在多节点、跨节点带宽受限的环境下,通信瓶颈越突出,TBO 隐藏延迟的效果越好。
足够的批次大小(Batch Size):TBO 只有在批次足够大时才能有效覆盖通信开销。在 SGLang 的测试中,当每个 DP rank 的批次大小超过 320 时,TBO 能带来约 5%~10% 的性能提升;但在小批次(如 bs=60)下,TBO 反而可能导致性能显著下降。

默认策略:
在多数实现中,TBO 默认仅用于 Prefill 阶段。因为 Prefill 阶段计算密集,适合用计算来隐藏通信。若强行用于 Decode 阶段,由于 Decode 本身计算量小,拆分批次带来的开销可能超过收益,导致吞吐量下降。

3.4.3.2 SBO(Single-Batch Overlap)

SBO 工作原理:

TBO 在超高并发下也可能遇到瓶颈。有实践发现,在特定硬件(如 H20)上,当并发极高时,TBO 的双批次策略可能导致延迟“爆炸”,反而变慢。为此,社区(如蚂蚁集团与 SGLang、DeepSeek 合作)提出了 SBO(Single-Batch Overlap) 作为改进方案。SBO 的优化在于,它不再等待整个批次完全准备好,而是在单个批次内部进行更细粒度的重叠,例如在计算 shared experts 的同时就启动 dispatch 通信,或者在计算完一个 block 后就立即发送,从而实现更极致的流水线。

在单 batch 的 MoE 层内部,用多 stream 将通信(dispatch/combine)与计算(gemm/shared expert)重叠。

  • combine 操作与 down_gemm 重叠
  • dispatch 与 shared expert 计算重叠

8
combine 以 block_m=64 为单位,down_gemm 每产出一块,combine 立刻消费一块,两者在 GPU 上并行执行(不同 SM 集合,不同 CUDA stream)

上面是非 blackwell 的实现,Blackwell 在 _pre_combine_hook 中运行 shared expert:

Flag H 卡值 含义
enable_combine_down_gemm_two_stream_overlap True(如果 backend 是 deep_gemm) combine 和 down_gemm 双 stream 重叠
enable_dispatch_shared_one_stream_overlap True dispatch 和 shared experts 同 stream 重叠
enable_combine_shared_two_stream_overlap False —
fuse_shared_experts_inside_sbo True shared experts 在 SBO 窗口内执行

3.4.4 EPLB(Expert Parallelism Load Balancer)

在 MoE 的专家并行中,不同的“专家”被放置在不同的 GPU 上。但由于输入数据的差异,某些“热门专家”会被大量 token 选中,其所在的 GPU 就会过载;而“冷门专家”所在的 GPU 则相对空闲。整个系统的速度取决于最慢的那块 GPU,这会造成严重的资源浪费。

EPLB 的核心思想是:
动态地为高负载的“热门专家”创建冗余副本,并将这些副本智能地分配到不同 GPU 上,从而将负载从过载的 GPU 分摊出去,最终最小化所有 GPU 中的最大负载。

EPLB 的工作流程分为三个主要步骤:

  1. 收集负载数据:在模型运行时,持续统计每个专家被路由到的 token 数量,作为其“热度”指标。这些数据通常通过历史统计的移动平均值来估计。
  2. 执行均衡算法:基于收集到的负载数据,EPLB 算法会计算出一个新的专家放置方案。这个方案决定了哪些专家需要被复制,以及所有专家(包括副本)应该被放置在哪些 GPU 上。
  3. 动态重排专家:根据算法计算出的方案,在运行时动态地重新排列专家及其副本在 GPU 上的分布,以实现负载均衡。

我们引用 DeepSeek 官方仓库中的一个经典例子来具体说明。

场景设定:
一个两层的 MoE 模型。
每层有 12 个逻辑专家(编号 0-11)。
部署在 2 个节点上,每个节点有 4 个 GPU,共 8 个 GPU。
为每层引入 4 个冗余专家,因此每层总共有 12 + 4 = 16 个“物理专家副本”需要放置。
假设的负载情况:经过数据收集,假设在第1层中,专家 4、5、1、10 是“热门专家”,负载显著高于其他专家。

EPLB 的分配过程(采用分层负载均衡):
节点间均衡 (Step 1):EPLB 先将 12 个逻辑专家分成 4 组,例如 (0-2), (3-5), (6-8), (9-11)。然后,将这 4 个组分配到 2 个节点上,确保每个节点分到的专家总“基础负载”大致相等。这是一个“背包问题”,可用贪心算法解决。
复制热门专家 (Step 2):在节点内部,EPLB 识别出热门专家。假设节点A分到了专家组 (3-5),其中专家 4 和 5 是热门的;节点B分到了专家组 (9-11),其中专家 10 是热门的(专家 1 可能被分到了另一个组)。EPLB 会为这些热门专家创建副本。
节点内均衡 (Step 3):最后,EPLB 将所有物理专家副本(包括原始专家和冗余副本)打包到该节点内的 4 个 GPU 上。它同样将此视为“背包问题”,目标是最小化每个 GPU 上的总负载。例如,它可能会将专家 4 的一个副本放在 GPU0,另一个副本放在 GPU2,从而分摊其负载,避免任何单个 GPU 过载。

通过逻辑专家和实际物理专家的映射,原本因专家 4 和 5 而可能过载的某个 GPU,其负载被多个副本分摊到了不同 GPU 上,实现了节点内和节点间的负载均衡。

动态统计专家热度 → 为热门专家复制多份完整权重副本 → 把所有专家副本重新排布到不同 GPU → 让路由把请求分散到副本上,从而最小化最大 GPU 负载。
缺点:需要事先申请冗余专家显存,降低了显存上限。

3.5 CP:Context Parallelism

将长上下文分发到多个 GPU 上计算,所以用了特殊的attention:

9

  1. 基础版:Ring Attention(环状注意力)
    QKV序列分割:把长序列切成 N 块,分给 N 个 GPU(Rank)。
    设备环状结构:GPU 之间形成一个环。每个 GPU 计算当前块的 Q 与当前块的 KV 的注意力,然后把自己的 KV 块发给下一个 GPU,同时接收上一个 GPU 发来的 KV 块。
    完成条件:经过 N 次传递,每个 GPU 都看到了序列中所有的 KV 部分,计算完成。
    开销隐藏:由于是异步通信,GPU 可以在等待接收下一个 KV 块的同时,计算当前块的注意力,从而把通信时间藏在计算时间背后(Overlap),实现近零开销。
    核心痛点:Ring Attention 假设序列是均匀切分的(比如序列长 12,4 个 GPU,则 GPU0 拿 0-2,GPU1 拿 3-5,以此类推)。但在因果注意力(Causal Attention) 中,前面的 token 只能看自己,后面的 token 要看前面所有 token。这就导致排在后面的 GPU 计算量远大于排在前面的 GPU,产生严重的负载不均衡,环上的其他 GPU 必须等最慢的那个 GPU,流水线效率大打折扣。(上图中白色部分随着轮次越来越少)

10
11
2. 改进版一:Striped Attention(条纹注意力)
Striped Attention 是为了解决上述因果掩码导致的负载不均衡而提出的。
它的做法不是按连续块切分,而是按“条纹”交错切分:
假设序列长 16,4 个 GPU。
GPU0 拿 token:0, 4, 8, 12
GPU1 拿 token:1, 5, 9, 13
GPU2 拿 token:2, 6, 10, 14
GPU3 拿 token:3, 7, 11, 15
效果:在因果注意力下,每个 GPU 拿到的 token 分布在序列的各个位置,每个 GPU 需要计算的注意力范围(即它需要看到的前序 KV 数量)变得大致相等。这就极大地缓解了负载不均衡,让环上的计算更加同步,提升了整体吞吐量。(可以看到图上白色部分每一轮大致相等)
任意两张卡上的 Q 与 K 之间的 attention mask,都是一个完整的方形,而且每张卡上的计算量也是相对平均的
每个 DP 组分配 chunked_prefill_size // dp_size,每个 CP rank 在组内又拿到 1/cp_size token。

12
3. 改进版二:Zigzag Attention(锯齿注意力)
Zigzag Attention 是另一种解决因果掩码负载不均衡的切分策略,常见于后来的一些长上下文优化框架中。
原始块顺序:block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7
Zigzag 重排后:block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4

3.5.1 两种 CP 方式

主要为了降低 TTFT

  • CP 只用于 prefill(extend)阶段,decode 阶段不使用 CP(每次只生成 1 个 token,无法按 seq 切分)。
  • Decode 时每个 rank 需要读取完整的历史 KV cache,因此 prefill 结束后每个 rank 的 kv_pool 必须包含完整序列的 KV。

权重在 CP 维度上被复制了 cp_size 倍;但 attn 计算时每张 rank 只算 seq/cp_size 份 token(按 seq 切),CP 组内的 ranks 处理不同的 token(seq 按 CP 切分),它们的 o_proj 输出通过最后的 AllGather 合并,而不是 AllReduce
支持两种 CP 方式

round-robin-split:token_idx % cp_size,rank i拿到所有下标模 cp_size 等于 i 的 token
* batch 里所有 req 的 token 拼接成一维后的张量进行取模
13

in-seq-split:序列切成 2*cp_size 个连续 block,rank i 取 block i 和 block 2*cp_size-1-i 只支持单个 req 进行切分
14

约束差异

round-robin-split in-seq-split
batch 约束 支持 B ≥ 1 强制 B=1
切分约束 S % N == 0 S ≥ 2N
负载均衡 token 均衡 使用 zigzag 保证 Attention 计算均衡

4. 总结

4.1 策略对指标的影响

TTFT / Prefill

并行 通常影响 原因
TP 降低 TTFT,但收益与通信开销相关 **优势:**GEMM 权重读取和计算大致分到多个 GPU;
**代价:**o_proj/MLP 结果需要 all-reduce/gather。
长 prefill 中 GEMM 足够大,TP 通常有收益;小 batch/过大 TP 时通信和小 GEMM 效率会反噬。
PP 单请求 TTFT 不一定下降,常可能上升 **优势:**GEMM 权重读取和计算大致分到多个 GPU;
**代价:**stage 边界通信和排队可能抵消收益。
DP Attention 高并发 prefill 更有价值,单请求收益有限 **优势:**多请求 continuous batching 下,每个 attention DP rank 处理更少 KV/token ;
**代价:**dp_gather_partial / dp_scatter / reduce_scatter 会引入同步。
EP MoE prefill 常降低显存和专家权重读压力 **优势:**experts 分布到不同 GPU,每张卡只放/读部分 expert 权重。
代价:all-to-all 通信;若 batch 很小,dispatch/combine 延迟会更明显。
CP 长上下文 prefill 开启降低每个 rank 的 kv cache 压力 **优势:**Q token 被切到不同 CP rank,单 rank attention query 数下降;SGLang 用 zigzag/round-robin 缓解 causal mask 下前后块负载不均。
代价: K/V/index K 的 CP all-gather。

TPOT / Decode

并行 通常影响 原因
TP 降低 TPOT,但收益与通信开销相关 **优势:**GEMM 权重读取和计算大致分到多个 GPU;
**代价:**o_proj/MLP 结果需要 all-reduce/gather。
PP 不适合降低单请求 TPOT 每生成一个 token 都要按 stage 顺序过完整模型,并跨 stage 传 hidden。对单请求 decode latency,PP 往往增加 pipeline stage 等待和通信;它更偏容量扩展/多请求吞吐。
DP Attention 对高并发 TPS 提升有帮助,对单请求有限 **优势:**多请求 continuous batching 下,每个 attention DP rank 处理更少 KV/token ;
**代价:**dp_gather_partial / dp_scatter / reduce_scatter 会引入同步。
EP 高并发/大 batch decode 有收益,使用 tbo/sbo 实现通信和计算重叠 **优势:**experts 分布到不同 GPU,每张卡只放/读部分 expert 权重。
代价:all-to-all 通信;若 batch 很小,dispatch/combine 延迟会更明显。

decode cp 不使用,设置参数也会被忽略

4.2 选用何种并行组合

4.2.1 Dense 模型

  • 选型推荐:TP (节点内) + PP (跨节点)
  • 依据:
    • Dense 模型的计算分布均匀,没有专家路由机制。首要解决的是单卡显存(装不下权重和 KV Cache)与单层算力不足的问题。
    • TP 是切分 Dense 模型最直接的方式,有效降低单次矩阵运算的维度。

4.2.2 MoE + NSA 模型

  • 选型推荐:EP + TP + DP Attention
  • 依据:
    • **EP:**通过 EP 将海量专家参数分散到不同 GPU,避免全量复制带来的显存问题,同时提高 expert 的计算并行度。
    • DP Attention: 在多请求并发时大幅降低单个 Rank 的 KV Cache 压力。
    • TP: 用于处理 MoE 模型中未被路由的稠密层(如 QKV 投影、Dense MLP 层)。

4.2.3 降低 TPOT

  • 选型推荐:最大化单机 TP + (MoE场景下引入 DP Attention) + 大并发下 EP
  • 依据:
    • TP 提升聚合显存带宽,显著压低单次 GEMV 耗时。
    • DP Attention 允许更大的 Batch Size,摊薄每个 Token 的访存成本。
    • 大并发下 EP 可以分摊权重读取成本

4.2.4 降低 TTFT

  • **选型推荐:TP + EP + 按需引入 CP (针对长文本) **
  • 核心依据:
    • **EP **将 Token 分发到不同 GPU 的特定专家上并行计算,实现通信和计算的 overlap
    • TP 能够将大矩阵切块,利用多卡算力直接并行计算
    • CP 可以多卡并行计算 req。