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

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.

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. See NOTICE.