Eliminating Hidden Serialization in Multi-Node Megakernel Communication

发表时间: 2026-05 · arXiv:2605.00686 (Cornell)

作者/机构: Byungsoo Oh (Cornell University), Rachee Singh (Cornell University)

A1 主要贡献

本文研究的核心问题是,为专家混合(MoE)推理设计的、在单节点上表现优异的“巨核”(megakernel)架构,在扩展到多节点环境时出现的显著性能衰退。这些巨核通过将专家计算与细粒度的、由GPU发起的通信融合成一个持久化的GPU内核,在单节点上通过在tile粒度上重叠数据传输和计算,超越了基于集合通信的MoE方法。然而,当专家分布在由RDMA网络连接的多个节点上时,这种优势无法平滑地延续。对于通信密集型的MoE模型,在8个节点上的性能衰退高达10倍,并且随着节点数量的增加,衰退情况会加剧。

研究目标与根本原因分析
本文旨在解决多节点巨核通信中的性能瓶颈。研究发现,性能衰退的根源在于基于代理(proxy-based)的RDMA传输中存在的“隐藏序列化”(hidden serialization)。巨核在每次tile传输后都需要发送一个完成信号,以实现细粒度的计算与通信重叠。为了保证数据在信号到达前完成传输,每次“传输-信号”对之间都需要一个“栅栏”(fence)。在代理模式下,这个栅栏会强制排空(drain)NIC流水线中所有正在进行的数据传输,然后代理才会提交信号。这种操作的开销随着并发传输数量的增加而增长。因此,对于那些单专家计算量过小、不足以吸收这种膨胀的网络延迟的模型,通信开销便暴露在关键路径上,导致性能下降,如图1所示。


图 1: 顶图:弱扩展性(理想=1.0x)。Llama4使用16个专家,将仅专家并行(EP-only)的扩展限制在4个节点(16个GPU)。底图:4个节点上的GPU流式多处理器(SM)追踪图。Llama4显示出密集的重叠;Qwen3由于信号延迟显示出频繁的停顿。

核心创新点 (Perseus系统)
为了消除这种序列化瓶颈,本文提出了Perseus系统,它包含两项互补的技术(如图2所示):

  1. 解耦信令 (Decoupled Signaling):在巨核协议层,该技术将数据传输(PUTs)与宣告其到达的信号(SIGNALs)分离开来。GPU线程块连续地、非阻塞地发出数据传输请求。然后,每个目标组的指定领导者(leader)线程发出一个单独的栅栏,随后是该组的所有信号。通过按目标GPU对信号进行分组,栅栏的数量从每个专家一次减少到每个远程GPU一次(在文中的Qwen3示例中减少了8倍),同时保持了传输的并发性。

  2. NIC侧排序 (NIC-side Ordering):在传输层,该技术解决了剩余的单次栅栏开销问题。Perseus不再通过阻塞代理来排空NIC流水线,而是在信号提交时附加一个硬件栅栏标志,将排序的执行委托给NIC硬件。这样,代理永远不会阻塞,并且只有每组的第一个信号需要携带栅栏标志。


图 2: 3次带信号传输的代理-NIC交互。(a)原生(Vanilla):每次耦合的put+signal插入一个代理栅栏(3次代理停止)。(b)解耦信令:PUT连续提交;一个代理栅栏先于所有信号(1次代理停止)。(c)NIC侧排序:代理通过每个信号的栅栏标志(siF)将排序委托给NIC;代理从不阻塞,但每个栅栏标志会序列化NIC流水线(0次代理停止,3次NIC停顿)。(d)Perseus(两者结合):只有第一个信号携带栅栏标志(0次代理停止,1次NIC停顿)。

主要成果
在基于代理的传输(Libfabric和IBRC)上,Perseus实现了高达10.3倍的端到端加速。值得注意的是,在IBRC(代理模式)上运行的Perseus性能媲美甚至超越了原生的IBGDA GPU-direct(GPU直通)模式,最高超出1.2倍。这表明,是序列化问题,而非代理模式与GPU直通模式的选择,限制了多节点巨核的性能。该优化具有通用性,无需修改应用代码即可将Triton-distributed的ALLTOALL基准测试性能提升高达79倍。

A3 背景知识与关键观察

2 分布式巨核的兴起

分布式专家混合(MoE)推理暴露了机器学习框架所依赖的CPU驱动执行模型的低效性。在该模型中,主机CPU一次一个地向GPU流中启动称为内核(kernel)的GPU任务,每个内核在所有GPU的流式多处理器(SM)上运行至完成,然后下一个内核才能开始。GPU间的通信使用主机启动的集合通信操作,这在每个参与的GPU上都施加了全局同步屏障。这些在计算和通信内核中的粗粒度同步点导致SM在大量时间内处于空闲状态。作为回应,一类新的巨核设计应运而生。巨核将整个MoE层,甚至整个模型,融合成一个持久化的GPU内核,并在tile粒度上将计算与GPU发起的通信重叠起来【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】【9, Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs, 2025, arXiv】【48, Designing a Low-Latency Megakernel for Llama1B, 2025, Hazy Research Blog】【55, Comet: Fine-grained computation-communication overlapping for mixtureof-experts, 2025, MLSys】【58, Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler, 2025, arXiv】。本节提供了关于巨核如何解决分布式MoE推理性能瓶颈的背景知识。

2.1 专家混合与专家并行

MoE架构。在大型语言模型中,一个标准的Transformer块由一个自注意力模块和一个密集的前馈网络(FFN)组成【51, Attention is all you need, 2017, NeurIPS】。专家混合(MoE)块用E个形状相同的FFN(称为专家)集合以及一个轻量级的门控网络替换了单个FFN,该门控网络为每个token选择一小部分专家【24, GShard: Scaling giant models with conditional computation and automatic sharding, 2021, ICLR】【28, DeepSeek-V3 technical report, 2024, arXiv】【33, Hugging Face model card: Llama-4-Scout17B-16E, 2025, Meta】【45, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017, arXiv】【53, Qwen3 technical report, 2025, arXiv】。对于每个token,门控网络会生成对所有专家的亲和度分数,并将其路由到得分最高的k个专家,其中k远小于E。MoE层的输出是所选专家输出的加权组合。这种条件计算将参数数量与每个token的浮点运算量(FLOPs)解耦,因此增加专家可以增长模型容量,而不会相应地增加推理成本。大多数现代前沿模型都使用MoE架构,例如,DeepSeek-V3使用E=256个专家和top-8路由【28, DeepSeek-V3 technical report, 2024, arXiv】,Qwen3-235B使用E=128【53, Qwen3 technical report, 2025, arXiv】,Llama4 Maverick使用E=128【33, Hugging Face model card: Llama-4-Scout17B-16E, 2025, Meta】。

专家并行 (EP)。前沿MoE模型中的总专家权重远超单个GPU的内存,因此专家通过专家并行(EP)被划分到多个GPU上【24, GShard: Scaling giant models with conditional computation and automatic sharding, 2021, ICLR】【44, DeepSpeed-MoE: Advancing mixture-of-experts inference and training to power next-generation AI scale, 2022, ICML】。设P为专家并行组中的GPU(处理单元或PE)数量;每个PE托管E/P个专家。由于门控网络在运行时动态分配token,任何PE都可能向任何其他PE发送token。因此,每个MoE层包括四个步骤:(1)本地门控计算,为每个token分配k个专家;(2)一个分发(dispatch)ALLTOALL操作,将每个token移动到托管其所选专家的PE上;(3)在接收到的token上进行专家计算,包括两次通用矩阵乘法(GEMM)和一个激活函数;(4)一个合并(combine)ALLTOALL操作,将专家输出返回给源PE进行加权聚合。分发和合并操作位于每一层的关键路径上,其延迟随E和P的增加而增长。

2.2 传统的CPU驱动执行模型

基于内核序列的执行。在现代ML框架如PyTorch【39, PyTorch: An imperative style, highperformance deep learning library, 2019, NeurIPS】中,GPU工作负载被表示为主机CPU提交到GPU流的一系列短时内核。例如,矩阵乘法、激活函数、归约操作的PyTorch代码会编译成一个或多个GPU内核启动。主机CPU将每个内核加入GPU流队列后立即返回,而GPU则按FIFO顺序执行队列中的内核。这种执行模型具有吸引力,因为它是可组合的(算子可以作为库调用组合)和可调试的(每个内核是独立的故障单元)。对于规模适中的密集模型,它也是高效的,因为内核足够大,启动开销可以被分摊。


图 3: 巨核执行模型。每个GPU运行一个由主机一次性启动的持久化内核。CTA专门化为处理器或操作系统(订阅者+调度器)。两个GPU上的对称内存形成一个PGAS地址空间(绿色虚线),支持单边通信。工作流程:❶ 一个GPU上的处理器向对等方的对称内存发出PUTWITH-SIGNAL;❷ 订阅者轮询标志并通知调度器;❸–❹ 调度器入队并分配任务;❺–❻ 处理器出队,读取到达的tile,并进行计算。为清晰起见,左侧GPU的操作系统内部、任务队列和主机CPU被省略;两个GPU运行相同的巨核代码。

CPU驱动模型的局限性。CPU驱动模型的两个特性限制了MoE的推理性能。首先是模型对每个内核施加的启动开销。一个朴素实现的MoE层会分发数百个短时内核。最近的研究报告称,DeepSpeedMoE每层有多达550次内核启动,Megatron-LM有261次【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】。每次启动都带有固定的主机提交成本,当内核如此短且数量如此之多时,主机可能无法在GPU消耗之前填充其流,导致GPU SM空闲。其次是CPU驱动模型在集合通信边界强制执行的“全有或全无”依赖。包括ALLTOALL在内的集合通信原语作为一个整体完成,因此整个分发必须完成后,任何token才能进入专家计算;整个专家计算必须完成后,合并才能开始。

瓶颈的复合效应。这两个特性对于像MoE推理这样的通信密集型工作负载会产生复合效应。关键路径由集合通信主导,而这些集合通信本身又穿插在数百个短内核之间,其启动间隙导致GPU SM空闲。最近在两台A100 GPU上的测量显示,在MoE前向传播过程中,GPU有60-90%的时间处于空闲状态,其中内核启动间隙和未重叠的ALLTOALL各自贡献了大部分空闲时间【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】。

2.3 巨核:持久化GPU内核

融合与持久化。巨核通过将一系列算子(传统上每个都是一个独立的内核)融合成一个单一的、持久化的GPU内核,在计算中途从不将控制权返回给主机,从而同时解决了CPU驱动执行模型的两种成本【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】【9, Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs, 2025, arXiv】【11, Flashattention: Fast and memoryefficient exact attention with io-awareness, 2022, NeurIPS】【48, Designing a Low-Latency Megakernel for Llama1B, 2025, Hazy Research Blog】【55, Comet: Fine-grained computation-communication overlapping for mixtureof-experts, 2025, MLSys】【58, Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler, 2025, arXiv】。例如,最近的工作用一个内核替换了传统MoE层中的260-550个内核【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】。

Tile级并行。巨核的操作粒度是tile,即在编译时选择的输入矩阵的固定大小子矩阵(例如,128×64【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】),而不是整个输入张量。在门控网络将每个token路由到其top-k专家后,巨核将token分组为tile,每个tile遍历由分发、GEMM、激活、第二个GEMM、合并和累加组成的依赖图。巨核按tile跟踪依赖关系,因此一个tile在其输入准备就绪的瞬间就进入下一阶段,而不管任何其他tile的进度如何。

管理Tile依赖。由于巨核从不将控制权返回给主机,它必须在GPU上管理自己的执行。它通过将协作线程阵列(CTA),即GPU计算的可并行化单元,专门化为不同的角色来实现这一点(图3)。大多数CTA作为处理器,对单个tile执行计算(例如GEMM)。少数CTA或CTA内的warp扮演操作系统的角色,协调依赖流水线。这个操作系统包括一个订阅者,用于解码来自对等GPU的tile到达通知,以及一个调度器,用于维护tile级任务的就绪队列并将其分配给空闲的处理器【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】【9, Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs, 2025, arXiv】。因为调度器在GPU上运行并直接观察就绪状态,一个tile从“已到达”到“计算中”的转换无需主机往返。由此产生的巨核设计报告称,在这种组织下SM利用率超过93%,而CPU驱动的基线为9-60%【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】。

GPU发起的通信。巨核完全在GPU上运行,直到完成才将控制权返回给主机,因此包括ALLTOALL在内的主机启动的集合通信从巨核内部是无法访问的【50, The landscape of gpu-centric communication, 2024, arXiv】。即使可以访问,tile级调度也只有在通信可以在tile粒度上发起和观察时才有用,而集合通信的单位是模型的整个层,其完成是一个全局屏障,这正是巨核旨在消除的粗粒度同步。因此,巨核需要既是设备发起(device-initiated)又是单边(one-sided)的通信。设备发起意味着由GPU线程而不是主机CPU发起数据传输,这使得通信可以在持久化内核内部表达。单边意味着发送方可以直接写入接收方的内存,而无需接收方发布匹配的接收请求,这使得巨核中的生产者CTA能够在tile准备好的瞬间将其交给消费者,而无需等待消费者到达会合点。

全局寻址支持单边通信。单边通信要求发送方知道数据应落在接收方GPU内存的哪个位置,因为接收方不参与传输,无法提供目标地址。分区全局地址空间(PGAS)模型【8, Introducing openshmem: Shmem for the pgas community, 2010, PGAS】为发送方提供了这种知识。每个PE在相同的虚拟偏移量处分配一个对称内存区域,因此这些区域的联合形成一个单一的逻辑地址空间,发送方可以在其中使用普通的指针运算来命名任何远程PE上的任何位置,无需事先协调。包括NVSHMEM【38, NVIDIA OpenSHMEM Library (NVSHMEM), NVIDIA Corporation, GitHub】、rocSHMEM【5, ROCm OpenSHMEM (rocSHMEM), AMD, GitHub】、NCCL设备API【13, NCCL EP: Towards a Unified Expert Parallel Communication API for NCCL, 2026, arXiv】【15, GPU-Initiated Networking for NCCL, 2025, arXiv】、NCCLX【47, Collective communication for 100k+ GPUs, 2025, arXiv】和PyTorch Symmetric Memory【42, PyTorch Symmetric Memory, PyTorch Docs】在内的GPU发起通信库向代码公开了此模型,使PE上的CTA能够通过单个库调用写入另一个PE上的对称缓冲区。

单边通信中的信令。将PGAS与巨核中的tile级调度联系起来的通信原语称为“带信号的放置”(put-with-signal)。一个完成了为一个远程专家暂存token的处理器CTA会调用putmem_signal_nbi()(图3 ❶),该调用执行两个从程序员角度看是原子的操作:(1)将token非阻塞地写入远程对称缓冲区(PUT),以及(2)在数据保证在远程可见后,写入远程PE上的一个小标志变量(SIGNAL)。远程订阅者CTA轮询该标志;当标志翻转时,订阅者将依赖的计算任务入队,处理器CTA会拾取它们。整个交换过程无需主机参与或集合通信屏障,并且发送方CTA在调用返回的瞬间就移至其下一个任务,使得这种方法非常高效。

3 将巨核扩展到单节点之外

最先进的分布式MoE巨核在单台拥有8个H100 GPU的多GPU服务器上,相比基于内核的MoE基线,实现了高达6倍的延迟降低和9倍的SM利用率提升【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】,相关的巨核系统在单节点配置中也报告了类似的增益【9, Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs, 2025, arXiv】【55, Comet: Fine-grained computation-communication overlapping for mixtureof-experts, 2025, MLSys】【58, Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler, 2025, arXiv】。然而,前沿的MoE模型无法在单个节点上服务。总专家权重超过了任何单个8-GPU服务器的HBM,其中DeepSeek-V3每个token激活671B参数中的37B【28, DeepSeek-V3 technical report, 2024, arXiv】,而Qwen3-235B和Llama4 Maverick每个层都使用128个专家【33, Hugging Face model card: Llama-4-Scout17B-16E, 2025, Meta】【53, Qwen3 technical report, 2025, arXiv】。因此,生产部署通常在多个节点上运行专家并行,通常是跨越2-16个节点的16-64个GPU【27, RDMA Point-to-Point Communication for LLM Systems, 2025, arXiv】【28, DeepSeek-V3 technical report, 2024, arXiv】【37, Scaling Large MoE Models with Wide Expert Parallelism on NVL72 Rack Scale Systems, 2025, NVIDIA Blog】。现有研究没有回答单节点巨核的增益是否能延续到多节点部署的问题。

3.1 多节点弱扩展中的性能崩溃

实验设置与观察。我们通过将一个最先进的MoE巨核【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】分布到NERSC Perlmutter超级计算机【34, Perlmutter, NERSC Docs】的八个节点上来回答这个问题。每个Perlmutter GPU节点包含四个NVIDIA A100 GPU,每个GPU都有一个专用的HPE Cassini NIC,通过libfabric【14, A brief introduction to the openfabrics interfaces-a new network api for maximizing high performance application efficiency, 2015, HOTI】暴露RDMA。这些NIC连接到一个三级蜻蜓(dragonfly)网络结构,该结构在节点之间提供全互联,最坏情况路径最多三跳。为了表征巨核在部署规模增长时的行为,我们采用弱扩展(weak scaling),这是一种标准技术,即在增加GPU的同时保持每个GPU的工作负载不变,从而使总问题大小与配置的资源成比例增长。在理想的弱扩展下,每次迭代的延迟应保持近乎平坦。我们通过固定每个GPU的工作负载(即每个设备发送的token数)并在一到八个节点上进行扫描,将此方法应用于我们的Perlmutter部署。我们评估了三种最先进的MoE模型的扩展巨核,这些模型涵盖了计算-通信谱系:Qwen3-30B【53, Qwen3 technical report, 2025, arXiv】是通信密集型,GPT-OSS-120B【1, gpt-oss-120b & gptoss-20b model card, 2025, arXiv】是平衡型,Llama4-Scout-17B【33, Hugging Face model card: Llama-4-Scout17B-16E, 2025, Meta】是计算密集型。

通信密集型模型的性能下降。图1(顶部)显示,随着节点数的增加,巨核性能下降,并且对于通信与计算比率较高的模型,性能下降更严重。通信密集型的Qwen3-30B在8个节点上相对于单节点基线减速高达10倍。平衡型的GPT-OSS在节点数每次翻倍时性能下降1.3-1.7倍。计算密集型的Llama4扩展接近线性,仅有1.1-1.3倍的减速。性能下降与通信计算比之间的这种相关性是我们第一个迹象,表明瓶颈在于节点间通信路径而非计算流水线。SM追踪图证实了这一假设(图1,底部)。处理器CTA在tile计算之间停顿,等待迟到的信号,而本应通过tile级并行保持繁忙的SM最终却在等待网络传输。

性能崩溃的意外性。这种性能崩溃是令人惊讶的,因为巨核抽象本身在部署从一个节点扩展到多个节点时并未改变。调度器仍然根据标志翻转进行分派,发送方仍然向对称缓冲区发出PUT-WITH-SIGNAL。变化的是PUT在底层是如何实现的。在节点内部,PUT是通过NVLink的写入,以NVLink的延迟完成,并且先前的工作已显示NVLink传输随并发性近乎线性扩展【25, Evaluating modern gpu interconnect: Pcie, nvlink, nv-sli, nvswitch and gpudirect, 2019, IEEE TPDS】。跨节点时,PUT遍历一个RDMA传输,而该传输的结构是我们接下来要研究的。

3.2 设备发起的RDMA及其提交路径

RDMA基础。远程直接内存访问(RDMA)允许发送方将数据传输到远程机器的内存中,而无需在数据路径上涉及任何一台机器的CPU。发送方指定本地源和远程目标内存区域,两个NIC协作在它们之间复制字节。发送NIC从本地内存读取源区域并发出网络传输,接收NIC将到达的字节写入远程主机的目标区域。在ML服务器中,这些内存区域位于GPU HBM中,每个NIC通过PCIe访问其本地GPU的HBM【36, GPUDirect RDMA, 2024, NVIDIA Docs】。


图 4: 基于代理的RDMA提交路径。GPU线程将网络请求排入一个主机常驻缓冲区,一个专用的CPU代理线程将该缓冲区的内容排空到NIC。NVSHMEM为每个PE使用一个代理通道,来自一个GPU上所有CTA的请求都通过一个共享的提交路径进行汇集。

设备发起的RDMA的提交路径。巨核使用设备发起的RDMA,其中GPU线程通过像putmem_signal_nbi()这样的调用从GPU内核内部触发传输。设备发起是编程模型的一个属性,它允许GPU成为数据传输的语义发起者。另一个维度决定了哪个处理器将请求转发给NIC硬件,我们称之为RDMA传输的提交路径。设备发起的RDMA有两种提交路径。在GPU-direct提交路径中,以NVIDIA在InfiniBand上的IBGDA【32, Improving Network Performance of HPC Systems Using NVIDIA Magnum IO NVSHMEM and GPUDirect Async, 2022, NVIDIA Blog】为例,GPU线程通过NIC向GPU暴露的硬件接口向NIC提交请求,没有CPU线程参与此过程。在基于代理的提交路径中,GPU线程无法到达NIC,而是将请求写入一个主机常驻队列。一个专用的CPU代理线程排空队列并将每个请求代表GPU转发给NIC。

生产环境中的代理提交。大多数生产网络依赖于基于代理的提交。GPU-direct提交路径需要NIC硬件支持,而这种支持在包括AWS EFA【6, Elastic Fabric Adapter, AWS】、HPE Slingshot-11【12, An in-depth analysis of the slingshot interconnect, 2020, SC】和基于Broadcom的云部署【31, UCCL-EP: Portable Expert-Parallel Communication, 2025, arXiv】在内的几种网络结构上是缺失的。NVSHMEM通信库支持两种基于代理的传输后端:Libfabric和InfiniBand可靠连接(IBRC)【14, A brief introduction to the openfabrics interfaces-a new network api for maximizing high performance application efficiency, 2015, HOTI】【32, Improving Network Performance of HPC Systems Using NVIDIA Magnum IO NVSHMEM and GPUDirect Async, 2022, NVIDIA Blog】。一个可移植的多节点巨核必须在基于代理和GPU-direct的RDMA传输上都能正确工作并表现良好。

巨核流量下的代理提交路径。图4显示了基于代理的传输如何处理设备发起的请求。巨核通常使用的通信库NVSHMEM为每个PE分配一个代理通道,GPU上的每个CTA都将其请求写入该通道的主机常驻队列。代理线程按FIFO顺序排空队列,并一次一个地将请求转发给NIC。通过这个单一通道的并发请求数量取决于巨核如何分配专家计算。在P个PE上分片E个专家,每个PE托管E/P个专家。代理只承载节点间流量,因为节点内传输使用NVLink而不是NIC,所以每个PE为每个远程专家每次分发发出一次传输。用$P_{local}$表示发送方自己节点上的PE数量,一个PE在最坏情况下通过其代理通道发出$(P - P_{local}) \cdot (E/P)$次并发传输。在我们的Qwen3-30B配置中,在4个节点和16个GPU上,每个Perlmutter节点托管4个PE,剩下12个远程PE各持有8个专家,因此一个PE每次分发通过单一代理通道发送96次并发传输。


图 5: 不同消息大小、节点数和并发级别下的信令开销。(a) 信令效率,定义为耦合的put+signal吞吐量归一化到流水线化的put-only吞吐量,随着并发度和节点数的增加而崩溃。(b) 在96次并发传输下,栅栏成本随节点数和消息大小的增加而增长。(c) 在96次并发传输下,栅栏开销占总通信时间的比例高达98%,并随节点数增加而增长。信号成本可忽略不计(<总时间的0.2%),并包含在传输部分。

3.3 根本原因:栅栏引起的序列化

putmem_signal_nbi的实现。如果代理可以流水线化处理巨核所需的大量并发数据传输(例如,每次分发96次传输到一个代理通道),这并不会构成性能瓶颈。我们发现它不能,原因在于代理如何遵守putmem_signal_nbi的排序语义。每个putmem_signal_nbi调用都保证,当信号在接收方可见时,相应PUT的有效载荷已经传输完毕。代理通过将该调用扩展为其通道上的一个三步序列来强制执行此保证:一个PUT,一个等待通道上所有先前PUT完成的FENCE,以及SIGNAL。FENCE使SIGNAL在NIC上保持在其PUT之后,以确保数据传输完成。在节点内部,代理不在路径上,这个FENCE不适用,这就是为什么PUT-WITH-SIGNAL在NVLink上没有可观察到的成本。跨节点时,FENCE强制代理停止排空新请求,直到通道上每个在途的PUT都从NIC返回一个完成信号。

隔离栅栏成本。我们用一个微基准测试来测量FENCE的成本,该测试剥离了图1实验中的应用级复杂性。每个PE向远程PE发出N次并发RDMA传输,我们比较了一个仅发出PUT而不带信号的put-only上限,与巨核调用的耦合PUT-WITH-SIGNAL原语。我们在2-8个Perlmutter节点上,扫描了1到128次并发传输和4KB到4MB的消息大小。

结果分析。图5总结了结果。在8个节点上进行96次并发传输时,耦合的PUT-WITH-SIGNAL吞吐量下降到仅PUT上限的2%(图5a)。这个旨在实现细粒度重叠的机制,反而成了阻止重叠的瓶颈,性能崩溃源于我们接下来将要分离的两个复合效应。

并发栅栏在共享通道中交错。并发的CTA将PUT→FENCE→SIGNAL三元组提交到一个共享的代理FIFO队列中(图6a)。每个FENCE等待其前面的所有工作请求,包括其他CTA发布的PUT,因此随着并发度的增长,每个FENCE会遇到更多在前的PUT,总的排空成本随之攀升(图5a)。这种效应在小消息尺寸时达到峰值,此时FENCE成本主导了传输时间,并且在1MB时仍然存在,此时在8个节点上的效率保持在50%以下(图5b)。

单个栅栏的排空成本随节点数增长。一个FENCE会等待其前面的所有PUT都返回一个完成信号。随着部署增加节点,这些PUT会指向更多不同的远程目标,排空操作会等待所有目标中最慢的那个完成,这是一个随目标数量增长的尾部延迟。图5b量化了这种效应。在96次并发传输时,对于4KB消息,总FENCE时间从2个节点时的0.96毫秒增长到8个节点时的6.1毫秒;对于1MB消息,从3.5毫秒增长到9.2毫秒。图5c显示,在小消息尺寸下,FENCE开销占总通信时间的比例高达98%,即使在4MB时也保持在19%以上。巨核正是在这个区域内操作,发送许多并发的、按专家的传输,以便接收方可以在tile粒度上尽早开始依赖的计算(例如,128 × 64 BF16的tile,每个16KB)。因此,巨核期望获得最大收益的区域,恰恰是基于代理的信令性能下降最严重的区域。

3.4 为何初步解决方案失败

现有方法的局限性。两种自然的方法都无法在保留细粒度重叠的同时消除序列化瓶颈。批量合并(Bulk batching)通过将许多小传输合并为少数大传输来减少FENCE数量,但这样做是回到了巨核旨在消除的粗粒度同步。提高代理并行度,通过多个代理线程或通道,可以消除单一FIFO瓶颈,但需要一种跨通道的排序机制来保持put-with-signal语义,这会重新引入同步并增加复杂性(§8)。

核心矛盾。这些初步方案暴露了问题的核心矛盾。细粒度的、按专家的信令是让接收方调度器在CTA粒度上了解tile依赖性的关键,而将其粗粒度化则会丧失巨核所追求的重叠效益。然而,每个细粒度的信号都需要一个FENCE,而在基于代理的RDMA传输上,每个FENCE都会序列化整个提交通道。因此,要消除这个瓶颈,需要重新设计信令协议和代理传输,以消除序列化并保留细粒度的重叠。

A2 方法细节

4 Perseus 设计

Perseus通过两种互补的机制消除了在§3中识别出的由FENCE引起的序列化,这两种机制在不同的层面上解决问题。在协议层,解耦信令(§4.1)重构了CTA提交PUT和SIGNAL的方式,使得代理FIFO不再将FENCE与数据传输交错。在传输层,NIC侧排序(§4.2)将排序执行委托给NIC硬件,因此代理从不因FENCE而阻塞。这两种机制共同消除了序列化的频率和每次FENCE的成本,同时保留了巨核所依赖的细粒度PUT-WITH-SIGNAL接口。

4.1 解耦信令

问题分析。在原生PUT-WITH-SIGNAL发出的耦合PUT → FENCE → SIGNAL序列中,FENCE是序列化点。每个FENCE都会阻塞代理,直到通道上所有先前的PUT完成,因此在一次FENCE排空期间,代理无法提交新的PUT。在我们运行的例子中(Qwen3-30B,4个节点,16个GPU),每个PE向96个远程专家分发token,因此每次分发会排空代理96次(图6a)。将多个小传输合并成大传输会分摊FENCE成本,但也会减少传输并发性并延迟信号,从而消除了巨核所追求的tile级重叠(§3)。

核心洞察。FENCE只在SIGNAL之前是必需的,而不是在每个PUT之后。如果所有CTA首先提交它们的PUT,并将FENCE推迟到由一个单一的领导者CTA发出SIGNAL时再执行,那么代理FIFO看到的是一连串的PUT,然后是一个FENCE,再然后是一批SIGNAL,而不是紧密交错的三元组序列。这种重新排序带来了两个好处。它允许NIC并发地流水线化处理PUT,因为它们之间没有FENCE。它还将FENCE的数量从每个专家一个减少到每组专家一个,因为一个FENCE现在为组中的所有信号排序。


图 6: 12个远程PE上96个远程专家的CTA到代理交互。(a) 原生:每个CTA发出PUT + FENCE + SIGNAL,序列化所有96次传输。(b) Perseus:阶段1流水线化处理96个PUT,NIC可以自由流水线化。阶段2发出96个信号,其中只有12个携带NIC栅栏标志(每组一个)。领导者CTA(★)发出带标志的SIGNAL;其他CTA继续进行计算。

算法 1 解耦信令

要求: 专家索引 e, 组 g = ⌊e/group_size⌋

# 阶段 1: 数据传输 (所有 CTA)
1: 将专家 e 的 token 暂存到本地缓冲区
2: PUT_NBI(dst, src, size, peer) ▷ 无信号
3: ATOMIC_ADD(counter[g], 1) ▷ 通知领导者
4: if ¬is_leader(g) then
5:   恢复为可调度状态 ▷ 释放以进行计算
6: end if
7: WAIT_UNTIL(counter[g] = group_size)
8: FENCE( ) ▷ 每组一个
9: for each expert e′ in group g do
10:  SIGNAL(sig[e′], peer)
11: end for

机制。Perseus通过一个两阶段协议实现了这一洞察(算法1,图6b)。阶段1执行数据传输。每个CTA暂存其token,发出一个不带SIGNAL的非阻塞PUT,并通过递增一个每组的原子计数器来向领导者宣告完成(第1-3行)。非领导者CTA随后返回到巨核调度器,并可用于计算任务,包括为本地或通过NVLink路由的token进行专家计算(第5行)。这种通信与计算的重叠提高了CTA的利用率(§6)。阶段2执行信令。领导者CTA等待计数器达到组大小,这保证了组中的每个PUT都已进入代理FIFO(第7行)。然后它发出一个单一的FENCE,随后是组中每个专家的信号(第8-11行)。

正确性。算法1通过构造保留了PUT-WITH-SIGNAL的排序。原子计数器确保组中的所有PUT在领导者提交FENCE之前都进入代理FIFO。然后,单一的代理线程保留了提交顺序,因此FENCE跟随组中的每个PUT。FENCE保证所有PUT在任何SIGNAL被处理之前在NIC上完成。通过传递性,组中的每个SIGNAL都在组中的每个PUT之后排序。

选择组粒度。组的大小暴露了FENCE分摊和FENCE内部等待之间的权衡。小组保留了许多FENCE,而大组则迫使每个FENCE等待更多未完成的PUT。为了表征这种权衡,我们将组大小作为一个可调参数,并扫描了远程专家数量的所有除数。图7显示了在S=1K、8个Perlmutter节点上的结果,并将总加速分解为两个部分。第一个是提交重排序(❶)。即使在组大小为1时,FENCE数量与耦合基线相同,仅解耦就能将延迟降低12%(22.7→19.9 ms),因为NIC可以背靠背地流水线化处理PUT,而不是在交错的FENCE后处理它们。第二个是FENCE数量减少(❷)。将组大小从1增加到28,将FENCE从112个减少到4个,并带来了额外的38%延迟降低(19.9→12.3 ms)。进一步粗化粒度的收益递减,与FENCE数量曲线的拐点一致(图7,左)。Perseus默认为按PE分组,其中一个组由发往一个远程PE的专家组成。按PE分组位于解耦收益曲线的拐点处(图7,左),并捕获了大部分FENCE减少的好处,同时限制了每个FENCE必须等待的传输数量。我们的扫描证实,这个默认值在1K到64K的序列长度范围内,与最优组大小的差距在2%以内。暴露的配置允许针对特定工作负载和硬件环境进行进一步调整。

4.2 NIC侧排序

问题分析。虽然解耦信令减少了FENCE的频率,但每个剩余的FENCE仍然要支付排空代理的全部成本。FENCE会轮询一个完成计数器,直到通道上所有未完成的请求都返回一个完成信号,在整个排空过程中,代理是空闲的,无法提交新的请求(图2a-b)。这种排空成本随着在途传输的数量和目标数量的增加而增长(§3)。

核心洞察。代理侧排空并不是强制执行PUT-WITH-SIGNAL排序的唯一方法。现代RDMA NIC暴露了一个每请求的FENCE标志(例如,Libfabric中的FI_FENCE和InfiniBand verbs中的IBV_SEND_FENCE),它将排序执行从软件轮询转移到NIC硬件中。当NIC遇到一个带标志的请求时,它会推迟该请求,直到同一连接上的所有先前请求都已完成,这是通过内部硬件寄存器而不是软件计数器来跟踪的。通过这种方式,NIC在硬件中吸收了排序成本,而代理则继续提交请求。


图 7: 解耦信令的效果(S=1K, 8个节点)。左图:FENCE数量与组大小的关系。右图:延迟与组大小的关系,以及耦合基线(红色虚线)。对加速的贡献可以分解为:❶ 提交重排序(从耦合到组大小为1的解耦),❷ FENCE数量从112(组大小1)减少到4(组大小28)。


图 8: 解耦信令和NIC侧排序的综合效果(4个节点,Qwen3-30B)。基线:原生(红色)和耦合+NIC排序(绿色)作为水平线;解耦曲线在启用NIC侧排序的情况下扫描不同组大小。(a) S=1K:NIC侧排序占主导;组大小影响极小。(b) S=64K:所有三种效应都有贡献。

机制。Perseus用NIC FENCE标志替换了代理的阻塞式排空。当请求FENCE时,代理设置一个待处理标志而不是轮询完成情况,并继续提交工作。发送到NIC的下一个SIGNAL会携带这个标志,指示NIC推迟该SIGNAL,直到其连接上的每个先前请求都已完成(图2c)。因此,代理不会停顿,排序保证由NIC本身来维护。

正确性。NIC FENCE标志提供了与代理侧排空相同的排序保证。在标志之后提交的任何操作都不能在NIC上先于该连接上的所有先前操作处理。由于代理在提交顺序上将PUT放在带标志的SIGNAL之前,NIC会在处理SIGNAL之前完成所有PUT。在多队列对(multi-queue pair)传输上,基于对等方哈希的路由(§5)将一个对等方的PUT及其SIGNAL固定到同一个连接上,因此每个连接的排序足以保证每个对等方的正确性。

4.3 为何两种机制是互补的

独立与组合效果。单独应用于原生的耦合信令时,NIC侧排序消除了所有的代理停止;NIC吸收了排序成本而不是阻塞代理(图2c)。与解耦信令结合使用时,只有每组的第一个信号需要该标志,因为代理顺序提交信号,NIC会保留该顺序(图2d)。

机制分工。图8揭示了两种机制之间的分工。对于小消息(S=1K),NIC侧排序将每次FENCE的成本降至接近零,并使对组大小的敏感度变得平坦;加速来自NIC排序加上提交重排序,FENCE数量无关紧要(图8a)。对于大消息(S=64K),NIC必须推迟每个带标志的SIGNAL,直到前面的PUT完成,而较大的传输需要更长的时间才能完成,因此每次FENCE的成本不再可以忽略不计,减少FENCE数量重新变得有价值(图8b)。此时,所有三种效应都对组合加速做出了贡献。

互补性总结。这两种机制是互补的。解耦减少了FENCE的频率,在FENCE数量占主导时最为重要。NIC侧排序消除了每次FENCE的阻塞,在每次FENCE的开销占主导时最为重要。它们共同覆盖了从开销主导到传输主导的整个工作负载范围。

5 实现

解耦信令。我们在FlashMoE【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】(一个支持多节点的先进开源MoE巨核)之上实现了分离信令,仅修改了RDMA请求逻辑,而保留了核心的巨核组件(对称张量布局、订阅者、调度器、工作者模块)不变。对节点内专家的分发保持不变。对于节点间通信,我们实现了两个过程。在第一个过程中,每个CTA将token暂存到本地堆中,发出不带信令的nvshmem_putmem_nbi(),并通过递增一个原子计数器(cuda::atomic::fetch_add())来通知每个PE的组领导者。在第二个过程中,每个组领导者CTA调用nvshmem_fence(),然后调用nvshmemx_signal_op()

NIC侧排序。我们的NIC侧排序修改了NVSHMEM传输模块(v3.5.21)中不到100行的代码。它用每个请求的NIC栅栏标志(Libfabric的FI_FENCE,IBRC的IBV_SEND_FENCE)替换了代理基于静默的排空操作(Libfabric的fi_cntr_wait,IBRC的check_poll_avail)。这一更改完全在传输共享库(nvshmem_transport_*.so)中。用户只需替换新的库文件,无需修改应用程序代码或重新编译主机应用程序。任何使用NVSHMEM作为其通信后端的系统都可以通过这种替换继承该优化。例如,我们在Triton-distributed【58, Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler, 2025, arXiv】的ALLTOALL基准测试上实现了高达79倍的加速,而没有触及其应用程序代码(§6.5)。

多QP传输适配。在具有多个队列对(QP)的IBRC上,默认的轮询策略会将到同一对等方的操作分散到不同的QP上,但IBV_SEND_FENCE仅在单个QP内强制排序。因此,依赖的PUT和signal操作可能会落在不同的QP上,这破坏了预期的排序。Perseus将所有针对同一对等方的操作固定到一个确定的QP上(qp = pe % num_qps)。这样,依赖的操作现在共享一个QP并继承其FIFO排序,而不同的对等方仍然分散在不同的QP上以保持总带宽。这种设计使得IBV_SEND_FENCE在多QP传输上对NIC侧排序有效。

A4 实验环境

  • 硬件配置:

    • 平台1 (Perlmutter): NERSC Perlmutter超级计算机,最多16个节点(64个A100 GPU)。每个节点包含4个NVIDIA A100 GPU,通过NVLink连接。节点间通过HPE Slingshot-11互联,每个GPU配备一个200 Gb/s的Cassini NIC,构成一个蜻蜓(dragonfly)拓扑。
    • 平台2 (商业云): 一个商业GPU云平台,最多4个节点(32个H100 GPU)。每个节点包含8个NVIDIA H100 GPU。节点间通过400 Gb/s的InfiniBand NDR网络连接,每个GPU配备一个ConnectX-7 NIC。
    • 传输后端: 实验涵盖了三种RDMA传输配置:
      1. Libfabric: 基于代理,在Perlmutter的Slingshot-11上运行。
      2. IBRC (InfiniBand Reliable Connection): 基于代理,在商业云的ConnectX-7上运行,配置了4个QP。
      3. IBGDA (GPU-direct): GPU直通,在商业云的ConnectX-7上运行。
  • 软件配置:

    • 基础系统: 原型基于FlashMoE【3, FlashMoE: Fast Distributed MoE in a Single Kernel, 2025, NeurIPS】和NVSHMEM v3.5.21【38, NVIDIA OpenSHMEM Library (NVSHMEM), NVIDIA Corporation, GitHub】。未修改的版本称为“vanilla”。
    • 通用性验证: 使用Triton-distributed【58, Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler, 2025, arXiv】的ALLTOALL基准测试来验证通用性。
    • 对比基线: 与NCCL集合通信进行比较。
  • 工作负载:

    • 模型: 评估了三个具有不同计算通信比的MoE模型,具体参数见表1。

      • Qwen3-30B: 通信密集型
      • GPT-OSS-120B: 平衡型
      • DeepSeek-V3: 计算密集型
    • 数据集/参数: 序列长度(S)从256到64K tokens,覆盖了开销主导(小S,类似解码)和传输主导(大S,类似预填充)的场景。默认使用均衡的专家路由,专家容量EC = S × k/E

表 1: MoE模型配置。H:隐藏维度,I:中间维度,E:专家数,k:top-k。

A4 实验结果

6.2 端到端性能


图 9: 端到端前向延迟:原生FlashMoE vs. Perseus。(a)–(c) Perlmutter(A100, Slingshot-11, Libfabric),2–16个节点(8–64个GPU):Perseus实现高达10.3倍的加速。(d)–(e) 商业GPU云(H100, ConnectX-7, InfiniBand),2–4个节点(16–32个GPU):Perseus在IBRC上实现高达2.47倍的加速,同时性能与IBGDA GPU-direct持平或超越。标注显示了Perseus相对于原生的加速倍数。

总体性能:如图9所示,Perseus在基于代理的传输上取得了显著的端到端加速。在Libfabric上,加速最高可达10.3倍;在IBRC上,加速最高可达2.47倍。一个关键发现是,在IBRC(代理模式)上运行的Perseus,其性能能够媲美甚至超越IBGDA(GPU直通模式),最高超出1.2倍。这表明,代理提交的开销本身并非瓶颈,真正的瓶颈是栅栏引起的序列化。一旦消除了这个瓶颈,代理模式可以与GPU直通模式竞争。所有实验结果均通过了正确性验证,误差率低于1%。

跨节点扩展性:Perseus解决了原生系统在弱扩展性测试中的性能衰退问题(图1),实现了近乎线性的扩展。随着节点数增加,原生系统的开销因栅栏频率和代理停顿的增加而恶化,而Perseus的性能保持稳定,因此在更大规模的节点上能获得更高的加速比。

跨工作负载性能:加速效果在通信密集型配置上更明显。例如,在Libfabric上,Qwen3的峰值加速为10.3倍,而GPT-OSS和DeepSeek-V3分别为2.8倍和2.2倍。随着序列长度S的增加,传输成本占主导,各模型的加速比趋于收敛。不同传输后端的加速趋势也不同:Libfabric上的加速在小S时达到峰值,而IBRC上的加速随S增长,这反映了它们不同的栅栏成本结构(见附录A)。

6.3 消融实验


8节点消融实验分解


图 10: 左图:在Perlmutter上,Qwen3-30B相对于原生系统的加速消融实验。在较大节点数下,单次栅栏成本增加,栅栏数量减少的比例相对较小,因此NIC侧排序贡献更大;在较小节点数下,情况相反。右图:8节点分解图,显示了栅栏数量和单次栅栏成本(↑ =