Running Gemma4 on Apple Neural Engine
3 hours ago
- #CoreML Optimization
- #Apple Neural Engine
- #LLM Inference
- Apple Neural Engine (ANE) is an efficient computing unit in Apple devices, introduced with A11 in 2017, offering ~19 TFLOPS FP16 performance in A19 Pro (2026).
- ANE is optimized for convolution, not matrix multiplication (matmul), leading to hardware inefficiencies (~1/3 utilization) for LLMs like Gemma 4, requiring workarounds.
- Gemma 4 E2B (2B params, INT4 quantized) and E4B (4B params) were split into 4 CoreML chunks to fit under 1 GB memory on iPhone 17 Pro (8 GB RAM), enabling ANE-only inference.
- Chunking avoids CoreML conversion failures due to graph depth limits (8-10 Attention layers), but introduces dispatch overhead (~2.3 ms per CoreML call, totaling ~9.2 ms baseline).
- Optimizations via MLComputePlan diagnostics achieved 99.78% ANE operation placement for E2B, boosting speed from 11 tok/s to 31 tok/s, with calculations per chunk taking 10-16 ms.
- Quantization experiments showed INT4 works, but INT3 (W3A16) fails with garbled output due to insufficient codebook resolution (8 vs. 16 points) for Transformer weight distributions.
- Anomalies like +196608 offset in argmax outputs revealed ANE's internal 4-parallel argmax structure, highlighting inference challenges from undisclosed hardware behaviors.
- Reducing chunks from 4 to 2 only increased speed to 31.7 tok/s, showing dispatch tax is minor compared to computation costs, with c4 (lm_head) being the heaviest chunk (~30% of time).
- Comparisons: LiteRT-LM on Metal GPU achieves 56.5 tok/s, 1.8x faster than ANE, due to factors like single command buffer, kernel fusion, and native matmul support, which CoreML can't match.
- Future prospects include porting models like Qwen 3.5 (with RNN/attention hybrid) and potential iOS 27 'Core AI' framework opening ANE instruction queue, which could overcome current constraints.
- The project documents failures (e.g., W2A16 gibberish, W8A8 crashes, INT8 KV cache no speed gain) as valuable logs for optimization, emphasizing experiential understanding over specifications.