Static Batching of Irregular Workloads on GPUs: Framework and Application to Efficient MoE Model Inference

Yinghan Li; Yifei Li; Jiejing Zhang; Bujiao Chen; Xiaotong Chen; Lian Duan; Yejun Jin; Zheng Li; Xuanyu Liu; Haoyu Wang; Wente Wang; Yajie Wang; Jiacheng Yang; Peiyang Zhang; Laiwen Zheng; Wenyuan Yu

Static Batching of Irregular Workloads on GPUs: Framework and Application to Efficient MoE Model Inference

Distributed, Parallel, and Cluster Computing 2025-01-28 v1 Machine Learning

Authors: Yinghan Li , Yifei Li , Jiejing Zhang , Bujiao Chen , Xiaotong Chen , Lian Duan , Yejun Jin , Zheng Li , Xuanyu Liu , Haoyu Wang , Wente Wang , Yajie Wang , Jiacheng Yang , Peiyang Zhang , Laiwen Zheng , Wenyuan Yu

View on arXiv ↗ PDF ↗

Abstract

It has long been a problem to arrange and execute irregular workloads on massively parallel devices. We propose a general framework for statically batching irregular workloads into a single kernel with a runtime task mapping mechanism on GPUs. We further apply this framework to Mixture-of-Experts (MoE) model inference and implement an optimized and efficient CUDA kernel. Our MoE kernel achieves up to 91% of the peak Tensor Core throughput on NVIDIA H800 GPU and 95% on NVIDIA H20 GPU.

Keywords

gpu computing large language model inference computer architecture

Cite

@article{arxiv.2501.16103,
  title  = {Static Batching of Irregular Workloads on GPUs: Framework and Application to Efficient MoE Model Inference},
  author = {Yinghan Li and Yifei Li and Jiejing Zhang and Bujiao Chen and Xiaotong Chen and Lian Duan and Yejun Jin and Zheng Li and Xuanyu Liu and Haoyu Wang and Wente Wang and Yajie Wang and Jiacheng Yang and Peiyang Zhang and Laiwen Zheng and Wenyuan Yu},
  journal= {arXiv preprint arXiv:2501.16103},
  year   = {2025}
}

Comments

11 pages

Static Batching of Irregular Workloads on GPUs: Framework and Application to Efficient MoE Model Inference

Abstract

Keywords

Cite

Comments

Related papers