Why architecture matters here
The architectural fact that everything rests on is that attention is memory-bound, and on small devices the memory wall is much closer than the compute wall. Modern accelerators can do far more arithmetic per second than they can feed data from main memory, so an algorithm's real cost is often the bytes it moves, not the flops it performs. The naive attention implementation moves the N-by-N score matrix to and from main memory several times, and that traffic — quadratic in sequence length — is what dominates the runtime and the energy. Optimizing attention therefore means optimizing memory access, which is a different problem from optimizing arithmetic.
Tiling is the standard answer to a memory-bound problem, and FlashAttention is tiling applied to attention. The full computation is broken into blocks small enough that the queries, keys, values, and intermediate scores for one block fit entirely in fast on-chip memory. Within a block you do a lot of arithmetic on data that is already close to the compute units, and you only touch main memory to stream in the next block of inputs and stream out the accumulated result. This turns a quadratic number of slow memory accesses into a linear number, which is the whole source of the speedup and the memory savings.
The online softmax is the clever piece that makes tiling possible, because softmax is normally a global operation over an entire row and tiling only gives you one block of that row at a time. FlashAttention computes softmax incrementally: as each block of scores arrives, it maintains a running maximum and a running sum, and rescales the accumulated output so that when the last block is processed the result is bit-for-bit the same as if the whole row had been softmaxed at once. This running-max, running-sum trick is what lets attention be computed block by block without ever holding the full row, and it is the mathematical heart of the method.
Fusion is the second efficiency lever. Naive attention is several separate kernels — score, softmax, weight — each reading and writing memory between them. FlashAttention fuses the whole sequence into a single kernel that keeps intermediates in registers and on-chip memory and never spills them to main memory. Fewer kernel launches and no intermediate spills mean less overhead and less bandwidth, which on a small device with limited memory and a simple scheduler is often as important as the asymptotic memory win. The output is accumulated in place as the blocks stream by.
On-device, a set of constraints reshapes all of this. Fast on-chip memory (SRAM) on a phone or NPU is measured in kilobytes, so block sizes must be small and chosen to fit; the device may use unified memory shared with the rest of the system, so the distinction between fast and slow memory is real but different from a discrete GPU; and the accelerator is often an NPU or mobile GPU with its own kernel programming model, so the fused kernel must be expressible on that hardware or approximated by whatever the runtime supports. The algorithm is the same, but the tile sizes, memory targets, and kernel implementation are all device-specific decisions rather than universal constants — which is precisely why deploying it well is an architecture problem and not a library import.
The architecture: every piece explained
The tiled Q/K/V blocks are the fundamental unit. Queries, keys, and values are partitioned into blocks sized so that a block of Q, a block of K, a block of V, and the intermediate score tile all fit together in on-chip memory at once. The outer loop walks query blocks; the inner loop streams key/value blocks past each query block. Block size is the master tuning knob: too large and it spills out of fast memory, defeating the purpose; too small and you pay excess loop overhead and under-utilize the compute units.
The online softmax state is carried per query block: a running maximum of the scores seen so far and a running sum of exponentials, plus the partially accumulated output. As each key/value block contributes its scores, the running max may increase, in which case the previously accumulated output and sum are rescaled to the new max so numerical stability is preserved without ever having seen all scores at once. This small piece of state is what replaces the full N-by-N matrix, and keeping it in registers is essential to the method's efficiency.
The fused accumulation is where the output is built. Rather than producing a full softmax matrix and then multiplying by V, FlashAttention multiplies each block of normalized scores by its corresponding V block and adds the result into the running output, rescaling as the softmax normalization evolves. The entire score-softmax-weight chain happens inside one kernel with intermediates living on-chip, so no attention matrix and no intermediate softmax ever reach main memory. This fusion is what converts the algorithmic idea into an actual bandwidth reduction on hardware.
The KV cache is the component that dominates the on-device decode phase. During autoregressive generation, each new token attends to all previous tokens, so the keys and values of past tokens are cached to avoid recomputation — and that cache grows with every generated token. On a memory-starved device the KV cache, not the attention matrix, often becomes the binding constraint, which is why on-device attention is usually paired with KV-cache techniques: paging, quantizing the cached keys/values to lower precision, or windowing so only recent tokens are retained.
Around these sits the device kernel and memory targeting: the actual implementation that maps tiling, online softmax, and fusion onto a specific NPU, mobile GPU, or CPU vector unit, choosing block sizes for that chip's SRAM and using its supported instructions. On a discrete GPU this is a hand-tuned CUDA kernel; on a phone it may be a vendor NPU kernel, a Metal or Vulkan compute shader, or a fallback. Getting the memory targeting right — what counts as fast memory, how big the tiles can be — is what determines whether the theoretical win shows up in practice. The diagram shows how tiles, online softmax, fused accumulation, and the KV cache relate during a forward pass.
End-to-end flow
Consider a prefill pass over a prompt of length N on an on-device SLM. The runtime does not allocate an N-by-N score buffer. Instead it partitions the queries into blocks and, for the first query block, begins streaming key/value blocks through fast memory. For each key/value block it computes the block of scores, updates the running max and sum, and folds the weighted values into the query block's accumulating output.
As successive key/value blocks arrive, the online softmax state evolves: whenever a new block contains a larger score than any seen before, the accumulated output and normalization sum are rescaled so the final result matches a full-row softmax exactly. When the last key/value block for this query block has been processed, the query block's output rows are complete and correct, and they are written out once — the only main-memory write for that block's result. The full attention matrix for these rows was never materialized anywhere.
The runtime then advances to the next query block and repeats, reusing the same on-chip buffers. Because the working set is one block's worth of Q, K, V, and softmax state — not the whole sequence — memory usage is linear in N rather than quadratic, and main-memory traffic is dominated by streaming the inputs once rather than by repeatedly reading and writing an N-by-N matrix. On a bandwidth-bound device this is exactly where the wall-clock and battery savings come from.
During the decode phase the flow shifts to one query at a time. Each newly generated token produces a single query that must attend to all previous tokens' keys and values, which are read from the KV cache. FlashAttention-style tiling still applies — the single query streams past blocks of cached K/V with an online softmax — but now the dominant memory cost is the KV cache itself, which grew by one token's worth of K and V on the previous step. The runtime's job becomes keeping that cache affordable: paging it, quantizing it, or windowing it so it fits the device's memory budget as generation continues.
Throughout both phases the same principle holds: fast on-chip memory holds the active working set, main memory is touched a linear number of times, and the N-by-N matrix never exists. The end result the user experiences is a small model that can accept a longer prompt and generate within the tight memory and bandwidth envelope of a phone or embedded accelerator — context lengths and responsiveness that naive attention simply could not reach on that hardware. The correctness is identical to standard attention; only the memory choreography changed, and on constrained devices that choreography is the whole ballgame.