Dao et al. (2023) — ARXIV-2307.08691
Key takeaways
- The diagnosis. FlashAttention reaches only 25–40% of theoretical peak FLOPs/s — far short of an optimised GEMM — and the cause is “suboptimal work partitioning between different thread blocks and warps”, producing low occupancy or unnecessary shared-memory traffic. Not the algorithm
- Three changes: reduce non-matmul FLOPs, which cost out of proportion because tensor cores run matmul far faster than anything else; parallelise across thread blocks even for a single head, adding the sequence-length dimension so occupancy survives a small batch × heads; and distribute work between warps inside a block to cut shared-memory communication
- ~2× over FlashAttention, reaching 50–73% of peak
- Forward-pass FLOPs are
4 · seqlen² · head_dim · n_heads; the backward pass runs 5 matmuls to the forward’s 2, because of recomputation - Block sizes are hand-tuned, not derived:
{64,128} × {64,128}, chosen per head dimension against device shared memory — “essentially only 4 choices”. This is a departure from FlashAttention’sB_c = ⌈M/4d⌉
Standing in the anthology
The source of SOTA-106.
Read in full for #114; the reading is NOTE-007. Two of the three practices drawn from the replaced takeaways are retired.
SOTA-107’s “multiple
of 128” attaches a real constant to the wrong quantity — every 128 here is a
head dimension or a block size — and SOTA-108’s vocabulary (padding,
block boundary, divisible) does not occur in the paper at all.
Both of those bodies contained correct reasoning about how blocked kernels behave. That is a distinct failure from the earlier clusters: not a takeaway belonging to another paper, but a sound inference written where a citation belongs.