Inference Engineering · Part iii of vi

Kernels and Memory in LLM Inference

FlashAttention through four generations, paged KV, GQA and MLA, flash-decoding, CUDA graphs. Same attention math. Different routes for the bytes.

Talal Ahmed20 min read

The last article ended at the machine.

Streaming multiprocessors. Tensor cores. The four-level ladder of memory that feeds them.

This one is about the code that runs on that machine.

That code comes in units called kernels.

A kernel is one small program. The CPU hands it to the GPU as a single piece of work: multiply two matrices, normalize the rows of one, run attention for a head. The GPU then runs that one program across thousands of threads at once.

Everything the chip ever does happens this way. When PyTorch multiplies two matrices, what actually runs on the silicon is a kernel.

A kernel is one program, launched once

CPU hands the GPU a single job · thousands of threads run that same program

CPUlaunchGPU · one kernelsame math · different data movement · different wall clock
CPU launching one kernel so the GPU runs the same program across thousands of threads

Here is the itch that makes kernels worth studying as their own subject.

Two kernels can compute the same attention on the same GPU, with identical inputs and identical outputs. One of them can finish up to 7.6× faster than the other.

The arithmatic cannot be the difference. Both kernels perform exactly the same operations.

What seperates them is everything the mathematics leaves open: which threads load which bytes, from which level of memory, in what order, at what precision, and where each result is written.

Every one of those is a choice about moving data.

The whole gap lives there.


What we will follow

We will follow one lineage of kernels, FlashAttention, through four generations.

It computes attention, the operation at the heart of the transformer. Once context grows long, attention's memory cost grows with it, and it becomes a dominant draw on memory bandwidth in the forward pass. That is what makes it one of the kernels that matter most for inference.

Across all four generations, the function never changes. Same scores. Same softmax. Same weighted sum.

What changes is the fraction of the chip's peak arithmetic the kernel actually delivers. It starts low, tensor cores mostly waiting on memory, and climbs toward roughly three quarters of peak.

A gap that wide, with identical mathematics, has to come from somewhere physical. From where the bytes sit.

After FlashAttention come the other levers of serving:

  • a model built to keep a smaller cache
  • a cache packed into memory without waste
  • decode work split, so the chip is not idle on one token
  • the launch of each kernel paid for once
  • every number carried in fewer bits

Two tiers decide the craft

Two tiers of the memory ladder from last time decide almost everything that follows.

HBM. Large. Slow relative to the math units. An H100 holds about 80 GB and streams it at 3.35 TB/s. For a chip this hungry, even that counts as slow.

SRAM. Up to ~228 KB per SM, and the chip has 132 of those. Sits right beside the arithmetic units. Roughly 10× faster than HBM.

Compute throughput has outgrown memory bandwidth for generations. Moving a byte across that gap costs more than doing arithmetic on it once it arrives.

So the whole craft of writing a kernel reduces to one habit:

Keep the bytes out of HBM and next to the math.

Which side a kernel lands on, memory or compute, is still decided by arithmetic intensity, and the roofline we already built.

On the H100 the ridge sits near ~295 ops/byte.

A large dense matmul does a lot of arithmetic on every byte it loads. It sits past the ridge. Compute-bound. Tensor cores set the pace.

Attention over long sequences is the opposite. Very little arithmetic per byte touched. Deep on the slanted ceiling. Memory-bound. Capped by bandwidth, no matter how fast the math units are.

For a kernel author, the roofline is a lever. The function fixes how many operations there are. Nothing fixes how many bytes move.

Cut HBM traffic and intensity rises. The same computation slides right, out from under the bandwidth ceiling.

That is the craft. Its measure is the fraction of peak arithmetic a kernel sustains. That is the number we will track.


Attention, said plainly again

Each token builds three things: a query (what it is looking for), a key (what it offers), a value (what it contributes).

Stack them over nn tokens and you get QQ, KK, VV.

  1. Compare every query to every key: S=QK⊤S = QK^{\top}. An n×nn \times n score matrix.
  2. Softmax each row into blending weights. Exponentiate, then normalize. In practice, subtract the row max first so the exp does not overflow.
  3. Multiply the weights by VV. One output row per token.

So attention is two matrix multiplies with a normalization between them.

It also runs as many independent heads. Everything below happens inside one head. Follow one and you follow them all.

The object in the middle is the problem

SS holds one score per pair of tokens. Double the sequence length and its size quadruples.

  • ~1k tokens → about a million entries
  • 8k tokens → ~67 million
  • 32k tokens → about a billion

A naive kernel builds every one of those entries and makes three seperate passes through HBM:

  1. compute SS and write it out in full
  2. read SS back, softmax row by row, write the weights
  3. read the weights again, multiply by VV

The arithmetic is mostly matmul. Exactly what tensor cores want.

The bottleneck is the traffic. Bytes grow as n2n^{2}. Work per byte stays small. On long sequences attention lands deep in the memory-bound corner. Bus saturated. Tensor cores waiting.

Naive attention: three HBM round trips

Write S · read S for softmax · read again for V · n² traffic every time

1. compute S = QKᵀ · write to HBM2. read S · softmax · write weights3. read weights · multiply by VHBMn × nparked twice
Naive attention writing and rereading the full n by n score matrix through HBM three times

The cure that the cost model points at

SS is scratch. Produced once. Consumed once. Then discarded.

Build it on chip. Never write it to HBM. The n2n^{2} traffic vanishes.

That cure has a name: FlashAttention.


FlashAttention-1: change the accounting

FlashAttention was published in 2022 by Tri Dao and collaborators.

It begins with a change of accounting rather than a new algorithm.

Attention is memory-bound. Wall-clock time is set by bytes moved, not operations performed. So the paper changes what it counts.

The cost model becomes the reads and writes that cross between HBM and SRAM. The authors call this IO awareness.

Under that accounting, the three-pass implementation is expensive for one reason. It parks two full n×nn \times n matrices in HBM, then pays the bus to read each one back.

SRAM is the cheap place to work. But SRAM holds only a couple hundred kilobytes per SM. The full score matrix will never fit at any realistic sequence length.

What rescues the plan is one property of matrix multiplication: it decomposes cleanly into blocks.

Tiling

FlashAttention cuts QQ, KK, and VV into blocks that fit in SRAM.

The kernel loads one block of KK and one block of VV on chip. For each block of QQ, it computes that tile of attention right there, folds the partial result into the output, then lets the next blocks arrive.

Each tile of SS lives briefly in SRAM, gets used, and gets overwritten. The full n×nn \times n matrix is built tile by tile on chip and never reaches HBM.

This is still exact attention. It does not skip token pairs. The output matches the three-pass version, identical up to the reordering of floating-point sums.

The function did not change. The route the bytes take did.

FlashAttention: tile in SRAM, never park S

Load K/V blocks · stream Q tiles · each score tile lives briefly on chip

Q tilesSRAM · on chipK blockV blockS tileused · overwritten · never HBMHBMno n×n dump
FlashAttention tiling with K and V blocks and score tiles living briefly in SRAM instead of HBM

The obstacle: softmax wants the whole row

Softmax still wants the whole row: the row maximum and the sum of exponentials.

A tile only ever sees a slice of a row. As softmax is normally written, nothing in the row can be finished until the whole row exists.

The fix predates FlashAttention by four years. Online softmax, NVIDIA, 2018 (Milakov and Gimelshein). Same result in a single streaming pass, by carrying two running values:

  1. a running maximum
  2. a running denominator (sum of exponentials so far)

Watch what happens when the maximum changes part way through.

First tile arrives. Largest score is 3. Exponentials use es−3e^{s-3}. Denominator accumulates to something like 1.59.

Next tile arrives carrying a 5. Running max updates to 5.

But every exponential already folded into that 1.59 subtracted the old maximum. They are all too large by the same factor, e5−3e^{5-3}.

That shared error is what saves the algorithm. One multiplication repairs the whole denominator at once. Scale 1.59 by e3−5e^{3-5}. Every accumulated term lands on the new scale.

FlashAttention pushes the same correction one step further. The running output is built from those same exponentials, so it drifts off scale the same way. Whenever the denominator rescales, the output rescales with it.

By the last tile, the running max is the true row max, and the running denominator is the true row sum. Exact attention. Becuase the correction was exact at every step.

Attention never needed the whole row at once. Only two running numbers per row.

That is what makes tiling legitimate.

Online softmax: two running numbers

Running max + running sum · when the max jumps, rescale everything at once

tiles of one row arriving over timemax=3max→5rescaleexactrunning max · running Σ expshared error → one multiply repairs the whole denominator
Online softmax updating a running maximum and denominator across tiles and rescaling when the max jumps

Training detour (brief)

Inference only runs the forward pass. Training also needs the backward pass, and the backward sweep wants the attention matrix a second time.

FlashAttention never stored it. It kept just enough to rebuild it: those two running numbers per row.

During backward, the kernel reloads QQ, KK, VV tile by tile and recomputes each attention tile on chip at the moment it is needed. Extra multiplies on data already in SRAM. The HBM traffic they avoid is the expensive kind.

Attention's memory footprint drops from quadratic in sequence length to linear.

Principle the whole kernel runs on:

Never spend HBM traffic on anything you can rebuild from what is already on chip.

On GPT-2, FlashAttention-1 ran attention up to 7.6× faster than standard PyTorch. That is the figure this article opened with.

Judged on utilization though, this kernel still sustains only ~25–40% of peak arithmetic.

The memory wall is handled. Bytes take the short route. And yet the tensor cores still sit idle most of the time for reasons that have nothing to do with HBM.

The next three generations exist becuase of that number.


FlashAttention-2: inventory the waits (2023)

Tri Dao's follow-up reads as an inventory of places where the chip was still waiting. Three of them.

1. Bookkeeping between matmuls

Online softmax does rescales and divisions between the matrix multiplies. That bookkeeping runs on CUDA cores. While it runs, tensor cores hold.

On an A100, tensor cores deliver ~312 TFLOP/s of matmul. The scalar FP path delivers ~19.5. Divide: about 16×. Every scalar op the kernel cannot avoid costs roughly 16 matmul-ops worth of machine time.

FA2 rearranges the bookkeeping so less of it stands between one matmul and the next. The output rides through the loop unnormalized. The division by the denominator happens once at the end. Result unchanged. Waiting shrinks.

2. Whole SMs with no work

A kernel is handed out as thread blocks. Each block lands on one SM and runs to completion there. The number of blocks is the number of SMs you can keep busy.

FA1 created one block per head per sequence in the batch. Serve a couple of long sequences and that product is small. Batch of 2 with 16 heads → 32 blocks on a chip with over a hundred SMs. The rest sit idle.

FA2 adds a third source of blocks: it splits the query rows of a single long sequence across many of them. One sequence alone can fan out and occupy the whole chip.

FlashAttention-2: fill the idle SMs

Split query rows so one long sequence can occupy the whole chip

FA1 · few blocksmost SMs idleFA2 · split Qone sequence fans out across the chip
FlashAttention-1 leaving most SMs idle versus FlashAttention-2 splitting query rows to fill the chip

3. Warp exchange inside a block

FA1 divided KK and VV across a block's warps. Each warp held partial results it had to exchange through shared memory.

FA2 divides QQ instead. Each warp carries its own rows all the way to the output. The exchange largely disappears.

Three waits. Three rearrangements. Utilization roughly doubles: ~50–73% of peak on the A100. End-to-end training numbers hit model FLOP utilization around 72% on that generation.

Then Hopper shipped. The kernel's assumptions did not move with the machine.

On the H100, FA2 fell to about ~35% of peak. Hopper had new features. FA2 used neither.


FlashAttention-3: overlap the two jobs (2024)

Hopper brought two features that matter here.

WGMMA. A tensor-core instruction issued by a warp group (four warps acting as one). It runs asynchronously: start a matmul, let the warps move on while it runs.

TMA. The Tensor Memory Accelerator from last article. A dedicated copy engine. Moves tiles between HBM and SM fast memory on its own. No threads required for the copy.

Put together, the chip can do its two jobs at the same time: move data and do math.

FA2 did them strictly in turn. Load a tile. Compute. Load the next. Copy engine idle during math. Math idle during copies. On hardware built so neither has to wait.

FA3 reorganizes around that idle time with warp specialization.

  • Producer warps do nothing but issue TMA loads
  • Consumer warps do nothing but compute on tiles that have already landed

Producers run ahead. Consumers work through what has arrived. Data movement and computation finally overlap.

FA3: producers load, consumers compute

TMA warps pull tiles · compute warps stay busy · copies and math overlap

Producer warpsissue TMA loads onlyConsumer warpsMMA + softmax on landed tiles
FlashAttention-3 producer warps issuing TMA loads while consumer warps compute on tiles already in SRAM

The same overlap continues one level down. Tensor cores run one tile's output multiply (AVAV) while the next tile's exponentials run on the special function unit (SFU / MuFu), a third unit that is neither tensor core nor CUDA core.

Precision: FP8 with a real failure mode

Hopper tensor cores can compute in FP8, roughly double the arithmetic rate of 16-bit. But 8 bits only express a couple hundred distinct values.

Using them means quantizing. That goes wrong in a specific way for transformers: values mostly huddle near zero, with a handful of outliers far out. One global scale stretches to reach the outliers, and the crowded majority near zero collapses onto a few levels.

FA3 answers twice:

  1. Per-block scales, so one region's outlier does not stretch everybody's grid
  2. Incoherent processing: multiply QQ and KK by a fixed random rotation before quantizing. Loud coordinates smear across dimensions. Dot products are unchanged becuase the rotations cancel inside QK⊤QK^{\top}

That combination cut numerical error by about 2.6× in the paper's measurements.

Payoff: about 1.5–2× over FA2. In FP16, up to ~740 TFLOP/s (~¾ of H100 FP16 peak). In FP8, climbing toward ~1.2 PFLOP/s.

Then Blackwell moved the bottleneck again.


FlashAttention-4: when the roof moves (2026)

From H100 to B200, BF16 tensor-core throughput rises by about 2.25×.

The special function units that compute the exponential do not grow at that rate. Neither does bandwidth into the SM's shared memory.

When one engine gets 2.25× faster and its neighbors stay fixed, the bottleneck does not disappear. It moves.

Tri Dao's group calls this asymmetric hardware scaling. Raising tensor-core utilization is no longer the main lever.

FA4 answers with three algorithm changes and one change of tools.

1. Split the exponential

On Blackwell, the multi-function unit delivers on the order of ~16 exponentials per cycle per SM, while BF16 tensor cores deliver thousands of ops per cycle beside it. Softmax production cannot keep up with softmax consumption.

FA4 therefore computes exponentials two ways at once:

  • the hardware SFU keeps running
  • a short polynomial runs on CUDA cores that would otherwise sit idle, approximating the exponential to the accuracy the kernel needs

2. Conditional rescales

Each rise in the running maximum used to force a rescale. FA4 sets a threshold τ\tau. The running max may drift by up to ~8 doublings (factor of 256) before a rescale fires. Smaller jumps ride as bounded drift. At the end, everything renormalizes once by the true final max and denominator. Answer unchanged. Roughly ~10× fewer rescales in Dao's Hot Chips telling.

3. Backward pass and TMEM

Blackwell adds tensor memory (TMEM), ~256 KB per SM wired into the tensor cores. A place to keep data, not just a copy engine.

FA4 pairs neighboring SMs in a 2-CTA mode so two cooperating blocks drive one matmul and share one copy of the data. Backward intermediates park in TMEM instead of crowding shared memory. Shared-memory traffic roughly halves for that path.

Tools: CuTe DSL

FA4 is written in CuTe DSL, NVIDIA's Python-embedded kernel language. Compiles ~20–30× faster than equivalent C++ templates. At this level you tune by rebuilding and measuring hundreds of times. Compiler speed becomes experiment speed.

Gains are real but regime specific. The paper claims the win for sequences of ~4k tokens and longer. Below that, not a consistent lead over every baseline.

Utilization lands around ~71% of B200 BF16 peak. Set beside FA3's ~75% on Hopper and the fraction has plateaued near three quarters, even as the ceiling underneath grew 2.25×.

Same attention. Climbing utilization.

FA1 ~25–40% · FA2 ~50–73% · FA3 ~75% · FA4 ~71% of a much taller roof

FA125–40%FA250–73%FA3~75%FA4~71%
FlashAttention generations climbing peak utilization from FA1 through FA4

Every headline number above was measured in BF16. The width of the numbers is a lever of its own. We will leave that for the next piece: quantization as a discipline.

Before that, a different byte problem. Kernel intermediates were never the only thing in HBM.


The KV cache was wasting half the machine

Every active request also keeps its keys and values in HBM.

The way serving systems laid that cache out turned out to waste more than half of the most expensive memory in the machine.

The reasonable-looking mistake

A serving system cannot know in advance how long a response will run. So systems before ~2023 played it safe.

For every request, reserve one contiguous stretch of HBM sized for the longest sequence that request might reach.

Limit 4096. Response stops at 200. The other 3896 slots sit reserved and empty. That waste inside a reservation is internal fragmentation.

There is a second kind. As requests of different lengths start and finish, free memory between reservations turns ragged. A new request can fit in the total free space and still find no single contiguous stretch large enough. External fragmentation.

OS textbooks have described both since the 1960s.

The vLLM team measured what the two cost together. In the serving systems of the time, only ~20–38% of KV cache memory held actual token states. More than half stored nothing.

PagedAttention (2023)

Operating systems also solved this decades ago. PagedAttention carries that solution over almost unchanged.

Break the cache into fixed-size blocks (16 or 32 tokens). Give every request a block table that maps logical blocks to physical block IDs.

Physically, a request's blocks can sit anywhere in HBM. Logically, the attention kernel follows the table and sees one contiguous sequence.

Internal fragmentation shrinks to at most one partial block per request. External fragmentation disappears, becuase every free block is the same size.

PagedAttention: blocks, not giant reservations

Logical sequence via a block table · physical blocks can sit anywhere

Contiguous reservationusedmostly empty · reservedPaged blocksblock table maps logical → physical
Contiguous KV reservation wasting most of HBM versus paged blocks mapped through a block table

The attention kernel still computes the same function. It just gathers each block by following the table.

Reclaimed memory turns directly into throughput. More concurrent requests fit in the same HBM. vLLM served on the order of 2–4× the throughput of prior systems at similar latency in their measurements.

The block table buys one more thing: sharing. System prompts, instruction preambles, beam-search candidates. Tables can point at the same physical blocks. Copy on write only when a sequence diverges. On beam search, shared portions reached up to ~55% of the cache in their numbers.

Paging controls where the cached bytes sit.

How many of them there need to be is a seperate lever. And that one is architectural.


Store less: MQA, GQA, MLA

In standard multi-head attention, every head carries its own keys and values. Cache stores two vectors per head per layer for every token.

For Llama-3-70B-class configs that lands around ~320 KB per token already with grouping. Without grouping, many times larger.

Multi-query attention (2019)

Noam Shazeer: generating each token means reloading all those cached K/V from HBM. Bandwidth-bound. Most of those bytes exist only becuase every head insists on a private copy.

MQA: keep all the query heads, but share one key head and one value head. Head-count term in the cache formula drops to 1. Cache shrinks by the full head count. Bandwidth bill shrinks with it.

Price: quality. 64 query heads trained to ask 64 different questions now consult the same single description of each token. Accuracy drops measurably.

Grouped-query attention (2023)

Turn the all-or-nothing choice into a dial. Split query heads into groups. Give each group its own K/V head.

One group = MQA. One group per head = full MHA. Useful settings sit in between. Around 8 groups often holds quality close to multi-head at close to multi-query speed.

Existing multi-head models can be converted for about ~5% of original pretraining compute ("uptraining"). Llama 2/3 and Mistral adopted this almost immediately.

Multi-head latent attention (DeepSeek V2, 2024)

Sharing heads is one way to store less. DeepSeek took another: compress what each token stores.

MLA does not keep full K/V in the cache. For each token it stores one short latent, a compressed summary the model learns to write. Keys and values can be rebuilt when attention needs them.

In DeepSeek V2: 128 heads × 128 dims → 32,768 elements per layer per token for full K/V. The latent stores 512, plus a small 64-element companion for position → 576 stored instead of 32,768. About ~57× smaller.

Arithmetic is cheap. HBM bytes are expensive. Decode leaves tensor cores mostly idle anyway. MLA spends that idle arithmetic to buy back cache space.

In practice even the rebuild multiplies largely vanish: two matmuls in a row fold into one. Decompression matrices absorb into the query/output projections. Position is the awkward part (RoPE cannot merge into a fixed matrix), so a small rotated channel rides beside the latent and keeps the merge legal.

Measured against their prior 67B model: KV cache down ~93%, max generation throughput up ~5.76× in their reporting. Quality claims are theirs, inside their recipe. Read carefully.

Store less: share heads or compress

GQA dials head sharing · MLA stores a latent and rebuilds K/V when needed

GQAKV1KV2KV3groups of queries share KVMLAlatent · 512+64rebuild K/V when attending
Grouped-query attention sharing KV heads versus multi-head latent attention storing a compressed latent

Paged, grouped, compressed. The cache is a fraction of its former size.

Shrinking the cache does not change how the cache is read.

And at decode, that reading has a problem of its own.


Flash-Decoding: one token, whole chip

FlashAttention-2 fills a chip by handing out many independent blocks of work. Prefill produces thousands. Decode produces almost none.

At decode the query is one new token. One row. The query dimension contributes exactly one block. Block count collapses to roughly batch × heads.

Serve one request on a 32-head model: 32 pieces of work. A hundred SMs stand idle. Each busy one walks the entire cache alone, tile after tile. Longer context, longer serial walk.

The work is all in one dimension nobody had split: the cache itself.

Flash-Decoding (2023) splits cached K/V into chunks. Every chunk gets its own block, which means its own SM. Each computes the new token's attention against its own chunk.

Softmax should forbid that. Normalization couples every position to every other.

What makes the split legal is the same online-softmax bookkeeping. Each chunk keeps its own running max and sum. A small second kernel merges the partials, rescaling each by how far its local max sits below the global one. Exact attention.

Online softmax ran that bookkeeping serially down one row. Flash-Decoding runs it in parallel across chunks of the cache.

Flash-Decoding: split the cache, not just Q

One decode token · many chunks · merge with online-softmax bookkeeping

qone new tokenchunk 1m · ℓchunk 2m · ℓchunk 3m · ℓchunk 4m · ℓmerge kernel rescales partials → exact attention
Flash-Decoding splitting the KV cache into chunks for one query token and merging partial softmax stats

Published benchmarks on long CodeLlama-style contexts: every attention implementation in the comparison slows as context grows. This one largely does not. At the long end, decoding up to ~8× faster on the back of an attention op up to ~50× faster than unsplit FlashAttention, concentrating exactly where decode hurt most: long context, small batch.

A large batch already fills the chip through the batch dimension. Flash-Decoding matters when that free lunch is gone.


CUDA graphs: stop paying launch tax every token

One cost at decode is still untouched: how many seperate kernels the CPU must launch to produce one token.

At batch size 1, each layer launches a handful: QKV, RoPE, attention, output projection, two big FFN matmuls, a few add/norm steps. Across a 32-layer model that is hundreds of kernel launches per generated token.

Each launch costs the CPU a few microseconds of bookkeeping. At batch size 1 the kernels themselves often run for only a few microseconds too. So the GPU finishes and then idles, waiting for the next launch.

At small batch that dispatch overhead can reach ~20–30% of the whole step.

CUDA graphs amortize it. During a warm-up step the runtime records the entire launch sequence into a graph. From then on a single call replays the whole graph. Hundreds of kernels for the price of one.

Constraint: a replayed graph must touch the same memory addresses every time. Awkward for a paged cache, so serving systems layout fixed buffers to keep decode graphs replayable.

Prefill gains little. Decode at small batch gains the most. Which is why major serving stacks all replay decode this way.

CUDA graphs: pay launch once, replay many

Hundreds of tiny decode kernels · one graph call instead of hundreds of CPU round trips

Without graphlaunchlaunchlaunchlaunchlaunchGPU idle between launchesWith graphreplay whole stepone call · hundreds of kernels
Decode without CUDA graphs paying a launch between every tiny kernel versus replaying a whole step as one graph

Recap the levers

LeverWhat it attacks
FlashAttention tiling + online softmaxn2n^{2} HBM traffic for scores
FA2 occupancy / warp layoutidle SMs and warp exchange
FA3 warp specialization + TMAcopy/compute taking turns
FA4 exp split / conditional rescaleasymmetric hardware bottlenecks
PagedAttentionfragmented KV reservations
GQA / MLAhow many KV bytes exist at all
Flash-Decodingserial cache walk at decode
CUDA graphsCPU launch overhead between kernels
Lower precision (next)width of every number on the bus

Kernel intermediates. Cache layout. Cache size. Parallel cache reads. Launch overhead.

Each one made the bytes cheaper to move, or rarer to need.

One lever remains. The one FA's BF16 headlines already hinted at: the width of each number itself.

Wall-clock tracks bytes moved. Cut 16 bits to 8 and you half the memory traffic outright. Compute gets cheaper too. Hopper runs FP8. Blackwell reaches FP4. Accuracy gives out before the hardware floor is reached. Navigating that is quantization, and that is where this series should go next.


What stayed fixed

Four generations of FlashAttention. Each in its headline precision. Exact attention.

Same S=QK⊤S = QK^{\top}. Same row-wise softmax. Same multiply by VV. Identical up to floating-point reordering.

Across four generations the function never moved.

What moved was the route the bytes took through the memory hierarchy, and which pieces of hardware carried them.

Tiling kept the n×nn \times n score matrix out of HBM entirely. Online softmax made that legal. Asynchrony let copies and math overlap. When exponential units saturated, a polynomial on idle multiply-add units routed work around it.

The same pattern holds beyond the kernel. Paging moved the cache without touching the math. Splitting it across the chip let one token read it in parallel. Graph replay cut the bookkeeping between kernels.

None of those changed the model's output.

Grouping and latent compression were the exception. They shrank the cache by changing the model itself: what attention stores in the first place.

We started with the tensor cores mostly idle, stalled waiting on memory. By the end they run near three quarters of peak.

Every point of that climb came from changing the route the bytes take.


Things to hold onto

  1. Kernels are the unit. Same math, different data movement, different wall clock.
  2. IO awareness beats "more FLOPs" when you are memory-bound.
  3. FlashAttention never changed attention. It changed where SS lives.
  4. Online softmax is what makes tiling exact.
  5. FA2 was mostly about occupancy and getting bookkeeping out of the tensor-core path.
  6. FA3 was overlap: producers, consumers, TMA, SFU.
  7. FA4 is life after asymmetric scaling: help the exponential, skip useless rescales, use TMEM.
  8. PagedAttention is virtual memory for the KV cache.
  9. GQA / MLA shrink how many bytes the cache even contains.
  10. Flash-Decoding parallelizes the cache walk when batch cannot fill the chip.
  11. CUDA graphs stop the CPU from taxing every microsecond kernel at decode.

One sentence:

Inference engineering after the roofline is mostly this: keep the mathematics fixed, and spend your creativity on the path the bytes take through a machine that hates leaving HBM.

The rabbit hole continues.