| 1 | # Structured reads on DiffusionGemma |
| 2 | |
| 3 | A discrete diffusion model denoises a whole canvas per forward pass. If the |
| 4 | canvas is seeded with the answer's fixed text and only the answer slots are |
| 5 | left as noise, one denoise step gives a distribution over each slot. These |
| 6 | `extra_args` fields (`vllm_xargs` on the OpenAI server) expose that: |
| 7 | |
| 8 | | field | type | meaning | |
| 9 | | --- | --- | --- | |
| 10 | | `diffusion_seed_canvas` | `list[int]`, exactly `canvas_length` ids | replaces the random initial canvas after prefill | |
| 11 | | `diffusion_pinned` | `list[int]` of canvas positions | held at their seed value on every denoise step, so a read past one step keeps its template | |
| 12 | | `diffusion_max_steps` | `int` | denoise steps before the canvas is emitted | |
| 13 | | `diffusion_read_only` | `bool` | emit the argmax canvas as soon as the cap is reached, end the request there, and return temperature-1 logprobs at every position | |
| 14 | | `diffusion_constrained` | `bool` | run the unembedding, sampler and self-conditioning over the request's `logprob_token_ids` only. Logprobs are normalized over that set. A step uses this only when every read in it has the same set | |
| 15 | |
| 16 | `structured_server.py` turns a question schema into those fields, with |
| 17 | `diffusion_constrained` on for every read (`--no-constrained` turns it off). It serves |
| 18 | `/v1/chat/completions`: the system message is the schema, the user message |
| 19 | is the state JSON, and the reply content is one distribution per question |
| 20 | with a standard error over a few noise draws. |
| 21 | |
| 22 | ```bash |
| 23 | vllm serve google/diffusiongemma-26B-A4B-it \ |
| 24 | --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 --enable-prefix-caching |
| 25 | python examples/features/structured_diffusion/structured_server.py \ |
| 26 | --upstream http://127.0.0.1:8000 --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64 |
| 27 | curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -d '{ |
| 28 | "messages": [ |
| 29 | {"role": "system", "content": "{\"questions\": [{\"id\": \"urgent\", \"type\": \"noul\", \"instructions\": \"Does the customer need a reply within the hour?\"}]}"}, |
| 30 | {"role": "user", "content": "{\"ticket\": \"Everything is down and we have a demo at noon.\"}"} |
| 31 | ]}' |
| 32 | ``` |
| 33 | |
| 34 | The attention backend is picked as for Gemma 4: FlashAttention 4 on every |
| 35 | layer when available, otherwise Triton. FlashInfer cannot serve this model (a |
| 36 | batch mixes causal prefill with bidirectional denoising), and |
| 37 | `--attention-backend FLASHINFER` is rejected. |
| 38 | |
| 39 | Per-request canvas widths may be smaller than the served canvas with either |
| 40 | synchronous or asynchronous scheduling. Omit `diffusion_canvas_length` to use |
| 41 | the served canvas width. |
| 42 | |
| 43 | Question types: `noul` (yes/no), `choice` with `options`, `score` with |
| 44 | ordered `levels`. Each label must be a single token in the answer template, |
| 45 | which the server checks with the tokenizer when a request uses the schema. |
| 46 | |
| 47 | `POST /v1/systemone` implements the Jev decision API. The body holds |
| 48 | `state`, `questions` (a map of id to `type`, `instructions` and `criteria`) |
| 49 | and `model`. Answers come back in that API's shapes: a `noul` probability, |
| 50 | a `choice` with `probabilities` and `confidence`, or a `score` with a |
| 51 | 0-indexed `legend`. The schema options above go in the same body as |
| 52 | extensions. |
| 53 | |
| 54 | A question may declare `depends_on` (read in a later stage with those |
| 55 | answers in its prompt), `ask_if` (asked only when a named question's answer |
| 56 | is among the listed ones, otherwise null) and `alone` (a read of its own). |
| 57 | |
| 58 | Images attach as `multipart/form-data`, with the JSON in a part named |
| 59 | `request` and each image as a file part, or as an `images` array of data |
| 60 | URLs. |
| 61 | |
| 62 | ```bash |
| 63 | curl -s localhost:8011/v1/systemone -H 'content-type: application/json' -d '{ |
| 64 | "model": "jev-latest", |
| 65 | "state": {"ticket": "Everything is down and we have a demo at noon."}, |
| 66 | "questions": {"urgent": {"type": "noul", "instructions": "Does the customer need a reply within the hour?"}}}' |
| 67 | ``` |
| 68 | |
| 69 | `"think": N` in the schema lets the model write up to N tokens in its |
| 70 | thought channel before the read. The thought is an ordinary generation with |
| 71 | the chat template's thinking marker on, and the read then runs with the |
| 72 | thought in its prompt, so the answer slots condition on it. The noise draws |
| 73 | of a decision share one thought. `diagnostics.thought` returns the text, its |
| 74 | length in tokens, whether the model closed the channel itself and the |
| 75 | generation time. |