Automatic Kernel Generation for Large Language Models on Deep Learning Accelerators

Large language model (LLM) is a promising trend to sustain accuracy growth with billions of parameters for various application domains. Deep learning (DL) accelerators that are typically designed in spatial architecture are potential platforms to handle the substantial computational demands of LLMs. Exploiting DL accelerators for LLMs require high-performance kernels, which are manually optimized or automatically generated. While prior automatic works reduce development costs, they fail to find optimal or near-optimal kernels. This is because the kernel design space is large containing many invalid kernels, and non-convex containing many local minima. In this paper, we propose an automatic kernel generation framework for large language models on deep learning accelerators. The key idea is the reinforcement learning (RL) formulation to generate a kernel with multi-step decision-making. We first develop a high-quality action space to satisfy architectural constraints of accelerator. Then, we provide a practical RL implementation by devising a policy network with Transformer and variance reduction techniques for gradients. Experimental results show our framework achieves average 3.5× speedup on TensorCore compared with exploration-based Ansor; 2.6× speedup on Simba compared with solver-based CoSA. Also, our framework achieves better energy efficiency compared to the state-of-the-art works.