AX Engine: A Deeply Optimized Local LLM Inference Engine for Apple Silicon

This article introduces AX Engine, a Rust-based inference engine natively built for Apple Silicon M3+ chips. Through Transformer-specific kernel fusion, deep optimization of Apple's UMA (Unified Memory Architecture), and model-aware execution plans, it achieves performance superior to general engines on supported models. AX Engine is positioned as an exclusive dedicated inference engine for Apple Silicon, complementing rather than competing with general engines like llama.cpp.

Local LLM inference is an important trend in AI application development, but existing general-purpose cross-platform engines (e.g., llama.cpp) are constrained by generality and struggle to fully unleash the performance of specific hardware. AX Engine adopts the approach of "Transformer-specialized, Apple-native", targeting the niche market of carefully curated Transformer model families on Apple Silicon Macs. It extracts hardware performance through Apple-exclusive decisions (such as integrated Metal scheduling, UMA buffer contracts, etc.) and does not aim to replace general engines.

AX Engine uses a Transformer-specific runtime and reduces scheduling overhead and memory traffic through selective fusion: 1. Attention preparation phase fusion (QKV splitting + bias addition + per-head QK normalization + RoPE + KV cache appending); 2. Residual + normalization fusion; 3. Activation + down-projection fusion; 4. Selective FFN pair kernels. Additionally, it leverages Apple's UMA architecture to implement mmap-supported zero-copy Metal buffer aliases, optimizing memory paths and reducing model loading copy overhead.

AX Engine is optimized for different model families: LLaMA3.1 benefits from deep fusion, Qwen3.5 supports hybrid attention + SSM, and Gemma4 benefits from per-head QK normalization fusion. Compared to llama.cpp: AX Engine is positioned as Apple Silicon-exclusive, prioritizes optimization over generality, supports a curated subset of Transformers, and can make Apple-exclusive decisions; while llama.cpp is a cross-platform portable runtime covering a wide range of GGUF models.

AX Engine offers experimental features: 1. Speculative decoding (accelerated via draft models, rolling back rejected tokens after batch verification); 2. Concurrent decoding mode (optional Metal concurrent scheduling); 3. Split-K decoding attention (improves long-context efficiency); 4. Performance analysis-based tuning (per-model heuristics and decoding mechanism routing).

Developing AX Engine requires macOS M3+, Xcode, and Rust 1.88+. The insights it brings: 1. Dedicated engines can complement general engines in specific scenarios; 2. Memory architecture should be a first-class optimization target; 3. Fusion should be chosen pragmatically to avoid over-optimization.

AX Engine represents the in-depth development of local LLM inference towards platform-specific optimization, filling the performance gap for specific models on Apple Silicon. It is an efficient option for Apple users and demonstrates the possibility of dedicated optimization for developers. In the future, deep optimization for different hardware architectures will be an important direction to improve AI computing efficiency.