--- library_name: mlx license: apache-2.0 license_link: https://ai.google.dev/gemma/docs/gemma_4_license pipeline_tag: image-text-to-text base_model: google/diffusiongemma-26B-A4B-it tags: - mlx - turboquant - moe - diffusion - block-diffusion --- # diffusiongemma-26B-A4B-it-tq3-g32 **3-bit** TurboQuant quantization of [google/diffusiongemma-26B-A4B-it](https://huggingface.co/google/diffusiongemma-26B-A4B-it) — Google DeepMind's **block-diffusion MoE** — produced with [TurboQuant-MLX](https://github.com/manjunathshiva/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](https://huggingface.co/mlx-community/diffusiongemma-26B-A4B-it-4bit) (~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`): ```bash 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: Recommended: ``` 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](https://huggingface.co/manjunathshiva/Qwen3.6-35B-A3B-tq3a-tqTe-g64) 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: 1. **Hadamard rotation** — random ±1 scaling + Walsh-Hadamard transform to gaussianize weights before quantization 2. **Lloyd-Max codebook** — optimal non-uniform scalar quantization 3. **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](https://ai.google.dev/gemma/docs/gemma_4_license) (same as the base model). Quantization tooling: TurboQuant-MLX. Copyright 2026 Manjunath Janardhan. ## 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} } ```