| --- |
| library_name: mlx |
| license: apache-2.0 |
| language: |
| - en |
| pipeline_tag: text-generation |
| base_model: |
| - google/diffusiongemma-26B-A4B-it |
| - nvidia/diffusiongemma-26B-A4B-it-NVFP4 |
| - Singularity-Compute/super-decider |
| tags: |
| - mlx |
| - nvfp4 |
| - diffusion-gemma |
| - diffusion-lm |
| - structured-output |
| - decision-making |
| --- |
| |
| # super-decider-mlx |
|
|
| `Singularity-Compute/super-decider` — an NVFP4 4-bit 26B-A4B DiffusionGemma fine-tune for |
| multiple-choice decision tasks — converted to run on Apple silicon under MLX, with the |
| checkpoint's 4-bit experts preserved bit for bit and no quality loss from re-quantisation. |
|
|
| Text only. MLX only. Expects single-token labels. |
|
|
| ## What is in the box |
|
|
| | | | |
| |---|---| |
| | Parameters | 25.82 B total (22.8 B of them expert weights at 4 bits) | |
| | Size on disk | 18.8 GB — 8 safetensors shards, 1167 tensors | |
| | Experts | NVFP4: packed uint32 codes, one E4M3 scale per 16 inputs, one F32 `global_scale` per expert matrix | |
| | Everything else | BF16 — attention, router, dense MLP, embeddings, `lm_head` | |
| | Architecture | 30 layers, 128 experts, top-8, hidden 2816, vocab 262144, bidirectional per-question blocks | |
| | Vision tower | **removed** (the text path never reads it) | |
|
|
| The experts arrive **pre-stacked and gate/up-fused**, exactly as the model consumes them: |
|
|
| ``` |
| model.decoder.layers.<i>.experts.gate_up_proj.{weight,scales,global_scale} |
| model.decoder.layers.<i>.experts.down_proj.{weight,scales,global_scale} |
| ``` |
|
|
| so loading copies nothing and `mx.gather_qmm` takes the tensors directly. `global_scale` is the |
| checkpoint's own per-expert scale (`448 * 6 * weight_scale_2`). `manifest.json` records the |
| shard list, the tensor → shard map and the conversion's verification result. |
|
|
| **This layout is not loadable by `transformers`.** Use the code in this repo (below). |
| `transformers` is used only to read the tokenizer. |
|
|
| ## Fidelity |
|
|
| - **Weights are bit-exact.** The conversion copies the NVFP4 codes and E4M3 block scales byte |
| for byte and keeps each expert's `global_scale`; the round-trip error is zero, and the |
| dequantisation check over all 30 layers reports a worst relative error of `9.0e-08`. |
| - **The forward pass matches the reference implementation numerically.** `tests/test_parity_tiny.py` |
| builds a tiny random DiffusionGemma in `transformers` 5.21.0, loads the same state dict into |
| this MLX model, and compares: denoiser logits (plain and self-conditioned), encoder hidden |
| states (causal and with per-question bidirectional blocks), final token ids of a multi-canvas |
| `generate()` including EOS and adaptive stopping, the NVFP4 expert path, the BF16 GPU path, |
| and the real checkpoint's configs, RoPE tables and tensor names. 68 tests pass. |
| - Note where the resemblance ends: the reference deployment serves with vLLM on an L40S |
| (Marlin W4A16, FP8 KV cache, FP8 dense). This is the same arithmetic class on different |
| kernels, so it is not bit-identical to that deployment. **The published JevBench scores were |
| measured on the vLLM deployment, not on this port**, and have not been re-run here. |
|
|
| ## Quickstart |
|
|
| ```bash |
| pip install -r requirements.txt |
| hf download <this-repo> --local-dir ./sdc # ~19 GB |
| cd ./sdc && MODEL=$(pwd) PY=python3 bash runtime/run.sh |
| # engine :8000 (FastAPI, vLLM-compatible subset), decision server :8011 |
| ``` |
|
|
| Talk to the decision server. It takes a state and a map of questions whose answers are |
| single-token labels. |
|
|
| ```bash |
| curl -s 127.0.0.1:8011/v1/systemone -H 'content-type: application/json' -d '{ |
| "state": "Invoice VAT DE811907980, 1.234,00 EUR, reverse charge noted.", |
| "questions": { |
| "vat": {"type": "noul", "instructions": "Is VAT correctly applied?", |
| "criteria": {"true": "reverse charge applied", "false": "VAT charged or missing"}}, |
| "tone": {"type": "choice", "instructions": "Tone?", |
| "criteria": {"formal": "bureaucratic", "neutral": "plain", "warm": "friendly"}}, |
| "risk": {"type": "score", "instructions": "Risk?", |
| "criteria": ["none", "low", "medium", "high"]} |
| } |
| }' |
| ``` |
|
|
| `criteria` shape depends on the type: `noul` a `{true, false}` object, `choice` an |
| option-name → description object, `score` an ordered list of levels, `span` omitted. A wrong |
| shape is a `422` that names the question. Question ids must be single tokens, and the |
| instruction text is part of the prompt — a label you never mention in the state still has to |
| be named in `criteria`. `span` questions return character offsets and are less proven than the |
| label types. |
|
|
| ## Performance (M5 Pro, 48 GB, macOS) |
|
|
| | | | |
| |---|---| |
| | Load | ~2 s, 17.5 GiB resident, one copy of the weights | |
| | Prefill | 1.8 s @ 2 k tokens, 11 s @ 6 k (attention grows quadratically) | |
| | One decision read | ~0.23 s — a single denoiser step over a 128-slot canvas | |
| | 8 draws of one read | 1 forward, not 8 (the canvases batch); identical answers to 8 serial reads | |
| | Repeated state, new questions | instruction prefix KV is reused; ~0.5 s warm vs ~1.4 s cold for a 1.6 k-token prompt | |
|
|
| `SD_LOG=1` prints a line per read: prompt tokens, reused prefix, canvas, draws, encode and |
| denoise seconds. If `reused=` stays at 0 your prompts share no prefix and every request pays |
| full prefill. |
|
|
| ## Limits, stated plainly |
|
|
| - **Text only.** Images are rejected, not silently ignored. |
| - **One request at a time.** A single mutex around the engine: MLX is not thread-safe and the |
| prompt prefix is cached in one slot. Concurrency above one just queues. |
| - **Canvas ≤ 128, prompt ≤ 12288 tokens** (`--max-prompt`, default in `run.sh`; the model |
| itself claims 262144, but prefill cost and KV size make that impractical here). Over the cap |
| you get a real HTTP 400 that says so. |
| - **Prefix-cache reuse is not bit-exact across requests.** Re-encoding a suffix over a shorter |
| key dimension shifts the last BF16 bits: label probabilities move ~0.02 median / ~0.05 worst |
| in 14 draws at 1.5 k tokens, with the top label unchanged. `SD_PREFIX_CACHE=0` disables it and |
| costs a full prefill per request. |
| - **Designed for one label at a time.** The decision path reads one token per question slot. It |
| is not a general chat model; `generate()` exists for the thought channel and is slower and |
| less exercised. |
|
|
| ## Attribution |
|
|
| Apache-2.0. Converted from `Singularity-Compute/super-decider`, which is a fine-tune of |
| `google/diffusiongemma-26B-A4B-it` (Google DeepMind) merged into NVIDIA's |
| `nvidia/diffusiongemma-26B-A4B-it-NVFP4` ModelOpt checkpoint. `runtime/structured_server.py` |
| is derived from [mmastrac/djev](https://github.com/mmastrac/djev). See `NOTICE`. |
|
|
