Optimizing On-Device Inference for Apple Silicon

gist

Perplexity built Lily, a Rust engine with custom Metal kernels written for one model, Qwen3.6-35B-A3B, on Apple silicon. On a 128 GB M5 Max, Perplexity measures it at 1.22× MLX-LM’s average prefill throughput and 1.32× its decode, faster at every length from 256 to 128K tokens, with neither PyTorch nor MLX in the execution path.

ideas

• One model, one chip, one engine. MLX-LM’s reusable operations must serve many architectures; Lily plans kernels, data movement, and scheduling around Qwen’s fixed mixture-of-experts and recurrent layers, in a single Rust process behind an OpenAI-compatible API.
• Prefill gains came from removing work around the experts. Fusing 4-bit dequantization into the grouped expert GEMM raised prefill 77.4% at a 512-token prompt, and keeping expert routing on the GPU raised it 89%, because expert GEMMs take about 90% of prefill time.
• Long-context decode is an attention problem. Packing the eight query heads that share a KV head raised decode 23.8% at 32K tokens, and a fixed-block attention path added 7.7% at 32K, 27.4% at 64K, and 40.2% at 128K.
• Speculative decoding did not help here. At batch 1 it made decode 18% slower: verifying two to five rows is an awkward shape for the chip, and the rows often picked different experts, so more weights were read.
• The kernels are near the hardware limit. Expert GEMMs and GEMVs reach 97.9% and 90.3% of the fastest sustained weight-read rates; teacher-forced perplexity was 0.04% higher than MLX-LM, with the same top token at 96.35% of positions.

quotes