Deploying DeepSeek with PD Disaggregation and Large-Scale Expert Parallelism on 96 H100 GPUs
Deploying DeepSeek with PD Disaggregation and Large-Scale Expert Parallelism on 96 H100 GPUs
发表时间: 2025-05 · Blog post by The SGLang Team (lmsys.org)
原文: https://lmsys.org/blog/2025-05-05-large-scale-ep
作者/机构:The SGLang Team
速读
一句话结论
本文通过在 SGLang 框架中引入预填充-解码分离(PD Disaggregation)与大规模专家并行(EP)等优化,在 96 张 H100 GPU 上成功部署 DeepSeek-V3,实现了单节点每秒 52.3k 输入和 22.3k 输出 token 的吞吐量,成为首个在大规模下几乎追平 DeepSeek 官方性能报告的开源实现。
要解决什么问题
DeepSeek 独特的混合专家(MoE)和多头潜在注意力(MLA)架构在传统推理系统中面临严重的扩展与显存卡点。首先,传统的预填充与解码混合调度机制会导致新来的预填充任务频繁打断解码任务,且无法兼容 DeepEP(专为 MoE 设计的通信库)的调度模式——DeepEP 在预填充期需要高吞吐的常规调度,在解码期需要支持 CUDA Graph 的低延迟调度,两者无法在同一通信组内共存。其次,在处理稠密 FFN 层时,若采用常规的张量并行(TP),由于其中间维度(18,432)无法被高并行度整除,会引发严重的显存碎片化和计算不对齐,且每层需要两次 All-Reduce,通信开销极大。最后,在稀疏 FFN 层采用专家并行(EP)时,由于不同 token 对专家的调用频率差异巨大,会导致严重的 GPU 负载不均衡,部分设备闲置等待,拖慢整体流水线并拉高显存峰值。
怎么做的
核心思路是通过架构解耦、定制化并行策略与计算通信重叠来彻底释放硬件算力。首先,系统实现了预填充-解码分离(PD Disaggregation),将两个阶段分配给独立的服务器。预填充服务器计算完 KV Cache 后,通过 RDMA(远程直接内存访问)非阻塞地传输给解码服务器。这不仅消除了任务打断,还让预填充和解码能分别使用 DeepEP 的常规调度和低延迟调度,且完美兼容 DP Attention(一种消除跨设备 KV Cache 冗余的数据并行注意力机制)。其次,针对稠密 FFN 和 LM Head,系统放弃了纯 TP,改用数据并行(DP)。作者推导了单层 Transformer 在 DP 和 TP 混合下的显存占用,当假设 $\text{DP}=\text{TP}$ 时,使显存最小化的最优 TP 维度为:
由于预填充期 $k=0$ 且 token 数量极大,最优 TP 往往小于 3。因此采用 DP 不仅更省显存、避免了碎片化,还将通信开销从两次 All-Reduce 降级为一次 Reduce-Scatter 和一次 All-Gather。第三,引入双批次重叠(TBO)机制,将一个批次拆分为两个微批次,让计算(如注意力层)与 DeepEP 的通信分发过程相互掩盖;为了防止预填充期的通信阻塞 CPU,系统强制在发起通信前先向 GPU 提交计算任务。最后,引入专家并行负载均衡器(EPLB),在原有的 256 个专家外增加 32 个冗余专家(共 288 个),通过复制高频专家或打包低频专家来抹平 GPU 间的负载差异;并辅以 DisposableTensor 工具,绕过 Python 引用计数,在张量用完后立即释放 CUDA 显存。
效果如何
实验在 12 个节点(共 96 张 H100 GPU,InfiniBand 网络)上部署 DeepSeek-V3。对比基线包括:SGLang 结合 TP16 的纯张量并行方案、SGLang 结合 PD 分离与默认 EPLB 方案、SGLang 结合模拟多 Token 预测(MTP)方案,以及 DeepSeek 官方公布的 Profile 数据。在预填充阶段(4 节点,EP32,4K 上下文),本方法单节点吞吐量达 50,302 tokens/s,比 TP16 基线提升 3.3 倍,若假设负载绝对均衡,仅落后官方数据 5.6%。在解码阶段(9 节点,EP72,2K 上下文,批大小 256),单节点吞吐量达 22,282 tokens/s,比 TP16 基线提升 5.2 倍;在模拟 MTP 场景下也达到了 17,373 tokens/s,仅落后官方 6.6%。本地部署成本折合每百万输出 token 仅 0.20 美元,约为官方 API 的五分之一。作者也指出了当前局限:为了追求极致吞吐,首字延迟(TTFT)高达 2 到 5 秒,token 间延迟约 100 毫秒,暂不适合对实时性要求极高的场景;此外,受限于 96 张 GPU 的总显存,目前无法支持超长上下文,且系统暂不支持灵活配置大于 1 的 TP 维度。
主要贡献
DeepSeek 作为一个备受欢迎的开源大型语言模型(LLM),因其强大的性能而受到赞誉。然而,其庞大的参数规模以及采用多头潜在注意力(Multi-head Latent Attention, MLA)和混合专家(Mixture of Experts, MoE)的独特架构,要求具备先进的系统才能实现高效的大规模服务。本文阐述了 SGLang 团队如何通过系统优化,使 SGLang 的推理系统性能与 DeepSeek 官方报告的性能相匹配。
本文的核心贡献与创新点如下:
- 系统部署与规模:在 Atlas Cloud 的 12 个节点上(每个节点配备 8 张 H100 GPU,共 96 张 H100)成功部署了该实现。
- 核心技术引入:系统采用了 Prefill-Decode(PD)分离架构以及大规模专家并行(Expert Parallelism, EP),并全面支持了 DeepEP、DeepGEMM 和 EPLB 的完整功能。
- 极致的吞吐性能:对于 2000 token 的输入序列,系统实现了每个节点每秒 52.3k 输入 token 和 22.3k 输出 token 的处理速度。这是目前已知首个在吞吐量上几乎媲美 DeepSeek 官方博客报告的大规模开源实现。
- 显著的成本效益与提升:本地部署该实现的成本转化为每百万输出 token 仅需 0.20 美元,约为官方 DeepSeek Chat API 成本的五分之一。与使用相同资源的传统张量并行(Tensor Parallelism, TP)相比,该优化策略将输出吞吐量提升了高达 5 倍。
- 全面开源:本文涉及的所有实验、并行设计、优化方法以及代码均已完全开源,供社区访问和进一步开发。
并行化设计
高效的并行化对于管理 DeepSeek 架构的计算复杂性和内存需求至关重要。本节概述了优化关键组件(注意力层、稠密前馈网络、稀疏前馈网络和语言模型头)的方法。每个组件都利用了定制的并行策略,以增强可扩展性、内存效率和性能。
注意力层
DeepSeek 采用多头潜在注意力(MLA)来有效建模输入序列中的复杂依赖关系。为了优化这一机制,作者实现了 DP Attention(数据并行注意力),这是一种消除跨设备 KV 缓存重复的数据并行策略,显著降低了内存开销。该方法在 SGLang v0.4 中首次引入,现已扩展为支持混合数据和张量并行,从而为高效处理小批量数据提供了灵活性。
稠密前馈网络(Dense FFNs)
尽管 DeepSeek-V3 仅使用了三个稠密 FFN 层,但其计算会显著增加峰值内存使用量,如果不加以精心管理,可能会导致系统崩溃。为了解决这个问题,作者采用了数据并行(DP)而非张量并行(TP),其优势如下:
- 增强的可扩展性:由于中间维度为 18,432,高 TP 度(例如 TP32)会导致低效的碎片化,将其分割成小单元片段(例如 576 个单元),而这些片段无法被 128 整除——128 是现代 GPU(如 H100)常见的对齐边界。这种不对齐会阻碍计算效率和内存利用率。DP 通过避免碎片化提供了一个更具可扩展性的解决方案,确保了跨设备的负载均衡。
- 优化的内存效率:传统上,随着 worker 数量的增加,TP 会减少内存使用,但这种优势在 DP attention 下会减弱。在纯 TP 设置中,单层 Transformer 模型的内存需求随 DP 大小扩展的公式为:$\text{Memory}=\frac{N_{\text{param}}}{\text{TP}}+(1+k)N_{\text{hidden\_state}}\cdot \text{DP}$。其中,$N_{\text{hidden\_state}}=n_\text{token}\times n_\text{hidden\_size}$ 是每个设备(DP rank)上的隐藏状态大小,$N_{\text{param}}=n_\text{intermediate\_size}\times n_\text{hidden\_size}$ 是模型参数的数量,$k$ 是表示来自 CUDA Graph 复制的额外内存开销系数。假设 $\text{DP}=\text{TP}$,当 $\text{TP}=\sqrt{\frac{N_{\text{param}}}{(1+k)N_{\text{hidden\_state}}}}$ 时,该内存使用函数最小化。DeepSeek-V3 使用的中间维度大小为 18,432。在 Prefill 阶段,CUDA Graph 通常被禁用,因此 $k=0$。然而,每个设备的 token 数量很容易超过 2,048,导致最佳 TP 大小为 3 或更小。在 Decode 阶段,一个实用的配置可能在每个设备上使用 128 个 token 并设置 $k=3$。在这种情况下,内存最优的 TP 大小为 6。在两个阶段中,较低的 TP 度都能最小化每个设备的内存使用。因此,与仅依赖 TP 相比,DP 可能是更具内存效率的扩展方法。
- 最小化通信开销:在纯 TP 中,每个 FFN 需要两次 all-reduce 操作,导致大量的通信开销。通过利用 DP,作者将此过程优化为在先前的注意力层之后进行一次 reduce-scatter,并在下一层之前进行一次 all-gather,从而将通信成本降低了 50%。此外,当注意力计算也完全在纯 DP 下进行时,设备间的通信被完全消除,显著提高了整体效率。
下图左侧展示了 DP 稠密 FFN 与 DP attention 的集成。用户可以通过设置 --moe-dense-tp-size=1 来启用此功能。
稀疏前馈网络(Sparse FFNs)
在 DeepSeek-V3 的混合专家(MoE)架构中,稀疏 FFN 需要大量的专家权重,造成了严重的内存瓶颈。为了解决这个问题,作者实现了专家并行(EP),将专家权重分布在多个设备上。这种方法在保持高性能的同时有效地扩展了内存容量,尽管它也带来了诸如不规则的 all-to-all 通信和工作负载不平衡等挑战。上图右侧展示了使用 DeepEP 框架的 EP 实现,后续章节将提供有关 EP 设计和优化的进一步细节。
语言模型头(LM Head)
LM head 计算大型词表上的输出概率,这是一项资源密集型操作,传统上通过词表并行(vocabulary parallelism)来处理,以聚合来自 TP 组的 token logits。为了增强可扩展性和效率,作者采用了数据并行(DP),这与处理稠密 FFN 的策略一致。这减少了内存开销并简化了跨设备的通信,提供了一个更精简的解决方案。
Prefill 与 Decode 分离 (PD Disaggregation)
LLM 推理包含两个截然不同的阶段:Prefill(预填充)和 Decode(解码)。Prefill 阶段是计算密集型的,处理整个输入序列;而 Decode 阶段是内存密集型的,管理用于 token 生成的键值(KV)缓存。传统上,这两个阶段在统一的引擎中处理,Prefill 和 Decode 批次的联合调度会引入低效性。为了应对这些挑战,SGLang 引入了 Prefill 和 Decode (PD) 分离。
统一调度的缺陷
将 Prefill 和 Decode 批次放在一起处理的传统统一引擎会导致三个重大问题:
- Prefill 中断:新到来的 Prefill 批次频繁中断正在进行的 Decode 批次,导致 token 生成出现严重延迟。
- DP Attention 不平衡:在 DP attention 中,一个 DP worker 可能会处理 Prefill 批次,而另一个 worker 同时处理 Decode 批次,导致 Decode 延迟增加。
- 与 DeepEP 不兼容:正如后文所述,DeepEP 为 Prefill 和 Decode 执行不同的调度模式,这使得统一调度与 DeepEP 不兼容。
PD 分离通过分离这两个阶段解决了上述问题,从而能够为每个阶段进行量身定制的优化。
实现细节
SGLang 中的 PD 分离设计如下图所示,它在 Prefill 服务器和 Decode 服务器之间交错执行:
在接收到输入请求后,工作流程如下进行:
- Prefill 服务器和 Decode 服务器通过握手进行配对,分别建立本地发送方和接收方。
- Decode 服务器预先分配 KV 缓存,并向 Prefill 服务器发出信号,指示其开始模型前向传递并计算 KV 缓存。
- 计算完成后,数据传输到 Decode 服务器,由其处理迭代的 token 生成。
这种分离确保了每个阶段都在最佳条件下运行,最大化了 GPU 资源的利用率。为了进一步提升性能,该实现包含了:
- 非阻塞传输:数据发送和接收操作在后台线程中运行,保持调度器的事件循环不被中断。
- 基于 RDMA 的传输:远程直接内存访问(RDMA)利用队列对(queue pairs)进行连接,并使用分散-聚集元素(scatter-gather elements, SGE)高效传输非连续的内存块。
- 灵活的 API 集成:SGLang 提供适应性强的 API,集成了 Mooncake 和 NIXL 等高性能 RDMA 库,简化了数据传输。
大规模专家并行
基于 DeepEP 的专家并行
DeepEP 是由 DeepSeek 团队实现的一个通信库,旨在简化 MoE 模型中的 EP。它解决了跨多个 GPU 将 token 高效路由到特定专家的挑战。通过提供优化的通信内核,DeepEP 降低了延迟并提高了吞吐量,使其成为大规模推理任务的理想选择。
DeepEP 提供两种专门的调度模式以应对不同的工作负载需求:
- 正常调度(Normal Dispatch):针对处理长输入序列(如 Prefill 阶段)进行了优化,此模式优先考虑最大的计算吞吐量。然而,它会生成与 CUDA Graph 不兼容的符号形状(symbolic shapes),使其在 Decode 阶段效果不佳,因为在 Decode 阶段内核启动开销会成为严重的瓶颈。
- 低延迟调度(Low-Latency Dispatch):专为在 Decode 阶段生成输出 token 而定制,此模式优先考虑最小化延迟以确保实时性能。它支持 CUDA Graph,但需要预先分配固定大小的内存。如果内存需求超过此预分配,则会发生运行时错误。
在 SGLang 中,DeepEP 的集成提供了自动模式(auto mode),可根据工作负载在两种调度模式之间动态选择。然而,如果没有 PD 分离,自动模式会面临一个限制:它无法在同一个通信组内同时支持正常调度(用于 Prefill)和低延迟调度(用于 Decode)。这种限制阻碍了其与 DP attention 的兼容性,而 DP attention 对于内存高效的推理至关重要。各模式的兼容性如下表所示:
| 模式 | 长输入 | 长输出 | DP Attention | CUDA Graph |
|---|---|---|---|---|
| Normal | ✅ | ❌ | ✅ | ❌ |
| Low-Latency | ❌ | ✅ | ✅ | ✅ |
| Auto | ✅ | ✅ | ❌ | ✅ |
PD 分离通过分离 Prefill 和 Decode 阶段解决了这一问题,允许在 Prefill 阶段使用正常调度,在 Decode 阶段使用低延迟调度,且两者都在 DP attention 下运行。这种集成优化了资源利用率,并通过使调度模式与每个阶段的特定需求保持一致来提高整体性能。
DeepGEMM 集成
DeepGEMM 是 DeepSeek 团队开发的另一个高效库,专门用于优化 MoE 模型中的计算。它提供了两个专门的函数来处理与 MoE 相关的矩阵乘法(Grouped GEMMs),每个函数都针对推理过程的不同阶段进行了定制。
- Grouped GEMMs(连续布局):该内核专为动态输入形状设计,非常适合 MoE 推理的 Prefill 阶段。它处理输入时,不同专家的数据被连续拼接在一起,从而能够灵活处理变化的输入大小。
- Grouped GEMMs(掩码布局):该内核假设输入形状固定,并使用掩码张量(mask tensor)仅计算输入的有效部分。它与 CUDA Graph 兼容,优化了内核启动,非常适合需要严格减少开销的 Decode 阶段。
DeepGEMM 与 DeepEP 的调度模式平滑集成:
- 对于在 Prefill 阶段与正常调度一起使用的连续布局内核,需要一个额外的步骤。由于正常调度输出的是符号形状,因此需要进行排列(permutation)以将输出转换为内核期望的连续格式。作者参考了 LightLLM 项目,并实现了一个自定义的 Triton 内核以实现高效排列。该内核确保正常调度的输出被正确重新排列,从而实现与连续 GEMM 内核的平滑集成。
- 掩码布局内核与 DeepEP 的低延迟调度无缝配合,因为两者都针对 Decode 阶段进行了优化并支持 CUDA Graph。
SGLang 还集成了张量并行下的 DeepGEMM 用于 MoE 计算。此外,DeepGEMM 提供了一个高效的通用 GeMM 内核,可以通过将环境变量 SGL_ENABLE_JIT_DEEPGEMM 设置为 1 在 SGLang 中激活,为非 MoE 操作提供更高的计算效率。
双批次重叠 (Two-batch Overlap, TBO)
在多节点环境中,有限的通信带宽会显著增加整体延迟。为了应对这一挑战,作者遵循 DeepSeek 的系统设计实现了双批次重叠 (TBO)。TBO 将单个批次拆分为两个微批次(micro-batches),允许计算和通信重叠,这也通过将有效批量大小减半来降低峰值内存使用。然而,将 TBO 投入实践带来了特定的实现困难。
实现挑战:
尽管 DeepSeek 发布了 TBO 的设计框架,但存在两个轻微的实现挑战。
- 代码复杂性:直接编写 TBO 代码会导致管理多个微批次的逻辑重复。这增加了代码库的复杂性,使其难以维护且容易出错,尤其是当微批次数量或重叠场景增加时。
- Prefill 阶段的同步问题:在 DeepEP 中的正常调度阻塞 CPU 时,要实现计算和通信的有效重叠需要仔细考量。这种阻塞行为会使流水线停滞,导致 GPU 闲置,从而破坏 TBO 的性能优势。
用于简洁实现的抽象化:
为了创建一个更易于维护和可重用的代码库,作者使用了一个由操作(operations)和让出点(yield points)组成的抽象层。这种方法简化了开发,允许编写代码时就像处理单个微批次一样,同时通过插入 yield points 战略性地暂停执行,以便让其他微批次继续。它消除了代码重复,减少了对变量后缀的潜在需求,并有效地管理了某些执行在层末尾完成而其他执行尚未完成的情况。此外,它还支持以最少的代码更改轻松适应不同的重叠区域选择或未来的增强功能(如三批次重叠)。以下是该方法的简明演示:
operations = [
self._forward_attn,
YieldOperation(), # 暂停执行以让出给其他微批次
self._forward_dispatch,
self._forward_mlp,
YieldOperation(), # 另一个暂停点
self._forward_combine,
]
# 处理单个微批次而无需重复代码
def _forward_attn(self, state):
state.hidden_states = self.self_attn(state.hidden_states, ...)
Prefill 重叠实现:
作者优化了 Prefill 阶段的启动顺序,以避免通过 DeepEP 中的 dispatch 操作阻塞 CPU,即使在使用其异步模式时也是如此。具体而言:
- dispatch 操作会阻塞 CPU,直到 GPU 接收到来自其他 rank 的元数据以分配正确大小的张量。
- 不当的实现会在此期间使计算流闲置,因为没有计算任务提交给 GPU。
为了优化,作者优先将计算任务提交给 GPU,然后再启动阻塞 CPU 的通信。这确保了 GPU 在通信期间保持活跃。如下图所示,具有正确启动顺序的 TBO(以粗边框标示)避免了由阻塞 CPU 的操作(即正常调度)引起的流水线气泡。
专家并行负载均衡器 (EPLB)
在 MoE 模型中,EP 通常会导致跨 GPU 的工作负载分布不均。这种不平衡迫使系统等待最慢的 GPU 计算或通信,浪费了计算周期,并由于专家激活而增加了内存使用。随着 GPU 数量(EP 大小)的增加,不平衡问题变得更加严重。
为了解决这个问题,DeepSeek 开发了 专家并行负载均衡器 (EPLB)。EPLB 将专家分布统计数据作为输入,并计算出专家的最佳排列以最小化不平衡。用户可以分配冗余专家(例如,32 个额外专家),当与原始的 256 个专家结合时,会创建一个包含 288 个专家的池。这个池允许 EPLB 战略性地放置或复制专家——例如,将最常用的专家复制多次,或者将一个中等使用的专家与极少使用的专家组合在一个 GPU 上。
除了平衡工作负载之外,EPLB 还为并行设计提供了更大的灵活性。使用原始的 256 个专家,并行大小被限制为 2 的幂。EPLB 使用 288 个专家,支持更多样化的配置,例如并行大小为 12 或 72。
在下图中,作者通过模拟展示了规模和 EPLB 算法对不平衡问题的影响。GPU 均衡度计算为 GPU 之间 MoE 层的平均计算时间与最大计算时间之比,并使用 GPU 的 token 数量来估计其计算时间。可以看出,当系统随节点数量扩展时,利用率会下降,而启用 EPLB 可显著提高利用率。
真实服务中的 EPLB:
为了使 EPLB 有效,输入分布必须与实际服务工作负载紧密匹配。两种策略增强了这种一致性:
- 增加批量大小:较大的批次减少了专家使用中的随机波动,从而改善了平衡。这可以通过扩展集群或使用多 Token 预测 (MTP) 等技术来实现。
- 定期重平衡:定期更新专家排列利用了时间局部性,但需要高效地重新加载专家。这需要将专家重新加载操作的成本降至最低。
即使使用了 EPLB,一些不平衡也是不可避免的,这使得进一步的优化成为一个有价值的未来方向。
重平衡的实现:
SGLang 分三个阶段实现专家重平衡,以确保效率并将中断降至最低:
- 系统加载阶段:权重可以选择从磁盘预加载到主内存以实现更快的重平衡,或者通过内存映射 (mmap) 保留在磁盘上以减少内存使用。
- 重平衡准备阶段:所需的权重在后台异步传输到设备内存,利用空闲的 DMA 硬件引擎而不中断正在进行的 GPU 操作。
- 重平衡执行阶段:设备到设备的拷贝更新权重。此步骤可以通过物理内存重新绑定技术进一步优化。
这种分阶段的方法确保了重平衡既高效又无中断,在更新期间保持系统性能。
实验环境
- 模型与配置:评估使用 DeepSeek-V3 模型。
- 硬件配置:集群包含 12 个节点,通过 InfiniBand 连接。每个节点配备 8 张 NVIDIA H100 GPU(共 96 张 H100 GPU)。
- 软件与对比基线:
- SGLang 结合 TP16 x 6:每两个节点配对成一个独立组,使用 TP=16 和 DP attention 运行推理。
- SGLang 结合 PD 分离:包含 PD 分离和完整的 EP 优化。EPLB 的分布匹配输入/输出数据。
- SGLang 结合 PD 分离及模拟 MTP:为模拟 MTP,将 batch size 翻倍并减半 KV cache 长度。在真实 attention 计算后插入 dummy kernels 以匹配 DeepSeek 的 profile 耗时。假设 MTP 接受率为 70%。
- DeepSeek Profile 结果:吞吐量估算源自 DeepSeek 官方的 profiling 数据。
实验结果
端到端性能
为了适应不同的工作负载需求,实验独立评估了 Prefill (P) 和 Decode (D) 阶段,假设非测试阶段资源无限,以隔离并最大化测试节点上的负载(与 DeepSeek 的设置一致)。
- Prefill 阶段:在 4 个节点(4×8×H100,EP32)上,系统对于 1K、2K 和 4K 的提示长度,分别实现了每节点 57,674、54,543 和 50,302 tokens/sec 的吞吐量。这代表了比 TP16 基线高达 3.3 倍的提升(主要归功于 DeepGEMM 和 TBO)。假设负载完美均衡,系统吞吐量与 DeepSeek 官方 profile 的差距在 5.6% 以内。
- Decode 阶段:在 9 个节点(9×8×H100,EP72;规模是 DeepSeek 的一半)上,系统对于 2K 输入实现了 22,282 tokens/sec 的单节点吞吐量(比 TP16 快 5.2 倍)。在模拟 MTP 条件下(注意内核被故意减速),系统对于 4K 输入维持了 17,373 tokens/sec 的吞吐量,仅比 DeepSeek 官方 profile 低 6.6%。性能提升主要归功于 EP 带来的 4 倍 batch size 提升,降低了权重的显存消耗。
Profile 结果分析
- 整体吞吐量 (Prefill):测试了每设备 16,384 tokens,输入长度 4,096 的场景。SGLang 默认专家分布 (50,302 tokens/sec) 比 DeepSeek profile (62,713 tokens/sec) 慢 20%。而模拟的完美 EPLB (59,337 tokens/sec) 将差距缩小到 6%。
- 整体吞吐量 (Decode):在 batch size 256、输入长度 2,000 时,SGLang 达到 22,282 tokens/sec,展现了强大的扩展性。在模拟 MTP 下(batch 128,长度 4000),吞吐量为 17,373 tokens/sec。
- 详细耗时分解 (Prefill):默认 EPLB 下,通信内核执行时间更长且方差更大(由于专家不平衡),导致计算流出现气泡。模拟完美 EPLB 则更接近 DeepSeek 的 profile。
- 详细耗时分解 (Decode):SGLang 的 combine 操作比 DeepSeek 慢 2 倍,原因是 attention 计算较短导致通信内核 busy-wait(在模拟慢速 attention 中得到证实,耗时对齐)。MoE 内核慢 25%,可能是因为 DeepSeek 使用 18 个节点分布专家更高效。Dispatch 时间仍有优化潜力(目前为 ~0.17ms,但单元测试显示理想可达 0.06ms)。
消融实验:双批次重叠 (TBO)
- Prefill 阶段:TBO 带来了两大优势:1) 支持更大的 Batch Size:原版配置在单卡 16,384 tokens 时会 OOM(上限 8,192)。TBO 优化了内存,支持 16,384 tokens 运行,使性能比其他最优配置提升 40.5%。2) 吞吐量提升:在相同 token 数量下,通过重叠计算与通信,TBO 相比原版提升了 27% 到 35% 的吞吐量。
- Decode 阶段:加速效果取决于 batch size。小于 64-128 时,TBO 收益微小甚至为负(如 32 tokens/device 时为 -27%)。在 256 tokens 时加速达 25.5%。在模拟 MTP 场景(长 attention 时间完美掩盖通信)中,128 序列下的加速达到最高 35%。
消融实验:EPLB
- 在大规模设置下,EPLB 带来了 1.49 倍 (Prefill) 和 2.54 倍 (Decode) 的显著加速,有效缓解了节点增加带来的不平衡。
- 案例研究显示,工作负载均衡度(平均 token 数量除以最大 token 数量)与吞吐量之间存在强相关性。
- 分布统计显示,少数专家被频繁使用,而多数使用较少。Prefill 和 Decode 的分布存在显著差异,这进一步证明了 PD 分离(允许阶段特定的专家放置)的必要性。
补充细节 (Toolkits)
一次性张量 (Disposable Tensor)
在 PyTorch 中,由于持久的对象引用,内存管理可能具有挑战性,特别是在 CUDA 内存稀缺的 GPU 密集型工作流中。例如,当一个张量在函数外部被引用时,函数内的 del 操作无法释放其内存,这会增加峰值内存使用,导致性能下降或 OOM。
SGLang 引入了 DisposableTensor 类(torch.Tensor 的子类)来解决此问题。它引入了一个 dispose() 方法来显式且立即释放张量的内存,绕过 Python 的引用计数限制。其工作原理如下:
def ffn(hidden_state: torch.Tensor, linear1: nn.Linear, linear2: nn.Linear):
intermediate_state = linear1(hidden_state)
hidden_state.dispose() # 立即释放 CUDA 内存
return linear2(nn.ReLU(intermediate_state))
# 将张量包装在 DisposableTensor 中
hidden_state = DisposableTensor(hidden_state)
hidden_state = ffn(hidden_state, linear1, linear2)
通过调用 dispose(),计算完成后 CUDA 内存会立即被释放,从而降低峰值内存使用率。
专家工作负载提取与模拟
SGLang 包含用于分析和模拟 MoE 模型中专家工作负载分布的工具集:
- 导出专家工作负载统计:提取累积统计数据(支持实时 EPLB 优化)或逐批次工作负载数据(用于深度分析)。
- 模拟专家利用率:无需昂贵的硬件即可模拟跨配置的专家平衡。例如,用户可以从 2x8xH100 收集数据,并推断模拟 22 个节点大规模部署的性能。这为调整重平衡频率或 batch size 提供了低成本的方法。
结论
通过利用 PD 分离、EP 以及精心设计的并行架构,团队在 SGLang 中以卓越的性能成功复现了 DeepSeek 的推理框架。该开源实现实现了每节点每秒 52.3k 输入 token 和 22.3k 输出 token 的性能,展示了 SGLang 在大规模 LLM 推理中的强大能力。
局限性与未来工作:
- 延迟优化:目前侧重于吞吐量,导致首字延迟(TTFT)在 2-5 秒,Token 间延迟(ITL)约为 100ms,需针对实时场景进一步优化。
- 序列长度限制:受限于 96 张 GPU 的规模,目前仅限于较短序列,扩展 GPU 资源将支持更长序列。
- 多 Token 预测 (MTP) 集成:SGLang 支持 MTP,但缺乏与 DP attention 的完全集成,降低了混合并行配置的效率。
- EPLB 分布漂移:实验使用了分布内数据,未来需要测试在存在分布漂移时的实际性能。
- 灵活的张量并行 (TP) 大小:DeepSeek-V3 稠密 FFN 的内存最优 TP 大小较小但大于 1。目前 SGLang 仅支持纯 TP 或 DP,需要灵活的 TP 选项来优化内存。
- Blackwell 架构支持:目前仅支持 NVIDIA Hopper 架构,正积极扩展对下一代 Blackwell 架构的兼容性。
附录
本文的附录部分列出了大量与 SGLang 核心库以及 DeepEP 通信库相关的 GitHub Pull Requests(PRs)。这些 PR 涵盖了从并行架构设计、PD 分离实现、DeepGEMM/DeepEP 集成到 TBO 和 EPLB 优化的所有代码提交记录(如 PR #1970, #2925, #4068 等,以及 DeepEP#142),为开发者和社区提供了完整的重现路径和技术实现细节。
💬 评论讨论
欢迎在这里分享您的想法和见解!