m1rkocasu/LLaDA2.2-mini-MLX-5bit

LLaDA2.2-mini for MLX, 5-bit

inclusionAI/LLaDA2.2-mini on Apple silicon: a 16B mixture-of-experts diffusion language model, 1.4B parameters active per token, quantized to 5 bits. It writes a block of 32 tokens at a time, filling masks in parallel and editing its own draft, instead of one token after another. Also available: 4-bit, 6-bit.

size on disk	11.2 GB (5.5 bits per weight; the router stays in full precision)
peak memory	10.5 GiB
speed, M4 Pro	53–91 tokens/s on code, 22–28 on prose
runs with	pip install mlx-vlm mlx-lm and the code in this repo, nothing else
verified against	the reference PyTorch implementation, layer by layer and token by token

Use

Everything needed is in this repo:

pip install -U mlx-vlm mlx-lm
hf download m1rkocasu/LLaDA2.2-mini-MLX-5bit --local-dir LLaDA2.2-mini-MLX-5bit
python LLaDA2.2-mini-MLX-5bit/generate.py --prompt "Write a Python function that checks for palindromes."

generate.py takes every option of mlx_vlm generate (--max-tokens, --temperature, --threshold, ...). From Python:

import sys
sys.path.insert(0, "LLaDA2.2-mini-MLX-5bit")
import llada22_mlx  # LLaDA2.2 support for mlx-vlm; import it before loading

from mlx_vlm import load, generate
from mlx_vlm.prompt_utils import apply_chat_template

model, processor = load("LLaDA2.2-mini-MLX-5bit", trust_remote_code=True)
prompt = apply_chat_template(processor, model.config, "What is 17 * 23?")
print(generate(model, processor, prompt, max_tokens=512, temperature=0.0).text)

trust_remote_code is for the model's own tokenizer class. The decoding defaults are the reference's: threshold 0.5, editing threshold 0.0, 16 refinement passes, blocks of 32. The block length is fixed by the model's routing and cannot be changed.

What llada22_mlx is. Two files of mlx-vlm's own LLaDA2 implementation (language.py, config.py; MIT, license included) with the 2.2 additions below. Importing the package makes mlx-vlm use them instead of its own; nothing inside the installed mlx-vlm is modified, and LLaDA2.0/2.1 models still load as before. It was written against mlx-vlm 0.7.1 and warns if the installed version's LLaDA2 files differ.

Where it does not run. LM Studio and oMLX decode autoregressively and cannot run a diffusion model.

Why a new conversion

mlx-vlm already runs LLaDA2.0 and 2.1. Loading 2.2 with it gives a model that answers, and answers wrong, without an error. Three things were missing:

1. Block routing. 2.2 picks experts per block of 32 tokens: the block first keeps 48 of the 256 experts, then each token takes its 8 among those. mlx-vlm used 2.1's group-limited routing, which on 2.2 chooses different experts for practically every token (64 of 64 in our control).
2. The 2.2 sampler. Besides filling masks, 2.2 rewrites tokens it has already written, and emits DELETE and INSERT tokens that shrink or grow its draft, with an anti-loop resampler and a few refinement passes per block. The old sampler left the edit tokens in the text as if they were words.
3. The tokenizer, after conversion. Converting rewrites tokenizer_config.json and drops trust_remote_code. On transformers 5 the reloaded tokenizer then rebuilds its BPE from the vocabulary alone, with no merges, and encodes one character per token: 76 tokens instead of 30 for the same prompt. The model receives text unlike anything it was trained on and replies with an empty turn, which looks like a broken model. The tokenizer files in this repo are correct.

How it was verified

The decoder, layer by layer. The bf16 checkpoint is 30 GB and fits neither in RAM nor under the GPU ceiling of a 24 GB Mac, so each layer is built, loaded, compared and freed on both sides. Run in float32 on both sides, so that rounding cannot hide logic:

	result
hidden-state error per layer	1e-9 to 8e-7, at most 2.6× the gap between PyTorch on CPU and PyTorch on MPS
experts chosen by the router	identical on 96 of 96 tokens, in all 19 MoE layers
argmax of the logits on the masked positions	identical, 100 %

In bf16 the two implementations differ more, and a few near-tied experts flip. That is rounding (MLX computes RoPE in float32, the reference in bf16), which the float32 run rules in.

The sampler, token by token. The reference generate is extracted from modeling_llada2_moe.py and driven with the logits of the MLX model, so both samplers see the same numbers. On three prompts (arithmetic, Italian, code) the two produce the same tokens, 167, 104 and 256 of them. With the 2.1 sampler put back, the same test diverges on all three prompts (at tokens 67, 29 and 47): the test can fail.

The download path. This repo's generate.py and the Python example above were run in a fresh environment with mlx-vlm 0.7.1 and transformers 5.17 from PyPI.

Compared with LLaDA2.1-mini

A short battery of 12 tasks, each asked in English and in Italian, with answers checked by code: exact numbers, generated functions run against test cases, required words, JSON that must parse to the expected object. Greedy decoding, blocks of 32, each model with its own family's default sampler. Speed is measured on six longer answers (explanation, history, code; 230–512 tokens each), where the fixed cost of a block matters less.

	LLaDA2.1-mini 4-bit (mlx-community)	LLaDA2.2-mini 4-bit	LLaDA2.2-mini 5-bit (this repo)	LLaDA2.2-mini 6-bit
correct, English	12 / 12	12 / 12	12 / 12	12 / 12
correct, Italian	10 / 12	11 / 12	10 / 12	10 / 12
tokens/s, long answers, English	45.6	60.7	39.0	38.8
tokens/s, long answers, Italian	33.7	41.1	30.6	29.3
garbled words in long prose	a few	several	few, mild	few, mild
size on disk	9.2 GB	9.2 GB	11.2 GB	13.2 GB
peak memory	8.6 GiB	8.6 GiB	10.5 GiB	12.4 GiB

The misses are informative: every build gets the four Swiss national languages right in English and wrong in Italian ("romano", "romances" or English in place of romancio; 2.1 also lists English instead of French). 2.1 answered 21 instead of 31 to "3 boxes of 12 apples, eat 5"; the 5- and 6-bit builds answered 195 minutes instead of 155 for a train trip. With 12 tasks per language, one task is noise: the table says all four work, not which is best by a point.

Which build. The 4-bit build is the fastest and the smallest, by a clear margin: a diffusion step reads the weights of every expert its block keeps, so bytes per weight set the pace. The 5-bit build writes cleaner long prose at about two thirds of the speed. The 6-bit build was not measurably better than the 5-bit one in these tests, and costs 2 GB more. For code and short, checkable answers, 4 bits; for long text, especially in languages other than English, 5 bits.

Speed depends on confidence. A diffusion model finishes a block when it is sure of it. Code, where the next tokens are predictable, runs fastest; prose is slower, and Italian prose slower than English, because the model is less certain of it and spends more refinement passes per block.

Limits

Garbled words in long prose. In answers of a few hundred words, a few words come out broken: Sveraera, costrì, tem di di di, a a in the 4-bit build. They are an artefact of deciding many tokens in parallel, and they are worse at 4 bits than at 5 or 6: the same Italian text had fewer and milder ones at 5 bits (la la, della trasporto) and at 6 (svizzi, tecnnologici), and the English one at 5 bits none. This is a reading of a handful of texts, not a measurement. Raising the decoding threshold from 0.5 to 0.9 did not remove them. Short answers, code and structured output were clean in every test. Whether the bf16 model shows them too could not be checked here: it does not fit on the GPU.

Facts. With 1.4B active parameters the model invents dates and lengths freely in open-ended prose (the Gotthard tunnel "completed in 1187"). Use it for tasks with checkable answers, not as a reference.

Credits

Model by inclusionAI, Apache-2.0; see the original card and the technical report. The MLX implementation is mlx-vlm's LLaDA2 support by Prince Canuma and contributors (MIT), extended here for 2.2.