Triton Row Sum Reduction

mediumFREEReductions & NormalizationTriton Programming

Problem Statement

Given a 2D tensor x of shape (M, N), compute the sum along each row, producing a length-M vector. Assume each row fits in one block.

Background

Launch one program per row (grid = (M,)). Load the whole row with BLOCK_SIZE = triton.next_power_of_2(N), mask columns >= N (load other=0.0), then reduce with tl.sum(row, axis=0).

Your Task

Implement row_sum_kernel and run(M=64, N=300) comparing to x.sum(dim=1).

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

M = 64, N = 300

Output

True

Reasoning

The input values are M = 64 and N = 300, representing the shape of the 2D tensor x.
We launch one program per row, resulting in a total of M programs, each loading a row of length N into a block of size BLOCK_SIZE = triton.next_power_of_2(N).
Within each block, we mask columns >= N by loading other=0.0 and then reduce the row using tl.sum(row, axis=0), effectively calculating the sum of each row.
The resulting row sums from the Triton kernel are compared to the reference solution x.sum(dim=1) using torch.allclose, yielding True if the outputs match within a certain tolerance.

Constraints

One program per row, grid = (M,)
BLOCK_SIZE = triton.next_power_of_2(N), pass as constexpr
Use tl.sum(row, axis=0); mask invalid columns with other=0.0

Learn this in

Introduction to TritonGPU Basics with Triton Kernel Fusion: Fused SoftmaxGPU 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 computing the sum along each row of a 2D tensor using Triton, a programming language for writing high-performance GPU code. To tackle this problem, it's essential to understand the basics of parallel computing and GPU programming. In parallel computing, tasks are divided into smaller sub-tasks that can be executed simultaneously, leveraging multiple processing units. GPU programming, in particular, involves writing code that can be executed on a Graphics Processing Unit (GPU), which has many more processing units than a Central Processing Unit (CPU).

In the context of this problem, we're dealing with a 2D tensor x of shape (M, N), where each row needs to be reduced to a single value by summing its elements. The reduction operation is a common pattern in parallel computing, where an array of values is reduced to a single value by applying a binary operation (in this case, addition). To achieve this, we'll be using Triton's programming model, which involves launching a program per row and using blocks to manage memory and computation.

The problem also mentions using triton.next_power_of_2(N) to determine the block size, which is a common technique in GPU programming to ensure efficient memory access and minimize memory fragmentation. Additionally, we'll be using masking to handle rows with fewer than BLOCK_SIZE elements, which is a common pattern in parallel computing to handle irregularly-shaped data.

Algorithm/Approach

The general approach to solving this problem involves:

Launching a program per row, using a grid of size (M,)
Loading each row into a block of size BLOCK_SIZE, which is a power of 2
Masking columns that are out of bounds (i.e., >= N)
Reducing each row using a sum operation, which can be implemented using Triton's tl.sum function

This approach leverages the parallel computing capabilities of the GPU to compute the sum along each row simultaneously, while also handling irregularly-shaped data using masking.

Step-by-Step Strategy

To implement the solution, follow these steps:

Define the row_sum_kernel function, which will be launched per row
Determine the block size using triton.next_power_of_2(N)
Load each row into a block, masking columns that are out of bounds
Reduce each row using a sum operation, storing the result in a register
Define the run function, which will allocate inputs on the GPU, launch the row_sum_kernel, and compare the result to the reference implementation using torch.allclose

Common Pitfalls

When implementing the solution, watch out for:

Incorrectly determining the block size, which can lead to inefficient memory access
Failing to mask columns that are out of bounds, which can result in incorrect results
Incorrectly implementing the sum operation, which can lead to incorrect results

Time & Space Complexity

The expected time complexity of the solution is $O(M \times N)$ , where $M$ is the number of rows and $N$ is the number of columns. The space complexity is $O(M)$ , as we need to store the sum along each row. Note that the actual time and space complexity may vary depending on the specific implementation and the characteristics of the input data.

solution.py

Python 3.13

GPU · T4

Test results