[Core] structured generation mode for DiffusionGemma model (Jev-like)

[mmastrac]

Prerequisite PRs

A few PRs that can go in on their own merits were split out:

• [Bugfix] DiffusionGemma: hand out stashed logprobs only on the committing step #57414: hand out stashed logprobs only on the committing step (bug fix)
• [Perf] Give a prefill-only batch the model state's number of logit rows #57416: give a prefill-only batch the model state's number of logit rows (perf)
• [Model] DiffusionGemma: honor logprob_token_ids on the converging step #57417: honor logprob_token_ids on the converging step (feature parity w/autoregressive)
• [Bugfix] DiffusionGemma: cast the self-conditioning soft embed to the buffer dtype #57462: type cast error if dynamo falls back to eager mid-run (bugfix)
• [Bugfix] DiffusionGemma: fix multimodal support #57589: fix multimodal support in DiffusionGemma (bugfix)

Purpose

DiffusionGemma (and likely other models in vLLM) contains all of the necessary components to build a Jev-like model that provides structured answers with bounded choices.

By fixing canvas positions in the given structured output format, you can ask the model to provide a token and logprobs (from which a client derives entropy) that can be used to determine the models most likely answer and its confidence in that answer. If entropy does not pass a given threshold, you can perform additional sampling to determine agreement between runs and therefore the actual confidence of the model's answer.

This PR contains a prototype example server for serving the /v1/systemone endpoint. It does not implement a standard endpoint in vLLM.

The model is sufficiently powerful to work with any structured output that follows one limitation: the token choices must be a single token, so that the canvas itself does not shift. The requesting client can always translate multi-token options into single token options, however (ie: "moderation_spam" -> "A"). The choices are always laid out in a system prompt and the model correctly diffuses single-token-slot answers (or makes a guess and flags it with uncertainty).

Why do we require the token IDs in the request? Currently we provide the token slot positions that are unfixed for the model to choose and this allows the two to be matched more accurately. A followup PR will likely add a template syntax.

The model, using this framework, can answer yes/no ("noul" in Jev parlance), scale, or multiple choice questions, and provides its certainty and error bars with the answer.

The only missing pieces for this to work in vLLM today are in this PR:

• A way to pre-seed the canvas with the fixed tokens
• A way to send read-only requests that emit the argmax canvas on the converging step and skip the commit forward (logprobs come from that step at temperature 1)
• A way to limit the number of steps run (we only want the first step)
• A bugfix for two parallel DiffusionGemma generations that stomp on each other in lockstep: two concurrent requests with logprobs on where one converges in the same step another commits loses its stash and the request fails with HTTP 500 "list index out of range".

This PR currently contains the machinery to make this model work and a sample interposer server is included to demonstrate end-to-end functionality that is sufficient to serve traffic at reasonable scale.

field	type	validation	effect
diffusion_seed_canvas	list[int], exactly the request's canvas width	ids in [0, vocab), length checked against diffusion_canvas_length, else the served --diffusion-config	replaces the random canvas after prefill
diffusion_max_steps	int >= 1	clamped to max_denoising_steps	denoise steps before the canvas is emitted
diffusion_pinned	list[int] of canvas positions	needs a diffusion_seed_canvas; positions must lie inside the request's canvas	held at their seed value on every denoise step, and their self-conditioning soft embed is zeroed, so a read past one step denoises the canvas the request seeded instead of one the model rewrote
diffusion_read_only	bool (the OpenAI server delivers JSON booleans in vllm_xargs as 0/1, which is accepted)	caps max_tokens at the canvas, sets ignore_eos	emit on the converging step, skip the commit forward, temperature-1 logprobs
diffusion_canvas_length	int, 1 to the served canvas_length	validated in verify; the seed must be this long	the request runs on a canvas this wide: the scheduler schedules this many draft tokens and the sampler tiles the read at this width, so a short schema on a wide server costs the same as on a server of its own width
logprob_token_ids (a request field, not in vllm_xargs)	list[int], up to 128	existing vLLM validation	exact logprobs for those ids at every canvas position. The sampler now honors it on the converging step, so a read asks for its label ids instead of hoping they rank in the top-k (with 26 options only 1-8 did)

Example request:

{
  "model": "dgemma",
  "messages": [
    {"role": "system", "content": "Answer a fixed set of questions about the state ... Question urgent: ...\n  yes\n  no\n..."},
    {"role": "user",   "content": "{\"ticket\": \"Since this morning the dashboard shows a blank page after login...\"}"}
  ],
  "chat_template_kwargs": {"enable_thinking": false},
  "max_tokens": 17,
  "logprobs": true,
  "top_logprobs": 20,
  "return_tokens_as_token_ids": true,
  "vllm_xargs": {
    "diffusion_seed_canvas": [100, 45518, 107, 101, 148436, 236787, 71203, 107, 36164, 236787, 9950, 107, 66387, 236787, 236743, 189012, 106, 0, 0, "... 64 ids total"],
    "diffusion_max_steps": 1,
    "diffusion_read_only": true
  }
}

Example Queries

These are raw example queries passed to the model:

Answer a fixed set of questions about the state the user provides. Each question lists its allowed answers; reply with exactly one label per question.

You are reviewing a user post on a community forum.

Question violates: Does the post break the community rules?
  yes
  no

Question category: Which rule does it come closest to breaking?
  A: none
  B: spam (unsolicited promotion)
  C: harassment (targets a person or group)
  D: self_harm

Question severity: How severe is it?
  1: none
  2: mild
  3: serious
  4: critical

Reply with one line per question, in this order, formatted as "id: label".

Answer template (@ standing in for a noised slot):

<|channel>thought\n<channel|>violates: @\ncategory: @\nseverity: @

Seeded canvas:

[100, 45518, 107, 101, 51584, 1090, 236787, 58369, 107, 10618, 236787, 13112, 107, 127243, 236787, 236743, 144194, 106, 0, 0, ... 0]
  pos 7:  violates <- 58369  (' Mell')
  pos 11: category <- 13112  (' laugh')
  pos 16: severity <- 144194 ('اهرة')

Test Plan

Run locally on a DGX spark, with a concurrency sweep to ensure it scales and does not corrupt parallel generations.

pytest tests/test_sampling_params.py tests/v1/sample/test_diffusion_gemma_reads.py

Run with:

vllm serve nvidia/diffusiongemma-26B-A4B-it-NVFP4 --diffusion-config '{"canvas_length": 32}' --max-logprobs 32 --enable-prefix-caching --async-scheduling --attention-backend TRITON_ATTN --max-num-seqs 32
python examples/features/structured_diffusion/structured_server.py --upstream http://127.0.0.1:8000 --tokenizer nvidia/diffusiongemma-26B-A4B-it-NVFP4 --canvas 32

Test Result

Hardware: 1 DGX Spark, sole active GPU process.

Setup: read-only single reads, 32-row canvas, three decisions per request, cache-busted state with the schema prefix cached

Results: 1-way 8.7 req/s at 0.12 s, 32-way 54.0 req/s at 0.58 s, about 162 decisions/s.

Using a sample request/response corpus to determine, from a mock support request, the urgency, tone and routing of an outage report, we get the correct variable number of samples per request (the auto policy re-reads only when the first read is uncertain):

H1: is the first read's entropy per question over the top-20 logprobs (note that reads=4 when H1>0.1)

== samples=auto
outage      527ms reads=4 urgent=0.88±0.04 bucket=outage(1.00) tone=furious(3.00) H1=urgent:0.87 bucket:0.36 tone:0.08
billing     301ms reads=4 urgent=0.00±0.00 bucket=billing(1.00) tone=calm(1.00) H1=urgent:0.65 bucket:0.36 tone:0.28
feature     117ms reads=1 urgent=0.00 bucket=feature(1.00) tone=calm(1.00) H1=urgent:0.06 bucket:0.07 tone:0.02
mixed       296ms reads=4 urgent=0.00±0.00 bucket=billing(1.00) tone=annoyed(2.02) H1=urgent:0.26 bucket:0.13 tone:0.25

Additional Results

== programming language
  ok  want=python       got=python       p(want)=1.00±0.00  reads=4  528ms
  ok  want=rust         got=rust         p(want)=0.99±0.00  reads=4  485ms
  ok  want=javascript   got=javascript   p(want)=1.00±0.00  reads=4  493ms
  ok  want=go           got=go           p(want)=0.97±0.01  reads=4  478ms
  ok  want=c            got=c            p(want)=1.00±0.00  reads=4  497ms
  ok  want=sql          got=sql          p(want)=1.00±0.00  reads=4  428ms
  ok  want=bash         got=bash         p(want)=1.00±0.00  reads=4  420ms
  ok  want=ruby         got=ruby         p(want)=0.99±0.00  reads=4  483ms
  ok  want=java         got=java         p(want)=1.00±0.00  reads=4  423ms
  ok  want=haskell      got=haskell      p(want)=0.90±0.02  reads=4  420ms
  accuracy 10/10
== human language
  ok  want=english      got=english      p(want)=1.00±0.00  reads=4  508ms
  ok  want=french       got=french       p(want)=1.00±0.00  reads=4  494ms
  ok  want=german       got=german       p(want)=1.00±0.00  reads=4  474ms
  ok  want=spanish      got=spanish      p(want)=1.00±0.00  reads=4  481ms
  ok  want=italian      got=italian      p(want)=0.95±0.01  reads=4  482ms
  MISS want=portuguese   got=french       p(want)=0.47±0.11  reads=4  496ms
  ok  want=dutch        got=dutch        p(want)=0.98±0.01  reads=4  512ms
  ok  want=japanese     got=japanese     p(want)=0.98±0.01  reads=4  514ms
  ok  want=swedish      got=swedish      p(want)=0.93±0.04  reads=4  508ms
  ok  want=polish       got=polish       p(want)=0.98±0.01  reads=4  504ms
  accuracy 9/10
== unit comparison
  ok  want=first        got=first        p(want)=0.98±0.00  reads=4  503ms
  ok  want=second       got=second       p(want)=0.69±0.03  reads=4  507ms
  MISS want=first        got=second       p(want)=0.20±0.03  reads=4  504ms
  ok  want=equal        got=equal        p(want)=1.00±0.00  reads=4  503ms
  ok  want=equal        got=equal        p(want)=1.00±0.00  reads=4  491ms
  ok  want=first        got=first        p(want)=0.86±0.01  reads=4  498ms
  ok  want=equal        got=equal        p(want)=1.00±0.00  reads=4  418ms
  ok  want=second       got=second       p(want)=0.99±0.00  reads=4  435ms
  MISS want=first        got=second       p(want)=0.19±0.02  reads=4  493ms
  ok  want=second       got=second       p(want)=0.98±0.00  reads=4  493ms
  ok  want=second       got=second       p(want)=0.80±0.01  reads=4  502ms
  ok  want=first        got=first        p(want)=0.77±0.01  reads=4  441ms
  accuracy 10/12

Solve an ASCII maze giving it only N/E/W/S options:

It even works with the vision tower!

== color, by square size (8 colors)
  size 8px        8/8   p(correct) 1.00    426 ms
  size 16px       8/8   p(correct) 1.00    459 ms
  size 32px       8/8   p(correct) 1.00    427 ms
  size 64px       8/8   p(correct) 1.00    452 ms
  size 128px      8/8   p(correct) 1.00    449 ms
== count, red 32px squares on a 4x4 grid
  1 squares       4/4   p(correct) 0.99    495 ms
  2 squares       4/4   p(correct) 0.98    487 ms
  3 squares       4/4   p(correct) 0.89    478 ms
  4 squares       4/4   p(correct) 0.94    492 ms
== quadrant, 32px square
  top_left        4/4   p(correct) 0.97    505 ms
  top_right       4/4   p(correct) 1.00    497 ms
  bottom_left     4/4   p(correct) 0.84    500 ms
  bottom_right    3/4   p(correct) 0.51    498 ms
== presence: is the square red? (half yes, half no)
  size 8px        8/8   p(correct) 0.94    441 ms
  size 16px       8/8   p(correct) 1.00    429 ms
  size 32px       8/8   p(correct) 1.00    407 ms
  size 64px       8/8   p(correct) 1.00    404 ms
== contrast: 48px gray square vs plain background, flat/square choice (half each)
  contrast +2     6/8   p(correct) 0.75    348 ms
  contrast +4     8/8   p(correct) 0.98    319 ms
  contrast +8     8/8   p(correct) 0.99    356 ms
  contrast +16    8/8   p(correct) 0.94    315 ms
  contrast +32    8/8   p(correct) 1.00    292 ms
  contrast +64    8/8   p(correct) 1.00    309 ms
  contrast +127   8/8   p(correct) 1.00    320 ms

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[mergify]

Documentation preview: https://vllm--57250.org.readthedocs.build/en/57250/

[mmastrac]

Some small updates to the PR: there's more work we can skip when not actually diffusing which shaves off ~20-40% off decision time with no impact to decision quality.

[kvnloo]

This is genuinely exciting work — Jev-style structured reads turn DiffusionGemma into a calibrated multiple-choice machine: bounded answers with real uncertainty bars, which is exactly what evals, moderation, and agentic tool-calling need instead of prose you have to parse and pray over. The machinery here is close. A few things I'd want resolved before it lands.

Merge order needs a decision first. The three prerequisite PRs (#57414, #57416, #57417) are stacked as commits inside this branch and open as standalone PRs. Both can't survive: if the standalone PRs land first, this branch conflicts; if this lands first, they become empty. Either rebase this after they land (dropping the three commits) or absorb them here and close the standalone PRs. Right now it's both, which guarantees a mess.

Should-fix:

1. No regression test for the claimed race. All 7 new sampler tests are single-slot state checks; nothing runs two lockstep requests where A converges (stash freshly written) while B commits (stash popped) — the exact case that 500'd with list index out of range. Every new test stays green while this regresses. A two-request mixed-phase test would pin it.

2. Example server has unbounded fan-out. auto_max is int-cast but never upper-bounded (fixed samples is capped at 32); read_many spawns one thread per sample and decide uses ThreadPoolExecutor(max_workers=len(groups)) with unbounded question count. {"samples": "auto", "auto_max": 100000} → ~10⁵ threads → can't start new thread or OOM. The server binds 0.0.0.0 with no auth. Mirror the min(..., 32) cap and bound question count.

3. _diffusion_read_in_flight is void without diffusion_max_steps, which validation permits. Under async scheduling a read-only request without a cap gets one extra speculative denoise step after its converging output is already in flight; the output is dropped and its +width placeholders leak on the dead request — a wasted forward pass per read. Either require the cap alongside diffusion_read_only or make the guard cap-independent.

4. diffusion_max_steps is silently clamped. Validation rejects < 1 but allows huge values, which the sampler clamps to max_denoising_steps with no signal (request 200 steps, get 48, read-only emits at 48). Validate the upper bound or log the clamp.

5. Read-only with max_tokens < canvas width truncates the canvas. The cap only lowers (min(self.max_tokens or expected_len, expected_len)), contradicting the comment "one canvas is the whole output": max_tokens=5, width=16 → the structured reader gets 5 of 16 slots and downstream slot alignment breaks. Reject or canonicalize.

6. verify() mutates the params it's validating. The read-only branch sets self.max_tokens and unconditionally stomps self.ignore_eos = True, overriding an explicit user ignore_eos=False. Idempotent on re-verify, but validation with side effects is a footgun — document the override instead of silently applying it.

Nits: README field table omits diffusion_canvas_length and says read-only emits "as soon as the cap is reached" — it emits at convergence, which can precede the cap. slot_distribution's entropy is computed over unnormalized top-set probabilities, systematically understating true entropy and biasing the auto policy toward single reads; the (properly normalized) label distribution's entropy would be the honest signal. set_seed_canvas only overwrites [:len(ids)], so a shorter seed on a reused slot leaves the previous request's tail in device memory past the new width — likely benign since those positions are never scheduled, but worth a zeroing for hygiene. When diffusion_config is None the seed-length check degrades to unenforced, re-opening the device-side assert the validation was built to prevent.

What's strong: the stash fix matches the race mechanism exactly — I traced placeholder add/drain across _update_after_schedule, the rejection path, and the async commit override, and the books balance for narrowed and full widths. Slot-reuse hygiene is careful (add_request resets seed/read-only/cap/width and evicts stale stashes — the exact leak vector for a preempt between converge and commit). And validating seed ids against vocab size before GPU indexing turns an unrecoverable device-side assert into a clean 400, including the bool-is-int exclusion most validators miss.

Design question: extra_args as the API surface. It works, but it's an untyped escape hatch — invisible to the OpenAPI schema, no typed client support, validation bolted on the side. You already anticipate a template syntax in a followup; I'd argue the three knobs (seed canvas, step cap, read-only) deserve first-class SamplingParams fields when that lands, with extra_args as the incubator, not the permanent home. Agree with keeping the JSON format itself outside vLLM — the interposer server is the right layer for it.