Research and ACM SRC Posters Archive

GPU Kernels for Mixture of Experts

Poster Type: Research Posters

Author: Arthur Feeney (University of California, Irvine), Ying Wai Li (Los Alamos National Laboratory (LANL)), Aparna Chandramowlishwaran (University of California, Irvine)

Supervisor: Aparna Chandramowlishwaran (University of California, Irvine)

Abstract: The Sparsely-Gated Mixture of Experts (MoE) has seen a surge in use over the last year. This is primarily motivated by a desire to increase the size of language models, without a proportional increase in the total number of FLOPs. Due to their popularity, there is a large volume of work studying distributed MoEs for large-scale training. In this work, we find that there is room to improve the performance of MoEs when run on a single GPU—an important case for inference and fine-tuning. In our efforts to improve single-GPU performance, we implement triton kernels for grouped matrix multiplications and gated linear units. These kernels support fusing operations for token routing, in order to reduce the number of access to slow off chip memory.

Best Poster Finalist (BP): no
Poster: PDF
Poster Summary: PDF

Back to Poster Archive Listing