FlashAttention: Skipping Masked Blocks
Intermediate
The problem, from first principles
Causal attention masks out every future position, so many key blocks are completely masked for a given query block. FlashAttention skips those blocks entirely, saving about half the compute in a causal model.
From theory to code
Implement causal_block_needed(qb, kb, bq, bk), deciding whether a key block must be computed.
Constraints
- Return a bool.
Hints
Hint 1
A key block is needed if its first key index is at most the last query index in the query block.
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.