Triton LeakyReLU Kernel

mediumFREEElementwise & Fused OpsTriton Programming

Problem Statement

Implement LeakyReLU: out = x if x >= 0 else slope * x, with slope a runtime scalar (default 0.01).

Background

tl.where(cond, a, b) selects elementwise. Combine it with the comparison x >= 0.

Your Task

Implement leaky_relu_kernel and run(n=1024, slope=0.01) comparing to torch.nn.functional.leaky_relu.

How it is tested

Your solution must define a top-level function run(...) that allocates inputs on the GPU, launches your Triton kernel, and returns a boolean from torch.allclose(triton_out, torch_reference, ...). The grader prints run(...); the expected output is True.

Example

Input

n = 1024, slope = 0.01

Output

True

Reasoning

The input values n = 1024 and slope = 0.01 are used to allocate an array of size n on the GPU and define the slope parameter for the LeakyReLU function.
The leaky_relu_kernel function is applied to the input array, using the tl.where function to element-wise select between the original value x and the scaled value slope * x, based on the condition x >= 0: $out = x$ if $x \geq 0$ else $out = slope \cdot x$ .
The resulting array from the leaky_relu_kernel function is compared to the output of torch.nn.functional.leaky_relu using the same input array and slope value.
The comparison is done using torch.allclose, which checks if the two arrays are element-wise equal within a certain tolerance, and returns True if they are equal, resulting in the output True.

Constraints

Use tl.where(x >= 0, x, slope * x)
slope is a runtime scalar argument

More from Triton Programming

Triton Single-Query Attention Kernelhard Triton Row Cumulative Sumhard Triton Stable LogSumExp Kernelhard Triton Seeded Fused Dropouthard Triton Autotuned Vector Addmedium

Concept primer

Background Knowledge

The LeakyReLU activation function is a variation of the ReLU function, which is widely used in deep neural networks. The standard ReLU function maps all negative values to 0 and all positive values to the same value. However, this can lead to the "dying ReLU" problem, where neurons with negative inputs become inactive and do not contribute to the network. The LeakyReLU function addresses this issue by allowing a small fraction of the input to pass through, even if it is negative. This is achieved by multiplying the negative input by a small slope value.

The Triton programming language is a Python-based library for writing high-performance, GPU-accelerated code. It provides a simple and intuitive API for launching kernels on the GPU, making it an ideal choice for implementing custom neural network layers and activation functions. In this problem, we will use Triton to implement the LeakyReLU function and compare it to the reference implementation provided by PyTorch.

The tl.where function in Triton is used to select elements based on a condition. It takes three arguments: a condition, a value to return if the condition is true, and a value to return if the condition is false. This function can be used to implement the LeakyReLU function by selecting between the input value and the scaled input value based on whether the input is positive or negative.

Algorithm/Approach

The general approach to solving this problem involves using the tl.where function to select between the input value and the scaled input value based on the condition x >= 0. We will also need to use the Triton API to launch a kernel on the GPU and perform the necessary computations.

Step-by-Step Strategy

To implement the LeakyReLU function in Triton, follow these steps:

Define a kernel function that takes the input tensor, slope value, and output tensor as arguments.
Use the tl.where function to select between the input value and the scaled input value based on the condition x >= 0.
Launch the kernel on the GPU using the Triton API.
Allocate the input and output tensors on the GPU and launch the kernel.
Compare the output of the Triton kernel to the reference implementation provided by PyTorch using torch.allclose.

Common Pitfalls

When implementing the LeakyReLU function in Triton, be careful to:

Use the correct data types for the input and output tensors.
Ensure that the kernel function is launched on the correct device (GPU).
Use the tl.where function correctly to select between the input value and the scaled input value.

Time & Space Complexity

The time complexity of the LeakyReLU function is $O(n)$ , where $n$ is the number of elements in the input tensor. This is because we need to perform a single pass over the input tensor to compute the output. The space complexity is also $O(n)$ , as we need to allocate space for the output tensor. Note that the Triton kernel will be launched on the GPU, which can provide significant performance benefits for large input tensors.

solution.py

Python 3.13

GPU · T4

Test results