3-bit TurboQuant quantization of google/diffusiongemma-26B-A4B-it — Google DeepMind's block-diffusion MoE — produced with TurboQuant-MLX. To our knowledge this is the first TurboQuant (Hadamard rotation + Lloyd-Max codebook) conversion of a diffusion LLM.

Status: experimental. Quantization quality is good (see below). As of turboquant-mlx-full 0.7.1 the runtime has a fused batched gather-GEMM for diffusion's 256-token canvas forwards (~4.6 tok/s, up from 1.6 at 0.7.0); the remaining gap to the mlx-community 4-bit (~35 tok/s) is partly per-op (attention kernel) and partly inherent — the 3-bit model needs more denoising iterations to converge than the 4-bit one. Use the 4-bit build if you want the fastest local DiffusionGemma; use this one for the smallest footprint or to follow TurboQuant's diffusion-LLM support.

Model Details

Base model: google/diffusiongemma-26B-A4B-it — block-diffusion Mixture-of-Experts (diffusion_gemma)
Architecture: 30 layers, hidden size 2816, 128 routed experts (top-8), ~25.2B total / ~3.8B active params; generates by iteratively denoising a 256-token canvas (≤48 steps, ~15–20 tokens per forward pass); multimodal (vision tower included)
Quantization: TurboQuant (Hadamard rotation + Lloyd-Max codebook), 3-bit, group size 32
- MoE experts (experts.gate_up_proj / experts.down_proj, ~90% of params) → 3-bit
- Attention (q/k/v/o_proj) → 3-bit
- Routers, dense per-layer MLP, self-conditioning block → bf16 (auto-skipped — the upstream model's own quantization recipe pins these at ≥8-bit, they are quant-sensitive)
- Vision tower, embeddings, norms → bf16
Size: 13.8 GB across 3 shards — vs ~48 GB BF16 (3.5x smaller) and 15.0 GB for MLX 4-bit

Will it fit my Mac?

Don't guess, and don't download 14.9 GB to find out — ask first (needs turboquant-mlx-full >= 0.15.1):

turboquant-plan --model manjunathshiva/diffusiongemma-26B-A4B-it-tq3-g32

It reads only this repo's safetensors headers over the network (a couple of hundred KB, a few seconds — the weights are never fetched), then projects the peak against your machine and prints the flags to use:

Model
  weights (exact)      14.9 GB
  KV cache             65.0 KB/token  (hybrid: 5/30 full-attention, 25 sliding (window 1024))

Verdict: <resident | needs a wired-limit raise | streaming>
Recommended: <the exact flags for your Mac>

Planning for a machine you're not sitting at — deciding whether a 16 GB mini or a 64 GB Mac can take this — is what --wired-gb / --ram-gb are for:

bash

turboquant-plan --model manjunathshiva/diffusiongemma-26B-A4B-it-tq3-g32 \
    --wired-gb 10.5 --ram-gb 16 --context 8192

turboquant-doctor runs the same projection plus a readiness check; both take --json. The projection is calibrated against real measurements on a 16 GB Mac mini rather than estimated from theory — on the 9.4 GB ternary 35B it predicts a 10.44 GB peak where that machine measures 10.42.

Quality & measured numbers (Apple Silicon, 64 GB)

	this repo (tq3-g32)	mlx-community 4-bit	BF16
Size on disk	13.8 GB	15.0 GB	~48 GB
Peak memory (text gen)	~18 GB	~19 GB	—
Generation speed	~4.6 tok/s (turboquant ≥0.7.1)	~35 tok/s	—

Spot checks (greedy decode):

Probe	Result
Multi-step math	✅ 17×23 = 391; 144/12 + 7×3 = 33 with clean PEMDAS steps
Code generation	✅ Correct iterative Fibonacci (minor wording artifact in a comment)
In-context recall	✅ Access code 7392 + server HELIOS recalled exactly
Free-form prose	✅ Coherent; occasional dropped/doubled word — same artifact class as the 4-bit build

Requirements

bash

# macOS with Apple Silicon (M1/M2/M3/M4/M5)
pip install "turboquant-mlx-full[vlm]>=0.7.1"   # pulls mlx-vlm >= 0.6.3

DiffusionGemma lives in mlx-vlm (not mlx-lm), and mlx-vlm's stock loader doesn't understand TurboQuant checkpoints — so generation goes through TurboQuant's generate_vlm entry point (shipped in 0.7.0). The two .py scripts bundled in this repo are a standalone alternative, but they predate the 0.7.1 gather-GEMM kernel and run slower — prefer the package entry point.

Quick Start

bash

python -m turboquant_mlx.generate_vlm \
    --model manjunathshiva/diffusiongemma-26B-A4B-it-tq3-g32 \
    --prompt "Write a short paragraph about the ocean." \
    --max-tokens 256 --temp 0.0

This loads the model with PolarQuantized layers and runs mlx-vlm's block-diffusion denoising sampler. An --image flag is available for multimodal prompts (the vision tower is kept at bf16).

How It Works

TurboQuant applies:

Hadamard rotation — random ±1 scaling + Walsh-Hadamard transform to gaussianize weights before quantization
Lloyd-Max codebook — optimal non-uniform scalar quantization
Group-wise scaling — per-group float16 scales (group size 32)

This achieves better quality than standard affine quantization at the same bit-width. Diffusion inference is a good fit for the scheme: every denoise step processes the whole 256-token canvas, so rotation overhead amortizes across the batch, and since 0.7.1 the expert matmuls run through a tiled batched gather-GEMM directly on the packed weights (nothing is dequantized to fp16).

Reproducing the conversion

bash

python -m turboquant_mlx.convert_vlm \
    --hf-path google/diffusiongemma-26B-A4B-it \
    --mlx-path ./diffusiongemma-26B-A4B-it-tq3-g32 \
    --bits 3 -g 32

Converts in-memory in ~20 s of quantize time on a 64 GB Mac (the 48 GB BF16 checkpoint is streamed lazily). Vision tower, routers, dense MLP and self-conditioning are skipped automatically for this architecture.

License

Apache 2.0, subject to the Gemma license terms (same as the base model). Quantization tooling: TurboQuant-MLX.

Citation

bibtex

@article{zandieh2025turboquant,
  title={TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate},
  author={Zandieh, Amir and Daliri, Majid and Hadian, Majid and Mirrokni, Vahab},
  year={2025},
  eprint={2504.19874},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2504.19874}
}