<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;600;700&display=swap">

rafal@arasz:~$cat blog/ai-raccoon-neural-engine.md

ai-raccoon on the Neural Engine: Twice as Fast at a Quarter of the Energy

Two days ago I wrote about making ai-raccoon quietThe previous part: how ai-raccoon 1.44.3 to 1.51.2 cut idle CPU by about 90%, went from 6.8 GB to 1.75 GB and moved its embeddings to the GPU.: less CPU, a quarter of the memory, and embeddings on the GPU. That post ended with the GPU as the good place for the model to live. It turns out the M4 has a better place, and it had been sitting there, power-gated, the whole time.

Every Apple Silicon chip has a Neural EngineApple's NPU: a block of the chip built only for the arithmetic neural networks do. On the M4 it has 16 cores. Apps cannot program it directly; they reach it through Core ML. next to the CPU and GPU. It is Apple's NPUNeural processing unit: a processor designed for one job, the multiply-and-add arithmetic of neural networks, at far less power than a CPU or GPU doing the same work., a piece of silicon that does one thing, the arithmetic inside neural networks, and does it on very little power. ai-raccoon 1.53.0 can now run its embedding model there. On an M4, embedding the same 4,914 chunks takes half the time the GPU needs, uses about a quarter of the energy, and searches come back faster. Recall did not move.

Getting there took two attempts. The first one failed on the measurements, and I wrote a decision record saying we would not ship it. The second one needed the model rebuilt in a different shape, a new file on Hugging Face, and a compile step that takes longer than most people would wait. This post covers both, and starts with a short primer for anyone who has never had a reason to care what an NPU is.

The full measurement reports live in the ai-raccoon repo: device benchmark on the M4The dated source report: time, energy, CPU and search latency for every device through the shipped product, three repeats each. and the ANE-layout re-exportThe research record for the re-exported model: op placement, the A/B against MLX, recall, the fp16 overflow and the powermetrics check..

Same corpus through the shipped product on an Apple M4 with 24 GB, medians of three interleaved repeats (1.53.0). Energy is net of the idle baseline. p95 search is the range of per-repeat p95s from a second run on 1.53.1.
Device on an M4Time to embed 4,914 chunksSystem energyEnergy per chunkServer CPU time (avg. cores busy)p95 search
Neural Engine (coreml, opt-in)33 s379 J77 mJ26 s (0.8 cores)35-39 ms
GPU through MLX (mlx, opt-in)49 s1,146 J233 mJ37 s (0.8 cores)not measured
GPU through WebGPU (auto, the default)70 s1,389 J283 mJ42 s (0.6 cores)51-88 ms
CPU only (cpu)170 s4,076 J829 mJ863 s (5.1 cores)not measured

Server CPU time adds up the time every core spent on ai-raccoon, so 863 seconds in a 170-second run means about five cores busy the whole time. The number in brackets is that average, CPU time divided by run time. The total is what the work costs. The Neural Engine's average (0.8 cores) sits slightly above WebGPU's (0.6) only because it does less CPU work in less than half the time.

First, what is an NPU and why would you want one?

If you already know what fp16, graph partitions and static shapes are, skip to the next section. If not, here is the short version.

What the model does. ai-raccoon stores notes and code and lets AI agents search them by meaning. To do that it turns each piece of text into an embeddingA list of 384 numbers that stands for the meaning of a piece of text. Texts with similar meaning get similar lists, so search becomes: find the nearest lists., a list of 384 numbers. The model that produces those numbers is IBM's granite-embedding-small-english-r2The small English embedding model ai-raccoon ships: 12 transformer layers, 47M parameters, 384-number vectors, Apache 2.0.. Running it means billions of multiplications and additions per piece of text, with a few other operations in between. Every chunk ingested and every search query goes through it.

Three places to run it. A laptop chip has three kinds of processor that can do that arithmetic:

The three processors on an Apple M4, from the point of view of an embedding model
Good atCost for this job
CPUAnything, one step after anotherSlow and power-hungry for matrix math. The cores you also need for your IDE and compiler.
GPUThousands of identical small calculations at onceFast, but it is a big block that draws a lot of power when it wakes up.
NPU (Neural Engine)Exactly the multiply-add patterns of neural networks, nothing elseFast and frugal, but very picky about what it accepts.

A useful picture: the CPU is a skilled craftsman, the GPU is a factory floor, and the NPU is a specialised machine that stamps out one part extremely efficiently, provided the raw material arrives in exactly the right size and shape.

Why picky matters. A model file is a graphAn ONNX model is a graph: a list of operations (nodes) wired together, each taking tensors in and producing tensors out, plus the weights they use. of operations: multiply this matrix, normalise that vector, add these two. ai-raccoon uses ONNX RuntimeMicrosoft's engine for running ONNX model files. It offers each part of the graph to its hardware backends in priority order, and the CPU takes whatever is left., which walks that graph and hands each operation to the first hardware backend, in priority order, that will take it, through a plug-in called an execution providerONNX Runtime's plug-in for one kind of hardware: CPU, WebGPU, CUDA, MLX, CoreML. Each one claims the operations it supports and leaves the rest.. On a Mac the only door to the Neural Engine is Apple's Core MLApple's framework for running models on-device. It alone decides whether work goes to the CPU, GPU or Neural Engine; there is no public API to the Neural Engine itself., so the door we used is ONNX Runtime's CoreML execution providerThe ONNX Runtime backend that converts supported parts of an ONNX graph to Core ML and lets Core ML run them, on the Neural Engine if it can.. Core ML decides on its own whether a piece of work goes to the Neural Engine. In practice it only took the whole model when the first three things below were true, and it ran best when the fourth was:

  1. The numbers are half precision. Models are usually stored in fp3232-bit floating point numbers: about 7 significant digits and a huge range. The default for training and for most exported models. (32-bit). The Neural Engine works in fp1616-bit floating point numbers: half the memory and much cheaper arithmetic, but the largest value is 65,504, so big intermediate numbers overflow., which uses half the memory but tops out at 65,504.
  2. The shapes are fixed. Text comes in any length, but the Neural Engine wants to know the exact size of every block of numbers (every tensorA multi-dimensional array of numbers. Everything a model reads, computes and outputs is a tensor.) up front (static shapesWith static shapes every input has a fixed size known when the model is compiled. Variable-length inputs must be padded up to one of those sizes.). So each piece of text gets padded up to one of a few fixed lengths, called buckets.
  3. The operations are ones it knows. Anything it doesn't support falls back to the CPU, and every switch between CPU and Neural Engine is a copy. The pieces Core ML accepts are called partitionsA contiguous piece of the graph handed to one execution provider. Every boundary between partitions is a copy between devices, so one partition is best and 25 means dozens of copies per inference.. One big partition is good. Twenty-five is bad.
  4. The data is laid out the way it likes. This is the one that is easiest to miss. Apple's own research paperDeploying Transformers on the Apple Neural Engine, Apple Machine Learning Research, June 2022: four principles for writing transformers the Neural Engine runs well, up to 10x faster and 14x less memory on Apple's own DistilBERT benchmark. says the Neural Engine prefers tensors shaped as (B, C, 1, S): batch, channels (the 384 numbers per token), a dummy 1, and sequence length. Ordinary transformer code, written for GPUs, is not shaped like that.

What we changed, in one paragraph. We could not change the Neural Engine, so we changed the model. We rebuilt the same model, with the same weights, in half precision and in Apple's preferred layout, so that Core ML accepts almost all of it as one piece. We then taught ai-raccoon to compile it for the Neural Engine in the background on first start, keep serving from the GPU while that happens, and switch over once a test row embedded on both devices comes out the same. The model's output is the same, and so are your stored vectors. It just runs somewhere else.

Attempt one: plug in CoreML and hope

The CoreML execution provider already ships inside ONNX Runtime's macOS package, so the first attempt was one line: ask for it. What came back was CPU in disguise. With the static-shape flag set but the model's variable-length dimensions left in, Core ML accepted 8 of 401 operations and the rest ran on the CPU. Allowing dynamic shapes got 305 of 401 operations across 37 partitions and ran about twice as slow as the plain CPU, around 200 ms against 111 ms at 512 tokens. Asking for the Neural Engine explicitly, with dynamic shapes, failed to build at all, with Error in declaring output ... with error -1.

Three fixes got it running:

  • Fixed shapes, including one nobody mentions. ONNX Runtime can pin symbolic dimensions with free-dimension overrides (batch_size=1, sequence_length=N). The build still failed with the same error -1 until the attention mask's own dimension, total_sequence_length, was pinned as well.
  • One compile cache per shape. Core ML compiles a model for each input shape and caches it on disk. A cache folder shared across shapes served the 64-token model to a 256-token row: MultiArray shape (1 x 256) does not match the shape (1 x 64). Every bucket now gets its own folder.
  • No spinning. ONNX Runtime's worker threads busy-waitA thread that loops checking whether work has arrived instead of sleeping. Low latency, but it burns a CPU core while waiting. by default. That roughly doubled CPU per row on every device. It is off everywhere now.

Fixed, it worked, and it lost. The shipped model still split into 25 partitions, and it cost 9 to 18 times the CPU-seconds of the MLX GPU path. Worse, the Neural Engine's own memory counter collapsed to about 8 MiB from the 704-token bucket up while the process footprint climbed past 6 GB, which looked like Core ML quietly running the long rows on the CPU. The rule I had written down before the first run said: if every CoreML repeat costs more than every MLX repeat, stop. It did. ADR-0117The decision record from the first attempt: the CoreML execution provider stays out, and the one lever that could change that, a re-export in Apple's ANE layout. recorded "CoreML stays out", with two re-open triggers written in, the likelier being a model re-exported in Apple's layout.

Attempt two: rebuild the model the way the Neural Engine likes it

Apple publishes a reference implementation of its paper, apple/ml-ane-transformersApple's reference PyTorch code for transformers laid out for the Neural Engine: Conv2d instead of Linear, channels-first (B, C, 1, S) tensors, per-head attention.. It is written for BERT-style models, and granite is a newer ModernBERTA 2024 redesign of BERT: rotary position embeddings, alternating local and global attention, longer context. granite-embedding-small-english-r2 is built on it., so I wrote a PyTorch module that loads granite's weights and computes the same thing in the ANE layout.

From the shipped model to one Neural Engine partition

flowchart
From the shipped model to one Neural Engine partition same weights half precision point at shipped file granite weights (Hugging Face) ANE-layout module (PyTorch) fp16 ONNX export 5.7 MB graph, shared weights CoreML EP, static buckets 1 partition, ~1405 of 1407 ops on ANE

The changes follow the paper's principles, plus two that only this model needed:

What the ANE-layout re-export changes, and why
ChangeWhy the Neural Engine cares
Tensors stay (B, C, 1, S) from the first norm to the lastIts native format is 4D, channels first. Anything else means reshapes and copies.
Every Linear layer becomes a 1x1 Conv2dSame maths, but a 1x1 convolutionA convolution with a 1x1 kernel multiplies every position by the same weight matrix, which is exactly what a Linear layer does, in the layout the ANE is built for. is the operation it runs best.
LayerNorm computed over the channel axisFollows from the layout change.
Attention split per head and written as MatMulThe paper uses einsum, but ONNX Runtime's CoreML provider does not take Einsum: all 288 of them stayed on the CPU, in 25 partitions. With MatMul, one partition.
RoPERotary position embedding: position is encoded by rotating pairs of numbers in the query and key vectors. Its usual implementation negates half the vector.'s rotate_half folded into extra conv output channelsThe usual code uses a negation (Neg), and the 24 Neg ops alone split a plain re-export into 13 partitions.

Two things surprised me. First, precision decides everything. The first exports kept fp32 weights, and Core ML placed zero operations on the Neural Engine. The same module exported in fp16 put 1,402 of its 1,407 operations there. No warning, no error, just a different answer.

Second, fp16 bites. On a 1,000-token row the layer norm's squared values reach about 3.2 million, far past fp16's 65,504 ceiling. The export now computes the norm's statistics on x/64, and a test on a 1,000-token row goes red without that scaling. ONNX Runtime's CPU provider can't catch this class of bug because it quietly upcasts fp16 to fp32, so the test runs the module in real half precision in PyTorch. Core ML ran the unscaled graph without overflowing, probably because it does not literally square in fp16, but I did not want to rely on that.

# loading chart…

Buckets. Static shapes mean every row is padded to a fixed length. MLX on the GPU used 64-token steps. On the Neural Engine that was a bad trade: 16 buckets meant 3.1 GiB of compiled cache and 1.6 GiB of memory. 256-token steps (256, 512, 768, 1024) need four compiled models, 808 MiB of cache and 0.49-0.57 GiB of memory, and cost no extra latency. Every chunk size ai-raccoon ships lands exactly on one of them.

Size. A standalone re-export is another 98 MB of weights next to the 97 MB the package already ships. The ANE graph now matches every one of its 74 weights by value to a tensor in the shipped model_fp16.onnx_data (exact, transposed, or transposed with rotate_half's sign flips) and rebuilds them with constant operations that ONNX Runtime folds away at load time. The graph file is 5.7 MB and the shipped weights file stays byte-identical. That is also why the stored vectors don't change: same weights, same maths, cosine similarity of at least 0.999999 against the CPU path.

The re-exported model is public on Hugging Face as Arraasz/granite-embedding-small-english-r2-aneOur fp16 ONNX re-export of granite-embedding-small-english-r2 in Apple's Neural Engine layout, ready for ONNX Runtime's CoreML execution provider. Apache 2.0., under the same Apache 2.0 license as IBM's original, if you want to run granite on a Neural Engine without ai-raccoon.

Is it really on the Neural Engine?

Core ML's static plan had already disagreed with the runtime memory counters once, so this time I asked the hardware. `powermetrics`macOS's built-in power sampler. With sudo it reports the power drawn by the CPU, GPU and Neural Engine separately, in milliwatts. reports the power of each block of the chip separately. Over a 10-second idle window and 30 seconds on each path:

# loading chart…

The Neural Engine's power rail rises from 29 mW to 1.28 W only while the re-exported model runs, and it stays dark on the CPU path. In the same 30 seconds it embedded 887 rows against the CPU's 46.

I also watched it live. After switching my own server over, the system monitor I have been using, TMOG, showed the Neural Engine powering up in bursts, while the models compiled and later while new notes were embedded, and power-gated in between. Click the image to see it full size.

The Neural Engine on my live ai-raccoon server after opting in: bursts of work, power-gated in between.

A caveat on reading that screenshot: "blocks powered" is the share of the Neural Engine that is not power-gated, not how busy it is. The powermetrics numbers above are the stronger evidence.

The waiting problem: 37 seconds before the first answer

The measurements before the screenshot came from a test harness. Shipping it hit one more wall. Core ML compiles a model for the Neural Engine the first time it sees it, and for four buckets that took 37 seconds on an M4, 30 CPU-seconds of Apple's ANE compiler, and between 350 J and 950 J of energy in two runs, up to two and a half times what a full Neural Engine embedding run of the corpus costs. After that, the compiled cache loads in 1.0 s.

Nobody should wait 37 seconds for a search. So the compile never blocks anything. On start, ai-raccoon serves from WebGPU as it always did, compiles the four Neural Engine sessions in the background, and only switches when all four have loaded and a parity probe passes: one fixed row embedded on both devices, and the two vectors must have a cosine similarityHow closely two vectors point the same way: 1.0 means identical direction. Embedding search ranks by it. 0.999 is the bar ai-raccoon uses for treating two vectors as the same. of at least 0.999. Any failure, timeout or probe miss drops back to WebGPU with a logged reason, and ai-raccoon doctor shows which state the engine is in and why.

How the server moves onto the Neural Engine

state-machine
How the server moves onto the Neural Engine server start 4 buckets loaded, probe >= 0.999 load failure, timeout or probe miss WebGpuServing CompilingNeuralEngine NeuralEngineServing WebGpuServing (refused)

A few smaller things turned up along the way. I found that Core ML keys its compile cache by the compiled model's absolute path, so the benchmark's first design, which copied the cache into a fresh folder for each run, silently recompiled every time. A cache folder left half-written by a restart mid-compile is deleted and rebuilt once instead of crashing the session. And once both MLX and CoreML had been loaded in one process, ONNX Runtime's own telemetry thread aborted the process at exit (recursive_mutex lock failed). ai-raccoon never meant to run that telemetry, so it is now switched off before the first session is created.

The numbers, through the real product

The harness numbers were good, but they measured the encoder alone. The question that matters is what the whole server does. A new benchmark script drives the shipped product on each device: a fresh copy of the bank, a server on a scratch port, 225 documentation files split into 4,914 chunks, three repeats per device, interleaved so that no device gets the coolest part of the run. powermetrics measures the chip's energy; the battery controller's system-load counter measures the whole machine's. Both are net of idle.

# loading chart…

The Neural Engine needs 77 millijoules per chunk. WebGPU, today's default, needs 283, MLX 233 and the CPU 829. On the chip's own rails the gap is wider: 214 J for the whole corpus against 936 J on WebGPU, so the Neural Engine run uses 23% of the energy.

# loading chart…

I expected to trade some speed for the energy saving. Instead it is faster too: 33 seconds against 70 for WebGPU, even though it runs one row at a time. Every Neural Engine repeat beat every WebGPU repeat on time, system energy and chip energy, with no overlap in the ranges.

Server CPU time fell from 42 to 26 CPU-secondsProcessor time summed over all cores. One CPU-second is one core fully busy for one second, so 60 CPU-seconds in a minute means one whole core.. Against MLX it is 70% of the CPU-seconds, not the harness's 6-8%, because now it counts the whole server: HTTP, chunking, SQLite writes and the vector index. Only the model moved.

Search latency is where I expected to lose. A search embeds one short query, and the Neural Engine pads it up to 256 tokens. A second run added 50 fixed searches after each ingest:

# loading chart…

Every Neural Engine p95The 95th percentile: 95 of every 100 searches were at least this fast. It shows the slow tail that an average hides. (35-39 ms) came in under every WebGPU p95 (51-88 ms), and the typical search (p50) went from about 40 ms to about 30 ms. Search quality did not move: on all four evaluation sets, the Neural Engine's recall matches the CPU baseline to the third decimal.

What it costs you

Nothing is free. Opting in costs about 808 MiB of compiled cache in ai-raccoon's data folder, and a one-time background compile of about 37 seconds on an M4, repeated after an ONNX Runtime upgrade. The package itself grew by 5.7 MB. It only applies to the bundled model; a custom model you downloaded keeps running on the CPU, because nobody has re-exported it.

One slow repeat is worth mentioning. The first warm Neural Engine run took 45.7 s instead of 33, with the same CPU and energy as the fast ones, so the extra time did no work. My guess is the first load of the compiled program onto the Neural Engine, but I haven't verified it.

Caveats

This is one chip. Everything above was measured on one M4 (Mac16,12), and newer chips have faster Neural Engines, so I expect the result to hold there, but expecting isn't measuring. That is why coreml is opt-in and auto still picks WebGPU. Making it the default needs a second chip, and a release of people using it without problems.

System energy and chip energy come from two different meters, the battery controller and powermetrics. They agree on the order of the devices and roughly on the gaps (27% and 23% of WebGPU), which is why I trust them together more than either alone.

Windows and Linux NPUs are untouched. Qualcomm's path needs int8 or int16, which hurt this model's quality in an earlier test. Intel's OpenVINO NPU path accepts fp16 and is the next one worth trying, once I have hardware to try it on.

Try it

If you run ai-raccoon on Apple Silicon:

bash
ai-raccoon settings model device coreml
ai-raccoon serve --restart

Then ai-raccoon doctor shows the compile going from CompilingNeuralEngine to NeuralEngineServing. If you have a different Mac, the repo has a benchmark script that measures your chip the same way; the README explains how to run it and where to send the result. I would like to know what an M1 does.

What I take from it

The Neural Engine had been in my laptop the whole time, and the first honest attempt to use it said no. The measurements were right. What they were measuring was the wrong model shape. The fix wasn't a setting or a newer library, it was rewriting the model in the form the hardware was designed for, which Apple had written down in 2022.

The bigger lesson for me is where energy goes. On today's numbers, moving the model from the CPU to the GPU, the change from the last article, cuts the energy per chunk by two thirds. Moving it from the GPU to the Neural Engine cuts what is left by almost three quarters more. For a server that embeds in the background all day, on a laptop that is often on battery, that is the difference that matters.

ai-raccoon is open source (MIT): github.com/Arasz/ai-raccoonai-raccoon on GitHub (MIT): a local-first memory server for AI agents, with hybrid keyword and vector search over notes and code.. Install it with dotnet tool install -g ai-raccoon. The Neural Engine model is on Hugging Face at Arraasz/granite-embedding-small-english-r2-aneThe fp16 ANE-layout ONNX re-export of granite-embedding-small-english-r2. Apache 2.0., and the design and its reasons are in ADR-0118The accepted decision: an opt-in coreml device, background compile, WebGPU fallback, and the three gates before it can become the default..