[Feat] Add Nemotron seeded decision reads to the structured gateway

docs/models/nemotron_labs_diffusion.md

@@ -125,3 +125,70 @@ NEMOTRON_DLM_MODEL_PATH=/path/to/Nemotron-Labs-Diffusion-8B \
     .venv/bin/python -m pytest --confcutdir=tests/models/language/generation \
     tests/models/language/generation/test_nemotron_dllm.py -v
 ```
+
+## Structured decision reads
+
+The masked-diffusion backend also supports seeded, read-only canvases. Prompt
+prefill is followed by one bidirectional masked forward, without an autoregressive
+verification/commit pass. Normal generation remains unchanged.
+
+Reuse the structured diffusion example to expose a TypeSafe-shaped
+`POST /v1/systemone` API. This is an optional **separate gateway process**, not
+an endpoint built into `vllm serve`:
+
+```bash
+VLLM_USE_V2_MODEL_RUNNER=1 vllm serve nvidia/Nemotron-Labs-Diffusion-8B \
+  --served-model-name jev-latest --port 8000 \
+  --diffusion-config '{"canvas_length":32}' --max-logprobs 128 \
+  --enable-prefix-caching
+
+API_KEY="$DECISION_API_KEY" python examples/features/structured_diffusion/structured_server.py \
+  --backend nemotron --tokenizer nvidia/Nemotron-Labs-Diffusion-8B \
+  --model jev-latest --canvas 32 --upstream http://localhost:8000 --port 8011
+```
+
+`jev-latest` is a local serving alias, not the TypeSafe Jev checkpoint.
+The gateway supports choice, Noul and Score questions over string or JSON state:
+
+```bash
+curl http://localhost:8011/v1/systemone \
+  -H "Authorization: Bearer $DECISION_API_KEY" \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"jev-latest","state":"Customer payouts have failed for three days.",
+       "questions":{"route":{"type":"choice","instructions":"Which team should handle this?",
+       "criteria":{"billing":"Payments and payouts","technical":"Infrastructure outages"}},
+       "urgent":{"type":"noul","instructions":"Has the problem lasted multiple days?"}}}'
+```
+
+Answers preserve the question IDs. Choice returns `choice`, `probabilities` and
+`confidence`; Noul returns `noul` (the probability of yes); Score returns a
+zero-indexed probability-weighted `score`, `legend`, `probabilities` and
+`confidence`. These are normalized model scores, not calibrated guarantees.
+
+This example accepts at most 64 questions and 26 alternatives per question.
+It chunks answer templates that exceed the canvas and keeps input order; there
+are no option- or position-ordering heuristics. Every option must resolve to one
+token at the same answer position. An incompatible template fails explicitly
+rather than approximating missing probabilities. Nemotron reads are text-only,
+deterministic (`samples=1`), and do not support the example's thinking or multi-step
+extensions. See the [OpenAPI schema](../../examples/features/structured_diffusion/systemone-openapi.yaml).
+
+For direct OpenAI-compatible calls, pass `logprob_token_ids` (up to 128 candidates),
+`return_tokens_as_token_ids=true`, `logprobs=true` and these `vllm_xargs`:
+
+```json
+{
+  "diffusion_read_only": true,
+  "diffusion_seed_canvas": [100, 100, 100, 100, 100, 100, 100, 100,
+                            100, 100, 100, 100, 100, 100, 100, 100,
+                            100, 100, 100, 100, 100, 100, 100, 100,
+                            100, 100, 100, 100, 100, 100, 100, 100]
+}
+```
+
+Use `temperature=0`, default `top_p=1`/`top_k=0`, and `max_tokens` no greater
+than the engine canvas length. The seed must contain exactly that many valid
+token IDs; `100` is the checkpoint's mask token. Non-mask positions remain fixed.
+The sampled tokens only transport the per-position logprobs: decisions should
+normalize the requested candidates at each mask position. Read-only scoring and
+indexed logprobs are not supported with `algorithm="linear_spec"`.

examples/features/structured_diffusion/README.md

@@ -72,3 +72,22 @@ 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.
+
+## Nemotron Labs Diffusion
+
+The same `/v1/systemone` gateway can serve Nemotron Labs Diffusion 3B/8B with
+`--backend nemotron`. See the [model guide](../../../docs/models/nemotron_labs_diffusion.md#structured-decision-reads)
+for launch commands and the [OpenAPI schema](systemone-openapi.yaml) for request shapes.
+This backend uses mask-token answer slots instead of random noise, one masked
+forward after prefill, and exact candidate logprobs. It does not reorder options
+or questions. Text/JSON state and choice, Noul and Score answers are supported;
+thinking, images and multiple samples are not.
+
+Verify a running gateway with:
+
+```bash
+python examples/features/structured_diffusion/nemotron_smoke.py --url http://localhost:8011
+```
+
+The smoke check validates answer shapes and finite, normalized probabilities for
+all three primitives and a 12-question request. It is not an accuracy benchmark.

examples/features/structured_diffusion/nemotron_smoke.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Exercise the existing System One gateway against a running Nemotron backend."""
+
+import argparse
+import json
+import math
+import os
+import time
+import urllib.request
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", default="http://localhost:8011")
+    args = parser.parse_args()
+    headers = {"Content-Type": "application/json"}
+    if key := os.environ.get("API_KEY"):
+        headers["Authorization"] = f"Bearer {key}"
+    cases = {
+        "choice": {
+            "route": {
+                "type": "choice",
+                "instructions": "Which team should handle this?",
+                "criteria": {
+                    "billing": "Payments and payouts",
+                    "technical": "Infrastructure outages",
+                },
+            }
+        },
+        "noul": {
+            "urgent": {"type": "noul", "instructions": "Has this lasted multiple days?"}
+        },
+        "score": {
+            "severity": {
+                "type": "score",
+                "instructions": "How disruptive is this problem?",
+                "criteria": [
+                    "No problem",
+                    "Minor inconvenience",
+                    "Blocked payments for multiple days",
+                ],
+            }
+        },
+        "twelve_questions": {
+            f"q{i}": {
+                "type": "noul",
+                "instructions": "Have payouts failed for multiple days?",
+            }
+            for i in range(12)
+        },
+    }
+    results = {}
+    for name, questions in cases.items():
+        body = {
+            "model": "jev-latest",
+            "state": "Customer payouts have failed for three days.",
+            "questions": questions,
+        }
+        started = time.perf_counter()
+        request = urllib.request.Request(
+            args.url.rstrip("/") + "/v1/systemone",
+            headers=headers,
+            data=json.dumps(body).encode(),
+        )
+        with urllib.request.urlopen(request, timeout=120) as response:
+            result = json.load(response)
+        elapsed = (time.perf_counter() - started) * 1000
+        assert result["answers"].keys() == questions.keys()
+        for answer in result["answers"].values():
+            if answer["type"] == "noul":
+                assert math.isfinite(answer["noul"]) and 0 <= answer["noul"] <= 1
+            else:
+                probabilities = answer["probabilities"]
+                assert all(
+                    math.isfinite(p) and 0 <= p <= 1 for p in probabilities.values()
+                )
+                assert math.isclose(sum(probabilities.values()), 1, abs_tol=1e-6)
+        results[name] = {"ms": round(elapsed, 2), "answers": result["answers"]}
+    print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()

examples/features/structured_diffusion/structured_server.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Structured decisions in front of a vLLM DiffusionGemma server.
+"""Structured decisions in front of a vLLM diffusion server.
 
 POST /v1/systemone takes Jev's request body: {"model", "state", "questions"}.
 "questions" maps an id to {"type", "instructions", "criteria"}, where "type"
@@ -114,6 +114,8 @@
 SCAFFOLD = None
 THOUGHT_OPEN = None
 THOUGHT_CLOSE = None
+BACKEND = "gemma"
+MASK_TOKEN_ID = 100
 
 
 # ----------------------------------------------------------------------------
@@ -213,13 +215,17 @@ def parse_schema(value):
                     f"must be among {names}"
                 )
     schedule(qs)  # refuses a cycle
-    samples = value.get("samples", "auto")
+    samples = value.get("samples", 1 if BACKEND == "nemotron" else "auto")
+    if BACKEND == "nemotron" and samples not in (1, "auto"):
+        raise SchemaError("Nemotron masked reads are deterministic; use samples=1")
     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)),
         }
+        if BACKEND == "nemotron":
+            policy["max"] = 1
     elif isinstance(samples, int) and samples >= 1:
         policy = {"mode": "fixed", "n": min(samples, MAX_SAMPLES)}
     else:
@@ -243,6 +249,8 @@ def parse_schema(value):
     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")
+    if BACKEND == "nemotron" and (think or int(value.get("steps", 1)) != 1):
+        raise SchemaError("Nemotron reads support one masked forward, without thinking")
     return {
         "questions": qs,
         "instructions": value.get("instructions"),
@@ -318,8 +326,16 @@ def enc(text):
 
 
 def init_tokenizer(tok):
-    global TOK, SCAFFOLD, THOUGHT_OPEN, THOUGHT_CLOSE
+    global TOK, SCAFFOLD, THOUGHT_OPEN, THOUGHT_CLOSE, TURN_CLOSE, PAD
     TOK = tok
+    _template_cache.clear()
+    if BACKEND == "nemotron":
+        SCAFFOLD, THOUGHT_OPEN, THOUGHT_CLOSE = [], [], []
+        TURN_CLOSE = TOK.eos_token_id
+        PAD = TOK.pad_token_id if TOK.pad_token_id is not None else TURN_CLOSE
+        if TURN_CLOSE is None:
+            raise ValueError("The Nemotron tokenizer must define eos_token_id")
+        return
     THOUGHT_OPEN = enc("<|channel>thought\n")
     THOUGHT_CLOSE = enc("<channel|>")
     SCAFFOLD = enc(SCAFFOLD_TEXT)
@@ -395,6 +411,8 @@ def template_for(schema, head, lead):
 
 def canvas_width(template):
     """Smallest multiple of CANVAS_STEP that holds the template and the turn close."""
+    if BACKEND == "nemotron":
+        return CANVAS_LEN
     need = len(template) + 1
     return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP)
 
@@ -415,13 +433,17 @@ def build_canvas(template, slots, seed):
     canvas = list(template) + [TURN_CLOSE]
     canvas += [PAD] * (canvas_width(template) - len(canvas))
     for s in slots:
-        canvas[s["pos"]] = rng.randrange(VOCAB)
+        canvas[s["pos"]] = (
+            MASK_TOKEN_ID if BACKEND == "nemotron" else 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.
+    if len(ids) > 128:
+        raise SchemaError("This read exceeds vLLM's 128 candidate-token limit")
+    return ids
 
 
 def upstream_chat(body, timeout=600):
@@ -528,6 +550,8 @@ def think_chat(sys_text, state_content, budget):
 def one_read(
     schema, template, slots, sys_text, state_content, seed, prefix=None, thinking=False
 ):
+    if BACKEND == "nemotron" and not isinstance(state_content, str):
+        raise SchemaError("Nemotron reads support text states only")
     if prefix is not None:
         return one_read_continuation(schema, template, slots, prefix, seed)
     messages = [
@@ -536,6 +560,7 @@ def one_read(
     ]
     body = {
         "model": ARGS.model,
+        **({"temperature": 0} if BACKEND == "nemotron" else {}),
         "messages": messages,
         "max_tokens": len(template) + 1,
         "logprobs": True,
@@ -571,6 +596,8 @@ def slot_distribution(top, label_ids):
     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."""
+    if BACKEND == "nemotron" and any(i not in top for i in label_ids):
+        raise SchemaError("Upstream omitted requested candidate logprobs")
     floor = min(top.values()) - 5.0
     lp_t = [top.get(i, floor) for i in label_ids]
     mx = max(lp_t)
@@ -591,6 +618,7 @@ def one_read_continuation(schema, template, slots, prompt_ids, seed):
     body = {
         "model": ARGS.model,
         "prompt": prompt_ids,
+        **({"temperature": 0} if BACKEND == "nemotron" else {}),
         "max_tokens": len(template) + 1,
         "logprobs": TOPK,
         "logprob_token_ids": label_id_union(slots),
@@ -1484,11 +1512,12 @@ def serve_tls(host, port, cert_dir):
 
 
 def main():
-    global ARGS, CANVAS_LEN, CANVAS_STEP
+    global ARGS, CANVAS_LEN, CANVAS_STEP, BACKEND
     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("--backend", choices=["gemma", "nemotron"], default="gemma")
     p.add_argument("--canvas", type=int, default=64, help="the served canvas length")
     p.add_argument(
         "--canvas-step",
@@ -1509,6 +1538,7 @@ def main():
     ARGS = p.parse_args()
     CANVAS_LEN = ARGS.canvas
     CANVAS_STEP = ARGS.canvas_step
+    BACKEND = ARGS.backend
     init_tokenizer(AutoTokenizer.from_pretrained(ARGS.tokenizer))
     if ARGS.tls_port:
         serve_tls(ARGS.host, ARGS.tls_port, ARGS.cert_dir)