diff --git a/examples/features/structured_diffusion/README.md b/examples/features/structured_diffusion/README.md new file mode 100644 index 000000000000..9840f9ea3f01 --- /dev/null +++ b/examples/features/structured_diffusion/README.md @@ -0,0 +1,74 @@ +# Structured reads on DiffusionGemma + +A discrete diffusion model denoises a whole canvas per forward pass. If the +canvas is seeded with the answer's fixed text and only the answer slots are +left as noise, one denoise step gives a distribution over each slot. These +`extra_args` fields (`vllm_xargs` on the OpenAI server) expose that: + +| field | type | meaning | +| --- | --- | --- | +| `diffusion_seed_canvas` | `list[int]`, exactly `canvas_length` ids | replaces the random initial canvas after prefill | +| `diffusion_pinned` | `list[int]` of canvas positions | held at their seed value on every denoise step, so a read past one step keeps its template | +| `diffusion_max_steps` | `int` | denoise steps before the canvas is emitted | +| `diffusion_read_only` | `bool` | emit the argmax canvas as soon as the cap is reached, end the request there, and return temperature-1 logprobs at every position | + +`structured_server.py` turns a question schema into those fields. It serves +`/v1/chat/completions`: the system message is the schema, the user message +is the state JSON, and the reply content is one distribution per question +with a standard error over a few noise draws. + +```bash +vllm serve google/diffusiongemma-26B-A4B-it \ + --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 --enable-prefix-caching +python examples/features/structured_diffusion/structured_server.py \ + --upstream http://127.0.0.1:8000 --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64 +curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -d '{ + "messages": [ + {"role": "system", "content": "{\"questions\": [{\"id\": \"urgent\", \"type\": \"noul\", \"instructions\": \"Does the customer need a reply within the hour?\"}]}"}, + {"role": "user", "content": "{\"ticket\": \"Everything is down and we have a demo at noon.\"}"} + ]}' +``` + +The attention backend is picked as for Gemma 4: FlashAttention 4 on every +layer when available, otherwise Triton. FlashInfer cannot serve this model (a +batch mixes causal prefill with bidirectional denoising), and +`--attention-backend FLASHINFER` is rejected. + +Per-request canvas widths smaller than the served canvas require async +scheduling. The diffusion async scheduler is selected automatically; no +`--scheduler-cls` argument is needed. Synchronous execution supports full-width +canvases only. + +Question types: `noul` (yes/no), `choice` with `options`, `score` with +ordered `levels`. Each label must be a single token in the answer template, +which the server checks with the tokenizer when a request uses the schema. + +`POST /v1/systemone` implements the Jev decision API. The body holds +`state`, `questions` (a map of id to `type`, `instructions` and `criteria`) +and `model`. Answers come back in that API's shapes: a `noul` probability, +a `choice` with `probabilities` and `confidence`, or a `score` with a +0-indexed `legend`. The schema options above go in the same body as +extensions. + +A question may declare `depends_on` (read in a later stage with those +answers in its prompt), `ask_if` (asked only when a named question's answer +is among the listed ones, otherwise null) and `alone` (a read of its own). + +Images attach as `multipart/form-data`, with the JSON in a part named +`request` and each image as a file part, or as an `images` array of data +URLs. + +```bash +curl -s localhost:8011/v1/systemone -H 'content-type: application/json' -d '{ + "model": "jev-latest", + "state": {"ticket": "Everything is down and we have a demo at noon."}, + "questions": {"urgent": {"type": "noul", "instructions": "Does the customer need a reply within the hour?"}}}' +``` + +`"think": N` in the schema lets the model write up to N tokens in its +thought channel before the read. The thought is an ordinary generation with +the chat template's thinking marker on, and the read then runs with the +thought in its prompt, so the answer slots condition on it. The noise draws +of a decision share one thought. `diagnostics.thought` returns the text, its +length in tokens, whether the model closed the channel itself and the +generation time. diff --git a/examples/features/structured_diffusion/structured_server.py b/examples/features/structured_diffusion/structured_server.py new file mode 100644 index 000000000000..a8f61da94e41 --- /dev/null +++ b/examples/features/structured_diffusion/structured_server.py @@ -0,0 +1,1528 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Structured decisions in front of a vLLM DiffusionGemma server. + +POST /v1/systemone takes Jev's request body: {"model", "state", "questions"}. +"questions" maps an id to {"type", "instructions", "criteria"}, where "type" +is "noul", "choice" or "score" and the criteria shape follows the type: + noul: optional {"true": ..., "false": ...} descriptions + choice: option name -> description or null + score: ordered list of levels +Answers take Jev's shapes, with this server's diagnostics alongside: + noul: {"noul": p} + choice: {"choice", "probabilities", "confidence"} + score: {"score", "legend", "probabilities", "confidence"} +The request body may also carry the schema keys "instructions", "samples", +"auto_max", "auto_threshold", "steps", "think", "ask", "chunk_rows", +"chunk_prompt" and "sequential" as extensions. Images go ahead of the +state, either as multipart/form-data with the JSON body in a part named +"request" and each image as a file part, or as an "images" array of data +URLs in the JSON body. + +POST /v1/chat/completions makes the same decision from an OpenAI-shaped +call. The system message is the schema JSON below and the user message is +the state. The reply's `content` is the JSON answer set. + +POST /v1/raw/chat/completions passes the body to vLLM's chat completions +unchanged, for plain generation through this port. + +With API_KEY set in the environment, every POST needs "Authorization: +Bearer ". --tls-port adds an HTTPS listener with a self-signed +certificate kept in --cert-dir, for clients that need a secure origin. + +Each answer is one distribution per question, from one denoise step over a +seeded canvas, averaged over a few noise draws. This server handles the +canvas, tokenizer, slot resolution, noise draws and averaging. + +Schema (system message): + {"questions": [ + {"id": "urgent", "type": "noul", "instructions": "..."}, + {"id": "bucket", "type": "choice", "instructions": "...", + "options": [{"name": "billing", "description": "..."}, ...]}, + {"id": "tone", "type": "score", "instructions": "...", + "levels": ["calm", "annoyed", "furious"]}], + "instructions": "optional context", + "samples": "auto" | N, "auto_threshold": 0.1, "auto_max": 4, + "steps": 1, "think": 0} + +A question may also declare: + "depends_on": [ids] answered after those, with their answers in + its prompt + "ask_if": {id: [answers]} asked only when that question's answer is + among them (a skipped answer is null) + "alone": true a read of its own +Questions run in stages by these dependencies. Each stage is one joint +read. Later stages continue the earlier answers, prefilled for a text +state and restated for an image. + +Up to ten questions answer as "id: label" lines. Past that the id runs +straight into the label, space separated, one row fewer per question. The +server splits a schema whose answer template does not fit the canvas into +chunks that run together, each with its own question list. "chunk_rows" +sets the rows per chunk, "ask" picks a subset of question ids for one +read, "sequential": true runs the chunks in order with the earlier answers +prefilled, and "chunk_prompt": "shared" lists every question in each +chunk's prompt. "think": N lets the model write up to N tokens in its +thought channel, as an ordinary generation, and the read then runs with +that thought in its prompt. With images the model writes the thought with +the image in view and the server seeds it into the canvas ahead of the +answer, so the canvas bounds it. The noise draws of a decision share one +thought. + +Serve the model with a canvas that holds the answer template, for example: + vllm serve google/diffusiongemma-26B-A4B-it \ + --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 \ + --enable-prefix-caching +then run this in front of it: + python structured_server.py --upstream http://127.0.0.1:8000 \ + --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64 --port 8011 +""" + +import argparse +import json +import math +import os +import random +import ssl +import subprocess +import threading +import time +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pybase64 as base64 +from transformers import AutoTokenizer + +ARGS = None +TOK = None +API_KEY = os.environ.get( + "API_KEY", "" +) # when set, POST routes need "Authorization: Bearer " +CANVAS_LEN = 64 # the served canvas length. A request may be narrower. +CANVAS_STEP = 16 # request widths are multiples of this +VOCAB = 262144 +TURN_CLOSE = 106 +PAD = 0 +TOPK = 20 +MAX_QUESTIONS = 64 # per request +MAX_SAMPLES = 32 # reads per question, fixed or auto +MAX_PARALLEL = 16 # question groups read at once +# the empty thought block the chat template leaves to the model +SCAFFOLD_TEXT = "<|channel>thought\n" +SCAFFOLD = None +THOUGHT_OPEN = None +THOUGHT_CLOSE = None + + +# ---------------------------------------------------------------------------- +# Schema +# ---------------------------------------------------------------------------- + + +class SchemaError(ValueError): + pass + + +def parse_schema(value): + if ( + not isinstance(value, dict) + or not isinstance(value.get("questions"), list) + or not value["questions"] + ): + raise SchemaError("schema: needs a non-empty questions array") + if len(value["questions"]) > MAX_QUESTIONS: + raise SchemaError(f"schema: at most {MAX_QUESTIONS} questions") + qs = [] + seen = set() + for q in value["questions"]: + qid = str(q.get("id", "")).strip() + if not qid or ":" in qid or "\n" in qid: + raise SchemaError( + f"question id {qid!r} must be non-empty, no ':' or newline" + ) + if qid in seen: + raise SchemaError(f"duplicate question id {qid!r}") + seen.add(qid) + kind = q.get("type") + if kind in ("noul", "bool", "boolean"): + kind = "noul" + crit = q.get("criteria") or {} + choices = [("yes", crit.get("true")), ("no", crit.get("false"))] + labels = ["yes", "no"] + elif kind == "choice": + opts = q.get("options") or [] + choices = [ + (o["name"], o.get("description")) + if isinstance(o, dict) + else (str(o), None) + for o in opts + ] + labels = [chr(ord("A") + i) for i in range(len(choices))] + elif kind == "score": + choices = [(str(level), None) for level in (q.get("levels") or [])] + labels = ( + [str(i + 1) for i in range(len(choices))] + if len(choices) <= 9 + else [chr(ord("A") + i) for i in range(len(choices))] + ) + else: + raise SchemaError(f"question {qid!r}: unknown type {kind!r}") + if len(choices) < 2: + raise SchemaError(f"question {qid!r}: needs at least two alternatives") + if len(choices) > 26: + raise SchemaError(f"question {qid!r}: at most 26 alternatives") + deps = q.get("depends_on") or [] + ask_if = q.get("ask_if") or {} + if not isinstance(deps, list) or not all(isinstance(d, str) for d in deps): + raise SchemaError( + f"question {qid!r}: depends_on must be a list of question ids" + ) + if not isinstance(ask_if, dict) or not all( + isinstance(v, list) and v for v in ask_if.values() + ): + raise SchemaError( + f"question {qid!r}: ask_if must map a question id to a " + "non-empty list of its answers" + ) + qs.append( + { + "id": qid, + "type": kind, + "instructions": str(q.get("instructions", "")), + "choices": choices, + "labels": labels, + "depends_on": list(dict.fromkeys(list(deps) + list(ask_if))), + "ask_if": ask_if, + "alone": bool(q.get("alone", False)), + } + ) + by_id = {q["id"]: q for q in qs} + for q in qs: + for dep in q["depends_on"]: + if dep not in by_id or dep == q["id"]: + raise SchemaError( + f"question {q['id']!r}: depends on unknown question {dep!r}" + ) + for dep, vals in q["ask_if"].items(): + names = [c[0] for c in by_id[dep]["choices"]] + if any(v not in names for v in vals): + raise SchemaError( + f"question {q['id']!r}: ask_if values for {dep!r} " + f"must be among {names}" + ) + schedule(qs) # refuses a cycle + samples = value.get("samples", "auto") + if samples == "auto": + policy = { + "mode": "auto", + "max": max(1, min(int(value.get("auto_max", 4)), MAX_SAMPLES)), + "threshold": float(value.get("auto_threshold", 0.1)), + } + elif isinstance(samples, int) and samples >= 1: + policy = {"mode": "fixed", "n": min(samples, MAX_SAMPLES)} + else: + raise SchemaError('schema: samples must be a positive count or "auto"') + ask = value.get("ask") + if ask is not None: + if not isinstance(ask, list) or not ask or any(a not in seen for a in ask): + raise SchemaError("schema: ask must list question ids from this schema") + for q in qs: + if q["id"] in ask and any(d not in ask for d in q["depends_on"]): + raise SchemaError( + f"schema: ask names {q['id']!r} but not everything it depends on" + ) + chunk_rows = value.get("chunk_rows") + if chunk_rows is not None and (not isinstance(chunk_rows, int) or chunk_rows < 8): + raise SchemaError("schema: chunk_rows must be an integer of at least 8") + chunk_prompt = value.get("chunk_prompt", "own") + if chunk_prompt not in ("shared", "own"): + raise SchemaError('schema: chunk_prompt must be "shared" or "own"') + sequential = bool(value.get("sequential", False)) + think = value.get("think", 0) + if isinstance(think, bool) or not isinstance(think, int) or not 0 <= think <= 4096: + raise SchemaError("schema: think must be a thought budget in tokens, 0 to 4096") + return { + "questions": qs, + "instructions": value.get("instructions"), + "policy": policy, + "steps": max(1, min(int(value.get("steps", 1)), 8)), + "think": think, + "ask": ask, + "chunk_rows": chunk_rows, + "chunk_prompt": chunk_prompt, + "sequential": sequential, + "format": "lines" if len(qs) <= 10 else "indexed", + } + + +# Answer template shape: (join between questions, what precedes the label, +# reply instruction). A small schema gets "lines", which is readable. +# "indexed" ("0yes 1no") costs three tokens a question against four or five +# and agreed with "lines" on every set tried: 42 booleans, ten 26-way +# choices, twenty 5-level scores. Past ten questions the saved rows keep a +# schema in one read. Two tokens a question, or no id at all, loses +# alignment beyond about twenty questions, because the id ties a label to +# its question. +FORMATS = { + "lines": ( + "\n", + "{id}: ", + 'Reply with one line per question, in this order, formatted as "id: label".', + ), + "indexed": ( + " ", + "{id}", + "Reply on one line with each question's id immediately followed by its " + "label, separated by single spaces.", + ), +} + + +def system_text(schema, chunked=False): + s = ( + "Answer a fixed set of questions about the state the user provides. " + "Each question lists its allowed answers; reply with exactly one label " + "per question.\n" + ) + if schema.get("instructions"): + s += "\n" + str(schema["instructions"]).strip() + "\n" + for q in schema["questions"]: + s += f"\nQuestion {q['id']}: {q['instructions'].strip()}\n" + for (name, desc), label in zip(q["choices"], q["labels"]): + if q["type"] == "noul": + s += f" {label}: {str(desc).strip()}\n" if desc else f" {label}\n" + elif desc: + s += f" {label}: {name} ({str(desc).strip()})\n" + else: + s += f" {label}: {name}\n" + s += "\n" + FORMATS[schema.get("format", "lines")][2] + if chunked: + s += ( + " A reply may cover only some of the questions; answer every line " + "that is present." + ) + return s + + +def answer_text(qs, labels, fmt="lines"): + join, lead, _ = FORMATS[fmt] + return join.join( + lead.format(id=q["id"]) + q["labels"][i] for q, i in zip(qs, labels) + ) + + +def enc(text): + return TOK.encode(text, add_special_tokens=False) + + +def init_tokenizer(tok): + global TOK, SCAFFOLD, THOUGHT_OPEN, THOUGHT_CLOSE + TOK = tok + THOUGHT_OPEN = enc("<|channel>thought\n") + THOUGHT_CLOSE = enc("") + SCAFFOLD = enc(SCAFFOLD_TEXT) + assert THOUGHT_OPEN + THOUGHT_CLOSE == SCAFFOLD, ( + "the thought tags must tokenize apart" + ) + + +def resolve_template(qs, head, lead, fmt): + """Tokenize the answer template and find each question's slot. Every label + must change exactly one token, at the same position for all of a question's + labels, or this raises SchemaError. ``head`` is the token run the canvas + starts with: the empty thought block for a plain read, and empty when the + prompt already ends the thought channel. ``lead`` is the text before the + first answer: the join when earlier answers are in the prompt, so the + tokens match one joint template.""" + base_labels = [0] * len(qs) + base = head + enc(lead + answer_text(qs, base_labels, fmt)) + if len(base) + 1 > CANVAS_LEN: + raise SchemaError( + f"answer template is {len(base)} tokens; the canvas holds {CANVAS_LEN - 1}" + ) + if len(qs) == 1 and len(base) + 1 > CANVAS_LEN: + raise SchemaError( + f"question {qs[0]['id']!r} alone needs {len(base) + 1} canvas rows" + ) + slots = [] + for qi, q in enumerate(qs): + pos = None + ids = [0] * len(q["labels"]) + for li in range(1, len(q["labels"])): + labels = list(base_labels) + labels[qi] = li + e = head + enc(lead + answer_text(qs, labels, fmt)) + if len(e) != len(base): + raise SchemaError( + f"question {q['id']!r}: label {q['labels'][li]!r} is not a " + "single token" + ) + diffs = [i for i in range(len(e)) if e[i] != base[i]] + if len(diffs) != 1 or (pos is not None and diffs[0] != pos): + raise SchemaError( + f"question {q['id']!r}: labels do not share one template slot" + ) + pos = diffs[0] + ids[li] = e[pos] + ids[0] = base[pos] + if len(set(ids)) != len(ids): + raise SchemaError( + f"question {q['id']!r}: two labels tokenize to the same id" + ) + slots.append({"pos": pos, "label_ids": ids}) + return base, slots + + +_template_cache = {} + + +def template_for(schema, head, lead): + fmt = schema.get("format", "lines") + key = json.dumps( + [head, lead, fmt] + [(q["id"], q["labels"]) for q in schema["questions"]] + ) + if key not in _template_cache: + _template_cache[key] = resolve_template(schema["questions"], head, lead, fmt) + return _template_cache[key] + + +# ---------------------------------------------------------------------------- +# Reads +# ---------------------------------------------------------------------------- + + +def canvas_width(template): + """Smallest multiple of CANVAS_STEP that holds the template and the turn close.""" + need = len(template) + 1 + return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP) + + +def pin_xargs(template, slots, steps): + """Past one denoise step the template must be held, or accept/renoise + rewrites it: pin every canvas position that is not an answer slot.""" + if steps <= 1: + return {} + free = {s["pos"] for s in slots} + return { + "diffusion_pinned": [p for p in range(canvas_width(template)) if p not in free] + } + + +def build_canvas(template, slots, seed): + rng = random.Random(seed) + canvas = list(template) + [TURN_CLOSE] + canvas += [PAD] * (canvas_width(template) - len(canvas)) + for s in slots: + canvas[s["pos"]] = rng.randrange(VOCAB) + return canvas + + +def label_id_union(slots): + ids = sorted({i for s in slots for i in s["label_ids"]}) + return ids[:128] # vLLM's cap per request. A schema needs far fewer. + + +def upstream_chat(body, timeout=600): + req = urllib.request.Request( + ARGS.upstream.rstrip("/") + "/v1/chat/completions", + data=json.dumps(body).encode(), + headers={"content-type": "application/json"}, + ) + return json.load(urllib.request.urlopen(req, timeout=timeout)) + + +def upstream_completions(body, timeout=600): + req = urllib.request.Request( + ARGS.upstream.rstrip("/") + "/v1/completions", + data=json.dumps(body).encode(), + headers={"content-type": "application/json"}, + ) + return json.load(urllib.request.urlopen(req, timeout=timeout)) + + +def chat_prompt_ids(sys_text, state_text, thinking=False): + """The prompt the chat endpoint would build, as token ids, ending after + the model turn marker. Text states only. ``thinking`` turns the chat + template's thinking marker on.""" + messages = [ + {"role": "system", "content": sys_text}, + {"role": "user", "content": state_text}, + ] + out = TOK.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True, enable_thinking=thinking + ) + ids = ( + out["input_ids"] if hasattr(out, "keys") else out + ) # newer transformers return a dict + return [int(t) for t in ids] + + +def think(sys_text, state_text, budget): + """A read prefix that ends with a thought the model wrote: the chat prompt + with thinking on, the open tag, up to ``budget`` generated tokens, the + close tag. Returns the prefix and a diagnostics dict for the thought.""" + prompt = chat_prompt_ids(sys_text, state_text, thinking=True) + THOUGHT_OPEN + started = time.time() + d = upstream_completions( + { + "model": ARGS.model, + "prompt": prompt, + "max_tokens": budget, + "logprobs": 0, + "return_tokens_as_token_ids": True, + "stop_token_ids": THOUGHT_CLOSE, + } + ) + ids = [int(t.split(":")[1]) for t in d["choices"][0]["logprobs"]["tokens"]] + closed = THOUGHT_CLOSE[0] in ids + if closed: + ids = ids[: ids.index(THOUGHT_CLOSE[0])] + info = { + "tokens": len(ids), + "closed": closed, + "ms": (time.time() - started) * 1e3, + "text": TOK.decode(ids), + } + return prompt + ids + THOUGHT_CLOSE, info + + +def think_chat(sys_text, state_content, budget): + """A thought written with the image in view: the chat endpoint with + thinking on, capped at ``budget`` tokens and cut at the close tag. + Returns the thought's token ids and a diagnostics dict.""" + body = { + "model": ARGS.model, + "messages": [ + {"role": "system", "content": sys_text}, + {"role": "user", "content": state_content}, + ], + "max_tokens": budget, + "logprobs": True, + "top_logprobs": 0, + "return_tokens_as_token_ids": True, + "stop_token_ids": THOUGHT_CLOSE, + "chat_template_kwargs": {"enable_thinking": True}, + } + started = time.time() + d = upstream_chat(body) + choice = d["choices"][0] + ids = [ + int(t["token"].split(":")[1]) + for t in (choice.get("logprobs") or {}).get("content") or [] + ] + if ids[: len(THOUGHT_OPEN)] == THOUGHT_OPEN: # the model opens the channel itself + ids = ids[len(THOUGHT_OPEN) :] + closed = THOUGHT_CLOSE[0] in ids + if closed: + ids = ids[: ids.index(THOUGHT_CLOSE[0])] + return ids, { + "tokens": len(ids), + "closed": closed, + "ms": (time.time() - started) * 1e3, + "text": TOK.decode(ids), + } + + +def one_read( + schema, template, slots, sys_text, state_content, seed, prefix=None, thinking=False +): + if prefix is not None: + return one_read_continuation(schema, template, slots, prefix, seed) + messages = [ + {"role": "system", "content": sys_text}, + {"role": "user", "content": state_content}, + ] + body = { + "model": ARGS.model, + "messages": messages, + "max_tokens": len(template) + 1, + "logprobs": True, + "top_logprobs": TOPK, + # Exact logprobs for every label at every position. With a long + # option list most labels never rank in the top-k, and the model's + # mass sits on tokens that spell the option name instead. + "logprob_token_ids": label_id_union(slots), + "return_tokens_as_token_ids": True, + "chat_template_kwargs": {"enable_thinking": thinking}, + "vllm_xargs": { + "diffusion_seed_canvas": build_canvas(template, slots, seed), + "diffusion_canvas_length": canvas_width(template), + "diffusion_max_steps": schema["steps"], + "diffusion_read_only": True, + **pin_xargs(template, slots, schema["steps"]), + }, + } + d = upstream_chat(body) + content = d["choices"][0]["logprobs"]["content"] + out = [] + for q, s in zip(schema["questions"], slots): + top = { + int(t["token"].split(":")[1]): t["logprob"] + for t in content[s["pos"]]["top_logprobs"] + } + out.append(slot_distribution(top, s["label_ids"])) + return out, d.get("usage", {}) + + +def slot_distribution(top, label_ids): + """Label probabilities at one slot from the returned logprobs: every + label's own value plus the argmax token. Read-only logprobs are at + temperature 1, so the label softmax uses them directly. The entropy is + over that returned set.""" + floor = min(top.values()) - 5.0 + lp_t = [top.get(i, floor) for i in label_ids] + mx = max(lp_t) + ex = [math.exp(x - mx) for x in lp_t] + probs = [e / sum(ex) for e in ex] + top_p = [math.exp(v) for v in top.values()] + return { + "probs": probs, + "label_mass": sum(math.exp(x) for x in lp_t), + "entropy": -sum(p * math.log(p) for p in top_p if p > 0), + "argmax_is_label": max(top, key=top.get) in label_ids, + } + + +def one_read_continuation(schema, template, slots, prompt_ids, seed): + """A read whose prompt already holds the thought scaffold and earlier + answer lines, sent as token ids so the chat template cannot alter it.""" + body = { + "model": ARGS.model, + "prompt": prompt_ids, + "max_tokens": len(template) + 1, + "logprobs": TOPK, + "logprob_token_ids": label_id_union(slots), + "return_tokens_as_token_ids": True, + "vllm_xargs": { + "diffusion_seed_canvas": build_canvas(template, slots, seed), + "diffusion_canvas_length": canvas_width(template), + "diffusion_max_steps": schema["steps"], + "diffusion_read_only": True, + **pin_xargs(template, slots, schema["steps"]), + }, + } + d = upstream_completions(body) + rows = d["choices"][0]["logprobs"]["top_logprobs"] + out = [] + for q, sl in zip(schema["questions"], slots): + top = {int(k.split(":")[1]): v for k, v in rows[sl["pos"]].items()} + out.append(slot_distribution(top, sl["label_ids"])) + return out, d.get("usage", {}) + + +def read_many( + schema, + template, + slots, + sys_text, + state_content, + seed, + n, + prefix=None, + thinking=False, +): + results = [None] * n + errors = [None] * n + usages = [None] * n + + def run(k): + try: + results[k], usages[k] = one_read( + schema, + template, + slots, + sys_text, + state_content, + seed + k * 7919, + prefix, + thinking, + ) + except Exception as e: # raised again below + errors[k] = e + + threads = [threading.Thread(target=run, args=(k,)) for k in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + for e in errors: + if e is not None: + raise e + return results, usages + + +def schedule(qs): + """Questions in stages: a question's stage comes after the stages of + everything it depends on. Declaration order is kept within a stage.""" + ids = {q["id"] for q in qs} + pending = list(qs) + done: set = set() + levels = [] + while pending: + level = [ + q + for q in pending + if all(d in done or d not in ids for d in q["depends_on"]) + ] + if not level: + raise SchemaError( + "schema: dependency cycle among " + ", ".join(q["id"] for q in pending) + ) + levels.append(level) + done |= {q["id"] for q in level} + pending = [q for q in pending if q["id"] not in done] + return levels + + +def chunk_groups(schema, qs): + """``qs`` split, in order, into the fewest groups whose answer templates + fit ``chunk_rows`` (the canvas by default). A question marked alone gets + its own group.""" + limit = schema.get("chunk_rows") or CANVAS_LEN + groups, group = [], [] + for q in qs: + if q["alone"]: + if group: + groups.append(group) + group = [] + groups.append([q]) + continue + trial = group + [q] + rows = ( + len(SCAFFOLD) + + len( + enc(answer_text(trial, [0] * len(trial), schema.get("format", "lines"))) + ) + + 1 + ) + if rows > limit and group: + groups.append(group) + group = [q] + else: + group = trial + if group: + groups.append(group) + return groups + + +def answer_name(q, a): + """The answer as the name ask_if compares against: yes or no, an option + name, or a level name.""" + if a is None: + return None + return a["label"] if q["type"] == "noul" else a.get("choice", a.get("level")) + + +def decide(schema, state_content, seed): + """One decision. Questions run in stages by their dependencies. A stage + is one joint read, chunked by the canvas, with a question marked alone + in its own read. Later stages condition on every earlier answer: a + prefilled continuation for a text state, or the answers restated in the + state for an image. A question whose ask_if condition failed is skipped + and its answer is null.""" + started = time.time() + qs = [ + q + for q in schema["questions"] + if not schema.get("ask") or q["id"] in schema["ask"] + ] + levels = schedule(qs) + text_state = isinstance(state_content, str) + fmt = schema["format"] + join = FORMATS[fmt][0] + # More than one read in sequence needs the full question list in every + # prompt, so that later reads continue one answer. + chained = len(levels) > 1 or schema["sequential"] + sys_full = system_text(schema, chunked=not text_state and chained) + base_ids, thought = None, None + if chained and text_state: + if schema["think"]: + base_ids, thought = think(sys_full, state_content, schema["think"]) + else: + base_ids = chat_prompt_ids(sys_full, state_content) + SCAFFOLD + shared = schema["chunk_prompt"] == "shared" + answered, lines, earlier = {}, [], [] + parts, stages, chunks, skipped = [], [], [], {} + by_id = {q["id"]: q for q in schema["questions"]} + + def run(group, k, conditioned): + sub = dict(schema, questions=group) + # The thought is written once: in the prefix of a chained text decision, + # or in the first read of anything else. + sub["think"] = 0 if (chained and text_state) or conditioned else schema["think"] + if text_state: + if conditioned: + prefix, lead, sys_text = ( + base_ids + enc(join.join(lines)), + join, + sys_full, + ) + elif chained: + prefix, lead, sys_text = (base_ids if thought else None), "", sys_full + else: + prefix, lead, sys_text = ( + None, + "", + (system_text(schema, chunked=True) if shared else system_text(sub)), + ) + return decide_group( + sub, sys_text, state_content, seed + 104729 * k, prefix, lead + ) + state = state_content + if conditioned: + text = next( + (p["text"] for p in state_content if p.get("type") == "text"), "" + ) + text += "\n\nAnswers so far:\n" + "\n".join(earlier) + state = [p for p in state_content if p.get("type") != "text"] + [ + {"type": "text", "text": text} + ] + sys_text = sys_full + else: + sys_text = sys_full if (chained or shared) else system_text(sub) + return decide_group(sub, sys_text, state, seed + 104729 * k) + + def absorb(group, body, rows): + parts.append((body, rows)) + chunks.append([q["id"] for q in group]) + answered.update(body["answers"]) + lines.append( + answer_text( + group, + [q["labels"].index(body["answers"][q["id"]]["label"]) for q in group], + fmt, + ) + ) + earlier.extend( + f"{q['id']}: {answer_name(q, body['answers'][q['id']])}" for q in group + ) + + k = 0 + for level in levels: + asked = [] + for q in level: + failed = next( + ( + (dep, vals) + for dep, vals in q["ask_if"].items() + if answer_name(by_id[dep], answered.get(dep)) not in vals + ), + None, + ) + if failed: + answered[q["id"]] = None + skipped[q["id"]] = { + "because": failed[0], + "was": answer_name(by_id[failed[0]], answered.get(failed[0])), + "wanted": failed[1], + } + continue + asked.append(q) + if not asked: + continue + stages.append([q["id"] for q in asked]) + groups = chunk_groups(schema, asked) + conditioned = bool(lines) + if schema["sequential"] or len(groups) == 1: + for group in groups: + body, rows = run( + group, k, conditioned or (schema["sequential"] and bool(lines)) + ) + absorb(group, body, rows) + k += 1 + else: + with ThreadPoolExecutor(max_workers=min(len(groups), MAX_PARALLEL)) as ex: + results = list( + ex.map( + lambda gk, conditioned=conditioned: run( + gk[1], gk[0], conditioned + ), + [(k + i, g) for i, g in enumerate(groups)], + ) + ) + for group, (body, rows) in zip(groups, results): + absorb(group, body, rows) + k += len(groups) + + answers = {q["id"]: answered.get(q["id"]) for q in qs} + diag_q = {} + for body, _ in parts: + diag_q.update(body["diagnostics"]["questions"]) + # A thought written here (a chained text decision) is outside every + # group's row count. One written inside a group is already counted there. + extra_rows = thought["tokens"] if thought else 0 + if thought is None: + thoughts = [b["diagnostics"].get("thought") for b, _ in parts] + thought = ( + thoughts[0] if len(parts) == 1 else ([t for t in thoughts if t] or None) + ) + one = len(parts) == 1 and not skipped + diagnostics = { + "steps": schema["steps"], + "stages": stages, + "skipped": skipped, + "chunks": chunks, + "chunk_prompt": "full" if chained else schema["chunk_prompt"], + "sequential": schema["sequential"], + "conditioning": ( + None + if len(stages) <= 1 and not schema["sequential"] + else ("prefill" if text_state else "restated") + ), + "thought": thought, + "samples": ( + parts[0][0]["diagnostics"]["samples"] + if one + else { + "n": [b["diagnostics"]["samples"]["n"] for b, _ in parts], + "tops": [b["diagnostics"]["samples"]["tops"] for b, _ in parts], + "policy": [b["diagnostics"]["samples"]["policy"] for b, _ in parts], + } + ), + "timing": { + "total_ms": (time.time() - started) * 1e3, + "reads": sum(b["diagnostics"]["timing"]["reads"] for b, _ in parts), + }, + "prompt_tokens": max( + (b["diagnostics"].get("prompt_tokens") or 0) for b, _ in parts + ) + or None, + "questions": diag_q, + "engine": "vllm", + } + return {"answers": answers, "diagnostics": diagnostics}, sum( + rows for _, rows in parts + ) + extra_rows + + +def decide_group(schema, sys_text, state_content, seed, prefix=None, lead=""): + started = time.time() + thought = None + head = SCAFFOLD if prefix is None else [] + thinking = False + if prefix is None and schema["think"]: + if isinstance(state_content, str): + prefix, thought = think(sys_text, state_content, schema["think"]) + head = [] + else: + # With images the model writes the thought with the image in view, and + # the server seeds it into the canvas ahead of the answer, so the read + # keeps the image. The canvas bounds the thought. + answer = enc( + answer_text( + schema["questions"], + [0] * len(schema["questions"]), + schema.get("format", "lines"), + ) + ) + fits = CANVAS_LEN - 1 - len(THOUGHT_OPEN) - len(THOUGHT_CLOSE) - len(answer) + if fits < 8: + raise SchemaError( + f"think: the canvas leaves {fits} rows for a thought " + "beside this template" + ) + ids, thought = think_chat( + sys_text, state_content, min(schema["think"], fits) + ) + thought["budget"] = min(schema["think"], fits) + head = THOUGHT_OPEN + ids + THOUGHT_CLOSE + thinking = True + template, slots = template_for(schema, head, lead) + policy = schema["policy"] + if policy["mode"] == "fixed": + reads, usages = read_many( + schema, + template, + slots, + sys_text, + state_content, + seed, + policy["n"], + prefix, + thinking, + ) + extended = None + first_entropy = None + else: + reads, usages = read_many( + schema, template, slots, sys_text, state_content, seed, 1, prefix, thinking + ) + first_entropy = { + q["id"]: r["entropy"] for q, r in zip(schema["questions"], reads[0]) + } + extended = ( + max(first_entropy.values()) > policy["threshold"] and policy["max"] > 1 + ) + if extended: + more, more_usages = read_many( + schema, + template, + slots, + sys_text, + state_content, + seed + 1, + policy["max"] - 1, + prefix, + thinking, + ) + reads += more + usages += more_usages + prompt_tokens = next( + (u["prompt_tokens"] for u in usages if u and u.get("prompt_tokens")), None + ) + elapsed_ms = (time.time() - started) * 1e3 + + answers = {} + diag_q = {} + n = len(reads) + for qi, q in enumerate(schema["questions"]): + per = [r[qi]["probs"] for r in reads] + mean = [sum(p[i] for p in per) / n for i in range(len(q["labels"]))] + top = max(range(len(mean)), key=lambda i: mean[i]) + a = { + "type": q["type"], + "label": q["labels"][top], + "confidence": mean[top], + "probabilities": {c[0]: m for c, m in zip(q["choices"], mean)}, + } + if q["type"] == "noul": + a["noul"] = mean[0] + elif q["type"] == "choice": + a["choice"] = q["choices"][top][0] + else: + a["score"] = sum((i + 1) * m for i, m in enumerate(mean)) + a["level"] = q["choices"][top][0] + if n > 1: + var = sum((p[top] - mean[top]) ** 2 for p in per) / (n - 1) + a["stderr"] = (var / n) ** 0.5 + a["agreement"] = ( + sum(1 for p in per if max(range(len(p)), key=lambda i: p[i]) == top) / n + ) + answers[q["id"]] = a + diag_q[q["id"]] = { + "pos": slots[qi]["pos"], + "entropy": [r[qi]["entropy"] for r in reads], + "label_mass": reads[0][qi]["label_mass"], + "argmax_is_label": reads[0][qi]["argmax_is_label"], + } + tops = [ + { + q["id"]: [ + q["labels"][ + max(range(len(r[qi]["probs"])), key=lambda i: r[qi]["probs"][i]) + ], + max(r[qi]["probs"]), + r[qi]["entropy"], + ] + for qi, q in enumerate(schema["questions"]) + } + for r in reads + ] + return { + "answers": answers, + "diagnostics": { + "steps": schema["steps"], + "samples": { + "n": n, + "tops": tops, + "policy": dict( + policy, extended=extended, first_read_entropy=first_entropy + ), + }, + "timing": {"total_ms": elapsed_ms, "reads": n}, + "thought": thought, + "prompt_tokens": prompt_tokens, + "questions": diag_q, + "engine": "vllm", + }, + }, len(template) + 1 + (thought["tokens"] if thought else 0) + + +# ---------------------------------------------------------------------------- +# Jev's contract +# ---------------------------------------------------------------------------- + +JEV_EXTENSIONS = ( + "instructions", + "samples", + "auto_max", + "auto_threshold", + "steps", + "think", + "ask", + "chunk_rows", + "chunk_prompt", + "sequential", +) + + +def jev_schema(body): + """This server's schema from a Jev request body.""" + qs = body.get("questions") + if not isinstance(qs, dict) or not qs: + raise SchemaError("questions: needs a non-empty map of id -> question") + out = [] + for qid, q in qs.items(): + if not isinstance(q, dict): + raise SchemaError(f"question {qid!r}: must be an object") + kind, crit, ins = q.get("type"), q.get("criteria"), q.get("instructions", "") + item = { + "id": qid, + "type": kind, + "instructions": ins if isinstance(ins, str) else json.dumps(ins), + } + if kind == "noul": + if crit is not None and not isinstance(crit, dict): + raise SchemaError( + f"question {qid!r}: noul criteria must be an object with " + "true and false" + ) + item["criteria"] = crit + elif kind == "choice": + if not isinstance(crit, dict) or not crit: + raise SchemaError( + f"question {qid!r}: choice criteria must map option names " + "to descriptions" + ) + item["options"] = [ + {"name": str(n), "description": d} for n, d in crit.items() + ] + elif kind == "score": + if not isinstance(crit, list): + raise SchemaError( + f"question {qid!r}: score criteria must be an ordered list " + "of levels" + ) + item["levels"] = crit + else: + raise SchemaError(f"question {qid!r}: unknown type {kind!r}") + for key in ("depends_on", "ask_if", "alone"): + if key in q: + item[key] = q[key] + out.append(item) + schema = {k: body[k] for k in JEV_EXTENSIONS if k in body} + schema["questions"] = out + return parse_schema(schema) + + +def image_part(content_type, data): + return { + "type": "image_url", + "image_url": { + "url": f"data:{content_type};base64," + base64.b64encode(data).decode() + }, + } + + +def jev_images(value): + """Image parts from the body's "images": data URLs, or objects with + content_type and base64.""" + parts = [] + for i, im in enumerate(value or []): + if isinstance(im, str) and im.startswith("data:image/"): + parts.append({"type": "image_url", "image_url": {"url": im}}) + elif ( + isinstance(im, dict) + and str(im.get("content_type", "")).startswith("image/") + and isinstance(im.get("base64"), str) + ): + parts.append( + { + "type": "image_url", + "image_url": { + "url": f"data:{im['content_type']};base64,{im['base64']}" + }, + } + ) + else: + raise SchemaError( + f"images[{i}]: a data:image/... URL or an object with " + "content_type and base64" + ) + return parts + + +def jev_state(body, image_parts=()): + """The user message: the state as text (as given, or as JSON), with any + images ahead of it.""" + state = body.get("state") + if state is None: + raise SchemaError("state: required") + text = state if isinstance(state, str) else json.dumps(state) + if not image_parts: + return text + return list(image_parts) + [{"type": "text", "text": text}] + + +def jev_answer(q, a): + if a is None: + return None + if q["type"] == "noul": + return {"type": "noul", "noul": a["noul"]} + if q["type"] == "choice": + return { + "type": "choice", + "choice": a["choice"], + "probabilities": a["probabilities"], + "confidence": a["confidence"], + } + names = [c[0] for c in q["choices"]] + probs = {str(i): a["probabilities"][n] for i, n in enumerate(names)} + return { + "type": "score", + "score": sum(i * p for i, p in enumerate(probs.values())), + "legend": {str(i): n for i, n in enumerate(names)}, + "probabilities": probs, + "confidence": a["confidence"], + } + + +# ---------------------------------------------------------------------------- +# HTTP +# ---------------------------------------------------------------------------- + + +def message_text(m): + c = m.get("content", "") + if isinstance(c, list): + return "".join(p.get("text", "") for p in c if isinstance(p, dict)) + return c if isinstance(c, str) else "" + + +class Handler(BaseHTTPRequestHandler): + def log_message(self, fmt, *args): + pass + + def _json(self, code, obj): + body = json.dumps(obj).encode() + self.send_response(code) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path == "/health": + return self._json(200, {"status": "ok"}) + return self._json(404, {"error": {"message": "unknown route"}}) + + def _read_request(self): + """-> (body, image parts): a JSON body, or multipart/form-data with the + JSON in a part named request and each image as a file part, in order.""" + raw = self.rfile.read(int(self.headers.get("content-length", "0"))) + ctype = self.headers.get("content-type", "") + if not ctype.lower().startswith("multipart/form-data"): + return json.loads(raw), [] + from email.parser import BytesParser + from email.policy import HTTP + + msg = BytesParser(policy=HTTP).parsebytes( + b"Content-Type: " + ctype.encode() + b"\r\n\r\n" + raw + ) + body, images = None, [] + for part in msg.iter_parts(): + name = part.get_param("name", header="content-disposition") + data = part.get_payload(decode=True) + if name == "request": + body = json.loads(data) + elif part.get_content_type().startswith("image/"): + images.append(image_part(part.get_content_type(), data)) + else: + raise ValueError( + f"part {name!r}: neither the request JSON nor an image" + ) + if body is None: + raise ValueError( + "multipart needs a part named request holding the JSON body" + ) + return body, images + + def do_POST(self): + if API_KEY and self.headers.get("authorization", "") != f"Bearer {API_KEY}": + return self._json( + 401, + { + "error": { + "message": "missing or wrong bearer token", + "type": "authentication_error", + } + }, + ) + if self.path == "/v1/raw/chat/completions": + return self._raw_chat() + try: + req, images = self._read_request() + except Exception as e: + return self._json( + 400, + { + "error": { + "message": f"invalid body: {e}", + "type": "invalid_request_error", + } + }, + ) + if self.path == "/v1/systemone": + return self._systemone(req, images) + if self.path == "/v1/chat/completions": + return self._chat(req) + return self._json(404, {"error": {"message": "unknown route"}}) + + def _raw_chat(self): + """Pass the body and status through to vLLM's chat completions.""" + raw = self.rfile.read(int(self.headers.get("content-length", "0"))) + req = urllib.request.Request( + ARGS.upstream.rstrip("/") + "/v1/chat/completions", + data=raw, + headers={ + "content-type": self.headers.get("content-type", "application/json") + }, + ) + try: + with urllib.request.urlopen(req, timeout=600) as r: + code, body = r.status, r.read() + except urllib.error.HTTPError as e: + code, body = e.code, e.read() + self.send_response(code) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _decide(self, schema, state, seed): + """-> (status, body) with the error body already shaped.""" + try: + return 200, decide(schema, state, seed) + except SchemaError as e: + return 422, {"error": {"message": str(e), "type": "validation_error"}} + except urllib.error.HTTPError as e: + return 502, { + "error": { + "message": ( + f"upstream {e.code}: {e.read()[:300].decode(errors='replace')}" + ), + "type": "server_error", + } + } + except Exception as e: + return 500, {"error": {"message": repr(e), "type": "server_error"}} + + def _systemone(self, req, images): + try: + schema = jev_schema(req) + state = jev_state(req, images + jev_images(req.get("images"))) + except SchemaError as e: + return self._json( + 422, {"error": {"message": str(e), "type": "validation_error"}} + ) + code, result = self._decide(schema, state, int(req.get("seed", 42))) + if code != 200: + return self._json(code, result) + body, completion_tokens = result + answers = { + q["id"]: jev_answer(q, body["answers"][q["id"]]) + for q in schema["questions"] + } + labels = " ".join( + f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items() + ) + print( + f"systemone: {labels} " + f"reads={body['diagnostics']['timing']['reads']} " + f"{body['diagnostics']['timing']['total_ms']:.0f}ms", + flush=True, + ) + self._json( + 200, + { + "model": ARGS.model, + "answers": answers, + # vLLM's prompt count when the reads reported one, since it covers + # images. Otherwise the tokenizer's count of the text prompt. + "usage": { + "input_tokens": body["diagnostics"].get("prompt_tokens") + or ( + len(chat_prompt_ids(system_text(schema), state)) + if isinstance(state, str) + else 0 + ), + "output_tokens": completion_tokens, + }, + "diagnostics": body["diagnostics"], + }, + ) + + def _chat(self, req): + msgs = req.get("messages") or [] + if ( + len(msgs) != 2 + or msgs[0].get("role") not in ("system", "developer") + or msgs[1].get("role") != "user" + ): + return self._json( + 400, + { + "error": { + "message": ( + "a structured request is exactly two messages: the " + "schema (system) and the state JSON (user)" + ), + "type": "invalid_request_error", + } + }, + ) + try: + schema_value = json.loads(message_text(msgs[0])) + schema = parse_schema(schema_value) + content = msgs[1].get("content", "") + has_image = isinstance(content, list) and any( + isinstance(p, dict) and p.get("type") in ("image_url", "image") + for p in content + ) + if has_image: + # image parts pass through to vLLM unchanged, with text parts as context + state = content + else: + state = message_text(msgs[1]).strip() + json.loads(state) + except SchemaError as e: + return self._json( + 400, {"error": {"message": str(e), "type": "invalid_request_error"}} + ) + except Exception as e: + return self._json( + 400, + { + "error": { + "message": ( + "system must be a JSON question schema and user " + f"must be JSON state or image parts: {e}" + ), + "type": "invalid_request_error", + } + }, + ) + code, result = self._decide(schema, state, int(req.get("seed", 42))) + if code != 200: + if code == 422: + result["error"]["type"] = "invalid_request_error" + code = 400 + return self._json(code, result) + body, completion_tokens = result + content = json.dumps(body, indent=2) + labels = " ".join( + f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items() + ) + print( + f"structured: {labels} " + f"reads={body['diagnostics']['timing']['reads']} " + f"{body['diagnostics']['timing']['total_ms']:.0f}ms", + flush=True, + ) + self._json( + 200, + { + "id": f"chatcmpl-{int(time.time() * 1000)}", + "object": "chat.completion", + "created": int(time.time()), + "model": req.get("model", "dgemma-structured"), + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 0, + "completion_tokens": completion_tokens, + "total_tokens": completion_tokens, + }, + }, + ) + + +def self_signed(cert_dir): + """Paths of a self-signed certificate and key in cert_dir, made with + openssl on first use.""" + os.makedirs(cert_dir, exist_ok=True) + cert, key = os.path.join(cert_dir, "djev.crt"), os.path.join(cert_dir, "djev.key") + if not (os.path.exists(cert) and os.path.exists(key)): + subprocess.run( + [ + "openssl", + "req", + "-x509", + "-newkey", + "rsa:2048", + "-nodes", + "-days", + "3650", + "-subj", + "/CN=djev", + "-keyout", + key, + "-out", + cert, + ], + check=True, + capture_output=True, + ) + return cert, key + + +def serve_tls(host, port, cert_dir): + cert, key = self_signed(cert_dir) + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ctx.load_cert_chain(cert, key) + srv = ThreadingHTTPServer((host, port), Handler) + srv.socket = ctx.wrap_socket(srv.socket, server_side=True) + threading.Thread(target=srv.serve_forever, daemon=True).start() + return srv + + +def main(): + global ARGS, CANVAS_LEN, CANVAS_STEP + p = argparse.ArgumentParser() + p.add_argument("--upstream", default="http://127.0.0.1:8010") + p.add_argument("--model", default="dgemma") + p.add_argument("--tokenizer", default="/models/dgemma", help="HF id or local path") + p.add_argument("--canvas", type=int, default=64, help="the served canvas length") + p.add_argument( + "--canvas-step", + type=int, + default=16, + help="request widths round up to a multiple of this", + ) + p.add_argument("--host", default="0.0.0.0") + p.add_argument("--port", type=int, default=8011) + p.add_argument( + "--tls-port", type=int, default=0, help="also listen with HTTPS here (0 = off)" + ) + p.add_argument( + "--cert-dir", + default=os.path.expanduser("~/.cache/djev"), + help="directory for the self-signed certificate", + ) + ARGS = p.parse_args() + CANVAS_LEN = ARGS.canvas + CANVAS_STEP = ARGS.canvas_step + init_tokenizer(AutoTokenizer.from_pretrained(ARGS.tokenizer)) + if ARGS.tls_port: + serve_tls(ARGS.host, ARGS.tls_port, ARGS.cert_dir) + print( + f"structured server https on {ARGS.host}:{ARGS.tls_port} (self-signed)", + flush=True, + ) + print( + f"structured server on {ARGS.host}:{ARGS.port} -> {ARGS.upstream} " + f"(canvas {CANVAS_LEN})", + flush=True, + ) + ThreadingHTTPServer((ARGS.host, ARGS.port), Handler).serve_forever() + + +if __name__ == "__main__": + main() diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index b68820c0d51c..90dbbf278bd3 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -15,8 +15,10 @@ from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig from vllm.config.model_arch import ModelArchitectureConfig +from vllm.transformers_utils.configs.diffusion_gemma import DiffusionGemmaConfig from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from vllm.transformers_utils.model_arch_config_convertor import ( + MODEL_ARCH_CONFIG_CONVERTORS, Gemma4ModelArchConfigConvertor, ModelArchConfigConvertorBase, ) @@ -496,3 +498,26 @@ def test_deepseek_v4_convertor_splits_vision_architecture(): conv = DeepseekV4ModelArchConfigConvertor(draft_cfg, draft_cfg) assert conv.get_architectures() == [draft_arch] assert draft_cfg.architectures == [draft_arch] + + +def test_diffusion_gemma_head_dims_vary_by_layer_type(): + """Checkpoints carry model_type `diffusion_gemma` at the top level. Under + the base convertor every layer reports the sliding head dim, Gemma4Config + skips its head-dim handling, and the full-attention KV-cache group falls to + FlashInfer, which cannot take a per-request causal tensor.""" + text_config = dict( + num_hidden_layers=6, + hidden_size=64, + num_attention_heads=8, + num_key_value_heads=4, + head_dim=16, + global_head_dim=32, + layer_types=["sliding_attention"] * 5 + ["full_attention"], + ) + hf_config = DiffusionGemmaConfig(text_config=text_config) + + convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS[hf_config.model_type] + assert convertor_cls is Gemma4ModelArchConfigConvertor + arch = convertor_cls(hf_config, hf_config.text_config).convert() + + assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32] diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index 3c909f4969df..84ce45caa958 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -1,11 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from types import SimpleNamespace import pytest from vllm import SamplingParams +from vllm.config.diffusion import DiffusionConfig from vllm.exceptions import VLLMValidationError +from vllm.utils.diffusion import validate_diffusion_sampling_params +from vllm.v1.engine.input_processor import InputProcessor @dataclass @@ -13,6 +17,7 @@ class MockModelConfig: is_diffusion: bool = False max_logprobs: int = 20 logits_processors: list | None = None + return_sampling_mask: bool = False def get_vocab_size(self) -> int: return 1024 @@ -59,6 +64,183 @@ def test_verify_leaves_logits_processors_to_admission(): ) +def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None): + model_config = MockModelConfig(is_diffusion=True) + params.verify(model_config, None, None, None) + validate_diffusion_sampling_params( + params, + canvas_length=canvas_length, + vocab_size=model_config.get_vocab_size(), + async_scheduling=True, + ) + + +def test_diffusion_extra_args_are_validated_without_a_served_canvas(): + # No --diffusion-config: the canvas is unknown, the ids are still checked + # and a read-only request is still normalised. + params = SamplingParams( + max_tokens=64, + extra_args={"diffusion_seed_canvas": [0, 1], "diffusion_read_only": True}, + ) + _verify_diffusion(params, canvas_length=None) + assert params.ignore_eos is True + + bad = SamplingParams(extra_args={"diffusion_seed_canvas": [0, 10**9]}) + with pytest.raises(VLLMValidationError, match="ids must be in"): + _verify_diffusion(bad, canvas_length=None) + + +@pytest.mark.parametrize("async_scheduling", [False, True]) +@pytest.mark.parametrize( + "extra_args", + [ + {}, + {"diffusion_canvas_length": None}, + {"diffusion_canvas_length": 4}, + {"diffusion_canvas_length": 8}, + ], +) +def test_narrow_diffusion_canvas_requires_async_scheduling( + async_scheduling, extra_args +): + processor = SimpleNamespace( + model_config=MockModelConfig(is_diffusion=True), + vllm_config=SimpleNamespace( + scheduler_config=SimpleNamespace(async_scheduling=async_scheduling) + ), + speculative_config=None, + structured_outputs_config=None, + diffusion_config=DiffusionConfig(canvas_length=8), + tokenizer=None, + validate_logits_processors_params=lambda params: None, + ) + params = SamplingParams(extra_args=extra_args) + if not async_scheduling and extra_args.get("diffusion_canvas_length") == 4: + with pytest.raises(VLLMValidationError, match="requires --async-scheduling"): + InputProcessor._validate_params(processor, params, ("generate",)) + else: + InputProcessor._validate_params(processor, params, ("generate",)) + + +@pytest.mark.parametrize( + "extra_args, match", + [ + ({"diffusion_seed_canvas": "abc"}, "list of token ids"), + ({"diffusion_seed_canvas": [1, 2.5]}, "list of token ids"), + ({"diffusion_seed_canvas": [1, True]}, "list of token ids"), + ({"diffusion_seed_canvas": [0, 1024]}, r"in \[0, 1024\)"), + ({"diffusion_seed_canvas": [0, -1]}, r"in \[0, 1024\)"), + ({"diffusion_max_steps": 0}, "positive integer"), + ({"diffusion_max_steps": "1"}, "positive integer"), + ({"diffusion_max_steps": True}, "positive integer"), + ({"diffusion_read_only": "yes"}, "boolean"), + ({"diffusion_read_only": 2}, "boolean"), + ({"diffusion_pinned": "0,1"}, "list of canvas positions"), + ({"diffusion_pinned": [0, True]}, "list of canvas positions"), + ({"diffusion_pinned": [0]}, "needs a diffusion_seed_canvas"), + ], +) +def test_diffusion_rejects_bad_extra_args(extra_args: dict, match: str): + with pytest.raises(VLLMValidationError, match=match): + _verify_diffusion(SamplingParams(extra_args=extra_args)) + + +def test_diffusion_seed_canvas_must_fill_the_canvas(): + params = SamplingParams(extra_args={"diffusion_seed_canvas": [0] * 7}) + with pytest.raises(VLLMValidationError, match="exactly 8 ids, got 7"): + _verify_diffusion(params, canvas_length=8) + # Without the served diffusion config the canvas length is unknown. + _verify_diffusion(params) + + +@pytest.mark.parametrize("max_tokens, expected", [(100, 8), (5, 5), (None, 8)]) +@pytest.mark.parametrize("flag", [True, 1]) +def test_diffusion_read_only_ends_after_one_canvas(max_tokens, expected, flag): + params = SamplingParams( + max_tokens=max_tokens, extra_args={"diffusion_read_only": flag} + ) + _verify_diffusion(params, canvas_length=8) + assert params.max_tokens == expected + assert params.ignore_eos + + +@pytest.mark.parametrize( + "width, match", + [ + (0, "positive integer"), + ("4", "positive integer"), + (True, "positive integer"), + (9, "no larger"), + ], +) +def test_diffusion_rejects_bad_canvas_length(width, match): + with pytest.raises(VLLMValidationError, match=match): + _verify_diffusion( + SamplingParams(extra_args={"diffusion_canvas_length": width}), + canvas_length=8, + ) + + +def test_diffusion_canvas_length_sizes_the_seed_and_the_read(): + params = SamplingParams( + max_tokens=100, + extra_args={ + "diffusion_canvas_length": 4, + "diffusion_seed_canvas": [1, 2, 3, 4], + "diffusion_read_only": True, + }, + ) + _verify_diffusion(params, canvas_length=8) + assert params.max_tokens == 4 + + params = SamplingParams( + extra_args={"diffusion_canvas_length": 4, "diffusion_seed_canvas": [0] * 8} + ) + with pytest.raises(VLLMValidationError, match="exactly 4 ids, got 8"): + _verify_diffusion(params, canvas_length=8) + + +def test_diffusion_pinned_positions_stay_inside_the_canvas(): + seed = [0] * 8 + with pytest.raises(VLLMValidationError, match="inside the canvas"): + _verify_diffusion( + SamplingParams( + extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [7, 8]} + ), + canvas_length=8, + ) + with pytest.raises(VLLMValidationError, match="inside the canvas"): + _verify_diffusion( + SamplingParams( + extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [-1]} + ) + ) + # A narrower request canvas bounds the positions. + with pytest.raises(VLLMValidationError, match="inside the canvas"): + _verify_diffusion( + SamplingParams( + extra_args={ + "diffusion_canvas_length": 4, + "diffusion_seed_canvas": seed[:4], + "diffusion_pinned": [4], + } + ), + canvas_length=8, + ) + + +def test_diffusion_accepts_extra_args(): + params = SamplingParams( + extra_args={ + "diffusion_seed_canvas": list(range(8)), + "diffusion_pinned": [0, 1, 7], + "diffusion_max_steps": 4, + "diffusion_read_only": True, + } + ) + _verify_diffusion(params, canvas_length=8) + + @pytest.mark.parametrize("value", [-(2**63) - 1, 2**64]) def test_extra_args_rejects_nested_integer_overflow(value): """Reject extension values before they reach the engine transport.""" diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index e31163923b29..9a5d16eb714b 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -40,6 +40,10 @@ from vllm.v1.core.encoder_cache_manager import EncoderCacheManager from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash +from vllm.v1.core.sched.diffusion_scheduler import ( + DiffusionAsyncScheduler, + diffusion_canvas_width, +) from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler @@ -6872,3 +6876,129 @@ def test_update_draft_token_ids_in_output_strips_padding(): -1, ] assert scheduler_output.num_invalid_spec_tokens == {request.request_id: 2} + + +def test_diffusion_canvas_width_defaults_to_the_served_canvas(): + def req(extra): + return SimpleNamespace(sampling_params=SimpleNamespace(extra_args=extra)) + + assert diffusion_canvas_width(req(None), 64) == 64 + assert diffusion_canvas_width(req({}), 64) == 64 + assert diffusion_canvas_width(req({"diffusion_canvas_length": 16}), 64) == 16 + assert diffusion_canvas_width(SimpleNamespace(sampling_params=None), 64) == 64 + + +def _diffusion_request(req_id: str, extra_args: dict) -> Request: + (request,) = create_requests( + num_requests=1, num_tokens=8, max_tokens=64, req_ids=[req_id] + ) + request.sampling_params.extra_args = extra_args + return request + + +def _diffusion_scheduler(**kwargs) -> DiffusionAsyncScheduler: + scheduler = create_scheduler( + async_scheduling=True, + diffusion_canvas_length=8, + scheduler_cls=DiffusionAsyncScheduler, + **kwargs, + ) + assert isinstance(scheduler, DiffusionAsyncScheduler) + return scheduler + + +@pytest.mark.parametrize("async_scheduling", [False, True]) +def test_diffusion_scheduler_is_selected_by_default(async_scheduling): + config = create_scheduler( + async_scheduling=async_scheduling, diffusion_canvas_length=8 + ).vllm_config.scheduler_config + assert config.get_scheduler_cls() is ( + DiffusionAsyncScheduler if config.async_scheduling else Scheduler + ) + + +def test_diffusion_scheduler_narrows_the_canvas_per_request(): + scheduler = _diffusion_scheduler() + wide = _diffusion_request("wide", {}) + narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4}) + scheduler.add_request(wide) + scheduler.add_request(narrow) + + # The prefill step lays down the served canvas as placeholders. + scheduler.schedule() + output = scheduler.schedule() + assert output.scheduled_spec_decode_tokens["wide"] == [-1] * 8 + assert output.scheduled_spec_decode_tokens["narrow"] == [-1] * 4 + assert output.num_scheduled_tokens["narrow"] == 4 + assert narrow.num_output_placeholders == 4 + + +@pytest.mark.parametrize("structured", [False, True]) +def test_diffusion_scheduler_trims_full_width_worker_drafts(structured): + """Padded worker drafts must be narrowed before scheduling or grammar validation.""" + scheduler = _diffusion_scheduler() + wide = _diffusion_request("wide", {}) + narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4}) + for request in (wide, narrow): + scheduler.add_request(request) + prefill = scheduler.schedule() + _model_output(scheduler, prefill, [[], []]) + + if structured: + for request in (wide, narrow): + request.structured_output_request = SimpleNamespace( + grammar=_RecordingGrammar(), reasoning_ended=True + ) + tokens = list(range(8)) if structured else [-1] * 8 + drafts = DraftTokenIds(["wide", "narrow"], [tokens.copy(), tokens.copy()]) + output = scheduler.schedule() + scheduler.update_draft_token_ids_in_output(drafts, output) + + assert output.scheduled_spec_decode_tokens == { + "wide": tokens, + "narrow": tokens[:4], + } + assert output.num_scheduled_tokens == {"wide": 8, "narrow": 4} + if structured: + assert wide.structured_output_request.grammar.seen == [tokens] + assert narrow.structured_output_request.grammar.seen == [tokens[:4]] + + +def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight(): + scheduler = _diffusion_scheduler() + one = _diffusion_request( + "one", {"diffusion_read_only": True, "diffusion_max_steps": 1} + ) + two = _diffusion_request( + "two", {"diffusion_read_only": True, "diffusion_max_steps": 2} + ) + # Only reads are deferred. A capped generation keeps its extra async step. + gen = _diffusion_request("gen", {"diffusion_max_steps": 1}) + for request in (one, two, gen): + scheduler.add_request(request) + + scheduler.schedule() + assert set(scheduler.schedule().num_scheduled_tokens) == {"one", "two", "gen"} + assert one.num_output_placeholders == 8 + assert set(scheduler.schedule().num_scheduled_tokens) == {"two", "gen"} + assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"} + # The deferral is renewed every call. + assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"} + + +def test_diffusion_read_deferral_keeps_a_longer_pp_wait(): + scheduler = _diffusion_scheduler(pipeline_parallel_size=3, use_v2_model_runner=True) + read = _diffusion_request( + "read", {"diffusion_read_only": True, "diffusion_max_steps": 1} + ) + scheduler.add_request(read) + + scheduler.schedule() + assert read.next_decode_eligible_step == 4 + for _ in range(2): + assert "read" not in scheduler.schedule().num_scheduled_tokens + assert "read" in scheduler.schedule().num_scheduled_tokens + assert read.next_decode_eligible_step == 7 + # Deferring this step alone would ask for 6. The PP wait to 7 stands. + assert "read" not in scheduler.schedule().num_scheduled_tokens + assert read.next_decode_eligible_step == 7 diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index f69b10ced346..ebda4f092cdb 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -7,6 +7,7 @@ from tests.v1.kv_connector.unit.utils import MockKVConfig from vllm.config import ( CacheConfig, + DiffusionConfig, ECTransferConfig, KVTransferConfig, ModelConfig, @@ -77,6 +78,8 @@ def create_scheduler( kv_cache_spec: KVCacheSpec | None = None, per_request_spec_decode_metrics: str = "none", scheduling_policy: SchedulerPolicy = "fcfs", + diffusion_canvas_length: int | None = None, + scheduler_cls: type[Scheduler] | None = None, ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -178,6 +181,12 @@ def create_scheduler( else None ) + diffusion_config: DiffusionConfig | None = None + if diffusion_canvas_length is not None: + # A diffusion checkpoint declares its canvas in the HF config. + model_config.hf_config.canvas_length = diffusion_canvas_length + diffusion_config = DiffusionConfig(canvas_length=diffusion_canvas_length) + vllm_config = VllmConfig( scheduler_config=scheduler_config, model_config=model_config, @@ -188,6 +197,7 @@ def create_scheduler( ), kv_transfer_config=kv_transfer_config, speculative_config=speculative_config, + diffusion_config=diffusion_config, ec_transfer_config=ec_transfer_config, observability_config=ObservabilityConfig( per_request_spec_decode_metrics=per_request_spec_decode_metrics, @@ -207,7 +217,8 @@ def create_scheduler( ) cache_config.num_gpu_blocks = num_blocks register_all_kvcache_specs(vllm_config) - scheduler_cls = AsyncScheduler if async_scheduling else Scheduler + if scheduler_cls is None: + scheduler_cls = AsyncScheduler if async_scheduling else Scheduler scheduler = scheduler_cls( vllm_config=vllm_config, kv_cache_config=kv_cache_config, diff --git a/tests/v1/engine/test_input_processor_trace_replay.py b/tests/v1/engine/test_input_processor_trace_replay.py index 24302f6de616..957f70b71858 100644 --- a/tests/v1/engine/test_input_processor_trace_replay.py +++ b/tests/v1/engine/test_input_processor_trace_replay.py @@ -80,10 +80,12 @@ def _validate(enable_trace_replay: bool) -> None: model_config=SimpleNamespace( return_sampling_mask=False, enable_trace_replay=enable_trace_replay, + is_diffusion=False, ), vllm_config=SimpleNamespace(reasoning_config=None), speculative_config=None, structured_outputs_config=None, + diffusion_config=None, tokenizer=None, validate_logits_processors_params=lambda params: None, ) diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py new file mode 100644 index 000000000000..369af010f4c7 --- /dev/null +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -0,0 +1,319 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-request DiffusionGemma state behind structured reads: seed canvases, +pinned positions, read-only slots and the per-slot step cap.""" + +import numpy as np +import pytest +import torch + +from vllm.model_executor.models.diffusion_gemma import ( + DiffusionGemmaRequestStates, + _compiled_sample_step, + _concat_logprob_stashes, +) +from vllm.platforms import current_platform +from vllm.v1.outputs import LogprobsTensors + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="the sampler state lives on the GPU" +) + +CL = 8 +VOCAB = 64 +MAX_REQS = 4 +MAX_STEPS = 48 + + +def _states() -> DiffusionGemmaRequestStates: + return DiffusionGemmaRequestStates( + max_num_reqs=MAX_REQS, + canvas_length=CL, + vocab_size=VOCAB, + max_denoising_steps=MAX_STEPS, + device=torch.device("cuda"), + hidden_size=4, + stability_threshold=2, + ) + + +def _slots(*idx: int) -> tuple[np.ndarray, torch.Tensor]: + slots = np.array(idx, dtype=np.int64) + return slots, torch.tensor(slots, device="cuda") + + +def test_seed_canvas_replaces_only_seeded_slots(): + states = _states() + for slot in range(3): + states.add_request(slot) + seed = list(range(CL)) + states.set_seed_canvas(1, seed) + slots, slots_gpu = _slots(0, 1, 2) + states.init_canvas(slots_gpu) + before = states.canvas[slots_gpu].clone() + + states.apply_seed_canvases(slots, slots_gpu) + + after = states.canvas[slots_gpu] + assert after[1].tolist() == seed + assert torch.equal(after[0], before[0]) + assert torch.equal(after[2], before[2]) + + +def test_apply_seed_canvases_leaves_unseeded_batches_alone(): + states = _states() + states.add_request(0) + slots, slots_gpu = _slots(0) + states.init_canvas(slots_gpu) + before = states.canvas[0].clone() + + states.apply_seed_canvases(slots, slots_gpu) + + assert torch.equal(states.canvas[0], before) + + +def test_add_request_clears_seed_and_read_only(): + states = _states() + states.add_request(0) + states.set_seed_canvas(0, [1] * CL) + states.set_read_only(0) + assert states.seeded_slots == {0} + assert states.read_only_slots == {0} + + states.add_request(0) + + assert not states.seeded_slots + assert not states.read_only_slots + assert not bool(states.has_seed[0]) + assert not bool(states.read_only[0]) + + +def test_canvas_width_resets_with_the_slot(): + states = _states() + states.add_request(0) + states.canvas_width_np[0] = 4 + states.set_seed_canvas(0, [7, 7, 7, 7]) + assert states.seed_canvas[0, :4].tolist() == [7, 7, 7, 7] + + states.add_request(0) + + assert states.canvas_width_np[0] == CL + + +def test_remove_request_forgets_the_slot(): + states = _states() + states.add_request(0) + states.set_seed_canvas(0, [1] * CL) + states.set_read_only(0) + + states.remove_request(0) + + assert not states.seeded_slots + assert not states.read_only_slots + + +def _denoise_once( + states: DiffusionGemmaRequestStates, + slots: list[int], + compute_sc: bool = True, + width: int = CL, + embed_weight: torch.Tensor | None = None, + embed_dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """One compiled denoise step over ``slots`` with flat logits, so nothing + converges by stability or confidence and only the step cap can end it. + ``width`` below CL runs the step on [:, :width] views, as the sampler + does for a narrow tile.""" + n = len(slots) + device = states.device + decode_slots = torch.tensor(slots, dtype=torch.int64, device=device) + decode_idx = torch.arange(n, dtype=torch.int64, device=device) + sampled = torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width] + num_sampled = torch.zeros(n, dtype=torch.int32, device=device) + _compiled_sample_step( + torch.zeros(n * width, VOCAB, device=device), + decode_slots, + decode_idx, + decode_slots, + torch.full((n,), width, dtype=torch.int64, device=device), + states.canvas[:, :width], + states.argmax_canvas[:, :width], + states.step, + states.is_encoder_phase, + states.confident, + states.self_conditioning_embeds[:, :width], + ( + torch.zeros(VOCAB, 4, dtype=embed_dtype, device=device) + if embed_weight is None + else embed_weight + ), + torch.tensor(1.0, dtype=embed_dtype, device=device), + states.accepted_canvas_history[:, :, :width], + states.accepted_canvas_history_len, + states.max_steps, + states.pin_mask[:, :width], + states.seed_canvas[:, :width], + states.read_only, + sampled, + num_sampled, + torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device), + max_denoising_steps=float(MAX_STEPS), + t_min=0.5, + t_max=1.0, + confidence_threshold=0.1, + vocab_size=VOCAB, + CL=width, + ST=states.stability_threshold, + entropy_bound=0.1, + sc_vocab_start=0, + sc_vocab_end=VOCAB, + tp_size=1, + tp_group_name="", + compute_sc=compute_sc, + ) + return sampled, num_sampled + + +def test_pinned_positions_hold_their_seed_through_a_step(): + states = _states() + states.add_request(0) + states.is_encoder_phase[0] = False + seed = list(range(10, 10 + CL)) + states.set_seed_canvas(0, seed) + states.canvas[0] = torch.tensor(seed, device="cuda") + states.set_pins(0, [0, 1, 2, 3]) + + # Flat logits accept nothing, so every free position is renoised. + _denoise_once(states, [0], embed_weight=torch.ones(VOCAB, 4, device="cuda")) + + assert states.canvas[0, :4].tolist() == seed[:4] + # The soft embed is zero at pinned positions and non-zero elsewhere. + sc = states.self_conditioning_embeds[0] + assert not sc[:4].any() + assert sc[4:].any() + + +def test_add_request_clears_pins(): + states = _states() + states.add_request(0) + states.set_seed_canvas(0, [1] * CL) + states.set_pins(0, [2, 3]) + assert states.pin_mask[0].tolist() == [False, False, True, True] + [False] * 4 + + states.add_request(0) + + assert not states.pin_mask[0].any() + + +def test_single_step_tile_skips_self_conditioning(): + states = _states() + states.add_request(0) + states.is_encoder_phase[0] = False + states.self_conditioning_embeds[0] = 1.0 + + _denoise_once(states, [0], compute_sc=False) + + assert not states.self_conditioning_embeds[0].any() + + +@pytest.mark.parametrize("stance", ["default", "force_eager"]) +def test_self_conditioning_stores_a_bf16_model_in_the_fp32_buffer(stance): + # The model's embeddings are bf16 while the buffer is fp32. Compiled code + # casts on the store, but eager does not, and the step runs eager once + # torch.compile hits its recompile limit. + states = _states() + states.add_request(0) + states.is_encoder_phase[0] = False + + with torch.compiler.set_stance(stance): + _denoise_once( + states, + [0], + embed_weight=torch.ones(VOCAB, 4, dtype=torch.bfloat16, device="cuda"), + embed_dtype=torch.bfloat16, + ) + + assert states.self_conditioning_embeds.dtype == torch.float32 + # uniform probs @ an all-ones embedding + assert torch.allclose( + states.self_conditioning_embeds[0], torch.ones(CL, 4, device="cuda") + ) + + +def test_narrow_tile_leaves_the_rest_of_the_canvas_alone(): + states = _states() + states.add_request(0) + states.is_encoder_phase[0] = False + states.canvas[0] = 5 + states.argmax_canvas[0] = 5 + + _denoise_once(states, [0], width=4) + + # Columns past the width are untouched; the step counted. + assert states.canvas[0, 4:].tolist() == [5] * (CL - 4) + assert states.argmax_canvas[0, 4:].tolist() == [5] * (CL - 4) + assert states.step[0] == 1 + + +def test_step_cap_is_per_slot(): + states = _states() + for slot in (0, 1): + states.add_request(slot) + states.is_encoder_phase[slot] = False + states.max_steps[0] = 1 + + _denoise_once(states, [0, 1]) + + # Slot 0 hit its cap and moves to commit. Slot 1 keeps denoising. + assert states.is_encoder_phase[:2].tolist() == [True, False] + assert states.step[:2].tolist() == [1, 1] + + +def _stash(rows: int, width: int, first_id: int) -> LogprobsTensors: + ids = torch.arange(first_id, first_id + rows * width).reshape(rows, width) + return LogprobsTensors( + logprob_token_ids=ids, + logprobs=-ids.float(), + selected_token_ranks=torch.zeros(rows, dtype=torch.int64), + ) + + +def test_stashes_of_different_widths_join(): + # A read that asked for 10 label ids and a generation that asked for 2 + # logprobs converged in different steps, so their stashes differ in width. + wide, narrow = _stash(2, 11, 100), _stash(3, 3, 0) + + out = _concat_logprob_stashes([narrow, wide], [0, 3]) + + assert out.logprob_token_ids.shape == (5, 11) + assert out.logprobs.shape == (5, 11) + assert out.cu_num_generated_tokens == [0, 3] + # The narrow rows keep their values and pad with id 0 at -inf. + assert torch.equal(out.logprob_token_ids[:3, :3], narrow.logprob_token_ids) + assert not out.logprob_token_ids[:3, 3:].any() + assert torch.isneginf(out.logprobs[:3, 3:]).all() + assert torch.equal(out.logprobs[3:], wide.logprobs) + + +@pytest.mark.parametrize("width", [4, CL]) +@pytest.mark.parametrize("steps", [1, 2]) +def test_read_emits_at_convergence_while_generation_waits_for_commit(width, steps): + states = _states() + for slot in (2, 0): + states.add_request(slot) + states.is_encoder_phase[slot] = False + states.max_steps[slot] = steps + states.set_read_only(2) + + for step in range(steps): + sampled, counts = _denoise_once(states, [2, 0], width=width) + assert counts.tolist() == ([width, 0] if step == steps - 1 else [0, 0]) + + assert torch.equal(sampled[0], states.argmax_canvas[2, :width].int()) + assert not states.is_encoder_phase[2] + assert states.is_encoder_phase[0] + assert not states.self_conditioning_embeds[2].any() + + sampled, counts = _denoise_once(states, [0], width=width) + assert counts.tolist() == [width] + assert torch.equal(sampled[0], states.argmax_canvas[0, :width].int()) diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 9185c26c4d85..86f863ab3a23 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -232,13 +232,17 @@ def get_scheduler_cls(self) -> type["SchedulerInterface"]: # The first half of this warning can be removed once the Scheduler interface is # finalized and we can maintain support for scheduler classes that implement it - logger.warning_once( - "Using custom scheduler class %s. This scheduler interface is not public " - "and compatibility may not be maintained. If you have subclassed Scheduler " - "instead of AsyncScheduler, you will see degraded performance due to async " - "scheduling being disabled.", - self.scheduler_cls, # type: ignore[arg-type] - ) + if not ( + isinstance(self.scheduler_cls, str) + and self.scheduler_cls.startswith("vllm.") + ): + logger.warning_once( + "Using custom scheduler class %s. This scheduler interface is not " + "public and compatibility may not be maintained. If you have " + "subclassed Scheduler instead of AsyncScheduler, you will see " + "degraded performance due to async scheduling being disabled.", + self.scheduler_cls, # type: ignore[arg-type] + ) if not isinstance(self.scheduler_cls, str): return cast(type["SchedulerInterface"], self.scheduler_cls) return resolve_obj_by_qualname(self.scheduler_cls) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index bf11255e8e0a..65f2f8df523d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1878,6 +1878,16 @@ def has_blocked_weights(): ) current_platform.check_and_update_config(self) + # After the platform hook, which has the last word on async scheduling. + if ( + self.diffusion_config is not None + and self.scheduler_config.async_scheduling + and self.scheduler_config.scheduler_cls is None + ): + self.scheduler_config.scheduler_cls = ( + "vllm.v1.core.sched.diffusion_scheduler.DiffusionAsyncScheduler" + ) + self._normalize_piecewise_cudagraph_mode( breakable_cudagraph_enabled=breakable_cudagraph_enabled ) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 8288eab236f3..3a0392c52322 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -280,11 +280,14 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: "FLASH_ATTN or TRITON_ATTN instead." ) if attention_config.backend is None and not attention_config.use_non_causal: + # Only the sm_100 default backend order reads this flag (it puts + # FlashInfer first for causal models). Elsewhere Gemma4Config above + # keeps FlashInfer out: FA4 on every layer, else TRITON_ATTN. attention_config.use_non_causal = True logger.info( - "DiffusionGemma uses mixed causal/bidirectional attention " - "within a batch; setting use_non_causal=True to exclude " - "FlashInfer from auto-selection." + "DiffusionGemma mixes causal and bidirectional attention in " + "one batch; setting use_non_causal=True so the default backend " + "order prefers FlashAttention over FlashInfer." ) # Auto-create DiffusionConfig from HF config if not provided. diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index aef9d816258b..a1b4eb3a0d73 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -22,6 +22,7 @@ import numpy as np import torch +import torch._dynamo from torch import nn from torch.nn import functional as F from transformers import AutoModel @@ -371,6 +372,34 @@ def _compute_num_rejected( return torch.where(is_denoise, query_lens, num_rejected) +def _concat_logprob_stashes( + parts: list[LogprobsTensors], cu_num_generated_tokens: list[int] +) -> LogprobsTensors: + """Join the logprobs stashed for the requests committing this step. + + Each stash is as wide as the widest logprobs request in the batch at the + step that request converged, so stashes from different steps can differ + in width. Pad the narrow ones the way compute_topk_scores pads a mixed + batch: token id 0 at -inf, which the output processor never reports. + """ + width = max(p.logprob_token_ids.shape[1] for p in parts) + + def pad(t: torch.Tensor, value: float) -> torch.Tensor: + return F.pad(t, (0, width - t.shape[1]), value=value) + + return LogprobsTensors( + logprob_token_ids=torch.cat([pad(p.logprob_token_ids, 0) for p in parts]), + logprobs=torch.cat([pad(p.logprobs, float("-inf")) for p in parts]), + selected_token_ranks=torch.cat([p.selected_token_ranks for p in parts]), + cu_num_generated_tokens=cu_num_generated_tokens, + ) + + +# Tiles specialize the graph on batch size 1, on a tile as wide as the canvas, +# on compute_sc and on sizes that coincide with the state buffers. That set +# is small but passes Dynamo's default of 8, after which every step would run +# eager: about twice as slow for a self-conditioned step. +@torch._dynamo.config.patch(recompile_limit=64) @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _compiled_sample_step( # Logits from the model [num_decode * CL, vocab] @@ -391,6 +420,10 @@ def _compiled_sample_step( normalizer: torch.Tensor, history: torch.Tensor, # [max_num_reqs, ST, CL] history_len_tensor: torch.Tensor, # [max_num_reqs] + max_steps_tensor: torch.Tensor, # [max_num_reqs] int32, per-slot step cap + pin_mask: torch.Tensor, # [max_num_reqs, CL] bool, positions held at the seed + seed_canvas: torch.Tensor, # [max_num_reqs, CL] + read_only: torch.Tensor, # [max_num_reqs] bool, emit without a commit forward # Output tensors (modified in-place) sampled: torch.Tensor, # [num_reqs, CL] num_sampled: torch.Tensor, # [num_reqs] @@ -412,6 +445,7 @@ def _compiled_sample_step( sc_vocab_end: int, tp_size: int, tp_group_name: str, + compute_sc: bool = True, ) -> torch.Tensor: """Compiled decode step: temperature → Gumbel sample → probs/confidence → accept/renoise → convergence, all as vectorized PyTorch ops. @@ -478,8 +512,11 @@ def _compiled_sample_step( 0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype ) - # Compute denoise canvas (accept/renoise) + # Compute denoise canvas (accept/renoise). Pinned positions keep the seed. denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens) + denoise_canvas = torch.where( + pin_mask[decode_slots], seed_canvas[decode_slots], denoise_canvas + ) # Canvas: commit → random reinit, denoise → accept/renoise result canvas[decode_slots] = torch.where( @@ -504,16 +541,6 @@ def _compiled_sample_step( new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode)) history_len_tensor[decode_slots] = new_hist_len - # Sampled output: commit → emit argmax_canvas, denoise → 0 (pre-zeroed) - sampled[decode_idx] = argmax_canvas[decode_slots].to( - sampled.dtype - ) * is_commit.unsqueeze(1).to(sampled.dtype) - # Commit only the real canvas length (== CL except for a canvas truncated - # near max_model_len); the padded tail positions are never emitted. - num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * valid_canvas_len.to( - num_sampled.dtype - ) - # ---- Phase 6: Stability + convergence ---- ref = history[decode_slots, 0] mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32) @@ -523,13 +550,17 @@ def _compiled_sample_step( step_after = step_tensor[decode_slots] converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | ( - step_after >= max_denoising_steps + step_after >= max_steps_tensor[decode_slots] ) # Commit done → denoise next (False); denoise converged → commit next (True) is_encoder_phase[decode_slots] = torch.where( is_commit, is_commit.new_zeros(num_decode), converged ) + emit = is_commit | (is_denoise & converged & read_only[decode_slots]) + sampled[decode_idx] = argmax_canvas[decode_slots].to(sampled.dtype) * emit[:, None] + num_sampled[decode_idx] = (emit * valid_canvas_len).to(num_sampled.dtype) + # SC soft embedding: store ``probs @ embed_weight`` (the value the next step's # self-conditioning MLP consumes) only for slots that will denoise next — i.e. # this step denoised AND it isn't about to commit (is_encoder_phase now False). @@ -537,24 +568,41 @@ def _compiled_sample_step( # sc_embeds directly. Storing the [.., hidden] soft embed instead of the full # [.., vocab] probs avoids a giant persistent buffer. sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None] - # Self-conditioning soft embed = probs @ embed_tokens.weight. Under tensor - # parallelism the embedding is vocab-sharded ([vocab/tp, hidden]) while - # probs spans the full vocab, so each rank multiplies its local vocab slice - # [sc_vocab_start, sc_vocab_end) and the partials are summed across ranks. - local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype) - soft_embeds = torch.matmul( - local_probs, embed_weight[: sc_vocab_end - sc_vocab_start] - ) - if tp_size > 1: - soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name) - soft_embeds = soft_embeds * normalizer - sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype) + if compute_sc: + # Self-conditioning soft embed = probs @ embed_tokens.weight. Under + # tensor parallelism the embedding is vocab-sharded ([vocab/tp, + # hidden]) while probs spans the full vocab, so each rank multiplies + # its local vocab slice [sc_vocab_start, sc_vocab_end) and the + # partials are summed across ranks. + local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype) + soft_embeds = torch.matmul( + local_probs, embed_weight[: sc_vocab_end - sc_vocab_start] + ) + if tp_size > 1: + soft_embeds = torch.ops.vllm.all_reduce( + soft_embeds, group_name=tp_group_name + ) + soft_embeds = soft_embeds * normalizer + # A pinned position holds its seed token as input. Zero its soft embed + # too, or the model's own prediction there reaches the next step + # through self-conditioning. The buffer is fp32 while the embedding is + # in the model dtype: compiled code casts on the store but eager does + # not, and the step runs eager once torch.compile hits its recompile + # limit. + sc_pin = (~pin_mask[decode_slots]).unsqueeze(-1) + sc_embeds[decode_slots] = (soft_embeds * sc_keep * sc_pin).to(sc_embeds.dtype) + else: + # Every slot in this tile ends after this step, so the soft embed + # would never be read. The matmul is a full pass over the vocabulary + # matrix, so skip it. + sc_embeds[decode_slots] = 0 # Overwrite canvas with argmax for newly converged denoise requests newly_converged = (converged & is_denoise).unsqueeze(1) canvas[decode_slots] = torch.where( newly_converged, argmax_canvas[decode_slots], canvas[decode_slots] ) + is_encoder_phase[decode_slots] &= ~read_only[decode_slots] # ---- Phase 7: Copy canvas → draft_tokens for all slots ---- draft_tokens[all_slots, :CL] = canvas[all_slots] @@ -628,6 +676,32 @@ def __init__( # Per-slot confidence flag, set by the sampler each step. self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + # Per-slot step cap, lowered per request from extra_args. + self.max_steps = torch.full( + (max_num_reqs,), max_denoising_steps, dtype=torch.int32, device=device + ) + # Seed canvases replace the random canvas after prefill. The host-side + # sets gate the seed and read-only work so plain generation skips it. + self.seed_canvas = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.int64, device=device + ) + self.has_seed = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + self.seeded_slots: set[int] = set() + # Pinned positions keep their seed value on every denoise step, so a + # multi-step read denoises the canvas the request seeded. + self.pin_mask = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.bool, device=device + ) + # Read-only slots emit on their converging step and skip the commit + # forward. + self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + self.read_only_slots: set[int] = set() + # Slots capped at one denoise step never consume a soft embed. + self.single_step_slots: set[int] = set() + # Per-slot canvas width, at most canvas_length. The scheduler schedules + # this many draft tokens for the slot and the sampler pads the rest. + self.canvas_width_np = np.full(max_num_reqs, canvas_length, dtype=np.int32) + # Per-slot self-conditioning soft embedding (probs @ embed_weight) from # the previous denoise step. Storing the [.., hidden] soft embed instead # of the full [.., vocab] distribution shrinks this buffer by @@ -657,11 +731,57 @@ def add_request(self, slot_idx: int) -> None: self.step[slot_idx].fill_(0) self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 + self.max_steps[slot_idx].fill_(self.max_denoising_steps) + self.has_seed[slot_idx].fill_(False) + self.seeded_slots.discard(slot_idx) + self.pin_mask[slot_idx].fill_(False) + self.read_only[slot_idx].fill_(False) + self.read_only_slots.discard(slot_idx) + self.single_step_slots.discard(slot_idx) + self.canvas_width_np[slot_idx] = self.canvas_length def remove_request(self, slot_idx: int) -> None: + # add_request resets the GPU flags before a slot is reused. The host + # sets gate work on every step, so they clear here. self.is_encoder_phase[slot_idx].fill_(False) self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 + self.seeded_slots.discard(slot_idx) + self.read_only_slots.discard(slot_idx) + self.single_step_slots.discard(slot_idx) + + def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None: + """``ids`` covers the slot's canvas width; positions past it are never + scheduled.""" + self.seed_canvas[slot_idx, : len(ids)] = async_tensor_h2d( + ids, dtype=torch.int64, device=self.device + ) + self.has_seed[slot_idx].fill_(True) + self.seeded_slots.add(slot_idx) + + def set_pins(self, slot_idx: int, positions: list[int]) -> None: + """Hold ``positions`` of the slot's seed canvas through every denoise + step. Validated against the request's canvas width upstream.""" + self.pin_mask[slot_idx].fill_(False) + self.pin_mask[ + slot_idx, async_tensor_h2d(positions, dtype=torch.int64, device=self.device) + ] = True + + def set_read_only(self, slot_idx: int) -> None: + self.read_only[slot_idx].fill_(True) + self.read_only_slots.add(slot_idx) + + def apply_seed_canvases( + self, slots_np: np.ndarray, slots_gpu: torch.Tensor + ) -> None: + """Replace the canvas of every seeded slot among ``slots_gpu``.""" + if self.seeded_slots.isdisjoint(slots_np.tolist()): + return + self.canvas[slots_gpu] = torch.where( + self.has_seed[slots_gpu, None], + self.seed_canvas[slots_gpu], + self.canvas[slots_gpu], + ) class DiffusionGemmaModelState(ModelState): @@ -1030,6 +1150,33 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None: self._pending_logprobs.pop(req_idx, None) self.sampling_states.add_request(req_idx, sampling_params) self.logprob_token_ids_state.add_request(req_idx, sampling_params) + extra = getattr(sampling_params, "extra_args", None) or {} + states = self.diffusion_states + cap = extra.get("diffusion_max_steps") + if cap is not None: + cap = max(1, min(int(cap), states.max_denoising_steps)) + states.max_steps[req_idx].fill_(cap) + if cap == 1: + states.single_step_slots.add(req_idx) + width = extra.get("diffusion_canvas_length") + if width: + states.canvas_width_np[req_idx] = max( + 1, min(int(width), self.canvas_length) + ) + width = int(states.canvas_width_np[req_idx]) + seed = extra.get("diffusion_seed_canvas") + if seed is not None: + if len(seed) != width: + raise ValueError( + f"diffusion_seed_canvas must hold exactly {width} ids, " + f"got {len(seed)}" + ) + states.set_seed_canvas(req_idx, seed) + pins = extra.get("diffusion_pinned") + if pins and seed is not None: + states.set_pins(req_idx, [int(p) for p in pins]) + if extra.get("diffusion_read_only"): + states.set_read_only(req_idx) def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() @@ -1071,6 +1218,7 @@ def _finish_prefills( ps.astype(np.int64), device=states.is_encoder_phase.device ) states.init_canvas(ps_gpu) + states.apply_seed_canvases(ps, ps_gpu) self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[ ps_gpu ] @@ -1151,7 +1299,6 @@ def __call__( # --- CPU/NumPy setup (outside compile): split decode vs prefill, init # canvas for any new prefills, and stage decode slot indices to GPU. --- states = self.diffusion_states - CL = self.canvas_length slots_np = input_batch.idx_mapping_np[:num_reqs] per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) @@ -1194,18 +1341,12 @@ def __call__( if top_k is not None or top_p is not None: logits = apply_top_k_top_p(logits.float(), top_k, top_p) - # Pad any truncated canvas back to CL so the uniform-CL sampler math - # holds. Phantom (padded) positions are zeroed → uniform logits → high - # entropy (no premature convergence) and argmax 0 (stable); they are - # never committed (num_sampled == real length). masked_fill (not - # multiply) so -inf entries from top_k/top_p filtering above don't - # turn phantom rows into NaN. - if num_decode > 0 and valid_canvas_len_np.min() < CL: - ar = torch.arange(CL, device=device) - starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req - valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL] - src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1) - logits = logits[src.reshape(-1)].masked_fill_(~valid.reshape(-1, 1), 0) + # Where each decode request's rows start in the flat logits. Tiles + # below gather and pad a request to the tile's width. + row_starts_np = np.concatenate( + ([0], np.cumsum(valid_canvas_len_np)[:-1]) + ).astype(np.int64) + row_starts = async_tensor_h2d(row_starts_np, device=device) # Clear once: the tiled loop below only scatters its own decode slots, # so it must not re-clear earlier tiles' writes. @@ -1228,125 +1369,159 @@ def __call__( want_logprobs = max_num_logprobs >= 0 or max_token_ids > 0 num_logprobs = max(max_num_logprobs, 0) - # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in - # _compiled_sample_step keeps several live [group * CL, vocab] copies, so - # size each tile to a fraction of free memory to bound the transient at - # high concurrency. Tiling is bit-identical to a single pass. - group = max(num_decode, 1) - if num_decode > 0: - free, _ = current_platform.mem_get_info() - # ~10 transient fp32 copies of [group * CL, vocab] inside the step + # Sample per tile. Decode requests are grouped by canvas width and each + # tile runs the compiled step at that width over [:, :W] views of the + # state, so a narrow read pays for its own rows rather than the served + # canvas. Widths ascend, so the last tile's canvas-to-draft copy (over + # all slots) is the widest. The fp32 pipeline keeps several live + # [tile * W, vocab] copies, so a tile is also bounded by free memory. + widths_np = states.canvas_width_np[decode_slots_np] + order = np.argsort(widths_np, kind="stable") + free = current_platform.mem_get_info()[0] if num_decode > 0 else 0 + run_start = 0 + while run_start < num_decode: + W = int(widths_np[order[run_start]]) + run_end = run_start + while run_end < num_decode and widths_np[order[run_end]] == W: + run_end += 1 + # ~10 transient fp32 copies of [tile * W, vocab] inside the step # (eager peaks at ~8; pad for allocator overhead and small tensors). - bytes_per_req = CL * self.vocab_size * 4 * 10 - budget = int(free * 0.5) // max(bytes_per_req, 1) - group = max(1, min(num_decode, budget)) - - for start_req in range(0, num_decode, group): - end_req = min(start_req + group, num_decode) - tile = slice(start_req, end_req) - tile_slots = decode_slots[tile] - - scaled = _compiled_sample_step( - logits[start_req * CL : end_req * CL], - tile_slots, - decode_idx[tile], - all_slots, - valid_canvas_len[tile], - # State - states.canvas, - states.argmax_canvas, - states.step, - states.is_encoder_phase, - states.confident, - states.self_conditioning_embeds, - self.embed_weight, - self.normalizer, - states.accepted_canvas_history, - states.accepted_canvas_history_len, - # Output - sampled, - num_sampled, - self.req_states.draft_tokens, - # Config - max_denoising_steps=float(states.max_denoising_steps), - t_min=self.t_min, - t_max=self.t_max, - confidence_threshold=self.confidence_threshold, - vocab_size=self.vocab_size, - CL=CL, - ST=states.stability_threshold, - entropy_bound=self.entropy_bound, - sc_vocab_start=self.sc_vocab_start, - sc_vocab_end=self.sc_vocab_end, - tp_size=self.tp_size, - tp_group_name=self.tp_group_name, - ) + budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 10, 1)) + for t0 in range(run_start, run_end, budget): + sel_np = order[t0 : min(t0 + budget, run_end)] + n = len(sel_np) + sel = async_tensor_h2d(sel_np.astype(np.int64), device=device) + tile_slots = decode_slots[sel] + tile_valid = valid_canvas_len[sel] + tile_valid_np = valid_canvas_len_np[sel_np] + contiguous = bool((sel_np == np.arange(sel_np[0], sel_np[0] + n)).all()) + if contiguous and tile_valid_np.min() == W: + r0 = int(row_starts_np[sel_np[0]]) + tile_logits = logits[r0 : r0 + n * W] + else: + # Pad each request to W. Phantom positions are zeroed: + # uniform logits, high entropy, argmax 0, never committed. + # masked_fill, not multiply, so -inf from top_k/top_p + # filtering above cannot turn a phantom row into NaN. + ar = torch.arange(W, device=device) + src = (row_starts[sel].unsqueeze(1) + ar.unsqueeze(0)).clamp_max( + logits.shape[0] - 1 + ) + valid = ar.unsqueeze(0) < tile_valid.unsqueeze(1) + tile_logits = logits[src.reshape(-1)].masked_fill_( + ~valid.reshape(-1, 1), 0 + ) + compute_sc = ( + not states.single_step_slots + or not states.single_step_slots.issuperset( + decode_slots_np[sel_np].tolist() + ) + ) - # Logprobs for denoise steps that just converged (is_encoder_phase - # flipped False→True), stashed per tile so `scaled` is freed each tile. - if want_logprobs: - converged_mask = states.is_encoder_phase[tile_slots] - just_converged = converged_mask & ~is_committing[tile] - if just_converged.any(): - flat_logits = scaled.reshape(-1, scaled.shape[-1]) - argmax_tokens = scaled.argmax(dim=-1) - for local_idx in just_converged.nonzero(as_tuple=True)[0]: - li = local_idx.item() - slot = tile_slots[local_idx] - # Stash only the real canvas positions (== CL unless this - # canvas was truncated near max_model_len); padded tail - # positions are never emitted. - k_i = int(valid_canvas_len_np[start_req + li]) - pos = li * CL - per_req_ids = max_token_ids > 0 - self._pending_logprobs[slot.item()] = compute_topk_scores( - flat_logits[pos : pos + k_i], - num_logprobs, - argmax_tokens[local_idx][:k_i], - logprob_token_ids_state=( - self.logprob_token_ids_state if per_req_ids else None - ), - # every row of this stash belongs to one slot - expanded_idx_mapping=( - torch.full( - (k_i,), - slot.item(), - dtype=torch.int32, - device=flat_logits.device, - ) - if per_req_ids - else None - ), - max_per_req_token_ids=max_token_ids, - logits_mode=self.logits_mode, - ) - - # Commit steps: is_committing was True at entry. Reassemble previously - # stashed logprobs and attach to SamplerOutput. + scaled = _compiled_sample_step( + tile_logits, + tile_slots, + decode_idx[sel], + all_slots, + tile_valid, + # State, viewed at this tile's width + states.canvas[:, :W], + states.argmax_canvas[:, :W], + states.step, + states.is_encoder_phase, + states.confident, + states.self_conditioning_embeds[:, :W], + self.embed_weight, + self.normalizer, + states.accepted_canvas_history[:, :, :W], + states.accepted_canvas_history_len, + states.max_steps, + states.pin_mask[:, :W], + states.seed_canvas[:, :W], + states.read_only, + # Output + sampled[:, :W], + num_sampled, + self.req_states.draft_tokens, + # Config + max_denoising_steps=float(states.max_denoising_steps), + t_min=self.t_min, + t_max=self.t_max, + confidence_threshold=self.confidence_threshold, + vocab_size=self.vocab_size, + CL=W, + ST=states.stability_threshold, + entropy_bound=self.entropy_bound, + sc_vocab_start=self.sc_vocab_start, + sc_vocab_end=self.sc_vocab_end, + tp_size=self.tp_size, + tp_group_name=self.tp_group_name, + compute_sc=compute_sc, + ) + + # Stash newly converged logprobs, including reads that emit now. + if want_logprobs: + converged_mask = states.is_encoder_phase[tile_slots] | ( + num_sampled[decode_idx[sel]] > 0 + ) + just_converged = converged_mask & ~is_committing[sel] + if just_converged.any(): + flat_logits = scaled.reshape(-1, scaled.shape[-1]) + argmax_tokens = scaled.argmax(dim=-1) + raw_flat: torch.Tensor | None = None + for local_idx in just_converged.nonzero(as_tuple=True)[0]: + li = local_idx.item() + slot = tile_slots[local_idx].item() + # Stash only the real canvas positions; padded tail + # positions are never emitted. + k_i = int(tile_valid_np[li]) + pos = li * W + src = flat_logits + if slot in states.read_only_slots: + # Read-only slots report logprobs at temperature 1. + # The schedule-tempered logits share the argmax. + if raw_flat is None: + raw_flat = tile_logits.float() + src = raw_flat + per_req_ids = max_token_ids > 0 + self._pending_logprobs[slot] = compute_topk_scores( + src[pos : pos + k_i], + num_logprobs, + argmax_tokens[local_idx][:k_i], + logprob_token_ids_state=( + self.logprob_token_ids_state + if per_req_ids + else None + ), + # every row of this stash belongs to one slot + expanded_idx_mapping=( + torch.full( + (k_i,), slot, dtype=torch.int32, device=device + ) + if per_req_ids + else None + ), + max_per_req_token_ids=max_token_ids, + logits_mode=self.logits_mode, + ) + run_start = run_end + + # Only emitting requests consume their stashed logprobs. logprobs_tensors = None - if want_logprobs and is_committing.any() and self._pending_logprobs: - committing_slots = set( - decode_slots_np[is_committing.cpu().numpy()].tolist() - ) - parts_ids, parts_lp, parts_ranks = [], [], [] + if want_logprobs and self._pending_logprobs: + emitting_slots = set(slots_np[num_sampled.cpu().numpy() > 0].tolist()) + parts: list[LogprobsTensors] = [] cu_gen: list[int] = [] flat_offset = 0 for i in range(num_reqs): cu_gen.append(flat_offset) slot = int(slots_np[i]) - if slot in committing_slots and slot in self._pending_logprobs: + if slot in emitting_slots and slot in self._pending_logprobs: lp = self._pending_logprobs.pop(slot) - parts_ids.append(lp.logprob_token_ids) - parts_lp.append(lp.logprobs) - parts_ranks.append(lp.selected_token_ranks) + parts.append(lp) flat_offset += lp.logprobs.shape[0] - if parts_ids: - logprobs_tensors = LogprobsTensors( - logprob_token_ids=torch.cat(parts_ids), - logprobs=torch.cat(parts_lp), - selected_token_ranks=torch.cat(parts_ranks), - cu_num_generated_tokens=cu_gen, - ) + if parts: + logprobs_tensors = _concat_logprob_stashes(parts, cu_gen) return self._build_output( input_batch, diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 72f2920480a1..f0126ca34581 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -813,6 +813,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: "dbrx": DbrxModelArchConfigConvertor, "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor, "deepseek_v4": DeepseekV4ModelArchConfigConvertor, + "diffusion_gemma": Gemma4ModelArchConfigConvertor, "diffusion_gemma_text": Gemma4ModelArchConfigConvertor, "ernie_mtp": ErnieMTPModelArchConfigConvertor, "falcon": FalconModelArchConfigConvertor, diff --git a/vllm/utils/diffusion.py b/vllm/utils/diffusion.py new file mode 100644 index 000000000000..278c17d83823 --- /dev/null +++ b/vllm/utils/diffusion.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Request validation for discrete diffusion models.""" + +from typing import TYPE_CHECKING + +from vllm.exceptions import VLLMValidationError + +if TYPE_CHECKING: + from vllm.sampling_params import SamplingParams + + +def validate_diffusion_sampling_params( + params: "SamplingParams", + *, + canvas_length: int | None, + vocab_size: int, + async_scheduling: bool | None, +) -> None: + """Validate request options and normalize one-canvas reads at admission.""" + extra = params.extra_args + if not extra: + return + + def check(valid: bool, message: str) -> None: + if not valid: + raise VLLMValidationError(message, parameter="extra_args") + + width = extra.get("diffusion_canvas_length") + if width is not None: + check( + isinstance(width, int) + and not isinstance(width, bool) + and width > 0 + and (canvas_length is None or width <= canvas_length), + "diffusion_canvas_length must be a positive integer no larger " + "than the served canvas.", + ) + # The request's canvas: its own width, else the served one. + expected_len = width or canvas_length + + seed = extra.get("diffusion_seed_canvas") + if seed is not None: + check( + isinstance(seed, (list, tuple)) + and all(isinstance(t, int) and not isinstance(t, bool) for t in seed), + "diffusion_seed_canvas must be a list of token ids.", + ) + check( + all(0 <= t < vocab_size for t in seed), + f"diffusion_seed_canvas ids must be in [0, {vocab_size}).", + ) + check( + expected_len is None or len(seed) == expected_len, + f"diffusion_seed_canvas must hold exactly {expected_len} ids, " + f"got {len(seed)}.", + ) + + pins = extra.get("diffusion_pinned") + if pins is not None: + check( + isinstance(pins, (list, tuple)) + and all(isinstance(p, int) and not isinstance(p, bool) for p in pins), + "diffusion_pinned must be a list of canvas positions.", + ) + check( + seed is not None, + "diffusion_pinned needs a diffusion_seed_canvas to hold.", + ) + check( + all(p >= 0 and (expected_len is None or p < expected_len) for p in pins), + "diffusion_pinned positions must lie inside the canvas.", + ) + + cap = extra.get("diffusion_max_steps") + if cap is not None: + check( + isinstance(cap, int) and not isinstance(cap, bool) and cap > 0, + "diffusion_max_steps must be a positive integer.", + ) + + # The OpenAI server's vllm_xargs narrows JSON booleans to 0/1. + read_only = extra.get("diffusion_read_only") + if read_only is not None: + check( + isinstance(read_only, (bool, int)) and read_only in (0, 1), + "diffusion_read_only must be a boolean (or 0/1).", + ) + if read_only: + # One canvas is the whole output. Cap max_tokens so the scheduler + # ends the request there, and ignore EOS so an end-of-turn token + # drawn into a noise slot does not end it early. + if expected_len is not None: + params.max_tokens = min(params.max_tokens or expected_len, expected_len) + params.ignore_eos = True + + check( + bool(async_scheduling) + or width is None + or canvas_length is None + or width == canvas_length, + "A diffusion_canvas_length smaller than the served canvas " + "requires --async-scheduling.", + ) diff --git a/vllm/v1/core/sched/diffusion_scheduler.py b/vllm/v1/core/sched/diffusion_scheduler.py new file mode 100644 index 000000000000..96462aa6abe1 --- /dev/null +++ b/vllm/v1/core/sched/diffusion_scheduler.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Async scheduling for diffusion requests. + +These rules live here so Scheduler.schedule() and AsyncScheduler stay +unchanged. VllmConfig selects this class for a diffusion model under async +scheduling. A sync scheduler creates no output placeholders, which both rules +read. +""" + +from vllm.v1.core.sched.async_scheduler import AsyncScheduler +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.request import Request + + +def diffusion_canvas_width(request: Request, canvas_length: int) -> int: + """The canvas width a diffusion request asked for, else the served one.""" + params = request.sampling_params + extra = params.extra_args if params is not None else None + width = extra.get("diffusion_canvas_length") if extra else None + return int(width) if width else canvas_length + + +def _read_in_flight(request: Request, width: int) -> bool: + """True when a read-only request has all its denoise steps in flight. + + The request emits its canvas on the last step and ends, so a further + step is discarded. + """ + params = request.sampling_params + extra = params.extra_args if params is not None else None + if not extra or not extra.get("diffusion_read_only"): + return False + steps = extra.get("diffusion_max_steps") + if not steps: + return False + # Each in-flight denoise step holds one canvas of placeholders. + return request.num_output_placeholders >= steps * width + + +class DiffusionAsyncScheduler(AsyncScheduler): + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + for request in self.running: + width = diffusion_canvas_width(request, self.num_spec_tokens) + # The placeholders are as wide as the served canvas. + if len(request.spec_token_ids) > width: + request.spec_token_ids = request.spec_token_ids[:width] + if _read_in_flight(request, width): + # Scheduler.schedule()'s max_tokens guard cannot be reached + # from a subclass. schedule() advances current_step before it + # checks decode eligibility, so current_step + 2 skips this + # call alone. max keeps a longer pipeline-parallel wait. + request.next_decode_eligible_step = max( + request.next_decode_eligible_step, self.current_step + 2 + ) + return super().schedule(throttle_prefills) diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 582e1abac685..199f8898531f 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -30,6 +30,7 @@ from vllm.tokenizers import TokenizerLike from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid from vllm.utils.async_utils import make_async +from vllm.utils.diffusion import validate_diffusion_sampling_params from vllm.utils.jsontree import json_iter_leaves from vllm.v1.engine import EngineCoreRequest from vllm.v1.kv_hints import KvHintsEnvelope @@ -53,6 +54,7 @@ def __init__( self.speculative_config = vllm_config.speculative_config self.structured_outputs_config = vllm_config.structured_outputs_config self.observability_config = vllm_config.observability_config + self.diffusion_config = vllm_config.diffusion_config # Load the custom logits processor classes once; the returned callable # runs their validate_params hooks per request at admission. self.validate_logits_processors_params = ( @@ -129,6 +131,20 @@ def _validate_params( self.validate_logits_processors_params(params) + if self.model_config.is_diffusion: + # Without --diffusion-config the served canvas is unknown here; + # the ids and the read-only normalisation are still checked. + validate_diffusion_sampling_params( + params, + canvas_length=( + self.diffusion_config.canvas_length + if self.diffusion_config is not None + else None + ), + vocab_size=self.model_config.get_vocab_size(), + async_scheduling=self.vllm_config.scheduler_config.async_scheduling, + ) + if self.model_config.return_sampling_mask: if params.temperature <= 0: raise ValueError(