Compute a numerically stable row-wise log-sum-exp in Triton. For each row i of the input:
out[i]=logj=0∑N−1exp(x[i,j])Direct evaluation overflows for large logits. The stable form subtracts the row max before exponentiating and adds it back after the log:
out[i]=mi+logj=0∑N−1exp(x[i,j]−mi),mi=jmaxx[i,j]Pattern: one program per row, structurally identical to fused softmax but with m+log(∑exp(⋅)) instead of dividing by the sum.
Input: x = [[0.0, 0.0, 0.0, 0.0]]
Output: [1.3863] (which is log4)
Input: x = [[1.0, 2.0, 3.0], [0.0, 0.0, 0.0]]
Output: [3.4076, 1.0986]
Same structure as fused softmax for the reduction half: one program per row, load with other=-float(’inf’), compute row_max then ∑exp(row−row_max).
The final scalar is lse=row_max+log(s). Adding back the max is what restores the value lost during the stability shift.
The output is one scalar per row, so the store is a single scalar per program; no column mask needed on the write.
torch.float32 on cuda and row-contiguous; output has shape (M,).atol=1e-3. The max-subtract identity is exact up to float32; no large-N widening is required for these shapes.Accepts: array
Submit a correct solution to see your kernel performance.
Compute a numerically stable row-wise log-sum-exp in Triton. For each row i of the input:
out[i]=logj=0∑N−1exp(x[i,j])Direct evaluation overflows for large logits. The stable form subtracts the row max before exponentiating and adds it back after the log:
out[i]=mi+logj=0∑N−1exp(x[i,j]−mi),mi=jmaxx[i,j]Pattern: one program per row, structurally identical to fused softmax but with m+log(∑exp(⋅)) instead of dividing by the sum.
Input: x = [[0.0, 0.0, 0.0, 0.0]]
Output: [1.3863] (which is log4)
Input: x = [[1.0, 2.0, 3.0], [0.0, 0.0, 0.0]]
Output: [3.4076, 1.0986]
Same structure as fused softmax for the reduction half: one program per row, load with other=-float(’inf’), compute row_max then ∑exp(row−row_max).
The final scalar is lse=row_max+log(s). Adding back the max is what restores the value lost during the stability shift.
The output is one scalar per row, so the store is a single scalar per program; no column mask needed on the write.
torch.float32 on cuda and row-contiguous; output has shape (M,).atol=1e-3. The max-subtract identity is exact up to float32; no large-N widening is required for these shapes.Accepts: array
Submit a correct solution to see your kernel performance.