通过细粒度调度实现专家计算与通信重叠,显著提升MoE模型推理速度。
Fine-grained Computation-Communication Overlap via Tile-level Signaling and Scheduling for Mixture-of-Experts

- 采用分块信号与调度机制,让计算与通信并行进行
- 在4张A100上实现最高2.64倍端到端加速
- 无需修改底层算子,适合部署于多卡分布式系统
Mixture-of-Experts(MoE)架构在不线性增加计算量的前提下提升模型容量,已成为构建万亿参数大语言模型的关键。其高效部署依赖多GPU分布式执行,每个MoE层包含两次all-to-all通信:将令牌分发至专家显卡,再将结果返回源显卡。传统实现中,返回通信在专家计算完成后启动,导致通信延迟暴露在关键路径上,降低GPU利用率。本文提出一种细粒度方法,通过分块级信号与调度,将专家计算与第二次all-to-all通信重叠。设计包括:(1) 持续运行的每卡计算核(生产者),覆盖该卡所有本地专家,消除重复内核启动开销,并优先处理远程关键分块;(2) 在少量专用流式多处理器(SM)上运行持续通信核(消费者),分段粒度发出数据传输,分块就绪即传。该协同设计不侵入底层算子或通信原语,适用于提升多卡系统下分布式MoE执行效率。在4张A100平台上,对三种MoE模型与四种先进MoE系统对比测试,本方法实现最高2.64倍端到端加速和2.74倍MoE层加速。相比传统非重叠基线,在不同GEMM形状、路由器模式及生产者/消费者SM分区配置下,持续提升算子级与MoE层性能,且保证正确性。
原文摘要 · Abstract (English)
Mixture-of-Experts (MoE) architectures increase model capacity without proportionally increasing computation cost and have become a key building block for scaling large language models (LLMs) to trillion-parameter regimes. Efficient deployment of these MoE models relies on distributed execution across multiple GPUs, where each MoE layer involves two all-to-all communications: dispatching tokens to expert ranks and returning the expert outputs to their source ranks. Conventional MoE implementations launch this return all-to-all after expert compute completes, exposing communication latency on the critical path and reducing GPU utilization. We present a fine-grained approach that overlaps expert compute with the second all-to-all via tile-level signaling and scheduling. Our producer-consumer co-design combines: (1) a persistent per-rank computation kernel (producer) that covers all local experts on the rank to eliminate repeated kernel launch overhead and prioritizes remote-critical tiles, and (2) a persistent communication kernel (consumer) on a small dedicated partition of streaming multiprocessors (SMs) that issues segment-granular transfers as tiles become ready. Our co-design avoids intrusive changes to the underlying computation operators or communication primitives, making it practical for improving distributed MoE execution efficiency on multi-GPU systems. On a 4-A100 GPU platform, evaluated on three MoE models against four state-of-the-art MoE systems, our approach achieves up to 2.64x end-to-end speedup and 2.74x MoE-layer speedup. Compared with a conventional non-overlap baseline, our approach consistently improves both operator- and MoE-layer-level performance across varying GEMM shapes, router modes, and a broad range of producer/consumer SM partitions, while preserving correctness.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。