FlashAttention-2: Deferred Normalization
Intermediate
The problem, from first principles
FlashAttention-2 (Dao, 2023) reduces non-matmul work by keeping the output unnormalized while streaming over key blocks, and dividing by the softmax denominator only once at the end.
From theory to code
Implement finalize_output(acc, l), dividing the accumulator by the denominator row by row.
Constraints
lhas one entry per query row.
Hints
Hint 1
Divide each row of the accumulator by its denominator, broadcasting over the feature axis.
Keywords
research-paperssystemsattentionkernels
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.