Implement the forward pass of FlashAttention as a single fused Triton kernel. Given query, key, and value matrices Q, K, V all of shape (N,D):
out=softmax(DQK⊤)VThe standard three-pass formulation writes the full N×N attention matrix to HBM, dominating bandwidth for large N. FlashAttention tiles Q into row-blocks and streams over K, V in column-blocks, holding the per-row running max mi and running denominator li in registers. Each inner step rescales the accumulator oi when the running max changes, so the output is always normalized with respect to the maximum seen so far.
After the last K, V block, dividing oi by li produces the exact softmax-weighted output. This is the capstone of the section: it composes tile loops, online softmax, masking, and reduction primitives into one kernel.
Input: Q = [[1.0, 0.0], [0.0, 1.0]], K = [[1.0, 0.0], [0.0, 1.0]], V = [[1.0, 2.0], [3.0, 4.0]]
Output: [[1.6604, 2.6604], [2.3396, 3.3396]]
Input: Q = [[0.0, 0.0]], K = [[1.0, 1.0], [2.0, 2.0]], V = [[5.0, 5.0], [10.0, 10.0]]
Output: [[7.5, 7.5]]
Keep three accumulators per query row: running max (init to −∞), running denominator (init to 0), running output (init to 0). Update them together on each K, V block.
When a new block produces a higher max mnew, multiply both the running denominator and the running output by exp(m−mnew) before adding the block's contribution. The exponential of the difference is the rescaling factor.
After the loop divide the running output by the running denominator once and store. Both the last query tile and the last key tile may be partial, so the row and column masks both matter.
torch.float32 on cuda. No batch dimension, single head.atol=1e-2 + rtol=1e-2 * abs(expected) for N≥64 to absorb the drift from online softmax rescaling.Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Implement the forward pass of FlashAttention as a single fused Triton kernel. Given query, key, and value matrices Q, K, V all of shape (N,D):
out=softmax(DQK⊤)VThe standard three-pass formulation writes the full N×N attention matrix to HBM, dominating bandwidth for large N. FlashAttention tiles Q into row-blocks and streams over K, V in column-blocks, holding the per-row running max mi and running denominator li in registers. Each inner step rescales the accumulator oi when the running max changes, so the output is always normalized with respect to the maximum seen so far.
After the last K, V block, dividing oi by li produces the exact softmax-weighted output. This is the capstone of the section: it composes tile loops, online softmax, masking, and reduction primitives into one kernel.
Input: Q = [[1.0, 0.0], [0.0, 1.0]], K = [[1.0, 0.0], [0.0, 1.0]], V = [[1.0, 2.0], [3.0, 4.0]]
Output: [[1.6604, 2.6604], [2.3396, 3.3396]]
Input: Q = [[0.0, 0.0]], K = [[1.0, 1.0], [2.0, 2.0]], V = [[5.0, 5.0], [10.0, 10.0]]
Output: [[7.5, 7.5]]
Keep three accumulators per query row: running max (init to −∞), running denominator (init to 0), running output (init to 0). Update them together on each K, V block.
When a new block produces a higher max mnew, multiply both the running denominator and the running output by exp(m−mnew) before adding the block's contribution. The exponential of the difference is the rescaling factor.
After the loop divide the running output by the running denominator once and store. Both the last query tile and the last key tile may be partial, so the row and column masks both matter.
torch.float32 on cuda. No batch dimension, single head.atol=1e-2 + rtol=1e-2 * abs(expected) for N≥64 to absorb the drift from online softmax rescaling.Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.