Standard attention stores the full S×S score matrix in high-bandwidth memory (HBM). Here, S is the sequence length. Flash Attention processes small blocks in fast on-chip memory (SRAM). This reduces peak memory use.
Sequence Length
Head Dimension (d_head)
Block Size (Flash)
HBM Traffic
Standard: —
Flash: —
Reduction: —
Peak Memory (score matrix)
Standard: —
Flash: —
Standard Attention
HBM Traffic
—
Peak Memory
—
Flash Attention
HBM Traffic
—
Peak Memory
—
Block 0 / 0
Click a preset to change the settings.
Compare memory traffic at S=64 and S=4096. Flash's peak-memory advantage grows quadratically.
Change block size and watch how the number of tiles and SRAM usage changes.
At what sequence length does Flash Attention's memory reduction become large?
The FLOPs are identical. For some tile shapes, Flash Attention also transfers less data through HBM.
Try d_head=128, block_size=32. Flash uses more HBM traffic because it reads K and V blocks many times. Now try d_head=32, block_size=128. Fewer, larger tiles require fewer repeated reads.
The HBM traffic tradeoff depends on the ratio d_head/block_size. When d is large relative to block size, Flash re-reads more data per tile than standard writes to its score matrix.
Flash Attention uses less peak memory for every setting in this model. Does it always use less HBM traffic? Try different d_head and block_size values.
What happens if the block size is too large for SRAM? What if it is too small?
Why does standard attention's memory grow as O(S²) while Flash stays O(block_size²)?
When d_head > block_size, Flash's HBM traffic can exceed standard's. Why? (Hint: count how many times each K,V block is re-read from HBM.)