Run a 20B Ternary MoE in 0.63 GB of RAM

Maple is a 20B-parameter A1B ternary MoE: 24 layers, 256 experts, top-8 routing. Serving it normally costs ~5.9–6.5 GB of resident memory. With file-backed expert offload and KV-cache quantization, it runs in 0.63 GB — bit-exact against the full model.

Maple comes from the DeepGrove fork of ml-explore/mlx-lm; the architecture and the 2-bit ternary weights are theirs (deepgrove/maple-2bit-mlx on Hugging Face). This post is about serving that exact model on less RAM — the offload adds no approximation.

Why file-backed expert offload is very cool, and how to cut a model’s RAM ~10× without losing a single bit

Authors

Zenith Rifle

Affiliation

Independent (fork of ml-explore/mlx-lm)

Published

Aug. 29, 2026

Table of Contents

1. Classical MoE Serving
2. Are All Those Experts Really Necessary?
   • Why this matters on unified memory
3. The Goal
4. Building the Offloader
   • A file index without loading the checkpoint
   • Byte-range reads without the loops
   • Fused projections and BF16 reinterpretation
   • IFP: masking the gate to the active set
   • Per-token reselect
   • LRU paging
5. Quantizing the KV Cache
6. Wiring Into mlx-lm
7. Actual Testing!!!
   • Verifying losslessness
   • Memory numbers
   • Why sampling matters
8. Results
9. Limitations
10. Run It
11. Donations
12. Reference List

1. Classical MoE Serving

Maple is a Mixture-of-Experts transformer: every one of its 24 layers routes each token through 8 of 256 experts (num_experts_per_tok=8), with a 512-token sliding window on 3 of every 4 layers and full attention on the remaining 6. The weights are 2-bit packed {-α, 0, +α} ternary values, one α per row, so a “20B” model is really a ~1B-actively-used model inside a ~20B-resident package.

In the stock mlx-lm build, serving Maple means loading all 256 expert tensors for all 24 layers into DRAM and keeping them there:

Peak resident memory on the test machine — a MacBook Air, Apple M2, 16 GB unified memory (passively cooled, no discrete GPU): ~5.9–6.5 GB. That is the baseline we want to beat. Every measurement in this post is from that M2 Air.

2. Are All Those Experts Really Necessary?

The key observation: per token, the gate only ever uses 8 of the 256 experts. The other 248 experts do zero compute for that token — yet they occupy (256 − 8)/256 ≈ 97% of the expert memory, sitting there in case a future token routes to them. That is roughly a 24× memory overhead for one token’s worth of compute.

This is not the same tradeoff as weight quantization — the weights stay at full 2-bit ternary precision. We are only deciding where bytes live: RAM vs. disk. And because we keep it bit-exact, we get the memory win without touching quality.

This buys us two things:

• One thing to reason about. No separate component to manage — just the model and the file system. When the gate wants an expert, the offloader fetches it.
• Much smaller memory footprint. The full checkpoint stays on disk; DRAM holds only the activated subset plus the KV cache.

2.1 Why this matters on unified memory

Apple Silicon doesn’t have dedicated VRAM. CPU and GPU share one physical memory pool — on this MacBook Air, 16 GB total. There is no “video memory” bucket to hide the model in; every byte the model pins is a byte taken from the same pool that macOS, the browser, the editor, the terminal, and every other process on the machine are fighting over.

So a ~6 GB model on a 16 GB machine isn’t “fits with room to spare” — it is ~38% of the entire machine’s memory handed to one process. macOS plus a browser plus an editor plus a couple of agents easily eat the other 10 GB, and then macOS starts swapping to disk: memory pressure, beachballs, the whole system crawling. Even though the full model technically fits, in practice you cannot multitask, and you cannot run multiple model instances or agents side by side — the fanless M2 Air thrashes.

The flash-MoE result reframes that: at 0.63 GB the model reclaims ~5 GB for the rest of the system. The same 16 GB machine that choked on one full Maple can now hold the model and a browser and an editor and several concurrent agents. “It fits” was never the bar on unified memory — “how much of the shared pool is left for everything else” is.

3. The Goal

Concretely:

• Serve Maple on Apple Silicon (MLX, Metal) using the stock mlx-lm build — no custom kernels, no custom metal library, fully portable.
• Cut peak resident memory from ~6 GB toward ~1 GB.
• Be mathematically lossless: with active = top_k and shared = 0, the resident set is the model’s true top-8 routing, so output must be bit-exact vs. the full model.
• Keep every feature working: generate, chat, and the OpenAI-compatible server, with live throughput metrics.

Where do we start? We need a way to keep the checkpoint file-backed and read only the rows the gate asks for.

4. Building the Offloader

What we have is a 5 GB checkpoint in safetensors format. What we want is a resident buffer of active expert rows that can be swapped out and paged in as the gate’s routing changes — while never materializing the full 256-expert tensor in MLX.

The cardinal rule, learned the hard way:

MLX materializes an entire lazy tensor on any access — even a single-row slice. So the full (256, ...) expert tensor must never exist as an MLX array.

That constraint shapes everything below.

4.1 A file index without loading the checkpoint

A safetensors file is 8 bytes of header length || JSON header || raw bytes. The header maps each tensor name to {dtype, shape, data_offsets}. So we can build a full index of every expert tensor — file path, absolute byte offset, dtype, shape — by reading only the JSON header, never any weights:

def _build_file_index(model_path):
    index = {}
    for fp in sorted(glob.glob(os.path.join(model_path, "*.safetensors"))):
        with open(fp, "rb") as f:
            n = struct.unpack("<Q", f.read(8))[0]
            hdr = json.loads(f.read(n))
        data_start = 8 + n
        for name, meta in hdr.items():
            if name == "__metadata__":
                continue
            dt = meta["dtype"]
            if dt not in _SAFETENSORS_DTYPE:
                continue
            off0, _ = meta["data_offsets"]
            index[name] = (fp, data_start + off0, dt, tuple(meta["shape"]))
    return index

For Maple layer 0 the interesting entries are:

(The weights are 2-bit packed ternary values stored as uint32, hence the tiny last dimension.)

4.2 Byte-range reads without the loops

With the index, reading a specific expert e of a tensor is a pure arithmetic jump: row_bytes = prod(shape[1:]) * dtype_itemsize, then f.seek(abs_off + e * row_bytes). We read only the requested rows into a numpy buffer and hand it to MLX. The whole 256-expert tensor is never touched:

class _DiskHolder:
    def read(self, sf_key, ids):
        fp, abs_off, dt_str, shape = self.index[sf_key]
        np_dt, mlx_dt = _SAFETENSORS_DTYPE[dt_str]
        row_elems = int(np.prod(shape[1:]))
        row_bytes = row_elems * np.dtype(np_dt).itemsize
        out = np.empty((len(ids), *shape[1:]), dtype=np_dt)
        with open(fp, "rb") as f:
            for i, e in enumerate(ids):
                f.seek(abs_off + int(e) * row_bytes)
                f.readinto(out[i].reshape(-1))
        ...
        return mx.array(out, dtype=mlx_dt)

To keep the active set in sync with the gate, _flash_set_active(ids) computes each expert’s resident slot once and reads all active rows:

ids = mx.array(ids, dtype=mx.int32)
self.active_ids = ids
# slot_of[e] = resident slot for expert e, or -1 if not resident
...
id_list = [int(e) for e in ids]
w_parts  = [self.disk.read(sk, id_list) for sk in self._src_weight]
self.weight = mx.concatenate(w_parts, axis=1) if self._concat else w_parts[0]

4.3 Fused projections and BF16 reinterpretation

The checkpoint stores expert projections unfused: up_proj and gate_proj as separate tensors. But after sanitize() the model wants a fused up_gate_proj. So the reader treats up_gate_proj as the concatenation of two source tensors and reads + concatenates both source rows per active expert:

proj_parts = ["up_proj", "gate_proj"] if name == "up_gate_proj" else ["down_proj"]
w_keys = [f"{base}.{p}.weight"   for p in proj_parts if ...]
...
self.weight = mx.concatenate(w_parts, axis=1) if self._concat else w_parts[0]

The per-row ternary scale α (row_alpha, BF16, one scalar per row) is expanded into per-group scales/biases at load time (ngroups = packed * 16 // group_size with group_size = 128), so downstream code sees the same affine-quantized layout it would from a fully resident checkpoint.

A subtle correctness trap: row_alpha is stored as raw 16-bit BF16 patterns. Casting the integer bit values to float16/float32 corrupts them. The fix is to reinterpret the top half of each 32-bit word as a float32 — which is exactly what BF16 is — and then convert to mx.bfloat16:

# Reinterpret the 16-bit patterns as bfloat16 (bf16 == top half of float32)
out = (out.astype(np.uint32) << 16).view(np.float32)
return mx.array(out, dtype=mx.float32).astype(mx.bfloat16)

A cast would change every α; a reinterpret changes nothing. This is one of those bit-exactness requirements that never shows up in a smoke test but shows up in a max_diff comparison against the full model.

4.4 IFP: masking the gate to the active set

IFP (inactive-expert-free policy) means: the gate is not allowed to route to experts that aren’t resident. After set_active, the gate masks inactive logits to -inf before the softmax, so all probability mass lands on the activated set. This guarantees every routed expert is in DRAM — no disk fetch inside the hot forward pass:

class MapleGate:
    def set_active(self, ids):
        ...
        self.active = ids
    def __call__(self, x):
        scores = x @ self.gate.weight.T
        if self.flash_active is not None:
            mask = (mx.arange(num_experts)[:, None] != self.active[None, :]).all(1)
            scores = mx.where(mask, -inf, scores)   # IFP
        ...

4.5 Per-token reselect

Which experts should be resident? The obvious first try — pick the active set from the prompt-average routing distribution — collapsed quality. Maple’s routing is highly token-specific; averaging over a long prompt is a poor predictor of what the next token needs.

The fix (_select in maple.py) computes the actual top-8 from the current token’s unmasked gate scores and makes that the active set:

gates = x @ self.gate.weight.T
inds, scores = group_expert_select(gates, self.gate.top_k)
unique_ids = [e for e in set(inds.reshape(-1).tolist())]
active = list(dict.fromkeys(shared + unique_ids))
self.switch_mlp.set_active(active)
self.gate.set_active(active)

By default this happens every token (--reselect-every 1), so the resident set tracks the routing exactly. With shared = 0 the resident set is the true top-8, which is the bit-exact regime.

Why does shared-experts default to 0? Quality loss is zero whenever active >= top_k + shared. Since active defaults to 8 and top_k is 8, any shared > 0 would evict a true top-8 expert and degrade output. shared only makes sense with a larger active.

4.6 LRU paging

Reselecting every token means the resident buffer is constantly changing. To avoid re-reading hot experts from disk, _flash_page implements a tiny LRU: an age counter per slot, bumped on every access; a miss pages the oldest slot with the requested expert (read from disk, written back into weight/scales via mx.concatenate slices):

def _flash_page(self, eid):
    slot = int(mx.argmin(self._age))     # least-recently-used slot
    wrow = self.disk.read(sk, [eid])      # one expert row from disk
    self.weight = mx.concatenate([self.weight[:slot], wrow, self.weight[slot+1:]])
    ...
    self._age = mx.where(aidx == slot, self._tick, self._age)

_flash_resolve translates the gate’s indices into resident slots and pages in anything missing. In the default IFP mode, every routed expert is already resident (no paging in the hot path); LRU is the fallback when the gate is allowed to route freely.

5. Quantizing the KV Cache

Expert weights aren’t the only thing that grows with the workload. The KV cache of the 6 full-attention layers grows with context length, and it’s the other source of peak memory.

mlx-lm gets a QuantizedKVCache with independent key and value bit-widths (k_bits / v_bits), threaded through base.py’s attention. We quantize keys to 8-bit and values to 4-bit (--kv-bits 8 --kv-v-bits 4) — roughly a 4× cut on the KV footprint. The sliding-window layers use RotatingKVCache, which does not support quantization in this build, so those layers keep full-precision KV (the window caps their size at 512 anyway).

Long contexts expose cache-growth paths that short prompts never hit. QuantizedKVCache.update_and_fetch grows the cache in 256-step chunks; on the first growth it called tree_map(expand_quant, ...) but expand_quant(x, el_per_int) required an argument that was never passed. A short prompt never grows past the initial allocation; a long one does. The fix is deleting the unused parameter:

def expand_quant(x):
    new_x = mx.zeros((*shape, x.shape[-1]), dtype=x.dtype)
    return mx.concatenate([x, new_x], axis=-2)

A second latent bug: RotatingKVCache.to_quantized() had a signature that didn’t accept v_bits, so quantizing raised TypeError (not the intended NotImplementedError) and crashed. Extending the signature so the “skip sliding-window caches” path works as designed.

6. Wiring Into mlx-lm

The offloader lives in mlx_lm/models/switch_layers.py, and prepare_flash_moe in mlx_lm/models/maple.py attaches it to every MoE layer:

• load(..., lazy=True) when flash-MoE is on, so un-selected experts are never faulted in by MLX’s lazy loader.
• prepare_flash_moe(model, active=8, shared=0, reselect_every=1, model_path=...) builds the file index once and initializes each layer’s resident buffer.

CLI flags (both generate and the server):

The OpenAI-compatible server (/v1/chat/completions) now also logs live throughput: prefill percentage + tok/s during prompt processing, then Prompt: N tokens, X tok/s and Generation: N tokens, Y tok/s per request.

7. Actual Testing!!!

7.1 Verifying losslessness

Losslessness is a testable claim, so we tested it:

• Load the full model (all 256 experts resident) and the flash-MoE model with --active-experts 8 --shared-experts 0.
• Compare every weight, scale, and bias tensor: max diff = 0.0. Bit-exact.
• Compare --active-experts 256 output to the baseline: identical.
• Generate end-to-end from both and confirm coherent, equivalent output.

Because active = top_k with shared = 0 makes the resident set the true top-8, the forward pass is mathematically identical to the full model — the 0.0 diff is the expected consequence, not a coincidence.

7.2 Memory numbers

Measured peak resident memory (MacBook Air M2, 16 GB unified, stock MLX):

That is a ~10× cut at the lossless setting. The resident footprint after load (0.62 GB) is essentially just the embedder, the non-expert layers, and the 8-expert buffers.

7.3 Why sampling matters

The 2-bit ternary weights make Maple’s greedy (temp=0) output loop or degenerate — this is a property of the 2-bit weights, not the offload. The same model is coherent with --temp 0.7 --top-p 0.9 --top-k 40 --min-p 0.05. Set those on the server; the default is greedy.

8. Results

What we ended up with:

• A ~10× RAM reduction at zero quality loss (0.63 GB, bit-exact).
• A tunable knob: raise --active-experts to trade memory for headroom (0.94 GB at 32, 5.89 GB at 256 = exact baseline).
• The full mlx-lm surface still works: generate, chat, OpenAI-compatible server with live tok/s metrics.
• The same code path found and fixed two latent long-context KV bugs.

The cost is real but well-defined: RAM is traded for disk I/O, and the trade is only a win when the model is memory-bound (which, on Apple Silicon, it is — the whole project exists because Maple can’t otherwise fit comfortably).

9. Limitations

• Prefill is disk-bound at long context. Batched prefill saturates the active set (up to all 256 experts per layer per step), so the offloader streams much of the checkpoint from disk. Flash-MoE wins on memory, not on prefill throughput.
• Generation is per-token disk-bound. Each reselect re-reads the active expert rows; short-run generation measured ~4 tok/s. Tune --reselect-every (at some cost in exactness) if disk latency dominates.
• KV quant skips sliding-window layers. RotatingKVCache stays full-precision (fine, since the window caps its size at 512).
• Long-context KV still grows. q8/q4 cuts it ~4×, but full-attention layers still scale with sequence length.
• 2-bit weights need non-greedy sampling. Not an offload artifact, but it must be configured to get usable output.
• Apple Silicon / MLX only. Built on the stock MLX runtime for portability, which is precisely why the byte-range approach was needed (MLX materializes whole tensors on any access).
• I/O-bound by design. The win assumes a reasonably fast local disk and a memory-constrained budget; on a RAM-rich box, keep experts resident.

10. Run It

python -m mlx_lm generate --model ./maple-2bit-mlx --trust-remote-code \
  --flash-moe --active-experts 8 --shared-experts 0 \
  --kv-bits 8 --kv-v-bits 4 \
  --prompt "Write a haiku about a grove." --temp 0.7 --top-p 0.9 --top-k 40

python -m mlx_lm server --model ./maple-2bit-mlx --trust-remote-code \
  --flash-moe --active-experts 8 --shared-experts 0 \
  --kv-bits 8 --kv-v-bits 4 --port 8080 \
  --temp 0.7 --top-p 0.9 --top-k 40 --min-p 0.05

11. Donations

If this work saves you money or a headache, a coffee (or a GPU-hour) is appreciated. Thanks for reading!

12. Reference List

• Repository: https://github.com/eli32-vlc/maple-flash-moe
• Maple (DeepGrove fork of mlx-lm): https://github.com/deepgrove-ai/mlx-lm-deepgrove
• Maple weights: https://huggingface.co/deepgrove/maple-2bit-mlx
• Base project: https://github.com/ml-explore/mlx-lm
• Safetensors file format (the byte-range index relies on it): https://github.com/huggingface/safetensors

Text and diagrams original. Maple and its 2-bit weights are from the DeepGrove fork of ml-explore/mlx-lm; this offload work is built on top of it.