Quantizing a mixture-of-experts model on MLX

Mixed-precision quantization on a dense model is straightforward. Every linear layer is one matrix, you measure how far the output moves when you quantize it, and the layers that move it most get more bits. Mixture-of-experts models do not have that shape, and the differences do not announce themselves. A quant can load, produce fluent text, and still be worse than it should have been. We also shipped a command-line flag that ran without errors for several releases and had no effect.

Notes from building the MoE path in OptiQ.

The experts are almost the whole model

Take gemma-4-26B-A4B. It has thirty decoder layers with 128 experts in each, and mlx-lm stores them fused, one tensor per projection holding all 128:

language_model.model.layers.0.experts.switch_glu.down_proj.weight
  shape [128, 2816, 176]   uint32   0.25 GB

There are ninety of those tensors, three per layer. They add up to 15.1 GB in a 17.6 GB file. Attention, the norms, the embeddings and the router account for the other 2.5 GB.

At that ratio the bit budget is almost entirely a set of choices about the experts. Another bit on attention barely registers.

Sensitivity sees one layer where there are 128 experts

OptiQ measures sensitivity by quantizing a single layer, running the calibration set through the model, and recording the KL divergence between the resulting output distribution and the reference. Repeat for every layer at every candidate bit-width and you have a table the knapsack can spend a budget against.

The fused tensor complicates the unit of measurement. Quantizing it means quantizing all 128 experts together, so the score describes the group rather than any individual expert. Each fused tensor gets one score and one bit-width.

Treating each expert as its own layer would multiply the sweep by 128, and a sweep on a 26B model already takes hours. The resulting bit map would also no longer match how the weights are stored on disk. We have not measured a case where the extra resolution pays for that, so the tensor is scored the way it is stored.

A flag that did nothing

OptiQ can stream expert weights off SSD rather than holding them resident. It identifies a routed expert by matching the tensor key against a list of path segments. The list read:

_EXPERT_SEGMENTS = (".switch_mlp.", ".mlp.experts.", ".ffn.experts.")

A comment above it claimed .mlp.experts. covered Gemma, which is wrong. Gemma-4 stores experts at layers.N.experts.switch_glu. and there is no mlp segment in the path at all. Qwen3 uses switch_mlp. The two names have no substring in common.

So --stream-experts matched nothing on any Gemma-4 MoE we publish. With no expert tensors found there was nothing to stream, the model loaded resident as before, and no error was raised. That went unnoticed across several releases.

The tests we had covered the decision, meaning: given this model and this mode, do we choose to stream. They passed, because that decision was correct. What none of them checked was whether the streaming path found any weights once it ran. The fix was one line. The test now reads real tensor names out of each family's published safetensors index, so an architecture that names things differently fails there rather than declining to stream in silence.

When the experts do not fit

A 4-bit Qwen3.6-35B-A3B is around 21 GB. On a 24 GB Mac there is then no room for a KV cache, so the model cannot really be used despite being nominally small enough.

Only eight experts out of 128 run per token. The routed weights are both the bulk of the file and the part least often touched, which makes them the obvious thing to leave on disk. The streaming path keeps scales and biases resident and reads only the expert rows the router selected. Those reads go through positional os.pread on a thread pool, so several experts arrive in parallel, and the compacted operand then goes through mx.gather_qmm as it would resident.

The scales and biases are the awkward part. On DeepSeek-V4, with 256 experts, they come to roughly 17 GB by themselves, which cancels out the benefit of leaving the weights on disk. Past a size budget they are streamed as well, from the same file handles.

Rounding shapes up to keep the kernel cache warm

The first version that worked was slower than expected, and the disk was not responsible.

Each step reads whichever experts the router selected, so the compacted operand has a different first dimension almost every time: seven experts, then twelve, then nine. MLX keys kernel compilation on shape. A single prefill was generating around 120 distinct shapes and recompiling for most of them.

Rounding the active-expert count up to one of a few fixed buckets fixes it. Pad the operand to 4, 8, 16, 32 and so on by repeating the last expert. The routing indices never reference the padded rows, so the output does not change and only the shape is stabilised. Around a hundred recompiles become a handful.

The prefetch that made decoding slower

Routing is reasonably stable between consecutive tokens, so warming the previous token's experts in the background while the GPU is busy seemed worth trying.

On Qwen3.6-35B-A3B on a 24 GB Mac, decoding went from 8.81 tok/s to 8.14 with prefetch enabled. The model largely fits the operating system's page cache, so the speculative reads were fetching pages that were already warm while competing with the reads that were needed.

It is disabled by default and remains behind OPTIQ_STREAM_PREFETCH=1. The reasoning still applies to a model genuinely larger than RAM, where the experts really would be cold, but we have not measured a case where it helps.

If you are doing this yourself

A few things worth checking on an architecture you have not handled before. Assert on how many expert tensors your matcher found rather than on whether it decided to stream, since the families name these tensors differently and a miss raises nothing. Work out what share of the file the routed experts occupy before tuning anything else. And measure any prefetch rather than reasoning about it; ours appeared sound and cost 8% of decode throughput.

Per-family behaviour is listed on the models page, and the streaming flags are documented under serve.