mlx-dllm

MLX (Apple Silicon) inference runtime for masked-diffusion language-model checkpoints produced by a2d, plus a native reference decoder for Qwen/Dream-family diffusion models.

a2d converts autoregressive checkpoints (GPT-2 today) into masked-diffusion models (MDLM). The saved artifact is a standard Hugging Face triple (config.json + tokenizer + model.safetensors) plus an "a2d" block in config.json. Bidirectional attention is a decode-time policy, not baked into the weights - this runtime supplies it.

What exists today

• mlx_dllm.load(path_or_repo) - loads GPT-2, Qwen2, Gemma (v1), and Gemma 3 (text) family HF checkpoints via mlx_lm (unmodified, as a library) and returns the parsed a2d config block when present. Per-family loading policy lives in the mlx_dllm/families/ adapter registry: a new model family is one new module with a register(...) call, no edits to shared dispatch (contract and worked example in mlx_dllm/families/__init__.py). Gemma 3 interleaves local (sliding-window) and global (full) attention layers; at a2d's alpha=1 decode policy the window is dropped and every layer runs full non-causal unwindowed attention, which the runtime reaches by neutralizing both the causal and sliding-window masks with no per-family code (see mlx_dllm/families/gemma3.py).
• mlx_dllm.bidirectional_forward(model, input_ids) - a full non-causal forward pass (no KV cache) returning logits for all positions, matching a2d's alpha=1 decode configuration.
• mlx_dllm.denoise(model, canvas, mask_token_id=..., steps=...) - the native Qwen/Dream correctness path: full bidirectional recomputation, greedy per-position predictions, and linearly scheduled confidence-ranked reveals with no KV cache or remasking. Predictions are read in-place (token for position i from logits at i), the a2d convention since a2d conversion drops the autoregressive next-token shift; pass logit_shift=True for published Dream checkpoints that keep the next-token head (token for position i from logits at i - 1).
• mlx_dllm.generate(...) - creates a fixed masked continuation canvas, denoises it with that reference path, and returns only the decoded continuation (prompt text excluded).
• Numerical parity gates prove the GPT-2, Qwen2, Gemma (v1), and Gemma 3 MLX forwards match PyTorch/HF eager bidirectional references. The Qwen and Gemma tests use tiny fixtures; the Gemma 3 gate also disables HF's sliding-window mask so both frameworks run full unwindowed attention. Full-size Dream/Gemma 3 validation is deliberately deferred to separate (Apple-silicon) hardware.

Acceleration (dual cache / confident-parallel decoding), an a2d-format bridge for Qwen, and a CLI are follow-on work.

Prior art

The Qwen/Dream diffusion techniques were informed by Fast-dLLM-mlx by MacPaw (Apache-2.0) and the original NVLabs/Fast-dLLM. They are reimplemented independently here: no Fast-dLLM-mlx code is copied, vendored, forked, imported, or added as a dependency. This implementation stays lean by reusing mlx-lm's own Qwen2 transformer and adding only the diffusion decoding layer.

Install

pip install -e .           # runtime: mlx + mlx-lm
pip install -e ".[test]"   # + torch/transformers for the parity gate
pytest

On macOS, mlx pulls the Metal backend automatically. On Linux the install resolves to mlx[cpu] (the ABI-matched CPU backend; the plain mlx wheel's libmlx.so is a Metal-only stub there), so the runtime is CPU-only and the Metal-gated parity tests skip.