flash-moe-inference

Run 397B parameter Mixture-of-Experts LLMs on a MacBook using pure C/Metal with SSD streaming

By reason-machines · 1,129 installs

npx skills add reason-machines/trending-skills --skill flash-moe-inference

Source repository · Upstream listing

Flash MoE Inference Engine Skill by [ara.so](https://ara.so) — Daily 2026 Skills collection. Flash MoE is a pure C/Objective C/Metal inference engine that runs Qwen3.5 397B A17B (397B parameter Mixture of Experts) on a MacBook Pro with 48GB RAM at 4.4+ tokens/second. It streams 209GB of expert weights from NVMe SSD on demand — no Python, no ML frameworks, just C, Objective C, and hand tuned Metal shaders. Requirements Hardware : Apple Silicon Mac (M3 Max or similar), 48GB+ unified memory, 1TB+ SSD with ~210GB free OS : macOS 26+ (Darwin 25+) Tools : Xcode Command Line Tools, Python 3.x (for weight extraction only) Model : Qwen3.5 397B A17B safetensors weights (download separately from HuggingFace) Installation & Build The Makefile compiles infer.m , chat.m , main.m with Metal shader compilation for shaders.metal . Weight Preparation Step 1: Extract non expert weights Step 2: Pack expert weights (4 bit, production) Step 3: Optional 2 bit requantization (faster but breaks JSON/tool calling) Key Commands Basic inference Interactive chat with tool calling MoE only benchmark (measures expert throughput) Project Structure Architecture Overview The model has 60 transformer layers : 45 GatedDeltaNet (linear attention) layers 15 standard full attention layers Each layer: 512 experts, K=4 activated per token + 1 shared expert Hidden dimension: 4096 Per layer pipeline (4.28ms average at 4 bit) Metal Shader Kernels The shaders.metal file contains hand written kernels. Key kernels: SSD Expert Streaming Pattern The core innovation — loading only K=4 active experts per layer from SSD: Why pread() not mmap() : mmap incurs per page fault overhead on cold data (~5x slower). Direct pread() with OS page cache achieves ~71% hit rate naturally. GatedDeltaNet Linear Attention (BLAS) The recurrence update uses Accelerate BLAS — 64% faster than scalar: Performance Configuration 4 bit (production default) Quality : Excellent — full tool calling, correct JSON Speed : 4.36 tok/s Disk : 209GB 2 bit (speed testing only) Quality : Good — but breaks JSON/tool calling ( \name\ instead of "name" ) Speed : 5.74 tok/s (7.05 peak single token with warm cache) Disk : 120GB Uses F NOCACHE flag to avoid page cache thrashing What NOT to Try (Learned from 58 Experiments) Approach Why it fails mmap() expert files Per page fault overhead: 5x slower than pread() dispatch io dispatch data management overhead: 70% F RDADVISE prefetch SSD DMA + GPU share memory controller — concurrent access: 73% GPU speed Custom Metal LRU cache GPU memory pressure: 38% vs OS page cache LZ4 expert compression Decompress overhead warm cache savings: 13% Temporal expert prediction 25% hit rate, wastes SSD bandwidth: 18% Speculative early routing Cache pollution: 38% MTP speculative decoding MoE I/O scales per token (unlike dense models): break even Spin poll GPU wait CPU thermal throttle competes with GPU: 23% Parallel SSD + GPU overlap Unified memory controller arbitration: net negative Key principle : On Apple Silicon, GPU DMA and SSD DMA share the same memory controller. The serial pipeline (GPU → SSD → GPU) is hardware optimal. Troubleshooting Build fails Out of memory The engine is designed to use ~6GB active: 5.5GB: model weights.bin (mmap'd, read only) ~200MB: Metal scratch buffers Remaining ~42GB: OS page cache for expert data If you see OOM, check for other processes consuming unified memory: Slow performance Wrong expert directory Tool calling broken Use 4 bit, not 2 bit. The 2 bit quantization corrupts quote characters in JSON output, making tool calling unreliable. Always use the default 4 bit configuration for agentic workloads. Memory Safety The engine explicitly manages all allocations: No unbounded caches Expert data never accumulates in GPU memory model weights.bin is mmap'd read only — kernel manages pages Expert files are opened/read/closed per inference step