ANE Prefill Research

ANE Prefill Research

Part of Technical Articles. Machine: Mac15,3 (M3, 24 GB), macOS 27.0, Swift 6.4, coremltools 9.0. Neural Engine behaviour is per-chip; these are measurements, not ceilings.

In one paragraph

TinyTitan can run the full-attention block of a prefill chunk on the Neural Engine from a Core ML sidecar. On the 35B MoE families and the dense Qwen 3.5 models that is a real win — 1.13× to 1.68× faster prefill. On Qwen 3.8-Flash-Next it is a loss (0.72× at 4-bit, 0.87× at 8-bit), and this page records both halves: the sparse-indexer fold that made 3.8 correct on the ANE, the two measured reasons it is still slower there, and the gather graph that was sized and rejected rather than built.

How the ANE prefill path works

• One Core ML sidecar per full-attention layer. layer_<L>.mlpackage is a multifunction program whose functions h0, h4096, h8192, … differ only in how much KV history they attend to; the fp16 weights are shared across them. The function name is the chunk's start position, and the runtime loads one function at a time (~/.venvs/coreml-py311/bin/python tools/export_ane_prefill.py).
• What crosses the boundary. In: the chunk's post-input-norm hidden rows, the token-major fp16 K/V history, and an additive mask. Out: the attention branch output plus the chunk's cache-layout K/V, so the KV cache and all of decode stay on the GPU and are untouched.
• When it runs. Only for a full chunk or a continuation of a long prompt; a short prompt is one partial chunk and deliberately stays on the GPU. The configured prefill chunk must equal the sidecar's width, and coverage is max(histories) + chunk.
• The degree of freedom that mattered. The mask input is arbitrary, not causal-by-contract. Almost everything below follows from that.

What worked

1. Offloading dense attention

Recorded in benchmark/ane-prefill/v5.5-ane-matrix-3.json (23,000-character prompt, ~4,333 tokens, chunk 4,096, two measured runs per arm after a warm-up each):

Model	GPU	ANE	Ratio
Qwen 3.5 2B 4-bit	49.8 s	36.0 s	1.38×
Qwen 3.5 2B 8-bit	104.1 s	91.6 s	1.14×
Qwen 3.5 4B 4-bit	140.6 s	102.2 s	1.38×
Qwen 3.5 4B 8-bit	275.6 s	218.1 s	1.26×
Qwen 3.5 9B 4-bit	210.7 s	168.2 s	1.25×
Qwen 3.5 9B 8-bit	442.4 s	392.9 s	1.13×
AgentWorld 35B-A3B 4-bit	106.5 s	63.3 s	1.68×
Qwen 3.8 125B-A6B 4-bit	175.1 s	—	see below

Two properties explain the shape of that table. The win scales with the full-attention share of the model (the 35B offloads 10 of 40 layers, the dense 2B only 6 of 24), and it grows with prompt length, because attention is the quadratic term. Ratios fall at 8-bit because the non-attention prefill the ANE never touches is heavier there.

A narrower sidecar reaches the band that 4,096 cannot serve at all: at ~2,500 tokens the dense 2B measured 23.33 s → 17.88 s, 1.30×, with --chunk 1024 (ane-chunk1024-2b-4bit.json).

2. Folding a sparse indexer's selection into the mask

Qwen 3.8's full-attention layers do not compute a different attention: they choose keys with a QSA indexer. Dense attention matches that choice only through keptBlocks × compressRatio + (compressRatio − 1) = 2,051 visible keys, and the smallest chunk the ANE accepts is 4,096 — so a sidecar fed the causal mask attends to keys the model drops, silently and plausibly.

The runtime now folds the indexer's compacted selection (keepIndices / keepCounts, the same keys the GPU's attention gathers) into the additive mask as -30000, which exp() underflows to zero exactly as an omitted key contributes nothing. The graph does not change at all; a sparse-indexed model is served by a per-layer mask buffer.

Correctness was measured rather than asserted, because the obvious check was not sensitive enough — see Two checks that were not sensitive enough below:

Check	Result
Graph vs NumPy reference under a QSA-shaped mask, both widths	0.47 % relative error
The causal-only mask against the same reference	7.6 % moved, 7.9 % off — 16× the fp16 noise
32-token greedy continuation, folded vs GPU	textually identical
Causal control's prefill_logits vs the folded arm	13.6–17.8 % apart, so the mask reaches the graph
QSA budget 2,048 → 8, folded arm	L3_after moves 10.2 %, logits 69 % — the mask is live

The row that took work is the last two: proving that the mask reaches the graph is a different claim from proving the graph honours a mask, and only the second is visible in the sidecar's own weights.

tools/verify_ane_sidecar.py now runs the folded check automatically for a sidecar that records selectionFolded, and fails as vacuous if a synthetic selection barely moves the reference. The runtime refuses a sparse-indexed model's sidecar that does not record the contract, and refuses to attend densely: a chunk past the dense-exact window with no selection throws.

3. Making "verified" mean something (issue #7)

The exporter records aneCompileVerified: true, and the runtime refuses a sidecar without it. That flag turned out to need three layers, because "Core ML stayed quiet" is not the same as "the Neural Engine got the graph":

1. Compile markers — Core ML reports an ANE compile refusal on the native stderr and still exits 0, so the exporter reads stderr back and refuses the export (ANECCompile() FAILED, MILCompilerForANE error, …). Now tested.
2. Load — every function the metadata records must load under its h<history> name. This is what catches a variant that loads but can never predict().
3. Assignment — MLComputePlan reports the device each operation is assigned to, so the exporter asks it per variant and fails on zero. Measured on the real 3.8 h12288: 0 of its 173 operations on the Neural Engine, against 64–74 for healthy variants. This is the layer that closes the silent case, where nothing is printed and the graph quietly runs on the CPU at roughly 38× the GPU cost.

One placement detail worth keeping: the assignment check runs per variant, before the multifunction merge, because the compute plan reports assignments only for a package's default function — the same 2B graph reads 74/173 standalone and 0/173 as a non-default function of the merge.

What didn't

1. Qwen 3.8: the ANE loses, and here is exactly why

Model	GPU	ANE	Ratio
Qwen 3.8 125B-A6B 4-bit	197.5 s	273.5 s	0.72×
Qwen 3.8 125B-A6B 8-bit	428.3 s	491.8 s	0.87×

Records: v5.6-ane-38-4bit.json, v5.6-ane-38-8bit.json.

• The GPU path is already sparse. It gathers ~2,051 selected keys per query; the ANE graph scores all history + chunk of them and masks the rest. At 4,333 tokens that is 4,096 × 8,192 score entries per layer-chunk against the GPU's 4,096 × 2,051 — about 4× the arithmetic, and the gap widens with context.
• Per-variant setup dominates. One MLModel(contentsOf:) measured 6.7 s for h0, 13.6 s for h4096, 37.5 s for h8192 on this geometry, against the ~0.5 s per layer-chunk the runtime's own note records for the 35B. A 4,333-token prompt visits 24 (layer, history) variants, so setup alone is minutes against a ~185 s GPU prefill.
• The load tracks the score arena, not the package. Loading the same function from different package compositions: 6.68 s for h0 alone, 6.94 s beside h4096, 6.84 s beside a variant that cannot load at all. Per-variant specialisation is what costs; packing is irrelevant.
• The ANE's per-FLOP edge on this graph is small: one h0 prediction measured 0.43 s on the ANE against 0.65 s on the CPU alone, where the 35B's blocks measured 26.7× the GPU's cost (recorded in the v4.5 design note).

Conclusion, and the shipped behaviour: 3.8 stays on the GPU, no sidecar is installed for it, and tools/ane_sidecars.sh skips the family rather than building one that would slow the default path.

2. The gather graph: sized, rejected

If the ANE scored only the selected keys, the extra arithmetic would disappear — so the variant was sized before it was built (docs/ane-gather-graph-sizing.md, benchmark/ane_gather_probe.py).

The catch is that a gathered key must be materialised once per query that selects it: the graph holds heads × chunk × budget × headDim where the dense graph holds heads × chunk × total — 64× more data, 103 GB at the real chunk 4,096, against a 1.6 GB dense score matrix. Measured at probe-sized chunks:

chunk	graph	ANE ops	load	predict
32	dense	13	0.133 s	0.042 s
32	gather	9	0.186 s	0.238 s (5.70×)
64	dense	13	0.144 s	0.043 s
64	gather	12	2.883 s	fails

The ANE accepts the gather — it is assigned 9–12 operations — so this is a cost problem, not a support limitation, which is the harder kind. It is 5.70× slower where it runs, it stops running one chunk size up once the gathered keys reach 1.6 GB, and it does not help the load either (1.4× slower at chunk 32, 20× at chunk 64).

3. Wide coverage on this geometry

--max-history 32768 on 3.8 converts all nine variants and then fails to load: the 24-head h12288 variant needs a 3.2 GB arena and is not an mlprogram at all once the ANE has refused it. A request that reached that history would have died at load rather than falling back. The exporter now refuses to record it (see the three layers above), and the practical cap is --max-history 8192 — coverage 12,288 tokens — for this model.

4. Two checks that were not sensitive enough

Worth recording because both looked convincing:

• A 32-token greedy comparison agreed with the wrong mask. On this prompt the causal-only control produced textually identical output to the folded arm, and to the GPU. Only the logits separated them (13.6–17.8 %). A digest at --max-new 1 is a determinism check and never a correctness one.
• A per-layer activation identity looked like a stale dump, and was not. The folded and causal arms were bit-identical at L3_after while their logits differed — which reads exactly like a dump taken before the layer's own work. It is not: the dumps are live and post-layer (each differs from the previous layer by tens of percent, and on the dense 2B the ANE arm's L3_after differs from the GPU's by 0.94 %), and a real mask change moves them (budget 2,048 → 8 moves L3_after by 10.2 %). The identity is a property of the masks: at the first full-attention layer the blocks this indexer drops contribute below fp16 resolution, so the selection and the causal mask compute the same output there, and the divergence begins at a later attention layer.

Lessons

• Where the ANE wins: dense, quadratic attention on long prompts, with a cheap per-variant setup. Offload share and prompt length set the ceiling.
• Where it loses: when the GPU path is already sparse — a budget-bounded gather is cheaper than scoring everything and masking — and when per-variant specialisation is large relative to the work it saves.
• The dominant ANE cost was setup, not FLOPs. 24 variant loads ≈ 245 s against a 185 s prefill. Optimising arithmetic there would have been wasted work.
• Prove the claim, not the absence of complaints. aneCompileVerified had to become "the compute plan assigns this variant to the Neural Engine", because a silent CPU fallback prints nothing.
• A correctness check needs a negative control. The mask was known to be load-bearing only because the causal-only arm could be run and measured, and the first check that distinguished them was the logits, not the text.

Reproducing

P=~/.venvs/coreml-py311/bin/python

# export (and gate) a sidecar for an installed model
$P tools/export_ane_prefill.py --model models/qwen3.5_2B_4Bit --chunk 4096 --max-history 12288

# the graph against an independent NumPy reference, including the fold check
$P tools/verify_ane_sidecar.py --model models/qwen3.5_2B_4Bit

# prefill A/B against the GPU path, and a 32-token correctness comparison
python3 benchmark/ane_prefill_ab_matrix.py --models qwen3.5_2B_4Bit --repeats 2 --label mine
python3 benchmark/ane_prefill_correctness.py --models qwen3.5_2B_4Bit --max-new 32

# per-arm activation diff, and the gather-graph probe
python3 benchmark/ane_prefill_layer_diff.py --models qwen3.5_2B_4Bit
$P benchmark/ane_gather_probe.py --chunk 32 --repeats 3

Records live under benchmark/ane-prefill/, the decision rule and family table in benchmark/ane-prefill/README.md, and the gather study in docs/ane-gather-graph-sizing.md. Golden baselines pin TINYTITAN_PREFILL_ANE=off, so byte-reproducible generation is unaffected by any of this.

What would re-open it

• A cheaper per-variant specialisation on a future chip — the term that dominates here is setup, so that is the number to beat.
• An attention implementation that does not materialise gathered keys (an indexed matmul at the compiler level, which MIL does not express today).
• A longer or differently-shaped workload where the offload share is larger.

Re-run the gather probe first: it takes minutes, and it is the cheapest way to find out whether the answer changed.