Source
Dao et al. (2022), LIT-074 — ARXIV-2205.14135.
introduced_by stays with LIT-074, and the reason is worth stating
because a predecessor now sits in the record. #342 flagged this as a likely
repeat of SOTA-192’s origin defect and asked for it to be verified first; it
is not one. That practice’s title was a generic technique credited to a paper
that adopted it. This one names an artifact — flash attention is Dao et
al.’s system, and the recommendation to run that kernel was first made by
them. The antecedent is LIT-655, which carries the memory-complexity
result and the condition below, and which is a source rather than an origin.
What “when hardware supports it” is carrying
The recommendation is easy because the result is exact. FlashAttention computes the same attention as the unfused implementation — not an approximation, not a sparsity pattern — so adopting it changes throughput and memory and nothing else about the model. That is unusual among the efficiency practices in this record and is why the title can be so unconditional.
Both the kernel and the recommendation to run it come from LIT-074. Its
Theorem 1 shows the tiled algorithm returns exact softmax attention with the
FLOP count unchanged at O(N²d) and only O(N) extra memory, and Proposition
3 proves no exact algorithm needs asymptotically fewer HBM accesses across SRAM
sizes from d to Nd. Measured on GPUs, that came to 15% over the MLPerf 1.1
BERT-large training record, 3× on GPT-2 and 2.4× on long-range arena.
It is the case SOTA-083 keeps once its framing is corrected: an operation no compiler covers well, hand-written because the algorithm is different and not just the schedule. That practice’s rule — hand-write the few critical operations where generated code falls materially short — is what this one applies to attention.
The condition in the title is doing real work, though. The kernel depends on enough on-chip memory per SM to hold a tile (SOTA-086), and on the head dimension and dtype being ones a compiled path exists for. Off that path the fallback is the unfused implementation, silently.
The memory saving and the speedup are different claims
Everything above is about whether the kernel runs. There is a second condition, about whether the benefit transfers, and it is measured in the work this line descends from.
Rabe and Staats (LIT-655) published the same algorithmic idea — exact
attention without materialising the N × N matrix — six months earlier, with
a JAX implementation on TPU. They get the memory result: 59× less overhead
at sequence length 16,384, 32× during differentiation. They do not get the
speedup:
is within a few percent of the runtime of the standard implementation of attention
and they say why: “standard self-attention already balances the available FLOPs and memory bandwidth of TPUs.”
So the two halves separate, and on hardware class rather than on kernel availability:
| what it depends on | where it holds | |
|---|---|---|
| less memory | the algorithm | anywhere it runs |
| more speed | HBM traffic being the bottleneck | measured on GPUs; near zero on TPU |
LIT-074’s 3× on GPT-2 is a GPU number, and Theorem 2’s HBM-access bound is
explicitly a claim about a memory hierarchy rather than about attention. Both
still hold. What does not follow, and what this practice implied by carrying
the two together, is that the speedup travels with the algorithm.
For a reader on GPUs nothing changes. For a reader choosing an accelerator, or
reading SOTA-086’s tile-size rule as general advice, it changes what to
expect.
Where the line goes from here
SOTA-106 supersedes the version of the kernel rather than the practice: FlashAttention-2 rebalances the work partitioning for the same exact result.
The one thing adopting it changes numerically is the order of accumulation, and that is not nothing at scale — SOTA-161 exists because the rounding bias of the fused reduction compounds rather than cancelling over a long run, and it is filed under stability rather than attention for exactly that reason.