https://mp.weixin.qq.com/s/DJpuqJnTCelMvNerDD2_Og

DeepSeek 是一个广受欢迎的开源大型语言模型 (LLM),因其强大的性能而备受赞誉。然而,由于其庞大的规模和独特的架构(采用多头潜在注意力 (MLA) 和混合专家 (MoE)),需要更先进的系统才能高效地大规模服务。在本篇博文中,我们将解释如何将 DeepSeek 的推理系统性能与 SGLang 进行匹配。

如上图所示,我们的实现在 Atlas Cloud 的 12 个节点上运行,每个节点配备 8 块 H100 GPU。它采用预填充解码分解和大规模专家并行 (EP) 技术,对于 2000 个 token 的输入序列,每个节点的速度达到每秒 52.3k 个输入 token 和每秒 22.3k 个输出 token。据我们所知,这是第一个在大规模环境下接近 DeepSeek 官方博客中报告的吞吐量的开源实现

通过在本地部署此实现,其成本为每 100 万个输出 token 0.20 美元,约为官方 DeepSeek Chat API 成本的五分之一。与使用相同资源的原生张量并行相比,这种优化策略可将输出吞吐量提高高达 5 倍。本博客深入探讨了我们的并行设计、优化方法和成果。我们工作的所有组件均完全开源,允许其他人探索和借鉴我们的成果。重现我们的实验的完整说明可在此处找到。

亮点

  • SGLang 现在支持预填充解码(PD)分解和大规模 EP,包括DeepEPDeepGEMMEPLB的全部功能。
  • 利用这些新功能,我们的团队成功地使用 12 个节点(每个节点配备 8 块 H100 GPU)复制了 DeepSeek 的推理系统。总体而言,对于 2000 个 token 的输入序列,SGLang 实现了每个节点每秒 52.3k 个输入 token 和每秒 22.3k 个输出 token 的吞吐量。
  • 本博客解释了我们方法的技术细节,重点关注效率优化、峰值内存使用量降低以及工作负载平衡。性能分析结果表明,我们的实现性能几乎与 DeepSeek 官方报告相当。
  • 所有实验和代码均完全开源,供社区访问和进一步开发。

并行设计

高效的并行性对于管理 DeepSeek 架构的计算复杂度和内存需求至关重要。本节概述了我们优化关键组件的方法:注意力层、密集前馈网络 (FFN)、稀疏 FFN 和语言模型 (LM) 头。每个组件都利用定制的并行策略来增强可扩展性、内存效率和性能。

注意力层

DeepSeek 采用多头潜在注意力 (MLA) 技术,有效地对输入序列中的复杂依赖关系进行建模。为了优化这一机制,我们实现了DP Attention,这是一种数据并行策略,可以消除跨设备的键值缓存重复,从而显著降低内存开销。这种方法在SGLang v0.4中引入,现已扩展至支持混合数据和张量并行,从而为高效处理小批量数据提供了灵活性。

稠密FFN

尽管仅使用三个密集 FFN 层,DeepSeek-V3 的计算仍会显著增加峰值内存使用量,如果不加以管理,可能会导致系统崩溃。为了解决这个问题,我们采用数据并行 (DP)而非张量并行 (TP),从而发挥以下优势:

增强的可扩展性:由于中等维度为 18,432,高 TP 度(例如 TP32)会导致低效地碎片化为小单元段(例如 576 个单元),这些单元段无法被 128 整除——这是 H100 等现代 GPU 常见的对齐边界。这种不对齐会影响计算效率和内存利用率。DP 通过避免碎片化,确保跨设备的工作负载均衡分配,从而提供更具可扩展性的解决方案。

优化的内存效率:传统上,随着工作器规模的增加,TP 会减少内存使用量,但在 DP 注意力机制下,这种优势会减弱。在纯 TP 设置中,单层 Transformer 模型的内存需求会随着 DP 规模的变化而变化,如下所示:

KaTeX parse error: Expected 'EOF', got '_' at position 71: …N_{\text{hidden_̲state}}\cdot \t…

这里,KaTeX parse error: Expected 'EOF', got '_' at position 16: N_{\text{hidden_̲state}}=n_\text… 是每个设备上隐藏状态的大小(DP等级),KaTeX parse error: Expected 'EOF', got '_' at position 38: …xt{intermediate_̲size}\times n_\… 是模型参数的数量, k k k 是表示 CUDA Graph 复制额外内存开销的系数。假设 DP = TP \text{DP}=\text{TP} DP=TP ,当 KaTeX parse error: Expected 'EOF', got '_' at position 61: …N_{\text{hidden_̲state}}}} DeepSeek-V3 使用 18,432 的中间大小。在预填充阶段,CUDA Graph 通常处于禁用状态,因此 k = 0 k = 0 k=0 。然而,每个设备的令牌大小很容易超过 2,048,因此最佳 TP 大小应为 3 或更少。在解码阶段,实际配置可能每个设备使用 128 个令牌,并设置 k = 3 k = 3 k=3 。在本例中,内存最优的 TP 大小为 6。在两个阶段中,较低的 TP 度数可最大限度地减少每个设备的内存使用量。因此,与单纯依赖 TP 相比,DP 可能提供一种更节省内存的扩展方法。

最小化通信开销:在纯 TP 中,每个 FFN 需要两次全归约 (all-reduce) 操作,这会导致巨大的通信开销。通过利用 DP,我们将此过程优化为在前一个注意力层之后进行一次归约散射 (reduce-scatter),并在下一个注意力层之前进行一次全归约 (all-gather),从而将通信成本降低 50%。此外,当在纯 DP 下计算注意力时,设备间通信完全被消除,从而显著提升了整体效率。

DP 密集 FFN 与 DP 注意力机制的集成如下图左图所示。用户可以通过设置 来启用此功能–moe-dense-tp-size=1。

稀疏FFN

在 DeepSeek-V3 的混合专家 (MoE) 架构中,稀疏 FFN 需要大量的专家权重,从而造成严重的内存瓶颈。为了解决这个问题,我们实现了专家并行 (EP),将专家权重分布在多个设备上。这种方法在保持高性能的同时有效地扩展了内存容量,尽管它确实带来了诸如不规则的全对全通信和工作负载不平衡等挑战。

上图右图展示了我们使用 DeepEP 框架的 EP 实现,有关我们的 EP 设计和优化的更多细节将在后续章节中提供。

LM 头

语言模型头 ( LM head) 计算大型词汇表的输出概率,这是一个资源密集型操作,传统上通常使用词汇并行来处理,以聚合来自 TP 组的 token logit。为了提高可扩展性和效率,我们采用了数据并行 (DP),这与我们的密集 FFN 策略相呼应。这减少了内存开销,简化了跨设备通信,从而提供了更精简的解决方案。

预填充和解码分解

LLM 推理包含两个不同的阶段:预填充和解码。预填充阶段计算密集型,需要处理整个输入序列;而解码阶段内存密集型,需要管理用于生成 token 的键值 (KV) 缓存。传统上,这两个阶段由一个统一的引擎处理,而预填充和解码批次的混合调度会导致效率低下。为了应对这些挑战,我们在 SGLang 中引入了预填充和解码 (PD) 分解功能。

统一调度的问题

传统的统一引擎将预填充和解码批次一起处理,这会导致三个重大问题:

  • 预填充中断:传入的预填充批次经常中断正在进行的解码批次,从而导致令牌生成的严重延迟。
  • DP 注意力不平衡:在 DP 注意力中,一个 DP 工作者可能处理预填充批次,而另一个 DP 工作者同时处理解码批次,从而导致解码延迟增加。
  • 与 DeepEP 不兼容:正如我们将在后面的部分讨论的那样,DeepEP 对预填充和解码执行不同的调度模式,这使得统一调度与 DeepEP 不兼容。

PD 分解通过分离两个阶段来解决这些问题,从而可以针对每个阶段进行定制优化。

实现细节

下图所示的 SGLang 中的 PD 分解设计在预填充服务器和解码服务器之间交错执行:

收到输入请求后,工作流程如下:

  • 预填充服务器和解码服务器通过握手配对,分别建立本地发送方和接收方。
  • 解码服务器预先分配 KV 缓存,向预填充服务器发出信号,开始模型前向传递并计算 KV 缓存。
  • 一旦计算完成,数据就会传输到解码服务器,由其处理迭代令牌生成。

这种分离确保每个阶段都在最佳条件下运行,从而最大限度地提高 GPU 资源利用率。为了进一步提升性能,我们的实现方案包括:

  • 非阻塞传输:数据发送和接收操作在后台线程中运行,保持调度程序的事件循环不间断。
  • 基于 RDMA 的传输:远程直接内存访问 (RDMA) 利用队列对进行连接,并利用分散-聚集元素 (SGE) 高效传输非连续内存块。
  • 灵活的 API 集成:SGLang 提供适应性 API,集成 Mooncake 和 NIXL 等高性能 RDMA 库,简化数据传输。

更多详细信息请参阅我们的设计文档

大规模专家并行

DeepEP 专家并行性

DeepEP是由 DeepSeek 团队实现的通信库,旨在简化 MoE 模型中的 EP 流程。它解决了如何高效地将 token 路由到跨多个 GPU 的特定专家的难题。通过提供优化的通信内核,DeepEP 降低了延迟并提高了吞吐量,使其成为大规模推理任务的理想选择。

DeepEP 提供两种专门的调度模式来满足不同的工作负载需求:

  • 正常调度:此模式针对处理长输入序列(例如在预填充阶段)进行了优化,优先考虑最大计算吞吐量。然而,它会生成与 CUDA Graph 不兼容的符号形状,导致其在解码阶段效率较低,而内核启动开销会成为显著的瓶颈。
  • 低延迟调度:此模式专为在解码阶段生成输出令牌而设计,优先考虑最小延迟以确保实时性能。它支持 CUDA Graph,但需要预先分配固定大小的内存。如果内存需求超过此预分配大小,则会发生运行时错误。

在 SGLang 中,DeepEP 集成提供了自动模式,可根据工作负载在两种调度模式之间动态选择。然而,如果没有 PD 分解,自动模式面临一个限制:它无法在同一通信组中同时支持正常调度(用于预填充)和低延迟调度(用于解码)。此限制阻碍了其与 DP 注意力机制的兼容性,而 DP 注意力机制对于高效内存推理至关重要。下表概述了每种模式的兼容性:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

PD 分解通过分离预填充和解码阶段解决了这个问题,允许预填充阶段进行正常调度,解码阶段进行低延迟调度,两者均在 DP 关注下进行。这种集成通过将调度模式与每个阶段的特定需求相结合,优化了资源利用率并提升了整体性能。

DeepGEMM 集成

DeepGEMM是 DeepSeek 团队开发的另一个高效库,专门用于优化 MoE 模型中的计算。它提供了两个专门用于处理 MoE 相关的矩阵乘法(分组 GEMM)的函数,每个函数都针对推理过程的不同阶段进行了定制。

  • 分组 GEMM(连续布局):此内核专为动态输入形状而设计,非常适合 MoE 推理的预填充阶段。它处理输入时,不同专家的数据会连续连接,从而能够灵活处理不同的输入大小。
  • 分组 GEMM(掩码布局):此内核假设输入形状固定,并使用掩码张量仅计算输入的有效部分。它与 CUDA Graph 兼容,CUDA Graph 优化了内核启动,使其非常适合降低开销至关重要的解码阶段。

DeepGEMM 与 DeepEP 的调度模式无缝集成:

  • 对于在预填充阶段与正常调度一起使用的连续布局内核,需要额外的步骤。由于正常调度输出的是符号形状,因此需要进行置换以将输出转换为内核所需的连续格式。我们参考了 LightLLM 项目,并实现了一个自定义的 Triton 内核以实现高效的置换。该内核确保正常调度的输出被正确地重新排列,从而能够与连续 GEMM 内核顺利集成。
  • 掩码布局内核与DeepEP 的低延迟调度无缝配对,因为两者都针对解码阶段进行了优化并支持 CUDA Graph。

SGLang 还集成了 DeepGEMM,用于在张量并行下进行 MoE 计算。此外,DeepGEMM 提供了一个高效的通用 GeMM 内核,可以通过在 SGLang 中将环境变量设置SGL_ENABLE_JIT_DEEPGEMM为 1 来激活该内核,从而为非 MoE 操作提供更高的计算效率。

两批次重叠

在多节点环境中,有限的通信带宽会显著增加整体延迟。为了应对这一挑战,我们遵循DeepSeek 的系统设计,实现了双批次重叠 (TBO)。TBO 将单个批次拆分为两个微批次,允许计算和通信重叠,同时通过将有效批次大小减半来降低峰值内存使用量。然而,将 TBO 付诸实践会带来一些具体的实现难题。

实施挑战

虽然DeepSeek发布了TBO的设计框架,但是在实现上还是存在两个小小的挑战。

  • 代码复杂性:直接编写 TBO 代码可能会导致管理多个微批次的逻辑重复。这会增加代码库的复杂性,使其更难维护且容易出错,尤其是在微批次数量或重叠场景增加的情况下。
  • 预填充阶段的同步问题:当 DeepEP 中的正常调度阻塞 CPU 时,需要考虑实现计算和通信之间的有效重叠。这种阻塞行为会导致流水线停滞,使 GPU 处于空闲状态,并削弱 TBO 的性能优势。
清晰实现的抽象

为了创建更易于维护和复用的代码库,我们使用了一个由操作和屈服点组成的抽象层。这种方法简化了开发过程,让我们能够像处理单个微批次一样编写代码,同时通过插入屈服点策略性地暂停执行,让其他微批次继续执行。它消除了代码重复,减少了对变量后缀的潜在需求,并有效地管理了某些执行在某一层结束时完成而其他执行尚未完成的情况。此外,它还支持轻松适应不同的重叠区域选择或未来的增强功能(例如三批次重叠),并且只需极少的代码更改。以下是此方法的简要演示:

operations = [
    self._forward_attn,
    YieldOperation(),  # Pause execution for other micro-batches
    self._forward_dispatch,
    self._forward_mlp,
    YieldOperation(),  # Another pause point
    self._forward_combine,
]

# Process a single micro-batch without duplicating code
def _forward_attn(self, state):
    state.hidden_states = self.self_attn(state.hidden_states, ...)
预填充重叠实施

我们在预填充阶段优化了启动顺序,以避免 DeepEP 中的调度操作阻塞 CPU,即使我们使用的是异步模式。具体如下:

  • 调度操作会阻塞 CPU,直到 GPU 从其他等级接收元数据以分配正确大小的张量。
  • 不正确的实现会导致计算流在此期间处于空闲状态,因为没有计算任务提交给 GPU。

为了进行优化,我们优先将计算任务提交给 GPU,然后再启动阻塞 CPU 的通信。这确保了 GPU 在通信期间保持活跃。如下图所示,采用适当启动顺序(以粗体边框表示)的 TBO 可以避免由阻塞 CPU 的操作(即正常调度)引起的气泡。

专家并行负载均衡器

在 MoE 模型中,EP 通常会导致 GPU 之间的工作负载分配不均。这种不平衡迫使系统等待最慢的 GPU 计算或通信,从而浪费计算周期,并因专家激活而增加内存占用。随着 GPU 数量(EP 大小)的增加,不平衡问题会变得更加严重。

为了解决这个问题,DeepSeek 开发了专家并行负载均衡器 (EPLB)。EPLB 将专家分布统计数据作为输入,并计算出最佳的专家分配方案,以最大限度地减少不平衡。用户可以分配冗余专家(例如,额外 32 位专家),这些专家与原有的 256 位专家合并后,将组成一个包含 288 位专家的专家池。该专家池允许 EPLB 策略性地放置或复制专家——例如,多次复制最常用的专家,或将中等使用率的专家与不常用的专家组合到单个 GPU 上。

除了平衡工作负载之外,EPLB 还在并行度设计方面提供了更大的灵活性。最初使用 256 位专家,并行度大小被限制为 2 的幂次方。而使用 288 位专家,EPLB 可以实现更多样化的配置,例如并行度大小为 12 或 72。

下图通过模拟演示了规模和 EPLB 算法对不平衡问题的影响。我们将 GPU 平衡性计算为 GPU 间 MoE 层的平均计算时间与最大计算时间的比率,并使用 GPU 的令牌数量来估算其计算时间。可以看出,当系统随着节点数量的增加而扩展时,利用率会下降,而启用 EPLB 可以显著提高利用率。

EPLB 在实际服务中的应用

为了使 EPLB 有效,输入分布必须与实际服务工作负载紧密匹配。以下两种策略可以增强这种一致性:

  • 增加批次大小:更大的批次可以减少专家使用中的随机波动,从而改善平衡,这可以通过扩展集群或使用多令牌预测(MTP)等技术来实现。
  • 定期重新平衡:定期更新专家系统布局可以利用时间局部性,但需要高效地重新加载专家系统。这需要最大限度地降低专家系统重新加载操作的成本。

即使有了 EPLB,一些不平衡也是不可避免的,因此进一步优化是一个有价值的未来方向。

重新平衡的实施

SGLang 分三个阶段实施专家重新平衡,以确保效率和最小干扰:

  • 系统加载阶段:权重可以选择从磁盘预加载到主内存以实现更快的重新平衡,或者使用内存映射(mmap)保存在磁盘上以减少内存使用量。
  • 重新平衡准备阶段:所需权重在后台异步传输到设备内存,利用免费的 DMA 硬件引擎,而不会中断正在进行的 GPU 操作。
  • 重新平衡执行阶段:设备到设备的复制更新权重。此步骤可以通过物理内存重新绑定技术进一步优化。

这种分阶段的方法可确保重新平衡既高效又无中断,从而在更新期间保持系统性能。

评估

端到端性能

实验设置

我们在一个由 12 个节点组成的集群上使用 DeepSeek-V3 评估了不同配置的 SGLang 的端到端性能。这些节点通过 InfiniBand 连接,每个节点配备 8 块 H100 GPU。本次评估重点突出了我们先进的优化技术所带来的吞吐量提升。我们比较了以下四种设置:

  • 具有 TP16 x 6 的 SGLang:每两个节点配对一个独立组,运行 TP 大小为 16 和 DP 注意力的 DeepSeek-V3 推理。
  • 带有 PD 分解的 SGLang:此版本包含 PD 分解和完整的 EP 优化。对于 EPLB,由于无法获取实时服务统计数据,我们采用了与输入/输出数据匹配的分布。
  • 带有 PD 分解和模拟 MTP 的 SGLang:为了模拟 MTP 的效果,我们首先将批次大小加倍,并将键值对 KV 缓存长度减半,以保持 GroupedGeMM 计算和内存访问的工作负载不变。此外,我们在实际注意力计算之后插入虚拟核,以确保注意力阶段的耗时与 DeepSeek 的配置文件相同,从而准确反映 MTP 注意力机制带来的性能下降。我们保守估计 MTP 下的接受率为 60%。
  • DeepSeek 分析结果:吞吐量估计值来自DeepSeek 的官方分析数据
预填充和解码阶段的性能分析

为了适应不同的工作负载需求,我们分别评估了预填充 § 和解码 (D) 阶段,并假设非测试阶段拥有无限资源,以隔离并最大化测试节点的负载——这与 DeepSeek 的设置类似。结果总结如下:

  • 预填充阶段:在 4 个节点(4×8×H100,EP32)上,系统在提示长度分别为 1K、2K 和 4K 的情况下,实现了每节点每秒 57,674、54,543 和 50,302 个令牌的吞吐量。如下方柱状图所示,这比 TP16 基线提升了 3.3 倍,这主要归功于优化的 GroupedGeMM 内核 (DeepGEMM) 和双批次重叠。假设工作负载完全平衡,我们系统的吞吐量与 DeepSeek 官方数据相差 5.6% 以内。
  • 解码阶段:在 9 个节点(9×8×H100,EP72;规模为 DeepSeek 的一半)上进行评估,系统在 2K 输入下实现了每节点 22,282 个令牌/秒的吞吐量,比 TP16 基准提升了 5.2 倍。在模拟 MTP 条件下(注意核的速度被故意降低以反映实际延迟),系统在 4K 输入下保持了每节点 17,373 个令牌/秒的高吞吐量,仅比 DeepSeek 的官方数据低 6.6%。如右图所示,这些性能提升主要归功于 EP 带来的 4 倍更大的批处理大小,这通过显著降低模型权重的每 GPU 内存消耗来增强可扩展性。

结果

本节将 SGLang 的性能与 DeepSeek 的推理系统进行比较,并将我们的实验设置尽可能地与 DeepSeek 的生产环境保持一致。我们分析了整体吞吐量和详细的内核细分,并以 DeepSeek 的博客和公开资料数据为基准。

  • 总吞吐量

对于预填充,我们测试了每台设备 16,384 个令牌、输入长度为 4,096 的场景。由于 DeepSeek 专家分布的不确定性,我们评估了两种情况:一种采用默认专家分布,另一种以模拟完美 EPLB(遵循组限制路由语义的随机专家选择)作为性能上限。

结果如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

DeepSeek 的配置文件显示其吞吐量大约是其生产环境的两倍。默认专家不平衡的 SGLang 比 DeepSeek 的配置文件慢 20%,而模拟的完美 EPLB 情况将差距缩小至 6%。

对于解码,结果显示如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

使用 DeepSeek 一半的节点,模拟 MTP 的 SGLang 仅比 DeepSeek 的配置文件慢一点。在更高的批次大小设置(256 个序列,2,000 输入长度)下,SGLang 实现了每节点每秒 22,282 个令牌,展现出强大的可扩展性。

详细分析

下图分解了预填充的内核执行时间,包括单元测试结果作为理论上限:

  • 默认 EPLB:与 DeepSeek 的配置文件相比,通信内核的执行时间更长,方差更大,这可能是由于专家不平衡性加剧造成的。这会导致计算流气泡延长,从而降低整体性能。
  • 模拟完美 EPLB:此设置与 DeepSeek 的配置文件更加接近,但仍然存在差异,表明存在潜在的优化空间。
  • 与单元测试的比较:DeepSeek 和 SGLang 的通信时间都比单元测试结果慢,而后者在禁用 TBO 时可以实现,如果通信是瓶颈,则揭示了潜在的优化方向。

SGLang 的解码内核细分与 DeepSeek 非常接近,如下所示:

主要观察结果包括:

  • 组合时间差异:由于注意力机制计算时间较短,SGLang 的组合操作似乎比 DeepSeek 慢 2 倍,导致通信内核处于忙等待状态。在模拟慢速注意力机制的实验中,组合时间与 DeepSeek 匹配,证实了这一假设。
  • MoE 性能:SGLang 的 MoE 内核速度慢了 25%,可能是因为 DeepSeek 的 18 个节点(而我们的只有 9 个)可以更有效地分配专家,从而减少 GEMM 操作的内存访问开销。
  • 调度优化潜力:DeepSeek 和 SGLang 均显示每层调度时间约为 0.17 毫秒,但 DeepEP 的单元测试显示,SM 占用时间可能高达 0.06 毫秒。目前,调度过程需要花费大量时间等待数据。在发送/接收操作之间插入慢速虚拟核可将调度时间缩短至 0.09 毫秒,并且使用单元测试数据进行的运行中持续时间分析表明,还有进一步改进的空间。

虽然仍存在一些细微的改进(主要是在“其他内核”下的内核融合),但 SGLang 的解码性能与 DeepSeek 基本一致,预填充优化是下一个重点。

消融研究:两批次重叠

批次大小和注意时间的影响

本节研究不同批次大小和模拟 MTP 场景下的 TBO 性能。

TBO 在预填充阶段提供了两个显著的优势,吞吐量比较和内存使用优化可以证明这一点:

  • 支持更大的批量大小:在原始配置中,每台设备最多可处理 8,192 个令牌,处理 16,384 个令牌时才会出现内存不足 (OOM) 错误。TBO 通过优化输入令牌的内存使用率来缓解这一问题,使每台设备能够以高达 16,384 个令牌的批量进行推理。与所有其他已优化的配置相比,启用 TBO 标记后,性能进一步提升至 40.5%。
  • 增强吞吐量:通过将计算(例如,注意力和 MLP 阶段)与通信(例如,DeepEP 组合和调度)重叠,即使每个设备处理相同的令牌数,TBO 也能比原始设置实现 27% 至 35% 的吞吐量提升。

TBO 在解码阶段的影响因场景而异,其性能与批次大小和注意力处理时间相关:

  • 实际测试用例:实际场景中的加速取决于批大小是否超过 64 到 128 个令牌之间的阈值。低于此阈值时,TBO 的提升很小甚至为负(例如,每设备 32 个令牌时为 -27%),因为较小的解码批大小会降低内核效率。在 256 个令牌时,加速比达到 25.5%,性能达到每秒 22,310 个令牌。
  • 模拟 MTP 场景:在模拟 MTP 场景中,当处理 128 个请求,每个解码步骤生成 256 个 token 时,TBO 提供了最显著的加速。这是由于注意力机制处理时间延长,使得计算(例如 DP Attention 层)与 DeepEP 通信开销(例如组合和调度步骤)保持一致。评估显示,在 128 个序列/设备的情况下,加速比为 35%,吞吐量为每秒 17,552 个 token,而未使用 TBO 时则为每秒 12,929 个 token。
详细分析

我们评估了三种预填充场景:每批次 16k 个令牌的 TBO、每批次 8k 个令牌的 TBO 以及每批次 8k 个令牌的非 TBO。下图揭示了关键洞察:

  • TBO 效率:与 8k 案例相比,TBO 通过重叠计算和通信提高了整体效率,正如预期的那样。
  • 批次大小影响:使用 TBO 将批次大小从 16k 减少到 8k 会导致速度略有下降,这反映出批次较小时内核效率会下降。
  • 内核性能:有趣的是,尽管无 TBO 8k 的情况和 TBO 16k 的情况的内核有效批大小均为 8k,但每个内核的速度却优于 TBO 16k 的情况。这可能是由于启用 TBO 后流多处理器 (SM) 数量减少、重叠期间潜在的噪声邻居效应,或计算和通信之间的内核不兼容所致。这些发现为 SGLang 的未来优化方向提供了参考。

对于解码阶段,我们分析了三种配置:批量大小为 256 的 TBO、批量大小为 256 的无 TBO 以及批量大小为 128 的无 TBO。时间细分如下所示:

  • TBO 与无 TBO(批次大小 256)对比:没有 TBO 时,由于缺乏重叠,通信时间会显著增加。然而,计算内核(尤其是 GEMM)受益于更大的有效批次大小,从而实现更快的执行速度。
  • TBO (256) vs. 无 TBO (128):与具有相同内核批次大小的案例相比,在无 TBO 设置下,只有非重叠通信会变慢,而计算速度保持不变。与预填充不同,解码通信内核要么充分利用 SM(发送/接收期间),要么完全不使用(在飞行等待期间),从而避免与计算内核发生资源争用。

消融研究:EPLB

本节通过总体吞吐量分析和详细的案例研究,评估 EPLB 对系统性能的影响。鉴于 EPLB 对生产环境中工作负载分布及其变化的敏感性,我们专注于定性和可推广的洞察,而非实际性能(后者需要生产数据)。

总体结果

下图展示了 EPLB 在大规模设置下对吞吐量的影响。由于 EPLB 能够缓解 GPU 之间的工作负载不平衡,其显著加速性能达到了预期的 1.49 倍(预填充)和 2.54 倍(解码)。随着 Rank 数量的增加,不平衡现象也随之加剧,而 EPLB 在我们的大规模实验中有效地解决了这一问题,从而显著提升了吞吐量。

案例研究:工作负载不平衡与总体吞吐量

为了探索工作负载不平衡和吞吐量之间的关系,我们进行了一个案例研究,使用解码实验进行了 1800 个输入令牌、100 个输出令牌和 256 的批次大小。吞吐量和平衡性(专家的平均令牌数除以最大令牌数)与解码步骤的关系被绘制成图:

结果显示平衡性和吞吐量之间存在很强的相关性,强调了保持高平衡性以实现最佳性能的重要性。

案例研究:专家分布统计

下图展示了预填充和解码样本数据的专家分布统计数据:

主要观察结果包括:

  • 专家使用不平衡:大多数专家很少被使用,而一小部分专家却被大量使用,这凸显了 MoE 模型固有的不平衡性。
  • 预填充与解码的差异:虽然预填充和解码分布有相似之处,但也存在显著差异。这支持使用 PD 分解,从而为每个阶段分配不同的专家,从而优化性能。

这些发现凸显了 EPLB 在解决工作量不平衡方面的作用以及根据特定阶段的需求定制专家安排的价值。

工具包

一次性张量

由于持久对象引用的存在,PyTorch 中的内存管理可能颇具挑战性,尤其是在 CUDA 内存资源稀缺的 GPU 密集型工作流程中。请考虑以下示例:

def ffn(hidden_state: torch.Tensor, linear1: nn.Linear, linear2: nn.Linear):
    intermediate_state = linear1(hidden_state)
    del hidden_state  # Attempt to free memory, but no effect due to external reference
    return linear2(nn.ReLU(intermediate_state))

hidden_state = ffn(hidden_state, linear1, linear2)

这段代码del hidden_state旨在释放计算完成hidden_state后占用的内存。然而,由于函数外部仍然引用了,因此该操作没有任何效果。这会增加峰值内存使用量,从而可能导致性能下降或内存不足错误。intermediate_statehidden_statedel

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()  # Immediately releases CUDA memory
    return linear2(nn.ReLU(intermediate_state))

# Wrap the tensor in DisposableTensor
hidden_state = DisposableTensor(hidden_state)
hidden_state = ffn(hidden_state, linear1, linear2)

通过包装hidden_state并在不再需要时DisposableTensor调用,CUDA 内存会被立即释放。这确保了张量在计算中的作用完成后立即释放内存,从而减少峰值内存使用量并提高整体效率。

专家工作负载提取与模拟

SGLang 还包含一套工具集,用于分析和模拟 MoE 模型中专家工作负载的分配。此功能使用户能够:

  • 转储专家工作负载统计数据:提取累积统计数据或每批次工作负载数据。累积统计数据支持 EPLB 管理器进行实时优化,而每批次数据则为分析和模拟提供细致的洞察。
  • 模拟专家利用率:无需昂贵的硬件或反复试验,即可在各种配置中模拟专家平衡。例如,用户可以从中等配置(例如 2x8xH100 或 8xH200)收集工作负载数据,并模拟大规模 22 节点部署的性能。

这项模拟功能允许用户评估诸如重新平衡频率、节点数量或批次大小等因素对系统性能的影响。这是一种在扩展之前进行配置微调的经济高效的方法。

局限性和未来工作

虽然我们为 DeepSeek-V3 推理实现的 SGLang 表现出了显著的吞吐量提升,但仍然存在一些限制和有待未来改进的领域:

延迟优化:当前对吞吐量的关注使得第一个令牌的时间(TTFT)为 2-5 秒,令牌间延迟(ITL)约为 100 毫秒,需要针对实时用例进行进一步优化。

序列长度限制:由于使用了 96 个 GPU,因此只能生成较短的序列。扩展 GPU 资源可以支持更长的序列,这对于特定应用至关重要。

多标记预测 (MTP) 集成:SGLang 支持 MTP,但缺乏与 DP 注意力的完全集成,从而降低了混合并行配置的效率。

EPLB 分布:本博客中的实验利用了专家并行负载均衡器 (EPLB) 的分布内数据,这可能无法反映现实世界的变化。未来的工作应该在分布发生变化时测试其性能。

灵活的张量并行 (TP) 大小:对于 DeepSeek-V3,密集 FFN 的内存最优 TP 大小较小但大于 1。目前,SGLang 仅支持纯 TP 或 DP,导致内存使用不理想。需要灵活的 TP 选项。

Blackwell 支持:目前,我们的实现仅支持 NVIDIA Hopper 架构。我们正在积极努力扩展对下一代 Blackwell 架构的兼容性。

结论

通过利用 PD 分解、EP 和精心设计的并行设计,我们在 SGLang 中复现了 DeepSeek 的推理框架,并实现了卓越的性能。我们的开源成果——实现了每秒 52.3k 个输入 token 和每秒 22.3k 个输出 token——展现了 SGLang 在大规模 LLM 推理方面的强大能力。我们诚邀社区探索、复制和扩展这项工作,以突破高效 AI 部署的界限。

Logo

欢迎加入DeepSeek 技术社区。在这里,你可以找到志同道合的朋友,共同探索AI技术的奥秘。

更多推荐