Partitioning MLPs Along the Intermediate Dimension Let the Apple Neural Engine Run Two of Three Guidance Branches

We ported Diffio 4.0, our diffusion-transformer speech restoration model, to a Mac mini with an M4 Pro (12-core CPU, 16-core GPU, 24 GB of unified memory). It restores audio at 7.8x realtime with the pro sampler and 10.7x with flash, against about 68x for pro on an RTX PRO 6000 in our CUDA deployment, and matches the CUDA run's quality within seed noise on word error rate, speaker similarity and UTMOS. Each guided sampler step runs on the GPU, the CPU's SME matrix units and the Neural Engine at once. The Neural Engine could take its share only after every MLP in its copy of the model was partitioned along the intermediate dimension into independent blocks, which more than doubled its MLP throughput.

A guided step is three independent forward passes

The model is a diffusion transformer with a Qwen2 trunk. It restores audio in overlapping 30 s windows, each conditioned on semantic features from the Granite speech encoder, on observation latents (the degraded input encoded by the Ming AudioVAE codec) and on the transcript text. The samplers are stochastic rectified-flow samplers (flow matching covers the velocity field they integrate) with classifier-free guidance. A guided step evaluates three branches: u (no conditions), so (semantic and observation, no text) and sot (all conditions). Late steps are unguided and evaluate sot alone.

The branches do not interact until the guidance arithmetic combines their velocities, so a step can be divided by branch across engines and lasts as long as the slowest engine's share.

guided step    (~235 ms)   GPU + CPU       sot    each MLP split between GPU (bf16) and CPU (fp32)
                           Neural Engine   so, u  one after the other
unguided step  (~215 ms)   GPU + CPU       sot
                           Neural Engine   codec work for neighbouring windows
per window                 GPU (PyTorch)   Granite, the condition encoders and the text encoder

MLX's GEMMs reach 89% of the GPU's peak, so kernels were not the bottleneck

Moving the transformer from PyTorch's MPS backend to MLX was the first large gain, because MLX fuses the elementwise work around each GEMM into fewer kernels; it also cut the relative L2 error of the output velocity against an fp32 forward pass from 6e-2 to 5e-3. Compiling each layer's attention block with mx.compile helped a little more. The GEMMs themselves had nothing left to give: Metal 4's tensor-op multiply (matmul2d) ran within 1% of MLX at our shapes, and PyTorch's torch.compile made the MPS forward pass 70% slower.

Splitting each MLP between the GPU and the CPU's SME units makes all CPU work compete with the model

MLX evaluates a GPU stream and a CPU stream concurrently. We partition each MLP along its intermediate dimension: the GPU computes 70% of the intermediate columns in bf16, the CPU's SME units compute the other 30% in fp32 (bf16 elementwise math measured four times slower than fp32 on the CPU), and the two partial down-projections are summed. In a microbenchmark this cut an unguided forward pass from 236 ms to 208 ms. The best share, found by sweep, gives the GPU more than the standalone rates would, because neither engine keeps its standalone rate while both run: they share one power budget, and the GPU's clock drops from 1578 MHz to 1.43–1.54 GHz while the CPU computes.

Because the CPU is part of every layer, any other CPU work lengthens the steps it overlaps. The codec decoder's final ISTFT stage cost 110 ms of numpy per window, and moving it to the GPU (11 ms) sped up the steps it had overlapped. A helper process computing the next window's Granite features on the CPU would have cost more than it saved, and on the efficiency cores at background priority the same work was far too slow.

Core ML runs the transformer on the Neural Engine only as fixed-shape fp16 chunks with overflow guards

Core ML, the only public route to the Neural Engine, compiles fixed-shape fp16 graphs for it, and when its compiler rejects a graph it runs it on the CPU without an error. The full 28-layer model and our first codec conversion (rank-5 tensors in its attention) both fell back; Core ML's compute plan (MLComputePlan, which reports each operation's device) and power readings (the Neural Engine at 0 W) catch it. The text-free branches run as chunks of consecutive layers. The small pieces around the chunks (input projection, timestep embedding, per-layer adaptive-norm vectors, output projection) run in numpy in a separate process, because Core ML's predict holds Python's GIL.

fp16 tops out at 65504, and this model exceeds it in attention-sink scores and in the squares inside RMSNorm. The conversion normalizes each RMSNorm row by its maximum absolute value before squaring; for attention it scales the queries by an extra 1/1024, subtracts each row's maximum, and scales back before the softmax. Both rewrites are exact in fp32. With them the Neural Engine's branches are as accurate as the GPU's: on inputs captured from real guided steps, relative L2 error against fp32 is 4.7–6.4e-3 on the Neural Engine and 5.2–7.5e-3 for GPU bf16.

The GPU's error comes mostly from its bf16 GEMM inputs: the same forward pass with fp16 GEMMs and an fp32 residual stream is 1.9–2.6e-3 from fp32. The Neural Engine gains fp16's three extra mantissa bits in its GEMMs and gives most of them back by keeping the residual stream in fp16. Quantizing its weights to int8 would have doubled the error for a 6% speedup.

Converting Granite exposed two fp16 precision losses: activations fused after a GEMM, and the overflow guard

A SiLU that directly follows a GEMM on the Neural Engine loses precision that the same SiLU alone does not: in a probe, relative error was 2.8e-4 standalone and 2.9e-3 after a linear layer. On Granite's own activations, writing the sigmoid as 1/(1+exp(-x)) brought a feed-forward block from 3.9e-3 to 2.4e-3, where tanh after a linear layer had reached 1.8e-2. The overflow guard costs precision where it is not needed: without it, Granite's attention error fell from 3.9e-2 to 1.6e-3. The likely cause is that the extra 1/1024 puts small query values near fp16's smallest normal number, where fewer mantissa bits remain. The transformer's text-free branches need the guard and did not change when their SiLU and softmax were rewritten.

A branch without text attends to one null token, so its cross-attention is a constant

A branch without transcript text still has a text cross-attention input: a learned null embedding, with every other position masked. A softmax over a single unmasked key is 1 for every query, so the cross-attention output is the null token's value projection passed through the output projection, a constant vector per layer. The text-free branches skip the query projection, the attention and the per-position output projection, and text keys and values are computed only for branches that have text. In a single forward pass the conditioned branch stayed bit-identical and the unconditional branch moved by at most one bf16 unit in the last place. The shortcut also makes so and u the same graph: their prompts have the same shape and differ only in token ids, whose embeddings are computed outside the Core ML chunks.

Partitioning the MLP along its intermediate dimension let the Neural Engine run two branches

Before partitioning, one text-free branch took 244 ms on the Neural Engine, and batching two branches into one call cost 2.17x the time of one, so the GPU kept two branches per guided step. A microbenchmark of one MLP at our shape showed where the time went: partitioned into independent blocks of intermediate columns, with the blocks' down-projections summed, the same MLP ran more than twice as fast at the same error against fp32.

Throughput rises until a block's weights plus its activations fit in about 32 MB, then flattens. That is the size of the Neural Engine's on-chip SRAM, where reverse-engineering measurements found throughput falling once a matrix multiply's operands no longer fit, although at our shapes the penalty was larger than the 30% they report. The ReLU MLP at 2 blocks, 27.5 MB of weights but about 33 MB with activations, sits between the extremes. We did not observe the compiler's schedule, so this reading fits the data without proving the cause. Apple's guidance for transformers on the Neural Engine recommends chunking large intermediate tensors, which the partition also does, for cache residency and to spread work across the engine's cores.

Partitioning alone was worth 1.7x; chunks longer than 7 layers and more than 8 blocks did not help. At 117 ms per branch, the Neural Engine runs so and then u while the GPU and CPU run sot, and guided steps fell from about 355 ms to about 235 ms. Clips gained less than that step saving predicts, flash least: the codec work that had filled the Neural Engine's idle time on guided steps lost it, and flash has the fewest unguided steps to absorb it.

Duplicate weight copies, not the model, made unified memory swap

Unified memory removes copies between engines, not copies within a program. With the Neural Engine running two branches, mid-run free memory fell to 13% and the system swapped 5–6 GB. The main process held 14 GB, mostly GPU memory (MLX's weights and buffers plus PyTorch's allocator), including original MLP weights that MLX kept after building the split GPU and CPU copies. Each set of Core ML chunks also holds its own weights, and u had its own set. Running u on so's chunks and dropping the original MLP weights kept swap under 0.7 GB and removed most of the run-to-run timing variance. Capping MLX's cache of freed buffers, and mx.set_wired_limit, which keeps MLX's buffers resident so the memory compressor never touches the weights, bought a few percent more.

Memory also blocked the remaining engine move. Unguided steps would take 156 ms on the Neural Engine instead of 215 ms on the GPU, but they condition on text and need a third set of chunks, which brought back 2–4 GB of swap. Core ML multifunction packages, which let two fixed-shape functions share one deduplicated weight file, reduced that to transient swap, and the configuration then only tied the default on this 24 GB machine.

The codec runs in Neural Engine idle time that the main process declares

The codec encodes each window's input into the observation condition and decodes each finished window into audio. Both run in the Neural Engine process as Core ML chunks of 30–60 ms, and a chunk starts only while the main process has declared the Neural Engine idle: for 0.18 s on each unguided step, when the GPU works alone, and for 0.5 s after each window. Each window is decoded while the next one samples, and the next window's observation is encoded ahead of time. A file's first window has no earlier window to overlap, so its encode starts at once and runs while the GPU computes the Granite features.

Declared idle time still costs something. Running Granite on the Neural Engine as well removed its work from the GPU but made clips slower overall. We did not isolate the cause; the most likely one is the CPU work every Neural Engine call brings (dispatch and conversion of large fp32 arrays), which slows the split MLP like any other CPU work. Splitting one branch's layers between the GPU and the Neural Engine failed the same way: guided steps got faster, but the per-step hand-off and the lost idle time for the codec left clip time unchanged.

Measured effect of each change on clip time

Each entry comes from an interleaved A/B at the time, repeated until the direction held (run-to-run spread was 0.1–0.3 s per clip); the entries are not strictly additive.

Quality matches the CUDA run except a ScoreQ offset that predates the Neural Engine

The benchmark is the one the CUDA production run uses. ScoreQ-NR is a no-reference speech-quality predictor on a 1–5 scale; WER is the word error rate of an ASR transcript of the restored speech against the reference transcript; speaker similarity is the cosine between speaker embeddings of the restored and clean speech; UTMOS is a predicted mean opinion score. The Mac draws its sampler noise from CPU generators, so it does not reproduce CUDA's noise, and per-clip differences include seed variation of about ±0.01 for ScoreQ and speaker similarity and ±0.02 for WER. Intervals are bootstrap intervals over per-clip differences.

Only the ScoreQ interval excludes zero, at one to two times the seed-to-seed noise. The first full run, before any branch ran on the Neural Engine, had the same offset (−0.025 pro, −0.018 flash), and it did not grow afterwards.

Every change was smoke-tested on two clips before its timing counted and re-scored on all 78 clips when it moved outputs beyond rounding. The final rewrite produces bit-identical audio to the benchmarked pipeline. From a cold start, a 3-minute file restores in 32 s, including compilation of the first window's graphs and excluding an 8 s model load.

Mac mini setup changes

The software was macOS 27, MLX 0.32, coremltools 9.0 and PyTorch 2.10. Core ML models were converted as fp16 ML programs for macOS 15 with CPU-and-Neural-Engine compute units and take fp32 NumPy inputs. Arrays cross the pipe to the Neural Engine process at under 1 MB per guided step, plus about 5 MB of conditions per window. These balances were measured on the M4 Pro; the M5 generation adds matrix units to the GPU, which would shift them.

Sources

• Apple Machine Learning Research, Deploying Transformers on the Apple Neural Engine (2022): the (B, C, 1, S) layout with 1x1 convolutions and chunking of intermediate tensors for cache residency and multicore use, in the partitioning section.
• Inside the M4 Apple Neural Engine, Part 2: ANE Benchmarks: the ~32 MB SRAM cliff and the ~30% penalty for matrix multiplies whose operands exceed it, in the partitioning section.
• Orion: Characterizing and Programming Apple's Neural Engine for LLM Training and Inference (arXiv:2603.06728): the 32 MB SRAM, the ~19 TFLOPS fp16 peak through private APIs and fp16's range, in the engine table and the Core ML section.
• Hello SME! Generating Fast Matrix Multiplication Kernels Using the Scalable Matrix Extension (arXiv:2409.18779): the M4 CPU's SME units, in the CPU section.
• MLX documentation: mlx.core.set_wired_limit and compilation: wired memory and mx.compile, in the memory and MLX sections.
• Core ML Tools: multifunction models: weight sharing between functions, in the memory section.
• Apple Machine Learning Research, Exploring LLMs with MLX and the Neural Accelerators in the M5 GPU: the M5 GPU's matrix units, in the setup section.

Related Diffio Posts

• Classifier-Free Guidance: The Guidance Scale Is an Exponent on an Implicit Classifier: why a guided step needs forward passes with and without conditions.
• Flow Matching: Regress the Field the Sampler Integrates: the velocity field the samplers integrate.