Skip to content

Add --constrained and --engine-samples for the engine perf branches #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mmastrac merged 1 commit into from
Sep 24, 2026
Merged

Add --constrained and --engine-samples for the engine perf branches #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
156 changes: 117 additions & 39 deletions structured_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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},
Expand All @@ -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):
Expand All @@ -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 = {
Expand All @@ -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(
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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")
Expand All @@ -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)"
)
Expand All @@ -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)
Expand Down