torch-ane

Run PyTorch models on the Apple Neural Engine through Apple's Core AI (macOS 27+).

Pass in a plain nn.Module, or a model from transformers or diffusers. torch-ane rewrites it into a form the ANE runs fast and correctly: fp16, static shapes, and the layout and ops from Apple's ANE guidance. It also works around the ANE's numeric quirks. You get back a module that behaves like the original.

import torch_ane

m = torch_ane.compile(model.eval(), example_inputs=(x,))  # any nn.Module -> ANE
y = m(x)                                                   # torch tensors in, torch tensors out

On an M4 MacBook, ANE prefill and convolutions are 2–3× faster than the GPU (MLX/MPS). Decoding is limited by memory bandwidth, and the ANE reads ~57 GB/s.

Install

Requires macOS 27+ on Apple silicon, and Python 3.11+.

pip install torch-ane              # core: compile()
pip install "torch-ane[llm]"       # + CausalLM / ChunkedCausalLM (transformers, gguf)
pip install "torch-ane[diffusion]" # + StableDiffusion (diffusers)

Core AI's Python packages (coreai-core, coreai-torch, coreai-opt) come from PyPI as dependencies.

What's inside

import torch_ane

torch_ane.available()          # is Core AI usable here

# Any module. The first call per input shape builds a program (seconds); then it's cached on disk.
m = torch_ane.compile(model, example_inputs=(x,), device="ane")   # or device="gpu"

# A whole LLM on the ANE (up to ~1-3B). The KV cache is Core AI state, so it is never copied.
lm = torch_ane.CausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", batch=True)
ids = lm.generate(prompt_ids, max_new_tokens=64, eos=tok.eos_token_id)
outs = lm.generate_batch([ids1, ids2, ids3], max_new_tokens=64)   # up to 8 streams per step
logits = lm.feed(more_ids)                                        # append to the cache -> last logits

# A large model from GGUF, in 4-layer chunks with 4-bit palettes (Devstral 24B etc.)
big = torch_ane.ChunkedCausalLM.from_gguf("model.gguf", "config_dir/")

# Stable Diffusion: text encoder, UNet and VAE on the ANE
sd = torch_ane.StableDiffusion("segmind/tiny-sd")
image, timings = sd("a cute robot", steps=20, seed=0)

What it does to your model

These passes need no knowledge of the specific model:

1. RMSNorm, found by a numerical check rather than by name, becomes a Core AI composite op.
2. Everything that doesn't depend on the input (masks, RoPE tables) is precomputed in fp32.
3. A causal mask becomes is_causal.
4. SiLU becomes x·0.5·(1 + tanh(x/2)). The native ANE SiLU is off by ~5% near zero, which wrecks GroupNorm→SiLU in diffusion UNets.
5. The model is converted to fp16, the only float type the ANE runs.
6. MLPs whose projection weight exceeds 4 MB are split along the hidden dimension. Each chunk of up feeds its own chunk of down, and the results are summed. Otherwise the weight plus the [S, hidden] activation doesn't fit the ANE's on-chip memory (~32 MB on M4), and the ANE stalls on the bus at ~50 GB/s. This cuts the Whisper large encoder by 24% and Qwen2.5-0.5B prefill by 14%. Disable it with rules=().

LLMs get two more passes:

• Per-projection 2^k scales, calibrated on real text, so small activations don't become fp16 subnormals while large ones don't overflow.
• int16 positions, embeddings on the CPU, and logits for one row only.

UNet/VAE attention uses Apple's layout: BC1S, heads split, einsum with no transposes.

Speed (MacBook M4, 24 GB)

Examples

• examples/compile_bert.py: BERT-base, eager vs ANE
• examples/generate.py: chat with Qwen2.5 on the ANE
• examples/chat_gguf.py: chat with a 24B GGUF model, with the engine in a restartable process
• examples/imagine.py: Stable Diffusion

Limitations

• Shapes are static, so a new input shape means a new build.
• The ANE computes in fp16 only (weights can be palettized or int8). -inf in masks is replaced by −40000.
• int8 weights in Conv1d make the ANE compiler fail silently, and the whole function moves to the GPU. compress="int8" therefore quantizes only Linear/Conv2d.
• The coreai-core Python wrapper aborts the process after 16 384 allocated buffers (apple/coreai-torch#75). ChunkedCausalLM.steps_left() lets you restart ahead of time; see examples/chat_gguf.py.
• The ANE unloads programs after ~10 s idle, which makes the next call ~3× slower. Use ChunkedCausalLM.keep_warm() to avoid this.
• Copying an input into Core AI runs at ~0.5 GB/s, about 10 ms per 4 MB input. The public API has no way around this.
• Code comments are in Russian.

Tests

pip install -e ".[llm,diffusion,dev]"
pytest

License

MIT