Triton Vector Addition Kernel

easyFREETriton FundamentalsTriton 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

n = 1024 is the length of the two random float vectors $x$ and $y$ ; the target is $\text{out}_i = x_i + y_i$ for every index $i = 0, \dots, n-1$ .
The vector is cut into blocks of BLOCK_SIZE $= B$ elements and each program instance owns one block: program $p$ covers indices $pB, pB+1, \dots, pB + B - 1$ .
With $B = 1024$ , the number of programs is $\lceil 1024 / 1024 \rceil = 1$ , and program $0$ covers indices $0$ to $1023$ — exactly the whole vector, so every lane is a real element.
For $n = 1000$ there is still $\lceil 1000/1024 \rceil = 1$ program, but its last $1024 - 1000 = 24$ lanes (indices $1000$ – $1023$ ) point past the end and must neither read nor write; for $n = 4096$ there are $4$ programs and no spare lanes.
The output True means every computed $\text{out}_i$ agrees with $x_i + y_i$ to within a tiny tolerance (about $10^{-5}$ ). Notice that a single wrongly covered or out-of-range index would make the comparison fail.

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),)

Learn this in

Why GPUs? The Parallel Computing MindsetGPU Basics with Triton Introduction to TritonGPU Basics with Triton Your First Kernel: Vector AdditionGPU Basics with Triton Memory Access, Autotuning & BenchmarkingGPU Basics with Triton

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 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