# APA: sloppy attention, exact where it counts

> A prefill kernel for the RTX 5090 that runs every attention tile in FP4, recomputes the 6 to 10 % that carry the softmax mass in exact FP16, and merges both: 2.50x faster than all-exact at 122880 keys.

Source: https://rfriedmann.de/blog/apa-adaptive-precision-attention/
Published: 2026-10-28 · Track: log · Level: Expert

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](https://github.com/kekzl/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](https://github.com/kekzl/apa/blob/main/paper/apa.pdf).

## Two camps, and the gap between them

There are two established ways to make attention cheaper.

| Approach | Examples | Keeps every key | Error |
|---|---|---|---|
| Low precision | SageAttention 1 to 3 | yes | a little on every score |
| Sparse | StreamingLLM, MInference, Quest | no | none on kept keys, total on dropped ones |
| APA | | yes | FP4 on low-mass tiles, none on high-mass tiles |

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](https://arxiv.org/abs/2309.17453), 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.

<figure class="chart">
<div class="chart-title">One query row, 64-key tiles, share of the softmax mass (schematic)</div>
[diagram omitted — see the page for the chart]
<figcaption class="chart-caption">Tiles above the threshold ε (accent) are hot and get recomputed in FP16. Everything below stays in FP4. At ε = 0.005 and 122880 keys, 6.0 to 10.3 % of the active tiles end up hot.</figcaption>
</figure>

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](/blog/online-softmax-and-flashattention/).
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](/blog/nvfp4-at-the-bit-level/). 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:

<ol class="pipeline">
<li data-n="1"><h4>Scores in FP4</h4><p>S = Q·Kᵀ on the FP4 tensor cores, causal mask applied.</p></li>
<li data-n="2"><h4>Probabilities in FP4</h4><p>Each 16-key group gets its own scale, folded into the exponent bias, so <code>ex2</code> directly produces the quantized P.</p></li>
<li data-n="3"><h4>Row sums from the same P</h4><p>P is multiplied by a tile of E2M1 ones. Numerator and denominator of the cold result carry the same quantization.</p></li>
<li data-n="4"><h4>The hot test</h4><p>If any of the warp's 16 rows has tile sum ℓ_t &gt; ε·λ, the tile is hot: set a bit in the hot mask, skip the PV product.</p></li>
<li data-n="5"><h4>Cold tiles accumulate</h4><p>o += P·V in FP4, the cold row sum grows by ℓ_t.</p></li>
</ol>

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:

| Quantity | Value |
|---|---:|
| Mass of the top tile (the sink) | 0.916 |
| Norm of the hot part of the reference | 0.2527 |
| Norm of the cold part of the reference | 0.3634 |
| Norm of the full reference | 0.1670 |
| Norm of APA's output | 0.2531 |
| Cosine | -0.4446 |

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.

| Dump | min cos 0.2.0 | pairs < 0 | min cos 0.3.0 | time 0.2.0 | time 0.3.0 |
|---|---:|---:|---:|---:|---:|
| lc_122880_0 | 0.634875 | 0 | 0.959676 | 4.555 ms | 4.490 ms |
| lc_122880_1 | 0.925502 | 0 | 0.975775 | 4.676 ms | 4.759 ms |
| lc_122880_2 | -0.638155 | 789 | 0.841302 | 4.758 ms | 4.840 ms |
| lc_32768_0 | 0.757021 | 0 | 0.927517 | 2.245 ms | 2.285 ms |
| lc_32768_1 | 0.927160 | 0 | 0.980919 | 1.881 ms | 1.889 ms |
| lc_32768_2 | -0.594803 | 110 | 0.969776 | 1.856 ms | 1.863 ms |

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:

| Only this operand in FP4 | Mean cosine, lc_65536_0 | lc_122880_2 |
|---|---:|---:|
| kernel (0.4.0) | 0.809680 | 0.888434 |
| emulation, all sources | 0.815728 | 0.890648 |
| P only | 0.996408 | 0.998214 |
| V only | 0.999817 | 0.999897 |
| K only | 0.998618 | 0.999879 |
| **Q only** | **0.823032** | **0.907227** |
| all, Q as two FP4 terms | 0.998386 | 0.998000 |

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.

| Twelve dumps | 0.4.0 | 0.6.1 |
|---|---:|---:|
| Pairs below cosine 0.9 | 137 | 0 |
| Pairs below cosine 0.99 | 5196 | 905 |
| Worst minimum cosine | 0.773144 | 0.984284 |

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:

| lc_122880_2, APA 0.3.0 | Hot | Mean cos | Min cos | Pairs < 0.9 | Time |
|---|---:|---:|---:|---:|---:|
| all exact | 100 % | 0.993600 | 0.834726 | 465 | 12.049 ms |
| ε = 0.001 | 19.0 % | 0.999274 | 0.926464 | 0 | 7.489 ms |
| ε = 0.003 | 9.1 % | 0.998755 | 0.857105 | 11 | 5.871 ms |
| ε = 0.01 | 3.2 % | 0.997936 | 0.834358 | 15 | 4.844 ms |
| ε = 0.02 | 1.7 % | 0.997372 | 0.624999 | 18 | 4.340 ms |

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:

<figure class="chart">
<div class="chart-title">Attention time, all-exact vs APA (ε = 0.005, APA 0.3.0, lc_122880_1)</div>
<div class="chart-legend"><span><i class="k2"></i>all tiles exact (FP16)</span><span><i class="k1"></i>APA</span></div>
[diagram omitted — see the page for the chart]
<figcaption class="chart-caption">Hot share falls from 36.0 % at 8192 keys to 9.6 % at 122880, and the speedup grows from 1.30x to 2.50x. Including the 0.730 ms quantization step that only APA needs, the ratio at 122880 keys is 2.196x. Attention kernel only, one run; times on this GPU drift by up to 10 % between runs, so every comparison is taken within one run.</figcaption>
</figure>

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](/blog/how-imp-works-model-to-token/) 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:

| Prefill mode | FlashAttention-2 | APA | Speedup |
|---|---:|---:|---:|
| default (sparse prefill) | 6661.95 ms | 5241.18 ms | 1.27x |
| dense | 11659.35 ms | 6433.23 ms | 1.81x |

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):

| Variant | Llama-3.2-3B | Qwen3-8B |
|---|---:|---:|
| FA2, f16 accumulation | 17.3644 | 10.7974 |
| FA2, f32 accumulation | 17.3732 | 10.8110 |
| APA ε = 0.002 | 17.3556 | 10.7855 |
| APA ε = 0.005 | 17.3423 | 10.7869 |
| APA ε = 0.01 | 17.3385 | 10.7728 |

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.