TilePipe

MoE Communication and Overlap Kernels with CuteDSL

Communication Kernels

Dispatch and combine, written as NVSHMEM push kernels over two copy backends: a SIMT copy loop and a TMA async-proxy path. Measured standalone, with the GEMM held out.

Note that as the number of communication SMs increase, both implementations can achieve a relatively high nvlink utilization. In low-SM rank, which is important for communication-compute overlap, TMA-based kernels can achieve higher bandwidth.

uv run torchrun --nproc-per-node 8 tilepipe/moe_comm.py \
  [--test-tma-dispatch | --test-push-combine] --comm-impl [simt | tma] \
  --token-sweep 2048,4096,8192,16384 --dispatch-ctas 8,12,16,24,36,48 \
  --iterations 20 --warmup_iterations 5 --results-dir bench_results/repro

Overlap

Fine-grained communication compute overlap is achieved with per-tile synchornization semaphores. The SM of each GPU is partitioned for the compute and communication kernel. In the producer kernel, after a tile finishes, it increments a semaphore that the consumer kernel waits on.

Because gemm kernels require a specific execution pattern to maximize L2 cache reuse such as expert-major, the schedule of the communication kernel is design to accomodate this order with earliest-deadline first scheduling that intuitively unblocks the consumer kernel as soon as possible to maximize overlap. Also, synchonization granularity (per tile/token/expert) can have noticable impact on performance: a finer graiend synchonization could potentially unblock communication sooner, but it adds more flag-setting overhead in the producer kernel and can disrupt the pipeline of gemm kernels.

The result of dispatch and combine kernel performance is shown below, and a time breakdown of the kernels are shown below. We included more detailed analysis in tilepipe/docs.

The commands to reproduce the results are below:

# GEMM -> combine
uv run torchrun --nproc-per-node 8 tilepipe/gemm_combine.py \
  --token-sweep 2048,4096,8192,16384 --results-dir bench_results/repro

# dispatch -> GEMM
uv run torchrun --nproc-per-node 8 tilepipe/dispatch_gemm.py \
  --gemm-n 4096 --experts 256 --copy simt --comm-warps 32 \
  --comm-sms 8,12,16,24,36 --token-sweep 2048,4096,8192,16384 \
  --results-dir bench_results/repro

GEMM kernel is based on QuACK (Apache-2.0), see ATTRIBUTION.md.