---
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`.