FLAT: An Optimized Dataflow for Mitigating Attention Bottlenecks

The attention score matrix grows with the square of sequence length and drowns accelerators in memory traffic — FLAT computes it in row tiles, runs softmax on-chip, and multiplies straight into the value matrix, so the quadratic tensor never leaves the chip and 64K-token transformers run on ordinary accelerators.

Why this problem is worth solving

An attention layer measures how strongly each token in a sequence relates to every other token. For a sequence of N tokens it computes an N×N matrix of scores — that pairwise, global view is exactly what makes transformers good, and it is also the scaling disaster. The paper’s motivating applications need long sequences because meaning emerges only from full context: two proteins can look identical fragment-by-fragment and differ only over the whole chain, and a summary can change meaning when the full document is considered. Demand for 8K to 69K token inputs already existed when the paper was written; 64K tokens is roughly the length of a short novel.

The quadratic term turns concrete fast. The intermediate score tensor for BERT is about 10 MB at 512 tokens, 136 MB at 2K tokens, and 8.2 GB at 16K tokens. For comparison, a cloud accelerator in the paper’s study (modeled on the TPU-v3) has 32 MB of on-chip memory, and an edge accelerator has 512 KB. At long sequence lengths the score matrix is not merely too big to keep on-chip — it is 250× too big, so it must round-trip through off-chip DRAM, whose bandwidth is roughly 20× lower than on-chip memory in both accelerator classes studied (50 GB/s vs 1 TB/s on edge; 400 GB/s vs 8 TB/s on cloud).

Hardware designers’ previous answer was to change the math — sparsify, compress, or approximate attention — trading model accuracy for speed. FLAT’s claim is that the bottleneck can be removed without touching the functionality at all: same numbers, same accuracy, different movement of data.

What people did before, and where it broke

To see why standard tools fail, split an attention layer’s six operators into two families.

Activation-weight operators (Q, K, V, O). These multiply the input by learned weight matrices — ordinary fully connected (FC) computations. Their

Activation-activation operators (L and A). The Logit operator (L) multiplies queries by keys to produce the N×N scores; the Attend operator (A) multiplies the normalized scores by values. Both inputs of each are freshly computed activations, not weights. Every new sequence brings entirely new operands, so batching adds zero reuse — the paper’s roofline analysis on a TPU-v3 shows L and A pinned deep in the memory-bound region at any batch size, while FC and convolution operators climb toward compute-bound as the batch grows. These two operators dominate attention latency and energy, and get worse as N grows.

Why not just fuse them? Operator fusion — feeding a producer’s output straight into a consumer without materializing the tensor in memory — is the classic cure. But ML compilers only fuse matrix operators with element-wise operators (ReLU, add). Between L and A sits softmax, a many-to-many operator: each output row requires reading all N scores of that row, reducing them, then scaling each one. Arbitrary tiling breaks that dependency. DNNFusion, a state-of-the-art fusion study from 2021, reported tensor-tensor fusion as “too complicated or unprofitable.” Meanwhile attention accelerators like A³, ELSA, and Sanger (see also the DOTA explainer elsewhere on this site) sidestepped the problem by pruning or approximating scores — buying speed with accuracy.

The idea

FLAT stands for Fused Logit Attention Tiling, and the name is the design: fuse L–softmax–A into one loop nest, and tile it at a granularity chosen to fit the on-chip memory.

Fuse exactly these two, and nothing else. The paper measures the operational intensity of every plausible fusion pair. Fusing L with A lifts the combined operator’s intensity above the accelerator’s compute-bound threshold — the quadratic intermediate tensor becomes internal traffic instead of memory traffic. Fusing an FC into the pair (Q with L, or A with O) makes things worse: the FC side loses weight reuse and drops below its stand-alone performance. FCs are already compute-bound; they need no help. So FLAT draws the fusion boundary tightly around the sick operators.

Respect softmax by making the row the atom. Softmax reduces across a full row of N scores, so the smallest legal unit of fused work is one complete row — computed from one query vector against all N keys. FLAT calls this row granularity and only ever tiles in multiples of it. That single constraint preserves exact data dependencies — no approximation, no online correction, bit-identical results.

Scale the tile to the memory you have. Above the row, FLAT builds coarser tiles: multiple rows (R-Gran), whole heads (H-Gran), whole sequences (B-Gran), or everything at once (M-Gran — the naive keep-it-all-on-chip approach that dies beyond 2K tokens). Here is the subtle enabling fact: for FC-FC fusion, shrinking tiles destroys weight reuse, so fine granularity is self-defeating. But L and A have no algorithmic reuse to lose — every score is used once. So FLAT can shrink tiles almost freely to fit a 20 KB edge buffer or grow them to exploit a 32 MB cloud buffer, paying only a modest cost of re-streaming keys and values per tile.

Search, don’t guess. Fusion adds new choices — shared outer-loop order, per-operator inner loops, tile sizes in three dimensions — on top of classic per-operator mapping. FLAT ships a map-space exploration framework that exhaustively searches loop orders and tile sizes against an analytical cost model of the accelerator (validated within 1% of the established MAESTRO model), so each hardware target gets its own optimal mapping.

traffic: O(N²) written, read, written, read
at N=16K: 8.2 GB × multiple trips over ~50–400 GB/s

Why it works

• Both operands of L and A are fresh activations → batching cannot add reuse → the only reuse left to create is temporal locality of the intermediate tensor → fusion is the uniquely-fitting lever.
• Softmax needs whole rows → FLAT makes the row the smallest tile → the dependency is satisfied structurally → fusion becomes legal with bit-identical results, where compilers had declared it impractical.
• Score tiles die on-chip → the N×N tensor generates no off-chip traffic → the fused operator’s arithmetic-per-byte rises above the compute-bound threshold → the processing array, previously starved, approaches full utilization (0.99 vs 0.55 for the best unfused dataflow at 20 KB in the paper’s sweep).
• L/A have no weight reuse to lose → tiles can shrink to fit any buffer → the same dataflow serves a 20 KB edge chip and a 2 GB research configuration, and footprint linear in N unlocks 64K-token sequences.
• Off-chip bandwidth demand falls 71–82% → the accelerator no longer competes with itself for the memory bus → performance survives on cheap, bandwidth-poor hardware.

What it costs

• Re-streaming. Keys and values are fetched again for every tile; finer tiles mean more re-fetches. This is the price of fitting small buffers, and the search framework balances it against footprint.
• Hardware assumptions. The accelerator needs a software-addressable scratchpad and a controller able to follow the fused schedule — common in modern designs, absent in fixed-function pipelines.
• Offline search. The exhaustive map-space exploration must be run per accelerator and per model shape before deployment.
• No help for the rest of the model. Q/K/V/O and feed-forward layers are already compute-bound; FLAT leaves their performance unchanged, so end-to-end gains (1.5× at 512 tokens) are smaller than L/A-only gains.
• Simulation-first evidence. Headline numbers come from an analytical cost model (validated within 1% of MAESTRO on single operators), not silicon — though a real-GPU implementation on a Tesla T4 corroborates the memory-footprint claims.

The evaluation

FLAT is evaluated on five transformer models (BERT, TrXL, FlauBERT, T5, XLM) across two accelerator classes — an edge design (32×32 processing elements, 512 KB buffer, 50 GB/s off-chip) and a cloud design modeled on TPU-v3 (256×256 elements, 32 MB, 400 GB/s) — against the strongest possible software baseline: FLEX-Opt, an exhaustive search over all existing per-operator dataflows, so any win is attributable to fusion rather than to better tuning.

A buffer-size sweep sharpens the story: with a luxurious 2 GB of on-chip memory, FLEX-Opt and FLAT-Opt both reach ~0.97–0.99 utilization — fusion is unnecessary when everything fits. At a realistic 200 KB, FLAT-Opt holds 0.98 while FLEX-Opt manages 0.58 (a 1.7× gap on L/A); at a scarce 20 KB, the gap is 0.55 vs 0.25. FLAT is precisely a technique for the memory-poor regime that real chips live in.