Accelerating MoE with Dynamic In-Switch Computing on Multi-GPUs

TL;DR

This paper introduces DySHARP, a dynamic in-switch computing solution for MoE models on multi-GPU systems, reducing redundant communication and achieving up to 1.79× speedup.

Contribution

DySHARP extends NVLink SHARP with dynamic communication primitives and token-centric kernel fusion to optimize MoE expert parallelism on multi-GPU systems.

Findings

DySHARP reduces redundant inter-GPU traffic in MoE.

DySHARP achieves up to 1.79× speedup over state-of-the-art.

Dynamic communication primitives improve MoE performance.

Abstract

Mixture-of-Experts (MoE) has been adopted by many leading large models to reduce computational requirements. However, frequent inter-GPU communication in MoE expert parallelism (EP) becomes a performance challenge. We observe substantial redundant inter-GPU data transfers in MoE that can be potentially addressed by in-switch computing. Unfortunately, the existing solution, NVLink SHARP (NVLS), can only support static collectives with regular patterns, incapable of dynamic communication with irregular patterns in MoE. To bridge the functionality gap, we propose DySHARP, an integral dynamic in-switch computing solution to accelerate MoE, encompassing both communication primitives and communication-aware scheduling: 1) Dynamic multimem addressing co-designs ISA, architecture, and runtime, as a dynamic extension to NVLS, reducing redundant traffic. However, the resulting traffic reduction is inherently asymmetric between two directions, preventing it from directly translating into speedup. 2) Token-centric kernel fusion deeply fuses the dispatch-computation-combine pipeline, resolving this asymmetry to translate traffic reduction into actual speedup. Compared with the state-of-the-art solution, DySHARP achieves up to 1.79$\times$ speedup.