Repository navigation
[Core] structured generation mode for DiffusionGemma model (Jev-like) #57250
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
LucasWilkinson
merged 26 commits into
vllm-project:main
from
mmastrac:structured-reads-main
Sep 22, 2026
+2,797
−157
Merged
Changes from all commits
Commits
Show all changes
26 commits
Select commit
Hold shift + click to select a range
cf9bfb9
[Model] DiffusionGemma: seeded canvas, step cap and read-only request…
mmastrac cf27f5d
[Core] Validate diffusion extra_args when the request is created
mmastrac 15eb9be
[Core] Read-only diffusion requests end on the canvas they emit
mmastrac a2908b2
[Perf] Keep seed and read-only bookkeeping off the diffusion hot path
mmastrac fa81b63
[Core] Read-only diffusion requests report temperature-1 logprobs
mmastrac 21562cc
[Misc] Keep the diffusion step cap in the step counter's dtype
mmastrac da4d83c
[Misc] Drop the redundant commit term from the read-only mask
mmastrac 5cc9b4f
[Misc] Tighten the comments in the structured-read changes
mmastrac 0c64f9f
[Model] DiffusionGemma: skip self-conditioning for one-step tiles
mmastrac 8dc04ec
[Core] Do not schedule a step past a read-only diffusion request's cap
mmastrac 5b6f205
[Core] Per-request canvas width for diffusion reads
mmastrac 9e9a23f
[Model] DiffusionGemma: tile the sampler by canvas width
mmastrac f2d852a
[Misc] Format the diffusion read changes for the pinned ruff
mmastrac 9b85cb4
[Docs] Example structured-decision server for DiffusionGemma reads
mmastrac 408f00d
[Core] Pin seeded canvas positions through diffusion denoise steps
mmastrac 55f03bb
Keep diffusion canvas widths in the scheduler
LucasWilkinson 2449995
Require async scheduling for narrower diffusion canvases
LucasWilkinson 084444a
[Docs] Drop the HTML pages from the diffusion reads example
mmastrac 340eca4
[Docs] Rename the example directory to structured_diffusion
mmastrac a68dab4
[Bugfix] DiffusionGemma: keep the sampler step compiled, and correct …
razorback16 d9af863
[Bugfix] DiffusionGemma: join logprob stashes of different widths
razorback16 58a1993
[Docs] Bound the example server's fan-out
mmastrac dd86395
Isolate diffusion request validation and simplify read emission
LucasWilkinson 7ba6675
[Core] Validate diffusion extra_args whenever the model is a diffusio…
mmastrac f792b82
Merge branch 'main' into structured-reads-main
mmastrac a7c23ac
[Bugfix] DiffusionGemma: register the top-level model type with the G…
mmastrac File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,74 @@ | ||
| # Structured reads on DiffusionGemma | ||
|
|
||
| A discrete diffusion model denoises a whole canvas per forward pass. If the | ||
| canvas is seeded with the answer's fixed text and only the answer slots are | ||
| left as noise, one denoise step gives a distribution over each slot. These | ||
| `extra_args` fields (`vllm_xargs` on the OpenAI server) expose that: | ||
|
|
||
| | field | type | meaning | | ||
| | --- | --- | --- | | ||
| | `diffusion_seed_canvas` | `list[int]`, exactly `canvas_length` ids | replaces the random initial canvas after prefill | | ||
| | `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 | | ||
| | `diffusion_max_steps` | `int` | denoise steps before the canvas is emitted | | ||
| | `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 | | ||
|
|
||
| `structured_server.py` turns a question schema into those fields. It serves | ||
| `/v1/chat/completions`: the system message is the schema, the user message | ||
| is the state JSON, and the reply content is one distribution per question | ||
| with a standard error over a few noise draws. | ||
|
|
||
| ```bash | ||
| vllm serve google/diffusiongemma-26B-A4B-it \ | ||
| --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 --enable-prefix-caching | ||
| python examples/features/structured_diffusion/structured_server.py \ | ||
| --upstream http://127.0.0.1:8000 --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64 | ||
| curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -d '{ | ||
| "messages": [ | ||
| {"role": "system", "content": "{\"questions\": [{\"id\": \"urgent\", \"type\": \"noul\", \"instructions\": \"Does the customer need a reply within the hour?\"}]}"}, | ||
| {"role": "user", "content": "{\"ticket\": \"Everything is down and we have a demo at noon.\"}"} | ||
| ]}' | ||
| ``` | ||
|
|
||
| The attention backend is picked as for Gemma 4: FlashAttention 4 on every | ||
| layer when available, otherwise Triton. FlashInfer cannot serve this model (a | ||
| batch mixes causal prefill with bidirectional denoising), and | ||
| `--attention-backend FLASHINFER` is rejected. | ||
|
|
||
| Per-request canvas widths smaller than the served canvas require async | ||
| scheduling. The diffusion async scheduler is selected automatically; no | ||
| `--scheduler-cls` argument is needed. Synchronous execution supports full-width | ||
| canvases only. | ||
|
|
||
| Question types: `noul` (yes/no), `choice` with `options`, `score` with | ||
| ordered `levels`. Each label must be a single token in the answer template, | ||
| which the server checks with the tokenizer when a request uses the schema. | ||
|
|
||
| `POST /v1/systemone` implements the Jev decision API. The body holds | ||
| `state`, `questions` (a map of id to `type`, `instructions` and `criteria`) | ||
| and `model`. Answers come back in that API's shapes: a `noul` probability, | ||
| a `choice` with `probabilities` and `confidence`, or a `score` with a | ||
| 0-indexed `legend`. The schema options above go in the same body as | ||
| extensions. | ||
|
|
||
| A question may declare `depends_on` (read in a later stage with those | ||
| answers in its prompt), `ask_if` (asked only when a named question's answer | ||
| is among the listed ones, otherwise null) and `alone` (a read of its own). | ||
|
|
||
| Images attach as `multipart/form-data`, with the JSON in a part named | ||
| `request` and each image as a file part, or as an `images` array of data | ||
| URLs. | ||
|
|
||
| ```bash | ||
| curl -s localhost:8011/v1/systemone -H 'content-type: application/json' -d '{ | ||
| "model": "jev-latest", | ||
| "state": {"ticket": "Everything is down and we have a demo at noon."}, | ||
| "questions": {"urgent": {"type": "noul", "instructions": "Does the customer need a reply within the hour?"}}}' | ||
| ``` | ||
|
|
||
| `"think": N` in the schema lets the model write up to N tokens in its | ||
| thought channel before the read. The thought is an ordinary generation with | ||
| the chat template's thinking marker on, and the read then runs with the | ||
| thought in its prompt, so the answer slots condition on it. The noise draws | ||
| of a decision share one thought. `diagnostics.thought` returns the text, its | ||
| length in tokens, whether the model closed the channel itself and the | ||
| generation time. |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.