Running Gemma4 on Apple Neural Engine
3 hours ago
- #CoreML Optimization
- #Apple Neural Engine
- #LLM Inference
- Apple Neural Engine（ANE）是苹果设备中的高效计算单元，随A11芯片于2017年推出，在A19 Pro（2026款）上可提供约19 TFLOPS的FP16性能。
- ANE针对卷积而非矩阵乘法（matmul）优化，导致在处理像Gemma 4这样的大型语言模型时硬件效率较低（约1/3利用率），需要变通方案。
- Gemma 4 E2B（20亿参数，INT4量化）和E4B（40亿参数）模型被分割成4个CoreML块，以适应iPhone 17 Pro（8GB RAM）中1GB内存的限制，从而实现仅使用ANE的推理。
- 分块避免了因图深度限制（8-10层Attention）导致的CoreML转换失败，但引入了调度开销（每次CoreML调用约2.3毫秒，总计约9.2毫秒基线）。
- 通过MLComputePlan诊断进行优化，E2B模型实现了99.78%的ANE操作分配，将速度从11 tok/s提升至31 tok/s，每个分块的计算耗时10-16毫秒。
- 量化实验表明INT4有效，但INT3（W3A16）因代码本分辨率不足（8个点 vs. 16个点）无法适应Transformer权重分布，导致输出乱码。
- 诸如argmax输出中出现+196608偏移等异常现象，揭示了ANE内部4并行argmax结构，突显了未公开硬件行为带来的推理挑战。
- 将分块从4个减少到2个仅将速度提升至31.7 tok/s，表明调度开销相对于计算成本较小，其中c4（lm_head）是最重的分块（约占总时间的30%）。
- 对比显示：LiteRT-LM在Metal GPU上达到56.5 tok/s，比ANE快1.8倍，这得益于单一命令缓冲区、内核融合和原生矩阵乘法支持等优势，而CoreML无法匹敌。
- 未来展望包括移植如Qwen 3.5（采用RNN/注意力混合架构）等模型，以及潜在的iOS 27“Core AI”框架开放ANE指令队列，这可能突破当前限制。
- 该项目记录了失败案例（如W2A16乱码、W8A8崩溃、INT8 KV缓存无速度提升）作为优化的重要日志，强调实践经验重于规格说明。