📘

Triton Vector Addition Kernel

EasyTriton Fundamentals 📚 Triton Programming

Problem Statement

Write a Triton kernel that adds two 1D float tensors element-wise: out = x + y.

Background

A Triton program instance handles one block of BLOCK_SIZE elements. Use tl.program_id(0) to get the block index, build per-element offsets with tl.arange, guard out-of-bounds lanes with a mask, then tl.load / tl.store.

Your Task

Implement add_kernel and a run(n=1024) that launches it over a 1D grid and returns whether the result matches x + y.

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

Output:

True

Reasoning:

The input n = 1024 represents the total number of elements in the 1D float tensors x and y.
The add_kernel function is launched over a 1D grid, with each block handling BLOCK_SIZE elements, and performs element-wise addition: $out_i = x_i + y_i$ .
The run function allocates x and y on the GPU, launches the add_kernel, and stores the result in triton_out.
The final output True indicates that triton_out matches the reference result torch_reference = x + y within a tolerance, as verified by torch.allclose.

Constraints:

Use @triton.jit and tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
Mask out-of-bounds lanes with mask = offsets < n
Launch grid = (triton.cdiv(n, BLOCK_SIZE),)

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

Background Knowledge

The problem involves writing a Triton kernel, which is a fundamental concept in Triton Programming. Triton is a programming language and framework for writing high-performance, GPU-accelerated code. In this context, a kernel refers to a small program that runs on the GPU, performing a specific computation. The Triton Vector Addition Kernel problem requires implementing a kernel that adds two 1D float tensors element-wise.

To understand this problem, it's essential to have a basic knowledge of GPU programming and parallel computing. In GPU programming, a block is a group of threads that can cooperate with each other. The BLOCK_SIZE variable determines the number of elements processed by each block. The tl.program_id(0) function returns the block index, which can be used to calculate the global thread ID. Additionally, Triton provides various functions, such as tl.arange and tl.load, to help with tensor operations and memory management.

The problem also involves tensor operations, which are a crucial aspect of deep learning and scientific computing. Tensors are multi-dimensional arrays used to represent data, and operations like element-wise addition are common in many applications. In this problem, the goal is to implement an efficient kernel that performs element-wise addition of two 1D float tensors.

Algorithm/Approach

The general approach to solving this problem involves:

Launching a 1D grid of blocks, where each block processes a subset of elements from the input tensors.
Using thread indices and block indices to calculate the global thread ID and access the corresponding elements in the input tensors.
Implementing element-wise addition using Triton functions like tl.load and tl.store to load and store data from the input and output tensors.
Using masking to guard against out-of-bounds accesses and ensure that the kernel only processes valid elements.

Step-by-Step Strategy

To implement the solution, follow these steps:

Define the add_kernel function, which will be launched as a Triton kernel.
Calculate the global thread ID using tl.program_id(0) and tl.arange.
Load the input elements from the x and y tensors using tl.load.
Perform element-wise addition and store the result in the out tensor using tl.store.
Implement masking to guard against out-of-bounds accesses.
Define the run function, which launches the add_kernel over a 1D grid and returns a boolean indicating whether the result matches the reference output.

Common Pitfalls

When implementing the solution, watch out for:

Out-of-bounds accesses, which can occur if the kernel tries to access elements outside the bounds of the input tensors.
Incorrect indexing, which can lead to incorrect results or crashes.
Inefficient memory access patterns, which can negatively impact performance.

Time & Space Complexity

The expected time complexity of the solution is $O(n)$ , where $n$ is the length of the input tensors. The space complexity is $O(n)$ as well, since the output tensor has the same length as the input tensors. The Triton kernel will be launched over a 1D grid, and the number of blocks and threads will depend on the BLOCK_SIZE and the length of the input tensors.

solution.py

Python 3.13

GPU · T4

Test Results

0/0

Run code to see test results.