Triton Fused Softmax Kernel

hardPROReductions & NormalizationTriton Programming

Problem Statement

Implement a numerically stable row-wise softmax over a (M, N) tensor, matching the official Triton fused-softmax tutorial. Each row fits in one block.

Background

Subtract the row max before exponentiating (x - tl.max(x)) to avoid overflow, then normalize by the sum of exponentials. Padded lanes load as -inf so exp gives 0.

Your Task

Implement softmax_kernel and run(M=64, N=300) comparing to torch.softmax(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 the softmax operation.
For each row in the tensor, we subtract the row maximum tl.max(x) from each element x to avoid overflow, resulting in x - tl.max(x), then apply the exponential function to get exp(x - tl.max(x)).
We normalize these exponentials by dividing by their sum, $\frac{exp(x - tl.max(x))}{\sum exp(x - tl.max(x))}$ , to obtain the softmax values for each row.
The run function compares the resulting Triton softmax output triton_out with the PyTorch reference output torch_reference using torch.allclose, returning True if the two outputs are close enough, which is the case for the given sample input.

Constraints

Subtract row max before exp (numerical stability)
Normalize by tl.sum of the exponentials
One program per row; mask with other=-inf

Learn this in

Kernel Fusion: Fused SoftmaxGPU 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 Autotuned Vector Addmedium

Concept primerPREVIEW

Background Knowledge

The softmax function is a common mathematical function used in machine learning models, particularly in the output layer of classification models. It takes a vector of real numbers as input and outputs a vector of values in the range (0, 1) that add up to 1. This is useful for representing a probability distribution over a set of classes. The softmax function is defined as $\frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}$ , where $x_i$ is the $i^{th}$ element of the input vector.

In the context of this problem, we are dealing with a (M, N) tensor, where each row represents a vector of real numbers. We need to apply the softmax function to each row independently. However, the naive implementation of the softmax function can be numerically unstable due to the possibility of overflow when computing the exponentials. To avoid this, we can use a trick called "log-sum-exp" or "subtraction of the maximum", where we subtract the maximum value of each row from all elements in that row before exponentiating. This ensures that the largest value is 0, which prevents overflow.

The Triton programming model is a Python-based programming model for writing high-performance GPU code. It allows developers to write custom GPU kernels using a Pythonic API. In this problem, we need to implement a Triton kernel that applies the softmax function to each row of a (M, N) tensor. We will also compare our implementation with the official PyTorch implementation of the softmax function.

Algorithm/Approach

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

Load the input tensor onto the GPU
Launch a Triton kernel that applies the softmax function to each row of the tensor
Compare the output of our implementation with the output of the official PyTorch implementation

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