Kernel development and optimization with Triton

Author: Jatinkumar parmar Knowledge level: Intermediate

OpenAI Triton is an open-source programming language that is supported by AMD GPUs and is designed to simplify GPU programming for high-performance tasks, particularly in AI applications. This tutorial demonstrates how to set up the Triton development environment and optimize Triton kernel performance on AMD GPUs.

Prerequisites

This tutorial was developed and tested using the following setup.

Operating system

Software

• ROCm 6.2, 6.3, or 6.4: Install and verify ROCm by following the ROCm install guide.
• After installation, confirm your setup using the command. AMD also provides prebuilt ROCm Docker images, for example, a ROCm PyTorch image, ROCm Ubuntu 22.04 image, and ROCm Ubuntu 24.04 image. You can use these prebuilt Docker images to reduce the effort required to set up a ROCm environment.
• Note: For ROCm 6.4 and earlier, use the command instead of .
• Docker: Ensure Docker is installed and configured correctly. Follow the Docker installation guide for your operating system.
• Note: Ensure the Docker permissions are correctly configured. To configure permissions to allow non-root access, run the following commands:

sudousermod-aGdocker$USER newgrpdocker

• Verify Docker is working correctly:

Set up the Triton development environment

This tutorial uses the prebuilt ROCm PyTorch image, but you can also try other ROCm environments as the base image.

Step 1: Launch the Docker image

Launch the Docker container. Replace /path/to/Triton_Sample with the full path to the directory on your host machine where the Triton sample code is located.

Note: This command mounts the current directory to the directory in the container. Ensure the notebook file is either copied to this directory before running the Docker command or uploaded into the Jupyter Notebook environment after it starts. Save the token or URL provided in the terminal output to access the notebook from your web browser. You can download this notebook from the AI Developer Hub GitHub repository.

Step 2: Launch Jupyter Notebooks in the container

Inside the Docker container, install Jupyter using the following command:

Start the Jupyter server:

Note: Ensure port is not already in use on your system before running the above command. If it is, you can specify a different port by replacing with another port number, for example, .

Note: The rest of this tutorial can run as interactive blocks in your Jupyter notebook after you upload this tutorial to your server.

Step 3: Install OpenAI Triton

Before you can install the correct version of OpenAI Triton, you must uninstall any old versions.

1. Uninstall the old version of Triton

It’s strongly recommended that you use the latest version of Triton in your project. AMD and other vendors frequently update their optimization passes and algorithms in OpenAI Triton. These updates can improve your Triton kernel performance.

2. Install OpenAI Triton from the source code

The detailed steps to install Triton are listed below.

Note: If you have any questions or issues when building Triton, submit them to Triton Issues.

Optimize the Triton code for AMD GPUs

The softmax function is often used in convolutional neural network (CNN) classification models and even Transformer-based LLM models. It converts raw output scores, or logits, into probabilities by taking the exponential of each value and normalizing these values by dividing by the sum of all the exponentials. This process ensures that the output values are in the range (0,1) and sum to 1 to allow them to be interpreted as probabilities. PyTorch implements this function as a standard API.

Naive version

According to the specification, you implemented the naive version of the softmax algorithm in the Triton kernel. To determine the maximum data point and the corresponding sum of all the exponentials, this kernel version uses two for-loops, along with one more for-loop to calculate the final softmax result, for a total of three loops in all.

The following example tests the kernel performance on an 8192 x 8192 tensor, with a block size for the column dimension of 256. After running the warmup section to avoid including the kernel compilation time in the final data, you can obtain the performance data for the naive version.

Online softmax version

It’s easy to use the Triton language to implement algorithms. To obtain better performance from the current kernel, first determine whether there’s a more efficient algorithm or solution. If so, try the new algorithm in your Triton Kernel. To reduce the memory accesses caused by the three for-loops in the naive softmax algorithm, a new online softmax algorithm has been proposed in the Online normalizer calculation for softmax paper.

In accordance with the online softmax algorithm, the following code makes a few modifications to the naive version kernel.

Fused-softmax version

OpenAI Triton provides the “fused-softmax” softmax reference example. Based on the online softmax algorithm, it continues to simplify the maximum data calculation by removing one for-loop. It also tells the compiler to use more threads per row by increasing the number of warps. This configuration is often tuned for better performance. Finally, it improves the kernel launching scheme by adjusting the GPU hardware configuration, which can lead to higher GPU kernel occupancy and better performance.

Summary

In this tutorial, you learned how to develop and optimize Triton kernels on AMD GPUs. To learn more about OpenAI Triton, see the official Triton documentation. To find out more about running Triton on AMD GPUs, see ROCm Triton optimization and the Kernel development optimization on Triton blog. Hopefully, this tutorial encourages you to tune, test, and contribute to Triton on AMD GPUs and help shape the future of AI acceleration.

Originally published at https://rocm.docs.amd.com.