[Feat][Frontend] structured decisions: DiffusionGemma read strategy (#60466)

Signed-off-by: Matt Mastracci <matthew@mastracci.com>
Signed-off-by: Simon Veitner <sveitner@redhat.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Simon Veitner <sveitner@redhat.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>

docs/serving/online_serving/structured_decisions.md

@@ -20,7 +20,10 @@ log probabilities.
 
 Thinking is off unless `chat_template_kwargs` turns it on.
 
-This endpoint does not serve diffusion models yet.
+Probabilities can differ between identical requests. They depend on which
+requests share a batch and on whether a read hits the prefix cache. Set
+`VLLM_BATCH_INVARIANT=1` for reads that repeat exactly, at a cost in
+throughput.
 
 Every read of a request starts with the state, so with
 `--enable-prefix-caching` the state is prefilled once and each further

tests/entrypoints/generate/structured_decisions/test_question_types.py

@@ -105,7 +105,7 @@ def test_noul_labels_and_prompt():
     assert q.labels == ("yes", "no")
     assert [a.name for a in q.options] == ["yes", "no"]
     assert q.type.prompt(q) == (
-        "Question: Reply within the hour?\nAnswer with yes or no only."
+        "Question: Reply within the hour?\nyes\nno\nAnswer with yes or no only."
     )
     described = noul(criteria={"true": "same day", "false": "later"})
     assert described.type.prompt(described) == (

tests/entrypoints/generate/structured_decisions/test_strategies.py

@@ -17,12 +17,16 @@ from vllm.entrypoints.generate.structured_decisions.protocol import (
     StructuredDecisionRequest,
     StructuredDecisionResponse,
 )
-from vllm.entrypoints.generate.structured_decisions.question_types import LABELS
+from vllm.entrypoints.generate.structured_decisions.question_types import (
+    LABELS,
+    StructuredDecisionError,
+)
 from vllm.entrypoints.generate.structured_decisions.serving import (
     ServingStructuredDecisions,
     state_text,
 )
 from vllm.entrypoints.generate.structured_decisions.strategies import (
+    DiffusionGemmaCanvasStrategy,
     NextTokenStrategy,
     ReadContext,
     reply_label_ids,
@@ -52,6 +56,10 @@ def test_strategy_selection():
     qwen = "Qwen3ForCausalLM"
     assert select_read_strategy(model(qwen)) is NextTokenStrategy
     assert select_read_strategy(model(qwen, "processed_logprobs")) is NextTokenStrategy
+    assert (
+        select_read_strategy(model("DiffusionGemmaForBlockDiffusion"))
+        is DiffusionGemmaCanvasStrategy
+    )
     with pytest.raises(ValueError, match="does not support LlamaForCausalLM"):
         select_read_strategy(model("LlamaForCausalLM"))
     with pytest.raises(ValueError, match="not raw_logits"):
@@ -220,3 +228,30 @@ async def test_request_metadata_reaches_every_question(
         assert call.kwargs["trace_headers"] == (
             TRACE_HEADERS if tracing_enabled and request_headers else None
         )
+
+
+def test_canvas_read():
+    strategy = DiffusionGemmaCanvasStrategy.__new__(DiffusionGemmaCanvasStrategy)
+    strategy.thought, strategy.end, strategy.pad = [10, 11, 12, 13], 106, 0
+    strategy.width, strategy.vocab_size = 16, 1000
+    strategy.max_model_len = 23
+    params = strategy._sampling_params([65, 66], prompt_ids=[1, 2, 3])
+    assert params.extra_args is not None
+    canvas = params.extra_args["diffusion_seed_canvas"]
+    # The label slot as noise, the end of the turn, padding.
+    assert canvas[1:] == [106] + [0] * 14
+    assert params.max_tokens == 2 and params.logprob_token_ids == [65, 66]
+    again = strategy._sampling_params([65, 66], prompt_ids=[1, 2, 3])
+    assert again.extra_args == params.extra_args
+    read_input = strategy._read_input(
+        {"type": "token", "prompt_token_ids": [1, 2, 3]}, [1, 2, 3]
+    )
+    assert read_input == {
+        "type": "token",
+        "prompt_token_ids": [1, 2, 3, 10, 11, 12, 13],
+    }
+    # One token more and the thought and the canvas no longer fit.
+    with pytest.raises(StructuredDecisionError, match="max_model_len=23"):
+        strategy._read_input(
+            {"type": "token", "prompt_token_ids": [1, 2, 3, 4]}, [1, 2, 3, 4]
+        )

vllm/entrypoints/generate/api_router.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any
 from fastapi import FastAPI
 
 import vllm.envs as envs
+from vllm.logger import init_logger
 
 if TYPE_CHECKING:
     from argparse import Namespace
@@ -17,6 +18,8 @@ if TYPE_CHECKING:
 else:
     RequestLogger = object
 
+logger = init_logger(__name__)
+
 
 def register_generate_api_routers(app: FastAPI):
     from vllm.entrypoints.openai.chat_completion.api_router import (
@@ -243,7 +246,9 @@ async def init_generate_state(
                     default_chat_template_kwargs=default_chat_template_kwargs,
                 )
             )
-        except ValueError:
+        except ValueError as e:
+            # Info, since every model without a read strategy lands here.
+            logger.info("/v1/systemone is disabled: %s", e)
             strategy = None
     state.serving_structured_decisions = (
         ServingStructuredDecisions(

vllm/entrypoints/generate/structured_decisions/question_types.py

@@ -55,8 +55,7 @@ class QuestionType(ABC):
         ``question.labels[i]`` among the labels, and the list sums to 1.
         ``label_mass`` is the labels' total probability over the vocabulary."""
 
-    def reply_line(self, label: str, option: Option) -> str | None:
-        """The option's line in the prompt, or None to leave it out."""
+    def reply_line(self, label: str, option: Option) -> str:
         return (
             f"{label}: {option.name} - {option.description}"
             if option.description
@@ -67,9 +66,7 @@ class QuestionType(ABC):
         """The question as the model reads it, after the state."""
         lines = [f"Question: {question.instructions}"] if question.instructions else []
         for label, o in zip(question.labels, question.options):
-            line = self.reply_line(label, o)
-            if line is not None:
-                lines.append(line)
+            lines.append(self.reply_line(label, o))
         lines.append(self.reply_instruction)
         return "\n".join(lines)
 
@@ -192,8 +189,10 @@ class NoulQuestion(QuestionType):
             Option("no", None if false is None else str(false)),
         ]
 
-    def reply_line(self, label: str, option: Option) -> str | None:
-        return f"{label}: {option.description}" if option.description else None
+    def reply_line(self, label: str, option: Option) -> str:
+        # A side without a description still shows its label. Without the
+        # label in the prompt, models answer "Yes" or "No" instead.
+        return f"{label}: {option.description}" if option.description else label
 
     def answer(
         self, question: Question, probs: list[float], label_mass: float

vllm/entrypoints/generate/structured_decisions/strategies.py

@@ -2,10 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Read strategies: how a model is asked for label probabilities.
 
-Only models in NEXT_TOKEN_ARCHITECTURES are currently supported.
+Only models in READ_STRATEGIES are currently supported.
 """
 
+import hashlib
+import json
 import math
+import random
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
@@ -15,7 +18,7 @@ from vllm.config import ModelConfig
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
 from vllm.entrypoints.generate.label_reads import next_token_label_reads
-from vllm.inputs import EngineInput
+from vllm.inputs import EngineInput, tokens_input
 from vllm.lora.request import LoRARequest
 from vllm.renderers.inputs.preprocess import extract_prompt_components
 from vllm.renderers.online_renderer import OnlineRenderer
@@ -80,6 +83,13 @@ class ReadStrategy(ABC):
         """One read per question, in order."""
 
 
+def prompt_tail(tokenizer: TokenizerLike, prompt_ids: Sequence[int]) -> list[int]:
+    """The prompt's ids after its last added token."""
+    added = set(tokenizer.get_added_vocab().values())
+    start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0)
+    return list(prompt_ids[start:])
+
+
 def reply_label_ids(
     tokenizer: TokenizerLike,
     prompt_ids: Sequence[int],
@@ -88,9 +98,7 @@ def reply_label_ids(
     """The prompt's ids after its last added token, and the token each of
     ``labels`` adds after them as the first token of the reply. Raises
     ValueError if a label is not one distinct token there."""
-    added = set(tokenizer.get_added_vocab().values())
-    start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0)
-    tail = list(prompt_ids[start:])
+    tail = prompt_tail(tokenizer, prompt_ids)
     text = tokenizer.decode(tail)
     ids: list[int] = []
     for label in labels:
@@ -104,6 +112,18 @@ def reply_label_ids(
     return tail, ids
 
 
+def type_label_ids(
+    tokenizer: TokenizerLike, prompt_ids: Sequence[int]
+) -> dict[str, int]:
+    """The reply's first token for each label of every registered question
+    type. Raises ValueError as reply_label_ids does."""
+    label_ids: dict[str, int] = {}
+    for qtype in QUESTION_TYPES.values():
+        _, ids = reply_label_ids(tokenizer, prompt_ids, qtype.label_set)
+        label_ids.update(zip(qtype.label_set, ids))
+    return label_ids
+
+
 class NextTokenStrategy(ReadStrategy):
     """Autoregressive models. Each question is one request: the state, then the
     question with its labeled options, and the logprobs of the label tokens as
@@ -112,7 +132,18 @@ class NextTokenStrategy(ReadStrategy):
 
     def __init__(self, context: ReadContext):
         super().__init__(context)
-        tokenizer = context.online_renderer.renderer.get_tokenizer()
+        # Every prompt ends with the same generation prompt, so a label's token
+        # is the same for every question.
+        prompt = self._generation_prompt()
+        self.tail = prompt_tail(self._tokenizer(), prompt)
+        self.label_ids = type_label_ids(self._tokenizer(), prompt)
+
+    def _tokenizer(self) -> TokenizerLike:
+        return self.context.online_renderer.renderer.get_tokenizer()
+
+    def _generation_prompt(self) -> list[int]:
+        """A chat prompt's ids up to the generation prompt, with thinking off."""
+        context, tokenizer = self.context, self._tokenizer()
         probe = tokenizer.apply_chat_template(
             [{"role": "user", "content": "x"}],
             **{
@@ -126,13 +157,20 @@ class NextTokenStrategy(ReadStrategy):
         )
         if isinstance(probe, str):
             probe = tokenizer.encode(probe, add_special_tokens=False)
-        # Every prompt ends with the same generation prompt, so a label's token
-        # is the same for every question.
-        self.tail, _ = reply_label_ids(tokenizer, probe, ())
-        self.label_ids: dict[str, int] = {}
-        for qtype in QUESTION_TYPES.values():
-            _, ids = reply_label_ids(tokenizer, probe, qtype.label_set)
-            self.label_ids.update(zip(qtype.label_set, ids))
+        return list(probe)
+
+    def _sampling_params(
+        self, label_ids: list[int], prompt_ids: list[int]
+    ) -> SamplingParams:
+        return SamplingParams(
+            max_tokens=1, temperature=0.0, logprob_token_ids=label_ids
+        )
+
+    def _read_input(
+        self, engine_input: EngineInput, prompt_ids: list[int]
+    ) -> EngineInput:
+        """The engine input for a read of the rendered prompt."""
+        return engine_input
 
     def limits(self) -> DecisionLimits:
         return DecisionLimits(max_questions=64, max_options=len(LABELS))
@@ -180,7 +218,7 @@ class NextTokenStrategy(ReadStrategy):
         ctx = self.context
         read_request = self._read_request(chat_template_kwargs, cache_salt)
 
-        slots, engine_inputs = [], []
+        slots, engine_inputs, params = [], [], []
         for q in questions:
             messages = [{"role": "user", "content": f"{state}\n\n{q.type.prompt(q)}"}]
             if instructions:
@@ -192,15 +230,13 @@ class NextTokenStrategy(ReadStrategy):
                     "labels' tokens are unknown"
                 )
             slots.append([self.label_ids[label] for label in q.labels])
-            engine_inputs.append(engine_input)
+            engine_inputs.append(self._read_input(engine_input, prompt_ids))
+            params.append(self._sampling_params(slots[-1], prompt_ids))
 
         label_reads = await next_token_label_reads(
             ctx.engine_client,
             engine_inputs,
-            [
-                SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=ids)
-                for ids in slots
-            ],
+            params,
             request_id,
             lora_request=lora_request,
             trace_headers=trace_headers,
@@ -223,13 +259,87 @@ class NextTokenStrategy(ReadStrategy):
         return reads
 
 
-NEXT_TOKEN_ARCHITECTURES = frozenset(
-    {
-        "Qwen3ForCausalLM",
-        "Qwen3_5ForConditionalGeneration",
-        "Qwen3_5MoeForConditionalGeneration",
-    }
-)
+class DiffusionGemmaCanvasStrategy(NextTokenStrategy):
+    """DiffusionGemma. Each question is one read-only request with one denoising
+    step."""
+
+    THOUGHT = "<|channel>thought\n<channel|>"
+    END_OF_TURN = "<turn|>"
+    CANVAS_STEP = 16
+
+    def __init__(self, context: ReadContext):
+        ReadStrategy.__init__(self, context)
+        tokenizer = self._tokenizer()
+        self.tail = prompt_tail(tokenizer, self._generation_prompt())
+        self.thought = tokenizer.encode(self.THOUGHT, add_special_tokens=False)
+        self.label_ids = type_label_ids(tokenizer, self.thought)
+        end = tokenizer.encode(self.END_OF_TURN, add_special_tokens=False)
+        self.pad = tokenizer.pad_token_id
+        if self.pad is None or len(end) != 1:
+            raise ValueError(
+                "Structured decisions need the model's "
+                f"{self.END_OF_TURN!r} and pad tokens"
+            )
+        self.end = end[0]
+        vllm_config = context.engine_client.vllm_config
+        served = (
+            vllm_config.diffusion_config.canvas_length
+            if vllm_config.diffusion_config
+            and vllm_config.diffusion_config.canvas_length
+            else vllm_config.model_config.hf_config.canvas_length
+        )
+        self.width = self.CANVAS_STEP
+        if self.width > served:
+            raise ValueError(
+                f"Structured decisions need a canvas of {self.width}, "
+                f"and the served canvas is {served}"
+            )
+        self.vocab_size = vllm_config.model_config.get_vocab_size()
+        self.max_model_len = vllm_config.model_config.max_model_len
+
+    def _sampling_params(
+        self, label_ids: list[int], prompt_ids: list[int]
+    ) -> SamplingParams:
+        # The noise at the label slot is drawn from the prompt, so a repeated
+        # read gets the same canvas.
+        rng = random.Random(hashlib.sha256(json.dumps(prompt_ids).encode()).digest())
+        canvas = [rng.randrange(self.vocab_size), self.end]
+        canvas += [self.pad] * (self.width - len(canvas))
+        return SamplingParams(
+            max_tokens=2,
+            logprob_token_ids=label_ids,
+            extra_args={
+                "diffusion_seed_canvas": canvas,
+                "diffusion_canvas_length": self.width,
+                "diffusion_max_steps": 1,
+                "diffusion_read_only": True,
+            },
+        )
+
+    def _read_input(
+        self, engine_input: EngineInput, prompt_ids: list[int]
+    ) -> EngineInput:
+        extra = len(self.thought) + self.width
+        if len(prompt_ids) + extra > self.max_model_len:
+            raise StructuredDecisionError(
+                f"the prompt has {len(prompt_ids)} tokens and a canvas read adds "
+                f"{extra}, more than max_model_len={self.max_model_len}"
+            )
+        salt = engine_input.get("cache_salt")
+        # The thought block _must_ appear outside of the canvas (far worse results
+        # otherwise)
+        return tokens_input(
+            prompt_ids + self.thought,
+            cache_salt=salt if isinstance(salt, str) else None,
+        )
+
+
+READ_STRATEGIES: dict[str, type[ReadStrategy]] = {
+    "Qwen3ForCausalLM": NextTokenStrategy,
+    "Qwen3_5ForConditionalGeneration": NextTokenStrategy,
+    "Qwen3_5MoeForConditionalGeneration": NextTokenStrategy,
+    "DiffusionGemmaForBlockDiffusion": DiffusionGemmaCanvasStrategy,
+}
 
 
 # label_mass sums the labels' full-vocabulary probabilities. The logprobs modes
@@ -239,15 +349,16 @@ LOGPROBS_MODES = frozenset({"raw_logprobs", "processed_logprobs"})
 
 def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]:
     """Raise ValueError when the model cannot serve structured decisions."""
-    if model_config.architecture not in NEXT_TOKEN_ARCHITECTURES:
+    strategy = READ_STRATEGIES.get(model_config.architecture)
+    if strategy is None:
         raise ValueError(
             "The structured decisions API does not support "
             f"{model_config.architecture}. Supported architectures: "
-            f"{sorted(NEXT_TOKEN_ARCHITECTURES)}"
+            f"{sorted(READ_STRATEGIES)}"
         )
     if model_config.logprobs_mode not in LOGPROBS_MODES:
         raise ValueError(
             "Structured decisions need --logprobs-mode raw_logprobs or "
             f"processed_logprobs, not {model_config.logprobs_mode}"
         )
-    return NextTokenStrategy
+    return strategy