From 3c802af3cf6263d0ef6fb356a7884680ecf15f04 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Fri, 9 Oct 2026 19:33:22 -0600 Subject: [PATCH] [Feat][Frontend] structured decisions: DiffusionGemma read strategy (#60466) Signed-off-by: Matt Mastracci Signed-off-by: Simon Veitner Co-authored-by: Claude Opus 5.5 Co-authored-by: Simon Veitner Co-authored-by: Misha Goin --- .../online_serving/structured_decisions.md | 5 +- .../test_question_types.py | 2 +- .../structured_decisions/test_strategies.py | 37 +++- vllm/entrypoints/generate/api_router.py | 7 +- .../structured_decisions/question_types.py | 13 +- .../structured_decisions/strategies.py | 169 +++++++++++++++--- 6 files changed, 193 insertions(+), 40 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index a1f58a3769..a9ed5a2ee5 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/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 diff --git a/tests/entrypoints/generate/structured_decisions/test_question_types.py b/tests/entrypoints/generate/structured_decisions/test_question_types.py index 3b9ad7ba60..26047b5d6c 100644 --- a/tests/entrypoints/generate/structured_decisions/test_question_types.py +++ b/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) == ( diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 5dcb8edb30..a63e258752 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/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] + ) diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 67a2561924..18dfb87a26 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/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( diff --git a/vllm/entrypoints/generate/structured_decisions/question_types.py b/vllm/entrypoints/generate/structured_decisions/question_types.py index 0e80694b14..18e7a60018 100644 --- a/vllm/entrypoints/generate/structured_decisions/question_types.py +++ b/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 diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 65b5a9dd7b..f79375e42a 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/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" + END_OF_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