Repository navigation
feat(api): add the TypeSafe System One API (/v1/systemone) on DiffusionGemma - #3897
MasakiMu319 wants to merge 1 commit into
Conversation
POST /v1/systemone answers typed noul/choice/score questions about a state with calibrated probabilities read from DiffusionGemma's diffusion decoder, following TypeSafe's OpenAPI 0.2.0. Where the contract is silent the read follows djev (the structured-read server of vllm-project/vllm#57250). A per-model setting, system_one_rereads_enabled (default on), keeps djev's auto noise re-reads or reads once for lower latency.
|
Thanks for this. I tried the branch locally and ran the same JevBench public set (231 items, I have been trying a different read path for the same endpoint: scoring typed questions with the chat models oMLX already serves, from the next-token logits of the answer labels, as featherless-ai/simple-jev does. The prototype is standalone, not wired into oMLX's engine: it runs as its own process outside oMLX's scheduler, one request at a time, and caches the system prompt and question list across requests. It uses Simple Jev's baseline prompt and scoring (thinking off;
The latencies are not a like-for-like comparison: the PR row has re-reads on (your own ablation shows roughly half the latency with them off, at the same accuracy), and the prototype's single-process, single-request setup and prompt caching described above also favor it. The models differ in size too (weight files 16.5 GB, 22.1 GB and 7.0 GB). So on this set the logits path with Qwen3.6 lands close to this PR, while the much smaller e4b trades hard-tier accuracy and calibration for speed. The logits path works with ordinary chat models, no dedicated checkpoint needed. The PR and the Qwen3.6 prototype give the same correct/incorrect outcome on 211 of 231 items. Happy to share the prototype code if that's useful. Would a follow-up be welcome that answers |
|
I would really like to see this land, with a path for Cloudflare’s two Apache-2.0 open-weight decision models, Clef (27B) and Clef-Flash (9B): These look like compelling candidates for everyday local SystemOne inference. They implement the same state + typed questions interface, but use a dedicated joint decision head to score the options across questions rather than diffusion reads or ordinary next-token label scoring. MLX conversions with a custom runner are already available, for example https://huggingface.co/mlx-community/clef-flash-4bit . An agent-assisted local smoke test on my M3 Max (64 GB), using current llama.cpp with Metal and Q4_K_M weights, passed all 30 checks for both models across choice, noul, and score. Median latency for three questions over short inputs was about 549 ms for Flash and 1.82 s for Clef. That is only a functional check, not a quality comparison with Jev or the DiffusionGemma results here. What I especially want is oMLX’s existing model management around this capability: download and configure a decision model in the UI, call /v1/systemone, and have it auto-load and unload under the same memory budget as my chat models. A separate server gets me the API, but misses that integration. Would it make sense to keep the SystemOne request/response contract backend-independent, with DiffusionGemma reads as the first implementation and Clef’s joint head as a follow-up? That seems compatible with the non-diffusion direction discussed above, while giving dedicated decision models their own execution path. Happy to help test the Clef side. |
3 participants
Summary
Adds
POST /v1/systemone: TypeSafe's System One API (the Jev API, OpenAPI 0.2.0), served by DiffusionGemma. Instead of generating text, oMLX reads each typednoul/choice/scoreanswer directly from the diffusion decoder. It prefills the prompt once and runs a read-only decoder pass over a canvas seeded with the answer template, noise tokens in the answer slots only. Each answer is a calibrated probability distribution over the question's labels, so it cannot go off-schema. TypeSafe's SDKs work unchanged whenTYPESAFE_BASE_URLpoints at oMLX.Related: #1805 (DiffusionGemma support).
Design decisions to review
confidenceis(K·max(p) − 1)/(K − 1), which reproduces every worked example in TypeSafe's docs; djev reportsmax(p).diffusion_gemma, because the read relies on DiffusionGemma's chat format and special tokens. Any other model gets a 400.images,steps,samples,think,sequential,span).GET /v1/modelskeeps its OpenAI shape, so the SDKs' model listing isn't supported. Aliasing the model asjev-latestmakes the SDKs' default model work.system_one_rereads_enabled(default on, which is djev's policy). When a read is uncertain it takes three more noise draws and averages them; turning it off reads once, for about half the latency. It is a universal (request-time) profile field, so a profile or exposed profile can switch it without reloading the engine. The web admin and the macOS app show it for diffusion models only.Questions for you: is the endpoint name/path OK, and should the re-read default stay on?
Implementation
omlx/api/systemone_models.py: the request and response models.omlx/api/systemone.py: request planning and the reads.VLMBatchedEngine.diffusion_read_session(): prefills once on the MLX executor while holding the diffusion lane's lock.session.read()runs read-only decoder passes against that prefill.omlx/server.py: the/v1/systemoneroute (engine lease, JSON keepalive, metrics).docs/TESTING.md.Benchmark
Setup: M5, 32 GB;
mlx-community/diffusiongemma-26B-A4B-it-4bit; mlx-vlmea79808.Accuracy: JevBench public items. 231 items (easy 48, standard 72, hard 111),
typesafeadapter, one request at a time. The other rows are JevBench's published per-task outcomes for the same items.Re-read ablation. Same setup with
system_one_rereads_enabledoff:Latency. HTTP p50 with the model warm:
Reproduction: start
omlx serve, then from a JevBench checkout runTYPESAFE_API_KEY=local python -m jevbench.cli run --tasks datasets/public/easy.jsonl,datasets/public/original.jsonl,datasets/public/hard.jsonl --adapter typesafe --endpoint http://127.0.0.1:8000 --model <model id> --price-in-per-m 0 --price-out-per-m 0 --results results.jsonlfollowed by
jevbench.cli summarize. Compare the results withresults/v1.2/jevbench-v1.2-per-task.json.Validation
python -m pyteston this branch: 14774 passed, 1 failed. The failure,test_cluster_rdma_mailbox.py::test_a_daemon_ending_the_registration_is_noticed, is timing-dependent: it passes 3 out of 3 times on its own and is unrelated to this change.node tests/system_one_rereads_ui.test.cjs: checks the toggle's state, visibility, PUT payload and profile values.OMLX_DIFFUSIONGEMMA_MODEL=… pytest tests/test_systemone.py -m slowpasses: all 255 labels are single tokens, long mixed question lists resolve, and one real read runs.TypeSafeBadRequestError.oMLXTestsruns 326 tests with 1 failure,ThemeTests.testLightWindowBackgroundUsesStandardWindowColor, which also fails on a cleanmainhere (macOS 27 beta system colors).moe_expert_offload_ui,network_auth_uiandqwen_oq_a8_uialso fail on a cleanmainhere (Node 26); this PR doesn't change them.Not tested
jev-latestalias path (the SDK was tested withTYPESAFE_DEFAULT_MODEL).Screenshots
Web admin, DiffusionGemma model settings: