Source
Dao et al. (2022), LIT-074 — ARXIV-2205.14135.
What the tiling is actually constrained by
The block sizes exist to make one tile of Q, K and V, plus the running softmax statistics, fit in the streaming multiprocessor’s on-chip memory at once. That is the whole mechanism: if the working set fits, the score block is produced and consumed without ever touching HBM (SOTA-087); if it does not, the kernel spills and the advantage disappears.
So “match hardware SRAM size” is not a tuning heuristic. It is the condition under which the algorithm is the algorithm.
The formula, which the title only gestures at
LIT-074 does not tune the tile — it derives it. Step 1 of Algorithm 1,
with M the SRAM size and d the head dimension:
B_c = ⌈M / 4d⌉ B_r = min(⌈M / 4d⌉, d)
Q is then split into T_r = ⌈N/B_r⌉ blocks and K, V into
T_c = ⌈N/B_c⌉. The 4 is the four B×d tiles that must be SRAM-resident
at once, and the min on B_r keeps the query tile from exceeding what the
running softmax statistics need.
This is why “match the SRAM size” is a real constraint rather than a slogan:
the relationship is stated, and Theorem 2’s Θ(N²d²M⁻¹) bound is what you
get by respecting it.
The successor tunes rather than derives
Worth carrying, because it qualifies the formula above. FlashAttention-2
(LIT-106) does not compute the block size — it says: “we manually tune for
each head dimension since there are essentially only 4 choices”, over
{64,128} × {64,128}, bounded by device shared memory beyond which the
kernel will not run at all.
Both papers agree on the constraint and disagree on how to satisfy it.
LIT-074’s ⌈M/4d⌉ is the analytically right tile for its schedule; once
LIT-106 changed the work partitioning, the best tile stopped following from
SRAM capacity alone and became something to measure over a small discrete
set. The record should not say the size is derived without saying that the
current kernel picks it by hand.
What that means in practice
Not that anybody sets it by hand. The block sizes are chosen per architecture and per head dimension by the kernel, because SRAM per SM, head dimension and occupancy all move together — a tile large enough to be efficient and small enough that enough tiles are resident to hide latency. Hard-coding a size tuned for one generation is how a kernel gets slower on the next.
The practice is best read as a constraint to respect when writing or selecting a kernel, and as a reason a fused attention implementation is hardware-specific in a way an unfused one is not.