Implement a numerically stable, fused row-wise softmax in Triton. For each row i and column j of the input matrix:
out[i,j]=∑k=0N−1exp(x[i,k]−maxkx[i,k])exp(x[i,j]−maxkx[i,k])The pattern is one program per row: each program loads its full row into registers, takes the row max, subtracts it for numerical stability, exponentiates, sums, and divides. The whole row pipeline happens in one kernel with one read of x and one write of out.
Input: x = [[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]]
Output: [[0.0900, 0.2447, 0.6652], [0.6652, 0.2447, 0.0900]]
Input: x = [[0.0, 0.0, 0.0, 0.0]]
Output: [[0.25, 0.25, 0.25, 0.25]]
One program per row. Use \texttt{tl.program_id} as the row index and address the row by row⋅stride.
Always max-subtract before exponentiating. Without it large positive logits overflow float32; with it the largest exponent argument is exactly 0.
For the row mask, load with other=-float(’inf’) so masked lanes lose the max and contribute zero to the sum after exp(−∞)=0.
torch.float32 on cuda and row-contiguous; output has the same shape.atol=1e-3. Max-subtract keeps the exponent in [−∞,0], so float32 holds.Accepts: array
Submit a correct solution to see your kernel performance.
Implement a numerically stable, fused row-wise softmax in Triton. For each row i and column j of the input matrix:
out[i,j]=∑k=0N−1exp(x[i,k]−maxkx[i,k])exp(x[i,j]−maxkx[i,k])The pattern is one program per row: each program loads its full row into registers, takes the row max, subtracts it for numerical stability, exponentiates, sums, and divides. The whole row pipeline happens in one kernel with one read of x and one write of out.
Input: x = [[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]]
Output: [[0.0900, 0.2447, 0.6652], [0.6652, 0.2447, 0.0900]]
Input: x = [[0.0, 0.0, 0.0, 0.0]]
Output: [[0.25, 0.25, 0.25, 0.25]]
One program per row. Use \texttt{tl.program_id} as the row index and address the row by row⋅stride.
Always max-subtract before exponentiating. Without it large positive logits overflow float32; with it the largest exponent argument is exactly 0.
For the row mask, load with other=-float(’inf’) so masked lanes lose the max and contribute zero to the sum after exp(−∞)=0.
torch.float32 on cuda and row-contiguous; output has the same shape.atol=1e-3. Max-subtract keeps the exponent in [−∞,0], so float32 holds.Accepts: array
Submit a correct solution to see your kernel performance.