GPU-Initiated EP Comm. for MoE Training

• Dispatch and Combine with PyTorch Symmetric Memory.
• An optimized implementation featuring memory-pool reuse and zero-copy paths.
• Benchmarked against host-initiated EP (NCCL), with a side-by-side Nsys profile comparison.

WIP Code cleanup and writeup in progress

Early Results on 8xH100 (sxm5), 1-layer MoE Transformer Layers.

Training-step profiles

Observe ranges of fwd, bwd, spot dispatch & combine.

NCCL-EP (dispatch.forward is wide (long) enough to be visible)

SymmMem-EP (dispatch.forward is harder to spot since it is compressed)

References:

• PyTorch Symmetric Memory Documentation
• PyTorch Symmetric Memory: A New Programming Paradigm for Distributed AI - Ke Wen & Chien-Chin Huang
• PyTorch APIs for High Performance MoE Training and Inference - D. Vega-Myhre; Ke Wen & N. Gimelshein