在AMD Instinct MI300X GPU上通过Expert Parallelism实现每秒21K输出token的DeepSeek推理

概述

大语言模型(LLM)的重大突破之一是采用了Mixture-of-Experts(MoE)架构。模型不再使用单一的大型feed-forward network(FFN)层,而是包含多个较小的FFN层,每个层被称为一个expert。对于每个token,一个轻量级的gating网络会动态地选择一小部分expert来处理每个步骤。这种条件计算使得MoE模型能够在保持高效计算利用率的同时扩展到极大的参数量。例如,2025年1月发布时引起轰动的DeepSeek-R1 671B模型和2025年8月发布的OpenAI热门开源GPT-OSS 120B模型都使用了MoE架构。

然而,由于模型规模庞大且设计稀疏,像DeepSeek-R1这样的模型需要高级优化技术才能在大规模环境下高效服务。特别是,通过应用Expert Parallelism(EP),我们可以实现高吞吐量(即每秒总输出token数)。这是因为(1)GPU内存中只存储一部分expert,从而允许更大的批量大小,以及(2)单个expert的参数在加载后可以更多地被复用,从而提高FLOPS per byte。然而,在高效实现称为dispatch/combine的细粒度通信模式——将激活值传输到负责被激活expert的GPU——以及解决expert之间的不平衡方面,仍然存在挑战。

AMD的软件合作伙伴Moreh最近展示了通过在ROCm软件栈上实现EP,可以高吞吐量地执行DeepSeek-R1推理。要在LLM推理中实际达到如此高的性能,需要在软件栈的多个层级上进行精心优化。Moreh团队在库和模型实现两个层面都应用了各种优化。最终,他们在配备8x AMD Instinct MI300X GPU的服务器上实现了>21,000 tokens/sec的解码吞吐量。

并行化方法

在DeepSeek R1模型中,每个decoder块包含256个expert,其中每个token激活8个。此外,还有一个无论路由结果如何都始终执行的shared expert,以及DeepSeek引入的新注意力机制Multi-Head Latent Attention(MLA)层。EP通常与data parallelism(DP)结合实现。每个GPU持有并执行不同的expert子集。同时,对于始终执行的组件——如shared expert和MLA——每个GPU并行处理单独的请求批次。因此,我们在每个节点上应用了DP=8和EP=8的并行化配置。即每个GPU负责32个expert。

此外,应用prefill-decode disaggregation——即在不同节点上分别执行prefill和decode阶段——可以进一步提高集群系统的整体吞吐量。prefill阶段一次性处理整个输入序列,属于计算密集型(compute-bound),而decode阶段以自回归方式逐个生成输出token,属于内存密集型(memory-bound)。不仅分离执行,还根据prefill和decode不同的计算特性定制并行化和优化配置,可以进一步提高GPU效率。在本文中,由于典型应用中大部分GPU用于解码,我们主要关注decode阶段的性能评估。

性能优化

仅仅实现一种新的并行化方法并不会自动产生最佳性能。要充分释放硬件的潜力,必须将从库到模型实现的整个软件栈优化以适应新的计算和通信模式。Moreh团队基于vLLM实现了以下多项优化。

GPU间负载均衡

实现EP的主要挑战之一是工作负载在GPU之间分配不均,且随路由结果而变化。虽然MoE模型在训练时尽量最小化这种不平衡,但仍存在固有的局限性。DeepSeek也承认了这一问题,并报告使用了他们自己的名为EPLB的负载均衡方法来解决它。在不应用任何显式负载均衡策略的情况下,当expert被简单地分为八个连续块时——即将expert 0-31分配给第一个GPU,32-63分配给第二个,……,224-255分配给最后一个——我们观察到GPU之间最高达2倍的工作负载不平衡。我们设计了一种算法,测量每个expert的激活频率,并将256个expert分为八组(每组32个),使得每组的总激活频率尽可能均衡。(这可以看作是均衡数字分区问题的变体。)该算法随后对expert重新排序,使每组连续放置,并在此基础上应用EP。由于expert排列发生了变化,gating网络也必须生成反映新expert排列的路由结果。最终,我们成功将GPU之间的工作负载不平衡降低到5%以内。据我们所知,这是在AMD GPU上的首个EP负载均衡实现。在实际服务场景中,可以通过定期重新测量激活频率并相应生成新的expert排列来实现持续的负载均衡。

计算-通信重叠

应用EP的另一个挑战是最小化dispatch/combine all-to-all通信带来的开销。将批次分为两个微批次并并发执行,可以使一个微批次的通信与另一个微批次的计算重叠。当批量大小足够大时——即在高吞吐量场景下——这种方法可以在不牺牲计算性能的情况下有效隐藏通信延迟。这种方法有多种实现变体,其中我们采用了vLLM的DBO(Dual Batch Overlap)系统。我们修改了DBO系统以使用MORI-EP,特别是增强了MORI-EP库,使其通信操作可以从DBO的两个工作线程中调用,并与其他计算正确同步。

库级优化

Moreh团队还应用了多项库级优化,以最大化GPU效率并改善吞吐量(tokens/sec)和延迟(首token时间和token间延迟)。以下是部分优化示例:

• 最优GEMM和Attention Kernel选择:无需在线分析和手动调优即可动态选择最优的GEMM和Attention内核。
• Fused MoE Kernel优化:实现了高度优化的fused MoE内核,特别是在小批量大小时性能优于AMD的AITER库。
• FP8 KV Cache支持:实现了Multi-head Latent Attention(MLA)内核,使KV cache能够以FP8格式存储和加载,特别是在长上下文场景中提高性能。
• 垂直和水平Kernel Fusion:同时采用垂直fusion(如fused RoPE内核)和水平fusion(如合并shared expert中的多个GEMM),以减少内核开销并提高计算效率。

有关这些库级优化及其独立于EP的性能改进的详细信息,请参阅Moreh的技术报告。

利用HIP Graphs

HIP Graphs是一种通过将多个GPU操作捕获为静态图并一次性发出来减少CPU运行时开销的技术。这在decode阶段尤为重要,因为单个操作相对较短,充分利用GPU至关重要。然而,构建静态图要求操作之间传递的张量大小必须固定。在EP中,每个expert的输入大小本质上是动态的,由路由结果决定,这使得将其表示为静态图极具挑战性。我们修改了vLLM中的DeepSeek模型实现,在支持EP的同时尽可能使张量大小静态化。这使我们能够充分利用HIP Graph,从而最小化CPU开销。将MoRI库操作整合到图中面临了多个技术挑战,但我们最终成功地将MORI-EP的dispatch/combine操作作为图执行的一部分纳入其中。

实验方法

我们的实验在以下配置的MI300X服务器上进行。

• CPU: 2x AMD EPYC 9474F 48-core 3.6 GHz
• Main memory: 2,304 GB
• GPU: 8x AMD Instinct MI300X OAM GPU 192 GB
• Server: Gigabyte G593-ZX1-AAX1
• Operating system: Ubuntu 22.04.4 LTS
• ROCm version: 6.4.1

我们参考LMSys博客文章设计了评估方法。在该文章中,评估EP性能的关键指标(假设prefill/decode disaggregation)是每个decode节点的总吞吐量(tokens/sec)和token间延迟。LMSys在12个节点上实现了prefill/decode disaggregation,而我们旨在即使在单节点环境中也能测量decode性能。为此,我们对vLLM进行了以下修改:

• 在服务器端,我们排除了prefill时间,仅累计decode时间和token数量,以单独测量decode吞吐量。token间延迟仍然可以使用标准基准测试方法在客户端测量。
• vLLM V1引擎的调度器原本将prefill和decode请求在单个批次中执行(称为mixed prefill),这使得无法排除prefill时间。我们对其进行了修改,使prefill和decode请求发送到不同的队列并独立执行。

可以使用以下命令向(修改后的)vLLM服务器发送请求并测量性能。

vllm bench serve \
  --backend vllm \
  --model "deepseek-ai/DeepSeek-R1" \
  --metric-percentiles "10,25,50,75,90" \
  --percentile-metrics "itl,tps,ttft,e2el" \
  --port 8001 \
  --num-prompts 2048 \
  --max-concurrency 2048 \
  --ignore-eos \
  --dataset-name sharegpt \
  --dataset-path /app/dataset/ShareGPT_V3_unfiltered_cleaned_split.json \
  --sharegpt-input-len 2000 \
  --sharegpt-output-len 100

实验结果

我们测量了四种不同并发级别(1024、2048、3072和4096)下的吞吐量(每节点)和token间延迟,以观察延迟和吞吐量之间的权衡关系。结果如下。

在最高并发级别下,我们实现了21,224.6 tokens/sec的吞吐量。这与SGLang团队在8x NVIDIA H100 GPU服务器上报告的22,282 tokens/sec几乎相当,证实了我们的实现提供了最先进的EP性能。同时,即使在针对最低延迟优化的配置下(中位数token间延迟为91.48 ms),吞吐量下降也不到15%,表明系统可以根据目标SLO(Service Level Objectives)灵活调整其最大并发量。

结论

Moreh通过为MoE模型应用和优化多种技术,成功在AMD Instinct MI300系列GPU上实现了EP高吞吐量推理。我们还提供MoAI Inference Framework,它能够在AMD Instinct GPU集群上自动应用包括本文所述高效EP实现在内的各种分布式推理技术。更多信息请访问Moreh网站。