bibtext: @inproceedings{
  tan2024scattered,
  title={Scattered Mixture-of-Experts Implementation},
  author={Shawn Tan and Yikang Shen and Rameswar Panda and Aaron Courville},
  booktitle={First Conference on Language Modeling},
  year={2024},
  url={https://openreview.net/forum?id=YDZ7GeFLxq}
  }
abstract: ScatterMoE is an implementation of Sparse Mixture-of-Experts (SMoE) on GPUs. ScatterMoE builds upon techniques in existing implementations, and overcoming some of the current limitations to improve batched inference, training speed, and memory footprint. This implementation achieves this by avoiding padding and making excessive copies of the input. We also fuse expert linear transforms and reordering operations with ParallelLinear, a module that can be used to extend the concept of SMoEs. We benchmark our implementation  against Megablocks, and show that it enables a higher throughput and lower memory footprint. We also show how ParallelLinear enables extension of the Mixture-of-Experts concept by demonstrating with an implementation of Mixture-of-Attention.
title: Scattered Mixture-of-Experts Implementation
authors: Shawn Tan Yikang Shen Rameswar Panda Aaron Courville