[Perf] DiffusionGemma: constrained reads over the request's logprob_token_ids (#58216)

Signed-off-by: Matt Mastracci <matthew@mastracci.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>

examples/features/structured_diffusion/README.md

@@ -11,8 +11,10 @@ left as noise, one denoise step gives a distribution over each slot. These
 | `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 |
+| `diffusion_constrained` | `bool` | run the unembedding, sampler and self-conditioning over the request's `logprob_token_ids` only. Logprobs are normalized over that set. A step uses this only when every read in it has the same set |
 
-`structured_server.py` turns a question schema into those fields. It serves
+`structured_server.py` turns a question schema into those fields, with
+`diffusion_constrained` on for every read (`--no-constrained` turns it off). 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.

examples/features/structured_diffusion/structured_server.py

@@ -399,6 +399,14 @@ def canvas_width(template):
     return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP)
 
 
+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):
     """Past one denoise step the template must be held, or accept/renoise
     rewrites it: pin every canvas position that is not an answer slot."""
@@ -552,6 +560,7 @@ def one_read(
             "diffusion_max_steps": schema["steps"],
             "diffusion_read_only": True,
             **pin_xargs(template, slots, schema["steps"]),
+            **constrained_xargs(),
         },
     }
     d = upstream_chat(body)
@@ -601,6 +610,7 @@ def one_read_continuation(schema, template, slots, prompt_ids, seed):
             "diffusion_max_steps": schema["steps"],
             "diffusion_read_only": True,
             **pin_xargs(template, slots, schema["steps"]),
+            **constrained_xargs(),
         },
     }
     d = upstream_completions(body)
@@ -1496,6 +1506,12 @@ def main():
         default=16,
         help="request widths round up to a multiple of this",
     )
+    p.add_argument(
+        "--no-constrained",
+        dest="constrained",
+        action="store_false",
+        help="read over the whole vocabulary instead of the labels",
+    )
     p.add_argument("--host", default="0.0.0.0")
     p.add_argument("--port", type=int, default=8011)
     p.add_argument(