APA: sloppy attention, exact where it counts

Long prompts are expensive because of attention. For a prompt of length N the attention work grows with N², the rest of the model with N. At 122880 keys, 2048 query rows of Llama-3.2-3B take 13.023 ms on an RTX 5090 with every tile computed exactly. APA does the same job in 5.200 ms.

The idea fits in one sentence: compute everything cheaply in FP4, find the few tiles where FP4 error would actually show up in the output, redo only those in exact FP16, and merge the two results without any approximation. This post walks through how that works, and through the bug that turned a clean idea into a cosine of -0.638 before it got fixed. The full write-up with every table is the paper.

Two camps, and the gap between them

There are two established ways to make attention cheaper.

Low precision spreads its error evenly, including over the handful of keys that decide the output. Sparsity is exact where it looks and blind where it doesn’t. APA keeps every key, like the first camp, and spends precision where the mass is, like the second. The split is not predicted up front: the FP4 pass decides it online, from numbers it computes anyway.

Why this can work at all

In long-context causal prefill, the softmax mass of a query row is not spread evenly. A few key tiles carry almost all of it: typically tile 0, the attention sink, and the tiles right before the query on the causal diagonal. The long middle carries a thin layer of probability spread over thousands of keys.

If FP4 error on a tile that holds 0.5 % of the row’s mass is the price for not computing it in FP16, that is a good trade. If FP4 error on a tile holding 90 % of the mass is the price, it is not. So the kernel needs a test that tells the two apart while it runs.

The merge that makes it exact

The glue is a property of online softmax. FlashAttention keeps three running values per query row: the maximum score m, the sum of exponentials l and the unnormalised output o. Two such states computed over disjoint sets of keys merge exactly:

M = max(m1, m2)
w1 = exp(m1 - M),  w2 = exp(m2 - M)
O = (w1*o1 + w2*o2) / (w1*l1 + w2*l2)

So APA runs two passes. Pass 1 covers the cold tiles in FP4 and yields one state. Pass 2 covers the hot tiles in FP16 and yields the other. One log-sum-exp merge gives the final output, with no extra error from the split itself. The only condition: both passes must share one logit frame, which matters below.

Pass 1: FP4 everywhere, with a hot test

The RTX 5090 has a block-scaled FP4 tensor core instruction, mma.sync kind::mxf4nvf4 (shape m16n8k64). Operands are E2M1, which can only represent 0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6, plus one UE4M3 scale per 16 values (max 448). That is NVFP4. APA adds one FP32 scale per head, so the largest value of a head maps to 6 × 448 = 2688.

Keys come in tiles of 64. One CTA of 12 warps handles 192 query rows; with GQA the rows of all query heads that share a KV head are packed together, so each KV tile is loaded once for all of them. Per tile and warp:

1. Scores in FP4

   S = Q·Kᵀ on the FP4 tensor cores, causal mask applied.

2. Probabilities in FP4

   Each 16-key group gets its own scale, folded into the exponent bias, so ex2 directly produces the quantized P.

3. Row sums from the same P

   P is multiplied by a tile of E2M1 ones. Numerator and denominator of the cold result carry the same quantization.

4. The hot test

   If any of the warp's 16 rows has tile sum ℓ_t > ε·λ, the tile is hot: set a bit in the hot mask, skip the PV product.

5. Cold tiles accumulate

   o += P·V in FP4, the cold row sum grows by ℓ_t.

The hot test costs one small extra MMA per tile, the one against the ones tile. Everything else it needs, the pass computes anyway.

Cold verdicts are safe. λ is the running sum over all tiles seen so far, hot ones included. Sums of exponentials only grow, so a tile that is below ε of the running sum when it is judged is also below ε of the final sum. A cold verdict is never revoked later. A hot verdict can be too conservative, which costs pass-2 work, not accuracy. One caveat the paper states itself: this bound is on the FP4 estimate of the mass. The true FP32 share of a cold tile reached 0.0118 at ε = 0.005.

Tile order. Early in the loop λ is small, which makes every tile look hot. So pass 1 does not walk the keys left to right. It visits tile 0 first (the sink), then the tiles from the causal diagonal backwards. The tiles that usually carry the mass enter λ first, and the long middle is judged against an almost final sum.

Three FP4 tricks keep the cold result usable:

• Mean-centred K. Subtracting the per-channel key mean shifts every score of a row by the same constant and leaves the softmax unchanged (taken from SageAttention). It frees FP4 range that a shared channel offset would eat. Pass 2 subtracts the same constant from its exact scores, so both passes share a frame.
• Lazy maximum. The running max is only updated, with a reduction and a rescale, when a score exceeds it by more than τ = 4 (in log2 units). Most tiles skip the rescale.
• A 9216-byte blob per tile. A prep stage writes K and Vᵀ in E2M1 plus their scales, swizzled for conflict-free ldmatrix, so one bulk copy fetches a tile. Pass 1 streams them through a 6-stage shared-memory ring with cp.async.bulk.

Pass 2: exact FP16 on the hot tiles

A pass-2 CTA takes the same 192 rows, collects the union of their hot tiles and streams them from the original FP16 K and V, flat or paged. Each warp computes only its own hot tiles with mma.sync.m16n8k16, does an FP32 online softmax and finally merges with its pass-1 state.

Both passes run in one persistent launch, one CTA per SM. CTAs pull tickets in order from an atomic counter: first pass-1 work, then whole query blocks for pass 2 from a second counter. Every wait is on an earlier ticket held by a CTA that is already running, so the scheme cannot deadlock. Pass-2 work fills the SMs that pass 1’s last wave would leave idle.

The bug: a sink that drowned the tail

APA 0.2.0 worked well on average. On one dump, a 512-pair sample already showed a cosine of -0.444615 against the FP32 reference. Over all 49152 (row, head) pairs the minimum was -0.638155, and 789 pairs pointed the wrong way.

A diagnostic mode splits the reference output of the worst sampled pair (ε = 0.01) into the share of the hot and of the cold tiles:

Two numbers explain everything. The cold part, carried by 8 % of the mass over 1890 tiles, has a larger norm than the whole reference: hot and cold nearly cancel each other. And APA’s output norm, 0.2531, equals the hot part alone, 0.2527. The cold contribution was gone.

The cause is the logit frame. Every FP4 probability is scaled relative to the running max. The per-group scale exponent is clamped at -6 to stay a normal UE4M3 value, so a group whose maximum lies more than about 2⁻¹⁰ below the frame rounds to E2M1 zero. In 0.2.0 the frame was the running max over all tiles. The sink comes first, is hot, and sets the frame. In rows where the remaining mass is spread thin, every cold group sat far enough below the sink to vanish. The tail that decided the direction of the output was quantized to zero.

The fix is a decoupled cold frame. Hot tiles don’t need the frame: their P is only used for the hot test, and pass 2 recomputes them anyway. Since 0.3.0 a hot tile may lift the frame only to max(m, c - 32); λ absorbs the hot sum with a factor of at most 2³², well within FP32. With the offset set to 0, the kernel reproduces 0.2.0 to all six printed digits, which isolates the change.

All 899 negative pairs gone, the minimum up on every dump, the time within run-to-run noise. Three other suspects were tested and rejected along the way: cold tiles above ε, heads with unusual channel spread, and the tile order (ascending order still leaves 787 pairs below 0).

The lesson from the sample is worth keeping too: on the worst dump the sample minimum was -0.444615, the all-pairs minimum -0.638155. A 1 % sample overstates the worst case on every dump.

Second hunt: which operand is to blame

After the frame fix the worst rows looked different: diffuse rows with 0.2263 to 0.3973 of their mass spread over 314 to 454 cold tiles. The FP4 estimate of the cold share matched FP32, so no mass was lost. Something else was rounding badly.

To find out what, the host emulates pass 1 with one quantization at a time on the cold tiles of the 20 worst pairs:

The emulation reproduces the kernel, and Q alone accounts for the error. The fix quantizes Q twice: Q₁ as before, Q₂ the FP4 residual Q - Q₁ in the same scale. Both terms accumulate into the same fragment against the same K operands, for 16 more FP4 MMAs per tile.

With two terms, P became the largest error source. Its group scales switched from powers of two to UE4M3, so the group maximum lands near 6 instead of anywhere in (3, 6]. At short context K turned out to decide the worst rows, fixed by picking every block scale from a small set of candidates by least squared error instead of taking the nearest one.

This is not free. On the nine long dumps attention takes 6.7 to 25.4 % longer: the second Q term doubles the QKᵀ work of pass 1, the UE4M3 P scale adds 5.1 to 5.2 %.

Speed: the longer the context, the better

The trade-off is controlled by ε. A larger ε makes fewer tiles hot and the kernel faster, at the cost of the worst rows:

The first row is the surprise: the all-exact run is not the most accurate one. Pass 2 accumulates its output in FP16 (imp’s FlashAttention-2 default), and next to a dominant sink each small tile’s contribution falls below half an FP16 ulp of the accumulator. An FP32 accumulator exists as an option, but costs 14.5 to 24.9 % time at ε = 0.005 through register spills, so it is off by default.

How much APA saves depends on context length, because the hot share shrinks as the context grows:

The accuracy fixes of 0.4.0 to 0.6.1 came after these measurements and cost the 6.7 to 25.4 % from the previous section. Separately, 0.8.0 cut prep plus attention by 1.55 to 14.32 % on synthetic inputs by scheduling pass 2 between KV segments (below).

End to end in imp

APA is integrated into imp behind the config key attention.apa_eps. These numbers were measured in the imp repository with APA 0.2.0 and ε = 0.005, Llama-3.2-3B Q8_0, a 106451-token prompt, whole prefill:

And quality, as perplexity (lower is better) on a 13k to 14k token corpus, APA 0.2.0 (the ε = 0.002 and 0.01 rows from 0.1.0, which 0.2.0 reproduces exactly at ε = 0.005):

Every APA setting with ε > 0 lies below both FlashAttention-2 variants. The safe reading is “no measurable loss”, not “FP4 improves the model”: with APA 0.3.0 and ε = 0.01 on a different corpus, perplexity moves slightly up, Llama-3.2-3B 18.2760 → 18.2836, Qwen3-8B 10.7522 → 10.7549.

What else is in the box

• Deterministic. All reductions run in a fixed order. In 0.1.0 float atomics made 5 reruns differ in up to 1329742 output elements; now they are bitwise identical.
• KV window. The FP4 copy holds the sink tile plus 2 ring slots of 96 tiles per KV head instead of every tile: 13.6 MiB at 8 KV heads for any context, 135.0 MiB at 122880 keys before. Quantization of the next segment runs inside the launch while pass 1 works on the current one, and pass 2 of a segment runs as soon as the next segment’s pass 1 is scheduled; idle SM time dropped from 7.6 to 1.4 % of the span. Outputs stay bitwise equal to one pass over all tiles.
• Sliding window. A 4096-key window at 122880 keys takes 0.804 ms against 6.053 ms causal.
• Paged KV and chunked prefill. Pass 2 reads through a block table; chunked prefill keeps the quantized tiles and only quantizes new keys, with a restat once the context has grown by 1.125x.
• Head dim 128 and 256, GQA, header-only CUDA library, MPL-2.0.

Honest limits

• Mass is not output. The hot test bounds each cold tile’s share of the softmax mass, not of the output. The worst 0.3.0 pair had 3 hot tiles, 0.0794 cold mass and a cold output component 1.918 times the reference norm. A test weighted by value norm would target such rows. The worst pair of the current kernel over twelve dumps: cosine 0.984284.
• Granularity. The verdict is per warp and tile: one of 16 rows makes the tile exact for all 16.
• Scope. sm_120a only, causal, FP16 inputs. Head dim 256 is tested on synthetic inputs only.
• Evidence. The dumps come from one model, Llama-3.2-3B. A head-to-head against SageAttention3 on the same inputs is still open.

The skew that makes APA pay is also its main hazard. Because a few tiles dominate, exact arithmetic on 6 to 10 % of them is enough. Because a few tiles dominate, the dominant one can push the informative tail out of FP4 range. Keeping the cold frame independent of the hot tiles removed that failure at unchanged speed.