Notes from moving a sparse spiking simulation to MLX: memory traffic, not dispatch count, was the bottleneck

[Kisame76]

I moved a leaky integrate-and-fire simulation of the fruit fly connectome to MLX with a few custom Metal kernels and wrote down what generalises beyond it. The workload is the published Shiu et al. model: 127,400 neurons, 14.7 M directed edges, 10,000 sequential steps per simulated second, and a data-dependent branch in every step. Nothing about the lessons depends on it being neuroscience.

Full notes with every table: docs/mlx-notes.md. Repository: Kisame76/drosophila-brain-mlx.

All numbers below come from those notes, measured on an Apple M4 Pro (24 GB) with MLX 0.32.2 on macOS 26.6.2. Across my own implementations the same model went from 22.0 s to 0.29 s per simulated second, and removing host synchronisation was a small part of that.

What I took away

1. Chunking with mx.async_eval helped less than I expected: 22.1 s to 19.4 s. Varying the network's activity by 5,000x moved the runtime by 2 %, which pointed at a fixed cost per step rather than at the data path.
2. The intermediates of elementwise chains dominated. A chain of 16 elementwise ops over 127,400 float32 elements took 0.1333 ms per step and moved 15.6 MiB; the same arithmetic in one mx.fast.metal_kernel took 0.0225 ms and moved 1.0 MiB. In the full engine, fusing the tick's chain took it from 0.75 s to 0.29 s per simulated second. I first attributed that to dispatch count, which was wrong:
3. mx.eval is the synchronisation unit. 1, 2, 4, 8 or 16 trivial kernels inside one eval all cost about 0.11 ms in total.
4. .at[].add() is the accumulating scatter, but it cannot skip work. arr[idx] = vals with repeated indices keeps the last write; arr.at[idx].add(vals) accumulates, through scatter_add with GPU atomics. My baseline ran it densely over all 14.7 M edges every step, so its cost did not depend on how many neurons fired. Skipping the inactive edges needs a compaction whose size depends on the data, and that means a host read. A Metal kernel dispatched densely with an early exit and integer atomics gets the sparsity without one, and took the step from 19.6 s to 0.75 s. (Corrected: this point first said there was no accumulating scatter.)
5. mx.compile changed results here, through its constants. It prints captured scalar constants into the kernel source with 7 significant digits (mx.compile inlines float scalar constants with 7 significant digits, so compiled results differ from eager by 1 ulp #4503), and two of mine came back 1 ulp off. With a strict threshold downstream the runs diverged (48,308 against 47,499 events at 10,000 steps); passing the constants as arguments makes the compiled run bitwise identical. (Corrected: this point first blamed FMA contraction, but mx.compile does not contract a*b + c. A handwritten kernel that writes it as one expression does, and #pragma clang fp contract(off) in the header avoids that.)
6. Integer atomics are reproducible, float atomics are not. Accumulating synaptic contact counts as int32 and scaling once afterwards made the kernel bit-reproducible regardless of thread order.
7. Skewed rows need a tunable unroll. Out-degree runs from 0 to 9,615. Letting K threads stride through one row helped, but the best K depends on the load (2 when the work is tiny, 8 to 16 when it is not), and choosing it at runtime would need exactly the host read the design avoids.

Two questions, in case I missed something

• Is there a supported way to turn off floating-point contraction in mx.compile? (Not needed after all; see point 5.)
• Is an accumulating scatter (scatter_add or bincount style) planned, or is .at[].add() the intended path for large index sets too? (It is the intended path; see point 4.)

Corrections are welcome. I would rather fix the notes than leave a wrong explanation in them, which has already happened twice (see points 2 and 5).

[Diacod-I]

Holy moly, I read the first line and got hooked. Really great write up! Also regarding your questions, here are my answers from reading the source on current main:

• FP contraction in mx.compile: There's no supported way right now. mx.compile has no compiler options, and its kernels use the default safe math mode, which still allows FMA contraction, judging by your results.
• Accumulating scatter: .at[idx].add() is the accumulating scatter. It lowers to scatter_add, which runs in parallel on the GPU with atomic adds, so it's the intended path for large index sets too. Float results aren't order-deterministic, consistent with your point 6. I suspect your custom kernel won at 14.7M edges less because of the scatter and more because using .at[].add() would need the active-edge set gathered first, which has a data-dependent size and forces a host read. Your dense dispatch with early exit avoids that. It might be worth rewording point 4 to say that in the notes.

[Kisame76]

Thanks for reading the source. Both answers hold on v0.32.2, and testing the premise of the first one showed that my point 5 was wrong.

Contraction. Right, there is nothing to pass: Compiled::eval_gpu builds its library with the default CompileOptions, which means MathMode::Safe. But mx.compile doesn't contract in the first place. Over 2^20 random float32 triples, mx.compile(lambda a, b, c: a * b + c) is bitwise identical to eager on all 245,246 elements where FMA and multiply then add round differently, because it emits each op as its own statement. A metal_kernel that writes a * b + c as one expression is contracted on all of them unless the header has #pragma clang fp contract(off).

My divergence was #4503 instead: captured scalar constants are printed into the kernel source with 7 significant digits, and two of my eight come back 1 ulp off. decay_v is 0.995012462 as a float32, becomes the literal 0.9950125 and parses back as 0.995012522; v0_term loses 1 ulp the same way. Passing the constants to the compiled function as arguments makes the compiled run bitwise identical to the reference over 10,000 steps, and the compiled lane in the repository now does that. #4511 should fix it upstream.

Scatter. Agreed, .at[].add() is the accumulating scatter, and "I found no accumulating scatter" was wrong. One detail on the case with 14.7 M edges: my baseline didn't gather the active edges. It ran .at[].add() densely over all 14.7 M edges every step (gather the delayed spike mask per edge, zero the inactive ones, scatter everything), with no host read, which is why its runtime was flat across a 5,000x range of activity. So your conclusion holds with a slightly different mechanism: the kernel wins by skipping the edges of neurons that didn't fire, and doing that with .at[].add() would need the compaction, which is where the host read comes in.

I also timed the scatter itself at equal work, over all 14.7 M edges: 1.161 ms per step for .at[].add() against 1.146 ms for a handwritten kernel with one integer atomic per edge, with identical sums. The dense propagation as a whole took 2.140 ms, so the scatter was only about half of it, and the speedup did not come from a faster scatter.

I've corrected both sections of the notes, and points 4 and 5 above.

[Diacod-I]

Thanks for the corrections and for the equal-work scatter timing. It's a cleaner explanation of where the speedup came from. And good catch on mx.compile not contracting in the first place. My bad on that one. Love this work, you get an upvote from my end ❤️.