Implement row-wise softmax with the online (one-pass) algorithm. For each row i and column j:
out[i,j]=liexp(x[i,j]−mi)The standard implementation makes two passes over each row: one for the max, one for the normalized sum. The online formulation walks the row in chunks, maintaining a running max m and a running denominator l, updating both per chunk:
mnew=max(m,j∈chunkmaxxj),lnew=l⋅em−mnew+j∈chunk∑exj−mnewThis is the core trick behind FlashAttention's softmax, isolated as a self-contained kernel. Each program handles one row in a single forward pass over the chunks, then a second pass writes the normalized output.
Input: x = [[1.0, 2.0, 3.0]]
Output: [[0.0900, 0.2447, 0.6652]]
Input: x = [[0.0, 0.0], [1.0, 1.0]]
Output: [[0.5, 0.5], [0.5, 0.5]]
First pass: walk the row in chunks. For each chunk, compute mnew=max(m,mchunk), rescale l by exp(m−mnew), add the chunk contribution, then update m.
Second pass: walk the row again with the same chunk loop and write exp(chunk−m)/l at the chunk offsets.
Mask both passes with offsets < N. For the max pass, load with other=-float(’inf’) so padded lanes never drag the max down.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for N≥64 due to multi-stage reduction drift.Accepts: array
Submit a correct solution to see your kernel performance.
Implement row-wise softmax with the online (one-pass) algorithm. For each row i and column j:
out[i,j]=liexp(x[i,j]−mi)The standard implementation makes two passes over each row: one for the max, one for the normalized sum. The online formulation walks the row in chunks, maintaining a running max m and a running denominator l, updating both per chunk:
mnew=max(m,j∈chunkmaxxj),lnew=l⋅em−mnew+j∈chunk∑exj−mnewThis is the core trick behind FlashAttention's softmax, isolated as a self-contained kernel. Each program handles one row in a single forward pass over the chunks, then a second pass writes the normalized output.
Input: x = [[1.0, 2.0, 3.0]]
Output: [[0.0900, 0.2447, 0.6652]]
Input: x = [[0.0, 0.0], [1.0, 1.0]]
Output: [[0.5, 0.5], [0.5, 0.5]]
First pass: walk the row in chunks. For each chunk, compute mnew=max(m,mchunk), rescale l by exp(m−mnew), add the chunk contribution, then update m.
Second pass: walk the row again with the same chunk loop and write exp(chunk−m)/l at the chunk offsets.
Mask both passes with offsets < N. For the max pass, load with other=-float(’inf’) so padded lanes never drag the max down.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for N≥64 due to multi-stage reduction drift.Accepts: array
Submit a correct solution to see your kernel performance.