Qwen3.6-35B MoE on M1 Max: pre-warming hot experts didn't beat the mmap cache

I got per-token expert selection logs out of Qwen’s MoE router, so I tried using them for a pre-inference warm-up: read the hot experts into memory before generating. Hot experts here are the experts the router picked especially often while choosing 8 out of 256 per token.

The Kimi K3 report stopped at the top-16-of-896 routing spec, and the SwiftLM hands-on never even reached SSD streaming: the OS mmap page cache covered everything.

If a log can tell me the hot experts ahead of time, is loading them into memory beforehand faster than leaving everything to the mmap cache? I benchmarked both under identical conditions.

Test environment

Item	Details
Machine	MacBook Pro M1 Max, 64GB unified memory
OS	macOS 26.5 (Darwin, build 25F71)
Python	3.13 (miniconda base)
mlx	0.31.2
mlx-lm	0.31.3
Target model	unsloth/Qwen3.6-35B-A3B-UD-MLX-4bit (dynamic quant based on 4-bit, 21.7GB / ~20.2GiB)
Larger model I also tried	mlx-community/Qwen3.5-122B-A10B-4bit (4-bit quant, 69.6GB / ~64.8GiB)
MoE architecture	qwen3_5_moe (shares Qwen3NextSparseMoeBlock with the Qwen3-Next family)
Generation	mlx_lm.stream_generate, default sampling parameters

35B-A3B routes top-8 out of 256 experts, and all 40 layers are MoE. 122B-A10B is the same qwen3_5_moe family with more layers and parameters.

Hooking the router

The router part of Qwen3NextSparseMoeBlock boils down to three lines.

gates = self.gate(x)
gates = mx.softmax(gates, axis=-1, precise=True)
inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]  # 選ばれたexpert index

Instead of editing site-packages, I monkey-patched Qwen3NextSparseMoeBlock.__call__ after loading the model: copy the original implementation verbatim, and add only the code that records inds per layer and token.

At first I called mx.eval() every time inds came out, but that creates one sync point per layer per token and blocks the memory reuse MLX gets from lazy evaluation. I ended up batching: when a token reaches the last MoE layer (layer 39 on 35B-A3B), evaluate that token’s inds in a single call, cutting the sync points by a factor of the layer count.

The router hook, the A/B experiment, and the safety wrapper are all on GitHub → LiltingChannelLabo

I prepared five prompts in deliberately different genres.

Key	Content
bst	Insert function for a binary search tree (Python)
bbs	Minimal BBS implementation (HTML)
kana_intro	Character self-introduction
math	Fast Fibonacci via matrix exponentiation
cn	Self-introduction in Chinese

122B-A10B died on Metal running out of memory

I tried the bigger model first. Putting a 69.6GB (~64.8GiB) 4-bit model on a 64GB unified-memory machine was obviously a stretch, so I loaded it with mlx_lm.load(model_id, lazy=True).

Loading finished in a second or two. mlx_lm.load() defaults to lazy=False, which eagerly evaluates every parameter inside the load call. But even with lazy=True, free memory dropped by tens of GB during the very first prefill (the phase that processes the whole prompt at once).

Judging a safe zone from outside the process using vm_stat free memory never stabilized, no matter how many times I changed the threshold.

Watch swap, not free memory

Across several runs, free memory swung between nearly 50GB and a few dozen MB, while swap usage never moved at all.

A read-only mmapped safetensors file consumes free memory as the OS pulls pages into the page cache, but unmodified pages need no write-back: the OS can drop them (without going through swap) and re-read them from the file when needed. A drop in free memory alone can’t tell you whether memory is actually under pressure. MLX’s own counter (mx.get_active_memory()) only counts MLX-managed buffers; right after model load it stayed at active=0.00GB, unrelated to the free-memory plunge.

Once I changed the danger criterion to swap growth only, the process stopped getting killed from outside and ran to completion.

Metal itself still reported out-of-memory

With the swap-based guard in place, generation now failed with a clear error from mlx_lm.

libc++abi: terminating due to uncaught exception of type std::runtime_error:
[METAL] Command buffer execution failed: Insufficient Memory
(00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory)

This wasn’t the guard script: Metal itself failed to allocate the memory needed to execute a GPU command buffer, and the process died on an uncaught C++ exception. The system never froze, and swap usage never moved.

Setting an explicit cap with mx.set_memory_limit() changed nothing. Per the documentation, the limit is a “guideline”: an allocation only raises an exception once the limit is exceeded and RAM including swap is exhausted. Lowering the value made no practical difference.

lazy=True only delays when weights get evaluated; it doesn’t promise a lower peak during the forward pass. In practice, the 4-bit 122B-A10B did not fit this 64GB unified-memory machine through mlx_lm’s plain generate implementation. SwiftLM had automatically classified this model under its “SSD STREAMING” strategy, reading it from SSD incrementally.

I switched to 35B-A3B, which reliably fits in 64GB.

Observing experts on 35B-A3B

35B-A3B ran without trouble. MLX’s counter read active=0.00GB right after load, then settled around 18.5GB as generation progressed, which matches the model size.

I ran the 5 prompts at max_tokens=150 and counted the unique experts selected per layer (out of 256).

Prompt	Generation time	Unique experts per layer (of 256)
bst	4.0s	158.3
bbs	3.3s	156.4
kana_intro	3.4s	149.3
math	3.2s	161.5
cn	3.2s	147.9

Every prompt used around 60% of the experts per layer (147.9 to 161.5). Each token only picks top-8, but the union over 150 tokens spread the selection across more than half of the experts. The premise that pre-reading just the top hot experts should be enough mostly didn’t hold at this generation length.

Pre-warm vs plain mmap cache

Reading SwitchLinear/QuantizedSwitchLinear (mlx_lm/models/switch_layers.py), each MoE layer stores its expert weights as a single (num_experts, output_dims, input_dims) array. Quantized models go through QuantizedSwitchLinear, where gather_qmm does the matrix multiply while selecting experts via rhs_indices.

x = mx.gather_qmm(x, self["weight"], self["scales"], ...,
                   rhs_indices=indices, transpose=True, ...)

Using this, I implemented the warm-up as one dummy-input gather_qmm call over the observed hot experts, executed before generation.

I compared two conditions in the same process with the same prompt. Condition A is plain stream_generate with no warm-up; condition B first pre-reads the per-layer top-20 experts from the observation log with this warm-up, then generates the same prompt.

Prompt	Condition A (plain)	Condition B (top-20 pre-warmed)	Diff
bst	61.83 tok/s	61.80 tok/s	-0.05%
kana_intro	62.02 tok/s	61.92 tok/s	-0.16%

The warm-up itself took 0.02 seconds, but in generation tok/s the A/B difference stayed within noise. Both prompts showed the same pattern, with B marginally lower; no meaningful improvement.

Since A and B run in the same process in that order, B does get a bias in its favor worth the first prefill’s cold start. But the generation speed being compared is a decode-side metric (tokens emitted one at a time), so prefill differences barely touch this number.

The 4-bit 35B-A3B is 21.7GB, and the plain mmap cache already delivered 61.8 to 62.0 tok/s. Even knowing in advance which experts would be selected, only noise-level cold-fault cost was left for the warm-up to shave off. Whether this changes at a scale where the model doesn’t fit in memory and SSD I/O becomes the real bottleneck, i.e. the 122B-A10B that wouldn’t run here, is unverified: plain mlx_lm generate couldn’t even start generating.