[Core] structured generation mode for DiffusionGemma model (Jev-like) (#57250)
Signed-off-by: Matt Mastracci <matthew@mastracci.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Razorback16 <razorback16@protonmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Razorback16 <razorback16@protonmail.com>
Matt Mastracci committed
Sep 22, 2026 at 14:51 UTC
1b3b88ec2b7457aa030db4d0e7d8aaf04f6d0fb8
16 files changed
+2797
-157
examples/features/structured_diffusion/README.md
new
+74
@@ -0,0 +1,74 @@
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
+
15
+`structured_server.py` turns a question schema into those fields. It serves
16
+`/v1/chat/completions`: the system message is the schema, the user message
17
+is the state JSON, and the reply content is one distribution per question
18
+with a standard error over a few noise draws.
19
+
20
+```bash
21
+vllm serve google/diffusiongemma-26B-A4B-it \
22
+ --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 --enable-prefix-caching
23
+python examples/features/structured_diffusion/structured_server.py \
24
+ --upstream http://127.0.0.1:8000 --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64
25
+curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -d '{
26
+ "messages": [
27
+ {"role": "system", "content": "{\"questions\": [{\"id\": \"urgent\", \"type\": \"noul\", \"instructions\": \"Does the customer need a reply within the hour?\"}]}"},
28
+ {"role": "user", "content": "{\"ticket\": \"Everything is down and we have a demo at noon.\"}"}
29
+ ]}'
30
+```
31
+
32
+The attention backend is picked as for Gemma 4: FlashAttention 4 on every
33
+layer when available, otherwise Triton. FlashInfer cannot serve this model (a
34
+batch mixes causal prefill with bidirectional denoising), and
35
+`--attention-backend FLASHINFER` is rejected.
36
+
37
+Per-request canvas widths smaller than the served canvas require async
38
+scheduling. The diffusion async scheduler is selected automatically; no
39
+`--scheduler-cls` argument is needed. Synchronous execution supports full-width
40
+canvases only.
41
+
42
+Question types: `noul` (yes/no), `choice` with `options`, `score` with
43
+ordered `levels`. Each label must be a single token in the answer template,
44
+which the server checks with the tokenizer when a request uses the schema.
45
+
46
+`POST /v1/systemone` implements the Jev decision API. The body holds
47
+`state`, `questions` (a map of id to `type`, `instructions` and `criteria`)
48
+and `model`. Answers come back in that API's shapes: a `noul` probability,
49
+a `choice` with `probabilities` and `confidence`, or a `score` with a
50
+0-indexed `legend`. The schema options above go in the same body as
51
+extensions.
52
+
53
+A question may declare `depends_on` (read in a later stage with those
54
+answers in its prompt), `ask_if` (asked only when a named question's answer
55
+is among the listed ones, otherwise null) and `alone` (a read of its own).
56
+
57
+Images attach as `multipart/form-data`, with the JSON in a part named
58
+`request` and each image as a file part, or as an `images` array of data
59
+URLs.
60
+
61
+```bash
62
+curl -s localhost:8011/v1/systemone -H 'content-type: application/json' -d '{
63
+ "model": "jev-latest",
64
+ "state": {"ticket": "Everything is down and we have a demo at noon."},
65
+ "questions": {"urgent": {"type": "noul", "instructions": "Does the customer need a reply within the hour?"}}}'
66
+```
67
+
68
+`"think": N` in the schema lets the model write up to N tokens in its
69
+thought channel before the read. The thought is an ordinary generation with
70
+the chat template's thinking marker on, and the read then runs with the
71
+thought in its prompt, so the answer slots condition on it. The noise draws
72
+of a decision share one thought. `diagnostics.thought` returns the text, its
73
+length in tokens, whether the model closed the channel itself and the
74
+generation time.
examples/features/structured_diffusion/structured_server.py
new
+1528
@@ -0,0 +1,1528 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Structured decisions in front of a vLLM DiffusionGemma server.
4
+
5
+POST /v1/systemone takes Jev's request body: {"model", "state", "questions"}.
6
+"questions" maps an id to {"type", "instructions", "criteria"}, where "type"
7
+is "noul", "choice" or "score" and the criteria shape follows the type:
8
+ noul: optional {"true": ..., "false": ...} descriptions
9
+ choice: option name -> description or null
10
+ score: ordered list of levels
11
+Answers take Jev's shapes, with this server's diagnostics alongside:
12
+ noul: {"noul": p}
13
+ choice: {"choice", "probabilities", "confidence"}
14
+ score: {"score", "legend", "probabilities", "confidence"}
15
+The request body may also carry the schema keys "instructions", "samples",
16
+"auto_max", "auto_threshold", "steps", "think", "ask", "chunk_rows",
17
+"chunk_prompt" and "sequential" as extensions. Images go ahead of the
18
+state, either as multipart/form-data with the JSON body in a part named
19
+"request" and each image as a file part, or as an "images" array of data
20
+URLs in the JSON body.
21
+
22
+POST /v1/chat/completions makes the same decision from an OpenAI-shaped
23
+call. The system message is the schema JSON below and the user message is
24
+the state. The reply's `content` is the JSON answer set.
25
+
26
+POST /v1/raw/chat/completions passes the body to vLLM's chat completions
27
+unchanged, for plain generation through this port.
28
+
29
+With API_KEY set in the environment, every POST needs "Authorization:
30
+Bearer <key>". --tls-port adds an HTTPS listener with a self-signed
31
+certificate kept in --cert-dir, for clients that need a secure origin.
32
+
33
+Each answer is one distribution per question, from one denoise step over a
34
+seeded canvas, averaged over a few noise draws. This server handles the
35
+canvas, tokenizer, slot resolution, noise draws and averaging.
36
+
37
+Schema (system message):
38
+ {"questions": [
39
+ {"id": "urgent", "type": "noul", "instructions": "..."},
40
+ {"id": "bucket", "type": "choice", "instructions": "...",
41
+ "options": [{"name": "billing", "description": "..."}, ...]},
42
+ {"id": "tone", "type": "score", "instructions": "...",
43
+ "levels": ["calm", "annoyed", "furious"]}],
44
+ "instructions": "optional context",
45
+ "samples": "auto" | N, "auto_threshold": 0.1, "auto_max": 4,
46
+ "steps": 1, "think": 0}
47
+
48
+A question may also declare:
49
+ "depends_on": [ids] answered after those, with their answers in
50
+ its prompt
51
+ "ask_if": {id: [answers]} asked only when that question's answer is
52
+ among them (a skipped answer is null)
53
+ "alone": true a read of its own
54
+Questions run in stages by these dependencies. Each stage is one joint
55
+read. Later stages continue the earlier answers, prefilled for a text
56
+state and restated for an image.
57
+
58
+Up to ten questions answer as "id: label" lines. Past that the id runs
59
+straight into the label, space separated, one row fewer per question. The
60
+server splits a schema whose answer template does not fit the canvas into
61
+chunks that run together, each with its own question list. "chunk_rows"
62
+sets the rows per chunk, "ask" picks a subset of question ids for one
63
+read, "sequential": true runs the chunks in order with the earlier answers
64
+prefilled, and "chunk_prompt": "shared" lists every question in each
65
+chunk's prompt. "think": N lets the model write up to N tokens in its
66
+thought channel, as an ordinary generation, and the read then runs with
67
+that thought in its prompt. With images the model writes the thought with
68
+the image in view and the server seeds it into the canvas ahead of the
69
+answer, so the canvas bounds it. The noise draws of a decision share one
70
+thought.
71
+
72
+Serve the model with a canvas that holds the answer template, for example:
73
+ vllm serve google/diffusiongemma-26B-A4B-it \
74
+ --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 \
75
+ --enable-prefix-caching
76
+then run this in front of it:
77
+ python structured_server.py --upstream http://127.0.0.1:8000 \
78
+ --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64 --port 8011
79
+"""
80
+
81
+import argparse
82
+import json
83
+import math
84
+import os
85
+import random
86
+import ssl
87
+import subprocess
88
+import threading
89
+import time
90
+import urllib.error
91
+import urllib.request
92
+from concurrent.futures import ThreadPoolExecutor
93
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
94
+
95
+import pybase64 as base64
96
+from transformers import AutoTokenizer
97
+
98
+ARGS = None
99
+TOK = None
100
+API_KEY = os.environ.get(
101
+ "API_KEY", ""
102
+) # when set, POST routes need "Authorization: Bearer <key>"
103
+CANVAS_LEN = 64 # the served canvas length. A request may be narrower.
104
+CANVAS_STEP = 16 # request widths are multiples of this
105
+VOCAB = 262144
106
+TURN_CLOSE = 106
107
+PAD = 0
108
+TOPK = 20
109
+MAX_QUESTIONS = 64 # per request
110
+MAX_SAMPLES = 32 # reads per question, fixed or auto
111
+MAX_PARALLEL = 16 # question groups read at once
112
+# the empty thought block the chat template leaves to the model
113
+SCAFFOLD_TEXT = "<|channel>thought\n<channel|>"
114
+SCAFFOLD = None
115
+THOUGHT_OPEN = None
116
+THOUGHT_CLOSE = None
117
+
118
+
119
+# ----------------------------------------------------------------------------
120
+# Schema
121
+# ----------------------------------------------------------------------------
122
+
123
+
124
+class SchemaError(ValueError):
125
+ pass
126
+
127
+
128
+def parse_schema(value):
129
+ if (
130
+ not isinstance(value, dict)
131
+ or not isinstance(value.get("questions"), list)
132
+ or not value["questions"]
133
+ ):
134
+ raise SchemaError("schema: needs a non-empty questions array")
135
+ if len(value["questions"]) > MAX_QUESTIONS:
136
+ raise SchemaError(f"schema: at most {MAX_QUESTIONS} questions")
137
+ qs = []
138
+ seen = set()
139
+ for q in value["questions"]:
140
+ qid = str(q.get("id", "")).strip()
141
+ if not qid or ":" in qid or "\n" in qid:
142
+ raise SchemaError(
143
+ f"question id {qid!r} must be non-empty, no ':' or newline"
144
+ )
145
+ if qid in seen:
146
+ raise SchemaError(f"duplicate question id {qid!r}")
147
+ seen.add(qid)
148
+ kind = q.get("type")
149
+ if kind in ("noul", "bool", "boolean"):
150
+ kind = "noul"
151
+ crit = q.get("criteria") or {}
152
+ choices = [("yes", crit.get("true")), ("no", crit.get("false"))]
153
+ labels = ["yes", "no"]
154
+ elif kind == "choice":
155
+ opts = q.get("options") or []
156
+ choices = [
157
+ (o["name"], o.get("description"))
158
+ if isinstance(o, dict)
159
+ else (str(o), None)
160
+ for o in opts
161
+ ]
162
+ labels = [chr(ord("A") + i) for i in range(len(choices))]
163
+ elif kind == "score":
164
+ choices = [(str(level), None) for level in (q.get("levels") or [])]
165
+ labels = (
166
+ [str(i + 1) for i in range(len(choices))]
167
+ if len(choices) <= 9
168
+ else [chr(ord("A") + i) for i in range(len(choices))]
169
+ )
170
+ else:
171
+ raise SchemaError(f"question {qid!r}: unknown type {kind!r}")
172
+ if len(choices) < 2:
173
+ raise SchemaError(f"question {qid!r}: needs at least two alternatives")
174
+ if len(choices) > 26:
175
+ raise SchemaError(f"question {qid!r}: at most 26 alternatives")
176
+ deps = q.get("depends_on") or []
177
+ ask_if = q.get("ask_if") or {}
178
+ if not isinstance(deps, list) or not all(isinstance(d, str) for d in deps):
179
+ raise SchemaError(
180
+ f"question {qid!r}: depends_on must be a list of question ids"
181
+ )
182
+ if not isinstance(ask_if, dict) or not all(
183
+ isinstance(v, list) and v for v in ask_if.values()
184
+ ):
185
+ raise SchemaError(
186
+ f"question {qid!r}: ask_if must map a question id to a "
187
+ "non-empty list of its answers"
188
+ )
189
+ qs.append(
190
+ {
191
+ "id": qid,
192
+ "type": kind,
193
+ "instructions": str(q.get("instructions", "")),
194
+ "choices": choices,
195
+ "labels": labels,
196
+ "depends_on": list(dict.fromkeys(list(deps) + list(ask_if))),
197
+ "ask_if": ask_if,
198
+ "alone": bool(q.get("alone", False)),
199
+ }
200
+ )
201
+ by_id = {q["id"]: q for q in qs}
202
+ for q in qs:
203
+ for dep in q["depends_on"]:
204
+ if dep not in by_id or dep == q["id"]:
205
+ raise SchemaError(
206
+ f"question {q['id']!r}: depends on unknown question {dep!r}"
207
+ )
208
+ for dep, vals in q["ask_if"].items():
209
+ names = [c[0] for c in by_id[dep]["choices"]]
210
+ if any(v not in names for v in vals):
211
+ raise SchemaError(
212
+ f"question {q['id']!r}: ask_if values for {dep!r} "
213
+ f"must be among {names}"
214
+ )
215
+ schedule(qs) # refuses a cycle
216
+ samples = value.get("samples", "auto")
217
+ if samples == "auto":
218
+ policy = {
219
+ "mode": "auto",
220
+ "max": max(1, min(int(value.get("auto_max", 4)), MAX_SAMPLES)),
221
+ "threshold": float(value.get("auto_threshold", 0.1)),
222
+ }
223
+ elif isinstance(samples, int) and samples >= 1:
224
+ policy = {"mode": "fixed", "n": min(samples, MAX_SAMPLES)}
225
+ else:
226
+ raise SchemaError('schema: samples must be a positive count or "auto"')
227
+ ask = value.get("ask")
228
+ if ask is not None:
229
+ if not isinstance(ask, list) or not ask or any(a not in seen for a in ask):
230
+ raise SchemaError("schema: ask must list question ids from this schema")
231
+ for q in qs:
232
+ if q["id"] in ask and any(d not in ask for d in q["depends_on"]):
233
+ raise SchemaError(
234
+ f"schema: ask names {q['id']!r} but not everything it depends on"
235
+ )
236
+ chunk_rows = value.get("chunk_rows")
237
+ if chunk_rows is not None and (not isinstance(chunk_rows, int) or chunk_rows < 8):
238
+ raise SchemaError("schema: chunk_rows must be an integer of at least 8")
239
+ chunk_prompt = value.get("chunk_prompt", "own")
240
+ if chunk_prompt not in ("shared", "own"):
241
+ raise SchemaError('schema: chunk_prompt must be "shared" or "own"')
242
+ sequential = bool(value.get("sequential", False))
243
+ think = value.get("think", 0)
244
+ if isinstance(think, bool) or not isinstance(think, int) or not 0 <= think <= 4096:
245
+ raise SchemaError("schema: think must be a thought budget in tokens, 0 to 4096")
246
+ return {
247
+ "questions": qs,
248
+ "instructions": value.get("instructions"),
249
+ "policy": policy,
250
+ "steps": max(1, min(int(value.get("steps", 1)), 8)),
251
+ "think": think,
252
+ "ask": ask,
253
+ "chunk_rows": chunk_rows,
254
+ "chunk_prompt": chunk_prompt,
255
+ "sequential": sequential,
256
+ "format": "lines" if len(qs) <= 10 else "indexed",
257
+ }
258
+
259
+
260
+# Answer template shape: (join between questions, what precedes the label,
261
+# reply instruction). A small schema gets "lines", which is readable.
262
+# "indexed" ("0yes 1no") costs three tokens a question against four or five
263
+# and agreed with "lines" on every set tried: 42 booleans, ten 26-way
264
+# choices, twenty 5-level scores. Past ten questions the saved rows keep a
265
+# schema in one read. Two tokens a question, or no id at all, loses
266
+# alignment beyond about twenty questions, because the id ties a label to
267
+# its question.
268
+FORMATS = {
269
+ "lines": (
270
+ "\n",
271
+ "{id}: ",
272
+ 'Reply with one line per question, in this order, formatted as "id: label".',
273
+ ),
274
+ "indexed": (
275
+ " ",
276
+ "{id}",
277
+ "Reply on one line with each question's id immediately followed by its "
278
+ "label, separated by single spaces.",
279
+ ),
280
+}
281
+
282
+
283
+def system_text(schema, chunked=False):
284
+ s = (
285
+ "Answer a fixed set of questions about the state the user provides. "
286
+ "Each question lists its allowed answers; reply with exactly one label "
287
+ "per question.\n"
288
+ )
289
+ if schema.get("instructions"):
290
+ s += "\n" + str(schema["instructions"]).strip() + "\n"
291
+ for q in schema["questions"]:
292
+ s += f"\nQuestion {q['id']}: {q['instructions'].strip()}\n"
293
+ for (name, desc), label in zip(q["choices"], q["labels"]):
294
+ if q["type"] == "noul":
295
+ s += f" {label}: {str(desc).strip()}\n" if desc else f" {label}\n"
296
+ elif desc:
297
+ s += f" {label}: {name} ({str(desc).strip()})\n"
298
+ else:
299
+ s += f" {label}: {name}\n"
300
+ s += "\n" + FORMATS[schema.get("format", "lines")][2]
301
+ if chunked:
302
+ s += (
303
+ " A reply may cover only some of the questions; answer every line "
304
+ "that is present."
305
+ )
306
+ return s
307
+
308
+
309
+def answer_text(qs, labels, fmt="lines"):
310
+ join, lead, _ = FORMATS[fmt]
311
+ return join.join(
312
+ lead.format(id=q["id"]) + q["labels"][i] for q, i in zip(qs, labels)
313
+ )
314
+
315
+
316
+def enc(text):
317
+ return TOK.encode(text, add_special_tokens=False)
318
+
319
+
320
+def init_tokenizer(tok):
321
+ global TOK, SCAFFOLD, THOUGHT_OPEN, THOUGHT_CLOSE
322
+ TOK = tok
323
+ THOUGHT_OPEN = enc("<|channel>thought\n")
324
+ THOUGHT_CLOSE = enc("<channel|>")
325
+ SCAFFOLD = enc(SCAFFOLD_TEXT)
326
+ assert THOUGHT_OPEN + THOUGHT_CLOSE == SCAFFOLD, (
327
+ "the thought tags must tokenize apart"
328
+ )
329
+
330
+
331
+def resolve_template(qs, head, lead, fmt):
332
+ """Tokenize the answer template and find each question's slot. Every label
333
+ must change exactly one token, at the same position for all of a question's
334
+ labels, or this raises SchemaError. ``head`` is the token run the canvas
335
+ starts with: the empty thought block for a plain read, and empty when the
336
+ prompt already ends the thought channel. ``lead`` is the text before the
337
+ first answer: the join when earlier answers are in the prompt, so the
338
+ tokens match one joint template."""
339
+ base_labels = [0] * len(qs)
340
+ base = head + enc(lead + answer_text(qs, base_labels, fmt))
341
+ if len(base) + 1 > CANVAS_LEN:
342
+ raise SchemaError(
343
+ f"answer template is {len(base)} tokens; the canvas holds {CANVAS_LEN - 1}"
344
+ )
345
+ if len(qs) == 1 and len(base) + 1 > CANVAS_LEN:
346
+ raise SchemaError(
347
+ f"question {qs[0]['id']!r} alone needs {len(base) + 1} canvas rows"
348
+ )
349
+ slots = []
350
+ for qi, q in enumerate(qs):
351
+ pos = None
352
+ ids = [0] * len(q["labels"])
353
+ for li in range(1, len(q["labels"])):
354
+ labels = list(base_labels)
355
+ labels[qi] = li
356
+ e = head + enc(lead + answer_text(qs, labels, fmt))
357
+ if len(e) != len(base):
358
+ raise SchemaError(
359
+ f"question {q['id']!r}: label {q['labels'][li]!r} is not a "
360
+ "single token"
361
+ )
362
+ diffs = [i for i in range(len(e)) if e[i] != base[i]]
363
+ if len(diffs) != 1 or (pos is not None and diffs[0] != pos):
364
+ raise SchemaError(
365
+ f"question {q['id']!r}: labels do not share one template slot"
366
+ )
367
+ pos = diffs[0]
368
+ ids[li] = e[pos]
369
+ ids[0] = base[pos]
370
+ if len(set(ids)) != len(ids):
371
+ raise SchemaError(
372
+ f"question {q['id']!r}: two labels tokenize to the same id"
373
+ )
374
+ slots.append({"pos": pos, "label_ids": ids})
375
+ return base, slots
376
+
377
+
378
+_template_cache = {}
379
+
380
+
381
+def template_for(schema, head, lead):
382
+ fmt = schema.get("format", "lines")
383
+ key = json.dumps(
384
+ [head, lead, fmt] + [(q["id"], q["labels"]) for q in schema["questions"]]
385
+ )
386
+ if key not in _template_cache:
387
+ _template_cache[key] = resolve_template(schema["questions"], head, lead, fmt)
388
+ return _template_cache[key]
389
+
390
+
391
+# ----------------------------------------------------------------------------
392
+# Reads
393
+# ----------------------------------------------------------------------------
394
+
395
+
396
+def canvas_width(template):
397
+ """Smallest multiple of CANVAS_STEP that holds the template and the turn close."""
398
+ need = len(template) + 1
399
+ return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP)
400
+
401
+
402
+def pin_xargs(template, slots, steps):
403
+ """Past one denoise step the template must be held, or accept/renoise
404
+ rewrites it: pin every canvas position that is not an answer slot."""
405
+ if steps <= 1:
406
+ return {}
407
+ free = {s["pos"] for s in slots}
408
+ return {
409
+ "diffusion_pinned": [p for p in range(canvas_width(template)) if p not in free]
410
+ }
411
+
412
+
413
+def build_canvas(template, slots, seed):
414
+ rng = random.Random(seed)
415
+ canvas = list(template) + [TURN_CLOSE]
416
+ canvas += [PAD] * (canvas_width(template) - len(canvas))
417
+ for s in slots:
418
+ canvas[s["pos"]] = rng.randrange(VOCAB)
419
+ return canvas
420
+
421
+
422
+def label_id_union(slots):
423
+ ids = sorted({i for s in slots for i in s["label_ids"]})
424
+ return ids[:128] # vLLM's cap per request. A schema needs far fewer.
425
+
426
+
427
+def upstream_chat(body, timeout=600):
428
+ req = urllib.request.Request(
429
+ ARGS.upstream.rstrip("/") + "/v1/chat/completions",
430
+ data=json.dumps(body).encode(),
431
+ headers={"content-type": "application/json"},
432
+ )
433
+ return json.load(urllib.request.urlopen(req, timeout=timeout))
434
+
435
+
436
+def upstream_completions(body, timeout=600):
437
+ req = urllib.request.Request(
438
+ ARGS.upstream.rstrip("/") + "/v1/completions",
439
+ data=json.dumps(body).encode(),
440
+ headers={"content-type": "application/json"},
441
+ )
442
+ return json.load(urllib.request.urlopen(req, timeout=timeout))
443
+
444
+
445
+def chat_prompt_ids(sys_text, state_text, thinking=False):
446
+ """The prompt the chat endpoint would build, as token ids, ending after
447
+ the model turn marker. Text states only. ``thinking`` turns the chat
448
+ template's thinking marker on."""
449
+ messages = [
450
+ {"role": "system", "content": sys_text},
451
+ {"role": "user", "content": state_text},
452
+ ]
453
+ out = TOK.apply_chat_template(
454
+ messages, tokenize=True, add_generation_prompt=True, enable_thinking=thinking
455
+ )
456
+ ids = (
457
+ out["input_ids"] if hasattr(out, "keys") else out
458
+ ) # newer transformers return a dict
459
+ return [int(t) for t in ids]
460
+
461
+
462
+def think(sys_text, state_text, budget):
463
+ """A read prefix that ends with a thought the model wrote: the chat prompt
464
+ with thinking on, the open tag, up to ``budget`` generated tokens, the
465
+ close tag. Returns the prefix and a diagnostics dict for the thought."""
466
+ prompt = chat_prompt_ids(sys_text, state_text, thinking=True) + THOUGHT_OPEN
467
+ started = time.time()
468
+ d = upstream_completions(
469
+ {
470
+ "model": ARGS.model,
471
+ "prompt": prompt,
472
+ "max_tokens": budget,
473
+ "logprobs": 0,
474
+ "return_tokens_as_token_ids": True,
475
+ "stop_token_ids": THOUGHT_CLOSE,
476
+ }
477
+ )
478
+ ids = [int(t.split(":")[1]) for t in d["choices"][0]["logprobs"]["tokens"]]
479
+ closed = THOUGHT_CLOSE[0] in ids
480
+ if closed:
481
+ ids = ids[: ids.index(THOUGHT_CLOSE[0])]
482
+ info = {
483
+ "tokens": len(ids),
484
+ "closed": closed,
485
+ "ms": (time.time() - started) * 1e3,
486
+ "text": TOK.decode(ids),
487
+ }
488
+ return prompt + ids + THOUGHT_CLOSE, info
489
+
490
+
491
+def think_chat(sys_text, state_content, budget):
492
+ """A thought written with the image in view: the chat endpoint with
493
+ thinking on, capped at ``budget`` tokens and cut at the close tag.
494
+ Returns the thought's token ids and a diagnostics dict."""
495
+ body = {
496
+ "model": ARGS.model,
497
+ "messages": [
498
+ {"role": "system", "content": sys_text},
499
+ {"role": "user", "content": state_content},
500
+ ],
501
+ "max_tokens": budget,
502
+ "logprobs": True,
503
+ "top_logprobs": 0,
504
+ "return_tokens_as_token_ids": True,
505
+ "stop_token_ids": THOUGHT_CLOSE,
506
+ "chat_template_kwargs": {"enable_thinking": True},
507
+ }
508
+ started = time.time()
509
+ d = upstream_chat(body)
510
+ choice = d["choices"][0]
511
+ ids = [
512
+ int(t["token"].split(":")[1])
513
+ for t in (choice.get("logprobs") or {}).get("content") or []
514
+ ]
515
+ if ids[: len(THOUGHT_OPEN)] == THOUGHT_OPEN: # the model opens the channel itself
516
+ ids = ids[len(THOUGHT_OPEN) :]
517
+ closed = THOUGHT_CLOSE[0] in ids
518
+ if closed:
519
+ ids = ids[: ids.index(THOUGHT_CLOSE[0])]
520
+ return ids, {
521
+ "tokens": len(ids),
522
+ "closed": closed,
523
+ "ms": (time.time() - started) * 1e3,
524
+ "text": TOK.decode(ids),
525
+ }
526
+
527
+
528
+def one_read(
529
+ schema, template, slots, sys_text, state_content, seed, prefix=None, thinking=False
530
+):
531
+ if prefix is not None:
532
+ return one_read_continuation(schema, template, slots, prefix, seed)
533
+ messages = [
534
+ {"role": "system", "content": sys_text},
535
+ {"role": "user", "content": state_content},
536
+ ]
537
+ body = {
538
+ "model": ARGS.model,
539
+ "messages": messages,
540
+ "max_tokens": len(template) + 1,
541
+ "logprobs": True,
542
+ "top_logprobs": TOPK,
543
+ # Exact logprobs for every label at every position. With a long
544
+ # option list most labels never rank in the top-k, and the model's
545
+ # mass sits on tokens that spell the option name instead.
546
+ "logprob_token_ids": label_id_union(slots),
547
+ "return_tokens_as_token_ids": True,
548
+ "chat_template_kwargs": {"enable_thinking": thinking},
549
+ "vllm_xargs": {
550
+ "diffusion_seed_canvas": build_canvas(template, slots, seed),
551
+ "diffusion_canvas_length": canvas_width(template),
552
+ "diffusion_max_steps": schema["steps"],
553
+ "diffusion_read_only": True,
554
+ **pin_xargs(template, slots, schema["steps"]),
555
+ },
556
+ }
557
+ d = upstream_chat(body)
558
+ content = d["choices"][0]["logprobs"]["content"]
559
+ out = []
560
+ for q, s in zip(schema["questions"], slots):
561
+ top = {
562
+ int(t["token"].split(":")[1]): t["logprob"]
563
+ for t in content[s["pos"]]["top_logprobs"]
564
+ }
565
+ out.append(slot_distribution(top, s["label_ids"]))
566
+ return out, d.get("usage", {})
567
+
568
+
569
+def slot_distribution(top, label_ids):
570
+ """Label probabilities at one slot from the returned logprobs: every
571
+ label's own value plus the argmax token. Read-only logprobs are at
572
+ temperature 1, so the label softmax uses them directly. The entropy is
573
+ over that returned set."""
574
+ floor = min(top.values()) - 5.0
575
+ lp_t = [top.get(i, floor) for i in label_ids]
576
+ mx = max(lp_t)
577
+ ex = [math.exp(x - mx) for x in lp_t]
578
+ probs = [e / sum(ex) for e in ex]
579
+ top_p = [math.exp(v) for v in top.values()]
580
+ return {
581
+ "probs": probs,
582
+ "label_mass": sum(math.exp(x) for x in lp_t),
583
+ "entropy": -sum(p * math.log(p) for p in top_p if p > 0),
584
+ "argmax_is_label": max(top, key=top.get) in label_ids,
585
+ }
586
+
587
+
588
+def one_read_continuation(schema, template, slots, prompt_ids, seed):
589
+ """A read whose prompt already holds the thought scaffold and earlier
590
+ answer lines, sent as token ids so the chat template cannot alter it."""
591
+ body = {
592
+ "model": ARGS.model,
593
+ "prompt": prompt_ids,
594
+ "max_tokens": len(template) + 1,
595
+ "logprobs": TOPK,
596
+ "logprob_token_ids": label_id_union(slots),
597
+ "return_tokens_as_token_ids": True,
598
+ "vllm_xargs": {
599
+ "diffusion_seed_canvas": build_canvas(template, slots, seed),
600
+ "diffusion_canvas_length": canvas_width(template),
601
+ "diffusion_max_steps": schema["steps"],
602
+ "diffusion_read_only": True,
603
+ **pin_xargs(template, slots, schema["steps"]),
604
+ },
605
+ }
606
+ d = upstream_completions(body)
607
+ rows = d["choices"][0]["logprobs"]["top_logprobs"]
608
+ out = []
609
+ for q, sl in zip(schema["questions"], slots):
610
+ top = {int(k.split(":")[1]): v for k, v in rows[sl["pos"]].items()}
611
+ out.append(slot_distribution(top, sl["label_ids"]))
612
+ return out, d.get("usage", {})
613
+
614
+
615
+def read_many(
616
+ schema,
617
+ template,
618
+ slots,
619
+ sys_text,
620
+ state_content,
621
+ seed,
622
+ n,
623
+ prefix=None,
624
+ thinking=False,
625
+):
626
+ results = [None] * n
627
+ errors = [None] * n
628
+ usages = [None] * n
629
+
630
+ def run(k):
631
+ try:
632
+ results[k], usages[k] = one_read(
633
+ schema,
634
+ template,
635
+ slots,
636
+ sys_text,
637
+ state_content,
638
+ seed + k * 7919,
639
+ prefix,
640
+ thinking,
641
+ )
642
+ except Exception as e: # raised again below
643
+ errors[k] = e
644
+
645
+ threads = [threading.Thread(target=run, args=(k,)) for k in range(n)]
646
+ for t in threads:
647
+ t.start()
648
+ for t in threads:
649
+ t.join()
650
+ for e in errors:
651
+ if e is not None:
652
+ raise e
653
+ return results, usages
654
+
655
+
656
+def schedule(qs):
657
+ """Questions in stages: a question's stage comes after the stages of
658
+ everything it depends on. Declaration order is kept within a stage."""
659
+ ids = {q["id"] for q in qs}
660
+ pending = list(qs)
661
+ done: set = set()
662
+ levels = []
663
+ while pending:
664
+ level = [
665
+ q
666
+ for q in pending
667
+ if all(d in done or d not in ids for d in q["depends_on"])
668
+ ]
669
+ if not level:
670
+ raise SchemaError(
671
+ "schema: dependency cycle among " + ", ".join(q["id"] for q in pending)
672
+ )
673
+ levels.append(level)
674
+ done |= {q["id"] for q in level}
675
+ pending = [q for q in pending if q["id"] not in done]
676
+ return levels
677
+
678
+
679
+def chunk_groups(schema, qs):
680
+ """``qs`` split, in order, into the fewest groups whose answer templates
681
+ fit ``chunk_rows`` (the canvas by default). A question marked alone gets
682
+ its own group."""
683
+ limit = schema.get("chunk_rows") or CANVAS_LEN
684
+ groups, group = [], []
685
+ for q in qs:
686
+ if q["alone"]:
687
+ if group:
688
+ groups.append(group)
689
+ group = []
690
+ groups.append([q])
691
+ continue
692
+ trial = group + [q]
693
+ rows = (
694
+ len(SCAFFOLD)
695
+ + len(
696
+ enc(answer_text(trial, [0] * len(trial), schema.get("format", "lines")))
697
+ )
698
+ + 1
699
+ )
700
+ if rows > limit and group:
701
+ groups.append(group)
702
+ group = [q]
703
+ else:
704
+ group = trial
705
+ if group:
706
+ groups.append(group)
707
+ return groups
708
+
709
+
710
+def answer_name(q, a):
711
+ """The answer as the name ask_if compares against: yes or no, an option
712
+ name, or a level name."""
713
+ if a is None:
714
+ return None
715
+ return a["label"] if q["type"] == "noul" else a.get("choice", a.get("level"))
716
+
717
+
718
+def decide(schema, state_content, seed):
719
+ """One decision. Questions run in stages by their dependencies. A stage
720
+ is one joint read, chunked by the canvas, with a question marked alone
721
+ in its own read. Later stages condition on every earlier answer: a
722
+ prefilled continuation for a text state, or the answers restated in the
723
+ state for an image. A question whose ask_if condition failed is skipped
724
+ and its answer is null."""
725
+ started = time.time()
726
+ qs = [
727
+ q
728
+ for q in schema["questions"]
729
+ if not schema.get("ask") or q["id"] in schema["ask"]
730
+ ]
731
+ levels = schedule(qs)
732
+ text_state = isinstance(state_content, str)
733
+ fmt = schema["format"]
734
+ join = FORMATS[fmt][0]
735
+ # More than one read in sequence needs the full question list in every
736
+ # prompt, so that later reads continue one answer.
737
+ chained = len(levels) > 1 or schema["sequential"]
738
+ sys_full = system_text(schema, chunked=not text_state and chained)
739
+ base_ids, thought = None, None
740
+ if chained and text_state:
741
+ if schema["think"]:
742
+ base_ids, thought = think(sys_full, state_content, schema["think"])
743
+ else:
744
+ base_ids = chat_prompt_ids(sys_full, state_content) + SCAFFOLD
745
+ shared = schema["chunk_prompt"] == "shared"
746
+ answered, lines, earlier = {}, [], []
747
+ parts, stages, chunks, skipped = [], [], [], {}
748
+ by_id = {q["id"]: q for q in schema["questions"]}
749
+
750
+ def run(group, k, conditioned):
751
+ sub = dict(schema, questions=group)
752
+ # The thought is written once: in the prefix of a chained text decision,
753
+ # or in the first read of anything else.
754
+ sub["think"] = 0 if (chained and text_state) or conditioned else schema["think"]
755
+ if text_state:
756
+ if conditioned:
757
+ prefix, lead, sys_text = (
758
+ base_ids + enc(join.join(lines)),
759
+ join,
760
+ sys_full,
761
+ )
762
+ elif chained:
763
+ prefix, lead, sys_text = (base_ids if thought else None), "", sys_full
764
+ else:
765
+ prefix, lead, sys_text = (
766
+ None,
767
+ "",
768
+ (system_text(schema, chunked=True) if shared else system_text(sub)),
769
+ )
770
+ return decide_group(
771
+ sub, sys_text, state_content, seed + 104729 * k, prefix, lead
772
+ )
773
+ state = state_content
774
+ if conditioned:
775
+ text = next(
776
+ (p["text"] for p in state_content if p.get("type") == "text"), ""
777
+ )
778
+ text += "\n\nAnswers so far:\n" + "\n".join(earlier)
779
+ state = [p for p in state_content if p.get("type") != "text"] + [
780
+ {"type": "text", "text": text}
781
+ ]
782
+ sys_text = sys_full
783
+ else:
784
+ sys_text = sys_full if (chained or shared) else system_text(sub)
785
+ return decide_group(sub, sys_text, state, seed + 104729 * k)
786
+
787
+ def absorb(group, body, rows):
788
+ parts.append((body, rows))
789
+ chunks.append([q["id"] for q in group])
790
+ answered.update(body["answers"])
791
+ lines.append(
792
+ answer_text(
793
+ group,
794
+ [q["labels"].index(body["answers"][q["id"]]["label"]) for q in group],
795
+ fmt,
796
+ )
797
+ )
798
+ earlier.extend(
799
+ f"{q['id']}: {answer_name(q, body['answers'][q['id']])}" for q in group
800
+ )
801
+
802
+ k = 0
803
+ for level in levels:
804
+ asked = []
805
+ for q in level:
806
+ failed = next(
807
+ (
808
+ (dep, vals)
809
+ for dep, vals in q["ask_if"].items()
810
+ if answer_name(by_id[dep], answered.get(dep)) not in vals
811
+ ),
812
+ None,
813
+ )
814
+ if failed:
815
+ answered[q["id"]] = None
816
+ skipped[q["id"]] = {
817
+ "because": failed[0],
818
+ "was": answer_name(by_id[failed[0]], answered.get(failed[0])),
819
+ "wanted": failed[1],
820
+ }
821
+ continue
822
+ asked.append(q)
823
+ if not asked:
824
+ continue
825
+ stages.append([q["id"] for q in asked])
826
+ groups = chunk_groups(schema, asked)
827
+ conditioned = bool(lines)
828
+ if schema["sequential"] or len(groups) == 1:
829
+ for group in groups:
830
+ body, rows = run(
831
+ group, k, conditioned or (schema["sequential"] and bool(lines))
832
+ )
833
+ absorb(group, body, rows)
834
+ k += 1
835
+ else:
836
+ with ThreadPoolExecutor(max_workers=min(len(groups), MAX_PARALLEL)) as ex:
837
+ results = list(
838
+ ex.map(
839
+ lambda gk, conditioned=conditioned: run(
840
+ gk[1], gk[0], conditioned
841
+ ),
842
+ [(k + i, g) for i, g in enumerate(groups)],
843
+ )
844
+ )
845
+ for group, (body, rows) in zip(groups, results):
846
+ absorb(group, body, rows)
847
+ k += len(groups)
848
+
849
+ answers = {q["id"]: answered.get(q["id"]) for q in qs}
850
+ diag_q = {}
851
+ for body, _ in parts:
852
+ diag_q.update(body["diagnostics"]["questions"])
853
+ # A thought written here (a chained text decision) is outside every
854
+ # group's row count. One written inside a group is already counted there.
855
+ extra_rows = thought["tokens"] if thought else 0
856
+ if thought is None:
857
+ thoughts = [b["diagnostics"].get("thought") for b, _ in parts]
858
+ thought = (
859
+ thoughts[0] if len(parts) == 1 else ([t for t in thoughts if t] or None)
860
+ )
861
+ one = len(parts) == 1 and not skipped
862
+ diagnostics = {
863
+ "steps": schema["steps"],
864
+ "stages": stages,
865
+ "skipped": skipped,
866
+ "chunks": chunks,
867
+ "chunk_prompt": "full" if chained else schema["chunk_prompt"],
868
+ "sequential": schema["sequential"],
869
+ "conditioning": (
870
+ None
871
+ if len(stages) <= 1 and not schema["sequential"]
872
+ else ("prefill" if text_state else "restated")
873
+ ),
874
+ "thought": thought,
875
+ "samples": (
876
+ parts[0][0]["diagnostics"]["samples"]
877
+ if one
878
+ else {
879
+ "n": [b["diagnostics"]["samples"]["n"] for b, _ in parts],
880
+ "tops": [b["diagnostics"]["samples"]["tops"] for b, _ in parts],
881
+ "policy": [b["diagnostics"]["samples"]["policy"] for b, _ in parts],
882
+ }
883
+ ),
884
+ "timing": {
885
+ "total_ms": (time.time() - started) * 1e3,
886
+ "reads": sum(b["diagnostics"]["timing"]["reads"] for b, _ in parts),
887
+ },
888
+ "prompt_tokens": max(
889
+ (b["diagnostics"].get("prompt_tokens") or 0) for b, _ in parts
890
+ )
891
+ or None,
892
+ "questions": diag_q,
893
+ "engine": "vllm",
894
+ }
895
+ return {"answers": answers, "diagnostics": diagnostics}, sum(
896
+ rows for _, rows in parts
897
+ ) + extra_rows
898
+
899
+
900
+def decide_group(schema, sys_text, state_content, seed, prefix=None, lead=""):
901
+ started = time.time()
902
+ thought = None
903
+ head = SCAFFOLD if prefix is None else []
904
+ thinking = False
905
+ if prefix is None and schema["think"]:
906
+ if isinstance(state_content, str):
907
+ prefix, thought = think(sys_text, state_content, schema["think"])
908
+ head = []
909
+ else:
910
+ # With images the model writes the thought with the image in view, and
911
+ # the server seeds it into the canvas ahead of the answer, so the read
912
+ # keeps the image. The canvas bounds the thought.
913
+ answer = enc(
914
+ answer_text(
915
+ schema["questions"],
916
+ [0] * len(schema["questions"]),
917
+ schema.get("format", "lines"),
918
+ )
919
+ )
920
+ fits = CANVAS_LEN - 1 - len(THOUGHT_OPEN) - len(THOUGHT_CLOSE) - len(answer)
921
+ if fits < 8:
922
+ raise SchemaError(
923
+ f"think: the canvas leaves {fits} rows for a thought "
924
+ "beside this template"
925
+ )
926
+ ids, thought = think_chat(
927
+ sys_text, state_content, min(schema["think"], fits)
928
+ )
929
+ thought["budget"] = min(schema["think"], fits)
930
+ head = THOUGHT_OPEN + ids + THOUGHT_CLOSE
931
+ thinking = True
932
+ template, slots = template_for(schema, head, lead)
933
+ policy = schema["policy"]
934
+ if policy["mode"] == "fixed":
935
+ reads, usages = read_many(
936
+ schema,
937
+ template,
938
+ slots,
939
+ sys_text,
940
+ state_content,
941
+ seed,
942
+ policy["n"],
943
+ prefix,
944
+ thinking,
945
+ )
946
+ extended = None
947
+ first_entropy = None
948
+ else:
949
+ reads, usages = read_many(
950
+ schema, template, slots, sys_text, state_content, seed, 1, prefix, thinking
951
+ )
952
+ first_entropy = {
953
+ q["id"]: r["entropy"] for q, r in zip(schema["questions"], reads[0])
954
+ }
955
+ extended = (
956
+ max(first_entropy.values()) > policy["threshold"] and policy["max"] > 1
957
+ )
958
+ if extended:
959
+ more, more_usages = read_many(
960
+ schema,
961
+ template,
962
+ slots,
963
+ sys_text,
964
+ state_content,
965
+ seed + 1,
966
+ policy["max"] - 1,
967
+ prefix,
968
+ thinking,
969
+ )
970
+ reads += more
971
+ usages += more_usages
972
+ prompt_tokens = next(
973
+ (u["prompt_tokens"] for u in usages if u and u.get("prompt_tokens")), None
974
+ )
975
+ elapsed_ms = (time.time() - started) * 1e3
976
+
977
+ answers = {}
978
+ diag_q = {}
979
+ n = len(reads)
980
+ for qi, q in enumerate(schema["questions"]):
981
+ per = [r[qi]["probs"] for r in reads]
982
+ mean = [sum(p[i] for p in per) / n for i in range(len(q["labels"]))]
983
+ top = max(range(len(mean)), key=lambda i: mean[i])
984
+ a = {
985
+ "type": q["type"],
986
+ "label": q["labels"][top],
987
+ "confidence": mean[top],
988
+ "probabilities": {c[0]: m for c, m in zip(q["choices"], mean)},
989
+ }
990
+ if q["type"] == "noul":
991
+ a["noul"] = mean[0]
992
+ elif q["type"] == "choice":
993
+ a["choice"] = q["choices"][top][0]
994
+ else:
995
+ a["score"] = sum((i + 1) * m for i, m in enumerate(mean))
996
+ a["level"] = q["choices"][top][0]
997
+ if n > 1:
998
+ var = sum((p[top] - mean[top]) ** 2 for p in per) / (n - 1)
999
+ a["stderr"] = (var / n) ** 0.5
1000
+ a["agreement"] = (
1001
+ sum(1 for p in per if max(range(len(p)), key=lambda i: p[i]) == top) / n
1002
+ )
1003
+ answers[q["id"]] = a
1004
+ diag_q[q["id"]] = {
1005
+ "pos": slots[qi]["pos"],
1006
+ "entropy": [r[qi]["entropy"] for r in reads],
1007
+ "label_mass": reads[0][qi]["label_mass"],
1008
+ "argmax_is_label": reads[0][qi]["argmax_is_label"],
1009
+ }
1010
+ tops = [
1011
+ {
1012
+ q["id"]: [
1013
+ q["labels"][
1014
+ max(range(len(r[qi]["probs"])), key=lambda i: r[qi]["probs"][i])
1015
+ ],
1016
+ max(r[qi]["probs"]),
1017
+ r[qi]["entropy"],
1018
+ ]
1019
+ for qi, q in enumerate(schema["questions"])
1020
+ }
1021
+ for r in reads
1022
+ ]
1023
+ return {
1024
+ "answers": answers,
1025
+ "diagnostics": {
1026
+ "steps": schema["steps"],
1027
+ "samples": {
1028
+ "n": n,
1029
+ "tops": tops,
1030
+ "policy": dict(
1031
+ policy, extended=extended, first_read_entropy=first_entropy
1032
+ ),
1033
+ },
1034
+ "timing": {"total_ms": elapsed_ms, "reads": n},
1035
+ "thought": thought,
1036
+ "prompt_tokens": prompt_tokens,
1037
+ "questions": diag_q,
1038
+ "engine": "vllm",
1039
+ },
1040
+ }, len(template) + 1 + (thought["tokens"] if thought else 0)
1041
+
1042
+
1043
+# ----------------------------------------------------------------------------
1044
+# Jev's contract
1045
+# ----------------------------------------------------------------------------
1046
+
1047
+JEV_EXTENSIONS = (
1048
+ "instructions",
1049
+ "samples",
1050
+ "auto_max",
1051
+ "auto_threshold",
1052
+ "steps",
1053
+ "think",
1054
+ "ask",
1055
+ "chunk_rows",
1056
+ "chunk_prompt",
1057
+ "sequential",
1058
+)
1059
+
1060
+
1061
+def jev_schema(body):
1062
+ """This server's schema from a Jev request body."""
1063
+ qs = body.get("questions")
1064
+ if not isinstance(qs, dict) or not qs:
1065
+ raise SchemaError("questions: needs a non-empty map of id -> question")
1066
+ out = []
1067
+ for qid, q in qs.items():
1068
+ if not isinstance(q, dict):
1069
+ raise SchemaError(f"question {qid!r}: must be an object")
1070
+ kind, crit, ins = q.get("type"), q.get("criteria"), q.get("instructions", "")
1071
+ item = {
1072
+ "id": qid,
1073
+ "type": kind,
1074
+ "instructions": ins if isinstance(ins, str) else json.dumps(ins),
1075
+ }
1076
+ if kind == "noul":
1077
+ if crit is not None and not isinstance(crit, dict):
1078
+ raise SchemaError(
1079
+ f"question {qid!r}: noul criteria must be an object with "
1080
+ "true and false"
1081
+ )
1082
+ item["criteria"] = crit
1083
+ elif kind == "choice":
1084
+ if not isinstance(crit, dict) or not crit:
1085
+ raise SchemaError(
1086
+ f"question {qid!r}: choice criteria must map option names "
1087
+ "to descriptions"
1088
+ )
1089
+ item["options"] = [
1090
+ {"name": str(n), "description": d} for n, d in crit.items()
1091
+ ]
1092
+ elif kind == "score":
1093
+ if not isinstance(crit, list):
1094
+ raise SchemaError(
1095
+ f"question {qid!r}: score criteria must be an ordered list "
1096
+ "of levels"
1097
+ )
1098
+ item["levels"] = crit
1099
+ else:
1100
+ raise SchemaError(f"question {qid!r}: unknown type {kind!r}")
1101
+ for key in ("depends_on", "ask_if", "alone"):
1102
+ if key in q:
1103
+ item[key] = q[key]
1104
+ out.append(item)
1105
+ schema = {k: body[k] for k in JEV_EXTENSIONS if k in body}
1106
+ schema["questions"] = out
1107
+ return parse_schema(schema)
1108
+
1109
+
1110
+def image_part(content_type, data):
1111
+ return {
1112
+ "type": "image_url",
1113
+ "image_url": {
1114
+ "url": f"data:{content_type};base64," + base64.b64encode(data).decode()
1115
+ },
1116
+ }
1117
+
1118
+
1119
+def jev_images(value):
1120
+ """Image parts from the body's "images": data URLs, or objects with
1121
+ content_type and base64."""
1122
+ parts = []
1123
+ for i, im in enumerate(value or []):
1124
+ if isinstance(im, str) and im.startswith("data:image/"):
1125
+ parts.append({"type": "image_url", "image_url": {"url": im}})
1126
+ elif (
1127
+ isinstance(im, dict)
1128
+ and str(im.get("content_type", "")).startswith("image/")
1129
+ and isinstance(im.get("base64"), str)
1130
+ ):
1131
+ parts.append(
1132
+ {
1133
+ "type": "image_url",
1134
+ "image_url": {
1135
+ "url": f"data:{im['content_type']};base64,{im['base64']}"
1136
+ },
1137
+ }
1138
+ )
1139
+ else:
1140
+ raise SchemaError(
1141
+ f"images[{i}]: a data:image/... URL or an object with "
1142
+ "content_type and base64"
1143
+ )
1144
+ return parts
1145
+
1146
+
1147
+def jev_state(body, image_parts=()):
1148
+ """The user message: the state as text (as given, or as JSON), with any
1149
+ images ahead of it."""
1150
+ state = body.get("state")
1151
+ if state is None:
1152
+ raise SchemaError("state: required")
1153
+ text = state if isinstance(state, str) else json.dumps(state)
1154
+ if not image_parts:
1155
+ return text
1156
+ return list(image_parts) + [{"type": "text", "text": text}]
1157
+
1158
+
1159
+def jev_answer(q, a):
1160
+ if a is None:
1161
+ return None
1162
+ if q["type"] == "noul":
1163
+ return {"type": "noul", "noul": a["noul"]}
1164
+ if q["type"] == "choice":
1165
+ return {
1166
+ "type": "choice",
1167
+ "choice": a["choice"],
1168
+ "probabilities": a["probabilities"],
1169
+ "confidence": a["confidence"],
1170
+ }
1171
+ names = [c[0] for c in q["choices"]]
1172
+ probs = {str(i): a["probabilities"][n] for i, n in enumerate(names)}
1173
+ return {
1174
+ "type": "score",
1175
+ "score": sum(i * p for i, p in enumerate(probs.values())),
1176
+ "legend": {str(i): n for i, n in enumerate(names)},
1177
+ "probabilities": probs,
1178
+ "confidence": a["confidence"],
1179
+ }
1180
+
1181
+
1182
+# ----------------------------------------------------------------------------
1183
+# HTTP
1184
+# ----------------------------------------------------------------------------
1185
+
1186
+
1187
+def message_text(m):
1188
+ c = m.get("content", "")
1189
+ if isinstance(c, list):
1190
+ return "".join(p.get("text", "") for p in c if isinstance(p, dict))
1191
+ return c if isinstance(c, str) else ""
1192
+
1193
+
1194
+class Handler(BaseHTTPRequestHandler):
1195
+ def log_message(self, fmt, *args):
1196
+ pass
1197
+
1198
+ def _json(self, code, obj):
1199
+ body = json.dumps(obj).encode()
1200
+ self.send_response(code)
1201
+ self.send_header("content-type", "application/json")
1202
+ self.send_header("content-length", str(len(body)))
1203
+ self.end_headers()
1204
+ self.wfile.write(body)
1205
+
1206
+ def do_GET(self):
1207
+ if self.path == "/health":
1208
+ return self._json(200, {"status": "ok"})
1209
+ return self._json(404, {"error": {"message": "unknown route"}})
1210
+
1211
+ def _read_request(self):
1212
+ """-> (body, image parts): a JSON body, or multipart/form-data with the
1213
+ JSON in a part named request and each image as a file part, in order."""
1214
+ raw = self.rfile.read(int(self.headers.get("content-length", "0")))
1215
+ ctype = self.headers.get("content-type", "")
1216
+ if not ctype.lower().startswith("multipart/form-data"):
1217
+ return json.loads(raw), []
1218
+ from email.parser import BytesParser
1219
+ from email.policy import HTTP
1220
+
1221
+ msg = BytesParser(policy=HTTP).parsebytes(
1222
+ b"Content-Type: " + ctype.encode() + b"\r\n\r\n" + raw
1223
+ )
1224
+ body, images = None, []
1225
+ for part in msg.iter_parts():
1226
+ name = part.get_param("name", header="content-disposition")
1227
+ data = part.get_payload(decode=True)
1228
+ if name == "request":
1229
+ body = json.loads(data)
1230
+ elif part.get_content_type().startswith("image/"):
1231
+ images.append(image_part(part.get_content_type(), data))
1232
+ else:
1233
+ raise ValueError(
1234
+ f"part {name!r}: neither the request JSON nor an image"
1235
+ )
1236
+ if body is None:
1237
+ raise ValueError(
1238
+ "multipart needs a part named request holding the JSON body"
1239
+ )
1240
+ return body, images
1241
+
1242
+ def do_POST(self):
1243
+ if API_KEY and self.headers.get("authorization", "") != f"Bearer {API_KEY}":
1244
+ return self._json(
1245
+ 401,
1246
+ {
1247
+ "error": {
1248
+ "message": "missing or wrong bearer token",
1249
+ "type": "authentication_error",
1250
+ }
1251
+ },
1252
+ )
1253
+ if self.path == "/v1/raw/chat/completions":
1254
+ return self._raw_chat()
1255
+ try:
1256
+ req, images = self._read_request()
1257
+ except Exception as e:
1258
+ return self._json(
1259
+ 400,
1260
+ {
1261
+ "error": {
1262
+ "message": f"invalid body: {e}",
1263
+ "type": "invalid_request_error",
1264
+ }
1265
+ },
1266
+ )
1267
+ if self.path == "/v1/systemone":
1268
+ return self._systemone(req, images)
1269
+ if self.path == "/v1/chat/completions":
1270
+ return self._chat(req)
1271
+ return self._json(404, {"error": {"message": "unknown route"}})
1272
+
1273
+ def _raw_chat(self):
1274
+ """Pass the body and status through to vLLM's chat completions."""
1275
+ raw = self.rfile.read(int(self.headers.get("content-length", "0")))
1276
+ req = urllib.request.Request(
1277
+ ARGS.upstream.rstrip("/") + "/v1/chat/completions",
1278
+ data=raw,
1279
+ headers={
1280
+ "content-type": self.headers.get("content-type", "application/json")
1281
+ },
1282
+ )
1283
+ try:
1284
+ with urllib.request.urlopen(req, timeout=600) as r:
1285
+ code, body = r.status, r.read()
1286
+ except urllib.error.HTTPError as e:
1287
+ code, body = e.code, e.read()
1288
+ self.send_response(code)
1289
+ self.send_header("content-type", "application/json")
1290
+ self.send_header("content-length", str(len(body)))
1291
+ self.end_headers()
1292
+ self.wfile.write(body)
1293
+
1294
+ def _decide(self, schema, state, seed):
1295
+ """-> (status, body) with the error body already shaped."""
1296
+ try:
1297
+ return 200, decide(schema, state, seed)
1298
+ except SchemaError as e:
1299
+ return 422, {"error": {"message": str(e), "type": "validation_error"}}
1300
+ except urllib.error.HTTPError as e:
1301
+ return 502, {
1302
+ "error": {
1303
+ "message": (
1304
+ f"upstream {e.code}: {e.read()[:300].decode(errors='replace')}"
1305
+ ),
1306
+ "type": "server_error",
1307
+ }
1308
+ }
1309
+ except Exception as e:
1310
+ return 500, {"error": {"message": repr(e), "type": "server_error"}}
1311
+
1312
+ def _systemone(self, req, images):
1313
+ try:
1314
+ schema = jev_schema(req)
1315
+ state = jev_state(req, images + jev_images(req.get("images")))
1316
+ except SchemaError as e:
1317
+ return self._json(
1318
+ 422, {"error": {"message": str(e), "type": "validation_error"}}
1319
+ )
1320
+ code, result = self._decide(schema, state, int(req.get("seed", 42)))
1321
+ if code != 200:
1322
+ return self._json(code, result)
1323
+ body, completion_tokens = result
1324
+ answers = {
1325
+ q["id"]: jev_answer(q, body["answers"][q["id"]])
1326
+ for q in schema["questions"]
1327
+ }
1328
+ labels = " ".join(
1329
+ f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items()
1330
+ )
1331
+ print(
1332
+ f"systemone: {labels} "
1333
+ f"reads={body['diagnostics']['timing']['reads']} "
1334
+ f"{body['diagnostics']['timing']['total_ms']:.0f}ms",
1335
+ flush=True,
1336
+ )
1337
+ self._json(
1338
+ 200,
1339
+ {
1340
+ "model": ARGS.model,
1341
+ "answers": answers,
1342
+ # vLLM's prompt count when the reads reported one, since it covers
1343
+ # images. Otherwise the tokenizer's count of the text prompt.
1344
+ "usage": {
1345
+ "input_tokens": body["diagnostics"].get("prompt_tokens")
1346
+ or (
1347
+ len(chat_prompt_ids(system_text(schema), state))
1348
+ if isinstance(state, str)
1349
+ else 0
1350
+ ),
1351
+ "output_tokens": completion_tokens,
1352
+ },
1353
+ "diagnostics": body["diagnostics"],
1354
+ },
1355
+ )
1356
+
1357
+ def _chat(self, req):
1358
+ msgs = req.get("messages") or []
1359
+ if (
1360
+ len(msgs) != 2
1361
+ or msgs[0].get("role") not in ("system", "developer")
1362
+ or msgs[1].get("role") != "user"
1363
+ ):
1364
+ return self._json(
1365
+ 400,
1366
+ {
1367
+ "error": {
1368
+ "message": (
1369
+ "a structured request is exactly two messages: the "
1370
+ "schema (system) and the state JSON (user)"
1371
+ ),
1372
+ "type": "invalid_request_error",
1373
+ }
1374
+ },
1375
+ )
1376
+ try:
1377
+ schema_value = json.loads(message_text(msgs[0]))
1378
+ schema = parse_schema(schema_value)
1379
+ content = msgs[1].get("content", "")
1380
+ has_image = isinstance(content, list) and any(
1381
+ isinstance(p, dict) and p.get("type") in ("image_url", "image")
1382
+ for p in content
1383
+ )
1384
+ if has_image:
1385
+ # image parts pass through to vLLM unchanged, with text parts as context
1386
+ state = content
1387
+ else:
1388
+ state = message_text(msgs[1]).strip()
1389
+ json.loads(state)
1390
+ except SchemaError as e:
1391
+ return self._json(
1392
+ 400, {"error": {"message": str(e), "type": "invalid_request_error"}}
1393
+ )
1394
+ except Exception as e:
1395
+ return self._json(
1396
+ 400,
1397
+ {
1398
+ "error": {
1399
+ "message": (
1400
+ "system must be a JSON question schema and user "
1401
+ f"must be JSON state or image parts: {e}"
1402
+ ),
1403
+ "type": "invalid_request_error",
1404
+ }
1405
+ },
1406
+ )
1407
+ code, result = self._decide(schema, state, int(req.get("seed", 42)))
1408
+ if code != 200:
1409
+ if code == 422:
1410
+ result["error"]["type"] = "invalid_request_error"
1411
+ code = 400
1412
+ return self._json(code, result)
1413
+ body, completion_tokens = result
1414
+ content = json.dumps(body, indent=2)
1415
+ labels = " ".join(
1416
+ f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items()
1417
+ )
1418
+ print(
1419
+ f"structured: {labels} "
1420
+ f"reads={body['diagnostics']['timing']['reads']} "
1421
+ f"{body['diagnostics']['timing']['total_ms']:.0f}ms",
1422
+ flush=True,
1423
+ )
1424
+ self._json(
1425
+ 200,
1426
+ {
1427
+ "id": f"chatcmpl-{int(time.time() * 1000)}",
1428
+ "object": "chat.completion",
1429
+ "created": int(time.time()),
1430
+ "model": req.get("model", "dgemma-structured"),
1431
+ "choices": [
1432
+ {
1433
+ "index": 0,
1434
+ "message": {"role": "assistant", "content": content},
1435
+ "finish_reason": "stop",
1436
+ }
1437
+ ],
1438
+ "usage": {
1439
+ "prompt_tokens": 0,
1440
+ "completion_tokens": completion_tokens,
1441
+ "total_tokens": completion_tokens,
1442
+ },
1443
+ },
1444
+ )
1445
+
1446
+
1447
+def self_signed(cert_dir):
1448
+ """Paths of a self-signed certificate and key in cert_dir, made with
1449
+ openssl on first use."""
1450
+ os.makedirs(cert_dir, exist_ok=True)
1451
+ cert, key = os.path.join(cert_dir, "djev.crt"), os.path.join(cert_dir, "djev.key")
1452
+ if not (os.path.exists(cert) and os.path.exists(key)):
1453
+ subprocess.run(
1454
+ [
1455
+ "openssl",
1456
+ "req",
1457
+ "-x509",
1458
+ "-newkey",
1459
+ "rsa:2048",
1460
+ "-nodes",
1461
+ "-days",
1462
+ "3650",
1463
+ "-subj",
1464
+ "/CN=djev",
1465
+ "-keyout",
1466
+ key,
1467
+ "-out",
1468
+ cert,
1469
+ ],
1470
+ check=True,
1471
+ capture_output=True,
1472
+ )
1473
+ return cert, key
1474
+
1475
+
1476
+def serve_tls(host, port, cert_dir):
1477
+ cert, key = self_signed(cert_dir)
1478
+ ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
1479
+ ctx.load_cert_chain(cert, key)
1480
+ srv = ThreadingHTTPServer((host, port), Handler)
1481
+ srv.socket = ctx.wrap_socket(srv.socket, server_side=True)
1482
+ threading.Thread(target=srv.serve_forever, daemon=True).start()
1483
+ return srv
1484
+
1485
+
1486
+def main():
1487
+ global ARGS, CANVAS_LEN, CANVAS_STEP
1488
+ p = argparse.ArgumentParser()
1489
+ p.add_argument("--upstream", default="http://127.0.0.1:8010")
1490
+ p.add_argument("--model", default="dgemma")
1491
+ p.add_argument("--tokenizer", default="/models/dgemma", help="HF id or local path")
1492
+ p.add_argument("--canvas", type=int, default=64, help="the served canvas length")
1493
+ p.add_argument(
1494
+ "--canvas-step",
1495
+ type=int,
1496
+ default=16,
1497
+ help="request widths round up to a multiple of this",
1498
+ )
1499
+ p.add_argument("--host", default="0.0.0.0")
1500
+ p.add_argument("--port", type=int, default=8011)
1501
+ p.add_argument(
1502
+ "--tls-port", type=int, default=0, help="also listen with HTTPS here (0 = off)"
1503
+ )
1504
+ p.add_argument(
1505
+ "--cert-dir",
1506
+ default=os.path.expanduser("~/.cache/djev"),
1507
+ help="directory for the self-signed certificate",
1508
+ )
1509
+ ARGS = p.parse_args()
1510
+ CANVAS_LEN = ARGS.canvas
1511
+ CANVAS_STEP = ARGS.canvas_step
1512
+ init_tokenizer(AutoTokenizer.from_pretrained(ARGS.tokenizer))
1513
+ if ARGS.tls_port:
1514
+ serve_tls(ARGS.host, ARGS.tls_port, ARGS.cert_dir)
1515
+ print(
1516
+ f"structured server https on {ARGS.host}:{ARGS.tls_port} (self-signed)",
1517
+ flush=True,
1518
+ )
1519
+ print(
1520
+ f"structured server on {ARGS.host}:{ARGS.port} -> {ARGS.upstream} "
1521
+ f"(canvas {CANVAS_LEN})",
1522
+ flush=True,
1523
+ )
1524
+ ThreadingHTTPServer((ARGS.host, ARGS.port), Handler).serve_forever()
1525
+
1526
+
1527
+if __name__ == "__main__":
1528
+ main()
tests/config/test_model_arch_config.py
+25
@@ -15,8 +15,10 @@ from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
15
16
from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
17
from vllm.config.model_arch import ModelArchitectureConfig
18
+from vllm.transformers_utils.configs.diffusion_gemma import DiffusionGemmaConfig
19
from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
20
from vllm.transformers_utils.model_arch_config_convertor import (
21
+ MODEL_ARCH_CONFIG_CONVERTORS,
22
Gemma4ModelArchConfigConvertor,
23
ModelArchConfigConvertorBase,
24
)
@@ -496,3 +498,26 @@ def test_deepseek_v4_convertor_splits_vision_architecture():
498
conv = DeepseekV4ModelArchConfigConvertor(draft_cfg, draft_cfg)
499
assert conv.get_architectures() == [draft_arch]
500
assert draft_cfg.architectures == [draft_arch]
501
+
502
+
503
+def test_diffusion_gemma_head_dims_vary_by_layer_type():
504
+ """Checkpoints carry model_type `diffusion_gemma` at the top level. Under
505
+ the base convertor every layer reports the sliding head dim, Gemma4Config
506
+ skips its head-dim handling, and the full-attention KV-cache group falls to
507
+ FlashInfer, which cannot take a per-request causal tensor."""
508
+ text_config = dict(
509
+ num_hidden_layers=6,
510
+ hidden_size=64,
511
+ num_attention_heads=8,
512
+ num_key_value_heads=4,
513
+ head_dim=16,
514
+ global_head_dim=32,
515
+ layer_types=["sliding_attention"] * 5 + ["full_attention"],
516
+ )
517
+ hf_config = DiffusionGemmaConfig(text_config=text_config)
518
+
519
+ convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS[hf_config.model_type]
520
+ assert convertor_cls is Gemma4ModelArchConfigConvertor
521
+ arch = convertor_cls(hf_config, hf_config.text_config).convert()
522
+
523
+ assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32]
tests/test_sampling_params.py
+182
@@ -1,11 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
+from types import SimpleNamespace
5
6
import pytest
7
8
from vllm import SamplingParams
9
+from vllm.config.diffusion import DiffusionConfig
10
from vllm.exceptions import VLLMValidationError
11
+from vllm.utils.diffusion import validate_diffusion_sampling_params
12
+from vllm.v1.engine.input_processor import InputProcessor
13
14
15
@dataclass
@@ -13,6 +17,7 @@ class MockModelConfig:
17
is_diffusion: bool = False
18
max_logprobs: int = 20
19
logits_processors: list | None = None
20
+ return_sampling_mask: bool = False
21
22
def get_vocab_size(self) -> int:
23
return 1024
@@ -59,6 +64,183 @@ def test_verify_leaves_logits_processors_to_admission():
64
)
65
66
67
+def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None):
68
+ model_config = MockModelConfig(is_diffusion=True)
69
+ params.verify(model_config, None, None, None)
70
+ validate_diffusion_sampling_params(
71
+ params,
72
+ canvas_length=canvas_length,
73
+ vocab_size=model_config.get_vocab_size(),
74
+ async_scheduling=True,
75
+ )
76
+
77
+
78
+def test_diffusion_extra_args_are_validated_without_a_served_canvas():
79
+ # No --diffusion-config: the canvas is unknown, the ids are still checked
80
+ # and a read-only request is still normalised.
81
+ params = SamplingParams(
82
+ max_tokens=64,
83
+ extra_args={"diffusion_seed_canvas": [0, 1], "diffusion_read_only": True},
84
+ )
85
+ _verify_diffusion(params, canvas_length=None)
86
+ assert params.ignore_eos is True
87
+
88
+ bad = SamplingParams(extra_args={"diffusion_seed_canvas": [0, 10**9]})
89
+ with pytest.raises(VLLMValidationError, match="ids must be in"):
90
+ _verify_diffusion(bad, canvas_length=None)
91
+
92
+
93
+@pytest.mark.parametrize("async_scheduling", [False, True])
94
+@pytest.mark.parametrize(
95
+ "extra_args",
96
+ [
97
+ {},
98
+ {"diffusion_canvas_length": None},
99
+ {"diffusion_canvas_length": 4},
100
+ {"diffusion_canvas_length": 8},
101
+ ],
102
+)
103
+def test_narrow_diffusion_canvas_requires_async_scheduling(
104
+ async_scheduling, extra_args
105
+):
106
+ processor = SimpleNamespace(
107
+ model_config=MockModelConfig(is_diffusion=True),
108
+ vllm_config=SimpleNamespace(
109
+ scheduler_config=SimpleNamespace(async_scheduling=async_scheduling)
110
+ ),
111
+ speculative_config=None,
112
+ structured_outputs_config=None,
113
+ diffusion_config=DiffusionConfig(canvas_length=8),
114
+ tokenizer=None,
115
+ validate_logits_processors_params=lambda params: None,
116
+ )
117
+ params = SamplingParams(extra_args=extra_args)
118
+ if not async_scheduling and extra_args.get("diffusion_canvas_length") == 4:
119
+ with pytest.raises(VLLMValidationError, match="requires --async-scheduling"):
120
+ InputProcessor._validate_params(processor, params, ("generate",))
121
+ else:
122
+ InputProcessor._validate_params(processor, params, ("generate",))
123
+
124
+
125
+@pytest.mark.parametrize(
126
+ "extra_args, match",
127
+ [
128
+ ({"diffusion_seed_canvas": "abc"}, "list of token ids"),
129
+ ({"diffusion_seed_canvas": [1, 2.5]}, "list of token ids"),
130
+ ({"diffusion_seed_canvas": [1, True]}, "list of token ids"),
131
+ ({"diffusion_seed_canvas": [0, 1024]}, r"in \[0, 1024\)"),
132
+ ({"diffusion_seed_canvas": [0, -1]}, r"in \[0, 1024\)"),
133
+ ({"diffusion_max_steps": 0}, "positive integer"),
134
+ ({"diffusion_max_steps": "1"}, "positive integer"),
135
+ ({"diffusion_max_steps": True}, "positive integer"),
136
+ ({"diffusion_read_only": "yes"}, "boolean"),
137
+ ({"diffusion_read_only": 2}, "boolean"),
138
+ ({"diffusion_pinned": "0,1"}, "list of canvas positions"),
139
+ ({"diffusion_pinned": [0, True]}, "list of canvas positions"),
140
+ ({"diffusion_pinned": [0]}, "needs a diffusion_seed_canvas"),
141
+ ],
142
+)
143
+def test_diffusion_rejects_bad_extra_args(extra_args: dict, match: str):
144
+ with pytest.raises(VLLMValidationError, match=match):
145
+ _verify_diffusion(SamplingParams(extra_args=extra_args))
146
+
147
+
148
+def test_diffusion_seed_canvas_must_fill_the_canvas():
149
+ params = SamplingParams(extra_args={"diffusion_seed_canvas": [0] * 7})
150
+ with pytest.raises(VLLMValidationError, match="exactly 8 ids, got 7"):
151
+ _verify_diffusion(params, canvas_length=8)
152
+ # Without the served diffusion config the canvas length is unknown.
153
+ _verify_diffusion(params)
154
+
155
+
156
+@pytest.mark.parametrize("max_tokens, expected", [(100, 8), (5, 5), (None, 8)])
157
+@pytest.mark.parametrize("flag", [True, 1])
158
+def test_diffusion_read_only_ends_after_one_canvas(max_tokens, expected, flag):
159
+ params = SamplingParams(
160
+ max_tokens=max_tokens, extra_args={"diffusion_read_only": flag}
161
+ )
162
+ _verify_diffusion(params, canvas_length=8)
163
+ assert params.max_tokens == expected
164
+ assert params.ignore_eos
165
+
166
+
167
+@pytest.mark.parametrize(
168
+ "width, match",
169
+ [
170
+ (0, "positive integer"),
171
+ ("4", "positive integer"),
172
+ (True, "positive integer"),
173
+ (9, "no larger"),
174
+ ],
175
+)
176
+def test_diffusion_rejects_bad_canvas_length(width, match):
177
+ with pytest.raises(VLLMValidationError, match=match):
178
+ _verify_diffusion(
179
+ SamplingParams(extra_args={"diffusion_canvas_length": width}),
180
+ canvas_length=8,
181
+ )
182
+
183
+
184
+def test_diffusion_canvas_length_sizes_the_seed_and_the_read():
185
+ params = SamplingParams(
186
+ max_tokens=100,
187
+ extra_args={
188
+ "diffusion_canvas_length": 4,
189
+ "diffusion_seed_canvas": [1, 2, 3, 4],
190
+ "diffusion_read_only": True,
191
+ },
192
+ )
193
+ _verify_diffusion(params, canvas_length=8)
194
+ assert params.max_tokens == 4
195
+
196
+ params = SamplingParams(
197
+ extra_args={"diffusion_canvas_length": 4, "diffusion_seed_canvas": [0] * 8}
198
+ )
199
+ with pytest.raises(VLLMValidationError, match="exactly 4 ids, got 8"):
200
+ _verify_diffusion(params, canvas_length=8)
201
+
202
+
203
+def test_diffusion_pinned_positions_stay_inside_the_canvas():
204
+ seed = [0] * 8
205
+ with pytest.raises(VLLMValidationError, match="inside the canvas"):
206
+ _verify_diffusion(
207
+ SamplingParams(
208
+ extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [7, 8]}
209
+ ),
210
+ canvas_length=8,
211
+ )
212
+ with pytest.raises(VLLMValidationError, match="inside the canvas"):
213
+ _verify_diffusion(
214
+ SamplingParams(
215
+ extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [-1]}
216
+ )
217
+ )
218
+ # A narrower request canvas bounds the positions.
219
+ with pytest.raises(VLLMValidationError, match="inside the canvas"):
220
+ _verify_diffusion(
221
+ SamplingParams(
222
+ extra_args={
223
+ "diffusion_canvas_length": 4,
224
+ "diffusion_seed_canvas": seed[:4],
225
+ "diffusion_pinned": [4],
226
+ }
227
+ ),
228
+ canvas_length=8,
229
+ )
230
+
231
+
232
+def test_diffusion_accepts_extra_args():
233
+ params = SamplingParams(
234
+ extra_args={
235
+ "diffusion_seed_canvas": list(range(8)),
236
+ "diffusion_pinned": [0, 1, 7],
237
+ "diffusion_max_steps": 4,
238
+ "diffusion_read_only": True,
239
+ }
240
+ )
241
+ _verify_diffusion(params, canvas_length=8)
242
+
243
+
244
@pytest.mark.parametrize("value", [-(2**63) - 1, 2**64])
245
def test_extra_args_rejects_nested_integer_overflow(value):
246
"""Reject extension values before they reach the engine transport."""
tests/v1/core/test_scheduler.py
+130
@@ -40,6 +40,10 @@ from vllm.utils.hashing import sha256
40
from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
41
from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
42
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
43
+from vllm.v1.core.sched.diffusion_scheduler import (
44
+ DiffusionAsyncScheduler,
45
+ diffusion_canvas_width,
46
+)
47
from vllm.v1.core.sched.interface import PauseState
48
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
49
from vllm.v1.core.sched.scheduler import Scheduler
@@ -6872,3 +6876,129 @@ def test_update_draft_token_ids_in_output_strips_padding():
6876
-1,
6877
]
6878
assert scheduler_output.num_invalid_spec_tokens == {request.request_id: 2}
6879
+
6880
+
6881
+def test_diffusion_canvas_width_defaults_to_the_served_canvas():
6882
+ def req(extra):
6883
+ return SimpleNamespace(sampling_params=SimpleNamespace(extra_args=extra))
6884
+
6885
+ assert diffusion_canvas_width(req(None), 64) == 64
6886
+ assert diffusion_canvas_width(req({}), 64) == 64
6887
+ assert diffusion_canvas_width(req({"diffusion_canvas_length": 16}), 64) == 16
6888
+ assert diffusion_canvas_width(SimpleNamespace(sampling_params=None), 64) == 64
6889
+
6890
+
6891
+def _diffusion_request(req_id: str, extra_args: dict) -> Request:
6892
+ (request,) = create_requests(
6893
+ num_requests=1, num_tokens=8, max_tokens=64, req_ids=[req_id]
6894
+ )
6895
+ request.sampling_params.extra_args = extra_args
6896
+ return request
6897
+
6898
+
6899
+def _diffusion_scheduler(**kwargs) -> DiffusionAsyncScheduler:
6900
+ scheduler = create_scheduler(
6901
+ async_scheduling=True,
6902
+ diffusion_canvas_length=8,
6903
+ scheduler_cls=DiffusionAsyncScheduler,
6904
+ **kwargs,
6905
+ )
6906
+ assert isinstance(scheduler, DiffusionAsyncScheduler)
6907
+ return scheduler
6908
+
6909
+
6910
+@pytest.mark.parametrize("async_scheduling", [False, True])
6911
+def test_diffusion_scheduler_is_selected_by_default(async_scheduling):
6912
+ config = create_scheduler(
6913
+ async_scheduling=async_scheduling, diffusion_canvas_length=8
6914
+ ).vllm_config.scheduler_config
6915
+ assert config.get_scheduler_cls() is (
6916
+ DiffusionAsyncScheduler if config.async_scheduling else Scheduler
6917
+ )
6918
+
6919
+
6920
+def test_diffusion_scheduler_narrows_the_canvas_per_request():
6921
+ scheduler = _diffusion_scheduler()
6922
+ wide = _diffusion_request("wide", {})
6923
+ narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4})
6924
+ scheduler.add_request(wide)
6925
+ scheduler.add_request(narrow)
6926
+
6927
+ # The prefill step lays down the served canvas as placeholders.
6928
+ scheduler.schedule()
6929
+ output = scheduler.schedule()
6930
+ assert output.scheduled_spec_decode_tokens["wide"] == [-1] * 8
6931
+ assert output.scheduled_spec_decode_tokens["narrow"] == [-1] * 4
6932
+ assert output.num_scheduled_tokens["narrow"] == 4
6933
+ assert narrow.num_output_placeholders == 4
6934
+
6935
+
6936
+@pytest.mark.parametrize("structured", [False, True])
6937
+def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
6938
+ """Padded worker drafts must be narrowed before scheduling or grammar validation."""
6939
+ scheduler = _diffusion_scheduler()
6940
+ wide = _diffusion_request("wide", {})
6941
+ narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4})
6942
+ for request in (wide, narrow):
6943
+ scheduler.add_request(request)
6944
+ prefill = scheduler.schedule()
6945
+ _model_output(scheduler, prefill, [[], []])
6946
+
6947
+ if structured:
6948
+ for request in (wide, narrow):
6949
+ request.structured_output_request = SimpleNamespace(
6950
+ grammar=_RecordingGrammar(), reasoning_ended=True
6951
+ )
6952
+ tokens = list(range(8)) if structured else [-1] * 8
6953
+ drafts = DraftTokenIds(["wide", "narrow"], [tokens.copy(), tokens.copy()])
6954
+ output = scheduler.schedule()
6955
+ scheduler.update_draft_token_ids_in_output(drafts, output)
6956
+
6957
+ assert output.scheduled_spec_decode_tokens == {
6958
+ "wide": tokens,
6959
+ "narrow": tokens[:4],
6960
+ }
6961
+ assert output.num_scheduled_tokens == {"wide": 8, "narrow": 4}
6962
+ if structured:
6963
+ assert wide.structured_output_request.grammar.seen == [tokens]
6964
+ assert narrow.structured_output_request.grammar.seen == [tokens[:4]]
6965
+
6966
+
6967
+def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight():
6968
+ scheduler = _diffusion_scheduler()
6969
+ one = _diffusion_request(
6970
+ "one", {"diffusion_read_only": True, "diffusion_max_steps": 1}
6971
+ )
6972
+ two = _diffusion_request(
6973
+ "two", {"diffusion_read_only": True, "diffusion_max_steps": 2}
6974
+ )
6975
+ # Only reads are deferred. A capped generation keeps its extra async step.
6976
+ gen = _diffusion_request("gen", {"diffusion_max_steps": 1})
6977
+ for request in (one, two, gen):
6978
+ scheduler.add_request(request)
6979
+
6980
+ scheduler.schedule()
6981
+ assert set(scheduler.schedule().num_scheduled_tokens) == {"one", "two", "gen"}
6982
+ assert one.num_output_placeholders == 8
6983
+ assert set(scheduler.schedule().num_scheduled_tokens) == {"two", "gen"}
6984
+ assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"}
6985
+ # The deferral is renewed every call.
6986
+ assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"}
6987
+
6988
+
6989
+def test_diffusion_read_deferral_keeps_a_longer_pp_wait():
6990
+ scheduler = _diffusion_scheduler(pipeline_parallel_size=3, use_v2_model_runner=True)
6991
+ read = _diffusion_request(
6992
+ "read", {"diffusion_read_only": True, "diffusion_max_steps": 1}
6993
+ )
6994
+ scheduler.add_request(read)
6995
+
6996
+ scheduler.schedule()
6997
+ assert read.next_decode_eligible_step == 4
6998
+ for _ in range(2):
6999
+ assert "read" not in scheduler.schedule().num_scheduled_tokens
7000
+ assert "read" in scheduler.schedule().num_scheduled_tokens
7001
+ assert read.next_decode_eligible_step == 7
7002
+ # Deferring this step alone would ask for 6. The PP wait to 7 stands.
7003
+ assert "read" not in scheduler.schedule().num_scheduled_tokens
7004
+ assert read.next_decode_eligible_step == 7
tests/v1/core/utils.py
+12
-1
@@ -7,6 +7,7 @@ import vllm.envs as envs
7
from tests.v1.kv_connector.unit.utils import MockKVConfig
8
from vllm.config import (
9
CacheConfig,
10
+ DiffusionConfig,
11
ECTransferConfig,
12
KVTransferConfig,
13
ModelConfig,
@@ -77,6 +78,8 @@ def create_scheduler(
78
kv_cache_spec: KVCacheSpec | None = None,
79
per_request_spec_decode_metrics: str = "none",
80
scheduling_policy: SchedulerPolicy = "fcfs",
81
+ diffusion_canvas_length: int | None = None,
82
+ scheduler_cls: type[Scheduler] | None = None,
83
) -> Scheduler | AsyncScheduler:
84
"""Create scheduler under test.
85
@@ -178,6 +181,12 @@ def create_scheduler(
181
else None
182
)
183
184
+ diffusion_config: DiffusionConfig | None = None
185
+ if diffusion_canvas_length is not None:
186
+ # A diffusion checkpoint declares its canvas in the HF config.
187
+ model_config.hf_config.canvas_length = diffusion_canvas_length
188
+ diffusion_config = DiffusionConfig(canvas_length=diffusion_canvas_length)
189
+
190
vllm_config = VllmConfig(
191
scheduler_config=scheduler_config,
192
model_config=model_config,
@@ -188,6 +197,7 @@ def create_scheduler(
197
),
198
kv_transfer_config=kv_transfer_config,
199
speculative_config=speculative_config,
200
+ diffusion_config=diffusion_config,
201
ec_transfer_config=ec_transfer_config,
202
observability_config=ObservabilityConfig(
203
per_request_spec_decode_metrics=per_request_spec_decode_metrics,
@@ -207,7 +217,8 @@ def create_scheduler(
217
)
218
cache_config.num_gpu_blocks = num_blocks
219
register_all_kvcache_specs(vllm_config)
210
- scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
220
+ if scheduler_cls is None:
221
+ scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
222
scheduler = scheduler_cls(
223
vllm_config=vllm_config,
224
kv_cache_config=kv_cache_config,
tests/v1/engine/test_input_processor_trace_replay.py
+2
@@ -80,10 +80,12 @@ def _validate(enable_trace_replay: bool) -> None:
80
model_config=SimpleNamespace(
81
return_sampling_mask=False,
82
enable_trace_replay=enable_trace_replay,
83
+ is_diffusion=False,
84
),
85
vllm_config=SimpleNamespace(reasoning_config=None),
86
speculative_config=None,
87
structured_outputs_config=None,
88
+ diffusion_config=None,
89
tokenizer=None,
90
validate_logits_processors_params=lambda params: None,
91
)
tests/v1/sample/test_diffusion_gemma_reads.py
new
+319
@@ -0,0 +1,319 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Per-request DiffusionGemma state behind structured reads: seed canvases,
4
+pinned positions, read-only slots and the per-slot step cap."""
5
+
6
+import numpy as np
7
+import pytest
8
+import torch
9
+
10
+from vllm.model_executor.models.diffusion_gemma import (
11
+ DiffusionGemmaRequestStates,
12
+ _compiled_sample_step,
13
+ _concat_logprob_stashes,
14
+)
15
+from vllm.platforms import current_platform
16
+from vllm.v1.outputs import LogprobsTensors
17
+
18
+pytestmark = pytest.mark.skipif(
19
+ not current_platform.is_cuda(), reason="the sampler state lives on the GPU"
20
+)
21
+
22
+CL = 8
23
+VOCAB = 64
24
+MAX_REQS = 4
25
+MAX_STEPS = 48
26
+
27
+
28
+def _states() -> DiffusionGemmaRequestStates:
29
+ return DiffusionGemmaRequestStates(
30
+ max_num_reqs=MAX_REQS,
31
+ canvas_length=CL,
32
+ vocab_size=VOCAB,
33
+ max_denoising_steps=MAX_STEPS,
34
+ device=torch.device("cuda"),
35
+ hidden_size=4,
36
+ stability_threshold=2,
37
+ )
38
+
39
+
40
+def _slots(*idx: int) -> tuple[np.ndarray, torch.Tensor]:
41
+ slots = np.array(idx, dtype=np.int64)
42
+ return slots, torch.tensor(slots, device="cuda")
43
+
44
+
45
+def test_seed_canvas_replaces_only_seeded_slots():
46
+ states = _states()
47
+ for slot in range(3):
48
+ states.add_request(slot)
49
+ seed = list(range(CL))
50
+ states.set_seed_canvas(1, seed)
51
+ slots, slots_gpu = _slots(0, 1, 2)
52
+ states.init_canvas(slots_gpu)
53
+ before = states.canvas[slots_gpu].clone()
54
+
55
+ states.apply_seed_canvases(slots, slots_gpu)
56
+
57
+ after = states.canvas[slots_gpu]
58
+ assert after[1].tolist() == seed
59
+ assert torch.equal(after[0], before[0])
60
+ assert torch.equal(after[2], before[2])
61
+
62
+
63
+def test_apply_seed_canvases_leaves_unseeded_batches_alone():
64
+ states = _states()
65
+ states.add_request(0)
66
+ slots, slots_gpu = _slots(0)
67
+ states.init_canvas(slots_gpu)
68
+ before = states.canvas[0].clone()
69
+
70
+ states.apply_seed_canvases(slots, slots_gpu)
71
+
72
+ assert torch.equal(states.canvas[0], before)
73
+
74
+
75
+def test_add_request_clears_seed_and_read_only():
76
+ states = _states()
77
+ states.add_request(0)
78
+ states.set_seed_canvas(0, [1] * CL)
79
+ states.set_read_only(0)
80
+ assert states.seeded_slots == {0}
81
+ assert states.read_only_slots == {0}
82
+
83
+ states.add_request(0)
84
+
85
+ assert not states.seeded_slots
86
+ assert not states.read_only_slots
87
+ assert not bool(states.has_seed[0])
88
+ assert not bool(states.read_only[0])
89
+
90
+
91
+def test_canvas_width_resets_with_the_slot():
92
+ states = _states()
93
+ states.add_request(0)
94
+ states.canvas_width_np[0] = 4
95
+ states.set_seed_canvas(0, [7, 7, 7, 7])
96
+ assert states.seed_canvas[0, :4].tolist() == [7, 7, 7, 7]
97
+
98
+ states.add_request(0)
99
+
100
+ assert states.canvas_width_np[0] == CL
101
+
102
+
103
+def test_remove_request_forgets_the_slot():
104
+ states = _states()
105
+ states.add_request(0)
106
+ states.set_seed_canvas(0, [1] * CL)
107
+ states.set_read_only(0)
108
+
109
+ states.remove_request(0)
110
+
111
+ assert not states.seeded_slots
112
+ assert not states.read_only_slots
113
+
114
+
115
+def _denoise_once(
116
+ states: DiffusionGemmaRequestStates,
117
+ slots: list[int],
118
+ compute_sc: bool = True,
119
+ width: int = CL,
120
+ embed_weight: torch.Tensor | None = None,
121
+ embed_dtype: torch.dtype = torch.float32,
122
+) -> tuple[torch.Tensor, torch.Tensor]:
123
+ """One compiled denoise step over ``slots`` with flat logits, so nothing
124
+ converges by stability or confidence and only the step cap can end it.
125
+ ``width`` below CL runs the step on [:, :width] views, as the sampler
126
+ does for a narrow tile."""
127
+ n = len(slots)
128
+ device = states.device
129
+ decode_slots = torch.tensor(slots, dtype=torch.int64, device=device)
130
+ decode_idx = torch.arange(n, dtype=torch.int64, device=device)
131
+ sampled = torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width]
132
+ num_sampled = torch.zeros(n, dtype=torch.int32, device=device)
133
+ _compiled_sample_step(
134
+ torch.zeros(n * width, VOCAB, device=device),
135
+ decode_slots,
136
+ decode_idx,
137
+ decode_slots,
138
+ torch.full((n,), width, dtype=torch.int64, device=device),
139
+ states.canvas[:, :width],
140
+ states.argmax_canvas[:, :width],
141
+ states.step,
142
+ states.is_encoder_phase,
143
+ states.confident,
144
+ states.self_conditioning_embeds[:, :width],
145
+ (
146
+ torch.zeros(VOCAB, 4, dtype=embed_dtype, device=device)
147
+ if embed_weight is None
148
+ else embed_weight
149
+ ),
150
+ torch.tensor(1.0, dtype=embed_dtype, device=device),
151
+ states.accepted_canvas_history[:, :, :width],
152
+ states.accepted_canvas_history_len,
153
+ states.max_steps,
154
+ states.pin_mask[:, :width],
155
+ states.seed_canvas[:, :width],
156
+ states.read_only,
157
+ sampled,
158
+ num_sampled,
159
+ torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device),
160
+ max_denoising_steps=float(MAX_STEPS),
161
+ t_min=0.5,
162
+ t_max=1.0,
163
+ confidence_threshold=0.1,
164
+ vocab_size=VOCAB,
165
+ CL=width,
166
+ ST=states.stability_threshold,
167
+ entropy_bound=0.1,
168
+ sc_vocab_start=0,
169
+ sc_vocab_end=VOCAB,
170
+ tp_size=1,
171
+ tp_group_name="",
172
+ compute_sc=compute_sc,
173
+ )
174
+ return sampled, num_sampled
175
+
176
+
177
+def test_pinned_positions_hold_their_seed_through_a_step():
178
+ states = _states()
179
+ states.add_request(0)
180
+ states.is_encoder_phase[0] = False
181
+ seed = list(range(10, 10 + CL))
182
+ states.set_seed_canvas(0, seed)
183
+ states.canvas[0] = torch.tensor(seed, device="cuda")
184
+ states.set_pins(0, [0, 1, 2, 3])
185
+
186
+ # Flat logits accept nothing, so every free position is renoised.
187
+ _denoise_once(states, [0], embed_weight=torch.ones(VOCAB, 4, device="cuda"))
188
+
189
+ assert states.canvas[0, :4].tolist() == seed[:4]
190
+ # The soft embed is zero at pinned positions and non-zero elsewhere.
191
+ sc = states.self_conditioning_embeds[0]
192
+ assert not sc[:4].any()
193
+ assert sc[4:].any()
194
+
195
+
196
+def test_add_request_clears_pins():
197
+ states = _states()
198
+ states.add_request(0)
199
+ states.set_seed_canvas(0, [1] * CL)
200
+ states.set_pins(0, [2, 3])
201
+ assert states.pin_mask[0].tolist() == [False, False, True, True] + [False] * 4
202
+
203
+ states.add_request(0)
204
+
205
+ assert not states.pin_mask[0].any()
206
+
207
+
208
+def test_single_step_tile_skips_self_conditioning():
209
+ states = _states()
210
+ states.add_request(0)
211
+ states.is_encoder_phase[0] = False
212
+ states.self_conditioning_embeds[0] = 1.0
213
+
214
+ _denoise_once(states, [0], compute_sc=False)
215
+
216
+ assert not states.self_conditioning_embeds[0].any()
217
+
218
+
219
+@pytest.mark.parametrize("stance", ["default", "force_eager"])
220
+def test_self_conditioning_stores_a_bf16_model_in_the_fp32_buffer(stance):
221
+ # The model's embeddings are bf16 while the buffer is fp32. Compiled code
222
+ # casts on the store, but eager does not, and the step runs eager once
223
+ # torch.compile hits its recompile limit.
224
+ states = _states()
225
+ states.add_request(0)
226
+ states.is_encoder_phase[0] = False
227
+
228
+ with torch.compiler.set_stance(stance):
229
+ _denoise_once(
230
+ states,
231
+ [0],
232
+ embed_weight=torch.ones(VOCAB, 4, dtype=torch.bfloat16, device="cuda"),
233
+ embed_dtype=torch.bfloat16,
234
+ )
235
+
236
+ assert states.self_conditioning_embeds.dtype == torch.float32
237
+ # uniform probs @ an all-ones embedding
238
+ assert torch.allclose(
239
+ states.self_conditioning_embeds[0], torch.ones(CL, 4, device="cuda")
240
+ )
241
+
242
+
243
+def test_narrow_tile_leaves_the_rest_of_the_canvas_alone():
244
+ states = _states()
245
+ states.add_request(0)
246
+ states.is_encoder_phase[0] = False
247
+ states.canvas[0] = 5
248
+ states.argmax_canvas[0] = 5
249
+
250
+ _denoise_once(states, [0], width=4)
251
+
252
+ # Columns past the width are untouched; the step counted.
253
+ assert states.canvas[0, 4:].tolist() == [5] * (CL - 4)
254
+ assert states.argmax_canvas[0, 4:].tolist() == [5] * (CL - 4)
255
+ assert states.step[0] == 1
256
+
257
+
258
+def test_step_cap_is_per_slot():
259
+ states = _states()
260
+ for slot in (0, 1):
261
+ states.add_request(slot)
262
+ states.is_encoder_phase[slot] = False
263
+ states.max_steps[0] = 1
264
+
265
+ _denoise_once(states, [0, 1])
266
+
267
+ # Slot 0 hit its cap and moves to commit. Slot 1 keeps denoising.
268
+ assert states.is_encoder_phase[:2].tolist() == [True, False]
269
+ assert states.step[:2].tolist() == [1, 1]
270
+
271
+
272
+def _stash(rows: int, width: int, first_id: int) -> LogprobsTensors:
273
+ ids = torch.arange(first_id, first_id + rows * width).reshape(rows, width)
274
+ return LogprobsTensors(
275
+ logprob_token_ids=ids,
276
+ logprobs=-ids.float(),
277
+ selected_token_ranks=torch.zeros(rows, dtype=torch.int64),
278
+ )
279
+
280
+
281
+def test_stashes_of_different_widths_join():
282
+ # A read that asked for 10 label ids and a generation that asked for 2
283
+ # logprobs converged in different steps, so their stashes differ in width.
284
+ wide, narrow = _stash(2, 11, 100), _stash(3, 3, 0)
285
+
286
+ out = _concat_logprob_stashes([narrow, wide], [0, 3])
287
+
288
+ assert out.logprob_token_ids.shape == (5, 11)
289
+ assert out.logprobs.shape == (5, 11)
290
+ assert out.cu_num_generated_tokens == [0, 3]
291
+ # The narrow rows keep their values and pad with id 0 at -inf.
292
+ assert torch.equal(out.logprob_token_ids[:3, :3], narrow.logprob_token_ids)
293
+ assert not out.logprob_token_ids[:3, 3:].any()
294
+ assert torch.isneginf(out.logprobs[:3, 3:]).all()
295
+ assert torch.equal(out.logprobs[3:], wide.logprobs)
296
+
297
+
298
+@pytest.mark.parametrize("width", [4, CL])
299
+@pytest.mark.parametrize("steps", [1, 2])
300
+def test_read_emits_at_convergence_while_generation_waits_for_commit(width, steps):
301
+ states = _states()
302
+ for slot in (2, 0):
303
+ states.add_request(slot)
304
+ states.is_encoder_phase[slot] = False
305
+ states.max_steps[slot] = steps
306
+ states.set_read_only(2)
307
+
308
+ for step in range(steps):
309
+ sampled, counts = _denoise_once(states, [2, 0], width=width)
310
+ assert counts.tolist() == ([width, 0] if step == steps - 1 else [0, 0])
311
+
312
+ assert torch.equal(sampled[0], states.argmax_canvas[2, :width].int())
313
+ assert not states.is_encoder_phase[2]
314
+ assert states.is_encoder_phase[0]
315
+ assert not states.self_conditioning_embeds[2].any()
316
+
317
+ sampled, counts = _denoise_once(states, [0], width=width)
318
+ assert counts.tolist() == [width]
319
+ assert torch.equal(sampled[0], states.argmax_canvas[0, :width].int())
vllm/config/scheduler.py
+11
-7
@@ -232,13 +232,17 @@ class SchedulerConfig:
232
233
# The first half of this warning can be removed once the Scheduler interface is
234
# finalized and we can maintain support for scheduler classes that implement it
235
- logger.warning_once(
236
- "Using custom scheduler class %s. This scheduler interface is not public "
237
- "and compatibility may not be maintained. If you have subclassed Scheduler "
238
- "instead of AsyncScheduler, you will see degraded performance due to async "
239
- "scheduling being disabled.",
240
- self.scheduler_cls, # type: ignore[arg-type]
241
- )
235
+ if not (
236
+ isinstance(self.scheduler_cls, str)
237
+ and self.scheduler_cls.startswith("vllm.")
238
+ ):
239
+ logger.warning_once(
240
+ "Using custom scheduler class %s. This scheduler interface is not "
241
+ "public and compatibility may not be maintained. If you have "
242
+ "subclassed Scheduler instead of AsyncScheduler, you will see "
243
+ "degraded performance due to async scheduling being disabled.",
244
+ self.scheduler_cls, # type: ignore[arg-type]
245
+ )
246
if not isinstance(self.scheduler_cls, str):
247
return cast(type["SchedulerInterface"], self.scheduler_cls)
248
return resolve_obj_by_qualname(self.scheduler_cls)
vllm/config/vllm.py
+10
@@ -1894,6 +1894,16 @@ class VllmConfig:
1894
)
1895
current_platform.check_and_update_config(self)
1896
1897
+ # After the platform hook, which has the last word on async scheduling.
1898
+ if (
1899
+ self.diffusion_config is not None
1900
+ and self.scheduler_config.async_scheduling
1901
+ and self.scheduler_config.scheduler_cls is None
1902
+ ):
1903
+ self.scheduler_config.scheduler_cls = (
1904
+ "vllm.v1.core.sched.diffusion_scheduler.DiffusionAsyncScheduler"
1905
+ )
1906
+
1907
self._normalize_piecewise_cudagraph_mode(
1908
breakable_cudagraph_enabled=breakable_cudagraph_enabled
1909
)
vllm/model_executor/models/config.py
+6
-3
@@ -280,11 +280,14 @@ class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig):
280
"FLASH_ATTN or TRITON_ATTN instead."
281
)
282
if attention_config.backend is None and not attention_config.use_non_causal:
283
+ # Only the sm_100 default backend order reads this flag (it puts
284
+ # FlashInfer first for causal models). Elsewhere Gemma4Config above
285
+ # keeps FlashInfer out: FA4 on every layer, else TRITON_ATTN.
286
attention_config.use_non_causal = True
287
logger.info(
285
- "DiffusionGemma uses mixed causal/bidirectional attention "
286
- "within a batch; setting use_non_causal=True to exclude "
287
- "FlashInfer from auto-selection."
288
+ "DiffusionGemma mixes causal and bidirectional attention in "
289
+ "one batch; setting use_non_causal=True so the default backend "
290
+ "order prefers FlashAttention over FlashInfer."
291
)
292
293
# Auto-create DiffusionConfig from HF config if not provided.
vllm/model_executor/models/diffusion_gemma.py
+321
-146
@@ -22,6 +22,7 @@ from typing import Any
22
23
import numpy as np
24
import torch
25
+import torch._dynamo
26
from torch import nn
27
from torch.nn import functional as F
28
from transformers import AutoModel
@@ -371,6 +372,34 @@ def _compute_num_rejected(
372
return torch.where(is_denoise, query_lens, num_rejected)
373
374
375
+def _concat_logprob_stashes(
376
+ parts: list[LogprobsTensors], cu_num_generated_tokens: list[int]
377
+) -> LogprobsTensors:
378
+ """Join the logprobs stashed for the requests committing this step.
379
+
380
+ Each stash is as wide as the widest logprobs request in the batch at the
381
+ step that request converged, so stashes from different steps can differ
382
+ in width. Pad the narrow ones the way compute_topk_scores pads a mixed
383
+ batch: token id 0 at -inf, which the output processor never reports.
384
+ """
385
+ width = max(p.logprob_token_ids.shape[1] for p in parts)
386
+
387
+ def pad(t: torch.Tensor, value: float) -> torch.Tensor:
388
+ return F.pad(t, (0, width - t.shape[1]), value=value)
389
+
390
+ return LogprobsTensors(
391
+ logprob_token_ids=torch.cat([pad(p.logprob_token_ids, 0) for p in parts]),
392
+ logprobs=torch.cat([pad(p.logprobs, float("-inf")) for p in parts]),
393
+ selected_token_ranks=torch.cat([p.selected_token_ranks for p in parts]),
394
+ cu_num_generated_tokens=cu_num_generated_tokens,
395
+ )
396
+
397
+
398
+# Tiles specialize the graph on batch size 1, on a tile as wide as the canvas,
399
+# on compute_sc and on sizes that coincide with the state buffers. That set
400
+# is small but passes Dynamo's default of 8, after which every step would run
401
+# eager: about twice as slow for a self-conditioned step.
402
+@torch._dynamo.config.patch(recompile_limit=64)
403
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
404
def _compiled_sample_step(
405
# Logits from the model [num_decode * CL, vocab]
@@ -391,6 +420,10 @@ def _compiled_sample_step(
420
normalizer: torch.Tensor,
421
history: torch.Tensor, # [max_num_reqs, ST, CL]
422
history_len_tensor: torch.Tensor, # [max_num_reqs]
423
+ max_steps_tensor: torch.Tensor, # [max_num_reqs] int32, per-slot step cap
424
+ pin_mask: torch.Tensor, # [max_num_reqs, CL] bool, positions held at the seed
425
+ seed_canvas: torch.Tensor, # [max_num_reqs, CL]
426
+ read_only: torch.Tensor, # [max_num_reqs] bool, emit without a commit forward
427
# Output tensors (modified in-place)
428
sampled: torch.Tensor, # [num_reqs, CL]
429
num_sampled: torch.Tensor, # [num_reqs]
@@ -412,6 +445,7 @@ def _compiled_sample_step(
445
sc_vocab_end: int,
446
tp_size: int,
447
tp_group_name: str,
448
+ compute_sc: bool = True,
449
) -> torch.Tensor:
450
"""Compiled decode step: temperature → Gumbel sample → probs/confidence →
451
accept/renoise → convergence, all as vectorized PyTorch ops.
@@ -478,8 +512,11 @@ def _compiled_sample_step(
512
0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype
513
)
514
481
- # Compute denoise canvas (accept/renoise)
515
+ # Compute denoise canvas (accept/renoise). Pinned positions keep the seed.
516
denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens)
517
+ denoise_canvas = torch.where(
518
+ pin_mask[decode_slots], seed_canvas[decode_slots], denoise_canvas
519
+ )
520
521
# Canvas: commit → random reinit, denoise → accept/renoise result
522
canvas[decode_slots] = torch.where(
@@ -504,16 +541,6 @@ def _compiled_sample_step(
541
new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode))
542
history_len_tensor[decode_slots] = new_hist_len
543
507
- # Sampled output: commit → emit argmax_canvas, denoise → 0 (pre-zeroed)
508
- sampled[decode_idx] = argmax_canvas[decode_slots].to(
509
- sampled.dtype
510
- ) * is_commit.unsqueeze(1).to(sampled.dtype)
511
- # Commit only the real canvas length (== CL except for a canvas truncated
512
- # near max_model_len); the padded tail positions are never emitted.
513
- num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * valid_canvas_len.to(
514
- num_sampled.dtype
515
- )
516
-
544
# ---- Phase 6: Stability + convergence ----
545
ref = history[decode_slots, 0]
546
mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32)
@@ -523,13 +550,17 @@ def _compiled_sample_step(
550
551
step_after = step_tensor[decode_slots]
552
converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | (
526
- step_after >= max_denoising_steps
553
+ step_after >= max_steps_tensor[decode_slots]
554
)
555
# Commit done → denoise next (False); denoise converged → commit next (True)
556
is_encoder_phase[decode_slots] = torch.where(
557
is_commit, is_commit.new_zeros(num_decode), converged
558
)
559
560
+ emit = is_commit | (is_denoise & converged & read_only[decode_slots])
561
+ sampled[decode_idx] = argmax_canvas[decode_slots].to(sampled.dtype) * emit[:, None]
562
+ num_sampled[decode_idx] = (emit * valid_canvas_len).to(num_sampled.dtype)
563
+
564
# SC soft embedding: store ``probs @ embed_weight`` (the value the next step's
565
# self-conditioning MLP consumes) only for slots that will denoise next — i.e.
566
# this step denoised AND it isn't about to commit (is_encoder_phase now False).
@@ -537,24 +568,41 @@ def _compiled_sample_step(
568
# sc_embeds directly. Storing the [.., hidden] soft embed instead of the full
569
# [.., vocab] probs avoids a giant persistent buffer.
570
sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None]
540
- # Self-conditioning soft embed = probs @ embed_tokens.weight. Under tensor
541
- # parallelism the embedding is vocab-sharded ([vocab/tp, hidden]) while
542
- # probs spans the full vocab, so each rank multiplies its local vocab slice
543
- # [sc_vocab_start, sc_vocab_end) and the partials are summed across ranks.
544
- local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype)
545
- soft_embeds = torch.matmul(
546
- local_probs, embed_weight[: sc_vocab_end - sc_vocab_start]
547
- )
548
- if tp_size > 1:
549
- soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name)
550
- soft_embeds = soft_embeds * normalizer
551
- sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype)
571
+ if compute_sc:
572
+ # Self-conditioning soft embed = probs @ embed_tokens.weight. Under
573
+ # tensor parallelism the embedding is vocab-sharded ([vocab/tp,
574
+ # hidden]) while probs spans the full vocab, so each rank multiplies
575
+ # its local vocab slice [sc_vocab_start, sc_vocab_end) and the
576
+ # partials are summed across ranks.
577
+ local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype)
578
+ soft_embeds = torch.matmul(
579
+ local_probs, embed_weight[: sc_vocab_end - sc_vocab_start]
580
+ )
581
+ if tp_size > 1:
582
+ soft_embeds = torch.ops.vllm.all_reduce(
583
+ soft_embeds, group_name=tp_group_name
584
+ )
585
+ soft_embeds = soft_embeds * normalizer
586
+ # A pinned position holds its seed token as input. Zero its soft embed
587
+ # too, or the model's own prediction there reaches the next step
588
+ # through self-conditioning. The buffer is fp32 while the embedding is
589
+ # in the model dtype: compiled code casts on the store but eager does
590
+ # not, and the step runs eager once torch.compile hits its recompile
591
+ # limit.
592
+ sc_pin = (~pin_mask[decode_slots]).unsqueeze(-1)
593
+ sc_embeds[decode_slots] = (soft_embeds * sc_keep * sc_pin).to(sc_embeds.dtype)
594
+ else:
595
+ # Every slot in this tile ends after this step, so the soft embed
596
+ # would never be read. The matmul is a full pass over the vocabulary
597
+ # matrix, so skip it.
598
+ sc_embeds[decode_slots] = 0
599
600
# Overwrite canvas with argmax for newly converged denoise requests
601
newly_converged = (converged & is_denoise).unsqueeze(1)
602
canvas[decode_slots] = torch.where(
603
newly_converged, argmax_canvas[decode_slots], canvas[decode_slots]
604
)
605
+ is_encoder_phase[decode_slots] &= ~read_only[decode_slots]
606
607
# ---- Phase 7: Copy canvas → draft_tokens for all slots ----
608
draft_tokens[all_slots, :CL] = canvas[all_slots]
@@ -628,6 +676,32 @@ class DiffusionGemmaRequestStates:
676
# Per-slot confidence flag, set by the sampler each step.
677
self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
678
679
+ # Per-slot step cap, lowered per request from extra_args.
680
+ self.max_steps = torch.full(
681
+ (max_num_reqs,), max_denoising_steps, dtype=torch.int32, device=device
682
+ )
683
+ # Seed canvases replace the random canvas after prefill. The host-side
684
+ # sets gate the seed and read-only work so plain generation skips it.
685
+ self.seed_canvas = torch.zeros(
686
+ max_num_reqs, canvas_length, dtype=torch.int64, device=device
687
+ )
688
+ self.has_seed = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
689
+ self.seeded_slots: set[int] = set()
690
+ # Pinned positions keep their seed value on every denoise step, so a
691
+ # multi-step read denoises the canvas the request seeded.
692
+ self.pin_mask = torch.zeros(
693
+ max_num_reqs, canvas_length, dtype=torch.bool, device=device
694
+ )
695
+ # Read-only slots emit on their converging step and skip the commit
696
+ # forward.
697
+ self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
698
+ self.read_only_slots: set[int] = set()
699
+ # Slots capped at one denoise step never consume a soft embed.
700
+ self.single_step_slots: set[int] = set()
701
+ # Per-slot canvas width, at most canvas_length. The scheduler schedules
702
+ # this many draft tokens for the slot and the sampler pads the rest.
703
+ self.canvas_width_np = np.full(max_num_reqs, canvas_length, dtype=np.int32)
704
+
705
# Per-slot self-conditioning soft embedding (probs @ embed_weight) from
706
# the previous denoise step. Storing the [.., hidden] soft embed instead
707
# of the full [.., vocab] distribution shrinks this buffer by
@@ -657,11 +731,57 @@ class DiffusionGemmaRequestStates:
731
self.step[slot_idx].fill_(0)
732
self.accepted_canvas_history_len[slot_idx].fill_(0)
733
self.self_conditioning_embeds[slot_idx] = 0
734
+ self.max_steps[slot_idx].fill_(self.max_denoising_steps)
735
+ self.has_seed[slot_idx].fill_(False)
736
+ self.seeded_slots.discard(slot_idx)
737
+ self.pin_mask[slot_idx].fill_(False)
738
+ self.read_only[slot_idx].fill_(False)
739
+ self.read_only_slots.discard(slot_idx)
740
+ self.single_step_slots.discard(slot_idx)
741
+ self.canvas_width_np[slot_idx] = self.canvas_length
742
743
def remove_request(self, slot_idx: int) -> None:
744
+ # add_request resets the GPU flags before a slot is reused. The host
745
+ # sets gate work on every step, so they clear here.
746
self.is_encoder_phase[slot_idx].fill_(False)
747
self.accepted_canvas_history_len[slot_idx].fill_(0)
748
self.self_conditioning_embeds[slot_idx] = 0
749
+ self.seeded_slots.discard(slot_idx)
750
+ self.read_only_slots.discard(slot_idx)
751
+ self.single_step_slots.discard(slot_idx)
752
+
753
+ def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None:
754
+ """``ids`` covers the slot's canvas width; positions past it are never
755
+ scheduled."""
756
+ self.seed_canvas[slot_idx, : len(ids)] = async_tensor_h2d(
757
+ ids, dtype=torch.int64, device=self.device
758
+ )
759
+ self.has_seed[slot_idx].fill_(True)
760
+ self.seeded_slots.add(slot_idx)
761
+
762
+ def set_pins(self, slot_idx: int, positions: list[int]) -> None:
763
+ """Hold ``positions`` of the slot's seed canvas through every denoise
764
+ step. Validated against the request's canvas width upstream."""
765
+ self.pin_mask[slot_idx].fill_(False)
766
+ self.pin_mask[
767
+ slot_idx, async_tensor_h2d(positions, dtype=torch.int64, device=self.device)
768
+ ] = True
769
+
770
+ def set_read_only(self, slot_idx: int) -> None:
771
+ self.read_only[slot_idx].fill_(True)
772
+ self.read_only_slots.add(slot_idx)
773
+
774
+ def apply_seed_canvases(
775
+ self, slots_np: np.ndarray, slots_gpu: torch.Tensor
776
+ ) -> None:
777
+ """Replace the canvas of every seeded slot among ``slots_gpu``."""
778
+ if self.seeded_slots.isdisjoint(slots_np.tolist()):
779
+ return
780
+ self.canvas[slots_gpu] = torch.where(
781
+ self.has_seed[slots_gpu, None],
782
+ self.seed_canvas[slots_gpu],
783
+ self.canvas[slots_gpu],
784
+ )
785
786
787
class DiffusionGemmaModelState(ModelState):
@@ -1030,6 +1150,33 @@ class DiffusionSampler:
1150
self._pending_logprobs.pop(req_idx, None)
1151
self.sampling_states.add_request(req_idx, sampling_params)
1152
self.logprob_token_ids_state.add_request(req_idx, sampling_params)
1153
+ extra = getattr(sampling_params, "extra_args", None) or {}
1154
+ states = self.diffusion_states
1155
+ cap = extra.get("diffusion_max_steps")
1156
+ if cap is not None:
1157
+ cap = max(1, min(int(cap), states.max_denoising_steps))
1158
+ states.max_steps[req_idx].fill_(cap)
1159
+ if cap == 1:
1160
+ states.single_step_slots.add(req_idx)
1161
+ width = extra.get("diffusion_canvas_length")
1162
+ if width:
1163
+ states.canvas_width_np[req_idx] = max(
1164
+ 1, min(int(width), self.canvas_length)
1165
+ )
1166
+ width = int(states.canvas_width_np[req_idx])
1167
+ seed = extra.get("diffusion_seed_canvas")
1168
+ if seed is not None:
1169
+ if len(seed) != width:
1170
+ raise ValueError(
1171
+ f"diffusion_seed_canvas must hold exactly {width} ids, "
1172
+ f"got {len(seed)}"
1173
+ )
1174
+ states.set_seed_canvas(req_idx, seed)
1175
+ pins = extra.get("diffusion_pinned")
1176
+ if pins and seed is not None:
1177
+ states.set_pins(req_idx, [int(p) for p in pins])
1178
+ if extra.get("diffusion_read_only"):
1179
+ states.set_read_only(req_idx)
1180
1181
def apply_staged_writes(self) -> None:
1182
self.sampling_states.apply_staged_writes()
@@ -1071,6 +1218,7 @@ class DiffusionSampler:
1218
ps.astype(np.int64), device=states.is_encoder_phase.device
1219
)
1220
states.init_canvas(ps_gpu)
1221
+ states.apply_seed_canvases(ps, ps_gpu)
1222
self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[
1223
ps_gpu
1224
]
@@ -1151,7 +1299,6 @@ class DiffusionSampler:
1299
# --- CPU/NumPy setup (outside compile): split decode vs prefill, init
1300
# canvas for any new prefills, and stage decode slot indices to GPU. ---
1301
states = self.diffusion_states
1154
- CL = self.canvas_length
1302
slots_np = input_batch.idx_mapping_np[:num_reqs]
1303
per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1])
1304
@@ -1194,18 +1341,12 @@ class DiffusionSampler:
1341
if top_k is not None or top_p is not None:
1342
logits = apply_top_k_top_p(logits.float(), top_k, top_p)
1343
1197
- # Pad any truncated canvas back to CL so the uniform-CL sampler math
1198
- # holds. Phantom (padded) positions are zeroed → uniform logits → high
1199
- # entropy (no premature convergence) and argmax 0 (stable); they are
1200
- # never committed (num_sampled == real length). masked_fill (not
1201
- # multiply) so -inf entries from top_k/top_p filtering above don't
1202
- # turn phantom rows into NaN.
1203
- if num_decode > 0 and valid_canvas_len_np.min() < CL:
1204
- ar = torch.arange(CL, device=device)
1205
- starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req
1206
- valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL]
1207
- src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1)
1208
- logits = logits[src.reshape(-1)].masked_fill_(~valid.reshape(-1, 1), 0)
1344
+ # Where each decode request's rows start in the flat logits. Tiles
1345
+ # below gather and pad a request to the tile's width.
1346
+ row_starts_np = np.concatenate(
1347
+ ([0], np.cumsum(valid_canvas_len_np)[:-1])
1348
+ ).astype(np.int64)
1349
+ row_starts = async_tensor_h2d(row_starts_np, device=device)
1350
1351
# Clear once: the tiled loop below only scatters its own decode slots,
1352
# so it must not re-clear earlier tiles' writes.
@@ -1228,125 +1369,159 @@ class DiffusionSampler:
1369
want_logprobs = max_num_logprobs >= 0 or max_token_ids > 0
1370
num_logprobs = max(max_num_logprobs, 0)
1371
1231
- # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in
1232
- # _compiled_sample_step keeps several live [group * CL, vocab] copies, so
1233
- # size each tile to a fraction of free memory to bound the transient at
1234
- # high concurrency. Tiling is bit-identical to a single pass.
1235
- group = max(num_decode, 1)
1236
- if num_decode > 0:
1237
- free, _ = current_platform.mem_get_info()
1238
- # ~10 transient fp32 copies of [group * CL, vocab] inside the step
1372
+ # Sample per tile. Decode requests are grouped by canvas width and each
1373
+ # tile runs the compiled step at that width over [:, :W] views of the
1374
+ # state, so a narrow read pays for its own rows rather than the served
1375
+ # canvas. Widths ascend, so the last tile's canvas-to-draft copy (over
1376
+ # all slots) is the widest. The fp32 pipeline keeps several live
1377
+ # [tile * W, vocab] copies, so a tile is also bounded by free memory.
1378
+ widths_np = states.canvas_width_np[decode_slots_np]
1379
+ order = np.argsort(widths_np, kind="stable")
1380
+ free = current_platform.mem_get_info()[0] if num_decode > 0 else 0
1381
+ run_start = 0
1382
+ while run_start < num_decode:
1383
+ W = int(widths_np[order[run_start]])
1384
+ run_end = run_start
1385
+ while run_end < num_decode and widths_np[order[run_end]] == W:
1386
+ run_end += 1
1387
+ # ~10 transient fp32 copies of [tile * W, vocab] inside the step
1388
# (eager peaks at ~8; pad for allocator overhead and small tensors).
1240
- bytes_per_req = CL * self.vocab_size * 4 * 10
1241
- budget = int(free * 0.5) // max(bytes_per_req, 1)
1242
- group = max(1, min(num_decode, budget))
1243
-
1244
- for start_req in range(0, num_decode, group):
1245
- end_req = min(start_req + group, num_decode)
1246
- tile = slice(start_req, end_req)
1247
- tile_slots = decode_slots[tile]
1248
-
1249
- scaled = _compiled_sample_step(
1250
- logits[start_req * CL : end_req * CL],
1251
- tile_slots,
1252
- decode_idx[tile],
1253
- all_slots,
1254
- valid_canvas_len[tile],
1255
- # State
1256
- states.canvas,
1257
- states.argmax_canvas,
1258
- states.step,
1259
- states.is_encoder_phase,
1260
- states.confident,
1261
- states.self_conditioning_embeds,
1262
- self.embed_weight,
1263
- self.normalizer,
1264
- states.accepted_canvas_history,
1265
- states.accepted_canvas_history_len,
1266
- # Output
1267
- sampled,
1268
- num_sampled,
1269
- self.req_states.draft_tokens,
1270
- # Config
1271
- max_denoising_steps=float(states.max_denoising_steps),
1272
- t_min=self.t_min,
1273
- t_max=self.t_max,
1274
- confidence_threshold=self.confidence_threshold,
1275
- vocab_size=self.vocab_size,
1276
- CL=CL,
1277
- ST=states.stability_threshold,
1278
- entropy_bound=self.entropy_bound,
1279
- sc_vocab_start=self.sc_vocab_start,
1280
- sc_vocab_end=self.sc_vocab_end,
1281
- tp_size=self.tp_size,
1282
- tp_group_name=self.tp_group_name,
1283
- )
1389
+ budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 10, 1))
1390
+ for t0 in range(run_start, run_end, budget):
1391
+ sel_np = order[t0 : min(t0 + budget, run_end)]
1392
+ n = len(sel_np)
1393
+ sel = async_tensor_h2d(sel_np.astype(np.int64), device=device)
1394
+ tile_slots = decode_slots[sel]
1395
+ tile_valid = valid_canvas_len[sel]
1396
+ tile_valid_np = valid_canvas_len_np[sel_np]
1397
+ contiguous = bool((sel_np == np.arange(sel_np[0], sel_np[0] + n)).all())
1398
+ if contiguous and tile_valid_np.min() == W:
1399
+ r0 = int(row_starts_np[sel_np[0]])
1400
+ tile_logits = logits[r0 : r0 + n * W]
1401
+ else:
1402
+ # Pad each request to W. Phantom positions are zeroed:
1403
+ # uniform logits, high entropy, argmax 0, never committed.
1404
+ # masked_fill, not multiply, so -inf from top_k/top_p
1405
+ # filtering above cannot turn a phantom row into NaN.
1406
+ ar = torch.arange(W, device=device)
1407
+ src = (row_starts[sel].unsqueeze(1) + ar.unsqueeze(0)).clamp_max(
1408
+ logits.shape[0] - 1
1409
+ )
1410
+ valid = ar.unsqueeze(0) < tile_valid.unsqueeze(1)
1411
+ tile_logits = logits[src.reshape(-1)].masked_fill_(
1412
+ ~valid.reshape(-1, 1), 0
1413
+ )
1414
+ compute_sc = (
1415
+ not states.single_step_slots
1416
+ or not states.single_step_slots.issuperset(
1417
+ decode_slots_np[sel_np].tolist()
1418
+ )
1419
+ )
1420
1285
- # Logprobs for denoise steps that just converged (is_encoder_phase
1286
- # flipped False→True), stashed per tile so `scaled` is freed each tile.
1287
- if want_logprobs:
1288
- converged_mask = states.is_encoder_phase[tile_slots]
1289
- just_converged = converged_mask & ~is_committing[tile]
1290
- if just_converged.any():
1291
- flat_logits = scaled.reshape(-1, scaled.shape[-1])
1292
- argmax_tokens = scaled.argmax(dim=-1)
1293
- for local_idx in just_converged.nonzero(as_tuple=True)[0]:
1294
- li = local_idx.item()
1295
- slot = tile_slots[local_idx]
1296
- # Stash only the real canvas positions (== CL unless this
1297
- # canvas was truncated near max_model_len); padded tail
1298
- # positions are never emitted.
1299
- k_i = int(valid_canvas_len_np[start_req + li])
1300
- pos = li * CL
1301
- per_req_ids = max_token_ids > 0
1302
- self._pending_logprobs[slot.item()] = compute_topk_scores(
1303
- flat_logits[pos : pos + k_i],
1304
- num_logprobs,
1305
- argmax_tokens[local_idx][:k_i],
1306
- logprob_token_ids_state=(
1307
- self.logprob_token_ids_state if per_req_ids else None
1308
- ),
1309
- # every row of this stash belongs to one slot
1310
- expanded_idx_mapping=(
1311
- torch.full(
1312
- (k_i,),
1313
- slot.item(),
1314
- dtype=torch.int32,
1315
- device=flat_logits.device,
1316
- )
1317
- if per_req_ids
1318
- else None
1319
- ),
1320
- max_per_req_token_ids=max_token_ids,
1321
- logits_mode=self.logits_mode,
1322
- )
1323
-
1324
- # Commit steps: is_committing was True at entry. Reassemble previously
1325
- # stashed logprobs and attach to SamplerOutput.
1421
+ scaled = _compiled_sample_step(
1422
+ tile_logits,
1423
+ tile_slots,
1424
+ decode_idx[sel],
1425
+ all_slots,
1426
+ tile_valid,
1427
+ # State, viewed at this tile's width
1428
+ states.canvas[:, :W],
1429
+ states.argmax_canvas[:, :W],
1430
+ states.step,
1431
+ states.is_encoder_phase,
1432
+ states.confident,
1433
+ states.self_conditioning_embeds[:, :W],
1434
+ self.embed_weight,
1435
+ self.normalizer,
1436
+ states.accepted_canvas_history[:, :, :W],
1437
+ states.accepted_canvas_history_len,
1438
+ states.max_steps,
1439
+ states.pin_mask[:, :W],
1440
+ states.seed_canvas[:, :W],
1441
+ states.read_only,
1442
+ # Output
1443
+ sampled[:, :W],
1444
+ num_sampled,
1445
+ self.req_states.draft_tokens,
1446
+ # Config
1447
+ max_denoising_steps=float(states.max_denoising_steps),
1448
+ t_min=self.t_min,
1449
+ t_max=self.t_max,
1450
+ confidence_threshold=self.confidence_threshold,
1451
+ vocab_size=self.vocab_size,
1452
+ CL=W,
1453
+ ST=states.stability_threshold,
1454
+ entropy_bound=self.entropy_bound,
1455
+ sc_vocab_start=self.sc_vocab_start,
1456
+ sc_vocab_end=self.sc_vocab_end,
1457
+ tp_size=self.tp_size,
1458
+ tp_group_name=self.tp_group_name,
1459
+ compute_sc=compute_sc,
1460
+ )
1461
+
1462
+ # Stash newly converged logprobs, including reads that emit now.
1463
+ if want_logprobs:
1464
+ converged_mask = states.is_encoder_phase[tile_slots] | (
1465
+ num_sampled[decode_idx[sel]] > 0
1466
+ )
1467
+ just_converged = converged_mask & ~is_committing[sel]
1468
+ if just_converged.any():
1469
+ flat_logits = scaled.reshape(-1, scaled.shape[-1])
1470
+ argmax_tokens = scaled.argmax(dim=-1)
1471
+ raw_flat: torch.Tensor | None = None
1472
+ for local_idx in just_converged.nonzero(as_tuple=True)[0]:
1473
+ li = local_idx.item()
1474
+ slot = tile_slots[local_idx].item()
1475
+ # Stash only the real canvas positions; padded tail
1476
+ # positions are never emitted.
1477
+ k_i = int(tile_valid_np[li])
1478
+ pos = li * W
1479
+ src = flat_logits
1480
+ if slot in states.read_only_slots:
1481
+ # Read-only slots report logprobs at temperature 1.
1482
+ # The schedule-tempered logits share the argmax.
1483
+ if raw_flat is None:
1484
+ raw_flat = tile_logits.float()
1485
+ src = raw_flat
1486
+ per_req_ids = max_token_ids > 0
1487
+ self._pending_logprobs[slot] = compute_topk_scores(
1488
+ src[pos : pos + k_i],
1489
+ num_logprobs,
1490
+ argmax_tokens[local_idx][:k_i],
1491
+ logprob_token_ids_state=(
1492
+ self.logprob_token_ids_state
1493
+ if per_req_ids
1494
+ else None
1495
+ ),
1496
+ # every row of this stash belongs to one slot
1497
+ expanded_idx_mapping=(
1498
+ torch.full(
1499
+ (k_i,), slot, dtype=torch.int32, device=device
1500
+ )
1501
+ if per_req_ids
1502
+ else None
1503
+ ),
1504
+ max_per_req_token_ids=max_token_ids,
1505
+ logits_mode=self.logits_mode,
1506
+ )
1507
+ run_start = run_end
1508
+
1509
+ # Only emitting requests consume their stashed logprobs.
1510
logprobs_tensors = None
1327
- if want_logprobs and is_committing.any() and self._pending_logprobs:
1328
- committing_slots = set(
1329
- decode_slots_np[is_committing.cpu().numpy()].tolist()
1330
- )
1331
- parts_ids, parts_lp, parts_ranks = [], [], []
1511
+ if want_logprobs and self._pending_logprobs:
1512
+ emitting_slots = set(slots_np[num_sampled.cpu().numpy() > 0].tolist())
1513
+ parts: list[LogprobsTensors] = []
1514
cu_gen: list[int] = []
1515
flat_offset = 0
1516
for i in range(num_reqs):
1517
cu_gen.append(flat_offset)
1518
slot = int(slots_np[i])
1337
- if slot in committing_slots and slot in self._pending_logprobs:
1519
+ if slot in emitting_slots and slot in self._pending_logprobs:
1520
lp = self._pending_logprobs.pop(slot)
1339
- parts_ids.append(lp.logprob_token_ids)
1340
- parts_lp.append(lp.logprobs)
1341
- parts_ranks.append(lp.selected_token_ranks)
1521
+ parts.append(lp)
1522
flat_offset += lp.logprobs.shape[0]
1343
- if parts_ids:
1344
- logprobs_tensors = LogprobsTensors(
1345
- logprob_token_ids=torch.cat(parts_ids),
1346
- logprobs=torch.cat(parts_lp),
1347
- selected_token_ranks=torch.cat(parts_ranks),
1348
- cu_num_generated_tokens=cu_gen,
1349
- )
1523
+ if parts:
1524
+ logprobs_tensors = _concat_logprob_stashes(parts, cu_gen)
1525
1526
return self._build_output(
1527
input_batch,
vllm/transformers_utils/model_arch_config_convertor.py
+1
@@ -813,6 +813,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
813
"dbrx": DbrxModelArchConfigConvertor,
814
"deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,
815
"deepseek_v4": DeepseekV4ModelArchConfigConvertor,
816
+ "diffusion_gemma": Gemma4ModelArchConfigConvertor,
817
"diffusion_gemma_text": Gemma4ModelArchConfigConvertor,
818
"ernie_mtp": ErnieMTPModelArchConfigConvertor,
819
"falcon": FalconModelArchConfigConvertor,
vllm/utils/diffusion.py
new
+104
@@ -0,0 +1,104 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Request validation for discrete diffusion models."""
4
+
5
+from typing import TYPE_CHECKING
6
+
7
+from vllm.exceptions import VLLMValidationError
8
+
9
+if TYPE_CHECKING:
10
+ from vllm.sampling_params import SamplingParams
11
+
12
+
13
+def validate_diffusion_sampling_params(
14
+ params: "SamplingParams",
15
+ *,
16
+ canvas_length: int | None,
17
+ vocab_size: int,
18
+ async_scheduling: bool | None,
19
+) -> None:
20
+ """Validate request options and normalize one-canvas reads at admission."""
21
+ extra = params.extra_args
22
+ if not extra:
23
+ return
24
+
25
+ def check(valid: bool, message: str) -> None:
26
+ if not valid:
27
+ raise VLLMValidationError(message, parameter="extra_args")
28
+
29
+ width = extra.get("diffusion_canvas_length")
30
+ if width is not None:
31
+ check(
32
+ isinstance(width, int)
33
+ and not isinstance(width, bool)
34
+ and width > 0
35
+ and (canvas_length is None or width <= canvas_length),
36
+ "diffusion_canvas_length must be a positive integer no larger "
37
+ "than the served canvas.",
38
+ )
39
+ # The request's canvas: its own width, else the served one.
40
+ expected_len = width or canvas_length
41
+
42
+ seed = extra.get("diffusion_seed_canvas")
43
+ if seed is not None:
44
+ check(
45
+ isinstance(seed, (list, tuple))
46
+ and all(isinstance(t, int) and not isinstance(t, bool) for t in seed),
47
+ "diffusion_seed_canvas must be a list of token ids.",
48
+ )
49
+ check(
50
+ all(0 <= t < vocab_size for t in seed),
51
+ f"diffusion_seed_canvas ids must be in [0, {vocab_size}).",
52
+ )
53
+ check(
54
+ expected_len is None or len(seed) == expected_len,
55
+ f"diffusion_seed_canvas must hold exactly {expected_len} ids, "
56
+ f"got {len(seed)}.",
57
+ )
58
+
59
+ pins = extra.get("diffusion_pinned")
60
+ if pins is not None:
61
+ check(
62
+ isinstance(pins, (list, tuple))
63
+ and all(isinstance(p, int) and not isinstance(p, bool) for p in pins),
64
+ "diffusion_pinned must be a list of canvas positions.",
65
+ )
66
+ check(
67
+ seed is not None,
68
+ "diffusion_pinned needs a diffusion_seed_canvas to hold.",
69
+ )
70
+ check(
71
+ all(p >= 0 and (expected_len is None or p < expected_len) for p in pins),
72
+ "diffusion_pinned positions must lie inside the canvas.",
73
+ )
74
+
75
+ cap = extra.get("diffusion_max_steps")
76
+ if cap is not None:
77
+ check(
78
+ isinstance(cap, int) and not isinstance(cap, bool) and cap > 0,
79
+ "diffusion_max_steps must be a positive integer.",
80
+ )
81
+
82
+ # The OpenAI server's vllm_xargs narrows JSON booleans to 0/1.
83
+ read_only = extra.get("diffusion_read_only")
84
+ if read_only is not None:
85
+ check(
86
+ isinstance(read_only, (bool, int)) and read_only in (0, 1),
87
+ "diffusion_read_only must be a boolean (or 0/1).",
88
+ )
89
+ if read_only:
90
+ # One canvas is the whole output. Cap max_tokens so the scheduler
91
+ # ends the request there, and ignore EOS so an end-of-turn token
92
+ # drawn into a noise slot does not end it early.
93
+ if expected_len is not None:
94
+ params.max_tokens = min(params.max_tokens or expected_len, expected_len)
95
+ params.ignore_eos = True
96
+
97
+ check(
98
+ bool(async_scheduling)
99
+ or width is None
100
+ or canvas_length is None
101
+ or width == canvas_length,
102
+ "A diffusion_canvas_length smaller than the served canvas "
103
+ "requires --async-scheduling.",
104
+ )
vllm/v1/core/sched/diffusion_scheduler.py
new
+56
@@ -0,0 +1,56 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Async scheduling for diffusion requests.
4
+
5
+These rules live here so Scheduler.schedule() and AsyncScheduler stay
6
+unchanged. VllmConfig selects this class for a diffusion model under async
7
+scheduling. A sync scheduler creates no output placeholders, which both rules
8
+read.
9
+"""
10
+
11
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
12
+from vllm.v1.core.sched.output import SchedulerOutput
13
+from vllm.v1.request import Request
14
+
15
+
16
+def diffusion_canvas_width(request: Request, canvas_length: int) -> int:
17
+ """The canvas width a diffusion request asked for, else the served one."""
18
+ params = request.sampling_params
19
+ extra = params.extra_args if params is not None else None
20
+ width = extra.get("diffusion_canvas_length") if extra else None
21
+ return int(width) if width else canvas_length
22
+
23
+
24
+def _read_in_flight(request: Request, width: int) -> bool:
25
+ """True when a read-only request has all its denoise steps in flight.
26
+
27
+ The request emits its canvas on the last step and ends, so a further
28
+ step is discarded.
29
+ """
30
+ params = request.sampling_params
31
+ extra = params.extra_args if params is not None else None
32
+ if not extra or not extra.get("diffusion_read_only"):
33
+ return False
34
+ steps = extra.get("diffusion_max_steps")
35
+ if not steps:
36
+ return False
37
+ # Each in-flight denoise step holds one canvas of placeholders.
38
+ return request.num_output_placeholders >= steps * width
39
+
40
+
41
+class DiffusionAsyncScheduler(AsyncScheduler):
42
+ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
43
+ for request in self.running:
44
+ width = diffusion_canvas_width(request, self.num_spec_tokens)
45
+ # The placeholders are as wide as the served canvas.
46
+ if len(request.spec_token_ids) > width:
47
+ request.spec_token_ids = request.spec_token_ids[:width]
48
+ if _read_in_flight(request, width):
49
+ # Scheduler.schedule()'s max_tokens guard cannot be reached
50
+ # from a subclass. schedule() advances current_step before it
51
+ # checks decode eligibility, so current_step + 2 skips this
52
+ # call alone. max keeps a longer pipeline-parallel wait.
53
+ request.next_decode_eligible_step = max(
54
+ request.next_decode_eligible_step, self.current_step + 2
55
+ )
56
+ return super().schedule(throttle_prefills)
vllm/v1/engine/input_processor.py
+16
@@ -30,6 +30,7 @@ from vllm.tasks import GENERATION_TASKS, POOLING_TASKS, SupportedTask
30
from vllm.tokenizers import TokenizerLike
31
from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
32
from vllm.utils.async_utils import make_async
33
+from vllm.utils.diffusion import validate_diffusion_sampling_params
34
from vllm.utils.jsontree import json_iter_leaves
35
from vllm.v1.engine import EngineCoreRequest
36
from vllm.v1.kv_hints import KvHintsEnvelope
@@ -53,6 +54,7 @@ class InputProcessor:
54
self.speculative_config = vllm_config.speculative_config
55
self.structured_outputs_config = vllm_config.structured_outputs_config
56
self.observability_config = vllm_config.observability_config
57
+ self.diffusion_config = vllm_config.diffusion_config
58
# Load the custom logits processor classes once; the returned callable
59
# runs their validate_params hooks per request at admission.
60
self.validate_logits_processors_params = (
@@ -129,6 +131,20 @@ class InputProcessor:
131
132
self.validate_logits_processors_params(params)
133
134
+ if self.model_config.is_diffusion:
135
+ # Without --diffusion-config the served canvas is unknown here;
136
+ # the ids and the read-only normalisation are still checked.
137
+ validate_diffusion_sampling_params(
138
+ params,
139
+ canvas_length=(
140
+ self.diffusion_config.canvas_length
141
+ if self.diffusion_config is not None
142
+ else None
143
+ ),
144
+ vocab_size=self.model_config.get_vocab_size(),
145
+ async_scheduling=self.vllm_config.scheduler_config.async_scheduling,
146
+ )
147
+
148
if self.model_config.return_sampling_mask:
149
if params.temperature <= 0:
150
raise ValueError(