Skipping dead neurons: a faster Nemotron 3.5 Lightning on Apple silicon
How we achieved 14% faster generation on an M5 Pro chip using the same weights at the same quality. The kernel is on GitHub.
LLMs write tokens one at a time. Each token is a complete trip through the model, where it does a small amount of arithmetic while needing to read gigabytes of model weights. Loading those weights into the chip is what takes time.
So a token can only arrive sooner if fewer bytes are read, and Nemotron 3.5 Lightning hands us a way to do exactly that. Inside its experts sits the ReLU² activation function, and most of the numbers passing through it become zero. Whatever a zero is about to be multiplied with never has to be fetched. Our kernel for Apple silicon skips those bytes.
Everything below runs on the mlx-community 4-bit conversion of Nemotron 3.5 Lightning.
stock mlx-lm
Heimr Sparse Down
Transposing to skip the zeros
Inside the model tokens are represented as a list of numbers. Most of what a layer does is multiply that list by a matrix of stored numbers, the weights. Here is such a list at the moment that matters for us: when ReLU² is applied.
- Twelve numbers from the token representation list. ReLU² makes negative numbers 0 while squaring positive numbers.
After applying ReLU² we multiply by the down projection matrix. Each row of the matrix gets multiplied by the corresponding element of our list, so we want to skip fetching all of the rows that are about to get multiplied by 0. This requires the matrix to be loaded transposed from how it is normally loaded, because of the direction the GPU reads the data.
- The data and weights are the same. The normal load on the left lets the GPU be granular on the output axis, while the transposed load gets granularity in input, which is the axis that contains the skippable 0s.
Storing the down projection weights transposed in unified memory gives us the possibility to skip moving weights into registers just to be multiplied by 0. This transposition does come with a caveat during prefill, which we will explore later.
Let's follow a token through the decode step
Everything from here follows one real token: layer 10 of a plain pass over 2048 tokens of a novel, with the routing, the zeros and the timings all its own.
52 layers, three kinds
Nemotron 3.5 Lightning is a hybrid stack:
23 Mamba-2 layers, 6 attention layers, 23 mixture-of-experts. The MoE layers hold most of the weights any one token touches.
Inside a MoE layer
A router picks 6 of 128 experts for this token. Above them sits the shared expert, which runs for every token and holds 3712 neurons, twice a routed expert's 1856, hence twice the height. Only it and the chosen six do any work; the other 122 are never read.
Two projections and a switch
An expert widens the 2688-dimensional token out to its neurons, applies ReLU², and narrows the result back down. Each projection is 5.6 MB of 4-bit weights, and a decode step reads each exactly once.
The standard way reads everything
The stock kernel cannot see the zeros: to produce each of the 2688 outputs it sweeps all 3712 values of h and the weights behind them, dead neurons included.
Our way reads only the rows that fired
With the weights stored by neuron, a live row is read once, straight through, for all 2688 outputs, and a dead row is never touched. The work is cut into 4 slices of 928 neurons and 3 bands of outputs, because one expert alone offers too few live rows to keep 20 GPU cores busy: each slice finds its own live rows and writes a partial sum, and a final small pass adds the partials together. Same up projection, same ReLU², same answer.
The whole decode step
Pulled all the way back: every layer of the stack as one segment, sized by how long it takes, in the order the token walks them. The MoE layers get faster, while the other layers stay the same.
The 4 bit checkpoint has a scale for every group of 64 neurons. The scales sit along the output axis, so we have to read a scale even if only one of the 64 neurons in a group are alive. Therefore the shared expert ends up reading 1.3 MB of the 5.6 MB of weights instead of the 0.6 that live rows alone would cost. Re-quantising with the scales along the neuron axis would remove this toll, but it cost a small, consistent amount of perplexity, so the kernel reads the mlx-community bytes as they are.
How the kernel finds the live rows, briefly. One cheap instruction asks each group of 32 neurons who fired, and a running count compacts the survivors into a short shared list; from that point on, the workers simply walk the list. A dead neuron's weights are never requested, and a group of 64 dead ones costs no scales either.
Most neurons are dead in every layer
The graph below shows share of live neurons for shared and routed experts by layer.
The down projection gets 2.9× faster, and the block a quarter shorter
Each MoE layer's decode time, standard against sparse. No layer comes out a loss, and summed over the 23 layers this is the cut the takeaway promised.
The weights now face the wrong way.
Storing the down projection by neuron is what makes the decode skip possible, and it raises a problem the moment a prompt arrives. Prefill pushes hundreds or thousands of tokens through a layer at once, and there the down projection is a matrix multiply, one that MLX runs on the M5's matrix units and one that expects the checkpoint's layout. Our copy is that matrix transposed, and MLX has no dense consumer for the transposed shape. There are three solutions:
- Keep both copies: This requires keeping 26.2 GB resident instead of 17.8.
- Transpose the weights back for every prompt: This takes about 1.9 seconds for all 23 layers, which is significant for short prompts.
- Write a prefill kernel that reads the transposed layout and uses some of the zeros too.
The zeros are less generous in prefill than in decode, and the storyboard below shows why. One token lights a few percent of an expert's neurons, but the roughly 96 tokens an expert sees in a 2048-token prompt light different ones, and between them they cover most of the expert. Per token, there are rows to skip; across the whole prompt, there are not. The way through is to cut the expert's tokens into small tiles and take the union of live neurons per tile: for 32 tokens at a time the union is 39 to 81% of the rows, and a matrix multiply over only those rows is exactly that much smaller.
A whole prompt at once
Prefill is the other half of inference: the prompt's 2048 tokens travel through the 52 layers together, as one tall matrix. The layer we watched during decode, layer 10, now sees all of them.
Sorted by expert
The router still picks 6 of 128 experts, but now once per token across all 2048. Grouped by expert, a routed one ends up with about 96 tokens to process, and the shared expert with all of them.
Per token sparse, per prompt dense
One expert, its tokens as columns, its neurons as rows. Each column lights only a few percent of the rows, but the columns disagree about which rows, and the union over all of them lights 65%. For the shared expert, which sees every token, it is 97%. Rows dead for the entire prompt are few, so skipping them alone buys little.
Two layouts, one problem
The checkpoint stores an expert's down weights one row per output, which is exactly what MLX's prefill matrix multiply reads. Decode needs the same numbers one row per neuron, the matrix transposed, so a dead neuron's row can be skipped in one piece. MLX offers no matrix multiply for the transposed shape: prefill cannot use the copy decode needs.
Two obvious answers, both costly
Keep a second copy in the checkpoint's layout for prefill: 8.4 GB more memory. Or transpose the weights back for every prompt: 1.9 seconds across the 23 layers, longer than the prefill itself for prompts under 3500 tokens.
What we do instead
A prefill kernel that reads the transposed copy directly. Cut the expert's tokens into tiles of 32; within one tile the union of live neurons is 38% of the rows, and only those rows are gathered and multiplied, on the matrix units, as one small dense product per tile. A row dead for all 32 tokens is never read, and there is just one copy of the weights.
Prefill, end to end
The whole model on real prompts: the neuron-major copy prefills as fast as the checkpoint layout does, and the model needs no more memory than stock mlx-lm. One copy of the weights now serves both phases.
Layer by layer, prefill
The down projection of a 2048-token chunk, per MoE layer: MLX's dense kernel on the checkpoint layout against ours on the neuron-major one. At the sparse edges of the model the tile union stays small and the kernel wins by a quarter or more; in the dense middle layers the union grows and some of that lead goes back. Over the whole model the two finish level, at 1918 against 1965 tokens per second for a 2048-token prompt, with 17.8 GB resident in both cases.
Speculative decoding.
NVIDIA's release of the model carries a multi-token prediction head: one more attention block and one more MoE block that, given the model's final hidden state and the token just chosen, guess the token after it. The mlx-community conversion dropped those tensors, and mlx-lm has no speculative decoding path for this model regardless, because its Mamba-2 layers are recurrent: verifying several drafted tokens in one step means running the recurrence forward over rows that may be rejected, and then putting its state back where it was. None of that is new. Speculative decoding with a verify step and a rollback is the standard recipe, and for a recurrent layer the rollback keeps the step's inputs and recomputes the state over the accepted prefix. We wrote that plumbing for MLX and published the head as a 772 MB 4-bit file next to the checkpoint.
Every draft is checked by the model itself, so the output is exactly what greedy decoding would have produced; only the speed changes, and the speed depends on the text. With two drafts per step, 1.3 to 1.7 of them survive on the two chat prompts below, which is why the sparse + MTP speeds in the table differ. The skip and the head compose cleanly: the verify step reads three rows' worth of activations through the same weights, and the neurons those rows leave dead are dead for all three.
| prompt | dense | dense + MTP 2 | sparse | sparse + MTP 2 | drafts accepted of 2 |
|---|---|---|---|---|---|
| short prose, chat · 27 prompt tokens | 106 | 118 | 121 | 139 | 1.26 |
| code, chat · 50 prompt tokens | 104 | 136 | 119 | 167 | 1.68 |
| a novel, raw continuation · 2048 prompt tokens | 99 | 92 | 114 | 105 | 0.84 |
| source code, raw continuation · 2048 prompt tokens | 100 | 102 | 115 | 118 | 1.15 |
| a novel, raw continuation · 8192 prompt tokens | 98 | 62 | 110 | 70 | 0.46 |
| source code, raw continuation · 8192 prompt tokens | 97 | 90 | 111 | 109 | 1.37 |
Measuring it
- Hardware
- Apple M5 Pro, 20-core GPU, 48 GB. The whole model fits, so the end-to-end numbers are measured, not summed.
- Software
- MLX 0.32.2, mlx-lm 0.31.3, custom Metal kernels through MLX's kernel interface
- Model
- NVIDIA Nemotron 3.5 Lightning 30B A3B