Static Batching of Irregular Workloads on GPUs: Framework and Application to Efficient MoE Model Inference
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866913666996109312 |
|---|---|
| author | Li, Yinghan Li, Yifei Zhang, Jiejing Chen, Bujiao Chen, Xiaotong Duan, Lian Jin, Yejun Li, Zheng Liu, Xuanyu Wang, Haoyu Wang, Wente Wang, Yajie Yang, Jiacheng Zhang, Peiyang Zheng, Laiwen Yu, Wenyuan |
| author_facet | Li, Yinghan Li, Yifei Zhang, Jiejing Chen, Bujiao Chen, Xiaotong Duan, Lian Jin, Yejun Li, Zheng Liu, Xuanyu Wang, Haoyu Wang, Wente Wang, Yajie Yang, Jiacheng Zhang, Peiyang Zheng, Laiwen Yu, Wenyuan |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2501_16103 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Static Batching of Irregular Workloads on GPUs: Framework and Application to Efficient MoE Model Inference Li, Yinghan Li, Yifei Zhang, Jiejing Chen, Bujiao Chen, Xiaotong Duan, Lian Jin, Yejun Li, Zheng Liu, Xuanyu Wang, Haoyu Wang, Wente Wang, Yajie Yang, Jiacheng Zhang, Peiyang Zheng, Laiwen Yu, Wenyuan Distributed, Parallel, and Cluster Computing Machine Learning D.1.3; I.2.6 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. |
| title | Static Batching of Irregular Workloads on GPUs: Framework and Application to Efficient MoE Model Inference |
| topic | Distributed, Parallel, and Cluster Computing Machine Learning D.1.3; I.2.6 |
| url | https://arxiv.org/abs/2501.16103 |