Run decoder on Apple Neural Engine

[chenqianhe]

TL;DR: I got an experimental Core ML decoder path running on ANE. The working design uses sharded MLProgram models, self-K/V in MLState, and explicit FP16 cross-K/V inputs. Correctness now matches the tested CPU path, but the end-to-end result is not yet faster than encoder-only Core ML.

---

This work started from the old limitation discussed in Core ML Support #548: the encoder is a mostly static graph and works well with Core ML/ANE, while the decoder has dynamic token-by-token behavior and KV-cache state, which historically made it hard or inefficient to place on ANE. Newer Core ML stateful model support changes part of that tradeoff. Apple’s Core ML documentation describes MLState as intended for recurrent/stateful workloads such as transformer KV caches, and Apple’s on-device Llama writeup also points out that stateful KV cache can avoid repeated KV I/O.

The goal here was specifically to test whether the decoder can run on ANE (CPU_AND_NE), not merely whether it can run through Core ML. That distinction matters because MLComputeUnits.all and CPU_AND_GPU often place the decoder on GPU/MPSGraph, which can be correct and fast but does not answer the ANE question.

Environment used in the main experiments

• model: ggml-large-v2.bin (For my personal needs, I chose the large model, and fortunately, this helped me discover earlier some issues that needed to be resolved and would not appear in decoders with fewer layers.)
• sample: samples/jfk.wav
• platform: Apple Macbook M4 Max (128G)
• Core ML decoder compute units tested: all, cpu_gpu, cpu_ne, cpu_only

(Failed) full stateful Core ML decoder graph with KV state

The first direct attempt was a full stateful Core ML decoder graph with KV state. That did not work on ANE. With CPU_AND_NE, Core ML failed ANE compilation with an error like:

MILCompilerForANE ... ANECCompile() FAILED
error -14

This was the first major design constraint. The full decoder graph was too large or too awkward for ANE compilation when self-attention, state reads/writes, cross-attention, MLPs, residuals, and logits were kept together.

I then compared placement behavior against the encoder. For large-v2, the encoder Compute Plan placed roughly 2443 ops on MLNeuralEngineComputeDevice in both cpu_ne and all. The decoder did not behave the same way. Full decoder under cpu_gpu or all generally preferred GPU/MPSGraph, while cpu_ne failed for the full graph.

Several reduced experiments were tried:

• A small update_union style graph that only updated state could run but preferred CPU.
• A read_cross_mlp style graph could place on ANE.
• A split attention-only plus MLP-only layer showed only a small improvement over CPU for a single layer, roughly 6 percent in one test, and the 32-layer estimate was still much worse than the GPU/MPSGraph decoder path.

This led to the current shard design.

Instead of one large decoder graph, the decoder is exported as no-write shards. Each shard reads self-K/V state and emits new K/V outputs, but the runtime writes those outputs back to the corresponding MLState. This avoids asking Core ML to compile the whole state-update decoder as one ANE graph.

8-layer shards working large-v2 experiment

The first working large-v2 experiment used four 8-layer shards:

s0-l8-token
s8-l8
s16-l8
s24-l8-logits

The first shard consumes token IDs plus position information. Middle shards consume hidden states. The final shard emits logits. The runtime chains the shards and copies final logits back into whisper.cpp.

This compiled, but the first true cpu_ne run produced wrong logits. A 1-token prompt diagnostic showed:

top_cpu=50259
top_coreml=50362

A direct shard diagnostic showed divergence starting in s0 x_out, with approximately:

max_abs=0.125
mean_abs=0.005137

At the same time, all mode matched much better. For example, a small s0-l1 comparison showed all with x_out max_abs=0.008269 mean_abs=0.001427, while cpu_ne had much larger error around max_abs=1.754883 mean_abs=0.288577. But all was not an ANE solution because the Compute Plan preferred GPU.

The root cause was cross-attention KV handling. Keeping cross V in MLState caused CPU_AND_NE divergence. The fix was to keep self-K/V in MLState, but pass cross-K/V as explicit FP16 MLMultiArray inputs to each shard. After that change, the cpu_ne path matched the CPU path.

Difference between MLMultiArray and MLState (My personal analysis and conjecture; may not be correct.)

The key difference between MLMultiArray and MLState is not the dtype, but how the Core ML runtime feeds them into the execution graph.

MLMultiArray is a regular input/output tensor. Each time predict is called, it is passed as an MLFeatureValue, and Core ML treats it as an input feature for that invocation. Apple’s documentation defines it as a multidimensional numeric array type, with support for data types such as FP16 and FP32.

MLState, on the other hand, is a persistent state handle for stateful models. Core ML allocates state buffers for the states declared by the model; during inference, the model is invoked via prediction(... using: state ...), and the graph reads those buffers internally through ops such as read_state. Both Apple and coremltools document MLState as the mechanism for preserving and updating model state across multiple predictions.

As a result, even if both eventually contain FP16 buffers underneath, they enter the ANE execution path differently:

• MLMultiArray input: a regular graph input, lowered according to input-feature shape, stride, and placement rules.
• MLState: a mutable persistent resource, read inside the graph via read_state, with its lifecycle, aliasing, caching, device-side state buffers, and possible state synchronization managed by Core ML.
• MLState also implies reuse of the same state across multiple prediction calls, which adds more constraints for the ANE compiler/runtime.

In our case, the issue does not look like "FP16 data was stored incorrectly". It looks more like a limitation or bug in how ANE lowers or places certain large MLState read patterns.

The main reasons I do not believe this is simply a mistake on our side are:

• The same data and logic match correctly under cpu_only.
• After moving cross K/V to regular FP16 inputs, cpu_ne produces aligned top-token/logit results.
• The fully stateful decoder fails under cpu_ne with an ANE compiler -14 error, suggesting that this state-graph pattern itself hits an ANE compiler limitation.
• In the state-backed no-write shard, the most obvious mismatch is concentrated around cross V. If this were merely an FP16 precision issue or an incorrect host copy, the CPU path would typically fail as well.

That said, this does not mean that MLState is generally broken in Core ML. A more accurate conclusion is:

Feeding very large cross-K/V tensors — especially the 1280 x 1500 cross V tensor — into an ANE shard as read-only MLState is semantically reasonable at the API level, but it is not a reliable execution path in the current Core ML/ANE runtime.

MLState is most naturally suited for decoder self-K/V caches, which are true cross-token accumulated states. Cross-K/V, however, is derived from the encoder output and is effectively a per-segment constant. It does not strictly need to be modeled as state. We originally placed it in MLState to unify KV-cache management and reduce the number of per-step input arguments. In practice, this ANE state-read path turned out to be unstable, so we moved cross-K/V back to regular inputs.

This defines the boundary of the current design:

• self-K/V remains in MLState, because it is genuinely cross-token state and has been verified to align under cpu_ne.
• cross-K/V is passed as MLMultiArray inputs, because it is fixed context for each audio segment, better matches regular input semantics, and avoids the problematic ANE read_state path.

In short, this is not a case of FP16 being unusable, nor is it necessarily an obvious API misuse. It is more likely an implementation limitation or edge-case bug in Core ML stateful MLProgram execution on ANE for large read-only cross states. The current implementation therefore routes cross-K/V through the regular input path, which has been empirically verified to remain stable and aligned during real greedy decoding.

References:

• Apple MLMultiArray documentation: https://developer.apple.com/documentation/CoreML/MLMultiArray
• Apple MLState documentation: https://developer.apple.com/documentation/CoreML/MLState?changes=_5
• Core ML Tools stateful models guide: https://apple.github.io/coremltools/docs-guides/source/stateful-models.html

Final parity checks after the explicit cross-K/V input fix:

8/8 top tokens matched between CPU_ONLY and CPU_AND_NE
mean abs range: 0.013824 - 0.037927
max abs range:  0.101562 - 0.177246

Compilation and placement results

All four decoder shards compile successfully with CPU_AND_NE. The resulting preferred placement is overwhelmingly on the Neural Engine:

Shard	Package	cpu_ne compile	Preferred placement
0–7	s0-l8-token	OK	NE 1409 ops, CPU 5 small token-index ops
8–15	s8-l8	OK	NE 1405 ops
16–23	s16-l8	OK	NE 1405 ops
24–31	s24-l8-logits	OK	NE 1409 ops

The 5 CPU ops in the first token-input shard are small embedding/position-indexing operations, such as greater_equal, add, select, gather, and cast. The main compute-heavy parts of the decoder — attention, MLP, layer normalization, softmax, and read_state — are all preferred on the Neural Engine.

With the corrected sharded decoder running under CPU_AND_NE, using four shards connected sequentially and a 16-token probe, the observed prediction timings are:

Metric	Time
First token	48.838 ms
Warm token p50	29.432 ms
Shard 0 p50	7.007 ms
Shard 1 p50	6.927 ms
Shard 2 p50	6.968 ms
Shard 3 p50	8.085 ms

The short JFK sample matched the CPU transcript. Representative decoder timings for the final ANE shard path were around:

samples/jfk.wav:
first token: ~34.962 ms
warm p50:    ~35.032 ms

longer local sample:
first token: ~35.930 ms
warm p50:    ~35.262 ms

Compute Plan inspection for the final four-shard path showed the important decoder operations, including einsum, softmax, conv, layer_norm, read_state, and final logits matmul/reshape, preferred MLNeuralEngineComputeDevice.

sequenceDiagram
    participant App as whisper.cpp decode loop
    participant CPU as Host CPU / ggml state prep
    participant S0 as Core ML shard0 L0-7
    participant S1 as Core ML shard1 L8-15
    participant S2 as Core ML shard2 L16-23
    participant S3 as Core ML shard3 L24-31
    participant Sampler as greedy sampler

    App->>CPU: prepare token, pos, slot_mask, self_mask
    App->>CPU: copy cross-K/V into shard inputs once per segment
    loop each prompt or generated token
        CPU->>S0: token/pos or hidden + masks + cross-K/V + MLState self-K/V
        S0-->>CPU: hidden + k_out/v_out for L0-7
        CPU->>CPU: write k_out/v_out back to shard0 MLState

        CPU->>S1: hidden + masks + cross-K/V + MLState self-K/V
        S1-->>CPU: hidden + k_out/v_out for L8-15
        CPU->>CPU: write k_out/v_out back to shard1 MLState

        CPU->>S2: hidden + masks + cross-K/V + MLState self-K/V
        S2-->>CPU: hidden + k_out/v_out for L16-23
        CPU->>CPU: write k_out/v_out back to shard2 MLState

        CPU->>S3: hidden + masks + cross-K/V + MLState self-K/V
        S3-->>CPU: logits + k_out/v_out for L24-31
        CPU->>CPU: write k_out/v_out back to shard3 MLState
        CPU->>Sampler: copy logits
        Sampler-->>App: next token
    end

Loading

Final shard plan

The shard plan also changed. The original fixed 8-layer plan was not the best fit across models. A dynamic shard planner now chooses by decoder layer count:

tiny:    [4]
base:    [6]
small:   [12]
medium:  [12,12]
large:   [12,8,12]
turbo:   [4]

Some measured partition data from the exploration:

tiny [2,2]:              p50 9.861 ms, first 12.395 ms
base [6]:                p50 17.536 ms, first 25.954 ms
small [8,4]:             p50 51.207 ms, first 77.908 ms
medium [4,4,4,4,4,4]:    p50 121.788 ms, first 193.150 ms
large-v2 [12,8,12]:      p50 205.893 ms, first 288.619 ms
large-v2 [8,8,8,8]:      p50 207.749 ms, first 340.526 ms
large-v3-turbo [4]:      p50 30.31 ms, first 42.87 ms
large-v3-turbo [2,2]:    p50 27.46 ms, first 42.63 ms

The current selected plan favors simpler artifact layout and practical runtime behavior rather than only the absolute best isolated microbenchmark in every case.

The other bottlenecks

Multi-candidate decoding

The first working Core ML decoder path was greedy-only. That was enough to prove that ANE execution and transcript parity were possible, but it was not enough for normal whisper.cpp behavior. The default whisper_full path may use beam search, best-of sampling, temperature fallback, grammar filtering, and multiple candidate decoders. Each candidate has its own token history, and therefore each candidate must also have its own self-attention KV state.

The ggml decoder already has logic to copy and reorder KV cache entries when candidates are cloned, pruned, or reassigned. The Core ML decoder needed equivalent state management. Reusing one Core ML MLState across candidates is incorrect: once candidate A and candidate B choose different next tokens, their self-K/V states diverge. If they share the same state, later logits are computed from a mixed history and beam search becomes semantically wrong.

The runtime now keeps an array of Core ML decoder contexts, one per active decoder/candidate. After the prompt is evaluated once, the prompt state is copied from candidate 0 into the other active candidates. During beam search, when the best candidates are selected and moved between decoder slots, the runtime copies only the self-state needed for the candidate transition. Cross-K/V does not need to be recopied in those self-state-only moves because it is fixed for the segment and shared by value in the explicit input arrays.

That distinction is reflected in the API shape:

• whisper_coreml_decoder_new_like(...) creates a full decoder context copy, including cross inputs;
• whisper_coreml_decoder_new_like_self(...) creates a decoder context for self-state use without duplicating cross inputs unnecessarily;
• whisper_coreml_decoder_copy_state(...) copies both self-state and cross inputs;
• whisper_coreml_decoder_copy_self_state(...) copies only self-attention state.

This made the Core ML decoder compatible with the normal logits-based decoding paths: beam search, best-of, temperature sampling, temperature fallback, and grammar penalties.

A related correctness issue was how Core ML state and input arrays were copied. The first implementation treated MLMultiArray memory too much like a flat C buffer. That is risky because MLMultiArray has shape and stride metadata, and the memory layout is not something the runtime should blindly assume in every case. A raw memcpy can be correct for a contiguous array with the exact expected layout, but it is not a safe general state-copy operation.

The fix was to make copying layout-aware. For state-to-state copies, the runtime uses Core ML’s transferToMultiArray: when available, after checking data type, element count, and shape. For F16 host data copied into Core ML arrays, the runtime checks the expected 4D ANE layout and writes by channel/time strides. The same care is used when writing the k_out_* and v_out_* outputs from a no-write shard back into the correct self-K/V state slot. This avoids silent corruption if the array is valid but not laid out as a trivial contiguous buffer.

No-speech probability

no_speech_prob is computed immediately after the prompt decode, before logit filtering. In the CPU path, the prompt decode produces logits for each prompt token, and the probability must be computed from the final prompt row. The Core ML prompt path also produces logits for the prompt, but the first version did not consistently read the same row that the CPU path uses. That can give a wrong no_speech_prob, which affects segment rejection/fallback decisions even if the generated transcript looks correct.

The fix was to align the Core ML prompt logits layout with the existing whisper_full logic: after prompt decode, state->decoders[0].i_batch is set to prompt.size() - 1, and no_speech_prob is computed from state->logits + i_batch * n_vocab. This keeps the Core ML decoder path semantically aligned with the CPU decoder path before any sampling, grammar, timestamp, or fallback logic runs.

End-to-end results

End-to-end results are more mixed. This path is a functional ANE decoder path, but it is not automatically faster than the existing encoder-only Core ML path.

With default-style decoding parameters (-bs 5 -bo 5 -tp 0.0) on samples/jfk.wav, 3-run median:

Pure CPU:
median total: 4172.87 ms
realtime:     2.64x

Core ML encoder only:
median total: 2638.28 ms
speedup:      1.58x vs CPU
realtime:     4.17x

Core ML encoder + Core ML decoder:
median total: 6939.46 ms
relative:     1.66x slower than CPU
realtime:     1.59x

With greedy settings:

-bs 1 -bo 1 -tp 0 -nf

the results were:

Pure CPU:
median total: 3313.39 ms
realtime:     3.32x

Core ML encoder only:
median total: 1775.75 ms
speedup:      1.87x vs CPU
realtime:     6.19x

Core ML encoder + Core ML decoder:
median total: 2809.98 ms
speedup:      1.18x vs CPU
realtime:     3.91x

So the decoder path can be correct and can run on ANE, but on this JFK benchmark it is still slower than the CPU decoder inside the encoder-only Core ML configuration. The main benefit demonstrated here is feasibility and placement correctness, not a clear end-to-end performance win yet.

Current known limitations:

• Requires full audio context.
• Requires the expected F32/F16 paired ggml/Core ML model setup.
• --dtw token timestamps are not supported yet because the Core ML decoder shards currently expose logits/state, not the cross-attention alignment tensors needed by DTW.
• -ac partial audio context is not supported because the shards currently use fixed full n_audio_ctx cross-K/V inputs. Supporting -ac correctly would need a cross-attention mask, flexible-shape model support, or separate artifacts per audio context.

To reproduce the current path for large-v2:

./models/generate-coreml-model.sh --decoder large-v2
cmake -B build -DWHISPER_COREML=1 -DWHISPER_COREML_DECODER=1
cmake --build build --target whisper-cli -j
./build/bin/whisper-cli -m models/ggml-large-v2.bin -f samples/jfk.wav -bs 1 -bo 1 -tp 0 -nf

Decoder-only build is also possible:

cmake -B build -DWHISPER_COREML_DECODER=1
cmake --build build --target whisper-cli -j

Code details

#3848

There was also a small but meaningful moment of reflection when I opened a new pull request: GitHub reminded me that my last contribution to the project was on February 2, 2023. Even more strikingly, I noticed that some of the ongoing maintenance today is still connected to changes I made back then. It is a quiet reminder that code can have a longer life than we often expect — shaped by new hardware, new tooling, and new constraints, yet still carrying traces of past decisions forward.

---

Created at 01:33 JUN 1, 2026 (UTC+8)