Triton Fused Softmax Kernel
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:
M = 64, N = 300
True
- The input values
M = 64andN = 300define 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 elementxto avoid overflow, resulting inx - tl.max(x), then apply the exponential function to getexp(x - tl.max(x)). - We normalize these exponentials by dividing by their sum, ∑exp(x−tl.max(x))exp(x−tl.max(x)), to obtain the softmax values for each row.
- The
runfunction compares the resulting Triton softmax outputtriton_outwith the PyTorch reference outputtorch_referenceusingtorch.allclose, returningTrueif 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
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 ∑j=1nexjexi, where xi is the ith 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
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
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.