The AI CUDA Engineer: Agentic CUDA Kernel Discovery, Optimization and Composition

Note: Updated on February 21, 2025.

At Sakana AI, we believe the path to develop much stronger AI systems is to automate the development of AI using AI. We aim to develop AI systems that can create even more capable and efficient AI systems.

In the past year, we introduced an AI system that can automate the creation of new AI foundation models, at a fraction of the cost. We showed that LLMs can invent more efficient methods to train LLMs. Recently, we proposed the first comprehensive agentic framework for fully automating the entire AI research process in The AI Scientist. This led us to the question: If AI can be used to conduct AI research, can we use AI to research ways to make AI run faster?

Introduction

Just like the human brain, modern AI systems also rely heavily on parallel processing, enabled by hardware accelerators such as GPUs. But unlike the human brain which is evolved (biologically and culturally) to operate efficiently under resource constraints, recent advances in AI foundation models have led to large-scale deployment and ever-growing inference time and energy demand, leading to exponentially increasing resource requirements to train and deploy AI models.

We believe that fundamentally, modern AI systems can and should be as efficient as the human brain, and that the best path to achieve this efficiency is to use AI to make AI more efficient! Inspired by our earlier work on The AI Scientist, we are proud to announce The AI CUDA Engineer, the first comprehensive agentic framework for fully automatic CUDA kernel discovery and optimization.

CUDA is a low-level software layer that gives direct access to the NVIDIA GPU’s hardware instruction set for parallel computation. CUDA kernels are functions written in the CUDA language that run on GPUs. By writing instructions directly at the CUDA kernel level, we can achieve much higher performance for AI algorithms. However, working with CUDA requires quite a bit of GPU knowledge, and in practice, most machine learning algorithms are written in a higher level abstraction layer such as PyTorch or JAX.

The AI CUDA Engineer is an agentic framework that leverages frontier LLMs with the goal of automating the conversion of standard PyTorch code into highly optimized CUDA kernels. Through the use of evolutionary optimization, and leveraging concepts in evolutionary computation, such as ‘crossover’ operations and ‘innovation archive’ to discover promising ‘stepping stone’ kernels, our proposed framework is able to not only automate the process of converting PyTorch modules to CUDA kernels, but our highly optimized CUDA kernels often achieve speedups that have significantly faster runtime.

We believe this technology can enable speedups that will accelerate both the training and running (inference) of foundation models like LLMs or other generative AI models, eventually making AI models run much faster on NVIDIA hardware.

The AI CUDA Engineer is able to generate CUDA Kernels with speedups of 10—100x over common PyTorch operations. Our framework is also able to produce highly optimized CUDA Kernels that are much faster than existing CUDA Kernels that are already commonly used in production (up to 5x speedups).

Stage 1 and 2 (Conversion and Translation):  The AI CUDA Engineer first translates PyTorch code into functioning CUDA kernels. We already observe initial runtime improvements without explicitly targeting these.

Stage 3 (Evolutionary Optimization):  Inspired by biological evolution, our framework utilizes evolutionary optimization (‘survival of the fittest’) to ensure only the best CUDA kernels are produced. Furthermore, we introduce a novel kernel crossover prompting strategy to combine multiple optimized kernels in a complementary fashion.

Stage 4 (Innovation Archive):  Just as how cultural evolution shaped our human intelligence with knowhow from our ancestors through millennia of civilization, The AI CUDA Engineer also takes advantage of what it learned from past innovations and discoveries it made (Stage 4), building an Innovation Archive from the ancestry of known high-performing CUDA Kernels, which uses previous stepping stones to achieve further translation and performance gains.

Kernel Runtime Speedups Discovered by the AI CUDA Engineer

The AI CUDA Engineer robustly discovered CUDA kernels used for common machine learning operations, with speedups as high as 10—100x faster than native and compiled kernels in PyTorch. Our approach is also able to convert entire machine learning architectures into optimized CUDA kernels. Here we highlight a couple of significant speedup discoveries made completely autonomously:

Our approach finds more efficient CUDA kernels for fundamental operations such as matrix multiplications to common deep learning operations, and as of writing, the performance of our discovered CUDA kernels achieved state-of-the-art performance on KernelBench.

Technical Report and Dataset Release

We believe that this is just the beginning of the great optimization of AI!

We’re excited to release our new paper, The AI CUDA Engineer: Agentic CUDA Kernel Discovery and Optimization.

In our report:

We introduce an end-to-end agentic workflow capable of translating PyTorch code to working CUDA kernels, optimizing CUDA runtime performance, and automatically fusing multiple kernels.

Furthermore, we construct various techniques for enhancing the consistency and performance of the pipeline including LLM ensembling, an iterative profiling feedback loop, local kernel code-editing, and crossover kernel optimization.

We show that The AI CUDA Engineer robustly translates more than 230 out of 250 considered torch operations and achieves strong runtime performance improvements for the majority of kernels. Furthermore, our approach is capable of efficiently fusing various kernel operations and can outperform several existing accelerated operations.

We release a dataset of over 17,000 verified kernels for operations covering a wide range of PyTorch operations.

We highlight some notable examples of discovered CUDA kernels that achieved significant speedups on key computation operations in AI models.

Highlighted AI CUDA Engineer-Discovered Kernels

Leveraging our novel LLM-driven evolutionary kernel optimization procedure we robustly obtain speedups for a diverse range of considerations. More specifically, we outperform PyTorch Native runtimes for 81% out of 229 considered tasks. Furthermore, 20% of all discovered CUDA kernels are at least twice as fast as their PyTorch implementations.

Below we show a subset of kernels. They highlight the diversity of different operations for which the AI CUDA Engineer can successfully be deployed. This includes normalization methods, loss functions, special matrix multiplications and even entire neural network architectures:

The AI CUDA Engineer Archive: A Dataset of 17,000+ Verified CUDA Kernels

A Text Embedding visualization of the AI CUDA Engineer Archive shows that the discovered kernels group into tasks (e.g. MatMul, Pooling, Convolution) and implementation strategies (unrolling, fusing, vectorization). The Archive is openly accessible and can be used for downstream fine-tuning of LLMs.

Along with this paper, we release The AI CUDA Engineer archive, a dataset consisting of more than 30,000 CUDA kernels generated by The AI CUDA Engineer. It is released under the CC-By-4.0 license and is accessible via HuggingFace. The dataset includes a torch reference implementation, torch, NCU and Clang-tidy profiling data, multiple kernels per task, error messages and speedup scores against torch native and compile runtimes.

Summary statistics of the AI CUDA Engineer Archive consisting of more than 30,000 kernels and more than 17,000 correct verified implementations. Approximately, 50% of all kernels improve over the torch native runtime.

We envision that this dataset can enable post-training of open-source models to perform better CUDA-enabling modules. This includes offline Reinforcement Learning, preference optimization, and standard supervised fine-tuning.

Explore 17,000+ Kernels in The AI CUDA Engineer Archive

We also published an interactive website for interactively inspecting more than 17,000 verified kernels and their profiles including torch, NCU and Clang-Tidy data. You can access our interactive website here.

The website allows you to explore various high-performing kernels across 230 tasks. It comes with a custom leaderboard that can be used to inspect related kernels across experiments and LLMs.

Furthermore, you can visualize the kernel, retrieve related kernels, download code to verify the implementation and speedup as well as view the obtained profiling data. Finally, you can take an in-depth look at the optimization experiment.

Detailed view of an optimized kernel including profiling data, downloading of evaluation scripts, related kernels and discovery experiment details.

Limitations and Bloopers

Combining evolutionary optimization with LLMs is powerful but can also find ways to trick the verification sandbox. We are fortunate to have readers, like @main_horse test our CUDA kernels, to identify that the system had found a way to “cheat”. For example, the system had found a memory exploit in the evaluation code which, in a number of cases, allowed it to avoid checking for correctness.

We have since made the evaluation harness more robust to eliminate this loophole and have updated our results.

Furthermore, we find the system could also find other novel exploits in the benchmark’s tasks. We are in the process of revising our paper and updating results, with further imporvements to the evaluation and runtime profiling harness, to reflect and discuss the effects, and mitigation of LLM reward hacking for CUDA kernel optimization.

In addition, we observed limitations in frontier LLMs’ ability to effectively utilize TensorCore WMMA capabilities. While LLMs could generate basic CUDA code, they often struggled to implement the specialized matrix multiplication acceleration features offered by modern GPU architectures. This suggests a potential gap in the training data or the models’ understanding of advanced hardware-specific optimizations.

As frontier LLMs, especially those with advanced coding reasoning capabilities become more capable, we expect code-optimization systems, such as ours, will continue to face these challenges. We envision a future where it is the role of human engineers to work with code optimization systems as tools, to produce the best and most reliable results.

Future Implications of The AI CUDA Engineer

The AI revolution is just getting started, and we are just at the very beginning of the transformation cycle. It is our view that today’s LLMs are our generation’s “Mainframe Computers”. We are still in the very early stages of AI, and it is inevitable, due to market competition and global innovation (especially from those innovating with resource constraints), that this technology will become a million times more efficient.

Currently, our AI systems consume immense resources, and if the technology continues to scale without thought for efficiency and energy consumption, the result will not be sustainable. There is no fundamental reason why our AI systems can’t be as efficient (or even more efficient) than human intelligence. We believe that the best path to achieve this greater efficiency is to use AI to make AI more efficient.

This is the direction that Sakana AI is pursuing, and this project is an important step towards making AI a million times faster. Just like the evolution of early clunky mainframe computers to modern computing, how we use AI today will look very different in a few years, compared to today’s ‘clunky’, inefficient LLMs.

Sakana AI

Want to make the AI that improves AI? Please see our Careers page for more information.


ROCm Documentation

Install

How to

Conceptual

Reference

Contribute

Optimizing with Composable Kernel

Contents

Optimizing with Composable Kernel#

2025-01-27

19 min read time

The AMD ROCm Composable Kernel (CK) library provides a programming model for writing performance-critical kernels for machine learning workloads. It generates a general-purpose kernel during the compilation phase through a C++ template, enabling developers to achieve operation fusions on different data precisions.

This article gives a high-level overview of CK General Matrix Multiplication (GEMM) kernel based on the design example of 03_gemm_bias_relu. It also outlines the steps to construct the kernel and run it. Moreover, the article provides a detailed implementation of running SmoothQuant quantized INT8 models on AMD Instinct MI300X accelerators using CK.

High-level overview: a CK GEMM instance#

GEMM is a fundamental block in linear algebra, machine learning, and deep neural networks. It is defined as the operation:
\(E = α \times (A \times B) + β \times (D)\), with A and B as matrix inputs, α and β as scalar inputs, and D as a pre-existing matrix.
Take the commonly used linear transformation in a fully connected layer as an example. These terms correspond to input activation (A), weight (B), bias (D), and output (E), respectively. The example employs a DeviceGemmMultipleD_Xdl_CShuffle struct from CK library as the fundamental instance to explore the compute capability of AMD Instinct accelerators for the computation of GEMM. The implementation of the instance contains two phases:

Template parameter definition

Instantiating and running the templated kernel

Template parameter definition#

The template parameters of the instance are grouped into four parameter types:

Parameters for determining matrix data precision

Parameters for determining matrix data layout

Parameters for determining extra operations on matrix elements

Performance-oriented tunable parameters

The template parameters of the selected GEMM kernel are classified into four groups. These template parameter groups should be defined properly before running the instance.#

Matrix data precision#

A, B, D, and E are defined as half-precision floating-point datatypes. The multiply-add results of matrix A and B are added with a pre-existing matrix D (half-precision), and the final GEMM results are also half-precision floating-points.

using ADataType        = F16;
using BDataType        = F16;
using AccDataType      = F32;
using CShuffleDataType = F16;
using DDataType        = F16;
using EDataType        = F16;

ADataType and BDataType denote the data precision of the A and B input matrices. AccDataType determines the data precision used for representing the multiply-add results of A and B elements. These results are stored in a CShuffle module in local data share (LDS), a low-latency and high-bandwidth explicitly-addressed memory used for synchronization within a workgroup LDS for later use.

CShuffleDataType denotes the data precision of CShuffle in LDS.

DDataType denotes the data precision of the pre-existing D matrix stored in GPU global memory, while EDatatype denotes the data precision of the final output. The CK kernel supports a fusion strategy so that CShuffle can be added with a single pre-existing matrix in the same GPU kernel for better performance.

Matrix data layout#

using ALayout = Row;
using BLayout = Col;
using DLayout = Row;
using ELayout = Row;

Following the convention of various linear algebra libraries, CK assumes that the input matrix A is an M x K matrix, meaning the matrix has M rows and K columns. Similarly, matrix B is assumed to be K x N, meaning it has K rows and N columns. In computing, row-major order and column-major order are commonly used ways to store matrices in linear storage. After understanding the matrix storage pattern, the underlying optimized memory access manner can be applied to achieve better performance depending on the storage ordering of these matrices.

Matrix element operation#

using AElementOp   = PassThrough;
using BElementOp   = PassThrough;
using CDEElementOp = AddRelu;

CK supports the pre-processing of the matrix before calculating GEMM, that is, C = AElementOp(A) * BElementOp(B). It similarly supports the post-processing of GEMM results the same way, that is, E = CDEElementOp(C, D).

AElementOp and BElementOp determine the operation applied to matrix A and B separately before GEMM, which is achieved by binding the operation with a C++ struct function.

The above PassThrough denotes no operations are performed on the target matrix. CDEELementOp determines the operations applied to CShuffle output and matrix D. The following binding struct AddRelu shows an example of adding the CShuffle output and matrix D, and ReLU (Rectified Linear Unit) operations to the addition result. It then passes the results to matrix E.

struct AddRelu
{
    __host__ __device__ void operator()(ck::half_t& e, const ck::half_t& c, const ck::half_t& d) const
    {
        const ck::half_t x = c + d;
        e = x > 0 ? x : 0;
    }
};

Tunable parameters#

The CK instance includes a series of tunable template parameters to control the parallel granularity of the workload to achieve load balancing on different hardware platforms.

These parameters include Block Size, M/N/K Per Block, M/N per XDL, AK1, BK1, etc.

Block Size determines the number of threads in the thread block.

M/N/K Per Block determines the size of tile that each thread block is responsible for calculating.

M/N Per XDL refers to M/N size for Instinct accelerator Matrix Fused Multiply Add (MFMA) instructions operating on a per-wavefront basis.

A/B K1 is related to the data type. It can be any value ranging from 1 to K Per Block. To achieve the optimal load/store performance, 128bit per load is suggested. In addition, the A/B loading parameters must be changed accordingly to match the A/B K1 value; otherwise, it will result in compilation errors.

Conditions for achieving computational load balancing on different hardware platforms can vary.

Instantiating and running the templated kernel#

After determining the template parameters, we instantiate the kernel with actual arguments. Do one of the following:

Use GetDeviceBuffer from CK’s custom struct DeviceMem to pass the element values of the matrices that need to be calculated.

Allocate device buffer via hipMalloc. Ensure the device buffer size can fit the matrix size.

Pass matrix elements through the data_ptr method in the Tensor object if the matrix to be calculated is of Tensor type.

The row and column, and stride information of input matrices are also passed to the instance. For batched GEMM, you must pass in additional batch count and batch stride values. The extra operations for pre and post-processing are also passed with an actual argument; for example, α and β for GEMM scaling operations. Afterward, the instantiated kernel is launched by the invoker, as illustrated in Figure 3.

Templated kernel launching consists of kernel instantiation, making arguments by passing in actual application parameters, creating an invoker, and running the instance through the invoker.#

Developing fused INT8 kernels for SmoothQuant models#

SmoothQuant (SQ) is a quantization algorithm that enables an INT8 quantization of both weights and activations for all the matrix multiplications in LLM. The required GPU kernel functionalities used to accelerate the inference of SQ models on Instinct accelerators are shown in the following table.

Functionality descriptions

Corresponding wrappers

\(E = α \times (A \times B) + β \times (D)\), where A, B, D, E are INT8 2-D tensors;

E = Linear_ABDE_I8(A, B, D, \(\alpha\), \(\beta\))

\(E = RELU (α \times (A \times B) + β \times (D))\), where A, B, D, E are INT8 2-D tensors;

E = Linear_ReLU_ABDE_I8(A, B, D, \(\alpha\), \(\beta\))

\(E = α \times (A \times B) + β \times (D)\), where A, B are INT8 2-D tensors, D and E are FP32 2-D tensors;

E = Linear_AB_I8_DE_F32(A, B, D, \(\alpha\), \(\beta\))

\(E = α \times (A \times B)\), where A, B, E are INT8 3-D tensors;

E = BMM_ABE_I8(A, B, \(\alpha\))

\(E = α \times (A \times B)\), where A, B are INT8 3-D tensors, E is FP32 3-D tensor;

E = BMM_AB_I8_E_F32(A, B, \(\alpha\))

Operation flow analysis#

The following section discusses the analysis of the operation flow of Linear_ReLU_ABDE_I8. The rest of the wrappers in Table 1 can be analyzed similarly.

The first operation in the process is to perform the multiplication of input matrices A and B. The resulting matrix C is then scaled with α to obtain T1. At the same time, the process performs a scaling operation on D elements to obtain T2. Afterward, the process performs matrix addition between T1 and T2, element activation calculation using ReLU, and element rounding sequentially. The operations to generate E1, E2, and E are encapsulated and completed by a user-defined template function in CK (given in the next sub-section). This template function is integrated into the fundamental instance directly during the compilation phase so that all these steps can be fused in a single GPU kernel.

Operation flow.#

The CK library contains many fundamental instances that implement different functions. Familiarize yourself with the names of various CK instances and determine whether they meet the target functional requirements.

Second, consider whether the format of input data meets your actual calculation needs. For SQ models, the 8-bit integer data format (INT8) is applied for matrix calculations.

Third, consider the platform for implementing CK instances. The instances suffixed with xdl only run on AMD Instinct accelerators after being compiled and cannot run on Radeon-series GPUs. This is due to the underlying device-specific instruction sets for implementing these basic instances.

Here, we use DeviceBatchedGemmMultiD_Xdl as the fundamental instance to implement the functionalities in the previous table.

Use the ‘DeviceBatchedGemmMultiD_Xdl’ instance as a root.#

The DeviceBatchedGemmMultiD_Xdl instance realizes the batched GEMM BMM_ABE_I8 and BMM_AB_I8_E_F32 kernels directly by using the proper input and output data precision types.

Based on the two batched GEMM kernels, GEMM kernel Linear_ABDE_I8 and Linear_AB_I8_DE_F32 can be implemented by expanding their input 2-D tensors to 3-D tensors. Then, the 3-D output tensors produced by the root instance are squeezed back to 2-D output tensors before returning back.

For example, unsqueeze A (M, K) to A (1, M, K) before assigning it into the root instance and squeeze E (1, M, N) to (M, N) after the calculations of the root instance return back. Linear_ReLU_ABDE_I8 is implemented by adding a ReLU operation on the result output of Linear_ABDE_I8.

Developing the complete function#

The inference of SQ quantized models relies on using PyTorch and Transformer libraries, and a tensor type is used to represent matrices and vectors in torch, the C++ data types in CK need to be replaced with the torch::tensor type. The data types of the input and output matrices should be a tensor type.

In GEMM, the A and B inputs are two-dimensional matrices, and the required input matrices of the selected fundamental CK instance are three-dimensional matrices. Therefore, we must convert the input 2-D tensors to 3-D tensors, by using tensor’s unsqueeze() method before passing these matrices to the instance. For batched GEMM in the preceding table, ignore this step.

// Function input and output
torch::Tensor linear_relu_abde_i8(
    torch::Tensor A_,
    torch::Tensor B_,
    torch::Tensor D_,
    float alpha,
    float beta)
{
  // Convert torch::Tensor A_ (M, K) to torch::Tensor A (1, M, K)
  auto A = A_.unsqueeze(0);

  // Convert torch::Tensor B_ (K, N) to torch::Tensor A (1, K, N)
  auto B = B_.unsqueeze(0);
...

As shown in the following code block, we obtain M, N, and K values using input tensor size values. This stride size information is used to reshape the input vector D and allocate the storage space of tensor E. Stride reflects the exact size of continuous elements in memory, which are passed as important parameters to the fundamental instance for GPU kernel use.

// Return the batch count from the size of dimension 0
  int batch_count = A.size(0);

  // Return the M, N, K from the size of dimension 1 & 2
  int M = A.size(1);
  int N = B.size(1);
  int K = A.size(2);

  // Initialize the stride size for A, B, D and E
  int stride_A = K;
  int stride_B = K;
  int stride_D0 = N;
  int stride_E = N;

  // Initialize the stride size for batched A, B, D and E
  long long int batch_stride_A = M * K;
  long long int batch_stride_B = K * N;
  long long int batch_stride_D0 = M * N;
  long long int batch_stride_E = M * N;

  // Convert the tensor of 2-D to 3-D
  auto D = D_.view({1,-1}).repeat({M, 1});

  // Allocate memory for E
  auto E = torch::empty({batch_count, M, N},
       torch::dtype(torch::kInt8).device(A.device()));

In the following code block, ADataType, BDataType and D0DataType are used to denote the data precision of the input tensors A, B and D, respectively. EDataType is used to denote the data precision of output tensor E. These parameters are specified to I8 data format (8-bit integer data format) to meet the kernel’s design requirements.

AccDataType determines the data precision used to represent the multiply-add results of A and B elements. Generally, a larger range data type is applied to store the multiply-add results of A and B to avoid result overflow; I32 is applied in this case. The CShuffleDataType I32 data type indicates that the multiply-add results continue to be stored in LDS as an I32 data format. All of this is implemented through the following code block.

// Data precision
  using ADataType        = I8;
  using BDataType        = I8;
  using AccDataType      = I32;
  using CShuffleDataType = I32;
  using D0DataType       = I8;
  using DsDataType       = ck::Tuple<D0DataType>;
  using EDataType        = I8;

Following the convention of various linear algebra libraries, row-major and column-major orders are used to denote the ways of storing matrices in linear storage. The advantage of specifying matrix B as column major is that all the relevant matrix elements are stored continuously in GPU global memory when a row in A is multiplied by a column in B, which can help GPU achieve data consistency access to improve access performance.

// Specify tensor order
  using ALayout  = RowMajor;
  using BLayout  = ColumnMajor;
  using D0Layout = RowMajor;
  using DsLayout = ck::Tuple<D0Layout>;
  using ELayout  = RowMajor;

In CK, PassThrough is a struct denoting if an operation is applied to the tensor it binds to. To fuse the operations between E1, E2, and E introduced in section Operation flow analysis, we define a custom C++ struct, ScaleScaleAddRelu, and bind it to CDEELementOp. It determines the operations that will be applied to CShuffle (A×B results), tensor D, α, and β.

// No operations bound to the elements of A and B
  using AElementOp   = PassThrough;
  using BElementOp   = PassThrough;

  // Operations bound to the elements of C, D and E
  using CDEElementOp = ScaleScaleAddRelu;

In the binding struct, operator() performs an addition operation between CShuffle and matrix D, a ReLU operation on the addition results, and a rounding operation on the output elements. It then returns the results to E.

struct ScaleScaleAddRelu {

  template <>
  __host__ __device__ constexpr void
  operator()<I8, I32, I8>(I8& e, const I32& c, const I8& d) const
  {
      // Scale AxB result with alpha
      const F32 c_scale = ck::type_convert<F32>(c) * alpha;

      // Scale D with beta
      const F32 d_scale = ck::type_convert<F32>(d) * beta;

      // Perform addition operation
      F32 temp = c_scale + d_scale;

      // Perform RELU operation
      temp = temp > 0 ? temp : 0;

      // Perform rounding operation
      temp = temp > 127 ? 127 : temp;

      // Return to E
      e = ck::type_convert<I8>(temp);
  }

  F32 alpha;
  F32 beta;
};

The original input tensors need to be padded to meet GPU tile-based parallelism.

static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::MNKPadding;

The template parameters of the target fundamental instance are initialized with the above parameters and includes default tunable parameters. For specific tuning methods, see Tunable parameters.

using DeviceOpInstance = ck::tensor_operation::device::DeviceBatchedGemmMultiD_Xdl<
    // Tensor layout
    ALayout, BLayout, DsLayout, ELayout,
    // Tensor data type
    ADataType, BDataType, AccDataType, CShuffleDataType, DsDataType, EDataType,
    // Tensor operation
    AElementOp,  BElementOp, CDEElementOp,
    // Padding strategy
    GemmDefault,
    // Tunable parameters
    tunable parameters>;

Return the address of the first element of tensors:

auto A_ref = A.data_ptr<ADataType>();
 auto B_ref = B.data_ptr<BDataType>();
 auto D0_ref = D.data_ptr<D0DataType>();
 auto E_ref = E.data_ptr<EDataType>();

The fundamental instance is then initialized and run with actual arguments:

auto device_op    = DeviceOpInstance{};
 auto invoker = device_op.MakeInvoker();
 auto argument = device_op.MakeArgument(
    A_ref, B_ref, {D0_ref}, E_ref,
    M, N, K,
    batch_count,
    stride_A, stride_B, {stride_D0}, stride_E,
    batch_stride_A, batch_stride_B, {batch_stride_D0}, batch_stride_E,
    AElementOp{}, BElementOp{}, CDEElementOp{alpha, beta});

invoker.Run(argument, StreamConfig{nullptr, 0});

The output of the fundamental instance is a calculated batched matrix E (batch, M, N). Before the return, it needs to be converted to a 2-D matrix if a normal GEMM result is required.

// Convert (1, M, N) to (M, N)
return E.squeeze(0);

Binding to Python#

Since these functions are written in C++ and torch::Tensor, you can use pybind11 to bind the functions and import them as Python modules. For the example, the necessary binding code for exposing the functions in the table spans but a few lines.

#include <torch/extension.h>

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
  m.def("linear_ab_i8_de_f32", &linear_ab_i8_de_f32);
  m.def("linear_relu_abde_i8", &linear_relu_abde_i8);
  m.def("linear_abde_i8", &linear_abde_i8);
  m.def("bmm_abe_i8", &bmm_abe_i8);
  m.def("bmm_ab_i8_e_f32", &bmm_ab_i8_e_f32);
}

Build the C++ extension by writing a setup.py script that uses setuptools to compile the C++ code. A reference implementation of the setup.py script is as follows.

import os
from setuptools import setup, find_packages
from torch.utils import cpp_extension
from torch.utils.cpp_extension import BuildExtension

os.environ["CC"] = "hipcc"
os.environ["CXX"] = "hipcc"

sources = [
    'torch_int/kernels/linear.cpp',
    'torch_int/kernels/bmm.cpp',
    'torch_int/kernels/pybind.cpp',
]

include_dirs = ['torch_int/kernels/include']
extra_link_args = ['libutility.a']
extra_compile_args = ['-O3','-DNDEBUG', '-std=c++17', '--offload-arch=gfx942', '-DCK_ENABLE_INT8', '-D__HIP_PLATFORM_AMD__=1']

setup(
    name='torch_int',
    ext_modules=[
        cpp_extension.CUDAExtension(
            name='torch_int.rocm',
            sources=sources,
            include_dirs=include_dirs,
            extra_link_args=extra_link_args,
            extra_compile_args=extra_compile_args
            ),
    ],
    cmdclass={
        'build_ext': BuildExtension.with_options(use_ninja=False)
    },
    packages=find_packages(
        exclude=['notebook', 'scripts', 'tests']),
)

Run python setup.py install to build and install the extension. It should look something like Figure 6:

Compilation and installation of the INT8 kernels.#

INT8 model inference and performance#

The implementation architecture of running SmoothQuant models on MI300X GPUs is illustrated in Figure 7, where (a) shows the decoder layer composition components of the target model, (b) shows the major implementation class for the decoder layer components, and (c) denotes the underlying GPU kernels implemented by CK instance.

The implementation architecture of running SmoothQuant models on AMD MI300X accelerators.#

For the target SQ quantized model, each decoder layer contains three major components: attention calculation, layer normalization, and linear transformation in fully connected layers.  The corresponding implementation classes for these components are:

Int8OPTAttention

W8A8B8O8LinearReLU

W8A8BF32OF32Linear

These classes’ underlying implementation logits will harness the functions in previous table. Note that for the example, the LayerNormQ module is implemented by the torch native module.

Testing environment:
The hardware platform used for testing equips with 256 AMD EPYC 9534 64-Core Processor, 8 AMD Instinct MI300X accelerators and 1.5T memory. The testing was done in a publicly available Docker image from Docker Hub:
rocm/pytorch:rocm6.1_ubuntu22.04_py3.10_pytorch_2.1.2

The tested models are OPT-1.3B, 2.7B, 6.7B and 13B FP16 models and the corresponding SmoothQuant INT8 OPT models were obtained from Hugging Face.

Note that since the default values were used for the tunable parameters of the fundamental instance, the performance of the INT8 kernel is suboptimal.

Figure 8 shows the performance comparisons between the original FP16 and the SmoothQuant-quantized INT8 models on a single MI300X accelerator. The GPU memory footprints of SmoothQuant-quantized models are significantly reduced. It also indicates the per-sample inference latency is significantly reduced for all SmoothQuant-quantized OPT models (illustrated in (b)). Notably, the performance of the CK instance-based INT8 kernel steadily improves with an increase in model size.

Performance comparisons between the original FP16 and the SmoothQuant-quantized INT8 models on a single MI300X accelerator.#

For accuracy comparisons between the original FP16 and INT8 models, the evaluation is done by using the first 1,000 samples from the LAMBADA dataset’s validation set. We employ the same Last Token Prediction Accuracy method introduced in SmoothQuant Real-INT8 Inference for PyTorch as our evaluation metric. The comparison results are shown in Table 2.

Models

Hugging Face FP16 model accuracy

SmoothQuant quantized INT8 model accuracy

opt-1.3B

0.72

0.70

opt-2.7B

0.76

0.75

opt-6.7B

0.80

0.79

opt-13B

0.79

0.77

Conclusion#

CK provides a rich set of template parameters for generating flexible accelerated computing kernels for difference application scenarios.

CK supports multiple instruction sets of AMD Instinct GPUs, operator fusion and different data precisions. Its composability helps users quickly construct operator performance verification.

With CK, you can build more effective AI applications with higher flexibility and better performance on different AMD accelerator platforms.

previous

Model acceleration libraries

next

Optimizing Triton kernels


The world’s leading publication for data science, AI, and ML professionals.

Unleashing the Power of Triton: Mastering GPU Kernel Optimization in Python

Accelerating AI/ML Model Training with Custom Operators – Part 2

According to Greek mythology, Triton, a god of the sea, would calm or stir the sea waters by using his conch shell to control its tides and waves. In one story, in particular, Triton is depicted as having used his powers to guide the Argonauts through particularly dangerous sea waters. In this post, we similarly call upon Triton for navigation through complex journeys, although this time we refer to the Triton language and compiler for writing Deep Learning (DL) kernels and to our journeys through the world of AI/ML development.

This is a sequel to a previous post on the topic of accelerating AI/ML applications with custom operators in which we demonstrated the potential for performance optimization by developing custom CUDA kernels. One of our intentions was to emphasize the accessibility of custom kernel development and the opportunities it provides even for non-expert CUDA developers. However, there are challenges to CUDA development that may prove insurmountable for some. For one, while many a modern-day AI/ML developer are well-versed in Python, they may not feel comfortable developing in C++. Furthermore, tuning a CUDA kernel to take full advantage of the GPU’s capabilities requires an intimate understanding of the underlying HW architecture and could take a non-trivial amount of work. This is particularly true if you want your kernel to run optimally on a variety of GPU architectures. Much of the complexity results from CUDA’s "thread-based" development model in which the developer is responsible for designing and optimizing all elements of the GPU kernel threads, including all details related to the use of GPU memory, thread-concurrency, TensorCore scheduling, and much more.

The Power of Triton

The Triton library aims to democratize and simplify Gpu kernel development in two primary ways. First, it provides an API for building custom operators in Python (rather than C++). Second, it enables kernel development at the block level (rather than the thread level) thereby abstracting away and automating all issues related to optimizing performance within CUDA thread blocks. Rather than taking the laborious steps of programming the details of the thread invocation, including the intricacies related to memory management, scheduling of on-chip acceleration engines, thread-synchronization, etc., kernel developers can rely on Triton to do it all for them. One important byproduct of the high-level API abstraction of Triton’s programming model is that it reduces the burden of needing to tune the kernel for multiple different GPU types and architectures.

Of course, as is usually the case when up-leveling an API, the Triton programming model does have its disadvantages. Some kernels might benefit from the thread-level control enabled by CUDA (e.g., they might benefit from the conditional execution flow discussed in our previous post). Other kernels might require very specialized and delicate treatment to reach peak performance and may suffer from the automated result of the Triton compiler. But even in cases such as these, where the development of a CUDA kernel may ultimately be required, the ability to quickly and easily create a temporary Triton kernel could greatly facilitate development and boost productivity.

For more on the motivations behind Triton and on the details of its programming model, see the Triton announcement, the official Triton documentation, and the original Triton white-paper.

Disclaimers

Similar to our [previous post](https://chaimrand.medium.com/accelerating-ai-ml-model-training-with-custom-operators-163ef2a04b12), our intention is to provide a simple demonstration of the opportunity offered by Triton. Please do not view this post as a replacement for the official Triton documentation or its associated tutorials. We will use the same face-detection model as in our previous post as a basis for our demonstration and perform our experiments in the same Google Cloud environment – a g2-standard-16 VM (with a single L4 GPU) with a dedicated deep learning VM image and Pytorch 2.4.0. As before, we make no effort to optimize our examples and/or verify their robustness, durability, or accuracy. It should be noted that although we will perform our experiments on a PyTorch model and on an NVIDIA GPU, Triton kernel development is supported by additional frameworks and underlying HWs.

Triton as a Component of Torch Compilation

In previous posts (e.g., here) we demonstrated the use of PyTorch compilation and its potential impact on runtime performance. The default compiler used by the [torch.compile](https://pytorch.org/docs/stable/generated/torch.compile.html)r is TorchInductor which relies heavily on Triton kernels for its GPU acceleration. Thus, it seems only appropriate that we begin our Triton exploration by assessing the automatic Triton-backed optimization afforded by torch.compile. The code block below includes the same forward pass of the face detection model we introduced in our previous post along with the compiled GIOU loss function. For the sake of brevity, we have omitted some of the supporting code. Please refer to our previous post for the full implementation.

def loss_with_padding(pred, targets):
    mask = (targets[...,3] > 0).to(pred.dtype)
    total_boxes = mask.sum()
    loss = generalized_box_iou(targets, pred)
    masked_loss = loss*mask
    loss_sum = masked_loss.sum()
    return loss_sum/torch.clamp(total_boxes, 1)

device = torch.device("cuda:0")
model = torch.compile(Net()).to(device).train()
loss_fn = torch.compile(loss_with_padding)

# forward portion of training loop wrapped with profiler object
with torch.profiler.profile(
   schedule=torch.profiler.schedule(wait=5, warmup=5, active=10, repeat=1)
) as prof:
    for step, data in enumerate(train_loader):

        with torch.profiler.record_function('copy data'):
            images, boxes = data_to_device(data, device)
            torch.cuda.synchronize(device)

        with torch.profiler.record_function('forward'):
            with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
                outputs = model(images)
            torch.cuda.synchronize(device)

        with torch.profiler.record_function('calc loss'):
            loss = loss_fn(outputs, boxes)
            torch.cuda.synchronize(device)
        prof.step()
        if step > 30:
            break

    # filter and print profiler results
    event_list = prof.key_averages()
    for i in range(len(event_list) - 1, -1, -1):
        if event_list[i].key not in ['forward', 'calc loss', 'copy data']:
            del event_list[i]
    print(event_list.table())

The performance results (averaged over multiple runs) are captured below:

-------------  ------------  ------------
         Name     CPU total  CPU time avg
-------------  ------------  ------------
    copy data      56.868ms       5.687ms
      forward        1.329s     132.878ms
    calc loss       8.282ms     828.159us
-------------  ------------  ------------

Recall that the average time of the original loss function (on padded input) was 1.844ms. Thus the performance boost resulting from torch compilation is greater than 2X(!!).

The Triton kernels automatically generated by torch.compile can actually be viewed by setting the TORCH_LOGS environment variable, as explained in this PyTorch tutorial. In fact, some have proposed the use of these kernels as a starting point for Triton development (e.g., see here). However, in our experience these kernels can be somewhat difficult to decipher.

In the next section we will attempt to further improve on the results of PyTorch compilation by implementing a GIOU Triton kernel.

Creating a Custom Triton Kernel

A great place to start your Triton development journey is with the official Triton tutorials. The tutorials are introduced in incremental order of complexity, with each one expanding on one or more of Triton’s unique features. Our GIOU Triton kernel most closely resembles the most basic vector addition example. As in our CUDA implementation, we assign a block to each sample in the input batch, and program it to operate on all of the bounding boxes in the sample. Note the use of tl.load and tl.store for reading and writing data from and to memory, as well as the block programs use of vectorized arithmetic.

import triton
import triton.language as tl

@triton.jit
def giou_kernel(preds_ptr,
                targets_ptr,
                output_ptr,
                valid_ptr,
                BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    box_id = tl.arange(0, BLOCK_SIZE)

    box_offsets = pid * BLOCK_SIZE + box_id

    preds_left = tl.load(preds_ptr + 0 + 4 * box_offsets)
    preds_top = tl.load(preds_ptr + 1 + 4 * box_offsets)
    preds_right = tl.load(preds_ptr + 2 + 4 * box_offsets)
    preds_bottom = tl.load(preds_ptr + 3 + 4 * box_offsets)

    gt_left = tl.load(targets_ptr + 0 + 4 * box_offsets)
    gt_top = tl.load(targets_ptr + 1 + 4 * box_offsets)
    gt_right = tl.load(targets_ptr + 2 + 4 * box_offsets)
    gt_bottom = tl.load(targets_ptr + 3 + 4 * box_offsets)

    epsilon = 1e-5

    # Compute the area of each box
    area1 = (preds_right - preds_left) * (preds_bottom - preds_top)
    area2 = (gt_right - gt_left) * (gt_bottom - gt_top)

    # Compute the intersection
    left = tl.maximum(preds_left, gt_left)
    top = tl.maximum(preds_top, gt_top)
    right = tl.minimum(preds_right, gt_right)
    bottom = tl.minimum(preds_bottom, gt_bottom)

    inter_w = tl.maximum(right - left, 0)
    inter_h = tl.maximum(bottom - top, 0)
    inter_area = inter_w * inter_h

    union_area = area1 + area2 - inter_area

    iou_val = inter_area / tl.maximum(union_area, epsilon)

    # Compute the smallest enclosing box
    enclose_left = tl.minimum(preds_left, gt_left)
    enclose_top = tl.minimum(preds_top, gt_top)
    enclose_right = tl.maximum(preds_right, gt_right)
    enclose_bottom = tl.maximum(preds_bottom, gt_bottom)

    enclose_w = tl.maximum(enclose_right - enclose_left, 0)
    enclose_h = tl.maximum(enclose_bottom - enclose_top, 0)
    enclose_area = enclose_w * enclose_h

    # Compute GIOU
    delta_area = (enclose_area - union_area)
    enclose_area = tl.maximum(enclose_area, epsilon)
    giou = iou_val - delta_area / enclose_area

    # Store results
    tl.store(output_ptr + (box_offsets),
             tl.where(gt_bottom > 0, giou, 0))
    tl.store(valid_ptr + (box_offsets), gt_bottom > 0)

def loss_with_triton(pred, targets):
    batch_size = pred.shape[0]
    n_boxes = pred.shape[1]

    # convert to float32 (remove to keep original dtypes)
    pred = pred.to(torch.float32)
    targets = targets.to(torch.float32)

    # allocate output tensors
    output = torch.empty_strided(pred.shape[0:2],
                                 stride=(n_boxes,1),
                                 dtype = pred.dtype,
                                 device = pred.device)
    valid = torch.empty_strided(pred.shape[0:2],
                                stride=(n_boxes,1),
                                dtype = torch.bool,
                                device = pred.device)

    # call Triton kernel
    giou_kernel[(batch_size,)](pred, targets, output, valid,
                               BLOCK_SIZE=n_boxes)

    total_valid = valid.sum()
    loss_sum = output.sum()
    return loss_sum/total_valid.clamp(1)

The results of running with our Triton kernel are captured below. While somewhat worse than in our previous experiment, this could be a result of additional optimizations performed by torch.compile.

-------------  ------------  ------------
         Name     CPU total  CPU time avg
-------------  ------------  ------------
    copy data      57.089ms       5.709ms
      forward        1.338s     133.771ms
    calc loss       8.908ms     890.772us
-------------  ------------  ------------

Following the recommendation of PyTorch’s documentation on the use of Triton kernels, we further assess the performance of our kernel, this time in combination with PyTorch compilation. The results (averaged over multiple runs) are slightly better than the auto-compiled loss of our first experiment.

-------------  ------------  ------------
         Name     CPU total  CPU time avg
-------------  ------------  ------------
    copy data      57.008ms       5.701ms
      forward        1.330s     132.951ms
    calc loss       7.189ms     718.869us
-------------  ------------  ------------

When developing our custom GIOU CUDA kernel, we noted the overhead of converting the input tensors to float32, and the need to enhance our kernel to support various input types in order to avoid this conversion. In the case of our Triton kernel this can be accomplished quite easily by simply removing the conversion operations. The custom kernel will be auto-generated (JIT-compiled) with the original types.

-------------  ------------  ------------
         Name     CPU total  CPU time avg
-------------  ------------  ------------
    copy data      57.034ms       5.703ms
      forward        1.325s     132.456ms
    calc loss       6.219ms     621.950us
-------------  ------------  ------------

Our final results are on par with CUDA kernel results that we saw in our previous post.

Results

The following table summarizes the results of our experimentation. The results were averaged over multiple runs due to some variance that we observed. We have included the results of our custom CUDA kernel from our previous post, for reference. Keep in mind that the comparative results are likely to vary greatly based on the details of the kernel and the runtime environment.

While our first Triton kernel experiment resulted in reduced performance, compared to our custom CUDA operator, by applying compilation and removing the data type conversions, we were able to match its speed.

These findings are in line with what one might expect from Triton: On the one hand, its high-level API abstraction implies a certain loss of control over the low-level flow which could result in reduced runtime performance. On the other hand, the (relative) simplicity and power of its APIs enable users to close the performance gap by implementing features with much greater ease than in CUDA.

One could make a strong argument that the Triton kernel we chose to evaluate is what the documentation would refer to as "embarrassingly parallel", i.e., comprised of element-wise operations, and that as such, is a terrible kernel on which to demonstrate the value of Triton. Indeed, a more complex program, requiring more sophisticated memory management, scheduling, synchronization, etc., may be required to showcase the full power of Triton.

Next Steps

Several additional steps are required to complete our task. These include tuning our custom kernel and implementing the backward function.

1. Kernel Optimization

Although, Triton abstracts away a lot of the low-level kernel optimization, there remain many controls that could greatly impact runtime performance. These include the size of each block, the number of thread warps to use (as demonstrated in the softmax tutorial), and how L2 memory is accessed (see the [matrix multiplication tutorial](https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html) for an example of swizzling). Triton includes an autotuning feature for optimizing the choice of hyper-parameters (as demonstrated in the matrix multiplication tutorial and in the PyTorch Triton example). Although we have omitted autotuning from our example, it is an essential step of Triton kernel development.

2. Backward Pass Implementation

We have limited our example to just the forward pass of the GIOU loss function. A full solution would require creating a kernel for the backward pass, as well (as demonstrated in the layer normalization tutorial). This is usually a bit more complicated than the forward pass. One may wonder why the high-level kernel development API exposed by Triton does not address this challenge by supporting automatic differentiation. As it turns out, for reasons that are beyond the scope of this post (e.g., see here), automatic differentiation of custom kernels is extremely difficult to implement. Nonetheless, this would be an absolute killer of a feature for Triton and we can only hope that this will be supported at some point in the future.

Summary

Triton is easily one of the most important and impactful AI/ML libraries of the past few years. While it is difficult to assess the amount of innovation and progress it has enabled in the field of AI, its footprints can be found everywhere – from the core implementation of PyTorch 2 and its dependencies, to the specialized attention layers within the advanced LLM models that are slowly perforating our every day lives.

Triton’s popularity is owed to its innovative programming model for kernel development. Once limited to the domain of CUDA experts, Triton makes creating customized DL primitives accessible to every Python developer.

In this post we have only touched the surface of Triton and its capabilities. Be sure to check out the Triton’s online documentation and other resources to learn more.

Written By

Topics:

Share this article:

Related Articles

Implementing Convolutional Neural Networks in TensorFlow

Step-by-step code guide to building a Convolutional Neural Network

What Do Large Language Models “Understand”?

A deep dive on the meaning of understanding and how it applies to LLMs

How to Forecast Hierarchical Time Series

A beginner’s guide to forecast reconciliation

Deep Dive into LSTMs & xLSTMs by Hand ✍

Explore the wisdom of LSTM leading into xLSTMs - a probable competition to the present-day LLMs

Does Your Company Have a Data Strategy?

This sophistication matrix can show you where you need to go

Speeding Up the Vision Transformer with BatchNorm

How integrating Batch Normalization in an encoder-only Transformer architecture can lead to reduced training time…

The Math Behind Keras 3 Optimizers: Deep Understanding and Application

This is a bit different from what the books say.

Build Your Own Modular Audio Course on AI Ethics and Safety

A hand-picked “listening list” on the questions and stakes at the forefront of artificial intelligence…

Latest picks: The ethics of AI

Your daily dose of data science

Your home for data science and Al. The world’s leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.

Sign up to our newsletter


Scaling Intelligence Lab

KernelBench: Can LLMs Write GPU Kernels?

Anne Ouyang*

                    Stanford

Simon Guo*

                    Stanford

Azalia Mirhoseini

                    Stanford

A benchmark designed to evaluate the ability of LLMs to generate efficient GPU kernels for optimizing neural network performance

TL;DR

We introduce KernelBench, a benchmark designed to evaluate the ability of large language models (LLMs) to generate efficient GPU kernels for optimizing neural network performance. With 250 well-defined neural network tasks spanning foundational operators, simple fusion patterns, and full ML architectures, the benchmark tasks LLMs to replace PyTorch implementations with custom kernels that are correct and performant. KernelBench highlights the potential for agentic optimization for computer systems with dense feedback signal, where systems iteratively refine kernel designs using profiling tools and tight feedback loops to achieve near-peak hardware utilization. As models scale, well-optimized kernels have far-reaching implications, from reducing the massive energy demands of AI systems to enabling fair and efficient comparisons of novel architectures. By providing aspirational tasks and focusing on agentic approaches, KernelBench envisions a future where LLMs can autonomously drive innovation in GPU programming and ML system optimization.

Kernels are the kernel of deep learning.

…but writing kernels sucks.

Consider a machine learning researcher with a promising new attention mechanism that could improve LLM efficiency by 30%. To actually test out this idea, they need to:

In an ideal world, you could:

This future (hopefully) isn’t science fiction: we think it’s possible. To measure progress, we’re introducing KernelBench, a dataset of 250 well-defined neural network operations with reference implementations given in Pytorch. KernelBench measures the ability of LLMs to write custom GPU kernels that implement and accelerate these operations. Beyond the 250 core tasks in KernelBench, we also introduce 20 aspirational tasks from HuggingFace models to benchmark the ability of LLM systems in not only just writing GPU kernels but also working on integrating GPU code optimizations in a software library setting.

Why Are Kernels Important?

As models grow larger and become more embedded into our daily lives, having fine-grained control over hardware resources to extract the most performance out of GPUs directly translates to significant energy and cost reductions. For example, ChatGPT alone is estimated to consume over half a million kilowatt-hours daily — roughly equivalent to the power usage of 180,000 U.S. households. At this scale, a 5% speedup isn’t just a number on a benchmark, it’s real energy and money saved. Beyond savings, optimized GPU kernels also allow machine learning researchers to fairly evaluate and compare new model architectures, and efficiency often means unlocking new capabilities that push the field of AI forward.

Big O is not all you need.

In algorithm classes we are taught to view Big O as the gold standard for measuring the efficiency of algorithms. In ML research, new model architectures may have better theoretical complexity, implying they should outperform traditional architectures in speed or efficiency, but when it comes down to real-world performance, these newer models can struggle to keep up with established architectures.

(Meme credit to Michael Zhang)

Why doesn’t Big O analysis match actual performance?

Established architectures benefit from years of optimization in their underlying kernels. These kernels are tailored to run efficiently on specific hardware, exploiting all the features of the hardware to maximize performance. On the flip side, newer models often lack this level of optimization and lack adequate hardware utilization, which can result in disappointing performance despite their appealing theoretical claims.

Optimized GPU kernels are important for designing ML architectures. The lack of well-written GPU kernels makes it difficult to do apples-to-apples model architecture comparisons given a fixed compute budget and a fixed hardware platform, so we cannot effectively determine the effectiveness of an architecture. Consider the following example scenarios:

Can we use LLMs to generate correct and performant GPU kernels?

Unlike many other coding tasks, writing efficient GPU kernels is challenging due to the need for parallelization scheme design, memory management, and hardware-specific optimizations. GPU programming isn’t just about writing syntactically correct code; it requires a deep understanding of GPU architecture to ensure that code is both correct (produces the right output) and performant (fully utilizes the GPU’s capabilities). These factors make GPU programming a rich problem for LLMs, as it involves a bigger optimization search space beyond basic syntax or logic generation.

Recent work on inference scaling laws shows that when you have automatic verifiers, throwing more compute at generation can dramatically improve success rates. Our lab’s recent work (Large Language Monkeys) showed that with coding tasks, going from 1 to 250 samples boosted the solve rate from 15.9% to 56% on SWE-Bench Lite with DeepSeek-Coder-V2-Instruct. GPU programming is a task with strong verification mechanisms and clear feedback signals. For correctness, ground truth is determined by running the generated code on random inputs and comparing the outputs with those of the baseline to check if they match. Performance is measured as the wallclock time and comparing speedup over the reference baseline.

GPU programming is also great for agentic and RL approaches, as the system can iteratively refine kernel designs with reliable, measurable outcomes. Profiling tools like NVIDIA’s Nsight Compute (NCU) provide in-depth feedback on performance bottlenecks, memory usage, and thread utilization, which gives the agent a lot of data to adjust optimizations and improve efficiency. Together, these qualities create a structured environment and tight feedback loop where an agent has verifiable correctness and performance metrics to iterate toward increasingly optimized and correct kernel code.

KernelBench

We introduce KernelBench, a collection of 250 PyTorch neural network operations that we think systems should be able to automatically write optimized kernels for. KernelBench provides the reference implementation in PyTorch, and the task for an LLM is to replace torch layers with custom implementations. We currently only focus on the forward pass as a first step. The core tasks in KernelBench are divided into three levels, with an additional level 4 of 20 aspirational tasks:

We only provide baseline evaluations for the first three levels. Level 4 is currently a far-reaching aim and we do not provide baseline evaluations for this level; however, we believe that this level could ultimately play a significant role in advancing the capabilities of LLMs to interact with complex, real-world codebases, where they not only assist with code generation but can also drive architectural improvements and optimizations in widely-used frameworks.

The tasks in KernelBench are a mix of written manually, generated by an LLM or script, and collected from Github. All tasks are manually cleaned up and verified. Each problem has a class named Model to denote which torch-based architecture we want optimized. The torch reference implementations are cleaned up to be self-contained in one file, with the modules containing only the init and forward functions (and helper functions called in the init and forward functions). In addition to the torch module, we also provide functions get_inputs() and get_init_inputs() for generating random parameters for the forward pass and the initialization, respectively. The shapes of random inputs for testing are manually chosen. We also modify the architecture manually to eliminate operations such as dropout to make the results deterministic (within a generous tolerance threshold).

While the tasks (architecture implementations) are given in Pytorch, KernelBench is language agnostic, allowing the solutions to use any libraries and DSLs (including Triton, ThunderKittens, CUTLASS, …) such that different levels of abstractions for GPU programming can be explored. It is also fully flexible for the LLMs to determine the optimizations to apply (e.g. making decisions such as kernel fusion).

Here’s a simple example of vector addition to illustrate our task format and a CUDA-based solution:

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b

def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]

def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

Here’s an example of a CUDA based solution using custom CUDA C++ operators in torch via load_inline(). This entire file is LLM-generated. The custom CUDA code is supplied as a string and JIT compiled. In this example, the torch addition expression for two vectors is swapped out with the custom elementwise_add_kernel.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name='elementwise_add',
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=['elementwise_add_cuda'],
    verbose=True,
    extra_cflags=[''],
    extra_ldflags=['']
)

class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

Evaluation

When evaluating GPU kernels, we focus on three criteria with each building upon the previous:

In our context, correctness specifically means, given randomized input values for a predefined set of shapes, the optimized kernel should yield outputs that are numerically equivalent (within an acceptable margin of error, if necessary for floating-point operations) to those produced by the baseline implementation. We choose our numerical equivalent threshold as having absolute and relative tolerances being 1e-02, a generous threshold enabling precision changes and alternative algorithms. A common tradeoff in GPU kernel design is specialization versus generality. Specialized kernels, tuned for particular input shapes or patterns, can often achieve significant performance gains; general-purpose kernels, by contrast, aim for broader compatibility but may sacrifice peak performance. For our purpose, since the aim of our project is to cheaply and quickly generate specialized kernels, we choose to constrain our correctness checks to specified input shapes without requiring broad generalization across all possible shapes. We generate 5 sets of random inputs with fixed shapes, and the kernel is considered to be correct if it produces the numerical equivalent outputs as the unoptimized baseline for all 5 inputs. It is possible to have a stricter measurement of correctness by using more random inputs (the number of correctness trials is a customizable parameter), but we capped it at 5 due to evaluation speed.

Caption: Illustration of KernelBench design

In KernelBench, we made the decision not to provide a predefined train/validation/test set split; however, users are welcome to create their own splits based on their specific needs and goals. Our benchmark doesn’t include additional information to distinguish between training and testing examples, as the focus is on real-world challenges that demand open-ended, high-performance solutions—writing custom GPU kernels and optimizing them to their absolute limits (you’re only constrained by the speed of light!). These tasks, which revolve around foundational operators for machine learning, are designed to have a meaningful impact. Improving these kernels, in any way, can potentially lead to substantial real-world benefits.

Initial Evaluation

Baseline Greedy Evaluation

We evaluate the 250 problems from Levels 1 to 3 from KernelBench on various frontier models, with greedy decoding parameters (temperature = 0).

Compilation and Correctness

Across the 3 levels of the problems, most models do generate compilable CUDA code. However, maintaining correctness (same output as torch reference code) becomes increasingly challenging as the reference torch code gets more complex (simple operators in level 1 to fused operators in level 2 to whole model architecture in level 3).

Comparing across various models, we note while some models do well on Level 1 tasks, correctness quickly drops off for Level 2 and 3 tasks. Larger models of the same family also seem to get more correct solutions. It is also particularly interesting that the o1 model does significantly better than gpt-4o on correctness for more challenging Level 2 and Level 3 problems, highlighting scaling inference time compute might have played a role here.

Caption: Percent of Correct Samples across 3 Levels of problems across models

Pass@k

Beyond greedy decoding, we are also interested in pass@k, having at least 1 correct (and successfully compiled) solution given k attempts, as introduced in the HumanEval paper. We sample models with high decoding temperature (deepseek-coder with temp=1.6, and Llama 3.1 70b-Instruct with temp=0.8) for more diverse samples, compute pass@1,3,5,10 with N=100 samples.

Pass@k is defined as
\(\text{pass@$k$} := \mathop{\mathbb{E}}_{\text{problems}} \left[ 1 - \frac{\binom{n - c}{k}}{\binom{n}{k}} \right]\)
where $n$ is the total number of samples and $c$ is the number of correct samples.

Caption: Pass@k performance for Deepseek-coder and Llama 3.1 70B Instruct

As we increase k, correctness improves, suggesting it might be easier to solve such tasks with more parallel samples (as introduced in the Large Language Monkeys paper). However, we see a stark difference between deepseek and llama 3.1 70b performance, highlighting the importance of base model capability even when conducting inference time scaling.

Tradeoff between Correctness and Performance

We only analyze correctness in the section above. However, in the case of kernel engineering, we care deeply about performance. Looking at the generated kernels, we found there is a tradeoff between correctness and performance, two objectives that are often at odds with each other. Code with more optimization could give better performance gain, but could also risk making more errors and hence likely to fail correctness. Optimizing for performant code while guaranteeing correctness creates a new direction for code generation, while most existing benchmarks and methodologies focus on passing correctness; we are excited to keep exploring that.

Performance: Percentiles of Speedups

When evaluating performance, we prioritize correctness, as incorrect but fast code is not useful. Therefore, speedups are calculated using only the correct samples. To present a comprehensive view of performance, we report speedups in percentiles. The count of correct samples for each model is indicated in parentheses after the model name in the table below.

In addition to the baseline PyTorch implementation, we also compare speedups against torch.compile() using its default mode. The speedup is defined as
\(\frac{t\_baseline}{t\_generated}\)

Caption: Percentile of Speedups vs. Baseline for both Torch and Torch Compile across 3 levels

Among the samples that are correct, we see that most generated kernels exhibit relatively slow speedups over torch and torch.compile baseline, but a few are notably faster as outliers! This piqued our interest and led us to the following investigations.

“Kernelsseum” –– A Per Problem Leaderboard

To better understand the LLM-generated kernels, we also present a leaderboard to inspect the kernels generated by greedy evaluation on KernelBench. This shows the top 5 LLM-generated kernels per problem, and some problems might lack any correct solutions. Note the performance result is hardware-dependent and currently evaluated on the Nvidia L40S GPU.

You can click on entries to see the generated code for each kernel.

Right now, the leaderboard only features solutions generated through greedy evaluation. In the future, we aim to make it an open submission leaderboard to allow contributions from the broader community.

Interesting Kernels

Diagonal Matrix Multiplication

Problem 13 in level 1 involves multiplying a matrix by another diagonal matrix:

torch.diag(A) @ B

torch.diag() takes in a vector of the diagonal elements of a matrix and returns a 2-D square tensor with the elements of input as the diagonal. The result is a matrix-matrix multiplication.

Mathematically, multiplying a matrix by a diagonal matrix is equivalent to scaling each row (or column, if the diagonal matrix is on the right side) of the original matrix by the corresponding diagonal element. As a result, the diagonal matrix doesn’t need to be explicitly constructed, reducing both memory usage and computational overhead.

This is the problem that gets the >12x speedup over torch and torch.compile() in level 1 for multiple models, one example of these generated CUDA kernel is below:

__global__ void diag_matmul_kernel(
    const float* diag,
    const float* mat,
    float* out,
    const int N,
    const int M) {

    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < N && col < M) {
        out[row * M + col] = diag[row] * mat[row * M + col];
    }
}

Kernel Fusion

Problem 14 in level 2 performs a matrix multiplication, division, summation, and then scaling:

x = torch.matmul(x, self.weight.T)  # Gemm
x = x / 2  # Divide
x = torch.sum(x, dim=1, keepdim=True) # Sum
x = x * self.scaling_factor  # Scaling

There’s a solution generated by claude-3.5-sonnet that has an approximately 3x speed up over both torch and torch compile:

// Fused kernel for matmul + divide + sum + scale
__global__ void fused_ops_kernel(
    const float* input,
    const float* weight,
    float* output,
    const float scaling_factor,
    const int batch_size,
    const int input_size,
    const int hidden_size
) {
    // Each thread handles one element in the batch
    const int batch_idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (batch_idx < batch_size) {
        float sum = 0.0f;

        // Compute matmul and divide for this batch element
        for(int h = 0; h < hidden_size; h++) {
            float elem = 0.0f;
            for(int i = 0; i < input_size; i++) {
                elem += input[batch_idx * input_size + i] *
                        weight[h * input_size + i];
            }
            // Divide by 2 as we go
            sum += (elem / 2.0f);
        }

        // Scale and store final result
        output[batch_idx] = sum * scaling_factor;
    }
}

The solution fuses all four operations into a single GPU kernel, eliminating the overhead of writing intermediate results to memory and reading them back for subsequent operations. Additionally, combining the matrix multiplication with the dimension-wise summation reduces the size of the final output, minimizing memory bandwidth usage.

Next steps and public involvement

In the Scaling Intelligence Lab, we plan to continue extending this work to enable LLMs to write efficient GPU kernels. The initial results show significant room for improvement, and we are optimistic about the potential for significant advancements in future iterations.

There is a lot of interest in the community in GPU programming and LLMs for GPU code generation. In particular, Project Popcorn of the GPU Mode Discord aims to build “an LLM that can actually write good GPU code”. There is also interest in running GPU programming competitions for humans as a way to collect high quality training tokens for the LLM. We look forward to seeing how KernelBench can contribute to these initiatives.

Our vision for the longer term future is to simplify the generation of high-performance kernels that seamlessly adapt to diverse hardware architectures, enabling developers to achieve optimal performance with minimal effort. By accelerating the iteration cycles for machine learning model architecture design, we aim to empower researchers and practitioners to explore, prototype, and optimize ideas faster than ever.

In addition, the ability to generate kernels quickly is very important for adapting to new hardware architectures, which is often a barrier for adoptions of new computing platforms. We think KernelBench and related techniques could enable faster development cycles for new hardware by lowering the amount of human engineering effort to write new kernels for architecture.

Let’s make writing high-performance kernels far more accessible and convenient!

FAQ

Why not a compiler?

The current development cycle—from efficient implementations to generalizations to compiler integration—is lengthy. Efficient compilers often lag behind new GPU architectures by over two years: approximately one year for CUDA experts to develop optimized implementations and another year to generalize these optimizations into compilers. Traditional compilers excel in generating provably-correct, robust, and general-purpose solutions, making them indispensable for a wide range of applications. However, developing compilers remains a labor-intensive and time-consuming process.

Many design patterns and optimizations are reusable across GPU kernels –– fundamental principles such as overlapping, fusion, efficient memory access, and maximizing occupancy. Our approach seeks to complement traditional compilers by focusing on a different objective. Rather than striving for general-purpose, provably-correct compiler solutions, we aim to distill human intuition directly into specialized, high-performance code with correctness tested empirically. This enables the generation of code highly optimized for specific input shapes and computational patterns, a level of specialization that would otherwise require extensive pattern-matching rules and manual engineering in traditional compilers.

Acknowledgements

We would like to thank Aaryan Singhal, AJ Root, Allen Nie, Anjiang Wei, Benjamin Spector, Bilal Khan, Bradley Brown, Dylan Patel, Genghan Zhang, Hieu Pham, Hugh Leather, John Yang, Jon Saad-Falcon, Jordan Juravsky, Mark Saroufim, Michael Zhang, Ryan Ehrlich, Sahan Paliskara, Sahil Jain, Shicheng (George) Liu, Simran Arora, Suhas Kotha, Vikram Sharma Mailthody, and Yangjun Ruan for insightful discussions and constructive feedback in shaping this work. We would also like to thank SWEBench for its inspiration and reference, which greatly contributed to the development of this work.

Citing

@misc{ouyang2024kernelbench,
      title={KernelBench: Can LLMs Write GPU Kernels?},
      author={Anne Ouyang and Simon Guo and Azalia Mirhoseini},
      year={2024},
      url={https://scalingintelligence.stanford.edu/blogs/kernelbench/},
}

Materials


ROCm blogs

GEMM Kernel Optimization For AMD GPUs

Contents

GEMM Kernel Optimization For AMD GPUs#

Matrix multiplication underlies critical computational pathways in AI, with General Matrix Multiplication (GEMM) operations serving as performance-critical kernels in neural network architectures. From fully connected layers to convolutions and transformer attention mechanisms, GEMMs consume substantial computational and memory resources in large language models (LLMs). This blog explores GEMM optimization techniques for AMD GPUs, demonstrating methodologies to significantly enhance computational efficiency and performance scaling.

ROCm software tools for GEMM Tuning#

To assist AMD GPU developers in efficiently discovering the best GEMM solutions, the ROCm software suite offers multiple tools designed to tune GEMM operation performance. Developers can select the appropriate tool based on their specific use case, as illustrated in the diagram below.

Let’s dive into the various GEMM tuning tools available for AMD GPU developers to use.

GEMM Tuning Techniques on AMD Instinct GPUs#

Technique 1: Optimizing Performance with Pre-Tuned GEMM Operations#

AMD provides an optimized ROCm docker for an out-of-the-box experience including a pre-tuned GEMM, in the vLLM docker. This rocm/vllm Docker has already integrated the pre-tuned GEMM solution with BLAS libraries, supporting most GEMM shapes for LLM inference. We highly recommend GPU developers to try AMD optimized Docker first.

To get started, follow the detail steps below:

1). Pull the optimized docker image from the ROCm/vLLM Docker Hub website.

docker pull rocm/vllm:rocm6.3.1_mi300_ubuntu22.04_py3.12_vllm_0.6.6

2). Run the LLM performance benchmark using the vLLM benchmarking tool. Since the pre-tuned GEMM configuration files (.csv) are integrated into the optimized Docker, use the vLLM benchmarking tool, it automatically utilize the pre-tuned GEMM for optimal performance. We use vllm latency benchmarking tool as the example, and the detailed info of vllm benchmarking tool can be found from vLLM benchmark.

python /app/vllm/benchmarks/benchmark_latency.py \
 --model ${model_path} \
 --trust-remote-code \
 --num-iters-warmup 3 \
 --num-iters 5 \
 --dtype float16 \
 --input-len {in_len} \
 --output-len {out_len} \
 --batch-size ${bs} \
 --tensor-parallel-size ${tp_nums} \
 --num-scheduler-steps 10

Technique 2: Optimizing Performance with PyTorch TunableOp (Framework Level GEMM Tuning)#

PyTorch TunableOp provides a GEMM tuning wrapper for both rocBLAS and hipBLASLt. Instead of relying on default GEMMs, TunableOp automatically searches for the optimal solution by querying the underlying BLAS library for all available solutions for a given GEMM, benchmarking each one, and selecting the fastest. The chosen solution is then stored on disk for use in subsequent runs.

For applications leveraging popular frameworks like PyTorch and vLLM, users can leverage PyTorch TunableOp online tuning. This process allows tuning to occur seamlessly while running training or inference workloads, requiring only a few environment setting adjustments. Detailed information about these environment variables can be found from PyTorch TunableOp.

To optimize performance with tuned GEMM operations at the framework level, follow below steps:

1). Configure the related settings to enable PyTorch TunableOp

export PYTORCH_TUNABLEOP_ENABLED=1
export PYTORCH_TUNABLEOP_TUNING=1
export PYTORCH_TUNABLEOP_VERBOSE=1
export PYTORCH_TUNABLEOP_FILENAME=/dockerx/tunableop-config.csv

2). GEMM tuning results will be saved to above tunableop-config.csv file. The GEMM tuning, described in the CSV file, will be integrated into the specific workload associated with your application.
3). Now, turn off tuning before running your application.

export PYTORCH_TUNABLEOP_ENABLED=1
export PYTORCH_TUNABLEOP_TUNING=0
export PYTORCH_TUNABLEOP_VERBOSE=1
export PYTORCH_TUNABLEOP_FILENAME=/dockerx/tunableop-config.csv

4). Run your application. The tuning result integration will work automatically. With native PyTorch support for AMD ROCm, developers can seamlessly leverage the PyTorch TuneableOps flow. In our experiments, this approach has yielded over 20% performance improvement in GEMM operations. Developers can check the details from TunableOp Blog.If developers meet questions or issues about TunableOp GEMM tuning,please submit them in PyTorch issues.

Technique 3: Optimizing Performance with Tuned GEMM Operations at Ops/Library Level#

AMD offers rocBLAS, the AMD library for Basic Linear Algebra Subprograms (BLAS), internally uses Tensile, which supplies the high-performance implementation of GEMM. Additionally, hipBLASLt is a library that provides general matrix-matrix operations.

Based on a developer’s preference they can choose either of the two Ops/Librariesfor GEMM tuning tools, rocBLAS tuning tool (rocblas-gemm-tune) or hipBLASLt tuning tool (hipblaslt-bench).

First, use the logging scheme of either rocBLAS or hipBLASLt (depending on the library in use) to capture the required GEMM shape information. Then, apply the respective GEMM tuning tools (rocblas-gemm-tune or hipblaslt-bench) to optimize performance.

GEMM Tuning with rocblas-gemm-tune#

The rocblas-gemm-tune tool works by using Tensile to heuristically search through various kernel parameters in order to find the optimal configuration that provides high GPU performance for performing GEMM operations.

The detail steps are as below:

1). Installing/rocBLAS Setup: In the ROCm Docker image, the rocBLAS library is pre-installed but if the rocBLAS client related executable bin files (rocblas-bench and rocblas-gemm-tune) are not pre-installed, you may need to build them from source code.

2). Generating GEMM Problem Sizes: rocBLAS provides the logging scheme to dump GEMM shapes info for further performance tuning, which is enabled by rocBLAS environment settings.

- Environment variable `ROCBLAS_LAYER=4` turns on log_profile, and outputs a YAML description of each rocBLAS function called, along with its arguments and number of times it is called. This list of entries can be used directly as input to `rocblas-gemm-tune` utility to do performance tuning.

- Use environment variable `ROCBLAS_LOG_PATH` to set the full path name for all logs, and store the grabbed GEMM shapes information into a YAML file, `ROCBLAS_LOG_PATH=~/dir/rocblas_gemms.YAML`

By using the two settings described above, developers can yield the GEMM shape information.

ROCBLAS_LAYER=4 ROCBLAS_LOG_PATH=./rocblas_gemm.YAML ./gemm-app

3). GEMM Tuning with rocblas-gemm-tune: At this stage, use the dumped YAML file to run GEMM tuning by running rocblas-gemm-tune. The sample command:

```bash
/opt/rocm/bin/rocblas-gemm-tune --YAML /home/rocblas_gemms.YAML
```

Running this will output the fastest solutions for each GEMM in the YAML file. Each solution is identified by an unique solutions index. It generates a CSV file by aggregating the output solution index, and the CSV file form looks like:

transA,transB,M,N,batch_count,K,alpha,beta,lda,ldb,ldc,input_type,output_type,comput_type,solution_index
N, N, 320,588,1,4096,1,0,320,6144,320,f32_r,f32_r,f32_r,3788
N, N, 512,3096,1,512,1,0,512,512,512,f16_r,f16_r,f16_r,4566

4). Integration: Now we have a list of faster solutions for all the GEMM problems, users can integrate this into the application, to pick these faster implementations in rocBLAS by setting the environment variable. Use below example command:

export ROCBLAS_TENSILE_GEMM_OVERRIDE_PATH = csv_file_path

If developers meet questions or issues about rocblas-gemm-tune,please submit them in rocBLAS issues.

GEMM Tuning with hipblaslt-bench#

hipBLASLt-bench is another GEMM tuning tool within hipBLASLt library and can be used to search the best-performing GEMM kernel for a given set of GEMM problems.

To use hipBLASLt, follow below steps:

1). Installing hipBLASLt: In the ROCm Docker image, the hipBLASLt library is pre-installed, however the hipBLASLt client executables, such as hipblaslt-bench, may not be included by default and you may need to build these executables from source.

2). Generating GEMM Problem Size: Similar with rocBLAS, hipBLASLt can also dump the required GEMM problem/shape sizes by its own logging scheme. Detailed info about hipBlASLt logging scheme in logging-heuristics. Use below sample command to generate the GEMM problem sized YAML file:

HIPBLASLT_LOG_MASK=32 HIPBLASLT_LOG_FILE=log_file_name.log ./application_bin

To organize the output logs further, you can get unique calls with call counts like below shell command:

cat log_file_name.log | sort | uniq -c > unique_log_file.log

3). GEMM Tuning with hipblaslt-bench: Set the environment variable HIPBLASLT_TUNING_FILE=<file_name> to tune and store the
tuning result of the best solution indices for the GEMM problems. The <file_name> points to the tuning file. GEMM tuning will be
completed by launching hipblaslt-bench, which input parameters can be set according to the log file of step 2.

A sample command to save file with below user-defined name in the current working directory:

export HIPBLASLT_TUNING_FILE=tuning.txt

/opt/rocm/bin/hipblaslt-bench --api_method c -m 28672 -n 8192 -k 8192 --lda 8192 --ldb 8192 --ldc 28672 --ldd 28672 --stride_a 0 --stride_b 0 --stride_c 0 --stride_d 0 --alpha 1.000000 --beta 0.000000 --transA T --transB N --batch_count 1 --scaleA 1 --scaleB 1 --a_type f8_r --b_type bf8_r --c_type bf16_r --d_type bf16_r --scale_type f32_r --bias_type f32_r --compute_type f32_r --initialization trig_float -i 100 -j 100 --flush --rotating 512 --algo_method all

4). Integration:

Unset tuning file name once tuning is complete:

unset HIPBLASLT_TUNING_FILE

Override the hipBLASLt library with the tuned file info:

export HIPBLASLT_TUNING_OVERRIDE_FILE=tuning.txt

Now we can replace the default GEMM kernel with the tuned GEMM kernel.If developers meet questions or issues about hipblaslt-bench GEMM tuning,please submit them in hipBLASLt issues.

Summary#

Given the pivotal role of GEMM operations in AI workloads, particularly for LLM applications, AMD offers a suite of powerful tuning tools, including rocblas-gemm-tune, hipblaslt-bench, and PyTorch TuneableOps. These tools provide GPU developers with the flexibility to optimize GEMM performance, allowing precise fine-tuning for maximum efficiency on AMD GPUs. By leveraging these resources, developers can enhance workload performance, ensuring optimal execution and superior results in AI-driven tasks.

Additional Resources#

Optimized docker hub: https://hub.docker.com/r/rocm/vllm/tags

Optimized docker image: rocm/vllm:rocm6.3.1_mi300_ubuntu22.04_py3.12_vllm_0.6.6

The opimized docker blog: https://www.amd.com/en/developer/resources/technical-articles/how-to-use-prebuilt-amd-rocm-vllm-docker-image-with-amd-instinct-mi300x-accelerators.html

PyTorch TunableOp: https://pytorch.org/docs/stable/cuda.tunable.html

Improve Performance :Accelerating models on ROCm using PyTorch TunableOp — ROCm Blogs

rocBLAS: https://rocm.docs.amd.com/projects/rocBLAS/en/latest/index.html

hipBLASLt: https://rocm.docs.amd.com/projects/hipBLASLt/en/latest/