Hugging Face and AMD collaborated to develop open‑source optimized kernels for the AMD MI300X GPU, targeting improved performance of Llama 3.1 405B in FP8 using VLLM. The partnership produced three custom kernels that delivered significant speedups on a node of eight MI300X GPUs, with results reproducible via the hf‑rocm‑kernels repository. The kernels are designed to replace default Nvidia‑centric implementations, enabling broader hardware support for large‑scale inference workloads.
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.
We are first going to do a quick refresher on the architecture of the device we are working on: the MI300X. Then, we are going to take a look at the state of our model’s inference before optimizing it. This will allow us to identify bottlenecks and know which custom kernels we need to write. Then, we will take a look at each kernel we have written, which will give us an opportunity to explore how kernel optimization is conducted through many angles.
In this blog post, we'll explore how we optimized performance for the MI300X and how each kernel was individually fine-tuned. But first, let’s look at the performance gains achieved using our custom kernels. By combining the following three optimized kernels:
we achieved significant speedups when running VLLM on a node powered by MI300X GPUs.
Measures were taken with input size 1 and output size 128 to mimic decoding regime. We measure decoding latency using the median over 30 iterations.
Optimizing kernel code for AMD GPUs expands high‑performance inference options beyond Nvidia, potentially reducing costs and improving efficiency for large‑scale AI deployments.