Creating custom kernels for the AMD MI300

AMD Kernels

Introduction

More than a billion per day: that’s a low estimate of how many requests ChatGPT handles daily, a number which is unlikely to go down soon. For each request and each generated token, we run an inference of a multi-billion parameters model. This is why model optimization is paramount at each and every level: when one deals with these kinds of scale, even a 1% latency or power gain can bring huge savings.

But where might that gain come from? Model architectures are already well established, and popular models have had quantized weight for a long time now. However, a crucial level at which we can optimize model inference remains: the kernel level. Kernels are the algorithms executed when you do any operation in your network: there are matrix multiplication kernels, convolution kernels, batch normalization kernels, etc. Kernels are low-level, highly-optimized algorithms, often tailored for the device they will be running on. They are notoriously long and hard to write, and require a good understanding of the inner working of the GPU.

Kernels are essential for running operations in neural networks—without a kernel, an operation effectively can't be used. Because of this, new innovations often launch with a "day 0" kernel, typically optimized only for the latest Nvidia hardware. This approach excludes many other devices, particularly AMD GPUs, which, despite offering comparable or superior specs, are often overlooked by kernel developers. Hugging Face collaborated with AMD to deliver state-of-the-art performance on AMD platforms and make it benefit the open source community. As part of this partnership, we decided with AMD to focus on delivering open-source optimized kernels to improve the performance of serving Llama 3.1 405B in FP8 on a node of 8 MI300X using VLLM.

• Fused residual connection, RMS norm and FP8 conversion kernel
• Fused SwiGLU activation and FP8 conversion kernel
• Skinny GEMM kernel