Triton Stable LogSumExp Kernel

hardPROAdvanced KernelsTriton Programming

Problem Statement

Compute the row-wise log-sum-exp of a (M, N) tensor in a numerically stable way: lse[m] = log(sum_j exp(x[m,j])).

Background

Subtract the row max before exponentiating, then add it back to the log: m + log(sum(exp(x - m))). This prevents overflow for large inputs and underflow for very negative ones.

Your Task

Implement logsumexp_kernel and run(M=64, N=300) comparing to torch.logsumexp(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 dimensions of the tensor.
For each row m in the tensor, we compute the row-wise max m and subtract it from each element x[m,j] to prevent overflow and underflow: $x'[m,j] = x[m,j] - m$ .
We then calculate the log-sum-exp for each row using the modified values: $lse[m] = m + \log\left(\sum_j \exp(x'[m,j])\right)$ .
The logsumexp_kernel function is compared to the PyTorch reference implementation torch.logsumexp(x, dim=1), and the output True indicates that the results are identical within a certain tolerance, verified by torch.allclose.

Constraints

Subtract row max before exp; add it back after log
One program per row; mask invalid lanes with other=-inf
Use tl.exp, tl.sum, tl.log

Learn this in

Kernel Fusion: Fused SoftmaxGPU Basics with Triton

More from Triton Programming

Triton Single-Query Attention Kernelhard Triton Row Cumulative Sumhard Triton Seeded Fused Dropouthard Triton Autotuned Vector Addmedium Triton Matrix-Vector Product (GEMV)medium

Concept primerPREVIEW

Background Knowledge

The log-sum-exp function is a fundamental component in many machine learning algorithms, particularly in the context of softmax and logistic regression. It is defined as $y = \log \sum_{j=1}^{N} \exp(x_j)$ , where $x_j$ are the input values. However, computing this directly can lead to numerical instability due to the large range of values that can be produced by the exponential function. To mitigate this, a common technique is to subtract the maximum value of the inputs before exponentiating, which helps prevent overflow for large inputs and underflow for very negative ones.

The mathematical justification for this technique lies in the properties of logarithms and exponentials. Specifically, the log-sum-exp function can be rewritten as $\log \sum_{j=1}^{N} \exp(x_j - m) + m$ , where $m = \max_j x_j$ . This transformation ensures that the values being exponentiated are centered around zero, reducing the risk of numerical instability. The log-sum-exp function is also closely related to the softmax function, which is commonly used in machine learning models to produce a probability distribution over a set of classes.

In the context of Triton Programming, the goal is to implement a custom kernel that computes the row-wise log-sum-exp of a tensor in a numerically stable way. This involves leveraging the capabilities of the Triton programming model to optimize the computation for performance and accuracy. The Triton framework provides a set of tools and APIs for building high-performance kernels, including support for GPU acceleration and automatic differentiation.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Compute the maximum value of each row in the input tensor
Subtract the maximum value from each element in the row to center the values around zero
Compute the exponential of each centered value
Compute the sum of the exponentials
Compute the logarithm of the sum and add back the maximum value

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