Creating Custom Kernels for the AMD MI300

Hugging Face and AMD have developed a set of open-source optimized kernels for the AMD MI300X to improve the performance of serving Llama 3.1 405B in FP8 using VLLM. By implementing three specific custom kernels—a fused residual connection/RMS norm/FP8 conversion kernel, a fused SwiGLU activation/FP8 conversion kernel, and a Skinny GEMM kernel—the team achieved significant latency reductions during the decoding regime (measured with input size 1 and output size 128).

Custom Kernel Implementations and Performance Gains

Fused RMS Norm Kernel

The RMS norm kernel optimizes the beginning of decoder blocks by fusing the residual connection, row-wise Root Mean Square (RMS) normalization, and FP8 quantization into a single operation.

Technical Optimizations:

  • Vectorized Memory Access: The kernel uses 128-bit wide loads to fetch 8 FP16 elements per instruction, ensuring memory accesses are coalesced and contiguous to maximize warp efficiency.
  • Shared Memory (SMEM) Utilization: To avoid repeated VRAM (global memory) access, the modified version of the hidden state $x$ is stored in shared memory. For Llama 405B, the dimension $d=16384$ fits within the 64KB shared memory available per compute unit.
  • Block-Level Reduction: The kernel assigns one thread block per row and uses shared memory to synchronize threads for the summation required by the RMS norm.

Results: The "Vectorized + SMEM" implementation outperformed both standard PyTorch and VLLM's existing implementation, providing a significant speedup across various batch sizes.

Fused SwiGLU Kernel

The SwiGLU kernel fuses the activation function and the subsequent FP8 quantization for the MLP block's "Gate / Up" projection.

Technical Optimizations:

  • Packed Instructions: The kernel utilizes MI300X packed instructions for FP16 addition and multiplication to increase the volume of work per instruction.
  • Fast Math Approximation: To reduce latency, the kernel replaces the standard exp instruction with a faster exp2 instruction by scaling the input by $\log(2)$, resulting in a negligible loss in precision.
  • Packed FP32 to FP8 Conversion: Since the MI300X only supports FP8 conversion from FP32, the kernel leverages packed conversion instructions to boost performance.

Results: The custom SwiGLU kernel is more than 14 times faster than PyTorch on average and between 27% and 100% faster than the VLLM kernel.

Skinny GEMM Kernel

General Matrix Multiplication (GEMM) kernels in standard libraries are often inefficient for "skinny" matrices—those with very few rows (typical during decoding with low batch sizes)—because they result in low GPU utilization due to limited tiling opportunities.

Technical Optimizations:

  • Split-K Algorithm: The kernel splits the GEMM along the shared K axis into several sub-GEMMs executed concurrently. This increases the number of active compute units (CUs) by distributing the workload, reducing the time each CU spends on the K axis.
  • Sparsity Trick for Padding Removal: To avoid wasteful padding when the number of rows is less than the minimum dense tensor core instruction size (e.g., 16), the kernel uses a 4:2 structured sparsity instruction. A dense 8-row matrix is mapped into a 16-row sparse matrix, allowing the use of the 16x16x64 sparse instruction, which has twice the depth of the smallest dense instruction.
  • Warp Specialization and Asynchronous Execution: To address low arithmetic intensity, the kernel separates warps into "producers" (dedicated to loading data from VRAM to shared memory) and "consumers" (dedicated to computation). A queue in shared memory coordinates these warps asynchronously, ensuring consumers are not idle while waiting for slow VRAM loads.

Results: The Skinny GEMM kernel shows notable speedups over PyTorch for low row counts (M = 1, 8, 16), particularly for QKV and Gate/Up projections, though gains diminish as the batch size increases to 32.

Implementation and Availability

All developed kernels are available in the hf-rocm-kernels GitHub repository, which includes source code, Python bindings, benchmarking scripts, and a test suite. The kernels are designed to be used independently or integrated into VLLM. To reproduce the results, Hugging Face recommends using the specific ROCm 6.3.1 container used during development.

Sources