ForgeMegakernel speeds up model decoding with tailor-made GPU kernels
ForgeMegakernel: A General Framework for Efficient Auto-Regressive Model Decode Megakernels
Distributed, Parallel, and Cluster Computing
Summary
Auto-regressive models generate text one token at a time, but decoding each token is slow because moving data in and out of memory limits speed. The authors created ForgeMegakernel, which automatically builds custom compute kernels for each model to decode tokens faster and correctly. ForgeMegakernel uses a step-by-step design guide and tests during creation to ensure the kernels work well and stay accurate. Their system improved decoding speed notably on many models, making text generation more efficient.
What this means in practice
- •For machine learning engineers: Build customized GPU kernels for faster decoding of auto-regressive models at inference time, improving throughput in deployment environments.
- •For gpu software developers: Develop efficient megakernels that manage parallelism and synchronization without global barriers, enabling more balanced GPU workloads for diverse models.
Authors
Leshan Li, Zhui Zhu, Xianglong Deng, Yaojian Chen, Qingfeng He, Yuxuan Li, Rong Zhao, Xu Han, Zhiyuan Liu
Abstract
Auto-regressive model decode is bandwidth-bound, since every weight and key/value-cache byte crosses high-bandwidth memory once per token. A megakernel is an ideal solution, but existing automatic megakernel generation approaches cannot achieve both generalization across models and correctness guarantees. We present ForgeMegakernel, which generates a per-model high-performance decode megakernel using coding agents. ForgeMegakernel pairs a universal knowledge base of ten progressive milestones with an independent mid-state test oracle. The milestones provide the megakernel's structural properties: a fine-grained instruction stream for each SM, dependency counters replacing the global synchronization, and a shared-memory buffer pool for workload balance across SMs and greater parallelism. The test oracle derives the mid-states of the megakernel and checks the performance, error and precision during the generation process, guaranteeing a correct and trustworthy forged megakernel. We evaluated ForgeMegakernel on 14 representative decoding operations across eight model families spanning 0.6B-13B parameters. The generated megakernels achieved 50.5-85.9% MBU and geometric mean speedups of 1.21x over SGLang 0.5.18 and 1.54x over a megakernel compiler under identical configurations. Inside SGLang, evaluated on GSM8K with ragged prompts, all 14 megakernels decoded faster than the SGLang engine at comparable answer accuracy.