Custom CUDA kernels juice Llama 3.1 405B on AMD MI300X, skipping Nvidia lock-in

A billion-plus requests a day. That’s the ballpark for ChatGPT, and it underscores why even a 1% latency shave in large language model inference translates to serious money. Most optimization efforts stop at model architecture or quantization. Hugging Face and AMD just went deeper, attacking the kernel level for the AMD MI300X GPU and wrestling Llama 3.1 405B into FP8 with some custom code that leaves generic solutions in the dust.

For a node packing 8 MI300X GPUs, the team wrote three kernels to fix specific bottlenecks they spotted during decoding with vLLM. The first fuses a residual connection, RMS normalization, and an FP8 type conversion into a single pass. The second tackles the SwiGLU activation function, again merging it with an FP8 cast. The third is a ‘skinny GEMM’ kernel, a matrix multiply tuned for the weirdly shaped matrices that pop up when batch sizes are small. These aren’t generic optimizations; they’re surgical strikes on the specific geometry of serving a massive model.

Using the median of 30 iterations with an input size of 1 and output size of 128, the benchmarks show a clear speedup in decoding latency. The exact numbers are in the post’s charts, but the takeaway is that AMD hardware, often stuck with ‘day-0’ kernels built only for Nvidia, can deliver competitive performance when given the same low-level love. This matters less for training, where raw teraflops dominate, and more in the brutal economics of production serving, where every microsecond of latency compounds across billions of tokens.

All the code is open-source and sitting in the hf-rocm-kernels repository on GitHub. It comes with Python bindings, benchmarking scripts, and a test suite, so any team can reproduce the results or use it as scaffolding for their own kernel experiments. The kernels are also slated for integration into AMD’s fork of vLLM. What makes this collaboration notable isn’t just the speedup—it’s a blueprint for how the open-source ecosystem can force a second hardware player into the conversation when the default is often a single green logo.