From 32f42b32d9511561b070b457260618c47d9d8aae Mon Sep 17 00:00:00 2001 From: Sam Snelling Date: Thu, 17 Sep 2026 06:33:28 -0500 Subject: [PATCH 1/4] Add structured read inference backends --- README.md | 11 +- docs/api.md | 8 +- docs/diffusion-gemma.md | 56 + docs/how-it-works.md | 12 +- docs/quickstart.md | 4 +- docs/vllm-metal.md | 62 + server/src/system_one_lite/api.py | 30 +- ...iongemma_26b_a4b_it_4bit_answer_codes.json | 2898 +++++++++++++++++ server/src/system_one_lite/engine.py | 64 +- .../src/system_one_lite/mlx_vlm_diffusion.py | 298 ++ server/src/system_one_lite/vllm_metal.py | 250 ++ server/tests/integration/test_api.py | 95 + server/tests/unit/test_mlx_vlm_diffusion.py | 242 ++ server/tests/unit/test_vllm_metal.py | 301 ++ server/tools/gen_answer_codes.py | 13 +- 15 files changed, 4314 insertions(+), 30 deletions(-) create mode 100644 docs/diffusion-gemma.md create mode 100644 docs/vllm-metal.md create mode 100644 server/src/system_one_lite/data/diffusiongemma_26b_a4b_it_4bit_answer_codes.json create mode 100644 server/src/system_one_lite/mlx_vlm_diffusion.py create mode 100644 server/src/system_one_lite/vllm_metal.py create mode 100644 server/tests/unit/test_mlx_vlm_diffusion.py create mode 100644 server/tests/unit/test_vllm_metal.py diff --git a/README.md b/README.md index 38fb555..961761a 100644 --- a/README.md +++ b/README.md @@ -48,9 +48,9 @@ For each question, the server: 4. Keeps only the logits for valid answer codes and applies softmax. 5. Maps those probabilities back to the answer names supplied by the caller. -Questions are evaluated independently. Each question gets its own prompt and -another copy of the state. This prevents one answer from affecting another, -but runtime and input-token use grow with the number and length of questions. +The default MLX and vllm-metal backends evaluate questions independently. The +DiffusionGemma backend instead places all answers in one shared canvas and +runs one denoising forward. See [How it works](docs/how-it-works.md) for the prompt format, answer-code registry, token limits, and confidence calculation. @@ -275,6 +275,11 @@ Select a profile before starting the server: SYSTEM_ONE_MODEL=larger uv run uvicorn system_one_lite.api:app --port 8010 ``` +To serve the same model through vllm-metal, see the +[vllm-metal backend guide](docs/vllm-metal.md). +For seeded, single-forward reads with the 4-bit DiffusionGemma model, see the +[DiffusionGemma backend guide](docs/diffusion-gemma.md). + The demo, benchmark, and evaluation tools also accept `--model default` or `--model larger`. Each supported model has a checked-in registry of answer codes that must remain single tokens with the pinned tokenizer. diff --git a/docs/api.md b/docs/api.md index 4ec1bbf..372c50c 100644 --- a/docs/api.md +++ b/docs/api.md @@ -99,8 +99,8 @@ How likely the statement is to be true. | --- | --- | --- | | `model` | `string` | The engine's own model ID. The default is `mlx-community/Qwen3-1.7B-4bit`. | | `answers` | `map` | One answer per question, keyed by your IDs, in request order. | -| `usage.input_tokens` | `integer` | Total tokens evaluated across the independent question prompts. | -| `usage.output_tokens` | `integer` | Always 0. Nothing is generated. | +| `usage.input_tokens` | `integer` | Prompt tokens evaluated by the selected backend. | +| `usage.output_tokens` | `integer` | 0 for local MLX and read-only DiffusionGemma. The vllm-metal backend uses one internal output token per question for each group of up to 128 options. | ### Choice answer @@ -152,5 +152,5 @@ You get an error, never a wrong answer. Only one inference request runs at a time. If the engine is already in use, the server returns `503` with `detail: "the inference engine is busy"`. -Each question uses a full independent prompt. `usage.input_tokens` counts -the state again for each question. +The default MLX and vllm-metal backends count one full prompt per question. +The DiffusionGemma backend counts one shared prompt for the whole request. diff --git a/docs/diffusion-gemma.md b/docs/diffusion-gemma.md new file mode 100644 index 0000000..0007ead --- /dev/null +++ b/docs/diffusion-gemma.md @@ -0,0 +1,56 @@ +# Run with 4-bit DiffusionGemma + +This backend uses `mlx-community/diffusiongemma-26B-A4B-it-4bit` at the pinned +revision in `engine.py`. It sends every question in one prompt and places one +`- label` answer slot per question in DiffusionGemma's active canvas. This shape +fits System One's limit of 64 questions. One read-only model pass returns the +requested answer-code scores. The canvas uses the same fixed seed for each +request, so the same input has the same canvas. + +The backend needs the `feature/diffusion-gemma-reads` branch of MLX-VLM. That +branch adds `/v1/diffusion/reads` without changing normal generation. + +Start the MLX-VLM server from its checkout: + +```bash +cd /path/to/mlx-vlm +uv run --with-editable . python -m mlx_vlm.server --port 8080 +``` + +Start System One in another terminal: + +```bash +cd server +SYSTEM_ONE_BACKEND=mlx-vlm-diffusion \ + uv run uvicorn system_one_lite.api:app --port 8010 +``` + +Use the compact profile for short, latency-sensitive questions: + +```bash +SYSTEM_ONE_BACKEND=mlx-vlm-diffusion-fast \ + uv run uvicorn system_one_lite.api:app --port 8010 +``` + +The compact profile accepts exactly one question. It removes the long prompt +instructions and uses one active canvas token. Use it only when the state, +question, and answer labels make the task clear without extra guidance. +It runs six full encoder layers. Later layers build lighter attention caches. +This approximation needs an accuracy check for each target task. + +The first System One startup downloads the pinned 4-bit model snapshot. The +weights are about 16.5 GB. Both processes use the same Hugging Face cache. + +The read endpoint accepts prompt tokens, an active seed canvas, and exact token +IDs for each answer slot. It checks every token and position before the model +runs. It scores only the requested answer codes and skips the full vocabulary +head. System One then applies its usual softmax over each question's valid +answer codes. + +This path differs from the vllm-metal backend in two ways: + +- All questions share one prompt and one denoising forward. +- The answer slots can attend to the fixed answer template and to one another. + +The public System One request and response formats do not change. Read-only +canvas work reports zero output tokens because no tokens are committed. diff --git a/docs/how-it-works.md b/docs/how-it-works.md index 3a060a2..310d92c 100644 --- a/docs/how-it-works.md +++ b/docs/how-it-works.md @@ -1,9 +1,13 @@ # How it works -The engine (`server/src/system_one_lite/engine.py`) never decodes text. It builds one full prompt -for each question. The next token must be an answer code. The engine runs a -forward pass and reads the code probabilities at that position. It generates -no tokens. +The default engine (`server/src/system_one_lite/engine.py`) never decodes text. +It builds one full prompt for each question. The next token must be an answer +code. The engine runs a forward pass and reads the code probabilities at that +position. It generates no tokens. + +The optional DiffusionGemma backend keeps the same public API but uses one +shared prompt and a seeded answer canvas. See [Run with 4-bit +DiffusionGemma](diffusion-gemma.md). This page walks through the pieces so you can read the code alongside it. diff --git a/docs/quickstart.md b/docs/quickstart.md index ea8591e..45cbeb6 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -149,8 +149,8 @@ Things worth noticing: never asks for money back. - These probabilities are model scores, not calibrated odds. A value near one does not prove that an answer is correct. -- `usage.output_tokens` is always 0. The engine never generates. Each question - uses a full independent prompt. +- `usage.output_tokens` is 0 for the default and read-only DiffusionGemma + backends. The default backend uses a full independent prompt per question. ## 4. Use the SDK (optional) diff --git a/docs/vllm-metal.md b/docs/vllm-metal.md new file mode 100644 index 0000000..30e64cb --- /dev/null +++ b/docs/vllm-metal.md @@ -0,0 +1,62 @@ +# vllm-metal backend + +System One can use a separate vllm-metal server instead of loading the model +in the API process. The HTTP API and answer shapes stay the same. + +Install vllm-metal in its own Homebrew environment: + +```bash +brew tap vllm-project/vllm-metal https://github.com/vllm-project/vllm-metal +brew install vllm-project/vllm-metal/vllm-metal +``` + +Start vllm-metal with the same pinned model revision as the answer-code +registry: + +```bash +vllm serve mlx-community/Qwen3-1.7B-4bit \ + --revision 3b1b1768f8f8cf8351c712464f906e86c2b8269e \ + --max-model-len 32769 \ + --port 8000 +``` + +The extra model-length slot is for vllM's internal output token. System One +still limits each prompt to 32,768 input tokens. + +Then start System One in another terminal: + +```bash +cd server +SYSTEM_ONE_BACKEND=vllm-metal \ +SYSTEM_ONE_VLLM_BASE_URL=http://127.0.0.1:8000 \ +uv run uvicorn system_one_lite.api:app --port 8010 +``` + +The backend builds each prompt with System One's pinned tokenizer. It sends +the exact prompt token IDs in completion batches. It also asks vllm-metal +for the exact token IDs of every allowed answer code, then normalizes those +scores at System One's temperature. It does not use a top-k cutoff or guess a +score for a missing answer. Question IDs are still used only to route answers +and are never sent to the model. + +vLLM accepts up to 128 requested token IDs in one read. A question with more +options uses more than one exact read. Each read makes vllm-metal sample one +internal token so it can return the requested log probabilities. The System +One response reports that work in `usage.output_tokens`. The local MLX backend +still reports zero output tokens. + +Before inference, the backend checks both token limits. It also asks +vllm-metal to decode the full answer-code registry. A server tokenizer +mismatch stops the request instead of reading the wrong token IDs. + +## Current limit + +This backend does not use the seeded parallel canvas from vLLM PR 57250. Use +the [4-bit DiffusionGemma backend](diffusion-gemma.md) for that execution +shape. MLX-VLM already implements DiffusionGemma on Apple Silicon, so the +seeded read endpoint lives there instead of adapting vllm-metal's +autoregressive scheduler. + +Keep vllm-metal outside the `server` uv environment. Its release pins an exact +MLX build for its native Metal extension, while System One has its own MLX +dependency. diff --git a/server/src/system_one_lite/api.py b/server/src/system_one_lite/api.py index 5144e80..b3ba23e 100644 --- a/server/src/system_one_lite/api.py +++ b/server/src/system_one_lite/api.py @@ -11,6 +11,7 @@ from fastapi import FastAPI, HTTPException from .engine import Engine, RequestContractError +from .mlx_vlm_diffusion import MlxVlmDiffusionEngine, MlxVlmError from .prompts import as_text, confidence, label from .schemas import ( ChoiceAnswer, @@ -22,11 +23,24 @@ ScoreAnswer, Usage, ) +from .vllm_metal import VllmMetalEngine, VllmMetalError def configured_engine(): """Build the process-wide engine from a named profile or exact model ID.""" - return Engine(os.environ.get("SYSTEM_ONE_MODEL")) + backend = os.environ.get("SYSTEM_ONE_BACKEND", "mlx") + if backend == "mlx": + return Engine(os.environ.get("SYSTEM_ONE_MODEL")) + if backend == "vllm-metal": + return VllmMetalEngine(os.environ.get("SYSTEM_ONE_MODEL")) + if backend == "mlx-vlm-diffusion": + return MlxVlmDiffusionEngine(os.environ.get("SYSTEM_ONE_MODEL")) + if backend == "mlx-vlm-diffusion-fast": + return MlxVlmDiffusionEngine( + os.environ.get("SYSTEM_ONE_MODEL"), + compact=True, + ) + raise ValueError(f"unknown SYSTEM_ONE_BACKEND {backend!r}") def labels_for(question): @@ -86,13 +100,20 @@ def evaluate(req: EvaluateRequest): raise HTTPException(503, "the inference engine is busy") runtime_engine = app.state.engine + engine_questions = [ + (as_text(question.instructions), labels_for(question)) for _, question in items + ] try: results, input_tokens, elapsed_ms = runtime_engine.evaluate( as_text(req.state), - [(as_text(question.instructions), labels_for(question)) for _, question in items], + engine_questions, ) except RequestContractError as error: raise HTTPException(422, str(error)) from error + except VllmMetalError as error: + raise HTTPException(502, str(error)) from error + except MlxVlmError as error: + raise HTTPException(502, str(error)) from error finally: app.state.inference_slot.release() @@ -107,7 +128,10 @@ def evaluate(req: EvaluateRequest): question_id: answer_for(question, probabilities) for (question_id, question), probabilities in zip(items, results) }, - usage=Usage(input_tokens=input_tokens, output_tokens=0), + usage=Usage( + input_tokens=input_tokens, + output_tokens=runtime_engine.output_tokens_for(engine_questions), + ), ) return app diff --git a/server/src/system_one_lite/data/diffusiongemma_26b_a4b_it_4bit_answer_codes.json b/server/src/system_one_lite/data/diffusiongemma_26b_a4b_it_4bit_answer_codes.json new file mode 100644 index 0000000..496ba52 --- /dev/null +++ b/server/src/system_one_lite/data/diffusiongemma_26b_a4b_it_4bit_answer_codes.json @@ -0,0 +1,2898 @@ +{ + "model": "mlx-community/diffusiongemma-26B-A4B-it-4bit", + "model_revision": "a7a81407613811e8ba63af92ac0d852b809e191f", + "tokenizer_sha256": "78d2ea3b75a7968d8d9c5dca2dcbddde6321191ff4c576acd6fe19b957b85be1", + "slot_context": "codes occupy one seeded DiffusionGemma canvas token after a list marker", + "codes": [ + { + "code": "A", + "token": " A", + "token_id": 562 + }, + { + "code": "B", + "token": " B", + "token_id": 603 + }, + { + "code": "C", + "token": " C", + "token_id": 565 + }, + { + "code": "D", + "token": " D", + "token_id": 622 + }, + { + "code": "E", + "token": " E", + "token_id": 645 + }, + { + "code": "F", + "token": " F", + "token_id": 633 + }, + { + "code": "G", + "token": " G", + "token_id": 667 + }, + { + "code": "H", + "token": " H", + "token_id": 640 + }, + { + "code": "I", + "token": " I", + "token_id": 564 + }, + { + "code": "J", + "token": " J", + "token_id": 730 + }, + { + "code": "K", + "token": " K", + "token_id": 751 + }, + { + "code": "L", + "token": " L", + "token_id": 639 + }, + { + "code": "M", + "token": " M", + "token_id": 587 + }, + { + "code": "N", + "token": " N", + "token_id": 646 + }, + { + "code": "O", + "token": " O", + "token_id": 708 + }, + { + "code": "P", + "token": " P", + "token_id": 593 + }, + { + "code": "Q", + "token": " Q", + "token_id": 1366 + }, + { + "code": "R", + "token": " R", + "token_id": 630 + }, + { + "code": "S", + "token": " S", + "token_id": 555 + }, + { + "code": "T", + "token": " T", + "token_id": 558 + }, + { + "code": "U", + "token": " U", + "token_id": 749 + }, + { + "code": "V", + "token": " V", + "token_id": 782 + }, + { + "code": "W", + "token": " W", + "token_id": 649 + }, + { + "code": "X", + "token": " X", + "token_id": 1684 + }, + { + "code": "Y", + "token": " Y", + "token_id": 895 + }, + { + "code": "Z", + "token": " Z", + "token_id": 1491 + }, + { + "code": "AA", + "token": " AA", + "token_id": 23952 + }, + { + "code": "AB", + "token": " AB", + "token_id": 14802 + }, + { + "code": "AC", + "token": " AC", + "token_id": 11659 + }, + { + "code": "AD", + "token": " AD", + "token_id": 7747 + }, + { + "code": "AE", + "token": " AE", + "token_id": 31616 + }, + { + "code": "AF", + "token": " AF", + "token_id": 18026 + }, + { + "code": "AG", + "token": " AG", + "token_id": 17692 + }, + { + "code": "AH", + "token": " AH", + "token_id": 46670 + }, + { + "code": "AI", + "token": " AI", + "token_id": 12498 + }, + { + "code": "AJ", + "token": " AJ", + "token_id": 43155 + }, + { + "code": "AK", + "token": " AK", + "token_id": 31194 + }, + { + "code": "AL", + "token": " AL", + "token_id": 9634 + }, + { + "code": "AM", + "token": " AM", + "token_id": 8151 + }, + { + "code": "AN", + "token": " AN", + "token_id": 7790 + }, + { + "code": "AO", + "token": " AO", + "token_id": 62715 + }, + { + "code": "AP", + "token": " AP", + "token_id": 9943 + }, + { + "code": "AQ", + "token": " AQ", + "token_id": 95915 + }, + { + "code": "AR", + "token": " AR", + "token_id": 8886 + }, + { + "code": "AS", + "token": " AS", + "token_id": 9447 + }, + { + "code": "AT", + "token": " AT", + "token_id": 9928 + }, + { + "code": "AU", + "token": " AU", + "token_id": 43777 + }, + { + "code": "AV", + "token": " AV", + "token_id": 28523 + }, + { + "code": "AW", + "token": " AW", + "token_id": 65292 + }, + { + "code": "AX", + "token": " AX", + "token_id": 74418 + }, + { + "code": "AY", + "token": " AY", + "token_id": 75343 + }, + { + "code": "AZ", + "token": " AZ", + "token_id": 45108 + }, + { + "code": "BA", + "token": " BA", + "token_id": 29619 + }, + { + "code": "BB", + "token": " BB", + "token_id": 23814 + }, + { + "code": "BC", + "token": " BC", + "token_id": 19339 + }, + { + "code": "BD", + "token": " BD", + "token_id": 41669 + }, + { + "code": "BE", + "token": " BE", + "token_id": 11219 + }, + { + "code": "BF", + "token": " BF", + "token_id": 79550 + }, + { + "code": "BG", + "token": " BG", + "token_id": 68637 + }, + { + "code": "BH", + "token": " BH", + "token_id": 51317 + }, + { + "code": "BI", + "token": " BI", + "token_id": 30778 + }, + { + "code": "BJ", + "token": " BJ", + "token_id": 104543 + }, + { + "code": "BK", + "token": " BK", + "token_id": 97121 + }, + { + "code": "BL", + "token": " BL", + "token_id": 15358 + }, + { + "code": "BM", + "token": " BM", + "token_id": 46647 + }, + { + "code": "BN", + "token": " BN", + "token_id": 53260 + }, + { + "code": "BO", + "token": " BO", + "token_id": 13609 + }, + { + "code": "BP", + "token": " BP", + "token_id": 39122 + }, + { + "code": "BR", + "token": " BR", + "token_id": 17494 + }, + { + "code": "BS", + "token": " BS", + "token_id": 33083 + }, + { + "code": "BT", + "token": " BT", + "token_id": 35554 + }, + { + "code": "BU", + "token": " BU", + "token_id": 27129 + }, + { + "code": "BV", + "token": " BV", + "token_id": 74853 + }, + { + "code": "BW", + "token": " BW", + "token_id": 76557 + }, + { + "code": "BX", + "token": " BX", + "token_id": 204893 + }, + { + "code": "BY", + "token": " BY", + "token_id": 13222 + }, + { + "code": "CA", + "token": " CA", + "token_id": 8253 + }, + { + "code": "CB", + "token": " CB", + "token_id": 27331 + }, + { + "code": "CC", + "token": " CC", + "token_id": 10914 + }, + { + "code": "CD", + "token": " CD", + "token_id": 8080 + }, + { + "code": "CE", + "token": " CE", + "token_id": 20200 + }, + { + "code": "CF", + "token": " CF", + "token_id": 24328 + }, + { + "code": "CG", + "token": " CG", + "token_id": 19790 + }, + { + "code": "CH", + "token": " CH", + "token_id": 5831 + }, + { + "code": "CI", + "token": " CI", + "token_id": 29302 + }, + { + "code": "CJ", + "token": " CJ", + "token_id": 56056 + }, + { + "code": "CK", + "token": " CK", + "token_id": 67216 + }, + { + "code": "CL", + "token": " CL", + "token_id": 9837 + }, + { + "code": "CM", + "token": " CM", + "token_id": 20991 + }, + { + "code": "CN", + "token": " CN", + "token_id": 26111 + }, + { + "code": "CO", + "token": " CO", + "token_id": 2671 + }, + { + "code": "CP", + "token": " CP", + "token_id": 23752 + }, + { + "code": "CQ", + "token": " CQ", + "token_id": 77921 + }, + { + "code": "CR", + "token": " CR", + "token_id": 11862 + }, + { + "code": "CS", + "token": " CS", + "token_id": 14923 + }, + { + "code": "CT", + "token": " CT", + "token_id": 21044 + }, + { + "code": "CU", + "token": " CU", + "token_id": 52957 + }, + { + "code": "CV", + "token": " CV", + "token_id": 25572 + }, + { + "code": "CW", + "token": " CW", + "token_id": 62496 + }, + { + "code": "CX", + "token": " CX", + "token_id": 52074 + }, + { + "code": "CY", + "token": " CY", + "token_id": 32931 + }, + { + "code": "CZ", + "token": " CZ", + "token_id": 97354 + }, + { + "code": "DA", + "token": " DA", + "token_id": 21119 + }, + { + "code": "DB", + "token": " DB", + "token_id": 18706 + }, + { + "code": "DC", + "token": " DC", + "token_id": 11829 + }, + { + "code": "DD", + "token": " DD", + "token_id": 38214 + }, + { + "code": "DE", + "token": " DE", + "token_id": 4216 + }, + { + "code": "DF", + "token": " DF", + "token_id": 52894 + }, + { + "code": "DG", + "token": " DG", + "token_id": 61205 + }, + { + "code": "DH", + "token": " DH", + "token_id": 41191 + }, + { + "code": "DI", + "token": " DI", + "token_id": 12271 + }, + { + "code": "DJ", + "token": " DJ", + "token_id": 24089 + }, + { + "code": "DK", + "token": " DK", + "token_id": 74672 + }, + { + "code": "DL", + "token": " DL", + "token_id": 40312 + }, + { + "code": "DM", + "token": " DM", + "token_id": 22529 + }, + { + "code": "DN", + "token": " DN", + "token_id": 50667 + }, + { + "code": "DO", + "token": " DO", + "token_id": 13391 + }, + { + "code": "DP", + "token": " DP", + "token_id": 41385 + }, + { + "code": "DQ", + "token": " DQ", + "token_id": 167505 + }, + { + "code": "DR", + "token": " DR", + "token_id": 19588 + }, + { + "code": "DS", + "token": " DS", + "token_id": 29066 + }, + { + "code": "DT", + "token": " DT", + "token_id": 32528 + }, + { + "code": "DU", + "token": " DU", + "token_id": 41349 + }, + { + "code": "DV", + "token": " DV", + "token_id": 99570 + }, + { + "code": "DW", + "token": " DW", + "token_id": 72185 + }, + { + "code": "DX", + "token": " DX", + "token_id": 70235 + }, + { + "code": "DY", + "token": " DY", + "token_id": 97135 + }, + { + "code": "DZ", + "token": " DZ", + "token_id": 174960 + }, + { + "code": "EA", + "token": " EA", + "token_id": 48005 + }, + { + "code": "EB", + "token": " EB", + "token_id": 72197 + }, + { + "code": "EC", + "token": " EC", + "token_id": 19521 + }, + { + "code": "ED", + "token": " ED", + "token_id": 17026 + }, + { + "code": "EE", + "token": " EE", + "token_id": 27198 + }, + { + "code": "EF", + "token": " EF", + "token_id": 48664 + }, + { + "code": "EG", + "token": " EG", + "token_id": 69228 + }, + { + "code": "EH", + "token": " EH", + "token_id": 113243 + }, + { + "code": "EI", + "token": " EI", + "token_id": 133121 + }, + { + "code": "EJ", + "token": " EJ", + "token_id": 107740 + }, + { + "code": "EK", + "token": " EK", + "token_id": 110764 + }, + { + "code": "EL", + "token": " EL", + "token_id": 17517 + }, + { + "code": "EM", + "token": " EM", + "token_id": 22950 + }, + { + "code": "EN", + "token": " EN", + "token_id": 9163 + }, + { + "code": "EO", + "token": " EO", + "token_id": 128569 + }, + { + "code": "EP", + "token": " EP", + "token_id": 23590 + }, + { + "code": "EQ", + "token": " EQ", + "token_id": 81912 + }, + { + "code": "ER", + "token": " ER", + "token_id": 17360 + }, + { + "code": "ES", + "token": " ES", + "token_id": 21005 + }, + { + "code": "ET", + "token": " ET", + "token_id": 20157 + }, + { + "code": "EU", + "token": " EU", + "token_id": 11514 + }, + { + "code": "EV", + "token": " EV", + "token_id": 15620 + }, + { + "code": "EW", + "token": " EW", + "token_id": 70800 + }, + { + "code": "EX", + "token": " EX", + "token_id": 8495 + }, + { + "code": "EY", + "token": " EY", + "token_id": 99707 + }, + { + "code": "EZ", + "token": " EZ", + "token_id": 105286 + }, + { + "code": "FA", + "token": " FA", + "token_id": 16154 + }, + { + "code": "FB", + "token": " FB", + "token_id": 65834 + }, + { + "code": "FC", + "token": " FC", + "token_id": 21267 + }, + { + "code": "FD", + "token": " FD", + "token_id": 52815 + }, + { + "code": "FE", + "token": " FE", + "token_id": 23185 + }, + { + "code": "FF", + "token": " FF", + "token_id": 50237 + }, + { + "code": "FG", + "token": " FG", + "token_id": 58000 + }, + { + "code": "FH", + "token": " FH", + "token_id": 78021 + }, + { + "code": "FI", + "token": " FI", + "token_id": 16068 + }, + { + "code": "FJ", + "token": " FJ", + "token_id": 185665 + }, + { + "code": "FK", + "token": " FK", + "token_id": 53192 + }, + { + "code": "FL", + "token": " FL", + "token_id": 16586 + }, + { + "code": "FM", + "token": " FM", + "token_id": 32684 + }, + { + "code": "FN", + "token": " FN", + "token_id": 88122 + }, + { + "code": "FO", + "token": " FO", + "token_id": 25880 + }, + { + "code": "FP", + "token": " FP", + "token_id": 40530 + }, + { + "code": "FR", + "token": " FR", + "token_id": 18309 + }, + { + "code": "FS", + "token": " FS", + "token_id": 45925 + }, + { + "code": "FT", + "token": " FT", + "token_id": 37483 + }, + { + "code": "FU", + "token": " FU", + "token_id": 57581 + }, + { + "code": "FV", + "token": " FV", + "token_id": 105235 + }, + { + "code": "FW", + "token": " FW", + "token_id": 91579 + }, + { + "code": "FX", + "token": " FX", + "token_id": 37113 + }, + { + "code": "FY", + "token": " FY", + "token_id": 29692 + }, + { + "code": "GA", + "token": " GA", + "token_id": 23639 + }, + { + "code": "GB", + "token": " GB", + "token_id": 26968 + }, + { + "code": "GC", + "token": " GC", + "token_id": 29013 + }, + { + "code": "GD", + "token": " GD", + "token_id": 46169 + }, + { + "code": "GE", + "token": " GE", + "token_id": 27133 + }, + { + "code": "GF", + "token": " GF", + "token_id": 74659 + }, + { + "code": "GG", + "token": " GG", + "token_id": 83079 + }, + { + "code": "GH", + "token": " GH", + "token_id": 48412 + }, + { + "code": "GI", + "token": " GI", + "token_id": 46577 + }, + { + "code": "GJ", + "token": " GJ", + "token_id": 182529 + }, + { + "code": "GK", + "token": " GK", + "token_id": 126979 + }, + { + "code": "GL", + "token": " GL", + "token_id": 14349 + }, + { + "code": "GM", + "token": " GM", + "token_id": 27506 + }, + { + "code": "GN", + "token": " GN", + "token_id": 85411 + }, + { + "code": "GO", + "token": " GO", + "token_id": 16861 + }, + { + "code": "GP", + "token": " GP", + "token_id": 27379 + }, + { + "code": "GQ", + "token": " GQ", + "token_id": 193930 + }, + { + "code": "GR", + "token": " GR", + "token_id": 17779 + }, + { + "code": "GS", + "token": " GS", + "token_id": 43080 + }, + { + "code": "GT", + "token": " GT", + "token_id": 24179 + }, + { + "code": "GU", + "token": " GU", + "token_id": 37518 + }, + { + "code": "GV", + "token": " GV", + "token_id": 125344 + }, + { + "code": "GW", + "token": " GW", + "token_id": 51667 + }, + { + "code": "GX", + "token": " GX", + "token_id": 154378 + }, + { + "code": "GY", + "token": " GY", + "token_id": 121205 + }, + { + "code": "HA", + "token": " HA", + "token_id": 39487 + }, + { + "code": "HB", + "token": " HB", + "token_id": 53261 + }, + { + "code": "HC", + "token": " HC", + "token_id": 50535 + }, + { + "code": "HD", + "token": " HD", + "token_id": 11876 + }, + { + "code": "HE", + "token": " HE", + "token_id": 13962 + }, + { + "code": "HF", + "token": " HF", + "token_id": 57330 + }, + { + "code": "HG", + "token": " HG", + "token_id": 100326 + }, + { + "code": "HH", + "token": " HH", + "token_id": 65091 + }, + { + "code": "HI", + "token": " HI", + "token_id": 53920 + }, + { + "code": "HJ", + "token": " HJ", + "token_id": 183707 + }, + { + "code": "HK", + "token": " HK", + "token_id": 45353 + }, + { + "code": "HL", + "token": " HL", + "token_id": 61611 + }, + { + "code": "HM", + "token": " HM", + "token_id": 59089 + }, + { + "code": "HN", + "token": " HN", + "token_id": 150602 + }, + { + "code": "HO", + "token": " HO", + "token_id": 30698 + }, + { + "code": "HP", + "token": " HP", + "token_id": 20761 + }, + { + "code": "HQ", + "token": " HQ", + "token_id": 71847 + }, + { + "code": "HR", + "token": " HR", + "token_id": 18341 + }, + { + "code": "HS", + "token": " HS", + "token_id": 32572 + }, + { + "code": "HT", + "token": " HT", + "token_id": 9794 + }, + { + "code": "HU", + "token": " HU", + "token_id": 50835 + }, + { + "code": "HV", + "token": " HV", + "token_id": 52446 + }, + { + "code": "HW", + "token": " HW", + "token_id": 74750 + }, + { + "code": "HX", + "token": " HX", + "token_id": 207837 + }, + { + "code": "HY", + "token": " HY", + "token_id": 41887 + }, + { + "code": "IA", + "token": " IA", + "token_id": 51279 + }, + { + "code": "IB", + "token": " IB", + "token_id": 42143 + }, + { + "code": "IC", + "token": " IC", + "token_id": 19403 + }, + { + "code": "ID", + "token": " ID", + "token_id": 5853 + }, + { + "code": "IE", + "token": " IE", + "token_id": 47780 + }, + { + "code": "IF", + "token": " IF", + "token_id": 14634 + }, + { + "code": "IG", + "token": " IG", + "token_id": 44772 + }, + { + "code": "IH", + "token": " IH", + "token_id": 69657 + }, + { + "code": "II", + "token": " II", + "token_id": 3664 + }, + { + "code": "IJ", + "token": " IJ", + "token_id": 115925 + }, + { + "code": "IK", + "token": " IK", + "token_id": 121834 + }, + { + "code": "IL", + "token": " IL", + "token_id": 15237 + }, + { + "code": "IM", + "token": " IM", + "token_id": 12607 + }, + { + "code": "IN", + "token": " IN", + "token_id": 3247 + }, + { + "code": "IO", + "token": " IO", + "token_id": 13421 + }, + { + "code": "IP", + "token": " IP", + "token_id": 10710 + }, + { + "code": "IQ", + "token": " IQ", + "token_id": 47304 + }, + { + "code": "IR", + "token": " IR", + "token_id": 16459 + }, + { + "code": "IS", + "token": " IS", + "token_id": 5246 + }, + { + "code": "IT", + "token": " IT", + "token_id": 9053 + }, + { + "code": "IU", + "token": " IU", + "token_id": 57150 + }, + { + "code": "IV", + "token": " IV", + "token_id": 10935 + }, + { + "code": "IW", + "token": " IW", + "token_id": 150007 + }, + { + "code": "IX", + "token": " IX", + "token_id": 34838 + }, + { + "code": "IZ", + "token": " IZ", + "token_id": 178335 + }, + { + "code": "JA", + "token": " JA", + "token_id": 48032 + }, + { + "code": "JB", + "token": " JB", + "token_id": 88172 + }, + { + "code": "JC", + "token": " JC", + "token_id": 71592 + }, + { + "code": "JD", + "token": " JD", + "token_id": 57469 + }, + { + "code": "JE", + "token": " JE", + "token_id": 47795 + }, + { + "code": "JF", + "token": " JF", + "token_id": 140634 + }, + { + "code": "JG", + "token": " JG", + "token_id": 164564 + }, + { + "code": "JH", + "token": " JH", + "token_id": 113593 + }, + { + "code": "JI", + "token": " JI", + "token_id": 190023 + }, + { + "code": "JJ", + "token": " JJ", + "token_id": 63669 + }, + { + "code": "JK", + "token": " JK", + "token_id": 112862 + }, + { + "code": "JL", + "token": " JL", + "token_id": 91474 + }, + { + "code": "JM", + "token": " JM", + "token_id": 53191 + }, + { + "code": "JN", + "token": " JN", + "token_id": 236659 + }, + { + "code": "JO", + "token": " JO", + "token_id": 22857 + }, + { + "code": "JP", + "token": " JP", + "token_id": 50603 + }, + { + "code": "JR", + "token": " JR", + "token_id": 56553 + }, + { + "code": "JS", + "token": " JS", + "token_id": 28551 + }, + { + "code": "JT", + "token": " JT", + "token_id": 94258 + }, + { + "code": "JU", + "token": " JU", + "token_id": 85240 + }, + { + "code": "JV", + "token": " JV", + "token_id": 119273 + }, + { + "code": "JW", + "token": " JW", + "token_id": 93191 + }, + { + "code": "JY", + "token": " JY", + "token_id": 234264 + }, + { + "code": "KA", + "token": " KA", + "token_id": 77878 + }, + { + "code": "KB", + "token": " KB", + "token_id": 50633 + }, + { + "code": "KC", + "token": " KC", + "token_id": 39684 + }, + { + "code": "KD", + "token": " KD", + "token_id": 92517 + }, + { + "code": "KE", + "token": " KE", + "token_id": 62574 + }, + { + "code": "KF", + "token": " KF", + "token_id": 150049 + }, + { + "code": "KG", + "token": " KG", + "token_id": 77173 + }, + { + "code": "KH", + "token": " KH", + "token_id": 58894 + }, + { + "code": "KI", + "token": " KI", + "token_id": 76148 + }, + { + "code": "KJ", + "token": " KJ", + "token_id": 161846 + }, + { + "code": "KK", + "token": " KK", + "token_id": 79838 + }, + { + "code": "KL", + "token": " KL", + "token_id": 58228 + }, + { + "code": "KM", + "token": " KM", + "token_id": 48878 + }, + { + "code": "KN", + "token": " KN", + "token_id": 67624 + }, + { + "code": "KO", + "token": " KO", + "token_id": 44479 + }, + { + "code": "KP", + "token": " KP", + "token_id": 74821 + }, + { + "code": "KQ", + "token": " KQ", + "token_id": 121555 + }, + { + "code": "KR", + "token": " KR", + "token_id": 59147 + }, + { + "code": "KS", + "token": " KS", + "token_id": 56122 + }, + { + "code": "KT", + "token": " KT", + "token_id": 89139 + }, + { + "code": "KU", + "token": " KU", + "token_id": 91408 + }, + { + "code": "KV", + "token": " KV", + "token_id": 132182 + }, + { + "code": "KW", + "token": " KW", + "token_id": 109600 + }, + { + "code": "KX", + "token": " KX", + "token_id": 206607 + }, + { + "code": "KY", + "token": " KY", + "token_id": 63346 + }, + { + "code": "KZ", + "token": " KZ", + "token_id": 189327 + }, + { + "code": "LA", + "token": " LA", + "token_id": 11404 + }, + { + "code": "LB", + "token": " LB", + "token_id": 67495 + }, + { + "code": "LC", + "token": " LC", + "token_id": 42110 + }, + { + "code": "LD", + "token": " LD", + "token_id": 36416 + }, + { + "code": "LE", + "token": " LE", + "token_id": 11618 + }, + { + "code": "LF", + "token": " LF", + "token_id": 84769 + }, + { + "code": "LG", + "token": " LG", + "token_id": 30310 + }, + { + "code": "LH", + "token": " LH", + "token_id": 86686 + }, + { + "code": "LI", + "token": " LI", + "token_id": 22979 + }, + { + "code": "LJ", + "token": " LJ", + "token_id": 148316 + }, + { + "code": "LK", + "token": " LK", + "token_id": 144505 + }, + { + "code": "LL", + "token": " LL", + "token_id": 32050 + }, + { + "code": "LM", + "token": " LM", + "token_id": 63457 + }, + { + "code": "LN", + "token": " LN", + "token_id": 112970 + }, + { + "code": "LO", + "token": " LO", + "token_id": 17962 + }, + { + "code": "LP", + "token": " LP", + "token_id": 35849 + }, + { + "code": "LQ", + "token": " LQ", + "token_id": 138139 + }, + { + "code": "LR", + "token": " LR", + "token_id": 45273 + }, + { + "code": "LS", + "token": " LS", + "token_id": 46450 + }, + { + "code": "LT", + "token": " LT", + "token_id": 56635 + }, + { + "code": "LU", + "token": " LU", + "token_id": 79740 + }, + { + "code": "LV", + "token": " LV", + "token_id": 63147 + }, + { + "code": "LW", + "token": " LW", + "token_id": 124651 + }, + { + "code": "LX", + "token": " LX", + "token_id": 82888 + }, + { + "code": "LY", + "token": " LY", + "token_id": 83331 + }, + { + "code": "LZ", + "token": " LZ", + "token_id": 177461 + }, + { + "code": "MA", + "token": " MA", + "token_id": 7883 + }, + { + "code": "MB", + "token": " MB", + "token_id": 18879 + }, + { + "code": "MC", + "token": " MC", + "token_id": 18768 + }, + { + "code": "MD", + "token": " MD", + "token_id": 16781 + }, + { + "code": "ME", + "token": " ME", + "token_id": 15932 + }, + { + "code": "MF", + "token": " MF", + "token_id": 71026 + }, + { + "code": "MG", + "token": " MG", + "token_id": 34815 + }, + { + "code": "MH", + "token": " MH", + "token_id": 60054 + }, + { + "code": "MI", + "token": " MI", + "token_id": 29067 + }, + { + "code": "MJ", + "token": " MJ", + "token_id": 67119 + }, + { + "code": "MK", + "token": " MK", + "token_id": 42372 + }, + { + "code": "ML", + "token": " ML", + "token_id": 20885 + }, + { + "code": "MM", + "token": " MM", + "token_id": 25228 + }, + { + "code": "MN", + "token": " MN", + "token_id": 36261 + }, + { + "code": "MO", + "token": " MO", + "token_id": 17231 + }, + { + "code": "MP", + "token": " MP", + "token_id": 13945 + }, + { + "code": "MQ", + "token": " MQ", + "token_id": 69825 + }, + { + "code": "MR", + "token": " MR", + "token_id": 24291 + }, + { + "code": "MS", + "token": " MS", + "token_id": 12501 + }, + { + "code": "MT", + "token": " MT", + "token_id": 27511 + }, + { + "code": "MU", + "token": " MU", + "token_id": 46490 + }, + { + "code": "MV", + "token": " MV", + "token_id": 53773 + }, + { + "code": "MW", + "token": " MW", + "token_id": 35920 + }, + { + "code": "MX", + "token": " MX", + "token_id": 45796 + }, + { + "code": "MY", + "token": " MY", + "token_id": 27714 + }, + { + "code": "MZ", + "token": " MZ", + "token_id": 154716 + }, + { + "code": "NA", + "token": " NA", + "token_id": 19674 + }, + { + "code": "NB", + "token": " NB", + "token_id": 55740 + }, + { + "code": "NC", + "token": " NC", + "token_id": 23316 + }, + { + "code": "ND", + "token": " ND", + "token_id": 36437 + }, + { + "code": "NE", + "token": " NE", + "token_id": 8874 + }, + { + "code": "NF", + "token": " NF", + "token_id": 42167 + }, + { + "code": "NG", + "token": " NG", + "token_id": 65799 + }, + { + "code": "NH", + "token": " NH", + "token_id": 22869 + }, + { + "code": "NI", + "token": " NI", + "token_id": 32426 + }, + { + "code": "NJ", + "token": " NJ", + "token_id": 31566 + }, + { + "code": "NK", + "token": " NK", + "token_id": 46314 + }, + { + "code": "NL", + "token": " NL", + "token_id": 51900 + }, + { + "code": "NM", + "token": " NM", + "token_id": 44285 + }, + { + "code": "NN", + "token": " NN", + "token_id": 60800 + }, + { + "code": "NO", + "token": " NO", + "token_id": 9424 + }, + { + "code": "NP", + "token": " NP", + "token_id": 40098 + }, + { + "code": "NR", + "token": " NR", + "token_id": 49428 + }, + { + "code": "NS", + "token": " NS", + "token_id": 10707 + }, + { + "code": "NT", + "token": " NT", + "token_id": 24505 + }, + { + "code": "NU", + "token": " NU", + "token_id": 93241 + }, + { + "code": "NV", + "token": " NV", + "token_id": 57484 + }, + { + "code": "NW", + "token": " NW", + "token_id": 45753 + }, + { + "code": "NX", + "token": " NX", + "token_id": 137707 + }, + { + "code": "NY", + "token": " NY", + "token_id": 13476 + }, + { + "code": "NZ", + "token": " NZ", + "token_id": 55990 + }, + { + "code": "OA", + "token": " OA", + "token_id": 76762 + }, + { + "code": "OB", + "token": " OB", + "token_id": 32610 + }, + { + "code": "OC", + "token": " OC", + "token_id": 58646 + }, + { + "code": "OD", + "token": " OD", + "token_id": 43680 + }, + { + "code": "OE", + "token": " OE", + "token_id": 115092 + }, + { + "code": "OF", + "token": " OF", + "token_id": 3639 + }, + { + "code": "OG", + "token": " OG", + "token_id": 75702 + }, + { + "code": "OH", + "token": " OH", + "token_id": 28963 + }, + { + "code": "OI", + "token": " OI", + "token_id": 116687 + }, + { + "code": "OJ", + "token": " OJ", + "token_id": 212412 + }, + { + "code": "OK", + "token": " OK", + "token_id": 16119 + }, + { + "code": "OL", + "token": " OL", + "token_id": 53418 + }, + { + "code": "OM", + "token": " OM", + "token_id": 69261 + }, + { + "code": "ON", + "token": " ON", + "token_id": 8203 + }, + { + "code": "OO", + "token": " OO", + "token_id": 129623 + }, + { + "code": "OP", + "token": " OP", + "token_id": 19577 + }, + { + "code": "OR", + "token": " OR", + "token_id": 8005 + }, + { + "code": "OS", + "token": " OS", + "token_id": 16080 + }, + { + "code": "OT", + "token": " OT", + "token_id": 53341 + }, + { + "code": "OU", + "token": " OU", + "token_id": 103247 + }, + { + "code": "OV", + "token": " OV", + "token_id": 128433 + }, + { + "code": "OW", + "token": " OW", + "token_id": 51986 + }, + { + "code": "OX", + "token": " OX", + "token_id": 72876 + }, + { + "code": "OZ", + "token": " OZ", + "token_id": 210259 + }, + { + "code": "PA", + "token": " PA", + "token_id": 13187 + }, + { + "code": "PB", + "token": " PB", + "token_id": 40920 + }, + { + "code": "PC", + "token": " PC", + "token_id": 7294 + }, + { + "code": "PD", + "token": " PD", + "token_id": 24313 + }, + { + "code": "PE", + "token": " PE", + "token_id": 20375 + }, + { + "code": "PF", + "token": " PF", + "token_id": 53839 + }, + { + "code": "PG", + "token": " PG", + "token_id": 43447 + }, + { + "code": "PH", + "token": " PH", + "token_id": 12462 + }, + { + "code": "PI", + "token": " PI", + "token_id": 6108 + }, + { + "code": "PJ", + "token": " PJ", + "token_id": 95080 + }, + { + "code": "PK", + "token": " PK", + "token_id": 32743 + }, + { + "code": "PL", + "token": " PL", + "token_id": 9123 + }, + { + "code": "PM", + "token": " PM", + "token_id": 8711 + }, + { + "code": "PN", + "token": " PN", + "token_id": 55561 + }, + { + "code": "PO", + "token": " PO", + "token_id": 13174 + }, + { + "code": "PP", + "token": " PP", + "token_id": 27115 + }, + { + "code": "PQ", + "token": " PQ", + "token_id": 106694 + }, + { + "code": "PR", + "token": " PR", + "token_id": 9206 + }, + { + "code": "PS", + "token": " PS", + "token_id": 11156 + }, + { + "code": "PT", + "token": " PT", + "token_id": 26912 + }, + { + "code": "PU", + "token": " PU", + "token_id": 25358 + }, + { + "code": "PV", + "token": " PV", + "token_id": 32985 + }, + { + "code": "PW", + "token": " PW", + "token_id": 50933 + }, + { + "code": "PX", + "token": " PX", + "token_id": 127287 + }, + { + "code": "PY", + "token": " PY", + "token_id": 70643 + }, + { + "code": "PZ", + "token": " PZ", + "token_id": 151048 + }, + { + "code": "QA", + "token": " QA", + "token_id": 73088 + }, + { + "code": "QB", + "token": " QB", + "token_id": 81378 + }, + { + "code": "QC", + "token": " QC", + "token_id": 66419 + }, + { + "code": "QD", + "token": " QD", + "token_id": 131260 + }, + { + "code": "QE", + "token": " QE", + "token_id": 179889 + }, + { + "code": "QH", + "token": " QH", + "token_id": 228576 + }, + { + "code": "QI", + "token": " QI", + "token_id": 215904 + }, + { + "code": "QM", + "token": " QM", + "token_id": 236116 + }, + { + "code": "QP", + "token": " QP", + "token_id": 176850 + }, + { + "code": "QQ", + "token": " QQ", + "token_id": 179904 + }, + { + "code": "QR", + "token": " QR", + "token_id": 44426 + }, + { + "code": "QS", + "token": " QS", + "token_id": 111800 + }, + { + "code": "QT", + "token": " QT", + "token_id": 132413 + }, + { + "code": "QU", + "token": " QU", + "token_id": 29809 + }, + { + "code": "QV", + "token": " QV", + "token_id": 190489 + }, + { + "code": "RA", + "token": " RA", + "token_id": 21242 + }, + { + "code": "RB", + "token": " RB", + "token_id": 44282 + }, + { + "code": "RC", + "token": " RC", + "token_id": 34023 + }, + { + "code": "RD", + "token": " RD", + "token_id": 53493 + }, + { + "code": "RE", + "token": " RE", + "token_id": 4769 + }, + { + "code": "RF", + "token": " RF", + "token_id": 28049 + }, + { + "code": "RG", + "token": " RG", + "token_id": 71373 + }, + { + "code": "RH", + "token": " RH", + "token_id": 55360 + }, + { + "code": "RI", + "token": " RI", + "token_id": 51720 + }, + { + "code": "RJ", + "token": " RJ", + "token_id": 69430 + }, + { + "code": "RK", + "token": " RK", + "token_id": 93622 + }, + { + "code": "RL", + "token": " RL", + "token_id": 61554 + }, + { + "code": "RM", + "token": " RM", + "token_id": 25543 + }, + { + "code": "RN", + "token": " RN", + "token_id": 47457 + }, + { + "code": "RO", + "token": " RO", + "token_id": 17006 + }, + { + "code": "RP", + "token": " RP", + "token_id": 48269 + }, + { + "code": "RQ", + "token": " RQ", + "token_id": 215785 + }, + { + "code": "RR", + "token": " RR", + "token_id": 50013 + }, + { + "code": "RS", + "token": " RS", + "token_id": 20949 + }, + { + "code": "RT", + "token": " RT", + "token_id": 23094 + }, + { + "code": "RU", + "token": " RU", + "token_id": 108257 + }, + { + "code": "RV", + "token": " RV", + "token_id": 36662 + }, + { + "code": "RW", + "token": " RW", + "token_id": 75067 + }, + { + "code": "RX", + "token": " RX", + "token_id": 39625 + }, + { + "code": "RY", + "token": " RY", + "token_id": 147326 + }, + { + "code": "SA", + "token": " SA", + "token_id": 15850 + }, + { + "code": "SB", + "token": " SB", + "token_id": 37122 + }, + { + "code": "SC", + "token": " SC", + "token_id": 15009 + }, + { + "code": "SD", + "token": " SD", + "token_id": 15336 + }, + { + "code": "SE", + "token": " SE", + "token_id": 7952 + }, + { + "code": "SF", + "token": " SF", + "token_id": 33152 + }, + { + "code": "SG", + "token": " SG", + "token_id": 46538 + }, + { + "code": "SH", + "token": " SH", + "token_id": 9586 + }, + { + "code": "SI", + "token": " SI", + "token_id": 24732 + }, + { + "code": "SJ", + "token": " SJ", + "token_id": 75248 + }, + { + "code": "SK", + "token": " SK", + "token_id": 22246 + }, + { + "code": "SL", + "token": " SL", + "token_id": 17700 + }, + { + "code": "SM", + "token": " SM", + "token_id": 15233 + }, + { + "code": "SN", + "token": " SN", + "token_id": 19818 + }, + { + "code": "SO", + "token": " SO", + "token_id": 12583 + }, + { + "code": "SP", + "token": " SP", + "token_id": 11446 + }, + { + "code": "SQ", + "token": " SQ", + "token_id": 55658 + }, + { + "code": "SR", + "token": " SR", + "token_id": 28155 + }, + { + "code": "SS", + "token": " SS", + "token_id": 18236 + }, + { + "code": "ST", + "token": " ST", + "token_id": 4439 + }, + { + "code": "SU", + "token": " SU", + "token_id": 15625 + }, + { + "code": "SV", + "token": " SV", + "token_id": 34854 + }, + { + "code": "SW", + "token": " SW", + "token_id": 19175 + }, + { + "code": "SX", + "token": " SX", + "token_id": 87863 + }, + { + "code": "SY", + "token": " SY", + "token_id": 20288 + }, + { + "code": "SZ", + "token": " SZ", + "token_id": 112930 + }, + { + "code": "TA", + "token": " TA", + "token_id": 39337 + }, + { + "code": "TB", + "token": " TB", + "token_id": 34949 + }, + { + "code": "TC", + "token": " TC", + "token_id": 41393 + }, + { + "code": "TD", + "token": " TD", + "token_id": 38952 + }, + { + "code": "TE", + "token": " TE", + "token_id": 17423 + }, + { + "code": "TF", + "token": " TF", + "token_id": 56165 + }, + { + "code": "TG", + "token": " TG", + "token_id": 72268 + }, + { + "code": "TH", + "token": " TH", + "token_id": 6925 + }, + { + "code": "TI", + "token": " TI", + "token_id": 43988 + }, + { + "code": "TJ", + "token": " TJ", + "token_id": 88129 + }, + { + "code": "TK", + "token": " TK", + "token_id": 63814 + }, + { + "code": "TL", + "token": " TL", + "token_id": 34147 + }, + { + "code": "TM", + "token": " TM", + "token_id": 38610 + }, + { + "code": "TN", + "token": " TN", + "token_id": 26493 + }, + { + "code": "TO", + "token": " TO", + "token_id": 6760 + }, + { + "code": "TP", + "token": " TP", + "token_id": 34352 + }, + { + "code": "TR", + "token": " TR", + "token_id": 8642 + }, + { + "code": "TS", + "token": " TS", + "token_id": 34654 + }, + { + "code": "TT", + "token": " TT", + "token_id": 40774 + }, + { + "code": "TU", + "token": " TU", + "token_id": 52263 + }, + { + "code": "TV", + "token": " TV", + "token_id": 6115 + }, + { + "code": "TW", + "token": " TW", + "token_id": 48704 + }, + { + "code": "TX", + "token": " TX", + "token_id": 27608 + }, + { + "code": "TY", + "token": " TY", + "token_id": 60057 + }, + { + "code": "TZ", + "token": " TZ", + "token_id": 216288 + }, + { + "code": "UA", + "token": " UA", + "token_id": 84168 + }, + { + "code": "UB", + "token": " UB", + "token_id": 90890 + }, + { + "code": "UC", + "token": " UC", + "token_id": 50606 + }, + { + "code": "UD", + "token": " UD", + "token_id": 74765 + }, + { + "code": "UE", + "token": " UE", + "token_id": 32646 + }, + { + "code": "UF", + "token": " UF", + "token_id": 93592 + }, + { + "code": "UG", + "token": " UG", + "token_id": 151515 + }, + { + "code": "UH", + "token": " UH", + "token_id": 134553 + }, + { + "code": "UI", + "token": " UI", + "token_id": 9711 + }, + { + "code": "UK", + "token": " UK", + "token_id": 6322 + }, + { + "code": "UL", + "token": " UL", + "token_id": 56997 + }, + { + "code": "UM", + "token": " UM", + "token_id": 67201 + }, + { + "code": "UN", + "token": " UN", + "token_id": 7866 + }, + { + "code": "UP", + "token": " UP", + "token_id": 24252 + }, + { + "code": "UR", + "token": " UR", + "token_id": 46412 + }, + { + "code": "US", + "token": " US", + "token_id": 2590 + }, + { + "code": "UT", + "token": " UT", + "token_id": 29297 + }, + { + "code": "UU", + "token": " UU", + "token_id": 109410 + }, + { + "code": "UV", + "token": " UV", + "token_id": 19296 + }, + { + "code": "UW", + "token": " UW", + "token_id": 57623 + }, + { + "code": "UX", + "token": " UX", + "token_id": 69873 + }, + { + "code": "VA", + "token": " VA", + "token_id": 30578 + }, + { + "code": "VB", + "token": " VB", + "token_id": 101537 + }, + { + "code": "VC", + "token": " VC", + "token_id": 51191 + }, + { + "code": "VD", + "token": " VD", + "token_id": 158142 + }, + { + "code": "VE", + "token": " VE", + "token_id": 33697 + }, + { + "code": "VF", + "token": " VF", + "token_id": 138963 + }, + { + "code": "VG", + "token": " VG", + "token_id": 95218 + }, + { + "code": "VH", + "token": " VH", + "token_id": 164123 + }, + { + "code": "VI", + "token": " VI", + "token_id": 18603 + }, + { + "code": "VK", + "token": " VK", + "token_id": 54609 + }, + { + "code": "VL", + "token": " VL", + "token_id": 86868 + }, + { + "code": "VM", + "token": " VM", + "token_id": 37231 + }, + { + "code": "VN", + "token": " VN", + "token_id": 90326 + }, + { + "code": "VO", + "token": " VO", + "token_id": 59054 + }, + { + "code": "VP", + "token": " VP", + "token_id": 42752 + }, + { + "code": "VR", + "token": " VR", + "token_id": 28644 + }, + { + "code": "VS", + "token": " VS", + "token_id": 38097 + }, + { + "code": "VT", + "token": " VT", + "token_id": 55503 + }, + { + "code": "VU", + "token": " VU", + "token_id": 171279 + }, + { + "code": "VV", + "token": " VV", + "token_id": 151030 + }, + { + "code": "VW", + "token": " VW", + "token_id": 70365 + }, + { + "code": "VX", + "token": " VX", + "token_id": 113564 + }, + { + "code": "VY", + "token": " VY", + "token_id": 220191 + }, + { + "code": "WA", + "token": " WA", + "token_id": 31582 + }, + { + "code": "WB", + "token": " WB", + "token_id": 72773 + }, + { + "code": "WC", + "token": " WC", + "token_id": 48380 + }, + { + "code": "WD", + "token": " WD", + "token_id": 100311 + }, + { + "code": "WE", + "token": " WE", + "token_id": 19716 + }, + { + "code": "WF", + "token": " WF", + "token_id": 113026 + }, + { + "code": "WG", + "token": " WG", + "token_id": 121848 + }, + { + "code": "WH", + "token": " WH", + "token_id": 16169 + }, + { + "code": "WI", + "token": " WI", + "token_id": 55927 + }, + { + "code": "WJ", + "token": " WJ", + "token_id": 193516 + }, + { + "code": "WK", + "token": " WK", + "token_id": 125879 + }, + { + "code": "WL", + "token": " WL", + "token_id": 80760 + }, + { + "code": "WM", + "token": " WM", + "token_id": 42668 + }, + { + "code": "WO", + "token": " WO", + "token_id": 42626 + } + ] +} diff --git a/server/src/system_one_lite/engine.py b/server/src/system_one_lite/engine.py index 8646d93..6c1abce 100644 --- a/server/src/system_one_lite/engine.py +++ b/server/src/system_one_lite/engine.py @@ -23,13 +23,19 @@ QWEN3_1_7B_MODEL = "mlx-community/Qwen3-1.7B-4bit" QWEN3_4B_MODEL = "mlx-community/Qwen3-4B-Instruct-2507-4bit" +DIFFUSION_GEMMA_MODEL = "mlx-community/diffusiongemma-26B-A4B-it-4bit" MODEL_REVISIONS = { QWEN3_1_7B_MODEL: "3b1b1768f8f8cf8351c712464f906e86c2b8269e", QWEN3_4B_MODEL: "50d427756c6b1b2fe0c0a10f67fbda1fc8e82c1b", + DIFFUSION_GEMMA_MODEL: "a7a81407613811e8ba63af92ac0d852b809e191f", } DEFAULT_MODEL = QWEN3_1_7B_MODEL LARGER_MODEL = QWEN3_4B_MODEL -MODEL_PROFILES = {"default": DEFAULT_MODEL, "larger": LARGER_MODEL} +MODEL_PROFILES = { + "default": DEFAULT_MODEL, + "larger": LARGER_MODEL, + "diffusion": DIFFUSION_GEMMA_MODEL, +} CODES_FILE = files("system_one_lite.data").joinpath("qwen3_1_7b_4bit_answer_codes.json") TEMPERATURE = 0.7 MAX_PROMPT_TOKENS = 32_768 @@ -47,6 +53,7 @@ ANSWER_CODE_CANDIDATES = tuple(string.ascii_uppercase) + tuple( a + b for a in string.ascii_uppercase for b in string.ascii_uppercase ) +MAX_ANSWER_CODES = 578 _REGISTERED_IDS = object() @@ -129,7 +136,25 @@ def answer_code_entries(tokenizer): ) if len(cuts) != 1: raise ValueError(f"codes land at different cut positions: {sorted(cuts)}") - return entries + return entries[:MAX_ANSWER_CODES] + + +def diffusion_answer_code_entries(tokenizer): + """Build codes that occupy one seeded-canvas token after a list marker.""" + entries = [] + for code in ANSWER_CODE_CANDIDATES: + token_ids = tokenizer.encode(" " + code, add_special_tokens=False) + if len(token_ids) != 1: + continue + token_id = token_ids[0] + entries.append( + { + "code": code, + "token": tokenizer.decode([token_id]), + "token_id": token_id, + } + ) + return entries[:MAX_ANSWER_CODES] def _registry_files(): @@ -143,8 +168,8 @@ def _registry_files(): ) -def load_code_registry(model_id, tokenizer, model_path, model_revision): - """Find and fully validate one model's answer-code registry.""" +def packaged_code_registry(model_id): + """Load one packaged registry without loading the model or tokenizer.""" matches = [] for path in _registry_files(): data = json.loads(path.read_text()) @@ -157,6 +182,19 @@ def load_code_registry(model_id, tokenizer, model_path, model_revision): ) path, data = matches[0] + entries = data.get("codes") or [] + codes = tuple(entry["code"] for entry in entries) + token_ids = tuple(entry["token_id"] for entry in entries) + if not codes or codes[0] != "A": + raise ValueError(f"{path.name} must start with the A answer code") + if len(codes) != len(set(codes)) or len(token_ids) != len(set(token_ids)): + raise ValueError(f"{path.name} contains duplicate codes or token ids") + return path, data, codes, token_ids + + +def load_code_registry(model_id, tokenizer, model_path, model_revision): + """Find and fully validate one model's answer-code registry.""" + path, data, codes, token_ids = packaged_code_registry(model_id) expected_hash = tokenizer_sha256(model_path) if data.get("model_revision") != model_revision: raise ValueError( @@ -166,16 +204,14 @@ def load_code_registry(model_id, tokenizer, model_path, model_revision): if data.get("tokenizer_sha256") != expected_hash: raise ValueError(f"{path.name} does not match the loaded tokenizer; regenerate it") - expected = answer_code_entries(tokenizer) + if model_id == DIFFUSION_GEMMA_MODEL: + expected = diffusion_answer_code_entries(tokenizer) + else: + expected = answer_code_entries(tokenizer) if data.get("codes") != expected: raise ValueError(f"{path.name} is incomplete or invalid for {model_id}; regenerate it") - codes = tuple(entry["code"] for entry in expected) - token_ids = tuple(entry["token_id"] for entry in expected) - if not codes or codes[0] != "A": - raise ValueError(f"{path.name} must start with the A answer code") - if len(codes) != len(set(codes)) or len(token_ids) != len(set(token_ids)): - raise ValueError(f"{path.name} contains duplicate codes or token ids") - validate_code_contexts(tokenizer, codes, token_ids) + if model_id != DIFFUSION_GEMMA_MODEL: + validate_code_contexts(tokenizer, codes, token_ids) return path, codes, token_ids @@ -244,6 +280,10 @@ def __init__(self, model_id=None): # warm up: compile Metal kernels before the first real request self.evaluate("warm up", [("pick one", ["yes", "no"])]) + def output_tokens_for(self, questions): + """Return the generated-token count for one prepared request.""" + return 0 + def evaluate(self, state, questions, template=None): """Run one independent full prompt per question. diff --git a/server/src/system_one_lite/mlx_vlm_diffusion.py b/server/src/system_one_lite/mlx_vlm_diffusion.py new file mode 100644 index 0000000..8bfe4cc --- /dev/null +++ b/server/src/system_one_lite/mlx_vlm_diffusion.py @@ -0,0 +1,298 @@ +"""System One backend for seeded DiffusionGemma reads through MLX-VLM.""" + +import json +import math +import os +import random +import time +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from mlx_lm.utils import load_tokenizer + +from .engine import ( + DIFFUSION_GEMMA_MODEL, + MAX_PROMPT_TOKENS, + MAX_TOTAL_INPUT_TOKENS, + TEMPERATURE, + RequestContractError, + TooManyOptions, + diffusion_answer_code_entries, + packaged_code_registry, + resolve_model_id, + resolve_model_snapshot, + tokenizer_sha256, +) + +DEFAULT_BASE_URL = "http://127.0.0.1:8080" +DEFAULT_TIMEOUT = 120.0 +CANVAS_PREFIX = "<|channel>thought\n" +CANVAS_SUFFIX = "\n" +DEFAULT_DIFFUSION_SEED = 42 +FAST_DIFFUSION_ENCODER_LAYERS = 6 + + +class MlxVlmError(RuntimeError): + """The MLX-VLM server could not complete a structured diffusion read.""" + + +def restricted_softmax(logprobs, temperature=TEMPERATURE): + """Renormalize vocabulary log probabilities over allowed answer codes.""" + scaled = [value / temperature for value in logprobs] + peak = max(scaled) + weights = [math.exp(value - peak) for value in scaled] + total = sum(weights) + return [value / total for value in weights] + + +def _encode(tokenizer, text): + return tokenizer.encode(text, add_special_tokens=False) + + +def build_prompt(tokenizer, state, questions, codes, compact=False): + """Build one encoder prompt containing every independent question.""" + if compact: + parts = [state.strip()] + for index, (instructions, labels) in enumerate(questions, 1): + prefix = "" if len(questions) == 1 else f"q{index} " + options = " ".join(f"{codes[i]} {item}" for i, item in enumerate(labels)) + parts.append(f" {prefix}{instructions} {options}") + text = tokenizer.apply_chat_template( + [{"role": "user", "content": "".join(parts)}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + return _encode(tokenizer, text) + + parts = [ + "Answer each question about the state. Each question lists its allowed answers.\n\n", + "State:\n", + state.strip(), + ] + for index, (instructions, labels) in enumerate(questions, 1): + parts.append(f"\n\nQuestion q{index}: {instructions}\n") + parts.append("\n".join(f" {codes[i]}: {item}" for i, item in enumerate(labels))) + parts.append('\n\nReply with one line per question, in order, formatted as "- label".') + text = tokenizer.apply_chat_template( + [{"role": "user", "content": "".join(parts)}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + return _encode(tokenizer, text) + + +def build_seed_canvas( + tokenizer, + question_count, + canvas_length, + vocab_size, + compact=False, +): + """Build the active answer template and randomize only its answer slots.""" + rng = random.Random(DEFAULT_DIFFUSION_SEED) + if compact: + if question_count > canvas_length: + raise RequestContractError( + f"{question_count} questions need {question_count} canvas tokens; " + f"DiffusionGemma provides {canvas_length}" + ) + return [rng.randrange(vocab_size) for _ in range(question_count)], list( + range(question_count) + ) + + canvas = _encode(tokenizer, CANVAS_PREFIX) + positions = [] + for _ in range(question_count): + canvas.extend(_encode(tokenizer, "-")) + positions.append(len(canvas)) + canvas.append(rng.randrange(vocab_size)) + canvas.extend(_encode(tokenizer, "\n")) + canvas.extend(_encode(tokenizer, CANVAS_SUFFIX)) + if len(canvas) > canvas_length: + raise RequestContractError( + f"{question_count} questions need {len(canvas)} canvas tokens; " + f"DiffusionGemma provides {canvas_length}" + ) + return canvas, positions + + +class MlxVlmDiffusionEngine: + """Read all question slots in one DiffusionGemma denoising forward.""" + + def __init__(self, model_id=None, base_url=None, transport=None, compact=False): + self.model_id = resolve_model_id(model_id or "diffusion") + if self.model_id != DIFFUSION_GEMMA_MODEL: + raise ValueError("the mlx-vlm-diffusion backend requires the diffusion model profile") + self.model_path, self.model_revision = resolve_model_snapshot(self.model_id) + self.tokenizer = load_tokenizer(self.model_path) + self.registry_file, registry, self.codes, self.code_token_ids = packaged_code_registry( + self.model_id + ) + if registry.get("model_revision") != self.model_revision: + raise ValueError(f"{self.registry_file.name} does not match the pinned model revision") + if registry.get("tokenizer_sha256") != tokenizer_sha256(self.model_path): + raise ValueError(f"{self.registry_file.name} does not match the pinned tokenizer") + if registry.get("codes") != diffusion_answer_code_entries(self.tokenizer): + raise ValueError(f"{self.registry_file.name} is incomplete or invalid") + + config = json.loads((self.model_path / "config.json").read_text()) + self.canvas_length = int(config["canvas_length"]) + self.vocab_size = int(config["text_config"]["vocab_size"]) + self.base_url = ( + base_url or os.environ.get("SYSTEM_ONE_MLX_VLM_BASE_URL") or DEFAULT_BASE_URL + ).rstrip("/") + self.timeout = float(os.environ.get("SYSTEM_ONE_MLX_VLM_TIMEOUT", DEFAULT_TIMEOUT)) + self.api_key = os.environ.get("SYSTEM_ONE_MLX_VLM_API_KEY") + self.transport = transport or self._post + self.compact = compact + + def output_tokens_for(self, questions): + """A read-only canvas does not commit generated tokens.""" + return 0 + + def evaluate(self, state, questions, template=None): + """Return one answer-code distribution per slot in a shared canvas.""" + if template is not None: + raise RequestContractError("custom templates are not supported by DiffusionGemma") + self._validate_questions(questions) + if self.compact and len(questions) != 1: + raise RequestContractError( + "the compact DiffusionGemma backend requires exactly one question" + ) + input_ids = build_prompt( + self.tokenizer, + state, + questions, + self.codes, + compact=self.compact, + ) + if len(input_ids) > MAX_PROMPT_TOKENS: + raise RequestContractError(f"prompt exceeds the {MAX_PROMPT_TOKENS} token limit") + if len(input_ids) > MAX_TOTAL_INPUT_TOKENS: + raise RequestContractError( + f"request exceeds the {MAX_TOTAL_INPUT_TOKENS} input token limit" + ) + seed_canvas, positions = build_seed_canvas( + self.tokenizer, + len(questions), + self.canvas_length, + self.vocab_size, + compact=self.compact, + ) + slots = [ + { + "position": position, + "token_ids": list(self.code_token_ids[: len(labels)]), + } + for position, (_, labels) in zip(positions, questions) + ] + payload = { + "model": str(self.model_path), + "input_ids": input_ids, + "seed_canvas": seed_canvas, + "slots": slots, + "candidate_only": True, + } + if self.compact: + payload["encoder_layers"] = FAST_DIFFUSION_ENCODER_LAYERS + + started = time.perf_counter() + response = self.transport("/v1/diffusion/reads", payload) + elapsed_ms = (time.perf_counter() - started) * 1000 + probabilities = self._parse_response( + response, + slots, + len(input_ids), + str(self.model_path), + expected_encoder_layers=(FAST_DIFFUSION_ENCODER_LAYERS if self.compact else None), + ) + return probabilities, len(input_ids), elapsed_ms + + def _validate_questions(self, questions): + if not questions: + raise RequestContractError("at least one question is required") + for instructions, labels in questions: + if not labels: + raise RequestContractError(f"{instructions!r} has no options") + if len(labels) > len(self.codes): + raise TooManyOptions( + f"{len(labels)} options for {instructions!r}; " + f"the engine has {len(self.codes)} single-token answer codes" + ) + + def _parse_response( + self, + response, + slots, + input_tokens, + expected_model, + expected_encoder_layers=None, + ): + try: + reads = response["reads"] + usage = response["usage"] + reported_tokens = usage["prompt_tokens"] + denoising_steps = usage["denoising_steps"] + candidate_only = usage["candidate_only"] + encoder_layers = usage["encoder_layers"] + model = response["model"] + except (KeyError, TypeError) as error: + raise MlxVlmError("MLX-VLM returned an invalid diffusion read response") from error + if model != expected_model: + raise MlxVlmError(f"MLX-VLM served {model!r}, expected pinned model {expected_model!r}") + if denoising_steps != 1: + raise MlxVlmError(f"MLX-VLM performed {denoising_steps!r} denoising steps, expected 1") + if candidate_only is not True: + raise MlxVlmError("MLX-VLM did not confirm candidate-only scoring") + if encoder_layers != expected_encoder_layers: + raise MlxVlmError( + "MLX-VLM used a different encoder layer count: " + f"expected {expected_encoder_layers!r}, got {encoder_layers!r}" + ) + if reported_tokens != input_tokens: + raise MlxVlmError( + "MLX-VLM token count differs from the pinned tokenizer: " + f"expected {input_tokens}, got {reported_tokens}" + ) + if len(reads) != len(slots): + raise MlxVlmError("MLX-VLM returned incomplete diffusion read slots") + + probabilities = [] + for read, slot in zip(reads, slots): + try: + token_ids = read["token_ids"] + logprobs = read["logprobs"] + position = read["position"] + except (KeyError, TypeError) as error: + raise MlxVlmError("MLX-VLM returned an invalid diffusion read slot") from error + if position != slot["position"] or token_ids != slot["token_ids"]: + raise MlxVlmError("MLX-VLM returned a mismatched diffusion read slot") + if len(logprobs) != len(token_ids) or not all( + isinstance(value, (int, float)) and math.isfinite(value) for value in logprobs + ): + raise MlxVlmError("MLX-VLM returned invalid token log probabilities") + probabilities.append(restricted_softmax(logprobs)) + return probabilities + + def _post(self, path, payload): + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + request = Request( + f"{self.base_url}{path}", + data=json.dumps(payload).encode("utf-8"), + headers=headers, + method="POST", + ) + try: + with urlopen(request, timeout=self.timeout) as response: + return json.load(response) + except HTTPError as error: + detail = error.read(4096).decode("utf-8", errors="replace") + raise MlxVlmError(f"MLX-VLM returned HTTP {error.code}: {detail}") from error + except (URLError, TimeoutError) as error: + raise MlxVlmError(f"could not reach MLX-VLM at {self.base_url}") from error + except (json.JSONDecodeError, UnicodeDecodeError) as error: + raise MlxVlmError("MLX-VLM returned invalid JSON") from error diff --git a/server/src/system_one_lite/vllm_metal.py b/server/src/system_one_lite/vllm_metal.py new file mode 100644 index 0000000..5941dcf --- /dev/null +++ b/server/src/system_one_lite/vllm_metal.py @@ -0,0 +1,250 @@ +"""System One backend for a separate vllm-metal OpenAI server.""" + +import json +import math +import os +import time +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from mlx_lm.utils import load_tokenizer + +from .engine import ( + MAX_PROMPT_TOKENS, + MAX_TOTAL_INPUT_TOKENS, + TEMPERATURE, + RequestContractError, + TooManyOptions, + common_prefix_len, + configured_model, + load_code_registry, + packaged_code_registry, + resolve_model_snapshot, +) +from .prompts import chat_filled + +DEFAULT_BASE_URL = "http://127.0.0.1:8000" +DEFAULT_TIMEOUT = 120.0 +MAX_LOGPROB_TOKEN_IDS = 128 + + +class VllmMetalError(RuntimeError): + """The vllm-metal server could not complete a structured read.""" + + +def restricted_softmax(logprobs, temperature=TEMPERATURE): + """Renormalize raw vocabulary log-probabilities over allowed codes.""" + scaled = [value / temperature for value in logprobs] + peak = max(scaled) + weights = [math.exp(value - peak) for value in scaled] + total = sum(weights) + return [value / total for value in weights] + + +class VllmMetalEngine: + """Read answer-code logits through vllm-metal's completions API. + + vllm-metal does not yet implement DiffusionGemma or seeded canvas reads. + This backend sends System One's exact prompt token IDs in one batch. Each + prompt stays independent. Requested answer codes are split into exact + groups of 128 because that is vLLM's per-read limit. + """ + + def __init__(self, model_id=None, base_url=None, transport=None, tokenizer=None): + self.model_id = configured_model(model_id) + if tokenizer is None: + model_path, model_revision = resolve_model_snapshot(self.model_id, tokenizer_only=True) + self.tokenizer = load_tokenizer(model_path) + self.registry_file, self.codes, self.code_token_ids = load_code_registry( + self.model_id, + self.tokenizer, + model_path, + model_revision, + ) + else: + self.tokenizer = tokenizer + self.registry_file, _, self.codes, self.code_token_ids = packaged_code_registry( + self.model_id + ) + self.base_url = ( + base_url or os.environ.get("SYSTEM_ONE_VLLM_BASE_URL") or DEFAULT_BASE_URL + ).rstrip("/") + self.timeout = float(os.environ.get("SYSTEM_ONE_VLLM_TIMEOUT", DEFAULT_TIMEOUT)) + self.api_key = os.environ.get("SYSTEM_ONE_VLLM_API_KEY") + self.transport = transport or self._post + + def output_tokens_for(self, questions): + """Count the internal one-token probes needed for this request.""" + return sum(math.ceil(len(labels) / MAX_LOGPROB_TOKEN_IDS) for _, labels in questions) + + def evaluate(self, state, questions, template=None): + """Return one exact code distribution per independent question.""" + prompts, input_tokens = self._prepare(state, questions, template) + started = time.perf_counter() + self._verify_server_tokenizer() + probabilities, reported_input_tokens = self._read_batches(prompts, questions) + elapsed_ms = (time.perf_counter() - started) * 1000 + if reported_input_tokens != input_tokens: + raise VllmMetalError( + "vllm-metal token count differs from the pinned tokenizer: " + f"expected {input_tokens}, got {reported_input_tokens}" + ) + return probabilities, input_tokens, elapsed_ms + + def _prepare(self, state, questions, template): + if template is not None: + raise RequestContractError("custom templates are not supported by vllm-metal") + if not questions: + raise RequestContractError("at least one question is required") + + prompts = [] + input_tokens = 0 + for instructions, labels in questions: + if not labels: + raise RequestContractError(f"{instructions!r} has no options") + if len(labels) > len(self.codes): + raise TooManyOptions( + f"{len(labels)} options for {instructions!r}; " + f"the engine has {len(self.codes)} single-token answer codes" + ) + prompt = self._prompt_token_ids(state, instructions, labels) + if len(prompt) > MAX_PROMPT_TOKENS: + raise RequestContractError( + f"prompt for {instructions!r} exceeds the {MAX_PROMPT_TOKENS} token limit" + ) + input_tokens += len(prompt) * math.ceil(len(labels) / MAX_LOGPROB_TOKEN_IDS) + if input_tokens > MAX_TOTAL_INPUT_TOKENS: + raise RequestContractError( + f"request exceeds the {MAX_TOTAL_INPUT_TOKENS} input token limit" + ) + prompts.append(prompt) + return prompts, input_tokens + + def _prompt_token_ids(self, state, instructions, labels): + text, marks = chat_filled( + self.tokenizer, + state, + [(instructions, labels)], + codes=self.codes, + ) + prefix = text[: marks[0]] + base = self.tokenizer.encode(prefix) + candidate = self.tokenizer.encode(prefix + self.codes[0]) + cut = common_prefix_len(candidate, base) + full = self.tokenizer.encode(text) + if len(candidate) - cut != 1 or candidate[: cut + 1] != full[: cut + 1]: + raise RequestContractError(f"slot for {instructions!r} is not a single token") + return full[:cut] + + def _verify_server_tokenizer(self): + token_ids = list(self.code_token_ids) + response = self.transport( + "/detokenize", + {"model": self.model_id, "tokens": token_ids}, + ) + try: + remote_text = response["prompt"] + except (KeyError, TypeError) as error: + raise VllmMetalError("vllm-metal returned an invalid detokenize response") from error + if remote_text != self.tokenizer.decode(token_ids): + raise VllmMetalError("vllm-metal does not match the pinned tokenizer") + + def _read_batches(self, prompts, questions): + max_options = max(len(labels) for _, labels in questions) + scores = [{} for _ in questions] + prompt_tokens = 0 + for start in range(0, max_options, MAX_LOGPROB_TOKEN_IDS): + end = min(start + MAX_LOGPROB_TOKEN_IDS, max_options) + active = [ + (index, prompt) + for index, (prompt, (_, labels)) in enumerate(zip(prompts, questions)) + if len(labels) > start + ] + token_ids = list(self.code_token_ids[start:end]) + payload = { + "model": self.model_id, + "prompt": [prompt for _, prompt in active], + "max_tokens": 1, + "temperature": 0, + "logprobs": 0, + "logprob_token_ids": token_ids, + "return_tokens_as_token_ids": True, + } + response = self.transport("/v1/completions", payload) + try: + choices = sorted(response["choices"], key=lambda choice: choice["index"]) + response_model = response["model"] + batch_prompt_tokens = response["usage"]["prompt_tokens"] + except (KeyError, TypeError, ValueError) as error: + raise VllmMetalError("vllm-metal returned an invalid logprobs response") from error + if response_model != self.model_id: + raise VllmMetalError( + f"vllm-metal served {response_model!r}, expected {self.model_id!r}" + ) + if ( + isinstance(batch_prompt_tokens, bool) + or not isinstance(batch_prompt_tokens, int) + or batch_prompt_tokens < 0 + ): + raise VllmMetalError("vllm-metal returned an invalid prompt token count") + prompt_tokens += batch_prompt_tokens + if [choice["index"] for choice in choices] != list(range(len(active))): + raise VllmMetalError("vllm-metal returned incomplete batch choices") + for choice, (question_index, _) in zip(choices, active): + option_count = len(questions[question_index][1]) + requested_ids = list(self.code_token_ids[start : min(end, option_count)]) + scores[question_index].update(self._choice_logprobs(choice, requested_ids)) + + probabilities = [ + restricted_softmax([score[token_id] for token_id in self.code_token_ids[: len(labels)]]) + for score, (_, labels) in zip(scores, questions) + ] + return probabilities, prompt_tokens + + def _choice_logprobs(self, choice, token_ids): + try: + top_logprobs = choice["logprobs"]["top_logprobs"][0] + if not isinstance(top_logprobs, dict): + raise TypeError + by_token_id = { + int(token.removeprefix("token_id:")): logprob + for token, logprob in top_logprobs.items() + if token.startswith("token_id:") + } + except (AttributeError, IndexError, KeyError, TypeError, ValueError) as error: + raise VllmMetalError("vllm-metal returned invalid token logprobs") from error + missing = [token_id for token_id in token_ids if token_id not in by_token_id] + if missing: + raise VllmMetalError( + "vllm-metal omitted requested answer-code token IDs: " + + ", ".join(map(str, missing)) + ) + if any( + isinstance(by_token_id[token_id], bool) + or not isinstance(by_token_id[token_id], (int, float)) + or not math.isfinite(by_token_id[token_id]) + for token_id in token_ids + ): + raise VllmMetalError("vllm-metal returned invalid token logprobs") + return {token_id: by_token_id[token_id] for token_id in token_ids} + + def _post(self, path, payload): + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + request = Request( + f"{self.base_url}{path}", + data=json.dumps(payload).encode("utf-8"), + headers=headers, + method="POST", + ) + try: + with urlopen(request, timeout=self.timeout) as response: + return json.load(response) + except HTTPError as error: + detail = error.read(4096).decode("utf-8", errors="replace") + raise VllmMetalError(f"vllm-metal returned HTTP {error.code}: {detail}") from error + except (URLError, TimeoutError) as error: + raise VllmMetalError(f"could not reach vllm-metal at {self.base_url}") from error + except (json.JSONDecodeError, UnicodeDecodeError) as error: + raise VllmMetalError("vllm-metal returned invalid JSON") from error diff --git a/server/tests/integration/test_api.py b/server/tests/integration/test_api.py index bfc7f5e..9d2e3ac 100644 --- a/server/tests/integration/test_api.py +++ b/server/tests/integration/test_api.py @@ -7,11 +7,13 @@ from system_one_lite import api as api_module from system_one_lite.api import app, create_app from system_one_lite.engine import RequestContractError +from system_one_lite.mlx_vlm_diffusion import MlxVlmError from system_one_lite.schemas import ( MAX_QUESTIONS, MAX_REQUEST_BYTES, MAX_STATE_CHARS, ) +from system_one_lite.vllm_metal import VllmMetalError client = None test_app = None @@ -49,6 +51,99 @@ def test_configured_engine_uses_selected_profile(monkeypatch): assert selected == ["larger"] +def test_configured_engine_can_select_vllm_metal(monkeypatch): + selected = [] + monkeypatch.setenv("SYSTEM_ONE_BACKEND", "vllm-metal") + monkeypatch.setenv("SYSTEM_ONE_MODEL", "larger") + monkeypatch.setattr(api_module, "VllmMetalEngine", selected.append) + + assert api_module.configured_engine() is None + assert selected == ["larger"] + + +def test_configured_engine_can_select_mlx_vlm_diffusion(monkeypatch): + selected = [] + monkeypatch.setenv("SYSTEM_ONE_BACKEND", "mlx-vlm-diffusion") + monkeypatch.delenv("SYSTEM_ONE_MODEL", raising=False) + monkeypatch.setattr(api_module, "MlxVlmDiffusionEngine", selected.append) + + assert api_module.configured_engine() is None + assert selected == [None] + + +def test_configured_engine_rejects_unknown_backend(monkeypatch): + monkeypatch.setenv("SYSTEM_ONE_BACKEND", "unknown") + + with pytest.raises(ValueError, match="unknown SYSTEM_ONE_BACKEND"): + api_module.configured_engine() + + +def test_vllm_metal_usage_counts_internal_read_tokens(): + class FakeVllmMetalEngine: + model_id = "test-model" + + def evaluate(self, state, questions): + return [[0.75, 0.25] for _ in questions], 20, 1.0 + + def output_tokens_for(self, questions): + return len(questions) + + request = { + "state": "state", + "questions": { + "first": {"type": "noul", "instructions": "first?"}, + "second": {"type": "noul", "instructions": "second?"}, + }, + } + with TestClient(create_app(engine=FakeVllmMetalEngine())) as vllm_client: + response = vllm_client.post("/evaluate", json=request) + + assert response.status_code == 200 + assert response.json()["usage"] == {"input_tokens": 20, "output_tokens": 2} + + +def test_vllm_metal_failure_becomes_bad_gateway(): + class FailedVllmMetalEngine: + model_id = "test-model" + + def evaluate(self, state, questions): + raise VllmMetalError("server unavailable") + + def output_tokens_for(self, questions): + return len(questions) + + request = { + "state": "state", + "questions": {"first": {"type": "noul", "instructions": "first?"}}, + } + with TestClient(create_app(engine=FailedVllmMetalEngine())) as vllm_client: + response = vllm_client.post("/evaluate", json=request) + + assert response.status_code == 502 + assert response.json()["detail"] == "server unavailable" + + +def test_mlx_vlm_failure_becomes_bad_gateway(): + class FailedDiffusionEngine: + model_id = "test-model" + + def evaluate(self, state, questions): + raise MlxVlmError("diffusion server unavailable") + + def output_tokens_for(self, questions): + return 0 + + request = { + "state": "state", + "questions": {"first": {"type": "noul", "instructions": "first?"}}, + } + with TestClient(create_app(engine=FailedDiffusionEngine())) as diffusion_client: + response = diffusion_client.post("/evaluate", json=request) + + assert response.status_code == 502 + assert response.json()["detail"] == "diffusion server unavailable" + + def test_quickstart_mixed_three_questions(): """The quickstart.md request: Choice + Score + Noul, one pass, full shape.""" r = evaluate( diff --git a/server/tests/unit/test_mlx_vlm_diffusion.py b/server/tests/unit/test_mlx_vlm_diffusion.py new file mode 100644 index 0000000..e877c8f --- /dev/null +++ b/server/tests/unit/test_mlx_vlm_diffusion.py @@ -0,0 +1,242 @@ +import math +from pathlib import Path + +import pytest + +from system_one_lite.engine import RequestContractError +from system_one_lite.mlx_vlm_diffusion import ( + MlxVlmDiffusionEngine, + MlxVlmError, + build_seed_canvas, +) + + +class FakeTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [ord(char) % 200 for char in text] + + def apply_chat_template(self, messages, **kwargs): + del kwargs + return "CHAT:" + messages[0]["content"] + + +def bare_engine(transport, canvas_length=256): + engine = MlxVlmDiffusionEngine.__new__(MlxVlmDiffusionEngine) + engine.model_id = "diffusion" + engine.model_path = Path("/pinned/diffusion") + engine.tokenizer = FakeTokenizer() + engine.codes = ("A", "B", "C") + engine.code_token_ids = (10, 11, 12) + engine.canvas_length = canvas_length + engine.vocab_size = 200 + engine.transport = transport + engine.compact = False + return engine + + +def test_evaluate_reads_all_questions_from_one_seeded_canvas(): + captured = {} + + def transport(path, payload): + captured.update(path=path, payload=payload) + return { + "model": "/pinned/diffusion", + "reads": [ + { + "position": slot["position"], + "token_id": slot["token_ids"][0], + "token_ids": slot["token_ids"], + "logprobs": [-0.1 * (index + 1) for index in range(len(slot["token_ids"]))], + } + for slot in payload["slots"] + ], + "usage": { + "prompt_tokens": len(payload["input_ids"]), + "denoising_steps": 1, + "candidate_only": payload["candidate_only"], + "encoder_layers": payload.get("encoder_layers"), + }, + } + + engine = bare_engine(transport) + questions = [("first", ["yes", "no"]), ("second", ["a", "b", "c"])] + + probabilities, input_tokens, elapsed_ms = engine.evaluate("state", questions) + + assert captured["path"] == "/v1/diffusion/reads" + assert captured["payload"]["model"] == "/pinned/diffusion" + assert len(captured["payload"]["input_ids"]) == input_tokens + assert len(captured["payload"]["seed_canvas"]) < 256 + assert captured["payload"]["candidate_only"] is True + assert [slot["token_ids"] for slot in captured["payload"]["slots"]] == [ + [10, 11], + [10, 11, 12], + ] + assert len(probabilities) == 2 + assert all(math.isclose(sum(row), 1.0) for row in probabilities) + assert elapsed_ms >= 0 + assert engine.output_tokens_for(questions) == 0 + + +def test_seed_canvas_rejects_more_questions_than_fit(): + with pytest.raises(RequestContractError, match="canvas tokens"): + build_seed_canvas(FakeTokenizer(), 2, canvas_length=5, vocab_size=200) + + +def test_seed_canvas_fits_the_public_64_question_limit(): + first, positions = build_seed_canvas(FakeTokenizer(), 64, canvas_length=256, vocab_size=200) + second, second_positions = build_seed_canvas( + FakeTokenizer(), 64, canvas_length=256, vocab_size=200 + ) + + assert first == second + assert positions == second_positions + assert len(first) <= 256 + assert len(positions) == 64 + + +def test_compact_read_uses_one_canvas_token_per_question(): + captured = {} + + def transport(path, payload): + captured.update(path=path, payload=payload) + return { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 0, + "token_id": 11, + "token_ids": [10, 11], + "logprobs": [-2.0, -0.1], + } + ], + "usage": { + "prompt_tokens": len(payload["input_ids"]), + "denoising_steps": 1, + "candidate_only": payload["candidate_only"], + "encoder_layers": payload.get("encoder_layers"), + }, + } + + engine = bare_engine(transport) + engine.compact = True + probabilities, _, _ = engine.evaluate("delivered order.", [("late?", ["yes", "no"])]) + + assert captured["payload"]["seed_canvas"] == [163] + assert captured["payload"]["slots"] == [{"position": 0, "token_ids": [10, 11]}] + assert captured["payload"]["encoder_layers"] == 6 + assert probabilities[0][1] > probabilities[0][0] + + +def test_compact_read_rejects_more_than_one_question(): + engine = bare_engine(None) + engine.compact = True + + with pytest.raises(RequestContractError, match="exactly one question"): + engine.evaluate( + "state", + [("first", ["yes", "no"]), ("second", ["yes", "no"])], + ) + + +def test_response_must_match_requested_slot(): + engine = bare_engine(None) + with pytest.raises(MlxVlmError, match="mismatched"): + engine._parse_response( + { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 2, + "token_ids": [10, 11], + "logprobs": [-1.0, -2.0], + } + ], + "usage": { + "prompt_tokens": 3, + "denoising_steps": 1, + "candidate_only": True, + "encoder_layers": None, + }, + }, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + ) + + +def test_response_must_confirm_model_and_one_step(): + engine = bare_engine(None) + base = { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 1, + "token_ids": [10, 11], + "logprobs": [-1.0, -2.0], + } + ], + "usage": { + "prompt_tokens": 3, + "denoising_steps": 1, + "candidate_only": True, + "encoder_layers": None, + }, + } + + wrong_model = {**base, "model": "another-model"} + with pytest.raises(MlxVlmError, match="expected pinned model"): + engine._parse_response( + wrong_model, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + ) + + wrong_steps = {**base, "usage": {**base["usage"], "denoising_steps": 2}} + with pytest.raises(MlxVlmError, match="expected 1"): + engine._parse_response( + wrong_steps, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + ) + + +def test_response_must_confirm_fast_read_settings(): + engine = bare_engine(None) + base = { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 1, + "token_ids": [10, 11], + "logprobs": [-1.0, -2.0], + } + ], + "usage": { + "prompt_tokens": 3, + "denoising_steps": 1, + "candidate_only": True, + "encoder_layers": 6, + }, + } + + with pytest.raises(MlxVlmError, match="invalid diffusion read response"): + engine._parse_response( + {**base, "usage": {"prompt_tokens": 3, "denoising_steps": 1}}, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + expected_encoder_layers=6, + ) + + with pytest.raises(MlxVlmError, match="different encoder layer count"): + engine._parse_response( + {**base, "usage": {**base["usage"], "encoder_layers": 30}}, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + expected_encoder_layers=6, + ) diff --git a/server/tests/unit/test_vllm_metal.py b/server/tests/unit/test_vllm_metal.py new file mode 100644 index 0000000..ac2f3e3 --- /dev/null +++ b/server/tests/unit/test_vllm_metal.py @@ -0,0 +1,301 @@ +import math + +import pytest + +from system_one_lite import vllm_metal as vllm_module +from system_one_lite.engine import DEFAULT_MODEL, RequestContractError, TooManyOptions +from system_one_lite.vllm_metal import ( + MAX_LOGPROB_TOKEN_IDS, + VllmMetalEngine, + VllmMetalError, + restricted_softmax, +) + + +class FakeTokenizer: + marker = '{"answer": "' + + def __init__(self, token_count=12): + self.token_count = token_count + + def apply_chat_template(self, messages, **kwargs): + return "" + messages[0]["content"] + "" + + def encode(self, text): + before, marker, suffix = text.rpartition(self.marker) + if not marker: + return list(range(self.token_count)) + tokens = list(range(self.token_count)) + if suffix: + tokens.append(32) + tokens.extend(ord(character) for character in suffix[1:]) + return tokens + + def decode(self, token_ids): + return ",".join(map(str, token_ids)) + + +def response_for(payload, scores=None, prompt_tokens=None): + token_ids = payload["logprob_token_ids"] + scores = scores or [-float(index) for index in range(len(token_ids))] + return { + "model": DEFAULT_MODEL, + "choices": [ + { + "index": choice_index, + "logprobs": { + "top_logprobs": [ + { + f"token_id:{token_id}": score + for token_id, score in zip(token_ids, scores) + } + ] + }, + } + for choice_index in range(len(payload["prompt"])) + ], + "usage": { + "prompt_tokens": prompt_tokens + if prompt_tokens is not None + else len(payload["prompt"]) * 12 + }, + } + + +def fake_engine(transport, token_count=12): + return VllmMetalEngine( + DEFAULT_MODEL, + transport=transport, + tokenizer=FakeTokenizer(token_count), + ) + + +def test_exact_requested_token_logprobs_are_normalized(): + calls = [] + tokenizer = FakeTokenizer() + + def transport(path, payload): + calls.append((path, payload)) + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return response_for(payload, [-2.0, -1.0, -3.0]) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + probabilities, input_tokens, elapsed_ms = engine.evaluate( + "state", [("pick", ["one", "two", "three"])] + ) + + assert probabilities[0] == pytest.approx(restricted_softmax([-2.0, -1.0, -3.0])) + assert input_tokens == 12 + assert elapsed_ms >= 0 + assert [path for path, _ in calls] == ["/detokenize", "/v1/completions"] + payload = calls[-1][1] + assert payload["model"] == DEFAULT_MODEL + assert payload["max_tokens"] == 1 + assert payload["logprob_token_ids"] == list(engine.code_token_ids[:3]) + assert payload["prompt"] == [list(range(12))] + assert engine.output_tokens_for([("pick", ["one", "two", "three"])]) == 1 + + +def test_questions_are_independent_and_token_usage_is_summed(): + payloads = [] + tokenizer = FakeTokenizer(token_count=7) + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + payloads.append(payload) + return response_for(payload, prompt_tokens=14) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + probabilities, input_tokens, _ = engine.evaluate( + "shared", [("first", ["yes", "no"]), ("second", ["up", "down"])] + ) + + assert len(probabilities) == 2 + assert input_tokens == 14 + assert len(payloads) == 1 + assert len(payloads[0]["prompt"]) == 2 + + +def test_more_than_128_options_are_read_in_exact_chunks(): + tokenizer = FakeTokenizer(token_count=5) + payloads = [] + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + payloads.append(payload) + return response_for(payload, prompt_tokens=5) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + labels = [f"option {index}" for index in range(MAX_LOGPROB_TOKEN_IDS + 2)] + probabilities, input_tokens, _ = engine.evaluate("state", [("pick", labels)]) + + assert [len(payload["logprob_token_ids"]) for payload in payloads] == [128, 2] + assert len(probabilities[0]) == 130 + assert sum(probabilities[0]) == pytest.approx(1.0) + assert input_tokens == 10 + assert engine.output_tokens_for([("pick", labels)]) == 2 + + +def test_mixed_option_counts_skip_finished_questions_on_later_chunks(): + tokenizer = FakeTokenizer(token_count=5) + batch_sizes = [] + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + batch_sizes.append(len(payload["prompt"])) + return response_for(payload, prompt_tokens=len(payload["prompt"]) * 5) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + questions = [ + ("small", ["yes", "no"]), + ("large", [f"option {index}" for index in range(129)]), + ] + _, input_tokens, _ = engine.evaluate("state", questions) + + assert batch_sizes == [2, 1] + assert input_tokens == 15 + assert engine.output_tokens_for(questions) == 3 + + +def test_server_tokenizer_mismatch_is_an_error_before_inference(): + calls = [] + + def transport(path, payload): + calls.append(path) + return {"prompt": "wrong"} + + engine = fake_engine(transport) + with pytest.raises(VllmMetalError, match="pinned tokenizer"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + assert calls == ["/detokenize"] + + +def test_missing_requested_token_is_an_error(): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + response = response_for(payload) + response["choices"][0]["logprobs"]["top_logprobs"][0].popitem() + return response + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="omitted requested"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +def test_wrong_response_model_is_an_error(): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return {**response_for(payload), "model": "wrong-model"} + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="expected"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +@pytest.mark.parametrize("prompt_tokens", ["12", -1, True]) +def test_invalid_prompt_token_count_is_an_error(prompt_tokens): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return response_for(payload, prompt_tokens=prompt_tokens) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="prompt token count"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), "-1.0", True]) +def test_invalid_requested_logprob_is_an_error(bad_value): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return response_for(payload, scores=[bad_value, -1.0]) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="invalid token logprobs"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +def test_non_mapping_logprobs_are_an_error(): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + response = response_for(payload) + response["choices"][0]["logprobs"]["top_logprobs"][0] = [] + return response + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="invalid token logprobs"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +def test_contract_checks_happen_before_transport(monkeypatch): + def transport(path, payload): + raise AssertionError("transport should not run") + + engine = fake_engine(transport) + with pytest.raises(RequestContractError, match="at least one"): + engine.evaluate("state", []) + with pytest.raises(RequestContractError, match="custom templates"): + engine.evaluate("state", [("pick", ["yes"])], template=object()) + with pytest.raises(TooManyOptions): + engine.evaluate("state", [("pick", ["x"] * (len(engine.codes) + 1))]) + + monkeypatch.setattr(vllm_module, "MAX_PROMPT_TOKENS", 11) + with pytest.raises(RequestContractError, match="prompt for"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +def test_total_token_limit_is_checked_before_transport(monkeypatch): + def transport(path, payload): + raise AssertionError("transport should not run") + + engine = fake_engine(transport, token_count=10) + monkeypatch.setattr(vllm_module, "MAX_TOTAL_INPUT_TOKENS", 19) + + with pytest.raises(RequestContractError, match="request exceeds"): + engine.evaluate( + "state", + [("first", ["yes", "no"]), ("second", ["up", "down"])], + ) + + +def test_prompt_at_input_limit_keeps_one_remote_output_slot(monkeypatch): + tokenizer = FakeTokenizer(token_count=12) + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return response_for(payload, prompt_tokens=12) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + monkeypatch.setattr(vllm_module, "MAX_PROMPT_TOKENS", 12) + + probabilities, input_tokens, _ = engine.evaluate("state", [("pick", ["yes", "no"])]) + + assert len(probabilities[0]) == 2 + assert input_tokens == 12 + + +def test_restricted_softmax_applies_system_one_temperature(): + probabilities = restricted_softmax([-2.0, -1.0]) + expected_second = 1 / (1 + math.exp(-1 / 0.7)) + + assert probabilities == pytest.approx([1 - expected_second, expected_second]) diff --git a/server/tools/gen_answer_codes.py b/server/tools/gen_answer_codes.py index dc90541..5b1e9dd 100644 --- a/server/tools/gen_answer_codes.py +++ b/server/tools/gen_answer_codes.py @@ -18,7 +18,9 @@ from system_one_lite.engine import ( DEFAULT_MODEL, + DIFFUSION_GEMMA_MODEL, answer_code_entries, + diffusion_answer_code_entries, resolve_model_id, resolve_model_snapshot, tokenizer_sha256, @@ -49,7 +51,10 @@ def main(): model_id, tokenizer_only=True, revision=args.revision ) tokenizer = load_tokenizer(model_path) - entries = answer_code_entries(tokenizer) + is_diffusion = model_id == DIFFUSION_GEMMA_MODEL + entries = ( + diffusion_answer_code_entries(tokenizer) if is_diffusion else answer_code_entries(tokenizer) + ) output = Path(args.output) if args.output else default_output(model_id) if not output.is_absolute(): output = Path.cwd() / output @@ -58,7 +63,11 @@ def main(): "model": model_id, "model_revision": revision, "tokenizer_sha256": tokenizer_sha256(model_path), - "slot_context": "codes appended after the assistant JSON answer prefix", + "slot_context": ( + "codes occupy one seeded DiffusionGemma canvas token after a list marker" + if is_diffusion + else "codes appended after the assistant JSON answer prefix" + ), "codes": entries, } output.write_text(json.dumps(out, indent=1) + "\n") From 1c4216784236dab285a24d2380368cdd31ecc069 Mon Sep 17 00:00:00 2001 From: Sam Snelling Date: Thu, 17 Sep 2026 18:35:55 -0500 Subject: [PATCH 2/4] Add structured read profiler Add a warmed concurrency profiler for structured reads. Report throughput, client latency, and request round-trip latency. Accept full-read responses without encoder metadata. --- docs/diffusion-gemma.md | 13 ++ .../src/system_one_lite/mlx_vlm_diffusion.py | 2 +- .../tests/tools/test_benchmark_diffusion.py | 30 +++ server/tests/unit/test_mlx_vlm_diffusion.py | 28 +++ server/tools/benchmark_diffusion.py | 189 ++++++++++++++++++ 5 files changed, 261 insertions(+), 1 deletion(-) create mode 100644 server/tests/tools/test_benchmark_diffusion.py create mode 100644 server/tools/benchmark_diffusion.py diff --git a/docs/diffusion-gemma.md b/docs/diffusion-gemma.md index 0007ead..774e2db 100644 --- a/docs/diffusion-gemma.md +++ b/docs/diffusion-gemma.md @@ -54,3 +54,16 @@ This path differs from the vllm-metal backend in two ways: The public System One request and response formats do not change. Read-only canvas work reports zero output tokens because no tokens are committed. + +## Profile the backend + +Start the MLX-VLM server, then run the warmed concurrency sweep from +`server/`: + +```bash +uv run python -m tools.benchmark_diffusion +``` + +Use `--json` to save machine-readable results. Set `--requests` and +`--concurrency` to change the load. The report includes request throughput, +decision throughput, client latency, and request round-trip latency. diff --git a/server/src/system_one_lite/mlx_vlm_diffusion.py b/server/src/system_one_lite/mlx_vlm_diffusion.py index 8bfe4cc..29e51c1 100644 --- a/server/src/system_one_lite/mlx_vlm_diffusion.py +++ b/server/src/system_one_lite/mlx_vlm_diffusion.py @@ -236,7 +236,7 @@ def _parse_response( reported_tokens = usage["prompt_tokens"] denoising_steps = usage["denoising_steps"] candidate_only = usage["candidate_only"] - encoder_layers = usage["encoder_layers"] + encoder_layers = usage.get("encoder_layers") model = response["model"] except (KeyError, TypeError) as error: raise MlxVlmError("MLX-VLM returned an invalid diffusion read response") from error diff --git a/server/tests/tools/test_benchmark_diffusion.py b/server/tests/tools/test_benchmark_diffusion.py new file mode 100644 index 0000000..aa66ba4 --- /dev/null +++ b/server/tests/tools/test_benchmark_diffusion.py @@ -0,0 +1,30 @@ +import math + +from tools.benchmark_diffusion import Sample, percentile, run_level + + +class FakeClient: + prompt_tokens = 100 + canvas_tokens = 15 + + def call(self, nonce): + return Sample(client_ms=float(nonce + 1), round_trip_ms=float(nonce) + 0.5) + + +def test_percentile_uses_nearest_rank(): + assert percentile([5, 1, 4, 2, 3], 0.95) == 5 + + +def test_run_level_reports_throughput_and_latency(): + result = run_level(FakeClient(), concurrency=2, request_count=3, nonce_start=0) + + assert result.concurrency == 2 + assert result.requests == 3 + assert result.decisions_per_request == 3 + assert result.prompt_tokens == 100 + assert result.canvas_tokens == 15 + assert math.isclose(result.decisions_per_second, result.requests_per_second * 3) + assert result.client_median_ms == 2.0 + assert result.client_p95_ms == 3.0 + assert result.round_trip_median_ms == 1.5 + assert result.round_trip_p95_ms == 2.5 diff --git a/server/tests/unit/test_mlx_vlm_diffusion.py b/server/tests/unit/test_mlx_vlm_diffusion.py index e877c8f..4159af7 100644 --- a/server/tests/unit/test_mlx_vlm_diffusion.py +++ b/server/tests/unit/test_mlx_vlm_diffusion.py @@ -240,3 +240,31 @@ def test_response_must_confirm_fast_read_settings(): "/pinned/diffusion", expected_encoder_layers=6, ) + + +def test_full_read_accepts_server_without_encoder_layer_metadata(): + engine = bare_engine(None) + + probabilities = engine._parse_response( + { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 1, + "token_ids": [10, 11], + "logprobs": [-1.0, -2.0], + } + ], + "usage": { + "prompt_tokens": 3, + "denoising_steps": 1, + "candidate_only": True, + }, + }, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + ) + + assert len(probabilities) == 1 + assert math.isclose(sum(probabilities[0]), 1.0) diff --git a/server/tools/benchmark_diffusion.py b/server/tools/benchmark_diffusion.py new file mode 100644 index 0000000..c732b65 --- /dev/null +++ b/server/tools/benchmark_diffusion.py @@ -0,0 +1,189 @@ +"""Profile warmed System One DiffusionGemma requests through MLX-VLM. + +Start the MLX-VLM server first. Then run this command from ``server/``: + + uv run python -m tools.benchmark_diffusion --json +""" + +import argparse +import json +import math +import platform +import statistics +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass + +from system_one_lite.mlx_vlm_diffusion import ( + DEFAULT_BASE_URL, + MlxVlmDiffusionEngine, + build_prompt, + build_seed_canvas, +) + +NONCE_START = 100_000 +STATE = "Everything is down and we have a demo at noon." +QUESTIONS = ( + ( + "Does the customer need a reply within the hour?", + ("yes", "no"), + ), + ( + "Which team should handle this request?", + ( + "outage — Service outage or availability issue", + "billing — Charges, refunds, or payment methods", + "feature — Feature request or product suggestion", + ), + ), + ( + "What is the customer's tone?", + ("calm", "annoyed", "furious"), + ), +) + + +@dataclass(frozen=True) +class Sample: + client_ms: float + round_trip_ms: float + + +@dataclass(frozen=True) +class Result: + concurrency: int + requests: int + decisions_per_request: int + prompt_tokens: int + canvas_tokens: int + wall_seconds: float + requests_per_second: float + decisions_per_second: float + client_median_ms: float + client_p95_ms: float + round_trip_median_ms: float + round_trip_p95_ms: float + + +def percentile(values, fraction): + ordered = sorted(values) + return ordered[math.ceil(fraction * len(ordered)) - 1] + + +class DiffusionClient: + def __init__(self, engine): + self.engine = engine + self.questions = tuple((instructions, list(labels)) for instructions, labels in QUESTIONS) + state = self._state(NONCE_START) + self.prompt_tokens = len( + build_prompt(engine.tokenizer, state, self.questions, engine.codes) + ) + canvas, _ = build_seed_canvas( + engine.tokenizer, + len(self.questions), + engine.canvas_length, + engine.vocab_size, + ) + self.canvas_tokens = len(canvas) + + @staticmethod + def _state(nonce): + return json.dumps({"ticket": STATE, "benchmark_nonce": nonce}) + + def call(self, nonce): + started = time.perf_counter() + probabilities, prompt_tokens, round_trip_ms = self.engine.evaluate( + self._state(nonce), self.questions + ) + client_ms = (time.perf_counter() - started) * 1000 + if prompt_tokens != self.prompt_tokens: + raise RuntimeError("prompt token count changed during the benchmark") + if len(probabilities) != len(self.questions): + raise RuntimeError("engine returned the wrong number of decisions") + if any(abs(sum(row) - 1.0) > 1e-5 for row in probabilities): + raise RuntimeError("engine returned an invalid probability distribution") + return Sample(client_ms=client_ms, round_trip_ms=round_trip_ms) + + +def run_level(client, concurrency, request_count, nonce_start): + started = time.perf_counter() + with ThreadPoolExecutor(max_workers=concurrency) as executor: + samples = list(executor.map(client.call, range(nonce_start, nonce_start + request_count))) + wall_seconds = time.perf_counter() - started + request_rate = request_count / wall_seconds + client_times = [sample.client_ms for sample in samples] + round_trip_times = [sample.round_trip_ms for sample in samples] + return Result( + concurrency=concurrency, + requests=request_count, + decisions_per_request=len(QUESTIONS), + prompt_tokens=client.prompt_tokens, + canvas_tokens=client.canvas_tokens, + wall_seconds=wall_seconds, + requests_per_second=request_rate, + decisions_per_second=request_rate * len(QUESTIONS), + client_median_ms=statistics.median(client_times), + client_p95_ms=percentile(client_times, 0.95), + round_trip_median_ms=statistics.median(round_trip_times), + round_trip_p95_ms=percentile(round_trip_times, 0.95), + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", default=DEFAULT_BASE_URL) + parser.add_argument("--timeout", type=float, default=600.0) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--requests", type=int, default=64) + parser.add_argument("--concurrency", default="1,2,4,8,16,32") + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + if args.warmup < 0 or args.requests < 1: + parser.error("warmup must be non-negative and requests must be positive") + concurrency = [int(value) for value in args.concurrency.split(",")] + if not concurrency or any(value < 1 for value in concurrency): + parser.error("concurrency values must be positive") + + engine = MlxVlmDiffusionEngine(base_url=args.base_url) + engine.timeout = args.timeout + client = DiffusionClient(engine) + timed_requests = args.requests * len(concurrency) + for nonce in range(args.warmup): + client.call(NONCE_START + timed_requests + nonce) + + results = [] + nonce = NONCE_START + for level in concurrency: + results.append(run_level(client, level, args.requests, nonce)) + nonce += args.requests + + payload = { + "backend": "mlx-vlm-diffusion", + "model": engine.model_id, + "model_revision": engine.model_revision, + "base_url": engine.base_url, + "host": { + "platform": platform.platform(), + "machine": platform.machine(), + "python": platform.python_version(), + }, + "state": STATE, + "warmup": args.warmup, + "results": [asdict(result) for result in results], + } + if args.json: + print(json.dumps(payload, indent=2)) + return + print(f"model: {engine.model_id}") + print("clients req/s decisions/s client p50 client p95 roundtrip p50 roundtrip p95") + for result in results: + print( + f"{result.concurrency:7} {result.requests_per_second:5.2f} " + f"{result.decisions_per_second:11.2f} " + f"{result.client_median_ms:10.1f} {result.client_p95_ms:10.1f} " + f"{result.round_trip_median_ms:13.1f} {result.round_trip_p95_ms:13.1f}" + ) + + +if __name__ == "__main__": + main() From ff7a5e3c28a9530479b080a25a053f8fbaa3d38f Mon Sep 17 00:00:00 2001 From: Sam Snelling Date: Thu, 17 Sep 2026 18:57:30 -0500 Subject: [PATCH 3/4] Improve structured read profiling Add question batch and compact profile controls. Keep the compact path exact on the current server. Use fixed prompts so every timed request has the same shape. --- docs/diffusion-gemma.md | 12 ++-- .../src/system_one_lite/mlx_vlm_diffusion.py | 12 ---- .../tests/tools/test_benchmark_diffusion.py | 46 +++++++++++- server/tests/unit/test_mlx_vlm_diffusion.py | 44 +----------- server/tools/benchmark_diffusion.py | 70 +++++++++++++------ 5 files changed, 100 insertions(+), 84 deletions(-) diff --git a/docs/diffusion-gemma.md b/docs/diffusion-gemma.md index 774e2db..c805677 100644 --- a/docs/diffusion-gemma.md +++ b/docs/diffusion-gemma.md @@ -35,8 +35,8 @@ SYSTEM_ONE_BACKEND=mlx-vlm-diffusion-fast \ The compact profile accepts exactly one question. It removes the long prompt instructions and uses one active canvas token. Use it only when the state, question, and answer labels make the task clear without extra guidance. -It runs six full encoder layers. Later layers build lighter attention caches. -This approximation needs an accuracy check for each target task. +It still runs the full model. Its speed comes only from the shorter prompt and +canvas, so it does not use an approximate encoder path. The first System One startup downloads the pinned 4-bit model snapshot. The weights are about 16.5 GB. Both processes use the same Hugging Face cache. @@ -64,6 +64,8 @@ Start the MLX-VLM server, then run the warmed concurrency sweep from uv run python -m tools.benchmark_diffusion ``` -Use `--json` to save machine-readable results. Set `--requests` and -`--concurrency` to change the load. The report includes request throughput, -decision throughput, client latency, and request round-trip latency. +Use `--json` to save machine-readable results. Set `--requests`, +`--concurrency`, and `--questions` to change the load. Use +`--profile fast --questions 1` to measure the compact profile. The report +includes request throughput, decision throughput, client latency, and request +round-trip latency. diff --git a/server/src/system_one_lite/mlx_vlm_diffusion.py b/server/src/system_one_lite/mlx_vlm_diffusion.py index 29e51c1..50fc1be 100644 --- a/server/src/system_one_lite/mlx_vlm_diffusion.py +++ b/server/src/system_one_lite/mlx_vlm_diffusion.py @@ -29,7 +29,6 @@ CANVAS_PREFIX = "<|channel>thought\n" CANVAS_SUFFIX = "\n" DEFAULT_DIFFUSION_SEED = 42 -FAST_DIFFUSION_ENCODER_LAYERS = 6 class MlxVlmError(RuntimeError): @@ -195,9 +194,6 @@ def evaluate(self, state, questions, template=None): "slots": slots, "candidate_only": True, } - if self.compact: - payload["encoder_layers"] = FAST_DIFFUSION_ENCODER_LAYERS - started = time.perf_counter() response = self.transport("/v1/diffusion/reads", payload) elapsed_ms = (time.perf_counter() - started) * 1000 @@ -206,7 +202,6 @@ def evaluate(self, state, questions, template=None): slots, len(input_ids), str(self.model_path), - expected_encoder_layers=(FAST_DIFFUSION_ENCODER_LAYERS if self.compact else None), ) return probabilities, len(input_ids), elapsed_ms @@ -228,7 +223,6 @@ def _parse_response( slots, input_tokens, expected_model, - expected_encoder_layers=None, ): try: reads = response["reads"] @@ -236,7 +230,6 @@ def _parse_response( reported_tokens = usage["prompt_tokens"] denoising_steps = usage["denoising_steps"] candidate_only = usage["candidate_only"] - encoder_layers = usage.get("encoder_layers") model = response["model"] except (KeyError, TypeError) as error: raise MlxVlmError("MLX-VLM returned an invalid diffusion read response") from error @@ -246,11 +239,6 @@ def _parse_response( raise MlxVlmError(f"MLX-VLM performed {denoising_steps!r} denoising steps, expected 1") if candidate_only is not True: raise MlxVlmError("MLX-VLM did not confirm candidate-only scoring") - if encoder_layers != expected_encoder_layers: - raise MlxVlmError( - "MLX-VLM used a different encoder layer count: " - f"expected {expected_encoder_layers!r}, got {encoder_layers!r}" - ) if reported_tokens != input_tokens: raise MlxVlmError( "MLX-VLM token count differs from the pinned tokenizer: " diff --git a/server/tests/tools/test_benchmark_diffusion.py b/server/tests/tools/test_benchmark_diffusion.py index aa66ba4..c92764a 100644 --- a/server/tests/tools/test_benchmark_diffusion.py +++ b/server/tests/tools/test_benchmark_diffusion.py @@ -1,22 +1,54 @@ +import json import math -from tools.benchmark_diffusion import Sample, percentile, run_level +from tools.benchmark_diffusion import DiffusionClient, Sample, percentile, run_level class FakeClient: prompt_tokens = 100 canvas_tokens = 15 + questions = (("first", ["yes", "no"]),) * 3 def call(self, nonce): return Sample(client_ms=float(nonce + 1), round_trip_ms=float(nonce) + 0.5) +class FakeTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return list(text.encode()) + + def apply_chat_template(self, messages, **kwargs): + del kwargs + return messages[0]["content"] + + +class FakeEngine: + tokenizer = FakeTokenizer() + codes = ("A", "B") + canvas_length = 256 + vocab_size = 200 + compact = True + + def test_percentile_uses_nearest_rank(): assert percentile([5, 1, 4, 2, 3], 0.95) == 5 +def test_state_is_fixed_between_requests(): + assert json.loads(DiffusionClient._state()) == { + "ticket": "Everything is down and we have a demo at noon." + } + + +def test_client_metadata_uses_the_selected_profile(): + client = DiffusionClient(FakeEngine(), (("first", ("yes", "no")),)) + + assert client.canvas_tokens == 1 + + def test_run_level_reports_throughput_and_latency(): - result = run_level(FakeClient(), concurrency=2, request_count=3, nonce_start=0) + result = run_level(FakeClient(), concurrency=2, request_count=3, request_index_start=0) assert result.concurrency == 2 assert result.requests == 3 @@ -28,3 +60,13 @@ def test_run_level_reports_throughput_and_latency(): assert result.client_p95_ms == 3.0 assert result.round_trip_median_ms == 1.5 assert result.round_trip_p95_ms == 2.5 + + +def test_run_level_uses_the_client_question_count(): + client = FakeClient() + client.questions = (("first", ["yes", "no"]), ("second", ["yes", "no"])) + + result = run_level(client, concurrency=1, request_count=1, request_index_start=0) + + assert result.decisions_per_request == 2 + assert math.isclose(result.decisions_per_second, result.requests_per_second * 2) diff --git a/server/tests/unit/test_mlx_vlm_diffusion.py b/server/tests/unit/test_mlx_vlm_diffusion.py index 4159af7..bf43a29 100644 --- a/server/tests/unit/test_mlx_vlm_diffusion.py +++ b/server/tests/unit/test_mlx_vlm_diffusion.py @@ -55,7 +55,6 @@ def transport(path, payload): "prompt_tokens": len(payload["input_ids"]), "denoising_steps": 1, "candidate_only": payload["candidate_only"], - "encoder_layers": payload.get("encoder_layers"), }, } @@ -115,7 +114,6 @@ def transport(path, payload): "prompt_tokens": len(payload["input_ids"]), "denoising_steps": 1, "candidate_only": payload["candidate_only"], - "encoder_layers": payload.get("encoder_layers"), }, } @@ -125,7 +123,7 @@ def transport(path, payload): assert captured["payload"]["seed_canvas"] == [163] assert captured["payload"]["slots"] == [{"position": 0, "token_ids": [10, 11]}] - assert captured["payload"]["encoder_layers"] == 6 + assert "encoder_layers" not in captured["payload"] assert probabilities[0][1] > probabilities[0][0] @@ -157,7 +155,6 @@ def test_response_must_match_requested_slot(): "prompt_tokens": 3, "denoising_steps": 1, "candidate_only": True, - "encoder_layers": None, }, }, [{"position": 1, "token_ids": [10, 11]}], @@ -181,7 +178,6 @@ def test_response_must_confirm_model_and_one_step(): "prompt_tokens": 3, "denoising_steps": 1, "candidate_only": True, - "encoder_layers": None, }, } @@ -204,44 +200,6 @@ def test_response_must_confirm_model_and_one_step(): ) -def test_response_must_confirm_fast_read_settings(): - engine = bare_engine(None) - base = { - "model": "/pinned/diffusion", - "reads": [ - { - "position": 1, - "token_ids": [10, 11], - "logprobs": [-1.0, -2.0], - } - ], - "usage": { - "prompt_tokens": 3, - "denoising_steps": 1, - "candidate_only": True, - "encoder_layers": 6, - }, - } - - with pytest.raises(MlxVlmError, match="invalid diffusion read response"): - engine._parse_response( - {**base, "usage": {"prompt_tokens": 3, "denoising_steps": 1}}, - [{"position": 1, "token_ids": [10, 11]}], - 3, - "/pinned/diffusion", - expected_encoder_layers=6, - ) - - with pytest.raises(MlxVlmError, match="different encoder layer count"): - engine._parse_response( - {**base, "usage": {**base["usage"], "encoder_layers": 30}}, - [{"position": 1, "token_ids": [10, 11]}], - 3, - "/pinned/diffusion", - expected_encoder_layers=6, - ) - - def test_full_read_accepts_server_without_encoder_layer_metadata(): engine = bare_engine(None) diff --git a/server/tools/benchmark_diffusion.py b/server/tools/benchmark_diffusion.py index c732b65..95d63a0 100644 --- a/server/tools/benchmark_diffusion.py +++ b/server/tools/benchmark_diffusion.py @@ -21,7 +21,7 @@ build_seed_canvas, ) -NONCE_START = 100_000 +REQUEST_INDEX_START = 0 STATE = "Everything is down and we have a demo at noon." QUESTIONS = ( ( @@ -71,29 +71,36 @@ def percentile(values, fraction): class DiffusionClient: - def __init__(self, engine): + def __init__(self, engine, questions=QUESTIONS): self.engine = engine - self.questions = tuple((instructions, list(labels)) for instructions, labels in QUESTIONS) - state = self._state(NONCE_START) + self.questions = tuple((instructions, list(labels)) for instructions, labels in questions) + state = self._state() self.prompt_tokens = len( - build_prompt(engine.tokenizer, state, self.questions, engine.codes) + build_prompt( + engine.tokenizer, + state, + self.questions, + engine.codes, + compact=engine.compact, + ) ) canvas, _ = build_seed_canvas( engine.tokenizer, len(self.questions), engine.canvas_length, engine.vocab_size, + compact=engine.compact, ) self.canvas_tokens = len(canvas) @staticmethod - def _state(nonce): - return json.dumps({"ticket": STATE, "benchmark_nonce": nonce}) + def _state(): + return json.dumps({"ticket": STATE}) - def call(self, nonce): + def call(self, _request_index): started = time.perf_counter() probabilities, prompt_tokens, round_trip_ms = self.engine.evaluate( - self._state(nonce), self.questions + self._state(), self.questions ) client_ms = (time.perf_counter() - started) * 1000 if prompt_tokens != self.prompt_tokens: @@ -105,10 +112,15 @@ def call(self, nonce): return Sample(client_ms=client_ms, round_trip_ms=round_trip_ms) -def run_level(client, concurrency, request_count, nonce_start): +def run_level(client, concurrency, request_count, request_index_start): started = time.perf_counter() with ThreadPoolExecutor(max_workers=concurrency) as executor: - samples = list(executor.map(client.call, range(nonce_start, nonce_start + request_count))) + samples = list( + executor.map( + client.call, + range(request_index_start, request_index_start + request_count), + ) + ) wall_seconds = time.perf_counter() - started request_rate = request_count / wall_seconds client_times = [sample.client_ms for sample in samples] @@ -116,12 +128,12 @@ def run_level(client, concurrency, request_count, nonce_start): return Result( concurrency=concurrency, requests=request_count, - decisions_per_request=len(QUESTIONS), + decisions_per_request=len(client.questions), prompt_tokens=client.prompt_tokens, canvas_tokens=client.canvas_tokens, wall_seconds=wall_seconds, requests_per_second=request_rate, - decisions_per_second=request_rate * len(QUESTIONS), + decisions_per_second=request_rate * len(client.questions), client_median_ms=statistics.median(client_times), client_p95_ms=percentile(client_times, 0.95), round_trip_median_ms=statistics.median(round_trip_times), @@ -136,6 +148,8 @@ def main(): parser.add_argument("--warmup", type=int, default=2) parser.add_argument("--requests", type=int, default=64) parser.add_argument("--concurrency", default="1,2,4,8,16,32") + parser.add_argument("--questions", type=int, default=len(QUESTIONS)) + parser.add_argument("--profile", choices=("full", "fast"), default="full") parser.add_argument("--json", action="store_true") args = parser.parse_args() if args.warmup < 0 or args.requests < 1: @@ -143,22 +157,30 @@ def main(): concurrency = [int(value) for value in args.concurrency.split(",")] if not concurrency or any(value < 1 for value in concurrency): parser.error("concurrency values must be positive") - - engine = MlxVlmDiffusionEngine(base_url=args.base_url) + if args.questions < 1: + parser.error("questions must be positive") + if args.profile == "fast" and args.questions != 1: + parser.error("the fast profile requires exactly one question") + + engine = MlxVlmDiffusionEngine( + base_url=args.base_url, + compact=args.profile == "fast", + ) engine.timeout = args.timeout - client = DiffusionClient(engine) + questions = tuple(QUESTIONS[index % len(QUESTIONS)] for index in range(args.questions)) + client = DiffusionClient(engine, questions) timed_requests = args.requests * len(concurrency) - for nonce in range(args.warmup): - client.call(NONCE_START + timed_requests + nonce) + for request_index in range(args.warmup): + client.call(REQUEST_INDEX_START + timed_requests + request_index) results = [] - nonce = NONCE_START + request_index = REQUEST_INDEX_START for level in concurrency: - results.append(run_level(client, level, args.requests, nonce)) - nonce += args.requests + results.append(run_level(client, level, args.requests, request_index)) + request_index += args.requests payload = { - "backend": "mlx-vlm-diffusion", + "backend": ("mlx-vlm-diffusion-fast" if args.profile == "fast" else "mlx-vlm-diffusion"), "model": engine.model_id, "model_revision": engine.model_revision, "base_url": engine.base_url, @@ -175,6 +197,10 @@ def main(): print(json.dumps(payload, indent=2)) return print(f"model: {engine.model_id}") + print( + f"profile: {args.profile}; decisions/request: {len(client.questions)}; " + f"prompt tokens: {client.prompt_tokens}; canvas tokens: {client.canvas_tokens}" + ) print("clients req/s decisions/s client p50 client p95 roundtrip p50 roundtrip p95") for result in results: print( From fdad9ff20532b34ff6592758392a7ecd9caf0f4f Mon Sep 17 00:00:00 2001 From: Sam Snelling Date: Thu, 17 Sep 2026 18:57:49 -0500 Subject: [PATCH 4/4] Speed up structured diffusion reads --- docs/diffusion-gemma.md | 18 ++++----- .../src/system_one_lite/mlx_vlm_diffusion.py | 37 ++++-------------- server/tests/unit/test_mlx_vlm_diffusion.py | 39 +++++++++++++++---- server/tools/benchmark_diffusion.py | 2 - 4 files changed, 46 insertions(+), 50 deletions(-) diff --git a/docs/diffusion-gemma.md b/docs/diffusion-gemma.md index c805677..6344294 100644 --- a/docs/diffusion-gemma.md +++ b/docs/diffusion-gemma.md @@ -1,13 +1,12 @@ # Run with 4-bit DiffusionGemma This backend uses `mlx-community/diffusiongemma-26B-A4B-it-4bit` at the pinned -revision in `engine.py`. It sends every question in one prompt and places one -`- label` answer slot per question in DiffusionGemma's active canvas. This shape -fits System One's limit of 64 questions. One read-only model pass returns the -requested answer-code scores. The canvas uses the same fixed seed for each -request, so the same input has the same canvas. +revision in `engine.py`. It sends every question in one prompt. It places one +active canvas token per question. This fits System One's limit of 64 questions. +One read-only model pass returns the requested answer-code scores. The canvas +uses the same fixed seed for each request. -The backend needs the `feature/diffusion-gemma-reads` branch of MLX-VLM. That +The backend needs the `feature/structured-reads` branch of MLX-VLM. That branch adds `/v1/diffusion/reads` without changing normal generation. Start the MLX-VLM server from its checkout: @@ -33,10 +32,9 @@ SYSTEM_ONE_BACKEND=mlx-vlm-diffusion-fast \ ``` The compact profile accepts exactly one question. It removes the long prompt -instructions and uses one active canvas token. Use it only when the state, -question, and answer labels make the task clear without extra guidance. -It still runs the full model. Its speed comes only from the shorter prompt and -canvas, so it does not use an approximate encoder path. +instructions. Use it only when the state, question, and answer labels make the +task clear without extra guidance. It still runs the full model. Its speed +comes only from the shorter prompt. The first System One startup downloads the pinned 4-bit model snapshot. The weights are about 16.5 GB. Both processes use the same Hugging Face cache. diff --git a/server/src/system_one_lite/mlx_vlm_diffusion.py b/server/src/system_one_lite/mlx_vlm_diffusion.py index 50fc1be..ff8517e 100644 --- a/server/src/system_one_lite/mlx_vlm_diffusion.py +++ b/server/src/system_one_lite/mlx_vlm_diffusion.py @@ -26,8 +26,6 @@ DEFAULT_BASE_URL = "http://127.0.0.1:8080" DEFAULT_TIMEOUT = 120.0 -CANVAS_PREFIX = "<|channel>thought\n" -CANVAS_SUFFIX = "\n" DEFAULT_DIFFUSION_SEED = 42 @@ -72,7 +70,7 @@ def build_prompt(tokenizer, state, questions, codes, compact=False): for index, (instructions, labels) in enumerate(questions, 1): parts.append(f"\n\nQuestion q{index}: {instructions}\n") parts.append("\n".join(f" {codes[i]}: {item}" for i, item in enumerate(labels))) - parts.append('\n\nReply with one line per question, in order, formatted as "- label".') + parts.append("\n\nReply with one answer code per question, in order.") text = tokenizer.apply_chat_template( [{"role": "user", "content": "".join(parts)}], tokenize=False, @@ -83,38 +81,19 @@ def build_prompt(tokenizer, state, questions, codes, compact=False): def build_seed_canvas( - tokenizer, question_count, canvas_length, vocab_size, - compact=False, ): - """Build the active answer template and randomize only its answer slots.""" - rng = random.Random(DEFAULT_DIFFUSION_SEED) - if compact: - if question_count > canvas_length: - raise RequestContractError( - f"{question_count} questions need {question_count} canvas tokens; " - f"DiffusionGemma provides {canvas_length}" - ) - return [rng.randrange(vocab_size) for _ in range(question_count)], list( - range(question_count) - ) - - canvas = _encode(tokenizer, CANVAS_PREFIX) - positions = [] - for _ in range(question_count): - canvas.extend(_encode(tokenizer, "-")) - positions.append(len(canvas)) - canvas.append(rng.randrange(vocab_size)) - canvas.extend(_encode(tokenizer, "\n")) - canvas.extend(_encode(tokenizer, CANVAS_SUFFIX)) - if len(canvas) > canvas_length: + """Build one deterministic active canvas token per question.""" + if question_count > canvas_length: raise RequestContractError( - f"{question_count} questions need {len(canvas)} canvas tokens; " + f"{question_count} questions need {question_count} canvas tokens; " f"DiffusionGemma provides {canvas_length}" ) - return canvas, positions + rng = random.Random(DEFAULT_DIFFUSION_SEED) + canvas = [rng.randrange(vocab_size) for _ in range(question_count)] + return canvas, list(range(question_count)) class MlxVlmDiffusionEngine: @@ -174,11 +153,9 @@ def evaluate(self, state, questions, template=None): f"request exceeds the {MAX_TOTAL_INPUT_TOKENS} input token limit" ) seed_canvas, positions = build_seed_canvas( - self.tokenizer, len(questions), self.canvas_length, self.vocab_size, - compact=self.compact, ) slots = [ { diff --git a/server/tests/unit/test_mlx_vlm_diffusion.py b/server/tests/unit/test_mlx_vlm_diffusion.py index bf43a29..8c71a13 100644 --- a/server/tests/unit/test_mlx_vlm_diffusion.py +++ b/server/tests/unit/test_mlx_vlm_diffusion.py @@ -7,6 +7,7 @@ from system_one_lite.mlx_vlm_diffusion import ( MlxVlmDiffusionEngine, MlxVlmError, + build_prompt, build_seed_canvas, ) @@ -21,6 +22,16 @@ def apply_chat_template(self, messages, **kwargs): return "CHAT:" + messages[0]["content"] +class TextTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return text + + def apply_chat_template(self, messages, **kwargs): + del kwargs + return messages[0]["content"] + + def bare_engine(transport, canvas_length=256): engine = MlxVlmDiffusionEngine.__new__(MlxVlmDiffusionEngine) engine.model_id = "diffusion" @@ -35,6 +46,18 @@ def bare_engine(transport, canvas_length=256): return engine +def test_full_prompt_asks_for_answer_codes(): + prompt = build_prompt( + TextTokenizer(), + "state", + [("pick", ["yes", "no"])], + ("A", "B"), + ) + + assert "one answer code per question" in prompt + assert '"- label"' not in prompt + + def test_evaluate_reads_all_questions_from_one_seeded_canvas(): captured = {} @@ -66,7 +89,9 @@ def transport(path, payload): assert captured["path"] == "/v1/diffusion/reads" assert captured["payload"]["model"] == "/pinned/diffusion" assert len(captured["payload"]["input_ids"]) == input_tokens - assert len(captured["payload"]["seed_canvas"]) < 256 + assert len(captured["payload"]["seed_canvas"]) == 2 + assert [slot["position"] for slot in captured["payload"]["slots"]] == [0, 1] + assert "encoder_layers" not in captured["payload"] assert captured["payload"]["candidate_only"] is True assert [slot["token_ids"] for slot in captured["payload"]["slots"]] == [ [10, 11], @@ -80,19 +105,17 @@ def transport(path, payload): def test_seed_canvas_rejects_more_questions_than_fit(): with pytest.raises(RequestContractError, match="canvas tokens"): - build_seed_canvas(FakeTokenizer(), 2, canvas_length=5, vocab_size=200) + build_seed_canvas(6, canvas_length=5, vocab_size=200) def test_seed_canvas_fits_the_public_64_question_limit(): - first, positions = build_seed_canvas(FakeTokenizer(), 64, canvas_length=256, vocab_size=200) - second, second_positions = build_seed_canvas( - FakeTokenizer(), 64, canvas_length=256, vocab_size=200 - ) + first, positions = build_seed_canvas(64, canvas_length=256, vocab_size=200) + second, second_positions = build_seed_canvas(64, canvas_length=256, vocab_size=200) assert first == second assert positions == second_positions - assert len(first) <= 256 - assert len(positions) == 64 + assert len(first) == 64 + assert positions == list(range(64)) def test_compact_read_uses_one_canvas_token_per_question(): diff --git a/server/tools/benchmark_diffusion.py b/server/tools/benchmark_diffusion.py index 95d63a0..320abb8 100644 --- a/server/tools/benchmark_diffusion.py +++ b/server/tools/benchmark_diffusion.py @@ -85,11 +85,9 @@ def __init__(self, engine, questions=QUESTIONS): ) ) canvas, _ = build_seed_canvas( - engine.tokenizer, len(self.questions), engine.canvas_length, engine.vocab_size, - compact=engine.compact, ) self.canvas_tokens = len(canvas)