Triton Autotuned Vector Add

mediumFREEAdvanced KernelsTriton Programming

Problem Statement

Add two vectors but let Triton pick the best BLOCK_SIZE using @triton.autotune.

Background

@triton.autotune benchmarks a list of triton.Configs (keyed on input sizes) and caches the fastest. The grid becomes a lambda over meta so it can read the chosen BLOCK_SIZE.

Your Task

Decorate add_kernel with @triton.autotune over at least two block sizes and implement run(n=4096) that launches with a meta-aware grid.

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

Output

True

Reasoning

n = 4096 is the length of two random float vectors $x$ and $y$ ; the target is $z_i = x_i + y_i$ for every index $i$ from $0$ to $n-1$ .
Each kernel instance handles one contiguous block of $B$ = BLOCK_SIZE elements (instance $p$ covers $pB$ to $pB + B - 1$ ). So the launch needs $\lceil n / B \rceil$ instances, and that count depends on which $B$ wins the benchmark.
With $B = 256$ there are $\lceil 4096/256 \rceil = 16$ instances; with $B = 1024$ there are $4$ . Since $16 \times 256 = 4 \times 1024 = 4096$ , each index is covered once.
The second test, $n = 10000$ , is not a multiple of either size: $B = 1024$ gives $10$ instances and $B = 256$ gives $40$ , both spanning $10240$ slots, so the last $240$ slots lie past the end of the vectors and must contribute nothing.
The output True says every $z_i$ equals $x_i + y_i$ to within $10^{-5}$ . Notice that $B$ changes speed, never the values.

Constraints

Use @triton.autotune with configs of differing BLOCK_SIZE, key=['n']
grid = lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),)
BLOCK_SIZE is supplied by the config, not passed explicitly

Learn this in

Your First Kernel: Vector AdditionGPU Basics with Triton Memory Access, Autotuning & BenchmarkingGPU Basics with Triton Advanced Kernels & OptimizationGPU 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 Matrix-Vector Product (GEMV)medium

Concept primer

Background Knowledge

The problem involves using Triton, a Python-based programming language and framework for writing high-performance GPU code. Triton allows users to define custom kernels, which are small programs that run on the GPU. In this problem, we need to define a kernel that adds two vectors. The Triton Autotune feature is used to automatically find the optimal block size for the kernel, which can significantly impact performance. The block size determines how many threads are executed in parallel on the GPU.

To understand this problem, it's essential to have a basic knowledge of GPU programming and parallel computing. In GPU programming, the GPU is divided into a grid of blocks, and each block is further divided into threads. The number of threads in a block is determined by the block size. A larger block size can lead to better performance, but it also increases the memory usage and can lead to slower performance if the block size is too large. Triton Autotune helps to find the optimal block size by benchmarking different block sizes and caching the fastest one.

The problem also involves using PyTorch, a popular deep learning framework, to allocate inputs on the GPU and launch the Triton kernel. The torch.allclose function is used to compare the output of the Triton kernel with the reference output. The goal is to define a run function that launches the kernel and returns a boolean indicating whether the output is correct.

Algorithm/Approach

The general approach to solve this problem involves defining a Triton kernel that adds two vectors, decorating the kernel with @triton.autotune to find the optimal block size, and launching the kernel using a meta-aware grid. The meta-aware grid allows the kernel to adapt to the optimal block size found by Triton Autotune.

Step-by-Step Strategy

To solve this problem, follow these steps:

Define a Triton kernel that adds two vectors.
Decorate the kernel with @triton.autotune and specify a list of block sizes to benchmark.
Define a run function that allocates inputs on the GPU, launches the Triton kernel, and returns a boolean indicating whether the output is correct.
Use torch.allclose to compare the output of the Triton kernel with the reference output.

Common Pitfalls

When implementing the solution, watch out for the following common pitfalls:

Incorrectly defining the Triton kernel or the run function.
Failing to decorate the kernel with @triton.autotune or specifying an incorrect list of block sizes.
Incorrectly allocating inputs on the GPU or launching the kernel.
Failing to use torch.allclose to compare the output of the Triton kernel with the reference output.

Time & Space Complexity

The time complexity of the solution depends on the block size and the number of threads launched. In general, the time complexity is $O(n)$ , where $n$ is the length of the input vectors. The space complexity is also $O(n)$ , as we need to allocate memory for the input and output vectors on the GPU. However, the actual time and space complexity may vary depending on the specific implementation and the hardware used.

solution.py

Python 3.13

GPU · T4

Test results