FlashAttention: Merging Softmax Blocks
Advanced
The problem, from first principles
FlashAttention (Dao et al., 2022) computes exact attention one block of keys at a time, so the full score matrix never reaches slow memory. Each block contributes a running max and a running sum. Two such summaries merge exactly by rescaling to the larger max.
From theory to code
Implement online_softmax_merge(m1, l1, m2, l2), merging two blocks' softmax statistics exactly.
Constraints
- The result must equal the softmax normalizer computed over the union.
Hints
Hint 1
Take the larger max as the new max, then scale each sum by exp(old_max - new_max) before adding.
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.