From e14af9bd7355921d37686cda838183840573b188 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 23 Sep 2026 17:54:54 -0600 Subject: [PATCH] Add --constrained and --engine-samples for the engine perf branches Ports the example server's side of two open vLLM changes, from the fork's constrained-vocab (vllm#58216) and diffusion-samples (vllm#58438) branches: - --constrained sends diffusion_constrained on every label read, so the engine runs the unembedding, sampler and self-conditioning over the labels only. - --engine-samples sends a fixed sample count as one diffusion_samples request with a seed, and reads one choice per draw. --max-samples caps the count. Both are opt-in, unlike the fork's example server, because a stock vLLM rejects seed on a diffusion request. The read path now returns one distribution list per draw, and diagnostics.timing gains "requests". Co-Authored-By: Claude Opus 5.5 --- README.md | 11 +++ structured_server.py | 156 ++++++++++++++++++++++++++++++++----------- 2 files changed, 128 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 710a1cf..8fd83dc 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,17 @@ scheduling. The diffusion async scheduler is selected automatically; no `--scheduler-cls` argument is needed. Synchronous execution supports full-width canvases only. +Two flags use engine changes that are still open against vLLM, so both are +off by default: + +| flag | engine change | effect | +|---|---|---| +| `--constrained` | [vllm#58216](https://github.com/vllm-project/vllm/pull/58216), `diffusion_constrained` | reads run the unembedding, sampler and self-conditioning over the labels only: the same argmax for about a quarter less GPU time. Logprobs are normalized over the labels, so `label_mass` is always 1 | +| `--engine-samples` | [vllm#58438](https://github.com/vllm-project/vllm/pull/58438), `diffusion_samples` | a fixed `samples: N` goes to vLLM as one request that fans out into N noise draws, instead of N requests. `--max-samples` caps N and must fit the engine's `diffusion_config.max_samples` | + +A stock vLLM rejects `seed` on a diffusion request, so `--engine-samples` +fails there with a 400; `--constrained` is ignored. + Question types: `noul` (yes/no), `choice` with `options`, `score` with ordered `levels`, and `span` / `spans` (below). Each label must be a single token in the answer template, which the server checks with the tokenizer diff --git a/structured_server.py b/structured_server.py index 588d349..98f647e 100644 --- a/structured_server.py +++ b/structured_server.py @@ -83,6 +83,13 @@ 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 + +Two flags use engine features that are not in vLLM main yet, so both are off +by default. --constrained reads over the labels only (vllm-project/vllm#58216): +the same argmax for about a quarter less GPU time, with logprobs normalized +over the labels, so label mass is always 1. --engine-samples sends a fixed +sample count as one diffusion_samples request (vllm-project/vllm#58438) +instead of one request per draw. """ import argparse @@ -115,7 +122,7 @@ PAD = 0 TOPK = 20 MAX_QUESTIONS = 64 # per request -MAX_SAMPLES = 32 # reads per question, fixed or auto +MAX_SAMPLES = 32 # cap on reads per question, set by --max-samples MAX_PARALLEL = 16 # question groups read at once SPAN_STEPS = ( 4 # denoise steps for a span blank; one step spells it position by position @@ -462,10 +469,20 @@ def canvas_width(template): return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP) -def pin_xargs(template, slots, steps): +def constrained_xargs(): + """Read over the labels only. The engine runs the unembedding, sampler and + self-conditioning over the request's logprob_token_ids instead of the whole + vocabulary. Same argmax, about a quarter less GPU time per read. Logprobs + come back normalized over the labels.""" + return {"diffusion_constrained": True} if ARGS.constrained else {} + + +def pin_xargs(template, slots, steps, samples=1): """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: + rewrites it: pin every canvas position that is not an answer slot. A + multi-sample request pins the same positions, which tells the engine + what to re-noise per sample.""" + if steps <= 1 and samples <= 1: return {} free = {s["pos"] for s in slots} return { @@ -588,11 +605,41 @@ def think_chat(sys_text, state_content, budget): } +def read_xargs(schema, template, slots, seed, samples): + x = { + "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"], samples), + **constrained_xargs(), + } + if samples > 1: + x["diffusion_samples"] = samples + return x + + +def samples_body(seed, samples): + """The engine draws the noise of a multi-sample request from the request + seed, so the draws reproduce for a schema seed.""" + return {"seed": seed} if samples > 1 else {} + + def one_read( - schema, template, slots, sys_text, state_content, seed, prefix=None, thinking=False + schema, + template, + slots, + sys_text, + state_content, + seed, + prefix=None, + thinking=False, + samples=1, ): + """One request, ``samples`` noise draws of the canvas: a list of one + distribution per question for each draw, and the usage.""" if prefix is not None: - return one_read_continuation(schema, template, slots, prefix, seed) + return one_read_continuation(schema, template, slots, prefix, seed, samples) messages = [ {"role": "system", "content": sys_text}, {"role": "user", "content": state_content}, @@ -609,24 +656,22 @@ def one_read( "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"]), - }, + "vllm_xargs": read_xargs(schema, template, slots, seed, samples), + **samples_body(seed, samples), } 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", {}) + draws = [] + for choice in d["choices"]: + content = choice["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"])) + draws.append(out) + return draws, d.get("usage", {}) def slot_distribution(top, label_ids): @@ -648,7 +693,7 @@ def slot_distribution(top, label_ids): } -def one_read_continuation(schema, template, slots, prompt_ids, seed): +def one_read_continuation(schema, template, slots, prompt_ids, seed, samples=1): """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 = { @@ -658,21 +703,19 @@ def one_read_continuation(schema, template, slots, prompt_ids, seed): "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"]), - }, + "vllm_xargs": read_xargs(schema, template, slots, seed, samples), + **samples_body(seed, samples), } 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", {}) + draws = [] + for choice in d["choices"]: + rows = choice["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"])) + draws.append(out) + return draws, d.get("usage", {}) def read_many( @@ -685,14 +728,22 @@ def read_many( n, prefix=None, thinking=False, + engine=False, ): + """``n`` noise draws of the read: one diffusion_samples request when + ``engine`` is set, otherwise one request per draw.""" + if n > 1 and engine and ARGS.engine_samples: + draws, usage = one_read( + schema, template, slots, sys_text, state_content, seed, prefix, thinking, n + ) + return draws, [usage] + [None] * (len(draws) - 1) results = [None] * n errors = [None] * n usages = [None] * n def run(k): try: - results[k], usages[k] = one_read( + (results[k],), usages[k] = one_read( schema, template, slots, @@ -974,6 +1025,7 @@ def absorb(group, body, rows): "timing": { "total_ms": (time.time() - started) * 1e3, "reads": sum(b["diagnostics"]["timing"]["reads"] for b, _ in parts), + "requests": sum(b["diagnostics"]["timing"]["requests"] for b, _ in parts), }, "prompt_tokens": max( (b["diagnostics"].get("prompt_tokens") or 0) for b, _ in parts @@ -1543,6 +1595,7 @@ def decide_group(schema, sys_text, state_content, seed, prefix=None, lead=""): policy["n"], prefix, thinking, + engine=True, ) extended = None first_entropy = None @@ -1632,7 +1685,11 @@ def decide_group(schema, sys_text, state_content, seed, prefix=None, lead=""): policy, extended=extended, first_read_entropy=first_entropy ), }, - "timing": {"total_ms": elapsed_ms, "reads": n}, + "timing": { + "total_ms": elapsed_ms, + "reads": n, + "requests": sum(1 for u in usages if u is not None), + }, "thought": thought, "prompt_tokens": prompt_tokens, "questions": diag_q, @@ -2177,7 +2234,7 @@ def serve_tls(host, port, cert_dir): def main(): - global ARGS, CANVAS_LEN, CANVAS_STEP + global ARGS, CANVAS_LEN, CANVAS_STEP, MAX_SAMPLES p = argparse.ArgumentParser() p.add_argument("--upstream", default="http://127.0.0.1:8010") p.add_argument("--model", default="dgemma") @@ -2189,8 +2246,28 @@ def main(): default=16, help="request widths round up to a multiple of this", ) + p.add_argument( + "--constrained", + action="store_true", + help="read over the labels only (diffusion_constrained). Needs an " + "engine with vllm-project/vllm#58216", + ) p.add_argument("--host", default="0.0.0.0") p.add_argument("--port", type=int, default=8011) + p.add_argument( + "--max-samples", + type=int, + default=32, + help="cap on reads per question. A fixed count must also fit the " + "engine's diffusion_config.max_samples", + ) + p.add_argument( + "--engine-samples", + action="store_true", + help="send a fixed sample count as one diffusion_samples request " + "instead of one request per draw. Needs an engine with " + "vllm-project/vllm#58438", + ) p.add_argument( "--tls-port", type=int, default=0, help="also listen with HTTPS here (0 = off)" ) @@ -2202,6 +2279,7 @@ def main(): ARGS = p.parse_args() CANVAS_LEN = ARGS.canvas CANVAS_STEP = ARGS.canvas_step + MAX_SAMPLES = max(1, ARGS.max_samples) init_tokenizer(AutoTokenizer.from_pretrained(ARGS.tokenizer)) if ARGS.tls_port: serve_tls(ARGS.host, ARGS.tls_port, ARGS.cert_dir)