Triton Row Cumulative Sum

hardPROAdvanced KernelsTriton Programming

Problem Statement

Compute the inclusive prefix sum (cumulative sum) along each row of a (M, N) tensor using tl.cumsum. Each row fits in one block.

Background

tl.cumsum(x, axis=0) returns the running sum along the block axis. Pad invalid lanes with 0 so they don't affect earlier prefixes, and mask them out on store.

Your Task

Implement cumsum_kernel and run(M=64, N=300) comparing to torch.cumsum(x, 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 M = 64 and N = 300 define the size of the tensor, with 64 rows and 300 columns.
The cumsum_kernel function computes the cumulative sum along each row using tl.cumsum, padding invalid lanes with 0 and masking them out on store.
The result from cumsum_kernel is compared to the result from torch.cumsum(x, dim=1) using torch.allclose, which checks if the two tensors are element-wise equal within a tolerance.
Since the implementation of cumsum_kernel is correct and matches the result from torch.cumsum, the output of run(M=64, N=300) is True, indicating that the two results are identical.

Constraints

Use tl.cumsum(row, axis=0)
Pad invalid lanes with other=0.0 and mask the store
One program per row

More from Triton Programming

Triton Single-Query Attention Kernelhard Triton Stable LogSumExp Kernelhard Triton Seeded Fused Dropouthard Triton Autotuned Vector Addmedium Triton Matrix-Vector Product (GEMV)medium

Concept primerPREVIEW

Background Knowledge

The problem involves computing the cumulative sum along each row of a 2D tensor using Triton, a Python-based programming language and framework for writing high-performance GPU code. The cumulative sum, also known as the prefix sum, is a fundamental operation in linear algebra and is used extensively in various fields, including machine learning, data analysis, and scientific computing. In this context, the cumulative sum is computed along the rows of a (M, N) tensor, where M is the number of rows and N is the number of columns.

The Triton framework provides a tl.cumsum function that computes the running sum along the block axis. However, this function needs to be used in conjunction with other Triton functions and programming constructs to achieve the desired result. The problem also involves comparing the result with the torch.cumsum function from the PyTorch library, which provides a reference implementation for computing the cumulative sum. Understanding the basics of GPU programming, parallel computing, and linear algebra is essential for tackling this problem.

The problem requires a good understanding of memory management, data parallelism, and synchronization techniques to ensure that the computation is performed efficiently and correctly on the GPU. The use of masks and padding is also crucial to handle invalid lanes and ensure that the computation is accurate. The problem is designed to test the learner's ability to write efficient and correct GPU code using the Triton framework.

Algorithm/Approach

The general approach to solving this problem involves using the tl.cumsum function in conjunction with other Triton functions and programming constructs to compute the cumulative sum along each row of the input tensor. The algorithm can be broken down into several steps, including:

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

GPU · T4

Test results