Skip to content

DP Attention ​

Data Parallelism Attention(DP Attention,也称 DPA)只在 Attention 区域沿请求或 Batch 维度做数据并行,而 MLP 或 MoE 仍使用 Tensor Parallelism(TP)或 Expert Parallelism(EP)。它的核心目标不是减少 Attention 计算量,而是让不同请求的 KV Cache 归属不同 DP Rank,避免低 KV Head Attention 在 TP 下重复保存缓存。

本文先说明 DP Attention 的通用数据流,再以 SGLang 为例介绍通信与调度实现。MLA、MoE 和 SGLang 参数属于具体应用;“Attention 按请求切分,其他层重新汇合计算”才是这一机制的核心。

TP 为什么会复制 KV Cache ​

TP 通常沿 Attention Head 和隐藏维度切分权重。对于 MHA,如果 KV Head 数量足够且可被 TP 度整除,每个 TP Rank 可以只保存本地 KV Heads;但对于 MQA 或 Multi-Head Latent Attention(MLA),可切分的 KV Head 或 Latent KV 数量很少,继续增加 TP 度可能只能复制缓存。

以 MLA 为例,设模型有 L 层,请求长度为 N,压缩后的 KV Latent Dimension 为 dc,RoPE Key Dimension 为 dR,每个元素占 s Byte。一次请求的 KV Cache 约为:

SKV=LN(dc+dR)s

如果该缓存无法沿 P 个 TP Rank 切分,同一请求的物理缓存占用约为 PSKV。这不会改变单卡看到的上下文,却会占用原本可以服务更多请求的显存。

DP Attention 改为把请求集合划分给 D 个 Attention DP Rank:

R=R0∪R1∪⋯∪RD−1,Ri∩Rj=∅

Rank i 只执行 Ri 中请求的 Attention,并只保存这些请求的 KV Cache。在常见的 Attention TP = 1 配置下,同一批请求的集群级物理缓存由近似 D∑rSKV(r) 降为 ∑rSKV(r);若负载均衡,每个 Rank 约保存其中的 1/D。

这里切分的是请求,不是单个请求的上下文。一个请求的 KV 仍完整保存在所属 DP Rank;若要把同一请求的上下文或 KV 沿设备切分,应使用 Context Parallelism。

DP Attention 也可用于 GQA 等标准 Attention,但收益取决于 KV Head 是否已经能被 TP 有效切分。MLA、MQA 或 KV Heads < TP Size 时,消除缓存复制的价值通常更明显。

混合并行拓扑 ​

DP Attention 不是整模型 DP,而是同一 Decoder Layer 内的混合并行。设一个执行组包含 P 个 Rank,Attention DP、Attention CP 的并行度分别为 DA、CA,则 Attention TP Size 为:

TA=PDACA

例如 SGLang 中设置 --tp 4 --dp-size 4 --ep 4 --enable-dp-attention,且 Attention CP 为 1,则产生 4 个 Attention DP Rank,每个 Rank 的 TA=1;MoE Experts 同时分片到 4 个 EP Rank。Attention 权重在 DP Rank 间复制,本地 Batch 和 KV Cache 彼此独立,MoE token 则跨 4 张 GPU 执行 All-to-All Dispatch 与 Combine。

DP Attention 与 EP 单层数据流:四个 Attention DP Rank 独立维护 Batch 和 KV Cache,MoE token 通过 All-to-All 进入四个 Expert Shard

因此,DP Attention 同时包含两种数据布局:

  • Attention 区域:每个 DP Rank 持有 local_tokens × hidden_size,请求和 KV Cache 互不重叠。
  • MoE 区域:Router 为本地 token 选择 Top-k Experts,All-to-All 将 token 发送到对应 EP Rank,计算后再 Combine 回原 DP Rank。

图中固定使用 DP Attention + EP:Attention 输出不做全量复制,而是由 MoE Router 生成 expert_id 和 Routing Weight,All-to-All 只把命中某个 Expert 的 token 发给其所属 EP Rank。Combine 再按 Routing Weight 聚合 Top-k Expert 输出,并恢复原 DP Rank 的本地 token 顺序。

与纯 TP 相比,DP Attention 用 Attention 权重复制和边界通信换取 KV Cache 容量;与普通 DP 相比,它没有复制一套完整的 MLP 或全部 Experts,而是让这些大权重继续跨设备分片。

单层执行过程 ​

以一个 Attention → MoE Decoder Layer 为例,一轮 Forward 按以下顺序执行:

  1. 请求分配:Router 或内置负载均衡器把请求固定到一个 Attention DP Rank,Prefix Cache 和 KV Cache 也随请求留在该 Rank。
  2. 独立组 Batch:各 Rank 使用自己的 Waiting Queue 和 Running Batch,可以分别处于 Prefill、Decode 或 Idle 状态。
  3. Collective 对齐:即使 Rank 3 没有本地请求,也必须调用同一轮 MoE All-to-All;它的本地发送量为 0,但作为 EP Rank 仍可能接收其他 Rank 路由到本地 Experts 的 token。这不是聚合 Batch,不会移动 KV Cache。
  4. 本地 Attention:每个 Rank 只读取和更新本地请求的 KV Cache,得到 Xi∈Rni×H。
  5. Dispatch:MoE Router 为 Xi 中的 token 选择 Top-k Experts,All-to-All 将 token 发送到对应 EP Rank。
  6. Expert 计算:每个 EP Rank 执行本地 Expert Shard,并按 Routing Weight 生成部分输出。
  7. Combine:All-to-All 将 Expert 输出送回请求所属 DP Rank,聚合后进入下一层本地 Attention。

这种设计允许 Rank 0 做 Prefill、Rank 1 做 Decode、Rank 3 暂时空闲,但“独立 Batch”不等于“完全异步前进”。所有 Rank 每层仍要在 MoE All-to-All 处会合;空闲 Rank 也必须生成 Idle Batch 参与相同的通信序列,否则会发生 Collective Mismatch 或死锁。

不等长 Batch 的 All-to-All ​

设 DP Rank i 当前有 ni 个 token,每个 token 选择 k 个 Experts。令 ci,j 表示 Rank i 发往 EP Rank j 的 token 副本数,则:

∑jci,j=kni,mj=∑ici,j

mj 是 EP Rank j 实际收到并交给本地 Experts 的 token 数。All-to-All Dispatch 只按路由结果移动这些 token,不会在每个 Rank 上复制完整的全局 Batch;Combine 再执行反向 All-to-All,把 Expert 输出送回原 DP Rank。

各 Attention DP Rank 可以有不同的 ni。图中 Rank 3 的 n3=0,因此它没有本地 Expert token 要发送,但仍必须调用本轮 All-to-All,并可能接收其他 Rank 发往本地 Experts 的 token。另一方面,即使各 ni 接近,Router 也可能让某些 Experts 更热门,使对应 mj 远大于其他 Rank;此时最忙的 EP Rank 决定 MoE 层完成时间。

SGLang Scheduler 会在 Forward 前用一个很小的控制张量交换各 Rank 的 token 数量和运行状态,以生成必要的 Idle Batch 并保证 Collective 调用顺序一致。这是控制面协调,不属于图中的 Hidden States 或 KV Cache 数据流。

📄 SGLang:Attention 并行度与 Rank 映射
python
def compute_dp_attention_world_info(
    enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size=1
):
    attn_dp_size = dp_size if enable_dp_attention else 1
    attn_tp_size = tp_size // attn_dp_size // attn_cp_size
    attn_tp_rank = tp_rank % attn_tp_size

    if not enable_dp_attention:
        attn_dp_rank = 0
    else:
        attn_dp_rank = tp_rank // (attn_tp_size * attn_cp_size)

    return attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size

与 DP、TP 和 CP 的区别 ​

方式请求如何分配权重布局KV Cache跨 Rank 通信
普通 DP不同 Replica 处理不同请求每个 Replica 有完整执行能力,可在内部继续使用 TP/PP每个请求只在所属 ReplicaReplica 间通常无逐层通信
Attention TP同一 Batch 由多个 Rank 共同处理Attention 和 MLP 均沿张量维切分尽量按 KV Head 切分;无法切分时复制几乎每层 All-Reduce 等
DP AttentionAttention Rank 处理不同请求,MLP/MoE 重新汇合Attention 可复制或使用较小 TP,MLP/MoE 使用 TP/EP按请求归属到 Attention DP RankAttention 与 MLP/MoE 边界 Gather/Scatter 或 All-to-All
CP多个 Rank 共同处理同一请求的不同上下文分片模型权重可复制或再叠加 TP单个请求的 KV 沿序列切分KV 交换或 Attention 结果归并

最容易混淆的是普通 DP 与 DP Attention。普通 DP 的每个 Replica 可以独立完成整次 Forward;DP Attention 的 Rank 只在 Attention 区域独立,进入共享 MLP/MoE 后仍属于同一个执行组。因此,普通 DP 更像复制多套引擎,DP Attention 更像在一层内部切换并行坐标系。

收益、代价与适用场景 ​

DP Attention 的主要收益是 KV Cache 不再随无效的 TP 复制倍增。释放出的显存可用于提高并发序列数或 Batch Size,尤其有利于 Decode 吞吐。对于 MoE 模型,Attention 使用 DP、Experts 使用 EP,还能让 KV 容量扩展与专家权重分片同时成立。

它不会自动带来以下收益:

  • 不降低单个请求的 Attention 复杂度,也不缩短必须执行的 Decoder 层数。
  • 不保证降低单请求延迟;低并发时,本地 GEMM 变小和额外通信可能使延迟上升。
  • 不一定减少权重显存;Attention 权重可能从 TP 分片变为在 DP Rank 间复制。
  • 不消除跨 Rank 同步;最慢的 Attention Rank 会决定共享 MLP/MoE 何时开始。

最后一点在 Prefill 与 Decode 混跑时尤其明显。长 Prefill 的 Attention 计算远慢于单步 Decode,Decode Rank 会在下一次共享层通信前等待 Prefill Rank。Prefill-Decode Disaggregation(PD 分离)可以把两种阶段放入不同资源池,使同一 DP Attention 组内的 Batch 形态更接近,但会新增 KV 传输成本。

因此,DP Attention 更适合以下条件:

  • MLA、MQA 或较少 KV Heads 导致 TP 下 KV Cache 明显复制。
  • 高并发 Decode 受 KV Cache 容量限制,而不是单请求计算延迟限制。
  • MoE Experts 需要 EP 分片,无法简单复制多套完整模型。
  • Rank 间具备高速互联,并能通过负载均衡避免 token 数量长期倾斜。

如果模型的 GQA 已能按 TP 完整切分 KV Heads、并发很低或跨节点通信较慢,普通 TP 或外层 Replica DP 可能更简单。最终选择应同时测量每卡 KV 占用、TTFT、TPOT、吞吐和集合通信时间。

SGLang 配置示例 ​

SGLang 使用 --enable-dp-attention 开启 DP Attention,并通过 --dp-size 指定 Attention DP Size。以下示例在 8 张 GPU 上让 Attention 使用 DP=8、MoE 使用 EP=8:

bash
python -m sglang.launch_server \
    --model-path deepseek-ai/DeepSeek-V3 \
    --tp 8 \
    --dp-size 8 \
    --ep 8 \
    --enable-dp-attention \
    --moe-a2a-backend deepep \
    --moe-runner-backend deep_gemm

当前实现要求 dp_size > 1,且 tp_size 能被 dp_size × attn_cp_size 整除。不要仅凭参数名假设所有模型和 Backend 都支持任意组合;模型权重布局、Attention Backend、MoE Backend、Speculative Decoding 和 PD 分离都可能增加额外约束,应以对应 SGLang 版本的启动校验为准。

总结 ​

DP Attention 本质上是一种分层混合并行:Attention 沿请求维度切分,让每个 Rank 独立维护 Batch 和 KV Cache;MLP 或 MoE 再切回 TP/EP 拓扑,共同处理这些 token。它用 Attention 权重复制、激活重分布和同步等待,换取低 KV Head 模型更高的 KV Cache 容量。

判断它是否合适,关键不是模型是否名为 MoE,而是三个问题:TP 是否正在复制 KV Cache、释放的显存能否转化为有效并发,以及 Attention Rank 的负载差异是否会被共享层同步放大。

参考 ​

  1. DeepSeek-V2:A Strong, Economical, and Efficient Mixture-of-Experts Language Model
  2. SGLang v0.4:Data Parallelism Attention for DeepSeek Models
  3. SGLang:DP、DPA 与 Model Gateway 指南
  4. SGLang:DP Attention 通信实现
  5. SGLang:DP Attention Collective 对齐