From cf9bfb9d27b210b566bb24ae041022282625ec75 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 16:34:06 -0600 Subject: [PATCH 01/25] [Model] DiffusionGemma: seeded canvas, step cap and read-only requests via extra_args SamplingParams.extra_args (vllm_xargs on the OpenAI server) gains three optional fields for structured reads on a DiffusionGemma model: diffusion_seed_canvas list[int], exactly canvas_length ids. Replaces the all-noise initial canvas after prefill; the caller puts its answer template in known positions and random ids in the slots it wants the model to fill. diffusion_max_steps int. Denoise steps before the request commits. diffusion_read_only bool. Commit the argmax canvas as soon as the step cap is reached instead of waiting for convergence, so a single-step request returns one forward's logprobs at every canvas position. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- vllm/model_executor/models/diffusion_gemma.py | 62 ++++++++++++++++++- 1 file changed, 59 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index aef9d816258b..7d0ecd10dac4 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -391,6 +391,7 @@ def _compiled_sample_step( normalizer: torch.Tensor, history: torch.Tensor, # [max_num_reqs, ST, CL] history_len_tensor: torch.Tensor, # [max_num_reqs] + max_steps_tensor: torch.Tensor, # [max_num_reqs] float, per-slot step cap # Output tensors (modified in-place) sampled: torch.Tensor, # [num_reqs, CL] num_sampled: torch.Tensor, # [num_reqs] @@ -523,7 +524,7 @@ def _compiled_sample_step( step_after = step_tensor[decode_slots] converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | ( - step_after >= max_denoising_steps + step_after.float() >= max_steps_tensor[decode_slots] ) # Commit done → denoise next (False); denoise converged → commit next (True) is_encoder_phase[decode_slots] = torch.where( @@ -628,6 +629,16 @@ def __init__( # Per-slot confidence flag, set by the sampler each step. self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + # Per-slot denoising step cap (structured reads run 1). Defaults to the + # global max; the sampler's add_request lowers it from extra_args. + self.max_steps = torch.full( + (max_num_reqs,), float(max_denoising_steps), dtype=torch.float32, device=device + ) + # Seeded initial canvases, by slot, applied when the prompt finishes. + self.seed_canvas: dict[int, list[int]] = {} + # Read-only slots finish on their converging step (no commit forward). + self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + # Per-slot self-conditioning soft embedding (probs @ embed_weight) from # the previous denoise step. Storing the [.., hidden] soft embed instead # of the full [.., vocab] distribution shrinks this buffer by @@ -657,11 +668,15 @@ def add_request(self, slot_idx: int) -> None: self.step[slot_idx].fill_(0) self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 + self.max_steps[slot_idx].fill_(float(self.max_denoising_steps)) + self.seed_canvas.pop(slot_idx, None) + self.read_only[slot_idx].fill_(False) def remove_request(self, slot_idx: int) -> None: self.is_encoder_phase[slot_idx].fill_(False) self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 + self.seed_canvas.pop(slot_idx, None) class DiffusionGemmaModelState(ModelState): @@ -1030,6 +1045,20 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None: self._pending_logprobs.pop(req_idx, None) self.sampling_states.add_request(req_idx, sampling_params) self.logprob_token_ids_state.add_request(req_idx, sampling_params) + extra = getattr(sampling_params, "extra_args", None) or {} + states = self.diffusion_states + cap = extra.get("diffusion_max_steps") + if cap is not None: + states.max_steps[req_idx].fill_(float(max(1, min(int(cap), states.max_denoising_steps)))) + canvas = extra.get("diffusion_seed_canvas") + if canvas is not None: + if len(canvas) != self.canvas_length: + raise ValueError( + f"diffusion_seed_canvas must hold exactly {self.canvas_length} ids, got {len(canvas)}" + ) + states.seed_canvas[req_idx] = [int(t) for t in canvas] + if extra.get("diffusion_read_only"): + states.read_only[req_idx].fill_(True) def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() @@ -1071,6 +1100,12 @@ def _finish_prefills( ps.astype(np.int64), device=states.is_encoder_phase.device ) states.init_canvas(ps_gpu) + for slot in ps.tolist(): + seed = states.seed_canvas.get(int(slot)) + if seed is not None: + states.canvas[slot] = torch.tensor( + seed, dtype=states.canvas.dtype, device=states.canvas.device + ) self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[ ps_gpu ] @@ -1263,6 +1298,7 @@ def __call__( self.normalizer, states.accepted_canvas_history, states.accepted_canvas_history_len, + states.max_steps, # Output sampled, num_sampled, @@ -1321,13 +1357,33 @@ def __call__( logits_mode=self.logits_mode, ) + # Read-only slots that converged this step: emit the argmax canvas now + # (as a commit would), keep them out of the encoder phase, and hand out + # their logprobs below. The request ends on these tokens. + emit_now: set[int] = set() + if num_decode > 0: + ro_mask = states.read_only[decode_slots] & states.is_encoder_phase[ + decode_slots + ] & ~is_committing + if bool(ro_mask.any()): + ro_idx = decode_idx[ro_mask] + ro_slots = decode_slots[ro_mask] + sampled[ro_idx] = states.argmax_canvas[ro_slots].to(sampled.dtype) + num_sampled[ro_idx] = valid_canvas_len[ro_mask].to(num_sampled.dtype) + states.is_encoder_phase[ro_slots] = False + emit_now = set(ro_slots.tolist()) + # Commit steps: is_committing was True at entry. Reassemble previously # stashed logprobs and attach to SamplerOutput. logprobs_tensors = None - if want_logprobs and is_committing.any() and self._pending_logprobs: + if ( + want_logprobs + and (emit_now or bool(is_committing.any())) + and self._pending_logprobs + ): committing_slots = set( decode_slots_np[is_committing.cpu().numpy()].tolist() - ) + ) | emit_now parts_ids, parts_lp, parts_ranks = [], [], [] cu_gen: list[int] = [] flat_offset = 0 From cf27f5dadada9875bf9dbde034f204c8ab3c0aab Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 16:58:41 -0600 Subject: [PATCH 02/25] [Core] Validate diffusion extra_args when the request is created The sampler read diffusion_seed_canvas, diffusion_max_steps and diffusion_read_only straight out of extra_args on the GPU worker. A seed id outside the vocabulary reached the embedding gather as a device-side assert, and a wrong type or length raised inside the engine loop, both of which take the engine down rather than the one request. SamplingParams.verify now checks the three fields and returns a 400. The input processor passes the served DiffusionConfig so the seed length is checked against the configured canvas, not the model's default. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/test_sampling_params.py | 57 +++++++++++++ .../test_input_processor_trace_replay.py | 1 + vllm/sampling_params.py | 79 ++++++++++++++++++- vllm/v1/engine/input_processor.py | 2 + 4 files changed, 136 insertions(+), 3 deletions(-) diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index 3c909f4969df..024cfc535122 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -59,6 +59,63 @@ def test_verify_leaves_logits_processors_to_admission(): ) +@dataclass +class MockDiffusionConfig: + canvas_length: int = 8 + + +def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None): + diffusion_config = ( + MockDiffusionConfig(canvas_length) if canvas_length is not None else None + ) + params.verify( + MockModelConfig(is_diffusion=True), + None, + None, + None, + diffusion_config=diffusion_config, + ) + + +@pytest.mark.parametrize( + "extra_args, match", + [ + ({"diffusion_seed_canvas": "abc"}, "list of token ids"), + ({"diffusion_seed_canvas": [1, 2.5]}, "list of token ids"), + ({"diffusion_seed_canvas": [1, True]}, "list of token ids"), + ({"diffusion_seed_canvas": [0, 1024]}, r"in \[0, 1024\)"), + ({"diffusion_seed_canvas": [0, -1]}, r"in \[0, 1024\)"), + ({"diffusion_max_steps": 0}, "positive integer"), + ({"diffusion_max_steps": "1"}, "positive integer"), + ({"diffusion_max_steps": True}, "positive integer"), + ({"diffusion_read_only": "yes"}, "boolean"), + ({"diffusion_read_only": 2}, "boolean"), + ], +) +def test_diffusion_rejects_bad_extra_args(extra_args: dict, match: str): + with pytest.raises(VLLMValidationError, match=match): + _verify_diffusion(SamplingParams(extra_args=extra_args)) + + +def test_diffusion_seed_canvas_must_fill_the_canvas(): + params = SamplingParams(extra_args={"diffusion_seed_canvas": [0] * 7}) + with pytest.raises(VLLMValidationError, match="exactly 8 ids, got 7"): + _verify_diffusion(params, canvas_length=8) + # Without the served diffusion config the canvas length is unknown. + _verify_diffusion(params) + + +def test_diffusion_accepts_extra_args(): + params = SamplingParams( + extra_args={ + "diffusion_seed_canvas": list(range(8)), + "diffusion_max_steps": 1, + "diffusion_read_only": True, + } + ) + _verify_diffusion(params, canvas_length=8) + + @pytest.mark.parametrize("value", [-(2**63) - 1, 2**64]) def test_extra_args_rejects_nested_integer_overflow(value): """Reject extension values before they reach the engine transport.""" diff --git a/tests/v1/engine/test_input_processor_trace_replay.py b/tests/v1/engine/test_input_processor_trace_replay.py index 24302f6de616..5fc1118167e5 100644 --- a/tests/v1/engine/test_input_processor_trace_replay.py +++ b/tests/v1/engine/test_input_processor_trace_replay.py @@ -84,6 +84,7 @@ def _validate(enable_trace_replay: bool) -> None: vllm_config=SimpleNamespace(reasoning_config=None), speculative_config=None, structured_outputs_config=None, + diffusion_config=None, tokenizer=None, validate_logits_processors_params=lambda params: None, ) diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 120dbaddfd80..715e37bc9701 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -15,7 +15,12 @@ from pydantic.dataclasses import dataclass import vllm.envs as envs -from vllm.config import ModelConfig, SpeculativeConfig, StructuredOutputsConfig +from vllm.config import ( + DiffusionConfig, + ModelConfig, + SpeculativeConfig, + StructuredOutputsConfig, +) from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike @@ -818,6 +823,7 @@ def verify( speculative_config: SpeculativeConfig | None, structured_outputs_config: StructuredOutputsConfig | None, tokenizer: TokenizerLike | None, + diffusion_config: DiffusionConfig | None = None, ) -> None: self._validate_logprobs(model_config) self._validate_logit_bias(model_config) @@ -825,7 +831,7 @@ def verify( self._validate_stop_token_ids(model_config) self._validate_allowed_token_ids(model_config) self._validate_spec_decode(speculative_config) - self._validate_diffusion(model_config) + self._validate_diffusion(model_config, diffusion_config) self._validate_structured_outputs( model_config, structured_outputs_config, tokenizer ) @@ -1037,7 +1043,11 @@ def _validate_spec_decode( "are not yet supported with speculative decoding." ) - def _validate_diffusion(self, model_config: ModelConfig) -> None: + def _validate_diffusion( + self, + model_config: ModelConfig, + diffusion_config: DiffusionConfig | None = None, + ) -> None: if not model_config.is_diffusion: return @@ -1058,6 +1068,69 @@ def _validate_diffusion(self, model_config: ModelConfig) -> None: "bad_words, and allowed_token_ids sampling parameters " "are not yet supported with diffusion models." ) + self._validate_diffusion_extra_args(model_config, diffusion_config) + + def _validate_diffusion_extra_args( + self, + model_config: ModelConfig, + diffusion_config: DiffusionConfig | None, + ) -> None: + """Check the per-request diffusion fields carried in ``extra_args``. + + ``diffusion_seed_canvas`` (list of token ids, one per canvas position) + replaces the random initial canvas once the prompt is prefilled, + ``diffusion_max_steps`` (int) caps the denoising steps per canvas, and + ``diffusion_read_only`` (bool) ends the request on the first canvas it + converges. The sampler consumes these on the GPU, where an id outside + the vocabulary is a device-side assert, so they are checked here. + """ + extra = self.extra_args + if not extra: + return + + seed = extra.get("diffusion_seed_canvas") + if seed is not None: + if not isinstance(seed, (list, tuple)) or not all( + isinstance(t, int) and not isinstance(t, bool) for t in seed + ): + raise VLLMValidationError( + "diffusion_seed_canvas must be a list of token ids.", + parameter="extra_args", + ) + vocab_size = model_config.get_vocab_size() + if any(t < 0 or t >= vocab_size for t in seed): + raise VLLMValidationError( + f"diffusion_seed_canvas ids must be in [0, {vocab_size}).", + parameter="extra_args", + ) + if ( + diffusion_config is not None + and len(seed) != diffusion_config.canvas_length + ): + raise VLLMValidationError( + "diffusion_seed_canvas must hold exactly " + f"{diffusion_config.canvas_length} ids, got {len(seed)}.", + parameter="extra_args", + ) + + cap = extra.get("diffusion_max_steps") + if cap is not None and ( + not isinstance(cap, int) or isinstance(cap, bool) or cap < 1 + ): + raise VLLMValidationError( + "diffusion_max_steps must be a positive integer.", + parameter="extra_args", + ) + + # The OpenAI server's vllm_xargs narrows JSON booleans to 0/1. + read_only = extra.get("diffusion_read_only") + if read_only is not None and ( + not isinstance(read_only, (bool, int)) or read_only not in (0, 1) + ): + raise VLLMValidationError( + "diffusion_read_only must be a boolean (or 0/1).", + parameter="extra_args", + ) def _validate_structured_outputs( self, diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 582e1abac685..c3abec1b41be 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -53,6 +53,7 @@ def __init__( self.speculative_config = vllm_config.speculative_config self.structured_outputs_config = vllm_config.structured_outputs_config self.observability_config = vllm_config.observability_config + self.diffusion_config = vllm_config.diffusion_config # Load the custom logits processor classes once; the returned callable # runs their validate_params hooks per request at admission. self.validate_logits_processors_params = ( @@ -125,6 +126,7 @@ def _validate_params( self.speculative_config, self.structured_outputs_config, self.tokenizer, + diffusion_config=self.diffusion_config, ) self.validate_logits_processors_params(params) From 15eb9be76f0397e2c153ac049108e8130f1d496a Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 16:58:58 -0600 Subject: [PATCH 03/25] [Core] Read-only diffusion requests end on the canvas they emit A read-only request emitted its argmax canvas and then, unless the client had set max_tokens at or below the canvas length and ignore_eos, rolled into a second block from fresh noise with the seed gone. The request now carries that contract itself: verify caps max_tokens at the served canvas length and sets ignore_eos, so the scheduler ends the request on the first canvas and an end-of-turn token drawn into a noise slot cannot cut it short. A smaller client max_tokens still trims the output. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/test_sampling_params.py | 11 +++++++++++ vllm/sampling_params.py | 8 ++++++++ 2 files changed, 19 insertions(+) diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index 024cfc535122..a66ed1d49ba9 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -105,6 +105,17 @@ def test_diffusion_seed_canvas_must_fill_the_canvas(): _verify_diffusion(params) +@pytest.mark.parametrize("max_tokens, expected", [(100, 8), (5, 5), (None, 8)]) +@pytest.mark.parametrize("flag", [True, 1]) +def test_diffusion_read_only_ends_after_one_canvas(max_tokens, expected, flag): + params = SamplingParams( + max_tokens=max_tokens, extra_args={"diffusion_read_only": flag} + ) + _verify_diffusion(params, canvas_length=8) + assert params.max_tokens == expected + assert params.ignore_eos + + def test_diffusion_accepts_extra_args(): params = SamplingParams( extra_args={ diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 715e37bc9701..c7a865184248 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1131,6 +1131,14 @@ def _validate_diffusion_extra_args( "diffusion_read_only must be a boolean (or 0/1).", parameter="extra_args", ) + if read_only: + # One canvas is the whole output. Cap max_tokens so the scheduler + # ends the request on it, and ignore EOS so a noise draw that lands + # an end-of-turn token in the canvas cannot cut it short. + if diffusion_config is not None: + canvas_length = diffusion_config.canvas_length + self.max_tokens = min(self.max_tokens or canvas_length, canvas_length) + self.ignore_eos = True def _validate_structured_outputs( self, From a2908b2353466ad7fcad4d70488513284e38739c Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 17:00:25 -0600 Subject: [PATCH 04/25] [Perf] Keep seed and read-only bookkeeping off the diffusion hot path Every denoise step ran a GPU sync for the read-only mask (`.any()` and `.tolist()`) whether or not any read-only request was in the batch, and each finished prefill uploaded its seed canvas in a Python loop, one host-to-device copy per slot. DiffusionGemmaRequestStates now keeps host-side sets of the seeded and read-only slots next to the GPU flags. The sampler checks the sets first, so ordinary generation pays nothing. Seeds live in a [max_num_reqs, canvas] tensor uploaded once at add_request and applied to a whole prefill batch with one masked copy. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/v1/sample/test_diffusion_gemma_reads.py | 95 +++++++++++++++++++ vllm/model_executor/models/diffusion_gemma.py | 64 +++++++++---- 2 files changed, 142 insertions(+), 17 deletions(-) create mode 100644 tests/v1/sample/test_diffusion_gemma_reads.py diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py new file mode 100644 index 000000000000..9cde7214f967 --- /dev/null +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-request DiffusionGemma state behind structured reads: seed canvases, +read-only slots and the per-slot step cap.""" + +import numpy as np +import pytest +import torch + +from vllm.model_executor.models.diffusion_gemma import DiffusionGemmaRequestStates +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="the sampler state lives on the GPU" +) + +CL = 8 +VOCAB = 64 +MAX_REQS = 4 +MAX_STEPS = 48 + + +def _states() -> DiffusionGemmaRequestStates: + return DiffusionGemmaRequestStates( + max_num_reqs=MAX_REQS, + canvas_length=CL, + vocab_size=VOCAB, + max_denoising_steps=MAX_STEPS, + device=torch.device("cuda"), + hidden_size=4, + stability_threshold=2, + ) + + +def _slots(*idx: int) -> tuple[np.ndarray, torch.Tensor]: + slots = np.array(idx, dtype=np.int64) + return slots, torch.tensor(slots, device="cuda") + + +def test_seed_canvas_replaces_only_seeded_slots(): + states = _states() + for slot in range(3): + states.add_request(slot) + seed = list(range(CL)) + states.set_seed_canvas(1, seed) + slots, slots_gpu = _slots(0, 1, 2) + states.init_canvas(slots_gpu) + before = states.canvas[slots_gpu].clone() + + states.apply_seed_canvases(slots, slots_gpu) + + after = states.canvas[slots_gpu] + assert after[1].tolist() == seed + assert torch.equal(after[0], before[0]) + assert torch.equal(after[2], before[2]) + + +def test_apply_seed_canvases_leaves_unseeded_batches_alone(): + states = _states() + states.add_request(0) + slots, slots_gpu = _slots(0) + states.init_canvas(slots_gpu) + before = states.canvas[0].clone() + + states.apply_seed_canvases(slots, slots_gpu) + + assert torch.equal(states.canvas[0], before) + + +def test_add_request_clears_seed_and_read_only(): + states = _states() + states.add_request(0) + states.set_seed_canvas(0, [1] * CL) + states.set_read_only(0) + assert states.seeded_slots == {0} + assert states.read_only_slots == {0} + + states.add_request(0) + + assert not states.seeded_slots + assert not states.read_only_slots + assert not bool(states.has_seed[0]) + assert not bool(states.read_only[0]) + + +def test_remove_request_forgets_the_slot(): + states = _states() + states.add_request(0) + states.set_seed_canvas(0, [1] * CL) + states.set_read_only(0) + + states.remove_request(0) + + assert not states.seeded_slots + assert not states.read_only_slots diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 7d0ecd10dac4..0342f77f4be3 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -634,10 +634,17 @@ def __init__( self.max_steps = torch.full( (max_num_reqs,), float(max_denoising_steps), dtype=torch.float32, device=device ) - # Seeded initial canvases, by slot, applied when the prompt finishes. - self.seed_canvas: dict[int, list[int]] = {} + # A seed canvas replaces the random initial canvas once the prompt is + # prefilled. The host-side slot sets let the sampler skip the seed and + # read-only work when no live request asked for it. + self.seed_canvas = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.int64, device=device + ) + self.has_seed = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + self.seeded_slots: set[int] = set() # Read-only slots finish on their converging step (no commit forward). self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + self.read_only_slots: set[int] = set() # Per-slot self-conditioning soft embedding (probs @ embed_weight) from # the previous denoise step. Storing the [.., hidden] soft embed instead @@ -669,14 +676,40 @@ def add_request(self, slot_idx: int) -> None: self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 self.max_steps[slot_idx].fill_(float(self.max_denoising_steps)) - self.seed_canvas.pop(slot_idx, None) + self.has_seed[slot_idx].fill_(False) + self.seeded_slots.discard(slot_idx) self.read_only[slot_idx].fill_(False) + self.read_only_slots.discard(slot_idx) def remove_request(self, slot_idx: int) -> None: self.is_encoder_phase[slot_idx].fill_(False) self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 - self.seed_canvas.pop(slot_idx, None) + self.seeded_slots.discard(slot_idx) + self.read_only_slots.discard(slot_idx) + + def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None: + self.seed_canvas[slot_idx] = async_tensor_h2d( + ids, dtype=torch.int64, device=self.device + ) + self.has_seed[slot_idx].fill_(True) + self.seeded_slots.add(slot_idx) + + def set_read_only(self, slot_idx: int) -> None: + self.read_only[slot_idx].fill_(True) + self.read_only_slots.add(slot_idx) + + def apply_seed_canvases( + self, slots_np: np.ndarray, slots_gpu: torch.Tensor + ) -> None: + """Replace the canvas of every seeded slot among ``slots_gpu``.""" + if self.seeded_slots.isdisjoint(slots_np.tolist()): + return + self.canvas[slots_gpu] = torch.where( + self.has_seed[slots_gpu, None], + self.seed_canvas[slots_gpu], + self.canvas[slots_gpu], + ) class DiffusionGemmaModelState(ModelState): @@ -1050,15 +1083,15 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None: cap = extra.get("diffusion_max_steps") if cap is not None: states.max_steps[req_idx].fill_(float(max(1, min(int(cap), states.max_denoising_steps)))) - canvas = extra.get("diffusion_seed_canvas") - if canvas is not None: - if len(canvas) != self.canvas_length: + seed = extra.get("diffusion_seed_canvas") + if seed is not None: + if len(seed) != self.canvas_length: raise ValueError( - f"diffusion_seed_canvas must hold exactly {self.canvas_length} ids, got {len(canvas)}" + f"diffusion_seed_canvas must hold exactly {self.canvas_length} ids, got {len(seed)}" ) - states.seed_canvas[req_idx] = [int(t) for t in canvas] + states.set_seed_canvas(req_idx, seed) if extra.get("diffusion_read_only"): - states.read_only[req_idx].fill_(True) + states.set_read_only(req_idx) def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() @@ -1100,12 +1133,7 @@ def _finish_prefills( ps.astype(np.int64), device=states.is_encoder_phase.device ) states.init_canvas(ps_gpu) - for slot in ps.tolist(): - seed = states.seed_canvas.get(int(slot)) - if seed is not None: - states.canvas[slot] = torch.tensor( - seed, dtype=states.canvas.dtype, device=states.canvas.device - ) + states.apply_seed_canvases(ps, ps_gpu) self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[ ps_gpu ] @@ -1361,7 +1389,9 @@ def __call__( # (as a commit would), keep them out of the encoder phase, and hand out # their logprobs below. The request ends on these tokens. emit_now: set[int] = set() - if num_decode > 0: + if states.read_only_slots and not states.read_only_slots.isdisjoint( + decode_slots_np.tolist() + ): ro_mask = states.read_only[decode_slots] & states.is_encoder_phase[ decode_slots ] & ~is_committing From fa81b63406a6eac927f27240afdcee84d9eae750 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 17:00:43 -0600 Subject: [PATCH 05/25] [Core] Read-only diffusion requests report temperature-1 logprobs The logprobs stashed on the converging step came from the sampler's schedule-tempered logits (temperature 0.8 on a first step), so a caller after calibrated probabilities had to know the schedule and rescale. A read-only request now stashes logprobs from the model's own logits at temperature 1. The argmax, and so the emitted canvas, is unchanged. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- vllm/model_executor/models/diffusion_gemma.py | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 0342f77f4be3..e676ce8a2dee 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -1354,17 +1354,26 @@ def __call__( if just_converged.any(): flat_logits = scaled.reshape(-1, scaled.shape[-1]) argmax_tokens = scaled.argmax(dim=-1) + raw_flat: torch.Tensor | None = None for local_idx in just_converged.nonzero(as_tuple=True)[0]: li = local_idx.item() - slot = tile_slots[local_idx] + slot = tile_slots[local_idx].item() # Stash only the real canvas positions (== CL unless this # canvas was truncated near max_model_len); padded tail # positions are never emitted. k_i = int(valid_canvas_len_np[start_req + li]) pos = li * CL + src = flat_logits + if slot in states.read_only_slots: + # A read reports the model's own distribution at + # temperature 1, not the schedule-tempered one the + # sampler draws from. The argmax is the same. + if raw_flat is None: + raw_flat = logits[start_req * CL : end_req * CL].float() + src = raw_flat per_req_ids = max_token_ids > 0 - self._pending_logprobs[slot.item()] = compute_topk_scores( - flat_logits[pos : pos + k_i], + self._pending_logprobs[slot] = compute_topk_scores( + src[pos : pos + k_i], num_logprobs, argmax_tokens[local_idx][:k_i], logprob_token_ids_state=( @@ -1372,12 +1381,7 @@ def __call__( ), # every row of this stash belongs to one slot expanded_idx_mapping=( - torch.full( - (k_i,), - slot.item(), - dtype=torch.int32, - device=flat_logits.device, - ) + torch.full((k_i,), slot, dtype=torch.int32, device=device) if per_req_ids else None ), From 21562ccfe18795cc2c0fd540c44ba6abf8db346e Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 17:01:24 -0600 Subject: [PATCH 06/25] [Misc] Keep the diffusion step cap in the step counter's dtype The per-slot cap was a float tensor compared against a float cast of the int32 step counter inside the compiled step. Both are now int32. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/v1/sample/test_diffusion_gemma_reads.py | 61 ++++++++++++++++++- vllm/model_executor/models/diffusion_gemma.py | 12 ++-- 2 files changed, 67 insertions(+), 6 deletions(-) diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py index 9cde7214f967..7699837ef885 100644 --- a/tests/v1/sample/test_diffusion_gemma_reads.py +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -7,7 +7,10 @@ import pytest import torch -from vllm.model_executor.models.diffusion_gemma import DiffusionGemmaRequestStates +from vllm.model_executor.models.diffusion_gemma import ( + DiffusionGemmaRequestStates, + _compiled_sample_step, +) from vllm.platforms import current_platform pytestmark = pytest.mark.skipif( @@ -93,3 +96,59 @@ def test_remove_request_forgets_the_slot(): assert not states.seeded_slots assert not states.read_only_slots + + +def _denoise_once(states: DiffusionGemmaRequestStates, slots: list[int]) -> None: + """One compiled denoise step over ``slots`` with flat logits, so nothing + converges by stability or confidence and only the step cap can end it.""" + n = len(slots) + device = states.device + decode_slots = torch.tensor(slots, dtype=torch.int64, device=device) + decode_idx = torch.arange(n, dtype=torch.int64, device=device) + _compiled_sample_step( + torch.zeros(n * CL, VOCAB, device=device), + decode_slots, + decode_idx, + decode_slots, + torch.full((n,), CL, dtype=torch.int64, device=device), + states.canvas, + states.argmax_canvas, + states.step, + states.is_encoder_phase, + states.confident, + states.self_conditioning_embeds, + torch.zeros(VOCAB, 4, device=device), + torch.tensor(1.0, device=device), + states.accepted_canvas_history, + states.accepted_canvas_history_len, + states.max_steps, + torch.zeros(n, CL, dtype=torch.int32, device=device), + torch.zeros(n, dtype=torch.int32, device=device), + torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device), + max_denoising_steps=float(MAX_STEPS), + t_min=0.5, + t_max=1.0, + confidence_threshold=0.1, + vocab_size=VOCAB, + CL=CL, + ST=states.stability_threshold, + entropy_bound=0.1, + sc_vocab_start=0, + sc_vocab_end=VOCAB, + tp_size=1, + tp_group_name="", + ) + + +def test_step_cap_is_per_slot(): + states = _states() + for slot in (0, 1): + states.add_request(slot) + states.is_encoder_phase[slot] = False + states.max_steps[0] = 1 + + _denoise_once(states, [0, 1]) + + # Slot 0 hit its cap and moves to commit; slot 1 keeps denoising. + assert states.is_encoder_phase[:2].tolist() == [True, False] + assert states.step[:2].tolist() == [1, 1] diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index e676ce8a2dee..aa8a8ce6ec0d 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -391,7 +391,7 @@ def _compiled_sample_step( normalizer: torch.Tensor, history: torch.Tensor, # [max_num_reqs, ST, CL] history_len_tensor: torch.Tensor, # [max_num_reqs] - max_steps_tensor: torch.Tensor, # [max_num_reqs] float, per-slot step cap + max_steps_tensor: torch.Tensor, # [max_num_reqs] int32, per-slot step cap # Output tensors (modified in-place) sampled: torch.Tensor, # [num_reqs, CL] num_sampled: torch.Tensor, # [num_reqs] @@ -524,7 +524,7 @@ def _compiled_sample_step( step_after = step_tensor[decode_slots] converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | ( - step_after.float() >= max_steps_tensor[decode_slots] + step_after >= max_steps_tensor[decode_slots] ) # Commit done → denoise next (False); denoise converged → commit next (True) is_encoder_phase[decode_slots] = torch.where( @@ -632,7 +632,7 @@ def __init__( # Per-slot denoising step cap (structured reads run 1). Defaults to the # global max; the sampler's add_request lowers it from extra_args. self.max_steps = torch.full( - (max_num_reqs,), float(max_denoising_steps), dtype=torch.float32, device=device + (max_num_reqs,), max_denoising_steps, dtype=torch.int32, device=device ) # A seed canvas replaces the random initial canvas once the prompt is # prefilled. The host-side slot sets let the sampler skip the seed and @@ -675,7 +675,7 @@ def add_request(self, slot_idx: int) -> None: self.step[slot_idx].fill_(0) self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 - self.max_steps[slot_idx].fill_(float(self.max_denoising_steps)) + self.max_steps[slot_idx].fill_(self.max_denoising_steps) self.has_seed[slot_idx].fill_(False) self.seeded_slots.discard(slot_idx) self.read_only[slot_idx].fill_(False) @@ -1082,7 +1082,9 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None: states = self.diffusion_states cap = extra.get("diffusion_max_steps") if cap is not None: - states.max_steps[req_idx].fill_(float(max(1, min(int(cap), states.max_denoising_steps)))) + states.max_steps[req_idx].fill_( + max(1, min(int(cap), states.max_denoising_steps)) + ) seed = extra.get("diffusion_seed_canvas") if seed is not None: if len(seed) != self.canvas_length: From da4d83c784e7cb976efb076431e689344ec2521f Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 17:01:39 -0600 Subject: [PATCH 07/25] [Misc] Drop the redundant commit term from the read-only mask A read-only slot never enters the encoder phase, so `~is_committing` was always true for it. Also state why remove_request leaves the GPU flags to add_request. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- vllm/model_executor/models/diffusion_gemma.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index aa8a8ce6ec0d..beae4b26416a 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -682,6 +682,9 @@ def add_request(self, slot_idx: int) -> None: self.read_only_slots.discard(slot_idx) def remove_request(self, slot_idx: int) -> None: + # The GPU flags (max_steps, has_seed, read_only) are reset by + # add_request before the slot is reused; only the host sets that gate + # the sampler's work need clearing now. self.is_encoder_phase[slot_idx].fill_(False) self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 @@ -1398,9 +1401,11 @@ def __call__( if states.read_only_slots and not states.read_only_slots.isdisjoint( decode_slots_np.tolist() ): + # A read-only slot never enters the encoder phase, so the flag + # being set here means it converged on this step. ro_mask = states.read_only[decode_slots] & states.is_encoder_phase[ decode_slots - ] & ~is_committing + ] if bool(ro_mask.any()): ro_idx = decode_idx[ro_mask] ro_slots = decode_slots[ro_mask] From 5cc9b4f822ca3ad7d84821310b118f5409fba5ea Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 17:45:33 -0600 Subject: [PATCH 08/25] [Misc] Tighten the comments in the structured-read changes Comments say the why once, in the active voice, without asides. Two long lines are rewrapped. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/v1/sample/test_diffusion_gemma_reads.py | 2 +- vllm/model_executor/models/diffusion_gemma.py | 36 +++++++++---------- vllm/sampling_params.py | 23 ++++++------ 3 files changed, 30 insertions(+), 31 deletions(-) diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py index 7699837ef885..c48ac2c5e44a 100644 --- a/tests/v1/sample/test_diffusion_gemma_reads.py +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -149,6 +149,6 @@ def test_step_cap_is_per_slot(): _denoise_once(states, [0, 1]) - # Slot 0 hit its cap and moves to commit; slot 1 keeps denoising. + # Slot 0 hit its cap and moves to commit. Slot 1 keeps denoising. assert states.is_encoder_phase[:2].tolist() == [True, False] assert states.step[:2].tolist() == [1, 1] diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index beae4b26416a..f2069ce90841 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -629,20 +629,19 @@ def __init__( # Per-slot confidence flag, set by the sampler each step. self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) - # Per-slot denoising step cap (structured reads run 1). Defaults to the - # global max; the sampler's add_request lowers it from extra_args. + # Per-slot step cap, lowered per request from extra_args. self.max_steps = torch.full( (max_num_reqs,), max_denoising_steps, dtype=torch.int32, device=device ) - # A seed canvas replaces the random initial canvas once the prompt is - # prefilled. The host-side slot sets let the sampler skip the seed and - # read-only work when no live request asked for it. + # Seed canvases replace the random canvas after prefill. The host-side + # sets gate the seed and read-only work so plain generation skips it. self.seed_canvas = torch.zeros( max_num_reqs, canvas_length, dtype=torch.int64, device=device ) self.has_seed = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) self.seeded_slots: set[int] = set() - # Read-only slots finish on their converging step (no commit forward). + # Read-only slots emit on their converging step and skip the commit + # forward. self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) self.read_only_slots: set[int] = set() @@ -682,9 +681,8 @@ def add_request(self, slot_idx: int) -> None: self.read_only_slots.discard(slot_idx) def remove_request(self, slot_idx: int) -> None: - # The GPU flags (max_steps, has_seed, read_only) are reset by - # add_request before the slot is reused; only the host sets that gate - # the sampler's work need clearing now. + # add_request resets the GPU flags before a slot is reused. The host + # sets gate work on every step, so they clear here. self.is_encoder_phase[slot_idx].fill_(False) self.accepted_canvas_history_len[slot_idx].fill_(0) self.self_conditioning_embeds[slot_idx] = 0 @@ -1092,7 +1090,8 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None: if seed is not None: if len(seed) != self.canvas_length: raise ValueError( - f"diffusion_seed_canvas must hold exactly {self.canvas_length} ids, got {len(seed)}" + "diffusion_seed_canvas must hold exactly " + f"{self.canvas_length} ids, got {len(seed)}" ) states.set_seed_canvas(req_idx, seed) if extra.get("diffusion_read_only"): @@ -1370,11 +1369,11 @@ def __call__( pos = li * CL src = flat_logits if slot in states.read_only_slots: - # A read reports the model's own distribution at - # temperature 1, not the schedule-tempered one the - # sampler draws from. The argmax is the same. + # Read-only slots report logprobs at temperature 1. + # The schedule-tempered logits share the argmax. if raw_flat is None: - raw_flat = logits[start_req * CL : end_req * CL].float() + tile_rows = slice(start_req * CL, end_req * CL) + raw_flat = logits[tile_rows].float() src = raw_flat per_req_ids = max_token_ids > 0 self._pending_logprobs[slot] = compute_topk_scores( @@ -1394,15 +1393,14 @@ def __call__( logits_mode=self.logits_mode, ) - # Read-only slots that converged this step: emit the argmax canvas now - # (as a commit would), keep them out of the encoder phase, and hand out - # their logprobs below. The request ends on these tokens. + # Read-only slots that converged this step emit their argmax canvas + # now, skip the encoder phase, and hand out their logprobs below. emit_now: set[int] = set() if states.read_only_slots and not states.read_only_slots.isdisjoint( decode_slots_np.tolist() ): - # A read-only slot never enters the encoder phase, so the flag - # being set here means it converged on this step. + # Read-only slots never enter the encoder phase, so the flag here + # means converged this step. ro_mask = states.read_only[decode_slots] & states.is_encoder_phase[ decode_slots ] diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index c7a865184248..2ccae5100488 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1075,14 +1075,13 @@ def _validate_diffusion_extra_args( model_config: ModelConfig, diffusion_config: DiffusionConfig | None, ) -> None: - """Check the per-request diffusion fields carried in ``extra_args``. - - ``diffusion_seed_canvas`` (list of token ids, one per canvas position) - replaces the random initial canvas once the prompt is prefilled, - ``diffusion_max_steps`` (int) caps the denoising steps per canvas, and - ``diffusion_read_only`` (bool) ends the request on the first canvas it - converges. The sampler consumes these on the GPU, where an id outside - the vocabulary is a device-side assert, so they are checked here. + """Reject bad diffusion fields in ``extra_args`` before the worker. + + On the GPU a seed id outside the vocabulary is a device-side assert. + ``diffusion_seed_canvas``: token ids, one per canvas position, that + replace the random canvas after prefill. ``diffusion_max_steps``: + denoise steps per canvas. ``diffusion_read_only``: end the request + on the first canvas that converges. """ extra = self.extra_args if not extra: @@ -1133,11 +1132,13 @@ def _validate_diffusion_extra_args( ) if read_only: # One canvas is the whole output. Cap max_tokens so the scheduler - # ends the request on it, and ignore EOS so a noise draw that lands - # an end-of-turn token in the canvas cannot cut it short. + # ends the request there, and ignore EOS so an end-of-turn token + # drawn into a noise slot does not end it early. if diffusion_config is not None: canvas_length = diffusion_config.canvas_length - self.max_tokens = min(self.max_tokens or canvas_length, canvas_length) + self.max_tokens = min( + self.max_tokens or canvas_length, canvas_length + ) self.ignore_eos = True def _validate_structured_outputs( From 0c64f9fedb4355ae701a6e0122da5e3df26f33fd Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 19:12:50 -0600 Subject: [PATCH 09/25] [Model] DiffusionGemma: skip self-conditioning for one-step tiles The compiled step computes the next step's self-conditioning soft embed, probs @ embed_weight, for every tile. That is a full pass over the vocabulary matrix (6.5 ms on a GB10) that no slot capped at one denoise step ever reads. The sampler keeps a host-side set of one-step slots and passes a flag that selects the variant without the matmul when a whole tile is in it. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/v1/sample/test_diffusion_gemma_reads.py | 16 +++++- vllm/model_executor/models/diffusion_gemma.py | 49 +++++++++++++------ 2 files changed, 49 insertions(+), 16 deletions(-) diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py index c48ac2c5e44a..28f92e4764e2 100644 --- a/tests/v1/sample/test_diffusion_gemma_reads.py +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -98,7 +98,9 @@ def test_remove_request_forgets_the_slot(): assert not states.read_only_slots -def _denoise_once(states: DiffusionGemmaRequestStates, slots: list[int]) -> None: +def _denoise_once( + states: DiffusionGemmaRequestStates, slots: list[int], compute_sc: bool = True +) -> None: """One compiled denoise step over ``slots`` with flat logits, so nothing converges by stability or confidence and only the step cap can end it.""" n = len(slots) @@ -137,9 +139,21 @@ def _denoise_once(states: DiffusionGemmaRequestStates, slots: list[int]) -> None sc_vocab_end=VOCAB, tp_size=1, tp_group_name="", + compute_sc=compute_sc, ) +def test_single_step_tile_skips_self_conditioning(): + states = _states() + states.add_request(0) + states.is_encoder_phase[0] = False + states.self_conditioning_embeds[0] = 1.0 + + _denoise_once(states, [0], compute_sc=False) + + assert not states.self_conditioning_embeds[0].any() + + def test_step_cap_is_per_slot(): states = _states() for slot in (0, 1): diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index f2069ce90841..6b3d39b6a770 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -413,6 +413,7 @@ def _compiled_sample_step( sc_vocab_end: int, tp_size: int, tp_group_name: str, + compute_sc: bool = True, ) -> torch.Tensor: """Compiled decode step: temperature → Gumbel sample → probs/confidence → accept/renoise → convergence, all as vectorized PyTorch ops. @@ -538,18 +539,27 @@ def _compiled_sample_step( # sc_embeds directly. Storing the [.., hidden] soft embed instead of the full # [.., vocab] probs avoids a giant persistent buffer. sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None] - # Self-conditioning soft embed = probs @ embed_tokens.weight. Under tensor - # parallelism the embedding is vocab-sharded ([vocab/tp, hidden]) while - # probs spans the full vocab, so each rank multiplies its local vocab slice - # [sc_vocab_start, sc_vocab_end) and the partials are summed across ranks. - local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype) - soft_embeds = torch.matmul( - local_probs, embed_weight[: sc_vocab_end - sc_vocab_start] - ) - if tp_size > 1: - soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name) - soft_embeds = soft_embeds * normalizer - sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype) + if compute_sc: + # Self-conditioning soft embed = probs @ embed_tokens.weight. Under + # tensor parallelism the embedding is vocab-sharded ([vocab/tp, + # hidden]) while probs spans the full vocab, so each rank multiplies + # its local vocab slice [sc_vocab_start, sc_vocab_end) and the + # partials are summed across ranks. + local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype) + soft_embeds = torch.matmul( + local_probs, embed_weight[: sc_vocab_end - sc_vocab_start] + ) + if tp_size > 1: + soft_embeds = torch.ops.vllm.all_reduce( + soft_embeds, group_name=tp_group_name + ) + soft_embeds = soft_embeds * normalizer + sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype) + else: + # Every slot in this tile ends after this step, so the soft embed + # would never be read. The matmul is a full pass over the vocabulary + # matrix, so skip it. + sc_embeds[decode_slots] = 0 # Overwrite canvas with argmax for newly converged denoise requests newly_converged = (converged & is_denoise).unsqueeze(1) @@ -644,6 +654,8 @@ def __init__( # forward. self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) self.read_only_slots: set[int] = set() + # Slots capped at one denoise step never consume a soft embed. + self.single_step_slots: set[int] = set() # Per-slot self-conditioning soft embedding (probs @ embed_weight) from # the previous denoise step. Storing the [.., hidden] soft embed instead @@ -679,6 +691,7 @@ def add_request(self, slot_idx: int) -> None: self.seeded_slots.discard(slot_idx) self.read_only[slot_idx].fill_(False) self.read_only_slots.discard(slot_idx) + self.single_step_slots.discard(slot_idx) def remove_request(self, slot_idx: int) -> None: # add_request resets the GPU flags before a slot is reused. The host @@ -688,6 +701,7 @@ def remove_request(self, slot_idx: int) -> None: self.self_conditioning_embeds[slot_idx] = 0 self.seeded_slots.discard(slot_idx) self.read_only_slots.discard(slot_idx) + self.single_step_slots.discard(slot_idx) def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None: self.seed_canvas[slot_idx] = async_tensor_h2d( @@ -1083,9 +1097,10 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None: states = self.diffusion_states cap = extra.get("diffusion_max_steps") if cap is not None: - states.max_steps[req_idx].fill_( - max(1, min(int(cap), states.max_denoising_steps)) - ) + cap = max(1, min(int(cap), states.max_denoising_steps)) + states.max_steps[req_idx].fill_(cap) + if cap == 1: + states.single_step_slots.add(req_idx) seed = extra.get("diffusion_seed_canvas") if seed is not None: if len(seed) != self.canvas_length: @@ -1312,6 +1327,9 @@ def __call__( end_req = min(start_req + group, num_decode) tile = slice(start_req, end_req) tile_slots = decode_slots[tile] + compute_sc = not states.single_step_slots or not states.single_step_slots.issuperset( + decode_slots_np[tile].tolist() + ) scaled = _compiled_sample_step( logits[start_req * CL : end_req * CL], @@ -1348,6 +1366,7 @@ def __call__( sc_vocab_end=self.sc_vocab_end, tp_size=self.tp_size, tp_group_name=self.tp_group_name, + compute_sc=compute_sc, ) # Logprobs for denoise steps that just converged (is_encoder_phase From 8dc04ec6f399fd605db9d08f6e73acde397f5081 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Fri, 18 Sep 2026 17:30:45 -0600 Subject: [PATCH 10/25] [Core] Do not schedule a step past a read-only diffusion request's cap Async scheduling schedules one step ahead and guards against stepping past max_tokens on the assumption that a step yields one token plus drafts that may be rejected. A diffusion denoise step yields nothing until it emits a whole canvas, so the guard never fired and every read-only request ran one extra 32-token forward that the scheduler then discarded, a third of the request's GPU time. A read-only request with diffusion_max_steps emits on exactly that denoise step and ends. Each in-flight denoise step holds one canvas of output placeholders, so once the cap's worth of placeholders is in flight the request is skipped. The rule lives in a DiffusionAsyncScheduler subclass, so Scheduler and AsyncScheduler stay untouched. The async guard is out of a subclass's reach. The PP eligibility step is in reach, and raising it by one skips the same call. VllmConfig selects the subclass for diffusion models once the platform hook has had the last word on async scheduling, and get_scheduler_cls warns about a custom scheduler only for classes outside vLLM. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/v1/core/test_scheduler.py | 75 +++++++++++++++++++++++ tests/v1/core/utils.py | 13 +++- vllm/config/scheduler.py | 18 +++--- vllm/config/vllm.py | 10 +++ vllm/v1/core/sched/diffusion_scheduler.py | 44 +++++++++++++ 5 files changed, 152 insertions(+), 8 deletions(-) create mode 100644 vllm/v1/core/sched/diffusion_scheduler.py diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index a58bcab7d5ea..0f4912be66da 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -40,6 +40,8 @@ from vllm.v1.core.encoder_cache_manager import EncoderCacheManager from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash +from vllm.v1.core.sched.async_scheduler import AsyncScheduler +from vllm.v1.core.sched.diffusion_scheduler import DiffusionAsyncScheduler from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler @@ -6842,3 +6844,76 @@ def test_update_draft_token_ids_in_output_strips_padding(): -1, ] assert scheduler_output.num_invalid_spec_tokens == {request.request_id: 2} + + +def _diffusion_request(req_id: str, extra_args: dict) -> Request: + (request,) = create_requests( + num_requests=1, num_tokens=8, max_tokens=64, req_ids=[req_id] + ) + request.sampling_params.extra_args = extra_args + return request + + +def _diffusion_scheduler(**kwargs) -> DiffusionAsyncScheduler: + scheduler = create_scheduler( + async_scheduling=True, + diffusion_canvas_length=8, + scheduler_cls=DiffusionAsyncScheduler, + **kwargs, + ) + assert isinstance(scheduler, DiffusionAsyncScheduler) + return scheduler + + +def test_diffusion_scheduler_is_selected_with_async_scheduling(): + def selected(**kwargs): + config = create_scheduler(**kwargs).vllm_config.scheduler_config + return config.get_scheduler_cls(), config.async_scheduling + + # The platform has the last word on async scheduling (CPU turns it off). + cls, is_async = selected(async_scheduling=True, diffusion_canvas_length=8) + assert cls is (DiffusionAsyncScheduler if is_async else Scheduler) + cls, _ = selected(async_scheduling=False, diffusion_canvas_length=8) + assert cls is Scheduler + cls, is_async = selected(async_scheduling=True) + assert cls is (AsyncScheduler if is_async else Scheduler) + + +def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight(): + scheduler = _diffusion_scheduler() + one = _diffusion_request( + "one", {"diffusion_read_only": True, "diffusion_max_steps": 1} + ) + two = _diffusion_request( + "two", {"diffusion_read_only": True, "diffusion_max_steps": 2} + ) + # Only reads are deferred. A capped generation keeps its extra async step. + gen = _diffusion_request("gen", {"diffusion_max_steps": 1}) + for request in (one, two, gen): + scheduler.add_request(request) + + scheduler.schedule() + assert set(scheduler.schedule().num_scheduled_tokens) == {"one", "two", "gen"} + assert one.num_output_placeholders == 8 + assert set(scheduler.schedule().num_scheduled_tokens) == {"two", "gen"} + assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"} + # The deferral is renewed every call. + assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"} + + +def test_diffusion_read_deferral_keeps_a_longer_pp_wait(): + scheduler = _diffusion_scheduler(pipeline_parallel_size=3, use_v2_model_runner=True) + read = _diffusion_request( + "read", {"diffusion_read_only": True, "diffusion_max_steps": 1} + ) + scheduler.add_request(read) + + scheduler.schedule() + assert read.next_decode_eligible_step == 4 + for _ in range(2): + assert "read" not in scheduler.schedule().num_scheduled_tokens + assert "read" in scheduler.schedule().num_scheduled_tokens + assert read.next_decode_eligible_step == 7 + # Deferring this step alone would ask for 6. The PP wait to 7 stands. + assert "read" not in scheduler.schedule().num_scheduled_tokens + assert read.next_decode_eligible_step == 7 diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index f69b10ced346..ebda4f092cdb 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -7,6 +7,7 @@ from tests.v1.kv_connector.unit.utils import MockKVConfig from vllm.config import ( CacheConfig, + DiffusionConfig, ECTransferConfig, KVTransferConfig, ModelConfig, @@ -77,6 +78,8 @@ def create_scheduler( kv_cache_spec: KVCacheSpec | None = None, per_request_spec_decode_metrics: str = "none", scheduling_policy: SchedulerPolicy = "fcfs", + diffusion_canvas_length: int | None = None, + scheduler_cls: type[Scheduler] | None = None, ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -178,6 +181,12 @@ def create_scheduler( else None ) + diffusion_config: DiffusionConfig | None = None + if diffusion_canvas_length is not None: + # A diffusion checkpoint declares its canvas in the HF config. + model_config.hf_config.canvas_length = diffusion_canvas_length + diffusion_config = DiffusionConfig(canvas_length=diffusion_canvas_length) + vllm_config = VllmConfig( scheduler_config=scheduler_config, model_config=model_config, @@ -188,6 +197,7 @@ def create_scheduler( ), kv_transfer_config=kv_transfer_config, speculative_config=speculative_config, + diffusion_config=diffusion_config, ec_transfer_config=ec_transfer_config, observability_config=ObservabilityConfig( per_request_spec_decode_metrics=per_request_spec_decode_metrics, @@ -207,7 +217,8 @@ def create_scheduler( ) cache_config.num_gpu_blocks = num_blocks register_all_kvcache_specs(vllm_config) - scheduler_cls = AsyncScheduler if async_scheduling else Scheduler + if scheduler_cls is None: + scheduler_cls = AsyncScheduler if async_scheduling else Scheduler scheduler = scheduler_cls( vllm_config=vllm_config, kv_cache_config=kv_cache_config, diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 321c76fe3058..fc1d77948e1b 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -229,13 +229,17 @@ def get_scheduler_cls(self) -> type["SchedulerInterface"]: # The first half of this warning can be removed once the Scheduler interface is # finalized and we can maintain support for scheduler classes that implement it - logger.warning_once( - "Using custom scheduler class %s. This scheduler interface is not public " - "and compatibility may not be maintained. If you have subclassed Scheduler " - "instead of AsyncScheduler, you will see degraded performance due to async " - "scheduling being disabled.", - self.scheduler_cls, # type: ignore[arg-type] - ) + if not ( + isinstance(self.scheduler_cls, str) + and self.scheduler_cls.startswith("vllm.") + ): + logger.warning_once( + "Using custom scheduler class %s. This scheduler interface is not " + "public and compatibility may not be maintained. If you have " + "subclassed Scheduler instead of AsyncScheduler, you will see " + "degraded performance due to async scheduling being disabled.", + self.scheduler_cls, # type: ignore[arg-type] + ) if not isinstance(self.scheduler_cls, str): return cast(type["SchedulerInterface"], self.scheduler_cls) return resolve_obj_by_qualname(self.scheduler_cls) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index bf11255e8e0a..65f2f8df523d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1878,6 +1878,16 @@ def has_blocked_weights(): ) current_platform.check_and_update_config(self) + # After the platform hook, which has the last word on async scheduling. + if ( + self.diffusion_config is not None + and self.scheduler_config.async_scheduling + and self.scheduler_config.scheduler_cls is None + ): + self.scheduler_config.scheduler_cls = ( + "vllm.v1.core.sched.diffusion_scheduler.DiffusionAsyncScheduler" + ) + self._normalize_piecewise_cudagraph_mode( breakable_cudagraph_enabled=breakable_cudagraph_enabled ) diff --git a/vllm/v1/core/sched/diffusion_scheduler.py b/vllm/v1/core/sched/diffusion_scheduler.py new file mode 100644 index 000000000000..3f4c1a31731d --- /dev/null +++ b/vllm/v1/core/sched/diffusion_scheduler.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Async scheduling for diffusion requests. + +This rule lives here so Scheduler.schedule() and AsyncScheduler stay +unchanged. VllmConfig selects this class for a diffusion model under async +scheduling. A sync scheduler creates no output placeholders, which the rule +reads. +""" + +from vllm.v1.core.sched.async_scheduler import AsyncScheduler +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.request import Request + + +def _read_in_flight(request: Request, width: int) -> bool: + """True when a read-only request has all its denoise steps in flight. + + The request emits its canvas on the last step and ends, so a further + step is discarded. + """ + params = request.sampling_params + extra = params.extra_args if params is not None else None + if not extra or not extra.get("diffusion_read_only"): + return False + steps = extra.get("diffusion_max_steps") + if not steps: + return False + # Each in-flight denoise step holds one canvas of placeholders. + return request.num_output_placeholders >= steps * width + + +class DiffusionAsyncScheduler(AsyncScheduler): + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + for request in self.running: + if _read_in_flight(request, self.num_spec_tokens): + # Scheduler.schedule()'s max_tokens guard cannot be reached + # from a subclass. schedule() advances current_step before it + # checks decode eligibility, so current_step + 2 skips this + # call alone. max keeps a longer pipeline-parallel wait. + request.next_decode_eligible_step = max( + request.next_decode_eligible_step, self.current_step + 2 + ) + return super().schedule(throttle_prefills) From 5b6f205d8ef7969a331f0ab990d2c0f33f539c37 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Fri, 18 Sep 2026 17:25:56 -0600 Subject: [PATCH 11/25] [Core] Per-request canvas width for diffusion reads `diffusion_canvas_length` in extra_args runs a request on a canvas narrower than the served one. The server keeps its canvas for buffers and graph capture. The scheduler schedules the request's width as draft tokens each step, and the sampler already pads a short canvas back to the served width and commits only the real positions. Under async scheduling the DiffusionAsyncScheduler cuts the placeholder spec-token list to the width before each step, and without it the draft-token handler reports the width per request, from a new model-state hook that the DiffusionGemma state answers from a per-slot width array. The example server picks the smallest multiple of 16 that holds a schema's answer template, so one wide server serves short reads at short-canvas cost. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/test_sampling_params.py | 30 ++++++++++++++++++ tests/v1/core/test_scheduler.py | 31 ++++++++++++++++++- tests/v1/sample/test_diffusion_gemma_reads.py | 12 +++++++ vllm/model_executor/models/diffusion_gemma.py | 20 +++++++++--- vllm/sampling_params.py | 30 ++++++++++++------ vllm/v1/core/sched/diffusion_scheduler.py | 20 +++++++++--- vllm/v1/worker/gpu/model_runner.py | 3 ++ vllm/v1/worker/gpu/model_states/interface.py | 5 +++ vllm/v1/worker/gpu/spec_decode/utils.py | 18 +++++++++-- 9 files changed, 148 insertions(+), 21 deletions(-) diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index a66ed1d49ba9..0980d2873efc 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -116,6 +116,36 @@ def test_diffusion_read_only_ends_after_one_canvas(max_tokens, expected, flag): assert params.ignore_eos +@pytest.mark.parametrize( + "width, match", + [(0, "positive integer"), ("4", "positive integer"), (True, "positive integer"), (9, "no larger")], +) +def test_diffusion_rejects_bad_canvas_length(width, match): + with pytest.raises(VLLMValidationError, match=match): + _verify_diffusion( + SamplingParams(extra_args={"diffusion_canvas_length": width}), canvas_length=8 + ) + + +def test_diffusion_canvas_length_sizes_the_seed_and_the_read(): + params = SamplingParams( + max_tokens=100, + extra_args={ + "diffusion_canvas_length": 4, + "diffusion_seed_canvas": [1, 2, 3, 4], + "diffusion_read_only": True, + }, + ) + _verify_diffusion(params, canvas_length=8) + assert params.max_tokens == 4 + + params = SamplingParams( + extra_args={"diffusion_canvas_length": 4, "diffusion_seed_canvas": [0] * 8} + ) + with pytest.raises(VLLMValidationError, match="exactly 4 ids, got 8"): + _verify_diffusion(params, canvas_length=8) + + def test_diffusion_accepts_extra_args(): params = SamplingParams( extra_args={ diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 0f4912be66da..0a8de5cf1679 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -41,7 +41,10 @@ from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash from vllm.v1.core.sched.async_scheduler import AsyncScheduler -from vllm.v1.core.sched.diffusion_scheduler import DiffusionAsyncScheduler +from vllm.v1.core.sched.diffusion_scheduler import ( + DiffusionAsyncScheduler, + diffusion_canvas_width, +) from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler @@ -6846,6 +6849,16 @@ def test_update_draft_token_ids_in_output_strips_padding(): assert scheduler_output.num_invalid_spec_tokens == {request.request_id: 2} +def test_diffusion_canvas_width_defaults_to_the_served_canvas(): + def req(extra): + return SimpleNamespace(sampling_params=SimpleNamespace(extra_args=extra)) + + assert diffusion_canvas_width(req(None), 64) == 64 + assert diffusion_canvas_width(req({}), 64) == 64 + assert diffusion_canvas_width(req({"diffusion_canvas_length": 16}), 64) == 16 + assert diffusion_canvas_width(SimpleNamespace(sampling_params=None), 64) == 64 + + def _diffusion_request(req_id: str, extra_args: dict) -> Request: (request,) = create_requests( num_requests=1, num_tokens=8, max_tokens=64, req_ids=[req_id] @@ -6879,6 +6892,22 @@ def selected(**kwargs): assert cls is (AsyncScheduler if is_async else Scheduler) +def test_diffusion_scheduler_narrows_the_canvas_per_request(): + scheduler = _diffusion_scheduler() + wide = _diffusion_request("wide", {}) + narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4}) + scheduler.add_request(wide) + scheduler.add_request(narrow) + + # The prefill step lays down the served canvas as placeholders. + scheduler.schedule() + output = scheduler.schedule() + assert output.scheduled_spec_decode_tokens["wide"] == [-1] * 8 + assert output.scheduled_spec_decode_tokens["narrow"] == [-1] * 4 + assert output.num_scheduled_tokens["narrow"] == 4 + assert narrow.num_output_placeholders == 4 + + def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight(): scheduler = _diffusion_scheduler() one = _diffusion_request( diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py index 28f92e4764e2..d7c37670b317 100644 --- a/tests/v1/sample/test_diffusion_gemma_reads.py +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -86,6 +86,18 @@ def test_add_request_clears_seed_and_read_only(): assert not bool(states.read_only[0]) +def test_canvas_width_resets_with_the_slot(): + states = _states() + states.add_request(0) + states.canvas_width_np[0] = 4 + states.set_seed_canvas(0, [7, 7, 7, 7]) + assert states.seed_canvas[0, :4].tolist() == [7, 7, 7, 7] + + states.add_request(0) + + assert states.canvas_width_np[0] == CL + + def test_remove_request_forgets_the_slot(): states = _states() states.add_request(0) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 6b3d39b6a770..6db1cbc5d2e3 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -656,6 +656,9 @@ def __init__( self.read_only_slots: set[int] = set() # Slots capped at one denoise step never consume a soft embed. self.single_step_slots: set[int] = set() + # Per-slot canvas width, at most canvas_length. The scheduler schedules + # this many draft tokens for the slot and the sampler pads the rest. + self.canvas_width_np = np.full(max_num_reqs, canvas_length, dtype=np.int32) # Per-slot self-conditioning soft embedding (probs @ embed_weight) from # the previous denoise step. Storing the [.., hidden] soft embed instead @@ -692,6 +695,7 @@ def add_request(self, slot_idx: int) -> None: self.read_only[slot_idx].fill_(False) self.read_only_slots.discard(slot_idx) self.single_step_slots.discard(slot_idx) + self.canvas_width_np[slot_idx] = self.canvas_length def remove_request(self, slot_idx: int) -> None: # add_request resets the GPU flags before a slot is reused. The host @@ -704,7 +708,8 @@ def remove_request(self, slot_idx: int) -> None: self.single_step_slots.discard(slot_idx) def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None: - self.seed_canvas[slot_idx] = async_tensor_h2d( + """``ids`` covers the slot's canvas width; positions past it are never scheduled.""" + self.seed_canvas[slot_idx, : len(ids)] = async_tensor_h2d( ids, dtype=torch.int64, device=self.device ) self.has_seed[slot_idx].fill_(True) @@ -842,6 +847,10 @@ def remove_request(self, req_id: str) -> None: if idx is not None: self.diffusion_states.remove_request(idx) + def num_draft_tokens_per_req(self, input_batch: Any) -> np.ndarray: + slots = input_batch.idx_mapping_np[: input_batch.num_reqs] + return self.diffusion_states.canvas_width_np[slots] + def prepare_inputs_embeds( self, scheduled_encoder_inputs: dict[str, list[int]], @@ -1101,12 +1110,15 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None: states.max_steps[req_idx].fill_(cap) if cap == 1: states.single_step_slots.add(req_idx) + width = extra.get("diffusion_canvas_length") + if width: + states.canvas_width_np[req_idx] = max(1, min(int(width), self.canvas_length)) + width = int(states.canvas_width_np[req_idx]) seed = extra.get("diffusion_seed_canvas") if seed is not None: - if len(seed) != self.canvas_length: + if len(seed) != width: raise ValueError( - "diffusion_seed_canvas must hold exactly " - f"{self.canvas_length} ids, got {len(seed)}" + f"diffusion_seed_canvas must hold exactly {width} ids, got {len(seed)}" ) states.set_seed_canvas(req_idx, seed) if extra.get("diffusion_read_only"): diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 2ccae5100488..b7660c930a27 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1087,6 +1087,22 @@ def _validate_diffusion_extra_args( if not extra: return + canvas_length = None if diffusion_config is None else diffusion_config.canvas_length + width = extra.get("diffusion_canvas_length") + if width is not None and ( + not isinstance(width, int) + or isinstance(width, bool) + or width < 1 + or (canvas_length is not None and width > canvas_length) + ): + raise VLLMValidationError( + "diffusion_canvas_length must be a positive integer no larger " + "than the served canvas.", + parameter="extra_args", + ) + # The request's canvas: its own width, else the served one. + expected_len = width or canvas_length + seed = extra.get("diffusion_seed_canvas") if seed is not None: if not isinstance(seed, (list, tuple)) or not all( @@ -1102,13 +1118,10 @@ def _validate_diffusion_extra_args( f"diffusion_seed_canvas ids must be in [0, {vocab_size}).", parameter="extra_args", ) - if ( - diffusion_config is not None - and len(seed) != diffusion_config.canvas_length - ): + if expected_len is not None and len(seed) != expected_len: raise VLLMValidationError( "diffusion_seed_canvas must hold exactly " - f"{diffusion_config.canvas_length} ids, got {len(seed)}.", + f"{expected_len} ids, got {len(seed)}.", parameter="extra_args", ) @@ -1134,11 +1147,8 @@ def _validate_diffusion_extra_args( # One canvas is the whole output. Cap max_tokens so the scheduler # ends the request there, and ignore EOS so an end-of-turn token # drawn into a noise slot does not end it early. - if diffusion_config is not None: - canvas_length = diffusion_config.canvas_length - self.max_tokens = min( - self.max_tokens or canvas_length, canvas_length - ) + if expected_len is not None: + self.max_tokens = min(self.max_tokens or expected_len, expected_len) self.ignore_eos = True def _validate_structured_outputs( diff --git a/vllm/v1/core/sched/diffusion_scheduler.py b/vllm/v1/core/sched/diffusion_scheduler.py index 3f4c1a31731d..96462aa6abe1 100644 --- a/vllm/v1/core/sched/diffusion_scheduler.py +++ b/vllm/v1/core/sched/diffusion_scheduler.py @@ -2,10 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Async scheduling for diffusion requests. -This rule lives here so Scheduler.schedule() and AsyncScheduler stay +These rules live here so Scheduler.schedule() and AsyncScheduler stay unchanged. VllmConfig selects this class for a diffusion model under async -scheduling. A sync scheduler creates no output placeholders, which the rule -reads. +scheduling. A sync scheduler creates no output placeholders, which both rules +read. """ from vllm.v1.core.sched.async_scheduler import AsyncScheduler @@ -13,6 +13,14 @@ from vllm.v1.request import Request +def diffusion_canvas_width(request: Request, canvas_length: int) -> int: + """The canvas width a diffusion request asked for, else the served one.""" + params = request.sampling_params + extra = params.extra_args if params is not None else None + width = extra.get("diffusion_canvas_length") if extra else None + return int(width) if width else canvas_length + + def _read_in_flight(request: Request, width: int) -> bool: """True when a read-only request has all its denoise steps in flight. @@ -33,7 +41,11 @@ def _read_in_flight(request: Request, width: int) -> bool: class DiffusionAsyncScheduler(AsyncScheduler): def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: for request in self.running: - if _read_in_flight(request, self.num_spec_tokens): + width = diffusion_canvas_width(request, self.num_spec_tokens) + # The placeholders are as wide as the served canvas. + if len(request.spec_token_ids) > width: + request.spec_token_ids = request.spec_token_ids[:width] + if _read_in_flight(request, width): # Scheduler.schedule()'s max_tokens guard cannot be reached # from a subclass. schedule() advances current_step before it # checks decode eligibility, so current_step + 2 skips this diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 92a37008cebe..f058eac2db19 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2169,6 +2169,9 @@ def sample_tokens( self.draft_tokens_handler.set_draft_tokens( input_batch, self.req_states.draft_tokens[input_batch.idx_mapping], + num_draft_tokens_per_req=self.model_state.num_draft_tokens_per_req( + input_batch + ), ) if self.pp_handler is not None: self.pp_handler.broadcast_drafts( diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index a5c8d6f4cfd7..70bb19c8d22f 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -237,3 +237,8 @@ def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: num_new_sampled_tokens_per_step: int = 1 """New tokens sampled on each decode step (excluding accepted draft tokens, a.k.a num bonus tokens).""" + + def num_draft_tokens_per_req(self, input_batch: Any) -> Any: + """Draft tokens each request in the batch carries into its next step, + as an int array, or None when every request carries the buffer width.""" + return None diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index f73acb2bdaed..465433f844c0 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -36,12 +36,20 @@ def __init__(self, device: torch.device | None = None): self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None self.num_draft_tokens: int = 0 + # Per-request draft counts when a model state narrows some requests. + self.num_draft_tokens_per_req: list[int] | None = None def set_draft_tokens( - self, input_batch: InputBatch, draft_tokens: torch.Tensor + self, + input_batch: InputBatch, + draft_tokens: torch.Tensor, + num_draft_tokens_per_req: np.ndarray | None = None, ) -> None: self.req_ids = input_batch.req_ids self.num_draft_tokens = draft_tokens.shape[1] + self.num_draft_tokens_per_req = ( + None if num_draft_tokens_per_req is None else num_draft_tokens_per_req.tolist() + ) if not input_batch.has_structured_output_reqs: # No draft token validation needs to be performed by # the scheduler for this batch. @@ -61,12 +69,18 @@ def set_draft_tokens( self.copy_event.record() def get_draft_tokens(self) -> DraftTokenIds | None: + per_req = self.num_draft_tokens_per_req if self.draft_tokens_np is not None: self.copy_event.synchronize() draft_token_ids = self.draft_tokens_np.tolist() + if per_req is not None: + draft_token_ids = [row[:n] for row, n in zip(draft_token_ids, per_req)] else: # This case only happens when async scheduling is disabled. - draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] + draft_token_ids = [ + [-1] * (self.num_draft_tokens if per_req is None else per_req[i]) + for i in range(len(self.req_ids)) + ] return DraftTokenIds(self.req_ids, draft_token_ids) From 9e9a23f21d931139525f58fd322eaa4733b615cf Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Wed, 16 Sep 2026 22:31:30 -0600 Subject: [PATCH 12/25] [Model] DiffusionGemma: tile the sampler by canvas width With per-request widths the sampler still padded every decode request to the served canvas and ran the compiled step over [tile * canvas, vocab], so a 32-row read on a 128-row server paid four times the fp32 vocabulary work of a native 32-row server (54 req/s fell to 38 at 32-way). Decode requests are now grouped by width and each tile runs the compiled step at its own width over [:, :W] views of the state, in ascending width so the final canvas-to-draft copy is the widest. Columns past a request's width are never scheduled, so they may hold stale values. A tile whose requests are contiguous and untruncated slices the flat logits directly, which is the whole batch on a server where every request uses the served width, so that case is unchanged. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Mastracci --- tests/v1/sample/test_diffusion_gemma_reads.py | 40 ++- vllm/model_executor/models/diffusion_gemma.py | 240 ++++++++++-------- 2 files changed, 159 insertions(+), 121 deletions(-) diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py index d7c37670b317..2165f6d77e66 100644 --- a/tests/v1/sample/test_diffusion_gemma_reads.py +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -111,32 +111,37 @@ def test_remove_request_forgets_the_slot(): def _denoise_once( - states: DiffusionGemmaRequestStates, slots: list[int], compute_sc: bool = True + states: DiffusionGemmaRequestStates, + slots: list[int], + compute_sc: bool = True, + width: int = CL, ) -> None: """One compiled denoise step over ``slots`` with flat logits, so nothing - converges by stability or confidence and only the step cap can end it.""" + converges by stability or confidence and only the step cap can end it. + ``width`` below CL runs the step on [:, :width] views, as the sampler + does for a narrow tile.""" n = len(slots) device = states.device decode_slots = torch.tensor(slots, dtype=torch.int64, device=device) decode_idx = torch.arange(n, dtype=torch.int64, device=device) _compiled_sample_step( - torch.zeros(n * CL, VOCAB, device=device), + torch.zeros(n * width, VOCAB, device=device), decode_slots, decode_idx, decode_slots, - torch.full((n,), CL, dtype=torch.int64, device=device), - states.canvas, - states.argmax_canvas, + torch.full((n,), width, dtype=torch.int64, device=device), + states.canvas[:, :width], + states.argmax_canvas[:, :width], states.step, states.is_encoder_phase, states.confident, - states.self_conditioning_embeds, + states.self_conditioning_embeds[:, :width], torch.zeros(VOCAB, 4, device=device), torch.tensor(1.0, device=device), - states.accepted_canvas_history, + states.accepted_canvas_history[:, :, :width], states.accepted_canvas_history_len, states.max_steps, - torch.zeros(n, CL, dtype=torch.int32, device=device), + torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width], torch.zeros(n, dtype=torch.int32, device=device), torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device), max_denoising_steps=float(MAX_STEPS), @@ -144,7 +149,7 @@ def _denoise_once( t_max=1.0, confidence_threshold=0.1, vocab_size=VOCAB, - CL=CL, + CL=width, ST=states.stability_threshold, entropy_bound=0.1, sc_vocab_start=0, @@ -166,6 +171,21 @@ def test_single_step_tile_skips_self_conditioning(): assert not states.self_conditioning_embeds[0].any() +def test_narrow_tile_leaves_the_rest_of_the_canvas_alone(): + states = _states() + states.add_request(0) + states.is_encoder_phase[0] = False + states.canvas[0] = 5 + states.argmax_canvas[0] = 5 + + _denoise_once(states, [0], width=4) + + # Columns past the width are untouched; the step counted. + assert states.canvas[0, 4:].tolist() == [5] * (CL - 4) + assert states.argmax_canvas[0, 4:].tolist() == [5] * (CL - 4) + assert states.step[0] == 1 + + def test_step_cap_is_per_slot(): states = _states() for slot in (0, 1): diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 6db1cbc5d2e3..f31c83a1fb14 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -1245,7 +1245,6 @@ def __call__( # --- CPU/NumPy setup (outside compile): split decode vs prefill, init # canvas for any new prefills, and stage decode slot indices to GPU. --- states = self.diffusion_states - CL = self.canvas_length slots_np = input_batch.idx_mapping_np[:num_reqs] per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) @@ -1288,18 +1287,10 @@ def __call__( if top_k is not None or top_p is not None: logits = apply_top_k_top_p(logits.float(), top_k, top_p) - # Pad any truncated canvas back to CL so the uniform-CL sampler math - # holds. Phantom (padded) positions are zeroed → uniform logits → high - # entropy (no premature convergence) and argmax 0 (stable); they are - # never committed (num_sampled == real length). masked_fill (not - # multiply) so -inf entries from top_k/top_p filtering above don't - # turn phantom rows into NaN. - if num_decode > 0 and valid_canvas_len_np.min() < CL: - ar = torch.arange(CL, device=device) - starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req - valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL] - src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1) - logits = logits[src.reshape(-1)].masked_fill_(~valid.reshape(-1, 1), 0) + # Where each decode request's rows start in the flat logits. Tiles + # below gather and pad a request to the tile's width. + row_starts_np = np.concatenate(([0], np.cumsum(valid_canvas_len_np)[:-1])).astype(np.int64) + row_starts = async_tensor_h2d(row_starts_np, device=device) # Clear once: the tiled loop below only scatters its own decode slots, # so it must not re-clear earlier tiles' writes. @@ -1322,107 +1313,134 @@ def __call__( want_logprobs = max_num_logprobs >= 0 or max_token_ids > 0 num_logprobs = max(max_num_logprobs, 0) - # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in - # _compiled_sample_step keeps several live [group * CL, vocab] copies, so - # size each tile to a fraction of free memory to bound the transient at - # high concurrency. Tiling is bit-identical to a single pass. - group = max(num_decode, 1) - if num_decode > 0: - free, _ = current_platform.mem_get_info() - # ~10 transient fp32 copies of [group * CL, vocab] inside the step + # Sample per tile. Decode requests are grouped by canvas width and each + # tile runs the compiled step at that width over [:, :W] views of the + # state, so a narrow read pays for its own rows rather than the served + # canvas. Widths ascend, so the last tile's canvas-to-draft copy (over + # all slots) is the widest. The fp32 pipeline keeps several live + # [tile * W, vocab] copies, so a tile is also bounded by free memory. + widths_np = states.canvas_width_np[decode_slots_np] + order = np.argsort(widths_np, kind="stable") + free = current_platform.mem_get_info()[0] if num_decode > 0 else 0 + run_start = 0 + while run_start < num_decode: + W = int(widths_np[order[run_start]]) + run_end = run_start + while run_end < num_decode and widths_np[order[run_end]] == W: + run_end += 1 + # ~10 transient fp32 copies of [tile * W, vocab] inside the step # (eager peaks at ~8; pad for allocator overhead and small tensors). - bytes_per_req = CL * self.vocab_size * 4 * 10 - budget = int(free * 0.5) // max(bytes_per_req, 1) - group = max(1, min(num_decode, budget)) - - for start_req in range(0, num_decode, group): - end_req = min(start_req + group, num_decode) - tile = slice(start_req, end_req) - tile_slots = decode_slots[tile] - compute_sc = not states.single_step_slots or not states.single_step_slots.issuperset( - decode_slots_np[tile].tolist() - ) + budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 10, 1)) + for t0 in range(run_start, run_end, budget): + sel_np = order[t0 : min(t0 + budget, run_end)] + n = len(sel_np) + sel = async_tensor_h2d(sel_np.astype(np.int64), device=device) + tile_slots = decode_slots[sel] + tile_valid = valid_canvas_len[sel] + tile_valid_np = valid_canvas_len_np[sel_np] + contiguous = bool((sel_np == np.arange(sel_np[0], sel_np[0] + n)).all()) + if contiguous and tile_valid_np.min() == W: + r0 = int(row_starts_np[sel_np[0]]) + tile_logits = logits[r0 : r0 + n * W] + else: + # Pad each request to W. Phantom positions are zeroed: + # uniform logits, high entropy, argmax 0, never committed. + # masked_fill, not multiply, so -inf from top_k/top_p + # filtering above cannot turn a phantom row into NaN. + ar = torch.arange(W, device=device) + src = (row_starts[sel].unsqueeze(1) + ar.unsqueeze(0)).clamp_max( + logits.shape[0] - 1 + ) + valid = ar.unsqueeze(0) < tile_valid.unsqueeze(1) + tile_logits = logits[src.reshape(-1)].masked_fill_( + ~valid.reshape(-1, 1), 0 + ) + compute_sc = ( + not states.single_step_slots + or not states.single_step_slots.issuperset( + decode_slots_np[sel_np].tolist() + ) + ) - scaled = _compiled_sample_step( - logits[start_req * CL : end_req * CL], - tile_slots, - decode_idx[tile], - all_slots, - valid_canvas_len[tile], - # State - states.canvas, - states.argmax_canvas, - states.step, - states.is_encoder_phase, - states.confident, - states.self_conditioning_embeds, - self.embed_weight, - self.normalizer, - states.accepted_canvas_history, - states.accepted_canvas_history_len, - states.max_steps, - # Output - sampled, - num_sampled, - self.req_states.draft_tokens, - # Config - max_denoising_steps=float(states.max_denoising_steps), - t_min=self.t_min, - t_max=self.t_max, - confidence_threshold=self.confidence_threshold, - vocab_size=self.vocab_size, - CL=CL, - ST=states.stability_threshold, - entropy_bound=self.entropy_bound, - sc_vocab_start=self.sc_vocab_start, - sc_vocab_end=self.sc_vocab_end, - tp_size=self.tp_size, - tp_group_name=self.tp_group_name, - compute_sc=compute_sc, - ) + scaled = _compiled_sample_step( + tile_logits, + tile_slots, + decode_idx[sel], + all_slots, + tile_valid, + # State, viewed at this tile's width + states.canvas[:, :W], + states.argmax_canvas[:, :W], + states.step, + states.is_encoder_phase, + states.confident, + states.self_conditioning_embeds[:, :W], + self.embed_weight, + self.normalizer, + states.accepted_canvas_history[:, :, :W], + states.accepted_canvas_history_len, + states.max_steps, + # Output + sampled[:, :W], + num_sampled, + self.req_states.draft_tokens, + # Config + max_denoising_steps=float(states.max_denoising_steps), + t_min=self.t_min, + t_max=self.t_max, + confidence_threshold=self.confidence_threshold, + vocab_size=self.vocab_size, + CL=W, + ST=states.stability_threshold, + entropy_bound=self.entropy_bound, + sc_vocab_start=self.sc_vocab_start, + sc_vocab_end=self.sc_vocab_end, + tp_size=self.tp_size, + tp_group_name=self.tp_group_name, + compute_sc=compute_sc, + ) - # Logprobs for denoise steps that just converged (is_encoder_phase - # flipped False→True), stashed per tile so `scaled` is freed each tile. - if want_logprobs: - converged_mask = states.is_encoder_phase[tile_slots] - just_converged = converged_mask & ~is_committing[tile] - if just_converged.any(): - flat_logits = scaled.reshape(-1, scaled.shape[-1]) - argmax_tokens = scaled.argmax(dim=-1) - raw_flat: torch.Tensor | None = None - for local_idx in just_converged.nonzero(as_tuple=True)[0]: - li = local_idx.item() - slot = tile_slots[local_idx].item() - # Stash only the real canvas positions (== CL unless this - # canvas was truncated near max_model_len); padded tail - # positions are never emitted. - k_i = int(valid_canvas_len_np[start_req + li]) - pos = li * CL - src = flat_logits - if slot in states.read_only_slots: - # Read-only slots report logprobs at temperature 1. - # The schedule-tempered logits share the argmax. - if raw_flat is None: - tile_rows = slice(start_req * CL, end_req * CL) - raw_flat = logits[tile_rows].float() - src = raw_flat - per_req_ids = max_token_ids > 0 - self._pending_logprobs[slot] = compute_topk_scores( - src[pos : pos + k_i], - num_logprobs, - argmax_tokens[local_idx][:k_i], - logprob_token_ids_state=( - self.logprob_token_ids_state if per_req_ids else None - ), - # every row of this stash belongs to one slot - expanded_idx_mapping=( - torch.full((k_i,), slot, dtype=torch.int32, device=device) - if per_req_ids - else None - ), - max_per_req_token_ids=max_token_ids, - logits_mode=self.logits_mode, - ) + # Logprobs for denoise steps that just converged (is_encoder_phase + # flipped False→True), stashed per tile so `scaled` is freed each tile. + if want_logprobs: + converged_mask = states.is_encoder_phase[tile_slots] + just_converged = converged_mask & ~is_committing[sel] + if just_converged.any(): + flat_logits = scaled.reshape(-1, scaled.shape[-1]) + argmax_tokens = scaled.argmax(dim=-1) + raw_flat: torch.Tensor | None = None + for local_idx in just_converged.nonzero(as_tuple=True)[0]: + li = local_idx.item() + slot = tile_slots[local_idx].item() + # Stash only the real canvas positions; padded tail + # positions are never emitted. + k_i = int(tile_valid_np[li]) + pos = li * W + src = flat_logits + if slot in states.read_only_slots: + # Read-only slots report logprobs at temperature 1. + # The schedule-tempered logits share the argmax. + if raw_flat is None: + raw_flat = tile_logits.float() + src = raw_flat + per_req_ids = max_token_ids > 0 + self._pending_logprobs[slot] = compute_topk_scores( + src[pos : pos + k_i], + num_logprobs, + argmax_tokens[local_idx][:k_i], + logprob_token_ids_state=( + self.logprob_token_ids_state if per_req_ids else None + ), + # every row of this stash belongs to one slot + expanded_idx_mapping=( + torch.full((k_i,), slot, dtype=torch.int32, device=device) + if per_req_ids + else None + ), + max_per_req_token_ids=max_token_ids, + logits_mode=self.logits_mode, + ) + run_start = run_end # Read-only slots that converged this step emit their argmax canvas # now, skip the encoder phase, and hand out their logprobs below. From f2d852a0a092d41cfbc52a17941da88b044e6cb3 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Fri, 18 Sep 2026 16:27:25 -0600 Subject: [PATCH 13/25] [Misc] Format the diffusion read changes for the pinned ruff pre-commit runs ruff-check and ruff-format, and upstream's copies of these files are clean, so every finding belongs to lines this branch adds. Co-Authored-By: Claude Opus 5 Signed-off-by: Matt Mastracci --- tests/test_sampling_params.py | 10 ++++-- vllm/model_executor/models/diffusion_gemma.py | 34 ++++++++++++------- vllm/sampling_params.py | 4 ++- vllm/v1/worker/gpu/spec_decode/utils.py | 4 ++- 4 files changed, 36 insertions(+), 16 deletions(-) diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index 0980d2873efc..ee3005ab71d2 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -118,12 +118,18 @@ def test_diffusion_read_only_ends_after_one_canvas(max_tokens, expected, flag): @pytest.mark.parametrize( "width, match", - [(0, "positive integer"), ("4", "positive integer"), (True, "positive integer"), (9, "no larger")], + [ + (0, "positive integer"), + ("4", "positive integer"), + (True, "positive integer"), + (9, "no larger"), + ], ) def test_diffusion_rejects_bad_canvas_length(width, match): with pytest.raises(VLLMValidationError, match=match): _verify_diffusion( - SamplingParams(extra_args={"diffusion_canvas_length": width}), canvas_length=8 + SamplingParams(extra_args={"diffusion_canvas_length": width}), + canvas_length=8, ) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index f31c83a1fb14..958f3da039de 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -708,7 +708,8 @@ def remove_request(self, slot_idx: int) -> None: self.single_step_slots.discard(slot_idx) def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None: - """``ids`` covers the slot's canvas width; positions past it are never scheduled.""" + """``ids`` covers the slot's canvas width; positions past it are never + scheduled.""" self.seed_canvas[slot_idx, : len(ids)] = async_tensor_h2d( ids, dtype=torch.int64, device=self.device ) @@ -1112,13 +1113,16 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None: states.single_step_slots.add(req_idx) width = extra.get("diffusion_canvas_length") if width: - states.canvas_width_np[req_idx] = max(1, min(int(width), self.canvas_length)) + states.canvas_width_np[req_idx] = max( + 1, min(int(width), self.canvas_length) + ) width = int(states.canvas_width_np[req_idx]) seed = extra.get("diffusion_seed_canvas") if seed is not None: if len(seed) != width: raise ValueError( - f"diffusion_seed_canvas must hold exactly {width} ids, got {len(seed)}" + f"diffusion_seed_canvas must hold exactly {width} ids, " + f"got {len(seed)}" ) states.set_seed_canvas(req_idx, seed) if extra.get("diffusion_read_only"): @@ -1289,7 +1293,9 @@ def __call__( # Where each decode request's rows start in the flat logits. Tiles # below gather and pad a request to the tile's width. - row_starts_np = np.concatenate(([0], np.cumsum(valid_canvas_len_np)[:-1])).astype(np.int64) + row_starts_np = np.concatenate( + ([0], np.cumsum(valid_canvas_len_np)[:-1]) + ).astype(np.int64) row_starts = async_tensor_h2d(row_starts_np, device=device) # Clear once: the tiled loop below only scatters its own decode slots, @@ -1429,11 +1435,15 @@ def __call__( num_logprobs, argmax_tokens[local_idx][:k_i], logprob_token_ids_state=( - self.logprob_token_ids_state if per_req_ids else None + self.logprob_token_ids_state + if per_req_ids + else None ), # every row of this stash belongs to one slot expanded_idx_mapping=( - torch.full((k_i,), slot, dtype=torch.int32, device=device) + torch.full( + (k_i,), slot, dtype=torch.int32, device=device + ) if per_req_ids else None ), @@ -1450,9 +1460,9 @@ def __call__( ): # Read-only slots never enter the encoder phase, so the flag here # means converged this step. - ro_mask = states.read_only[decode_slots] & states.is_encoder_phase[ - decode_slots - ] + ro_mask = ( + states.read_only[decode_slots] & states.is_encoder_phase[decode_slots] + ) if bool(ro_mask.any()): ro_idx = decode_idx[ro_mask] ro_slots = decode_slots[ro_mask] @@ -1469,9 +1479,9 @@ def __call__( and (emit_now or bool(is_committing.any())) and self._pending_logprobs ): - committing_slots = set( - decode_slots_np[is_committing.cpu().numpy()].tolist() - ) | emit_now + committing_slots = ( + set(decode_slots_np[is_committing.cpu().numpy()].tolist()) | emit_now + ) parts_ids, parts_lp, parts_ranks = [], [], [] cu_gen: list[int] = [] flat_offset = 0 diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index b7660c930a27..1f7c86b0faf5 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1087,7 +1087,9 @@ def _validate_diffusion_extra_args( if not extra: return - canvas_length = None if diffusion_config is None else diffusion_config.canvas_length + canvas_length = ( + None if diffusion_config is None else diffusion_config.canvas_length + ) width = extra.get("diffusion_canvas_length") if width is not None and ( not isinstance(width, int) diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 465433f844c0..0e673670d611 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -48,7 +48,9 @@ def set_draft_tokens( self.req_ids = input_batch.req_ids self.num_draft_tokens = draft_tokens.shape[1] self.num_draft_tokens_per_req = ( - None if num_draft_tokens_per_req is None else num_draft_tokens_per_req.tolist() + None + if num_draft_tokens_per_req is None + else num_draft_tokens_per_req.tolist() ) if not input_batch.has_structured_output_reqs: # No draft token validation needs to be performed by From 9b85cb4eb3e7463bc1b8e375e37675933df3c900 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Fri, 18 Sep 2026 16:30:44 -0600 Subject: [PATCH 14/25] [Docs] Example structured-decision server for DiffusionGemma reads The schema-to-canvas work (template rendering, single-token label slots, noise seeding, averaging over draws) sits behind a server so clients never build a canvas by hand. It speaks Jev's decision contract: POST /v1/systemone takes a state and a questions map with per-type criteria, and answers in Jev's shapes. Long option lists need two accommodations. Most labels never reach the top-k, and the model's mass sits on tokens spelling the option name, so the server sends the union of its label ids as logprob_token_ids. Past ten questions the answer template drops the separator ("0yes 1no") to fit more rows in one read. A schema whose answer lines exceed the canvas splits into chunks, each listing its own questions, because a chunk answering a subset of a longer list loses alignment. "sequential" runs them in order and prefills earlier answers, taking coupled questions from 17/24 to 24/24 consistent pairs at about twice the latency, so it is off by default. depends_on and ask_if generalise that into stages and skip a question whose condition failed. "think": N writes a thought first and reads with it in the prompt. Images arrive as OpenAI image parts, multipart file parts or data URLs; with TEST_PAGE=1, GET / serves a playground for sending them. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 Signed-off-by: Matt Mastracci --- examples/features/diffusion_reads/README.md | 65 + .../features/diffusion_reads/playground.html | 223 +++ .../diffusion_reads/structured_server.py | 1532 +++++++++++++++++ examples/features/diffusion_reads/walk.html | 193 +++ 4 files changed, 2013 insertions(+) create mode 100644 examples/features/diffusion_reads/README.md create mode 100644 examples/features/diffusion_reads/playground.html create mode 100644 examples/features/diffusion_reads/structured_server.py create mode 100644 examples/features/diffusion_reads/walk.html diff --git a/examples/features/diffusion_reads/README.md b/examples/features/diffusion_reads/README.md new file mode 100644 index 000000000000..025ffec25eae --- /dev/null +++ b/examples/features/diffusion_reads/README.md @@ -0,0 +1,65 @@ +# Structured reads on DiffusionGemma + +A discrete diffusion model denoises a whole canvas per forward pass. If the +canvas is seeded with the answer's fixed text and only the answer slots are +left as noise, one denoise step gives a distribution over each slot. These +`extra_args` fields (`vllm_xargs` on the OpenAI server) expose that: + +| field | type | meaning | +| --- | --- | --- | +| `diffusion_seed_canvas` | `list[int]`, exactly `canvas_length` ids | replaces the random initial canvas after prefill | +| `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 | + +`structured_server.py` turns a question schema into those fields. 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. + +```bash +vllm serve google/diffusiongemma-26B-A4B-it \ + --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 --enable-prefix-caching +python examples/features/diffusion_reads/structured_server.py \ + --upstream http://127.0.0.1:8000 --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64 +curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -d '{ + "messages": [ + {"role": "system", "content": "{\"questions\": [{\"id\": \"urgent\", \"type\": \"noul\", \"instructions\": \"Does the customer need a reply within the hour?\"}]}"}, + {"role": "user", "content": "{\"ticket\": \"Everything is down and we have a demo at noon.\"}"} + ]}' +``` + +Question types: `noul` (yes/no), `choice` with `options`, `score` with +ordered `levels`. Each label must be a single token in the answer template, +which the server checks with the tokenizer when a request uses the schema. + +`POST /v1/systemone` implements the Jev decision API. The body holds +`state`, `questions` (a map of id to `type`, `instructions` and `criteria`) +and `model`. Answers come back in that API's shapes: a `noul` probability, +a `choice` with `probabilities` and `confidence`, or a `score` with a +0-indexed `legend`. The schema options above go in the same body as +extensions. + +A question may declare `depends_on` (read in a later stage with those +answers in its prompt), `ask_if` (asked only when a named question's answer +is among the listed ones, otherwise null) and `alone` (a read of its own). + +Images attach as `multipart/form-data`, with the JSON in a part named +`request` and each image as a file part, or as an `images` array of data +URLs. With `TEST_PAGE=1` in the environment, `GET /` serves +`playground.html`, a page for sending requests with an image file or webcam +frames. + +```bash +curl -s localhost:8011/v1/systemone -H 'content-type: application/json' -d '{ + "model": "jev-latest", + "state": {"ticket": "Everything is down and we have a demo at noon."}, + "questions": {"urgent": {"type": "noul", "instructions": "Does the customer need a reply within the hour?"}}}' +``` + +`"think": N` in the schema lets the model write up to N tokens in its +thought channel before the read. The thought is an ordinary generation with +the chat template's thinking marker on, and the read then runs with the +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. diff --git a/examples/features/diffusion_reads/playground.html b/examples/features/diffusion_reads/playground.html new file mode 100644 index 000000000000..a29b0929c708 --- /dev/null +++ b/examples/features/diffusion_reads/playground.html @@ -0,0 +1,223 @@ + + + + + +djev playground + + + +

djev playground

+
+
+
Request body for POST /v1/systemone
+ +
+ + + + +
+
+ + + + + +
+ + + +
+
+
Answers
+
+
+
raw response

+  

+

+
+
+
diff --git a/examples/features/diffusion_reads/structured_server.py b/examples/features/diffusion_reads/structured_server.py
new file mode 100644
index 000000000000..d89de0c96638
--- /dev/null
+++ b/examples/features/diffusion_reads/structured_server.py
@@ -0,0 +1,1532 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Structured decisions in front of a vLLM DiffusionGemma server.
+
+POST /v1/systemone takes Jev's request body: {"model", "state", "questions"}.
+"questions" maps an id to {"type", "instructions", "criteria"}, where "type"
+is "noul", "choice" or "score" and the criteria shape follows the type:
+  noul:   optional {"true": ..., "false": ...} descriptions
+  choice: option name -> description or null
+  score:  ordered list of levels
+Answers take Jev's shapes, with this server's diagnostics alongside:
+  noul:   {"noul": p}
+  choice: {"choice", "probabilities", "confidence"}
+  score:  {"score", "legend", "probabilities", "confidence"}
+The request body may also carry the schema keys "instructions", "samples",
+"auto_max", "auto_threshold", "steps", "think", "ask", "chunk_rows",
+"chunk_prompt" and "sequential" as extensions. Images go ahead of the
+state, either as multipart/form-data with the JSON body in a part named
+"request" and each image as a file part, or as an "images" array of data
+URLs in the JSON body.
+
+POST /v1/chat/completions makes the same decision from an OpenAI-shaped
+call. The system message is the schema JSON below and the user message is
+the state. The reply's `content` is the JSON answer set.
+
+POST /v1/raw/chat/completions passes the body to vLLM's chat completions
+unchanged, for plain generation through this port.
+
+With API_KEY set in the environment, every POST needs "Authorization:
+Bearer ". With TEST_PAGE=1, GET / serves playground.html, a form for
+the request JSON with an image file or webcam frames attached, and GET
+/walk serves walk.html, a phone page that streams the back camera and
+reads one hazard label per frame. Browsers open the webcam only on
+a secure origin, so --tls-port adds an HTTPS listener with a self-signed
+certificate kept in --cert-dir.
+
+Each answer is one distribution per question, from one denoise step over a
+seeded canvas, averaged over a few noise draws. This server handles the
+canvas, tokenizer, slot resolution, noise draws and averaging.
+
+Schema (system message):
+  {"questions": [
+     {"id": "urgent", "type": "noul", "instructions": "..."},
+     {"id": "bucket", "type": "choice", "instructions": "...",
+      "options": [{"name": "billing", "description": "..."}, ...]},
+     {"id": "tone", "type": "score", "instructions": "...",
+      "levels": ["calm", "annoyed", "furious"]}],
+   "instructions": "optional context",
+   "samples": "auto" | N, "auto_threshold": 0.1, "auto_max": 4,
+   "steps": 1, "think": 0}
+
+A question may also declare:
+  "depends_on": [ids]        answered after those, with their answers in
+                             its prompt
+  "ask_if": {id: [answers]}  asked only when that question's answer is
+                             among them (a skipped answer is null)
+  "alone": true              a read of its own
+Questions run in stages by these dependencies. Each stage is one joint
+read. Later stages continue the earlier answers, prefilled for a text
+state and restated for an image.
+
+Up to ten questions answer as "id: label" lines. Past that the id runs
+straight into the label, space separated, one row fewer per question. The
+server splits a schema whose answer template does not fit the canvas into
+chunks that run together, each with its own question list. "chunk_rows"
+sets the rows per chunk, "ask" picks a subset of question ids for one
+read, "sequential": true runs the chunks in order with the earlier answers
+prefilled, and "chunk_prompt": "shared" lists every question in each
+chunk's prompt. "think": N lets the model write up to N tokens in its
+thought channel, as an ordinary generation, and the read then runs with
+that thought in its prompt. With images the model writes the thought with
+the image in view and the server seeds it into the canvas ahead of the
+answer, so the canvas bounds it. The noise draws of a decision share one
+thought.
+
+Serve the model with a canvas that holds the answer template, for example:
+  vllm serve google/diffusiongemma-26B-A4B-it \
+      --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 \
+      --enable-prefix-caching
+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
+"""
+
+import argparse
+import json
+import math
+import os
+import random
+import ssl
+import subprocess
+import threading
+import time
+import urllib.error
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import pybase64 as base64
+from transformers import AutoTokenizer
+
+ARGS = None
+TOK = None
+TEST_PAGE = os.environ.get("TEST_PAGE", "") == "1"  # serve the playground at /
+API_KEY = os.environ.get(
+    "API_KEY", ""
+)  # when set, POST routes need "Authorization: Bearer "
+PAGES = {  # served with TEST_PAGE=1
+    "/": "playground.html",
+    "/playground": "playground.html",
+    "/walk": "walk.html",
+}
+CANVAS_LEN = 64  # the served canvas length. A request may be narrower.
+CANVAS_STEP = 16  # request widths are multiples of this
+VOCAB = 262144
+TURN_CLOSE = 106
+PAD = 0
+TOPK = 20
+# the empty thought block the chat template leaves to the model
+SCAFFOLD_TEXT = "<|channel>thought\n"
+SCAFFOLD = None
+THOUGHT_OPEN = None
+THOUGHT_CLOSE = None
+
+
+# ----------------------------------------------------------------------------
+# Schema
+# ----------------------------------------------------------------------------
+
+
+class SchemaError(ValueError):
+    pass
+
+
+def parse_schema(value):
+    if (
+        not isinstance(value, dict)
+        or not isinstance(value.get("questions"), list)
+        or not value["questions"]
+    ):
+        raise SchemaError("schema: needs a non-empty questions array")
+    qs = []
+    seen = set()
+    for q in value["questions"]:
+        qid = str(q.get("id", "")).strip()
+        if not qid or ":" in qid or "\n" in qid:
+            raise SchemaError(
+                f"question id {qid!r} must be non-empty, no ':' or newline"
+            )
+        if qid in seen:
+            raise SchemaError(f"duplicate question id {qid!r}")
+        seen.add(qid)
+        kind = q.get("type")
+        if kind in ("noul", "bool", "boolean"):
+            kind = "noul"
+            crit = q.get("criteria") or {}
+            choices = [("yes", crit.get("true")), ("no", crit.get("false"))]
+            labels = ["yes", "no"]
+        elif kind == "choice":
+            opts = q.get("options") or []
+            choices = [
+                (o["name"], o.get("description"))
+                if isinstance(o, dict)
+                else (str(o), None)
+                for o in opts
+            ]
+            labels = [chr(ord("A") + i) for i in range(len(choices))]
+        elif kind == "score":
+            choices = [(str(level), None) for level in (q.get("levels") or [])]
+            labels = (
+                [str(i + 1) for i in range(len(choices))]
+                if len(choices) <= 9
+                else [chr(ord("A") + i) for i in range(len(choices))]
+            )
+        else:
+            raise SchemaError(f"question {qid!r}: unknown type {kind!r}")
+        if len(choices) < 2:
+            raise SchemaError(f"question {qid!r}: needs at least two alternatives")
+        if len(choices) > 26:
+            raise SchemaError(f"question {qid!r}: at most 26 alternatives")
+        deps = q.get("depends_on") or []
+        ask_if = q.get("ask_if") or {}
+        if not isinstance(deps, list) or not all(isinstance(d, str) for d in deps):
+            raise SchemaError(
+                f"question {qid!r}: depends_on must be a list of question ids"
+            )
+        if not isinstance(ask_if, dict) or not all(
+            isinstance(v, list) and v for v in ask_if.values()
+        ):
+            raise SchemaError(
+                f"question {qid!r}: ask_if must map a question id to a "
+                "non-empty list of its answers"
+            )
+        qs.append(
+            {
+                "id": qid,
+                "type": kind,
+                "instructions": str(q.get("instructions", "")),
+                "choices": choices,
+                "labels": labels,
+                "depends_on": list(dict.fromkeys(list(deps) + list(ask_if))),
+                "ask_if": ask_if,
+                "alone": bool(q.get("alone", False)),
+            }
+        )
+    by_id = {q["id"]: q for q in qs}
+    for q in qs:
+        for dep in q["depends_on"]:
+            if dep not in by_id or dep == q["id"]:
+                raise SchemaError(
+                    f"question {q['id']!r}: depends on unknown question {dep!r}"
+                )
+        for dep, vals in q["ask_if"].items():
+            names = [c[0] for c in by_id[dep]["choices"]]
+            if any(v not in names for v in vals):
+                raise SchemaError(
+                    f"question {q['id']!r}: ask_if values for {dep!r} "
+                    f"must be among {names}"
+                )
+    schedule(qs)  # refuses a cycle
+    samples = value.get("samples", "auto")
+    if samples == "auto":
+        policy = {
+            "mode": "auto",
+            "max": int(value.get("auto_max", 4)),
+            "threshold": float(value.get("auto_threshold", 0.1)),
+        }
+    elif isinstance(samples, int) and samples >= 1:
+        policy = {"mode": "fixed", "n": min(samples, 32)}
+    else:
+        raise SchemaError('schema: samples must be a positive count or "auto"')
+    ask = value.get("ask")
+    if ask is not None:
+        if not isinstance(ask, list) or not ask or any(a not in seen for a in ask):
+            raise SchemaError("schema: ask must list question ids from this schema")
+        for q in qs:
+            if q["id"] in ask and any(d not in ask for d in q["depends_on"]):
+                raise SchemaError(
+                    f"schema: ask names {q['id']!r} but not everything it depends on"
+                )
+    chunk_rows = value.get("chunk_rows")
+    if chunk_rows is not None and (not isinstance(chunk_rows, int) or chunk_rows < 8):
+        raise SchemaError("schema: chunk_rows must be an integer of at least 8")
+    chunk_prompt = value.get("chunk_prompt", "own")
+    if chunk_prompt not in ("shared", "own"):
+        raise SchemaError('schema: chunk_prompt must be "shared" or "own"')
+    sequential = bool(value.get("sequential", False))
+    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")
+    return {
+        "questions": qs,
+        "instructions": value.get("instructions"),
+        "policy": policy,
+        "steps": max(1, min(int(value.get("steps", 1)), 8)),
+        "think": think,
+        "ask": ask,
+        "chunk_rows": chunk_rows,
+        "chunk_prompt": chunk_prompt,
+        "sequential": sequential,
+        "format": "lines" if len(qs) <= 10 else "indexed",
+    }
+
+
+# Answer template shape: (join between questions, what precedes the label,
+# reply instruction). A small schema gets "lines", which is readable.
+# "indexed" ("0yes 1no") costs three tokens a question against four or five
+# and agreed with "lines" on every set tried: 42 booleans, ten 26-way
+# choices, twenty 5-level scores. Past ten questions the saved rows keep a
+# schema in one read. Two tokens a question, or no id at all, loses
+# alignment beyond about twenty questions, because the id ties a label to
+# its question.
+FORMATS = {
+    "lines": (
+        "\n",
+        "{id}: ",
+        'Reply with one line per question, in this order, formatted as "id: label".',
+    ),
+    "indexed": (
+        " ",
+        "{id}",
+        "Reply on one line with each question's id immediately followed by its "
+        "label, separated by single spaces.",
+    ),
+}
+
+
+def system_text(schema, chunked=False):
+    s = (
+        "Answer a fixed set of questions about the state the user provides. "
+        "Each question lists its allowed answers; reply with exactly one label "
+        "per question.\n"
+    )
+    if schema.get("instructions"):
+        s += "\n" + str(schema["instructions"]).strip() + "\n"
+    for q in schema["questions"]:
+        s += f"\nQuestion {q['id']}: {q['instructions'].strip()}\n"
+        for (name, desc), label in zip(q["choices"], q["labels"]):
+            if q["type"] == "noul":
+                s += f"  {label}: {str(desc).strip()}\n" if desc else f"  {label}\n"
+            elif desc:
+                s += f"  {label}: {name} ({str(desc).strip()})\n"
+            else:
+                s += f"  {label}: {name}\n"
+    s += "\n" + FORMATS[schema.get("format", "lines")][2]
+    if chunked:
+        s += (
+            " A reply may cover only some of the questions; answer every line "
+            "that is present."
+        )
+    return s
+
+
+def answer_text(qs, labels, fmt="lines"):
+    join, lead, _ = FORMATS[fmt]
+    return join.join(
+        lead.format(id=q["id"]) + q["labels"][i] for q, i in zip(qs, labels)
+    )
+
+
+def enc(text):
+    return TOK.encode(text, add_special_tokens=False)
+
+
+def init_tokenizer(tok):
+    global TOK, SCAFFOLD, THOUGHT_OPEN, THOUGHT_CLOSE
+    TOK = tok
+    THOUGHT_OPEN = enc("<|channel>thought\n")
+    THOUGHT_CLOSE = enc("")
+    SCAFFOLD = enc(SCAFFOLD_TEXT)
+    assert THOUGHT_OPEN + THOUGHT_CLOSE == SCAFFOLD, (
+        "the thought tags must tokenize apart"
+    )
+
+
+def resolve_template(qs, head, lead, fmt):
+    """Tokenize the answer template and find each question's slot. Every label
+    must change exactly one token, at the same position for all of a question's
+    labels, or this raises SchemaError. ``head`` is the token run the canvas
+    starts with: the empty thought block for a plain read, and empty when the
+    prompt already ends the thought channel. ``lead`` is the text before the
+    first answer: the join when earlier answers are in the prompt, so the
+    tokens match one joint template."""
+    base_labels = [0] * len(qs)
+    base = head + enc(lead + answer_text(qs, base_labels, fmt))
+    if len(base) + 1 > CANVAS_LEN:
+        raise SchemaError(
+            f"answer template is {len(base)} tokens; the canvas holds {CANVAS_LEN - 1}"
+        )
+    if len(qs) == 1 and len(base) + 1 > CANVAS_LEN:
+        raise SchemaError(
+            f"question {qs[0]['id']!r} alone needs {len(base) + 1} canvas rows"
+        )
+    slots = []
+    for qi, q in enumerate(qs):
+        pos = None
+        ids = [0] * len(q["labels"])
+        for li in range(1, len(q["labels"])):
+            labels = list(base_labels)
+            labels[qi] = li
+            e = head + enc(lead + answer_text(qs, labels, fmt))
+            if len(e) != len(base):
+                raise SchemaError(
+                    f"question {q['id']!r}: label {q['labels'][li]!r} is not a "
+                    "single token"
+                )
+            diffs = [i for i in range(len(e)) if e[i] != base[i]]
+            if len(diffs) != 1 or (pos is not None and diffs[0] != pos):
+                raise SchemaError(
+                    f"question {q['id']!r}: labels do not share one template slot"
+                )
+            pos = diffs[0]
+            ids[li] = e[pos]
+        ids[0] = base[pos]
+        if len(set(ids)) != len(ids):
+            raise SchemaError(
+                f"question {q['id']!r}: two labels tokenize to the same id"
+            )
+        slots.append({"pos": pos, "label_ids": ids})
+    return base, slots
+
+
+_template_cache = {}
+
+
+def template_for(schema, head, lead):
+    fmt = schema.get("format", "lines")
+    key = json.dumps(
+        [head, lead, fmt] + [(q["id"], q["labels"]) for q in schema["questions"]]
+    )
+    if key not in _template_cache:
+        _template_cache[key] = resolve_template(schema["questions"], head, lead, fmt)
+    return _template_cache[key]
+
+
+# ----------------------------------------------------------------------------
+# Reads
+# ----------------------------------------------------------------------------
+
+
+def canvas_width(template):
+    """Smallest multiple of CANVAS_STEP that holds the template and the turn close."""
+    need = len(template) + 1
+    return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP)
+
+
+def build_canvas(template, slots, seed):
+    rng = random.Random(seed)
+    canvas = list(template) + [TURN_CLOSE]
+    canvas += [PAD] * (canvas_width(template) - len(canvas))
+    for s in slots:
+        canvas[s["pos"]] = 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.
+
+
+def upstream_chat(body, timeout=600):
+    req = urllib.request.Request(
+        ARGS.upstream.rstrip("/") + "/v1/chat/completions",
+        data=json.dumps(body).encode(),
+        headers={"content-type": "application/json"},
+    )
+    return json.load(urllib.request.urlopen(req, timeout=timeout))
+
+
+def upstream_completions(body, timeout=600):
+    req = urllib.request.Request(
+        ARGS.upstream.rstrip("/") + "/v1/completions",
+        data=json.dumps(body).encode(),
+        headers={"content-type": "application/json"},
+    )
+    return json.load(urllib.request.urlopen(req, timeout=timeout))
+
+
+def chat_prompt_ids(sys_text, state_text, thinking=False):
+    """The prompt the chat endpoint would build, as token ids, ending after
+    the model turn marker. Text states only. ``thinking`` turns the chat
+    template's thinking marker on."""
+    messages = [
+        {"role": "system", "content": sys_text},
+        {"role": "user", "content": state_text},
+    ]
+    out = TOK.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True, enable_thinking=thinking
+    )
+    ids = (
+        out["input_ids"] if hasattr(out, "keys") else out
+    )  # newer transformers return a dict
+    return [int(t) for t in ids]
+
+
+def think(sys_text, state_text, budget):
+    """A read prefix that ends with a thought the model wrote: the chat prompt
+    with thinking on, the open tag, up to ``budget`` generated tokens, the
+    close tag. Returns the prefix and a diagnostics dict for the thought."""
+    prompt = chat_prompt_ids(sys_text, state_text, thinking=True) + THOUGHT_OPEN
+    started = time.time()
+    d = upstream_completions(
+        {
+            "model": ARGS.model,
+            "prompt": prompt,
+            "max_tokens": budget,
+            "logprobs": 0,
+            "return_tokens_as_token_ids": True,
+            "stop_token_ids": THOUGHT_CLOSE,
+        }
+    )
+    ids = [int(t.split(":")[1]) for t in d["choices"][0]["logprobs"]["tokens"]]
+    closed = THOUGHT_CLOSE[0] in ids
+    if closed:
+        ids = ids[: ids.index(THOUGHT_CLOSE[0])]
+    info = {
+        "tokens": len(ids),
+        "closed": closed,
+        "ms": (time.time() - started) * 1e3,
+        "text": TOK.decode(ids),
+    }
+    return prompt + ids + THOUGHT_CLOSE, info
+
+
+def think_chat(sys_text, state_content, budget):
+    """A thought written with the image in view: the chat endpoint with
+    thinking on, capped at ``budget`` tokens and cut at the close tag.
+    Returns the thought's token ids and a diagnostics dict."""
+    body = {
+        "model": ARGS.model,
+        "messages": [
+            {"role": "system", "content": sys_text},
+            {"role": "user", "content": state_content},
+        ],
+        "max_tokens": budget,
+        "logprobs": True,
+        "top_logprobs": 0,
+        "return_tokens_as_token_ids": True,
+        "stop_token_ids": THOUGHT_CLOSE,
+        "chat_template_kwargs": {"enable_thinking": True},
+    }
+    started = time.time()
+    d = upstream_chat(body)
+    choice = d["choices"][0]
+    ids = [
+        int(t["token"].split(":")[1])
+        for t in (choice.get("logprobs") or {}).get("content") or []
+    ]
+    if ids[: len(THOUGHT_OPEN)] == THOUGHT_OPEN:  # the model opens the channel itself
+        ids = ids[len(THOUGHT_OPEN) :]
+    closed = THOUGHT_CLOSE[0] in ids
+    if closed:
+        ids = ids[: ids.index(THOUGHT_CLOSE[0])]
+    return ids, {
+        "tokens": len(ids),
+        "closed": closed,
+        "ms": (time.time() - started) * 1e3,
+        "text": TOK.decode(ids),
+    }
+
+
+def one_read(
+    schema, template, slots, sys_text, state_content, seed, prefix=None, thinking=False
+):
+    if prefix is not None:
+        return one_read_continuation(schema, template, slots, prefix, seed)
+    messages = [
+        {"role": "system", "content": sys_text},
+        {"role": "user", "content": state_content},
+    ]
+    body = {
+        "model": ARGS.model,
+        "messages": messages,
+        "max_tokens": len(template) + 1,
+        "logprobs": True,
+        "top_logprobs": TOPK,
+        # Exact logprobs for every label at every position. With a long
+        # option list most labels never rank in the top-k, and the model's
+        # mass sits on tokens that spell the option name instead.
+        "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,
+        },
+    }
+    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", {})
+
+
+def slot_distribution(top, label_ids):
+    """Label probabilities at one slot from the returned logprobs: every
+    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."""
+    floor = min(top.values()) - 5.0
+    lp_t = [top.get(i, floor) for i in label_ids]
+    mx = max(lp_t)
+    ex = [math.exp(x - mx) for x in lp_t]
+    probs = [e / sum(ex) for e in ex]
+    top_p = [math.exp(v) for v in top.values()]
+    return {
+        "probs": probs,
+        "label_mass": sum(math.exp(x) for x in lp_t),
+        "entropy": -sum(p * math.log(p) for p in top_p if p > 0),
+        "argmax_is_label": max(top, key=top.get) in label_ids,
+    }
+
+
+def one_read_continuation(schema, template, slots, prompt_ids, seed):
+    """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 = {
+        "model": ARGS.model,
+        "prompt": prompt_ids,
+        "max_tokens": len(template) + 1,
+        "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,
+        },
+    }
+    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", {})
+
+
+def read_many(
+    schema,
+    template,
+    slots,
+    sys_text,
+    state_content,
+    seed,
+    n,
+    prefix=None,
+    thinking=False,
+):
+    results = [None] * n
+    errors = [None] * n
+    usages = [None] * n
+
+    def run(k):
+        try:
+            results[k], usages[k] = one_read(
+                schema,
+                template,
+                slots,
+                sys_text,
+                state_content,
+                seed + k * 7919,
+                prefix,
+                thinking,
+            )
+        except Exception as e:  # raised again below
+            errors[k] = e
+
+    threads = [threading.Thread(target=run, args=(k,)) for k in range(n)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+    for e in errors:
+        if e is not None:
+            raise e
+    return results, usages
+
+
+def schedule(qs):
+    """Questions in stages: a question's stage comes after the stages of
+    everything it depends on. Declaration order is kept within a stage."""
+    ids = {q["id"] for q in qs}
+    pending = list(qs)
+    done: set = set()
+    levels = []
+    while pending:
+        level = [
+            q
+            for q in pending
+            if all(d in done or d not in ids for d in q["depends_on"])
+        ]
+        if not level:
+            raise SchemaError(
+                "schema: dependency cycle among " + ", ".join(q["id"] for q in pending)
+            )
+        levels.append(level)
+        done |= {q["id"] for q in level}
+        pending = [q for q in pending if q["id"] not in done]
+    return levels
+
+
+def chunk_groups(schema, qs):
+    """``qs`` split, in order, into the fewest groups whose answer templates
+    fit ``chunk_rows`` (the canvas by default). A question marked alone gets
+    its own group."""
+    limit = schema.get("chunk_rows") or CANVAS_LEN
+    groups, group = [], []
+    for q in qs:
+        if q["alone"]:
+            if group:
+                groups.append(group)
+                group = []
+            groups.append([q])
+            continue
+        trial = group + [q]
+        rows = (
+            len(SCAFFOLD)
+            + len(
+                enc(answer_text(trial, [0] * len(trial), schema.get("format", "lines")))
+            )
+            + 1
+        )
+        if rows > limit and group:
+            groups.append(group)
+            group = [q]
+        else:
+            group = trial
+    if group:
+        groups.append(group)
+    return groups
+
+
+def answer_name(q, a):
+    """The answer as the name ask_if compares against: yes or no, an option
+    name, or a level name."""
+    if a is None:
+        return None
+    return a["label"] if q["type"] == "noul" else a.get("choice", a.get("level"))
+
+
+def decide(schema, state_content, seed):
+    """One decision. Questions run in stages by their dependencies. A stage
+    is one joint read, chunked by the canvas, with a question marked alone
+    in its own read. Later stages condition on every earlier answer: a
+    prefilled continuation for a text state, or the answers restated in the
+    state for an image. A question whose ask_if condition failed is skipped
+    and its answer is null."""
+    started = time.time()
+    qs = [
+        q
+        for q in schema["questions"]
+        if not schema.get("ask") or q["id"] in schema["ask"]
+    ]
+    levels = schedule(qs)
+    text_state = isinstance(state_content, str)
+    fmt = schema["format"]
+    join = FORMATS[fmt][0]
+    # More than one read in sequence needs the full question list in every
+    # prompt, so that later reads continue one answer.
+    chained = len(levels) > 1 or schema["sequential"]
+    sys_full = system_text(schema, chunked=not text_state and chained)
+    base_ids, thought = None, None
+    if chained and text_state:
+        if schema["think"]:
+            base_ids, thought = think(sys_full, state_content, schema["think"])
+        else:
+            base_ids = chat_prompt_ids(sys_full, state_content) + SCAFFOLD
+    shared = schema["chunk_prompt"] == "shared"
+    answered, lines, earlier = {}, [], []
+    parts, stages, chunks, skipped = [], [], [], {}
+    by_id = {q["id"]: q for q in schema["questions"]}
+
+    def run(group, k, conditioned):
+        sub = dict(schema, questions=group)
+        # The thought is written once: in the prefix of a chained text decision,
+        # or in the first read of anything else.
+        sub["think"] = 0 if (chained and text_state) or conditioned else schema["think"]
+        if text_state:
+            if conditioned:
+                prefix, lead, sys_text = (
+                    base_ids + enc(join.join(lines)),
+                    join,
+                    sys_full,
+                )
+            elif chained:
+                prefix, lead, sys_text = (base_ids if thought else None), "", sys_full
+            else:
+                prefix, lead, sys_text = (
+                    None,
+                    "",
+                    (system_text(schema, chunked=True) if shared else system_text(sub)),
+                )
+            return decide_group(
+                sub, sys_text, state_content, seed + 104729 * k, prefix, lead
+            )
+        state = state_content
+        if conditioned:
+            text = next(
+                (p["text"] for p in state_content if p.get("type") == "text"), ""
+            )
+            text += "\n\nAnswers so far:\n" + "\n".join(earlier)
+            state = [p for p in state_content if p.get("type") != "text"] + [
+                {"type": "text", "text": text}
+            ]
+            sys_text = sys_full
+        else:
+            sys_text = sys_full if (chained or shared) else system_text(sub)
+        return decide_group(sub, sys_text, state, seed + 104729 * k)
+
+    def absorb(group, body, rows):
+        parts.append((body, rows))
+        chunks.append([q["id"] for q in group])
+        answered.update(body["answers"])
+        lines.append(
+            answer_text(
+                group,
+                [q["labels"].index(body["answers"][q["id"]]["label"]) for q in group],
+                fmt,
+            )
+        )
+        earlier.extend(
+            f"{q['id']}: {answer_name(q, body['answers'][q['id']])}" for q in group
+        )
+
+    k = 0
+    for level in levels:
+        asked = []
+        for q in level:
+            failed = next(
+                (
+                    (dep, vals)
+                    for dep, vals in q["ask_if"].items()
+                    if answer_name(by_id[dep], answered.get(dep)) not in vals
+                ),
+                None,
+            )
+            if failed:
+                answered[q["id"]] = None
+                skipped[q["id"]] = {
+                    "because": failed[0],
+                    "was": answer_name(by_id[failed[0]], answered.get(failed[0])),
+                    "wanted": failed[1],
+                }
+                continue
+            asked.append(q)
+        if not asked:
+            continue
+        stages.append([q["id"] for q in asked])
+        groups = chunk_groups(schema, asked)
+        conditioned = bool(lines)
+        if schema["sequential"] or len(groups) == 1:
+            for group in groups:
+                body, rows = run(
+                    group, k, conditioned or (schema["sequential"] and bool(lines))
+                )
+                absorb(group, body, rows)
+                k += 1
+        else:
+            with ThreadPoolExecutor(max_workers=len(groups)) as ex:
+                results = list(
+                    ex.map(
+                        lambda gk, conditioned=conditioned: run(
+                            gk[1], gk[0], conditioned
+                        ),
+                        [(k + i, g) for i, g in enumerate(groups)],
+                    )
+                )
+            for group, (body, rows) in zip(groups, results):
+                absorb(group, body, rows)
+            k += len(groups)
+
+    answers = {q["id"]: answered.get(q["id"]) for q in qs}
+    diag_q = {}
+    for body, _ in parts:
+        diag_q.update(body["diagnostics"]["questions"])
+    # A thought written here (a chained text decision) is outside every
+    # group's row count. One written inside a group is already counted there.
+    extra_rows = thought["tokens"] if thought else 0
+    if thought is None:
+        thoughts = [b["diagnostics"].get("thought") for b, _ in parts]
+        thought = (
+            thoughts[0] if len(parts) == 1 else ([t for t in thoughts if t] or None)
+        )
+    one = len(parts) == 1 and not skipped
+    diagnostics = {
+        "steps": schema["steps"],
+        "stages": stages,
+        "skipped": skipped,
+        "chunks": chunks,
+        "chunk_prompt": "full" if chained else schema["chunk_prompt"],
+        "sequential": schema["sequential"],
+        "conditioning": (
+            None
+            if len(stages) <= 1 and not schema["sequential"]
+            else ("prefill" if text_state else "restated")
+        ),
+        "thought": thought,
+        "samples": (
+            parts[0][0]["diagnostics"]["samples"]
+            if one
+            else {
+                "n": [b["diagnostics"]["samples"]["n"] for b, _ in parts],
+                "tops": [b["diagnostics"]["samples"]["tops"] for b, _ in parts],
+                "policy": [b["diagnostics"]["samples"]["policy"] for b, _ in parts],
+            }
+        ),
+        "timing": {
+            "total_ms": (time.time() - started) * 1e3,
+            "reads": sum(b["diagnostics"]["timing"]["reads"] for b, _ in parts),
+        },
+        "prompt_tokens": max(
+            (b["diagnostics"].get("prompt_tokens") or 0) for b, _ in parts
+        )
+        or None,
+        "questions": diag_q,
+        "engine": "vllm",
+    }
+    return {"answers": answers, "diagnostics": diagnostics}, sum(
+        rows for _, rows in parts
+    ) + extra_rows
+
+
+def decide_group(schema, sys_text, state_content, seed, prefix=None, lead=""):
+    started = time.time()
+    thought = None
+    head = SCAFFOLD if prefix is None else []
+    thinking = False
+    if prefix is None and schema["think"]:
+        if isinstance(state_content, str):
+            prefix, thought = think(sys_text, state_content, schema["think"])
+            head = []
+        else:
+            # With images the model writes the thought with the image in view, and
+            # the server seeds it into the canvas ahead of the answer, so the read
+            # keeps the image. The canvas bounds the thought.
+            answer = enc(
+                answer_text(
+                    schema["questions"],
+                    [0] * len(schema["questions"]),
+                    schema.get("format", "lines"),
+                )
+            )
+            fits = CANVAS_LEN - 1 - len(THOUGHT_OPEN) - len(THOUGHT_CLOSE) - len(answer)
+            if fits < 8:
+                raise SchemaError(
+                    f"think: the canvas leaves {fits} rows for a thought "
+                    "beside this template"
+                )
+            ids, thought = think_chat(
+                sys_text, state_content, min(schema["think"], fits)
+            )
+            thought["budget"] = min(schema["think"], fits)
+            head = THOUGHT_OPEN + ids + THOUGHT_CLOSE
+            thinking = True
+    template, slots = template_for(schema, head, lead)
+    policy = schema["policy"]
+    if policy["mode"] == "fixed":
+        reads, usages = read_many(
+            schema,
+            template,
+            slots,
+            sys_text,
+            state_content,
+            seed,
+            policy["n"],
+            prefix,
+            thinking,
+        )
+        extended = None
+        first_entropy = None
+    else:
+        reads, usages = read_many(
+            schema, template, slots, sys_text, state_content, seed, 1, prefix, thinking
+        )
+        first_entropy = {
+            q["id"]: r["entropy"] for q, r in zip(schema["questions"], reads[0])
+        }
+        extended = (
+            max(first_entropy.values()) > policy["threshold"] and policy["max"] > 1
+        )
+        if extended:
+            more, more_usages = read_many(
+                schema,
+                template,
+                slots,
+                sys_text,
+                state_content,
+                seed + 1,
+                policy["max"] - 1,
+                prefix,
+                thinking,
+            )
+            reads += more
+            usages += more_usages
+    prompt_tokens = next(
+        (u["prompt_tokens"] for u in usages if u and u.get("prompt_tokens")), None
+    )
+    elapsed_ms = (time.time() - started) * 1e3
+
+    answers = {}
+    diag_q = {}
+    n = len(reads)
+    for qi, q in enumerate(schema["questions"]):
+        per = [r[qi]["probs"] for r in reads]
+        mean = [sum(p[i] for p in per) / n for i in range(len(q["labels"]))]
+        top = max(range(len(mean)), key=lambda i: mean[i])
+        a = {
+            "type": q["type"],
+            "label": q["labels"][top],
+            "confidence": mean[top],
+            "probabilities": {c[0]: m for c, m in zip(q["choices"], mean)},
+        }
+        if q["type"] == "noul":
+            a["noul"] = mean[0]
+        elif q["type"] == "choice":
+            a["choice"] = q["choices"][top][0]
+        else:
+            a["score"] = sum((i + 1) * m for i, m in enumerate(mean))
+            a["level"] = q["choices"][top][0]
+        if n > 1:
+            var = sum((p[top] - mean[top]) ** 2 for p in per) / (n - 1)
+            a["stderr"] = (var / n) ** 0.5
+            a["agreement"] = (
+                sum(1 for p in per if max(range(len(p)), key=lambda i: p[i]) == top) / n
+            )
+        answers[q["id"]] = a
+        diag_q[q["id"]] = {
+            "pos": slots[qi]["pos"],
+            "entropy": [r[qi]["entropy"] for r in reads],
+            "label_mass": reads[0][qi]["label_mass"],
+            "argmax_is_label": reads[0][qi]["argmax_is_label"],
+        }
+    tops = [
+        {
+            q["id"]: [
+                q["labels"][
+                    max(range(len(r[qi]["probs"])), key=lambda i: r[qi]["probs"][i])
+                ],
+                max(r[qi]["probs"]),
+                r[qi]["entropy"],
+            ]
+            for qi, q in enumerate(schema["questions"])
+        }
+        for r in reads
+    ]
+    return {
+        "answers": answers,
+        "diagnostics": {
+            "steps": schema["steps"],
+            "samples": {
+                "n": n,
+                "tops": tops,
+                "policy": dict(
+                    policy, extended=extended, first_read_entropy=first_entropy
+                ),
+            },
+            "timing": {"total_ms": elapsed_ms, "reads": n},
+            "thought": thought,
+            "prompt_tokens": prompt_tokens,
+            "questions": diag_q,
+            "engine": "vllm",
+        },
+    }, len(template) + 1 + (thought["tokens"] if thought else 0)
+
+
+# ----------------------------------------------------------------------------
+# Jev's contract
+# ----------------------------------------------------------------------------
+
+JEV_EXTENSIONS = (
+    "instructions",
+    "samples",
+    "auto_max",
+    "auto_threshold",
+    "steps",
+    "think",
+    "ask",
+    "chunk_rows",
+    "chunk_prompt",
+    "sequential",
+)
+
+
+def jev_schema(body):
+    """This server's schema from a Jev request body."""
+    qs = body.get("questions")
+    if not isinstance(qs, dict) or not qs:
+        raise SchemaError("questions: needs a non-empty map of id -> question")
+    out = []
+    for qid, q in qs.items():
+        if not isinstance(q, dict):
+            raise SchemaError(f"question {qid!r}: must be an object")
+        kind, crit, ins = q.get("type"), q.get("criteria"), q.get("instructions", "")
+        item = {
+            "id": qid,
+            "type": kind,
+            "instructions": ins if isinstance(ins, str) else json.dumps(ins),
+        }
+        if kind == "noul":
+            if crit is not None and not isinstance(crit, dict):
+                raise SchemaError(
+                    f"question {qid!r}: noul criteria must be an object with "
+                    "true and false"
+                )
+            item["criteria"] = crit
+        elif kind == "choice":
+            if not isinstance(crit, dict) or not crit:
+                raise SchemaError(
+                    f"question {qid!r}: choice criteria must map option names "
+                    "to descriptions"
+                )
+            item["options"] = [
+                {"name": str(n), "description": d} for n, d in crit.items()
+            ]
+        elif kind == "score":
+            if not isinstance(crit, list):
+                raise SchemaError(
+                    f"question {qid!r}: score criteria must be an ordered list "
+                    "of levels"
+                )
+            item["levels"] = crit
+        else:
+            raise SchemaError(f"question {qid!r}: unknown type {kind!r}")
+        for key in ("depends_on", "ask_if", "alone"):
+            if key in q:
+                item[key] = q[key]
+        out.append(item)
+    schema = {k: body[k] for k in JEV_EXTENSIONS if k in body}
+    schema["questions"] = out
+    return parse_schema(schema)
+
+
+def image_part(content_type, data):
+    return {
+        "type": "image_url",
+        "image_url": {
+            "url": f"data:{content_type};base64," + base64.b64encode(data).decode()
+        },
+    }
+
+
+def jev_images(value):
+    """Image parts from the body's "images": data URLs, or objects with
+    content_type and base64."""
+    parts = []
+    for i, im in enumerate(value or []):
+        if isinstance(im, str) and im.startswith("data:image/"):
+            parts.append({"type": "image_url", "image_url": {"url": im}})
+        elif (
+            isinstance(im, dict)
+            and str(im.get("content_type", "")).startswith("image/")
+            and isinstance(im.get("base64"), str)
+        ):
+            parts.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:{im['content_type']};base64,{im['base64']}"
+                    },
+                }
+            )
+        else:
+            raise SchemaError(
+                f"images[{i}]: a data:image/... URL or an object with "
+                "content_type and base64"
+            )
+    return parts
+
+
+def jev_state(body, image_parts=()):
+    """The user message: the state as text (as given, or as JSON), with any
+    images ahead of it."""
+    state = body.get("state")
+    if state is None:
+        raise SchemaError("state: required")
+    text = state if isinstance(state, str) else json.dumps(state)
+    if not image_parts:
+        return text
+    return list(image_parts) + [{"type": "text", "text": text}]
+
+
+def jev_answer(q, a):
+    if a is None:
+        return None
+    if q["type"] == "noul":
+        return {"type": "noul", "noul": a["noul"]}
+    if q["type"] == "choice":
+        return {
+            "type": "choice",
+            "choice": a["choice"],
+            "probabilities": a["probabilities"],
+            "confidence": a["confidence"],
+        }
+    names = [c[0] for c in q["choices"]]
+    probs = {str(i): a["probabilities"][n] for i, n in enumerate(names)}
+    return {
+        "type": "score",
+        "score": sum(i * p for i, p in enumerate(probs.values())),
+        "legend": {str(i): n for i, n in enumerate(names)},
+        "probabilities": probs,
+        "confidence": a["confidence"],
+    }
+
+
+# ----------------------------------------------------------------------------
+# HTTP
+# ----------------------------------------------------------------------------
+
+
+def message_text(m):
+    c = m.get("content", "")
+    if isinstance(c, list):
+        return "".join(p.get("text", "") for p in c if isinstance(p, dict))
+    return c if isinstance(c, str) else ""
+
+
+class Handler(BaseHTTPRequestHandler):
+    def log_message(self, fmt, *args):
+        pass
+
+    def _json(self, code, obj):
+        body = json.dumps(obj).encode()
+        self.send_response(code)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_GET(self):
+        if self.path == "/health":
+            return self._json(200, {"status": "ok"})
+        if self.path in PAGES and TEST_PAGE:
+            page_path = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), PAGES[self.path]
+            )
+            with open(page_path, "rb") as fh:
+                body = fh.read()
+            self.send_response(200)
+            self.send_header("content-type", "text/html; charset=utf-8")
+            self.send_header("content-length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+            return
+        return self._json(404, {"error": {"message": "unknown route"}})
+
+    def _read_request(self):
+        """-> (body, image parts): a JSON body, or multipart/form-data with the
+        JSON in a part named request and each image as a file part, in order."""
+        raw = self.rfile.read(int(self.headers.get("content-length", "0")))
+        ctype = self.headers.get("content-type", "")
+        if not ctype.lower().startswith("multipart/form-data"):
+            return json.loads(raw), []
+        from email.parser import BytesParser
+        from email.policy import HTTP
+
+        msg = BytesParser(policy=HTTP).parsebytes(
+            b"Content-Type: " + ctype.encode() + b"\r\n\r\n" + raw
+        )
+        body, images = None, []
+        for part in msg.iter_parts():
+            name = part.get_param("name", header="content-disposition")
+            data = part.get_payload(decode=True)
+            if name == "request":
+                body = json.loads(data)
+            elif part.get_content_type().startswith("image/"):
+                images.append(image_part(part.get_content_type(), data))
+            else:
+                raise ValueError(
+                    f"part {name!r}: neither the request JSON nor an image"
+                )
+        if body is None:
+            raise ValueError(
+                "multipart needs a part named request holding the JSON body"
+            )
+        return body, images
+
+    def do_POST(self):
+        if API_KEY and self.headers.get("authorization", "") != f"Bearer {API_KEY}":
+            return self._json(
+                401,
+                {
+                    "error": {
+                        "message": "missing or wrong bearer token",
+                        "type": "authentication_error",
+                    }
+                },
+            )
+        if self.path == "/v1/raw/chat/completions":
+            return self._raw_chat()
+        try:
+            req, images = self._read_request()
+        except Exception as e:
+            return self._json(
+                400,
+                {
+                    "error": {
+                        "message": f"invalid body: {e}",
+                        "type": "invalid_request_error",
+                    }
+                },
+            )
+        if self.path == "/v1/systemone":
+            return self._systemone(req, images)
+        if self.path == "/v1/chat/completions":
+            return self._chat(req)
+        return self._json(404, {"error": {"message": "unknown route"}})
+
+    def _raw_chat(self):
+        """Pass the body and status through to vLLM's chat completions."""
+        raw = self.rfile.read(int(self.headers.get("content-length", "0")))
+        req = urllib.request.Request(
+            ARGS.upstream.rstrip("/") + "/v1/chat/completions",
+            data=raw,
+            headers={
+                "content-type": self.headers.get("content-type", "application/json")
+            },
+        )
+        try:
+            with urllib.request.urlopen(req, timeout=600) as r:
+                code, body = r.status, r.read()
+        except urllib.error.HTTPError as e:
+            code, body = e.code, e.read()
+        self.send_response(code)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def _decide(self, schema, state, seed):
+        """-> (status, body) with the error body already shaped."""
+        try:
+            return 200, decide(schema, state, seed)
+        except SchemaError as e:
+            return 422, {"error": {"message": str(e), "type": "validation_error"}}
+        except urllib.error.HTTPError as e:
+            return 502, {
+                "error": {
+                    "message": (
+                        f"upstream {e.code}: {e.read()[:300].decode(errors='replace')}"
+                    ),
+                    "type": "server_error",
+                }
+            }
+        except Exception as e:
+            return 500, {"error": {"message": repr(e), "type": "server_error"}}
+
+    def _systemone(self, req, images):
+        try:
+            schema = jev_schema(req)
+            state = jev_state(req, images + jev_images(req.get("images")))
+        except SchemaError as e:
+            return self._json(
+                422, {"error": {"message": str(e), "type": "validation_error"}}
+            )
+        code, result = self._decide(schema, state, int(req.get("seed", 42)))
+        if code != 200:
+            return self._json(code, result)
+        body, completion_tokens = result
+        answers = {
+            q["id"]: jev_answer(q, body["answers"][q["id"]])
+            for q in schema["questions"]
+        }
+        labels = " ".join(
+            f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items()
+        )
+        print(
+            f"systemone: {labels} "
+            f"reads={body['diagnostics']['timing']['reads']} "
+            f"{body['diagnostics']['timing']['total_ms']:.0f}ms",
+            flush=True,
+        )
+        self._json(
+            200,
+            {
+                "model": ARGS.model,
+                "answers": answers,
+                # vLLM's prompt count when the reads reported one, since it covers
+                # images. Otherwise the tokenizer's count of the text prompt.
+                "usage": {
+                    "input_tokens": body["diagnostics"].get("prompt_tokens")
+                    or (
+                        len(chat_prompt_ids(system_text(schema), state))
+                        if isinstance(state, str)
+                        else 0
+                    ),
+                    "output_tokens": completion_tokens,
+                },
+                "diagnostics": body["diagnostics"],
+            },
+        )
+
+    def _chat(self, req):
+        msgs = req.get("messages") or []
+        if (
+            len(msgs) != 2
+            or msgs[0].get("role") not in ("system", "developer")
+            or msgs[1].get("role") != "user"
+        ):
+            return self._json(
+                400,
+                {
+                    "error": {
+                        "message": (
+                            "a structured request is exactly two messages: the "
+                            "schema (system) and the state JSON (user)"
+                        ),
+                        "type": "invalid_request_error",
+                    }
+                },
+            )
+        try:
+            schema_value = json.loads(message_text(msgs[0]))
+            schema = parse_schema(schema_value)
+            content = msgs[1].get("content", "")
+            has_image = isinstance(content, list) and any(
+                isinstance(p, dict) and p.get("type") in ("image_url", "image")
+                for p in content
+            )
+            if has_image:
+                # image parts pass through to vLLM unchanged, with text parts as context
+                state = content
+            else:
+                state = message_text(msgs[1]).strip()
+                json.loads(state)
+        except SchemaError as e:
+            return self._json(
+                400, {"error": {"message": str(e), "type": "invalid_request_error"}}
+            )
+        except Exception as e:
+            return self._json(
+                400,
+                {
+                    "error": {
+                        "message": (
+                            "system must be a JSON question schema and user "
+                            f"must be JSON state or image parts: {e}"
+                        ),
+                        "type": "invalid_request_error",
+                    }
+                },
+            )
+        code, result = self._decide(schema, state, int(req.get("seed", 42)))
+        if code != 200:
+            if code == 422:
+                result["error"]["type"] = "invalid_request_error"
+                code = 400
+            return self._json(code, result)
+        body, completion_tokens = result
+        content = json.dumps(body, indent=2)
+        labels = " ".join(
+            f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items()
+        )
+        print(
+            f"structured: {labels} "
+            f"reads={body['diagnostics']['timing']['reads']} "
+            f"{body['diagnostics']['timing']['total_ms']:.0f}ms",
+            flush=True,
+        )
+        self._json(
+            200,
+            {
+                "id": f"chatcmpl-{int(time.time() * 1000)}",
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": req.get("model", "dgemma-structured"),
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": content},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 0,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": completion_tokens,
+                },
+            },
+        )
+
+
+def self_signed(cert_dir):
+    """Paths of a self-signed certificate and key in cert_dir, made with
+    openssl on first use."""
+    os.makedirs(cert_dir, exist_ok=True)
+    cert, key = os.path.join(cert_dir, "djev.crt"), os.path.join(cert_dir, "djev.key")
+    if not (os.path.exists(cert) and os.path.exists(key)):
+        subprocess.run(
+            [
+                "openssl",
+                "req",
+                "-x509",
+                "-newkey",
+                "rsa:2048",
+                "-nodes",
+                "-days",
+                "3650",
+                "-subj",
+                "/CN=djev",
+                "-keyout",
+                key,
+                "-out",
+                cert,
+            ],
+            check=True,
+            capture_output=True,
+        )
+    return cert, key
+
+
+def serve_tls(host, port, cert_dir):
+    cert, key = self_signed(cert_dir)
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+    ctx.load_cert_chain(cert, key)
+    srv = ThreadingHTTPServer((host, port), Handler)
+    srv.socket = ctx.wrap_socket(srv.socket, server_side=True)
+    threading.Thread(target=srv.serve_forever, daemon=True).start()
+    return srv
+
+
+def main():
+    global ARGS, CANVAS_LEN, CANVAS_STEP
+    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("--canvas", type=int, default=64, help="the served canvas length")
+    p.add_argument(
+        "--canvas-step",
+        type=int,
+        default=16,
+        help="request widths round up to a multiple of this",
+    )
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=8011)
+    p.add_argument(
+        "--tls-port", type=int, default=0, help="also listen with HTTPS here (0 = off)"
+    )
+    p.add_argument(
+        "--cert-dir",
+        default=os.path.expanduser("~/.cache/djev"),
+        help="directory for the self-signed certificate",
+    )
+    ARGS = p.parse_args()
+    CANVAS_LEN = ARGS.canvas
+    CANVAS_STEP = ARGS.canvas_step
+    init_tokenizer(AutoTokenizer.from_pretrained(ARGS.tokenizer))
+    if ARGS.tls_port:
+        serve_tls(ARGS.host, ARGS.tls_port, ARGS.cert_dir)
+        print(
+            f"structured server https on {ARGS.host}:{ARGS.tls_port} (self-signed)",
+            flush=True,
+        )
+    print(
+        f"structured server on {ARGS.host}:{ARGS.port} -> {ARGS.upstream} "
+        f"(canvas {CANVAS_LEN})",
+        flush=True,
+    )
+    ThreadingHTTPServer((ARGS.host, ARGS.port), Handler).serve_forever()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/features/diffusion_reads/walk.html b/examples/features/diffusion_reads/walk.html
new file mode 100644
index 000000000000..d1883250be41
--- /dev/null
+++ b/examples/features/diffusion_reads/walk.html
@@ -0,0 +1,193 @@
+
+
+
+
+
+
+djev walk
+
+
+
+

+  
+  
↑

+  
+

+

+  

+    
+    
+  

+  
 

+  

+  

+    
+  

+  
 

+

+
+
+
+

From 408f00d7844ee265cee506d9a3e7789eebcef908 Mon Sep 17 00:00:00 2001
From: Matt Mastracci 
Date: Sun, 20 Sep 2026 14:48:18 -0600
Subject: [PATCH 15/25] [Core] Pin seeded canvas positions through diffusion
 denoise steps

With diffusion_max_steps above one, a seeded canvas does not survive the
step: accept/renoise replaces every position, the template included, with
the model's sample or a fresh random token, and self-conditioning carries
the model's own prediction at those positions into the next step. A
multi-step read therefore denoised a canvas the model had rewritten rather
than the one the request seeded, and its template positions drifted.

diffusion_pinned (extra_args) lists canvas positions to hold at their seed
value on every denoise step. The compiled step restores them after
accept/renoise and zeroes their self-conditioning soft embed, so a pinned
position looks the same to the model on every step as on the first.
Validation requires a seed canvas and positions inside the request's
canvas. The example server pins the template whenever a read asks for more
than one step, which makes its existing steps option sound.

Measured on a DGX Spark with two dependent choices in one canvas over 24
noise draws: consistent 18/24 at 8 unpinned steps, 24/24 pinned. A slot
seeded to a fixed label is followed by its dependent slot 6/24 times at 8
unpinned steps and 24/24 pinned.

Co-Authored-By: Claude Fable 5.1 
Signed-off-by: Matt Mastracci 
---
 examples/features/diffusion_reads/README.md   |  1 +
 .../diffusion_reads/structured_server.py      | 13 ++++++
 tests/test_sampling_params.py                 | 35 ++++++++++++++-
 tests/v1/sample/test_diffusion_gemma_reads.py | 44 ++++++++++++++++++-
 vllm/model_executor/models/diffusion_gemma.py | 36 ++++++++++++++-
 vllm/sampling_params.py                       | 30 +++++++++++--
 6 files changed, 151 insertions(+), 8 deletions(-)

diff --git a/examples/features/diffusion_reads/README.md b/examples/features/diffusion_reads/README.md
index 025ffec25eae..45cdf70fd84c 100644
--- a/examples/features/diffusion_reads/README.md
+++ b/examples/features/diffusion_reads/README.md
@@ -8,6 +8,7 @@ left as noise, one denoise step gives a distribution over each slot. These
 | field | type | meaning |
 | --- | --- | --- |
 | `diffusion_seed_canvas` | `list[int]`, exactly `canvas_length` ids | replaces the random initial canvas after prefill |
+| `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 |
 
diff --git a/examples/features/diffusion_reads/structured_server.py b/examples/features/diffusion_reads/structured_server.py
index d89de0c96638..439e3e8fa889 100644
--- a/examples/features/diffusion_reads/structured_server.py
+++ b/examples/features/diffusion_reads/structured_server.py
@@ -404,6 +404,17 @@ def canvas_width(template):
     return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP)
 
 
+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."""
+    if steps <= 1:
+        return {}
+    free = {s["pos"] for s in slots}
+    return {
+        "diffusion_pinned": [p for p in range(canvas_width(template)) if p not in free]
+    }
+
+
 def build_canvas(template, slots, seed):
     rng = random.Random(seed)
     canvas = list(template) + [TURN_CLOSE]
@@ -545,6 +556,7 @@ def one_read(
             "diffusion_canvas_length": canvas_width(template),
             "diffusion_max_steps": schema["steps"],
             "diffusion_read_only": True,
+            **pin_xargs(template, slots, schema["steps"]),
         },
     }
     d = upstream_chat(body)
@@ -593,6 +605,7 @@ def one_read_continuation(schema, template, slots, prompt_ids, seed):
             "diffusion_canvas_length": canvas_width(template),
             "diffusion_max_steps": schema["steps"],
             "diffusion_read_only": True,
+            **pin_xargs(template, slots, schema["steps"]),
         },
     }
     d = upstream_completions(body)
diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py
index ee3005ab71d2..472112d5cfcd 100644
--- a/tests/test_sampling_params.py
+++ b/tests/test_sampling_params.py
@@ -90,6 +90,9 @@ def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None):
         ({"diffusion_max_steps": True}, "positive integer"),
         ({"diffusion_read_only": "yes"}, "boolean"),
         ({"diffusion_read_only": 2}, "boolean"),
+        ({"diffusion_pinned": "0,1"}, "list of canvas positions"),
+        ({"diffusion_pinned": [0, True]}, "list of canvas positions"),
+        ({"diffusion_pinned": [0]}, "needs a diffusion_seed_canvas"),
     ],
 )
 def test_diffusion_rejects_bad_extra_args(extra_args: dict, match: str):
@@ -152,11 +155,41 @@ def test_diffusion_canvas_length_sizes_the_seed_and_the_read():
         _verify_diffusion(params, canvas_length=8)
 
 
+def test_diffusion_pinned_positions_stay_inside_the_canvas():
+    seed = [0] * 8
+    with pytest.raises(VLLMValidationError, match="inside the canvas"):
+        _verify_diffusion(
+            SamplingParams(
+                extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [7, 8]}
+            ),
+            canvas_length=8,
+        )
+    with pytest.raises(VLLMValidationError, match="inside the canvas"):
+        _verify_diffusion(
+            SamplingParams(
+                extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [-1]}
+            )
+        )
+    # A narrower request canvas bounds the positions.
+    with pytest.raises(VLLMValidationError, match="inside the canvas"):
+        _verify_diffusion(
+            SamplingParams(
+                extra_args={
+                    "diffusion_canvas_length": 4,
+                    "diffusion_seed_canvas": seed[:4],
+                    "diffusion_pinned": [4],
+                }
+            ),
+            canvas_length=8,
+        )
+
+
 def test_diffusion_accepts_extra_args():
     params = SamplingParams(
         extra_args={
             "diffusion_seed_canvas": list(range(8)),
-            "diffusion_max_steps": 1,
+            "diffusion_pinned": [0, 1, 7],
+            "diffusion_max_steps": 4,
             "diffusion_read_only": True,
         }
     )
diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py
index 2165f6d77e66..a705f4a30200 100644
--- a/tests/v1/sample/test_diffusion_gemma_reads.py
+++ b/tests/v1/sample/test_diffusion_gemma_reads.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Per-request DiffusionGemma state behind structured reads: seed canvases,
-read-only slots and the per-slot step cap."""
+pinned positions, read-only slots and the per-slot step cap."""
 
 import numpy as np
 import pytest
@@ -115,6 +115,7 @@ def _denoise_once(
     slots: list[int],
     compute_sc: bool = True,
     width: int = CL,
+    embed_weight: torch.Tensor | None = None,
 ) -> None:
     """One compiled denoise step over ``slots`` with flat logits, so nothing
     converges by stability or confidence and only the step cap can end it.
@@ -136,11 +137,17 @@ def _denoise_once(
         states.is_encoder_phase,
         states.confident,
         states.self_conditioning_embeds[:, :width],
-        torch.zeros(VOCAB, 4, device=device),
+        (
+            torch.zeros(VOCAB, 4, device=device)
+            if embed_weight is None
+            else embed_weight
+        ),
         torch.tensor(1.0, device=device),
         states.accepted_canvas_history[:, :, :width],
         states.accepted_canvas_history_len,
         states.max_steps,
+        states.pin_mask[:, :width],
+        states.seed_canvas[:, :width],
         torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width],
         torch.zeros(n, dtype=torch.int32, device=device),
         torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device),
@@ -160,6 +167,39 @@ def _denoise_once(
     )
 
 
+def test_pinned_positions_hold_their_seed_through_a_step():
+    states = _states()
+    states.add_request(0)
+    states.is_encoder_phase[0] = False
+    seed = list(range(10, 10 + CL))
+    states.set_seed_canvas(0, seed)
+    states.canvas[0] = torch.tensor(seed, device="cuda")
+    states.set_pins(0, [0, 1, 2, 3])
+
+    # Flat logits accept nothing, so every free position is renoised.
+    _denoise_once(states, [0], embed_weight=torch.ones(VOCAB, 4, device="cuda"))
+
+    assert states.canvas[0, :4].tolist() == seed[:4]
+    # The soft embed is zero at pinned positions and non-zero elsewhere.
+    sc = states.self_conditioning_embeds[0]
+    assert not sc[:4].any()
+    assert sc[4:].any()
+
+
+def test_add_request_clears_pins():
+    states = _states()
+    states.add_request(0)
+    states.set_seed_canvas(0, [1] * CL)
+    states.set_pins(0, [2, 3])
+    assert states.pinned_slots == {0}
+    assert states.pin_mask[0].tolist() == [False, False, True, True] + [False] * 4
+
+    states.add_request(0)
+
+    assert not states.pinned_slots
+    assert not states.pin_mask[0].any()
+
+
 def test_single_step_tile_skips_self_conditioning():
     states = _states()
     states.add_request(0)
diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py
index 958f3da039de..6d3e678e22bb 100644
--- a/vllm/model_executor/models/diffusion_gemma.py
+++ b/vllm/model_executor/models/diffusion_gemma.py
@@ -392,6 +392,8 @@ def _compiled_sample_step(
     history: torch.Tensor,  # [max_num_reqs, ST, CL]
     history_len_tensor: torch.Tensor,  # [max_num_reqs]
     max_steps_tensor: torch.Tensor,  # [max_num_reqs] int32, per-slot step cap
+    pin_mask: torch.Tensor,  # [max_num_reqs, CL] bool, positions held at the seed
+    seed_canvas: torch.Tensor,  # [max_num_reqs, CL]
     # Output tensors (modified in-place)
     sampled: torch.Tensor,  # [num_reqs, CL]
     num_sampled: torch.Tensor,  # [num_reqs]
@@ -480,8 +482,11 @@ def _compiled_sample_step(
         0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype
     )
 
-    # Compute denoise canvas (accept/renoise)
+    # Compute denoise canvas (accept/renoise). Pinned positions keep the seed.
     denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens)
+    denoise_canvas = torch.where(
+        pin_mask[decode_slots], seed_canvas[decode_slots], denoise_canvas
+    )
 
     # Canvas: commit → random reinit, denoise → accept/renoise result
     canvas[decode_slots] = torch.where(
@@ -554,7 +559,11 @@ def _compiled_sample_step(
                 soft_embeds, group_name=tp_group_name
             )
         soft_embeds = soft_embeds * normalizer
-        sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype)
+        # A pinned position holds its seed token as input. Zero its soft embed
+        # too, or the model's own prediction there reaches the next step
+        # through self-conditioning.
+        sc_pin = (~pin_mask[decode_slots]).unsqueeze(-1)
+        sc_embeds[decode_slots] = (soft_embeds * sc_keep * sc_pin).to(sc_embeds.dtype)
     else:
         # Every slot in this tile ends after this step, so the soft embed
         # would never be read. The matmul is a full pass over the vocabulary
@@ -650,6 +659,12 @@ def __init__(
         )
         self.has_seed = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
         self.seeded_slots: set[int] = set()
+        # Pinned positions keep their seed value on every denoise step, so a
+        # multi-step read denoises the canvas the request seeded.
+        self.pin_mask = torch.zeros(
+            max_num_reqs, canvas_length, dtype=torch.bool, device=device
+        )
+        self.pinned_slots: set[int] = set()
         # Read-only slots emit on their converging step and skip the commit
         # forward.
         self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
@@ -692,6 +707,8 @@ def add_request(self, slot_idx: int) -> None:
         self.max_steps[slot_idx].fill_(self.max_denoising_steps)
         self.has_seed[slot_idx].fill_(False)
         self.seeded_slots.discard(slot_idx)
+        self.pin_mask[slot_idx].fill_(False)
+        self.pinned_slots.discard(slot_idx)
         self.read_only[slot_idx].fill_(False)
         self.read_only_slots.discard(slot_idx)
         self.single_step_slots.discard(slot_idx)
@@ -704,6 +721,7 @@ def remove_request(self, slot_idx: int) -> None:
         self.accepted_canvas_history_len[slot_idx].fill_(0)
         self.self_conditioning_embeds[slot_idx] = 0
         self.seeded_slots.discard(slot_idx)
+        self.pinned_slots.discard(slot_idx)
         self.read_only_slots.discard(slot_idx)
         self.single_step_slots.discard(slot_idx)
 
@@ -716,6 +734,15 @@ def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None:
         self.has_seed[slot_idx].fill_(True)
         self.seeded_slots.add(slot_idx)
 
+    def set_pins(self, slot_idx: int, positions: list[int]) -> None:
+        """Hold ``positions`` of the slot's seed canvas through every denoise
+        step. Validated against the request's canvas width upstream."""
+        self.pin_mask[slot_idx].fill_(False)
+        self.pin_mask[
+            slot_idx, async_tensor_h2d(positions, dtype=torch.int64, device=self.device)
+        ] = True
+        self.pinned_slots.add(slot_idx)
+
     def set_read_only(self, slot_idx: int) -> None:
         self.read_only[slot_idx].fill_(True)
         self.read_only_slots.add(slot_idx)
@@ -1125,6 +1152,9 @@ def add_request(self, req_idx: int, sampling_params: Any) -> None:
                     f"got {len(seed)}"
                 )
             states.set_seed_canvas(req_idx, seed)
+        pins = extra.get("diffusion_pinned")
+        if pins and seed is not None:
+            states.set_pins(req_idx, [int(p) for p in pins])
         if extra.get("diffusion_read_only"):
             states.set_read_only(req_idx)
 
@@ -1386,6 +1416,8 @@ def __call__(
                     states.accepted_canvas_history[:, :, :W],
                     states.accepted_canvas_history_len,
                     states.max_steps,
+                    states.pin_mask[:, :W],
+                    states.seed_canvas[:, :W],
                     # Output
                     sampled[:, :W],
                     num_sampled,
diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py
index 1f7c86b0faf5..ceb35e8b1b8d 100644
--- a/vllm/sampling_params.py
+++ b/vllm/sampling_params.py
@@ -1079,9 +1079,11 @@ def _validate_diffusion_extra_args(
 
         On the GPU a seed id outside the vocabulary is a device-side assert.
         ``diffusion_seed_canvas``: token ids, one per canvas position, that
-        replace the random canvas after prefill. ``diffusion_max_steps``:
-        denoise steps per canvas. ``diffusion_read_only``: end the request
-        on the first canvas that converges.
+        replace the random canvas after prefill. ``diffusion_pinned``: canvas
+        positions held at their seed value on every denoise step.
+        ``diffusion_max_steps``: denoise steps per canvas.
+        ``diffusion_read_only``: end the request on the first canvas that
+        converges.
         """
         extra = self.extra_args
         if not extra:
@@ -1127,6 +1129,28 @@ def _validate_diffusion_extra_args(
                     parameter="extra_args",
                 )
 
+        pins = extra.get("diffusion_pinned")
+        if pins is not None:
+            if not isinstance(pins, (list, tuple)) or not all(
+                isinstance(p, int) and not isinstance(p, bool) for p in pins
+            ):
+                raise VLLMValidationError(
+                    "diffusion_pinned must be a list of canvas positions.",
+                    parameter="extra_args",
+                )
+            if seed is None:
+                raise VLLMValidationError(
+                    "diffusion_pinned needs a diffusion_seed_canvas to hold.",
+                    parameter="extra_args",
+                )
+            if any(
+                p < 0 or (expected_len is not None and p >= expected_len) for p in pins
+            ):
+                raise VLLMValidationError(
+                    "diffusion_pinned positions must lie inside the canvas.",
+                    parameter="extra_args",
+                )
+
         cap = extra.get("diffusion_max_steps")
         if cap is not None and (
             not isinstance(cap, int) or isinstance(cap, bool) or cap < 1

From 55f03bb5c862f8772d4bfe4838ec71e1045d8be4 Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson 
Date: Mon, 21 Sep 2026 04:08:16 +0000
Subject: [PATCH 16/25] Keep diffusion canvas widths in the scheduler

Trim synchronous worker drafts before grammar validation and select the
diffusion scheduler for either scheduling mode. Remove the model-state
count hook and the associated model-runner and draft-handler plumbing.

Co-authored-by: Codex 
Signed-off-by: Lucas Wilkinson 
---
 examples/features/diffusion_reads/README.md   |  4 ++
 tests/v1/core/test_scheduler.py               | 66 +++++++++++++++----
 vllm/config/vllm.py                           | 10 ++-
 vllm/model_executor/models/diffusion_gemma.py |  4 --
 vllm/v1/core/sched/diffusion_scheduler.py     | 22 +++++--
 vllm/v1/worker/gpu/model_runner.py            |  3 -
 vllm/v1/worker/gpu/model_states/interface.py  |  5 --
 vllm/v1/worker/gpu/spec_decode/utils.py       | 20 +-----
 8 files changed, 81 insertions(+), 53 deletions(-)

diff --git a/examples/features/diffusion_reads/README.md b/examples/features/diffusion_reads/README.md
index 45cdf70fd84c..f30c120527f6 100644
--- a/examples/features/diffusion_reads/README.md
+++ b/examples/features/diffusion_reads/README.md
@@ -29,6 +29,10 @@ curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -
   ]}'
 ```
 
+Diffusion models automatically select the diffusion scheduler for the active
+scheduling mode. Use `--no-async-scheduling` for synchronous execution;
+`--scheduler-cls` is only needed to override the default.
+
 Question types: `noul` (yes/no), `choice` with `options`, `score` with
 ordered `levels`. Each label must be a single token in the answer template,
 which the server checks with the tokenizer when a request uses the schema.
diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index 0a8de5cf1679..1c5ef1f61a65 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/tests/v1/core/test_scheduler.py
@@ -40,9 +40,9 @@
 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
 from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
-from vllm.v1.core.sched.async_scheduler import AsyncScheduler
 from vllm.v1.core.sched.diffusion_scheduler import (
     DiffusionAsyncScheduler,
+    DiffusionScheduler,
     diffusion_canvas_width,
 )
 from vllm.v1.core.sched.interface import PauseState
@@ -6878,18 +6878,14 @@ def _diffusion_scheduler(**kwargs) -> DiffusionAsyncScheduler:
     return scheduler
 
 
-def test_diffusion_scheduler_is_selected_with_async_scheduling():
-    def selected(**kwargs):
-        config = create_scheduler(**kwargs).vllm_config.scheduler_config
-        return config.get_scheduler_cls(), config.async_scheduling
-
-    # The platform has the last word on async scheduling (CPU turns it off).
-    cls, is_async = selected(async_scheduling=True, diffusion_canvas_length=8)
-    assert cls is (DiffusionAsyncScheduler if is_async else Scheduler)
-    cls, _ = selected(async_scheduling=False, diffusion_canvas_length=8)
-    assert cls is Scheduler
-    cls, is_async = selected(async_scheduling=True)
-    assert cls is (AsyncScheduler if is_async else Scheduler)
+@pytest.mark.parametrize("async_scheduling", [False, True])
+def test_diffusion_scheduler_is_selected_by_default(async_scheduling):
+    config = create_scheduler(
+        async_scheduling=async_scheduling, diffusion_canvas_length=8
+    ).vllm_config.scheduler_config
+    assert config.get_scheduler_cls() is (
+        DiffusionAsyncScheduler if config.async_scheduling else DiffusionScheduler
+    )
 
 
 def test_diffusion_scheduler_narrows_the_canvas_per_request():
@@ -6908,6 +6904,50 @@ def test_diffusion_scheduler_narrows_the_canvas_per_request():
     assert narrow.num_output_placeholders == 4
 
 
+@pytest.mark.parametrize("async_scheduling", [False, True])
+@pytest.mark.parametrize("structured", [False, True])
+def test_diffusion_scheduler_trims_full_width_worker_drafts(
+    async_scheduling, structured
+):
+    """Padded worker drafts must be narrowed before scheduling or grammar validation."""
+    scheduler = create_scheduler(
+        async_scheduling=async_scheduling,
+        diffusion_canvas_length=8,
+        scheduler_cls=(
+            DiffusionAsyncScheduler if async_scheduling else DiffusionScheduler
+        ),
+    )
+    wide = _diffusion_request("wide", {})
+    narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4})
+    for request in (wide, narrow):
+        scheduler.add_request(request)
+    prefill = scheduler.schedule()
+    _model_output(scheduler, prefill, [[], []])
+
+    if structured:
+        for request in (wide, narrow):
+            request.structured_output_request = SimpleNamespace(
+                grammar=_RecordingGrammar(), reasoning_ended=True
+            )
+    tokens = list(range(8)) if structured else [-1] * 8
+    drafts = DraftTokenIds(["wide", "narrow"], [tokens.copy(), tokens.copy()])
+    if async_scheduling:
+        output = scheduler.schedule()
+        scheduler.update_draft_token_ids_in_output(drafts, output)
+    else:
+        scheduler.update_draft_token_ids(drafts)
+        output = scheduler.schedule()
+
+    assert output.scheduled_spec_decode_tokens == {
+        "wide": tokens,
+        "narrow": tokens[:4],
+    }
+    assert output.num_scheduled_tokens == {"wide": 8, "narrow": 4}
+    if structured:
+        assert wide.structured_output_request.grammar.seen == [tokens]
+        assert narrow.structured_output_request.grammar.seen == [tokens[:4]]
+
+
 def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight():
     scheduler = _diffusion_scheduler()
     one = _diffusion_request(
diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py
index 65f2f8df523d..dead32af4984 100644
--- a/vllm/config/vllm.py
+++ b/vllm/config/vllm.py
@@ -1878,14 +1878,18 @@ def has_blocked_weights():
             )
         current_platform.check_and_update_config(self)
 
-        # After the platform hook, which has the last word on async scheduling.
+        # The platform has the last word on async scheduling.
         if (
             self.diffusion_config is not None
-            and self.scheduler_config.async_scheduling
             and self.scheduler_config.scheduler_cls is None
         ):
+            scheduler_cls = (
+                "DiffusionAsyncScheduler"
+                if self.scheduler_config.async_scheduling
+                else "DiffusionScheduler"
+            )
             self.scheduler_config.scheduler_cls = (
-                "vllm.v1.core.sched.diffusion_scheduler.DiffusionAsyncScheduler"
+                f"vllm.v1.core.sched.diffusion_scheduler.{scheduler_cls}"
             )
 
         self._normalize_piecewise_cudagraph_mode(
diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py
index 6d3e678e22bb..ce35049ceb3d 100644
--- a/vllm/model_executor/models/diffusion_gemma.py
+++ b/vllm/model_executor/models/diffusion_gemma.py
@@ -875,10 +875,6 @@ def remove_request(self, req_id: str) -> None:
         if idx is not None:
             self.diffusion_states.remove_request(idx)
 
-    def num_draft_tokens_per_req(self, input_batch: Any) -> np.ndarray:
-        slots = input_batch.idx_mapping_np[: input_batch.num_reqs]
-        return self.diffusion_states.canvas_width_np[slots]
-
     def prepare_inputs_embeds(
         self,
         scheduled_encoder_inputs: dict[str, list[int]],
diff --git a/vllm/v1/core/sched/diffusion_scheduler.py b/vllm/v1/core/sched/diffusion_scheduler.py
index 96462aa6abe1..6b3fb0b9675e 100644
--- a/vllm/v1/core/sched/diffusion_scheduler.py
+++ b/vllm/v1/core/sched/diffusion_scheduler.py
@@ -1,15 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Async scheduling for diffusion requests.
-
-These rules live here so Scheduler.schedule() and AsyncScheduler stay
-unchanged. VllmConfig selects this class for a diffusion model under async
-scheduling. A sync scheduler creates no output placeholders, which both rules
-read.
-"""
+"""Scheduling policy for per-request diffusion canvases."""
 
 from vllm.v1.core.sched.async_scheduler import AsyncScheduler
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.outputs import DraftTokenIds
 from vllm.v1.request import Request
 
 
@@ -38,6 +34,18 @@ def _read_in_flight(request: Request, width: int) -> bool:
     return request.num_output_placeholders >= steps * width
 
 
+class DiffusionScheduler(Scheduler):
+    def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
+        for req_id, tokens in zip(
+            draft_token_ids.req_ids, draft_token_ids.draft_token_ids
+        ):
+            request = self.requests.get(req_id)
+            if request is not None:
+                width = diffusion_canvas_width(request, self.num_spec_tokens)
+                del tokens[width:]
+        super().update_draft_token_ids(draft_token_ids)
+
+
 class DiffusionAsyncScheduler(AsyncScheduler):
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         for request in self.running:
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index f058eac2db19..92a37008cebe 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -2169,9 +2169,6 @@ def sample_tokens(
             self.draft_tokens_handler.set_draft_tokens(
                 input_batch,
                 self.req_states.draft_tokens[input_batch.idx_mapping],
-                num_draft_tokens_per_req=self.model_state.num_draft_tokens_per_req(
-                    input_batch
-                ),
             )
             if self.pp_handler is not None:
                 self.pp_handler.broadcast_drafts(
diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py
index 70bb19c8d22f..a5c8d6f4cfd7 100644
--- a/vllm/v1/worker/gpu/model_states/interface.py
+++ b/vllm/v1/worker/gpu/model_states/interface.py
@@ -237,8 +237,3 @@ def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
     num_new_sampled_tokens_per_step: int = 1
     """New tokens sampled on each decode step 
     (excluding accepted draft tokens, a.k.a num bonus tokens)."""
-
-    def num_draft_tokens_per_req(self, input_batch: Any) -> Any:
-        """Draft tokens each request in the batch carries into its next step,
-        as an int array, or None when every request carries the buffer width."""
-        return None
diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py
index 0e673670d611..f73acb2bdaed 100644
--- a/vllm/v1/worker/gpu/spec_decode/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/utils.py
@@ -36,22 +36,12 @@ def __init__(self, device: torch.device | None = None):
         self.req_ids: list[str] = []
         self.draft_tokens_np: np.ndarray | None = None
         self.num_draft_tokens: int = 0
-        # Per-request draft counts when a model state narrows some requests.
-        self.num_draft_tokens_per_req: list[int] | None = None
 
     def set_draft_tokens(
-        self,
-        input_batch: InputBatch,
-        draft_tokens: torch.Tensor,
-        num_draft_tokens_per_req: np.ndarray | None = None,
+        self, input_batch: InputBatch, draft_tokens: torch.Tensor
     ) -> None:
         self.req_ids = input_batch.req_ids
         self.num_draft_tokens = draft_tokens.shape[1]
-        self.num_draft_tokens_per_req = (
-            None
-            if num_draft_tokens_per_req is None
-            else num_draft_tokens_per_req.tolist()
-        )
         if not input_batch.has_structured_output_reqs:
             # No draft token validation needs to be performed by
             # the scheduler for this batch.
@@ -71,18 +61,12 @@ def set_draft_tokens(
             self.copy_event.record()
 
     def get_draft_tokens(self) -> DraftTokenIds | None:
-        per_req = self.num_draft_tokens_per_req
         if self.draft_tokens_np is not None:
             self.copy_event.synchronize()
             draft_token_ids = self.draft_tokens_np.tolist()
-            if per_req is not None:
-                draft_token_ids = [row[:n] for row, n in zip(draft_token_ids, per_req)]
         else:
             # This case only happens when async scheduling is disabled.
-            draft_token_ids = [
-                [-1] * (self.num_draft_tokens if per_req is None else per_req[i])
-                for i in range(len(self.req_ids))
-            ]
+            draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids]
         return DraftTokenIds(self.req_ids, draft_token_ids)
 
 

From 2449995bf8d73d255940e74d628f2f92446ca6e7 Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson 
Date: Mon, 21 Sep 2026 04:39:55 +0000
Subject: [PATCH 17/25] Require async scheduling for narrower diffusion
 canvases

Remove the synchronous diffusion scheduler and keep full-width synchronous
generation on the standard Scheduler. Reject narrower synchronous requests
at admission so their canvas width cannot be silently ignored.

Co-authored-by: Codex 
Signed-off-by: Lucas Wilkinson 
---
 examples/features/diffusion_reads/README.md |  7 +++--
 tests/test_sampling_params.py               | 35 +++++++++++++++++++++
 tests/v1/core/test_scheduler.py             | 24 +++-----------
 vllm/config/vllm.py                         | 10 ++----
 vllm/v1/core/sched/diffusion_scheduler.py   | 22 +++++--------
 vllm/v1/engine/input_processor.py           | 14 +++++++++
 6 files changed, 68 insertions(+), 44 deletions(-)

diff --git a/examples/features/diffusion_reads/README.md b/examples/features/diffusion_reads/README.md
index f30c120527f6..56f0ada2546d 100644
--- a/examples/features/diffusion_reads/README.md
+++ b/examples/features/diffusion_reads/README.md
@@ -29,9 +29,10 @@ curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -
   ]}'
 ```
 
-Diffusion models automatically select the diffusion scheduler for the active
-scheduling mode. Use `--no-async-scheduling` for synchronous execution;
-`--scheduler-cls` is only needed to override the default.
+Per-request canvas widths smaller than the served canvas require async
+scheduling. The diffusion async scheduler is selected automatically; no
+`--scheduler-cls` argument is needed. Synchronous execution supports full-width
+canvases only.
 
 Question types: `noul` (yes/no), `choice` with `options`, `score` with
 ordered `levels`. Each label must be a single token in the answer template,
diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py
index 472112d5cfcd..874237694ba0 100644
--- a/tests/test_sampling_params.py
+++ b/tests/test_sampling_params.py
@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
+from types import SimpleNamespace
 
 import pytest
 
 from vllm import SamplingParams
 from vllm.exceptions import VLLMValidationError
+from vllm.v1.engine.input_processor import InputProcessor
 
 
 @dataclass
@@ -13,6 +15,7 @@ class MockModelConfig:
     is_diffusion: bool = False
     max_logprobs: int = 20
     logits_processors: list | None = None
+    return_sampling_mask: bool = False
 
     def get_vocab_size(self) -> int:
         return 1024
@@ -77,6 +80,38 @@ def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None):
     )
 
 
+@pytest.mark.parametrize("async_scheduling", [False, True])
+@pytest.mark.parametrize(
+    "extra_args",
+    [
+        {},
+        {"diffusion_canvas_length": None},
+        {"diffusion_canvas_length": 4},
+        {"diffusion_canvas_length": 8},
+    ],
+)
+def test_narrow_diffusion_canvas_requires_async_scheduling(
+    async_scheduling, extra_args
+):
+    processor = SimpleNamespace(
+        model_config=MockModelConfig(is_diffusion=True),
+        vllm_config=SimpleNamespace(
+            scheduler_config=SimpleNamespace(async_scheduling=async_scheduling)
+        ),
+        speculative_config=None,
+        structured_outputs_config=None,
+        diffusion_config=MockDiffusionConfig(canvas_length=8),
+        tokenizer=None,
+        validate_logits_processors_params=lambda params: None,
+    )
+    params = SamplingParams(extra_args=extra_args)
+    if not async_scheduling and extra_args.get("diffusion_canvas_length") == 4:
+        with pytest.raises(VLLMValidationError, match="requires --async-scheduling"):
+            InputProcessor._validate_params(processor, params, ("generate",))
+    else:
+        InputProcessor._validate_params(processor, params, ("generate",))
+
+
 @pytest.mark.parametrize(
     "extra_args, match",
     [
diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index 1c5ef1f61a65..30d242209e01 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/tests/v1/core/test_scheduler.py
@@ -42,7 +42,6 @@
 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
 from vllm.v1.core.sched.diffusion_scheduler import (
     DiffusionAsyncScheduler,
-    DiffusionScheduler,
     diffusion_canvas_width,
 )
 from vllm.v1.core.sched.interface import PauseState
@@ -6884,7 +6883,7 @@ def test_diffusion_scheduler_is_selected_by_default(async_scheduling):
         async_scheduling=async_scheduling, diffusion_canvas_length=8
     ).vllm_config.scheduler_config
     assert config.get_scheduler_cls() is (
-        DiffusionAsyncScheduler if config.async_scheduling else DiffusionScheduler
+        DiffusionAsyncScheduler if config.async_scheduling else Scheduler
     )
 
 
@@ -6904,19 +6903,10 @@ def test_diffusion_scheduler_narrows_the_canvas_per_request():
     assert narrow.num_output_placeholders == 4
 
 
-@pytest.mark.parametrize("async_scheduling", [False, True])
 @pytest.mark.parametrize("structured", [False, True])
-def test_diffusion_scheduler_trims_full_width_worker_drafts(
-    async_scheduling, structured
-):
+def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
     """Padded worker drafts must be narrowed before scheduling or grammar validation."""
-    scheduler = create_scheduler(
-        async_scheduling=async_scheduling,
-        diffusion_canvas_length=8,
-        scheduler_cls=(
-            DiffusionAsyncScheduler if async_scheduling else DiffusionScheduler
-        ),
-    )
+    scheduler = _diffusion_scheduler()
     wide = _diffusion_request("wide", {})
     narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4})
     for request in (wide, narrow):
@@ -6931,12 +6921,8 @@ def test_diffusion_scheduler_trims_full_width_worker_drafts(
             )
     tokens = list(range(8)) if structured else [-1] * 8
     drafts = DraftTokenIds(["wide", "narrow"], [tokens.copy(), tokens.copy()])
-    if async_scheduling:
-        output = scheduler.schedule()
-        scheduler.update_draft_token_ids_in_output(drafts, output)
-    else:
-        scheduler.update_draft_token_ids(drafts)
-        output = scheduler.schedule()
+    output = scheduler.schedule()
+    scheduler.update_draft_token_ids_in_output(drafts, output)
 
     assert output.scheduled_spec_decode_tokens == {
         "wide": tokens,
diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py
index dead32af4984..65f2f8df523d 100644
--- a/vllm/config/vllm.py
+++ b/vllm/config/vllm.py
@@ -1878,18 +1878,14 @@ def has_blocked_weights():
             )
         current_platform.check_and_update_config(self)
 
-        # The platform has the last word on async scheduling.
+        # After the platform hook, which has the last word on async scheduling.
         if (
             self.diffusion_config is not None
+            and self.scheduler_config.async_scheduling
             and self.scheduler_config.scheduler_cls is None
         ):
-            scheduler_cls = (
-                "DiffusionAsyncScheduler"
-                if self.scheduler_config.async_scheduling
-                else "DiffusionScheduler"
-            )
             self.scheduler_config.scheduler_cls = (
-                f"vllm.v1.core.sched.diffusion_scheduler.{scheduler_cls}"
+                "vllm.v1.core.sched.diffusion_scheduler.DiffusionAsyncScheduler"
             )
 
         self._normalize_piecewise_cudagraph_mode(
diff --git a/vllm/v1/core/sched/diffusion_scheduler.py b/vllm/v1/core/sched/diffusion_scheduler.py
index 6b3fb0b9675e..96462aa6abe1 100644
--- a/vllm/v1/core/sched/diffusion_scheduler.py
+++ b/vllm/v1/core/sched/diffusion_scheduler.py
@@ -1,11 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Scheduling policy for per-request diffusion canvases."""
+"""Async scheduling for diffusion requests.
+
+These rules live here so Scheduler.schedule() and AsyncScheduler stay
+unchanged. VllmConfig selects this class for a diffusion model under async
+scheduling. A sync scheduler creates no output placeholders, which both rules
+read.
+"""
 
 from vllm.v1.core.sched.async_scheduler import AsyncScheduler
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.core.sched.scheduler import Scheduler
-from vllm.v1.outputs import DraftTokenIds
 from vllm.v1.request import Request
 
 
@@ -34,18 +38,6 @@ def _read_in_flight(request: Request, width: int) -> bool:
     return request.num_output_placeholders >= steps * width
 
 
-class DiffusionScheduler(Scheduler):
-    def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
-        for req_id, tokens in zip(
-            draft_token_ids.req_ids, draft_token_ids.draft_token_ids
-        ):
-            request = self.requests.get(req_id)
-            if request is not None:
-                width = diffusion_canvas_width(request, self.num_spec_tokens)
-                del tokens[width:]
-        super().update_draft_token_ids(draft_token_ids)
-
-
 class DiffusionAsyncScheduler(AsyncScheduler):
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         for request in self.running:
diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py
index c3abec1b41be..db71de57afa6 100644
--- a/vllm/v1/engine/input_processor.py
+++ b/vllm/v1/engine/input_processor.py
@@ -131,6 +131,20 @@ def _validate_params(
 
             self.validate_logits_processors_params(params)
 
+            if (
+                self.diffusion_config is not None
+                and self.model_config.is_diffusion
+                and not self.vllm_config.scheduler_config.async_scheduling
+                and params.extra_args
+            ):
+                width = params.extra_args.get("diffusion_canvas_length")
+                if width is not None and width != self.diffusion_config.canvas_length:
+                    raise VLLMValidationError(
+                        "A diffusion_canvas_length smaller than the served canvas "
+                        "requires --async-scheduling.",
+                        parameter="extra_args",
+                    )
+
             if self.model_config.return_sampling_mask:
                 if params.temperature <= 0:
                     raise ValueError(

From 084444aa853f1b27b00c8990f8ea9c098693c406 Mon Sep 17 00:00:00 2001
From: Matt Mastracci 
Date: Mon, 21 Sep 2026 09:43:23 -0600
Subject: [PATCH 18/25] [Docs] Drop the HTML pages from the diffusion reads
 example

The playground and walk pages were demo front ends for the example
server. Examples do not ship HTML, so the pages go, along with the
TEST_PAGE switch and the GET routes that served them. The server keeps
its JSON routes and the optional TLS listener.

Co-Authored-By: Claude Fable 5.1 
Signed-off-by: Matt Mastracci 
---
 examples/features/diffusion_reads/README.md   |   4 +-
 .../features/diffusion_reads/playground.html  | 223 ------------------
 .../diffusion_reads/structured_server.py      |  26 +-
 examples/features/diffusion_reads/walk.html   | 193 ---------------
 4 files changed, 3 insertions(+), 443 deletions(-)
 delete mode 100644 examples/features/diffusion_reads/playground.html
 delete mode 100644 examples/features/diffusion_reads/walk.html

diff --git a/examples/features/diffusion_reads/README.md b/examples/features/diffusion_reads/README.md
index 56f0ada2546d..2743b45c4840 100644
--- a/examples/features/diffusion_reads/README.md
+++ b/examples/features/diffusion_reads/README.md
@@ -51,9 +51,7 @@ is among the listed ones, otherwise null) and `alone` (a read of its own).
 
 Images attach as `multipart/form-data`, with the JSON in a part named
 `request` and each image as a file part, or as an `images` array of data
-URLs. With `TEST_PAGE=1` in the environment, `GET /` serves
-`playground.html`, a page for sending requests with an image file or webcam
-frames.
+URLs.
 
 ```bash
 curl -s localhost:8011/v1/systemone -H 'content-type: application/json' -d '{
diff --git a/examples/features/diffusion_reads/playground.html b/examples/features/diffusion_reads/playground.html
deleted file mode 100644
index a29b0929c708..000000000000
--- a/examples/features/diffusion_reads/playground.html
+++ /dev/null
@@ -1,223 +0,0 @@
-
-
-
-
-
-djev playground
-
-
-
-
djev playground

-

-  

-    
Request body for POST /v1/systemone

-    
-    

-      
-      
-      
-      
-    

-    

-      
-      
-      
-      
-      
-    

-    
-    
-    
-  

-  

-    
Answers

-    

-    

-    
raw response

-  

-

-
-
-
diff --git a/examples/features/diffusion_reads/structured_server.py b/examples/features/diffusion_reads/structured_server.py
index 439e3e8fa889..5e43d2b4c5ec 100644
--- a/examples/features/diffusion_reads/structured_server.py
+++ b/examples/features/diffusion_reads/structured_server.py
@@ -27,12 +27,8 @@
 unchanged, for plain generation through this port.
 
 With API_KEY set in the environment, every POST needs "Authorization:
-Bearer ". With TEST_PAGE=1, GET / serves playground.html, a form for
-the request JSON with an image file or webcam frames attached, and GET
-/walk serves walk.html, a phone page that streams the back camera and
-reads one hazard label per frame. Browsers open the webcam only on
-a secure origin, so --tls-port adds an HTTPS listener with a self-signed
-certificate kept in --cert-dir.
+Bearer ". --tls-port adds an HTTPS listener with a self-signed
+certificate kept in --cert-dir, for clients that need a secure origin.
 
 Each answer is one distribution per question, from one denoise step over a
 seeded canvas, averaged over a few noise draws. This server handles the
@@ -101,15 +97,9 @@
 
 ARGS = None
 TOK = None
-TEST_PAGE = os.environ.get("TEST_PAGE", "") == "1"  # serve the playground at /
 API_KEY = os.environ.get(
     "API_KEY", ""
 )  # when set, POST routes need "Authorization: Bearer "
-PAGES = {  # served with TEST_PAGE=1
-    "/": "playground.html",
-    "/playground": "playground.html",
-    "/walk": "walk.html",
-}
 CANVAS_LEN = 64  # the served canvas length. A request may be narrower.
 CANVAS_STEP = 16  # request widths are multiples of this
 VOCAB = 262144
@@ -1211,18 +1201,6 @@ def _json(self, code, obj):
     def do_GET(self):
         if self.path == "/health":
             return self._json(200, {"status": "ok"})
-        if self.path in PAGES and TEST_PAGE:
-            page_path = os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), PAGES[self.path]
-            )
-            with open(page_path, "rb") as fh:
-                body = fh.read()
-            self.send_response(200)
-            self.send_header("content-type", "text/html; charset=utf-8")
-            self.send_header("content-length", str(len(body)))
-            self.end_headers()
-            self.wfile.write(body)
-            return
         return self._json(404, {"error": {"message": "unknown route"}})
 
     def _read_request(self):
diff --git a/examples/features/diffusion_reads/walk.html b/examples/features/diffusion_reads/walk.html
deleted file mode 100644
index d1883250be41..000000000000
--- a/examples/features/diffusion_reads/walk.html
+++ /dev/null
@@ -1,193 +0,0 @@
-
-
-
-
-
-
-djev walk
-
-
-
-

-  
-  
↑

-  
-

-

-  

-    
-    
-  

-  
 

-  

-  

-    
-  

-  
 

-

-
-
-
-

From 340eca48e49b74000914ca70d3a479f430e042f6 Mon Sep 17 00:00:00 2001
From: Matt Mastracci 
Date: Mon, 21 Sep 2026 09:43:24 -0600
Subject: [PATCH 19/25] [Docs] Rename the example directory to
 structured_diffusion

Review suggestion: the name says what the example is for.

Co-Authored-By: Claude Fable 5.1 
Signed-off-by: Matt Mastracci 
---
 .../{diffusion_reads => structured_diffusion}/README.md         | 2 +-
 .../structured_server.py                                        | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename examples/features/{diffusion_reads => structured_diffusion}/README.md (98%)
 rename examples/features/{diffusion_reads => structured_diffusion}/structured_server.py (100%)

diff --git a/examples/features/diffusion_reads/README.md b/examples/features/structured_diffusion/README.md
similarity index 98%
rename from examples/features/diffusion_reads/README.md
rename to examples/features/structured_diffusion/README.md
index 2743b45c4840..317fa4d4bc04 100644
--- a/examples/features/diffusion_reads/README.md
+++ b/examples/features/structured_diffusion/README.md
@@ -20,7 +20,7 @@ with a standard error over a few noise draws.
 ```bash
 vllm serve google/diffusiongemma-26B-A4B-it \
     --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 --enable-prefix-caching
-python examples/features/diffusion_reads/structured_server.py \
+python examples/features/structured_diffusion/structured_server.py \
     --upstream http://127.0.0.1:8000 --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64
 curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -d '{
   "messages": [
diff --git a/examples/features/diffusion_reads/structured_server.py b/examples/features/structured_diffusion/structured_server.py
similarity index 100%
rename from examples/features/diffusion_reads/structured_server.py
rename to examples/features/structured_diffusion/structured_server.py

From a68dab46cad38871376835c6c59046fe807c42ec Mon Sep 17 00:00:00 2001
From: Razorback16 
Date: Fri, 18 Sep 2026 12:04:57 -0700
Subject: [PATCH 20/25] [Bugfix] DiffusionGemma: keep the sampler step
 compiled, and correct in eager

_compiled_sample_step stores the self-conditioning soft embed, computed in
the model dtype (bf16), into an fp32 buffer. Compiled code casts on the
store; eager raises "Index put requires the source and destination dtypes
match" and kills the EngineCore. The step runs eager once Dynamo hits its
recompile limit (8), which a server reaches after a few minutes of mixed
traffic: tiles specialize the graph on batch size 1, full-width tiles,
compute_sc and buffer-size coincidences. From then on every multi-step
read and every generation crashed the engine.

Cast the store to the buffer's dtype, and raise the recompile limit for
this function to 64 so it stays compiled (eager is about 2x slower per
self-conditioned step: 4.0 vs 1.9 ms at batch 1, 14 vs 6.9 ms at batch 8).

Test: the step with a bf16 embedding, compiled and forced eager.
Signed-off-by: Razorback16 
(cherry picked from commit b07a9bedb30c23710d89a14f30353d0352428b15)
---
 tests/v1/sample/test_diffusion_gemma_reads.py | 29 +++++++++++++++++--
 vllm/model_executor/models/diffusion_gemma.py | 11 ++++++-
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py
index a705f4a30200..36c77177f742 100644
--- a/tests/v1/sample/test_diffusion_gemma_reads.py
+++ b/tests/v1/sample/test_diffusion_gemma_reads.py
@@ -116,6 +116,7 @@ def _denoise_once(
     compute_sc: bool = True,
     width: int = CL,
     embed_weight: torch.Tensor | None = None,
+    embed_dtype: torch.dtype = torch.float32,
 ) -> None:
     """One compiled denoise step over ``slots`` with flat logits, so nothing
     converges by stability or confidence and only the step cap can end it.
@@ -138,11 +139,11 @@ def _denoise_once(
         states.confident,
         states.self_conditioning_embeds[:, :width],
         (
-            torch.zeros(VOCAB, 4, device=device)
+            torch.zeros(VOCAB, 4, dtype=embed_dtype, device=device)
             if embed_weight is None
             else embed_weight
         ),
-        torch.tensor(1.0, device=device),
+        torch.tensor(1.0, dtype=embed_dtype, device=device),
         states.accepted_canvas_history[:, :, :width],
         states.accepted_canvas_history_len,
         states.max_steps,
@@ -211,6 +212,30 @@ def test_single_step_tile_skips_self_conditioning():
     assert not states.self_conditioning_embeds[0].any()
 
 
+@pytest.mark.parametrize("stance", ["default", "force_eager"])
+def test_self_conditioning_stores_a_bf16_model_in_the_fp32_buffer(stance):
+    # The model's embeddings are bf16 while the buffer is fp32. Compiled code
+    # casts on the store, but eager does not, and the step runs eager once
+    # torch.compile hits its recompile limit.
+    states = _states()
+    states.add_request(0)
+    states.is_encoder_phase[0] = False
+
+    with torch.compiler.set_stance(stance):
+        _denoise_once(
+            states,
+            [0],
+            embed_weight=torch.ones(VOCAB, 4, dtype=torch.bfloat16, device="cuda"),
+            embed_dtype=torch.bfloat16,
+        )
+
+    assert states.self_conditioning_embeds.dtype == torch.float32
+    # uniform probs @ an all-ones embedding
+    assert torch.allclose(
+        states.self_conditioning_embeds[0], torch.ones(CL, 4, device="cuda")
+    )
+
+
 def test_narrow_tile_leaves_the_rest_of_the_canvas_alone():
     states = _states()
     states.add_request(0)
diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py
index ce35049ceb3d..5b6fc10c8e87 100644
--- a/vllm/model_executor/models/diffusion_gemma.py
+++ b/vllm/model_executor/models/diffusion_gemma.py
@@ -22,6 +22,7 @@
 
 import numpy as np
 import torch
+import torch._dynamo
 from torch import nn
 from torch.nn import functional as F
 from transformers import AutoModel
@@ -371,6 +372,11 @@ def _compute_num_rejected(
     return torch.where(is_denoise, query_lens, num_rejected)
 
 
+# Tiles specialize the graph on batch size 1, on a tile as wide as the canvas,
+# on compute_sc and on sizes that coincide with the state buffers. That set
+# is small but passes Dynamo's default of 8, after which every step would run
+# eager: about twice as slow for a self-conditioned step.
+@torch._dynamo.config.patch(recompile_limit=64)
 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
 def _compiled_sample_step(
     # Logits from the model [num_decode * CL, vocab]
@@ -561,7 +567,10 @@ def _compiled_sample_step(
         soft_embeds = soft_embeds * normalizer
         # A pinned position holds its seed token as input. Zero its soft embed
         # too, or the model's own prediction there reaches the next step
-        # through self-conditioning.
+        # through self-conditioning. The buffer is fp32 while the embedding is
+        # in the model dtype: compiled code casts on the store but eager does
+        # not, and the step runs eager once torch.compile hits its recompile
+        # limit.
         sc_pin = (~pin_mask[decode_slots]).unsqueeze(-1)
         sc_embeds[decode_slots] = (soft_embeds * sc_keep * sc_pin).to(sc_embeds.dtype)
     else:

From d9af863210e6f8112793d004ab0af56a97083f40 Mon Sep 17 00:00:00 2001
From: Razorback16 
Date: Fri, 18 Sep 2026 12:20:09 -0700
Subject: [PATCH 21/25] [Bugfix] DiffusionGemma: join logprob stashes of
 different widths

A request's logprobs are stashed at the step it converges, as wide as the
widest logprobs request in that step's batch, and joined at commit. Two
requests that converged in different steps (a read asking for 10 label
ids, a generation asking for 2 logprobs) left stashes of different widths,
and torch.cat raised "Sizes of tensors must match", killing the
EngineCore. Pad the narrow stashes the way compute_topk_scores pads a
mixed batch: token id 0 at -inf, which the output processor never reports.

Signed-off-by: Razorback16 
(cherry picked from commit 9bbf7418e85020dc76da9f60cdfe6c4e912ec048)
---
 tests/v1/sample/test_diffusion_gemma_reads.py | 28 ++++++++++++++
 vllm/model_executor/models/diffusion_gemma.py | 38 +++++++++++++------
 2 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py
index 36c77177f742..50fb75ed149b 100644
--- a/tests/v1/sample/test_diffusion_gemma_reads.py
+++ b/tests/v1/sample/test_diffusion_gemma_reads.py
@@ -10,8 +10,10 @@
 from vllm.model_executor.models.diffusion_gemma import (
     DiffusionGemmaRequestStates,
     _compiled_sample_step,
+    _concat_logprob_stashes,
 )
 from vllm.platforms import current_platform
+from vllm.v1.outputs import LogprobsTensors
 
 pytestmark = pytest.mark.skipif(
     not current_platform.is_cuda(), reason="the sampler state lives on the GPU"
@@ -263,3 +265,29 @@ def test_step_cap_is_per_slot():
     # Slot 0 hit its cap and moves to commit. Slot 1 keeps denoising.
     assert states.is_encoder_phase[:2].tolist() == [True, False]
     assert states.step[:2].tolist() == [1, 1]
+
+
+def _stash(rows: int, width: int, first_id: int) -> LogprobsTensors:
+    ids = torch.arange(first_id, first_id + rows * width).reshape(rows, width)
+    return LogprobsTensors(
+        logprob_token_ids=ids,
+        logprobs=-ids.float(),
+        selected_token_ranks=torch.zeros(rows, dtype=torch.int64),
+    )
+
+
+def test_stashes_of_different_widths_join():
+    # A read that asked for 10 label ids and a generation that asked for 2
+    # logprobs converged in different steps, so their stashes differ in width.
+    wide, narrow = _stash(2, 11, 100), _stash(3, 3, 0)
+
+    out = _concat_logprob_stashes([narrow, wide], [0, 3])
+
+    assert out.logprob_token_ids.shape == (5, 11)
+    assert out.logprobs.shape == (5, 11)
+    assert out.cu_num_generated_tokens == [0, 3]
+    # The narrow rows keep their values and pad with id 0 at -inf.
+    assert torch.equal(out.logprob_token_ids[:3, :3], narrow.logprob_token_ids)
+    assert not out.logprob_token_ids[:3, 3:].any()
+    assert torch.isneginf(out.logprobs[:3, 3:]).all()
+    assert torch.equal(out.logprobs[3:], wide.logprobs)
diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py
index 5b6fc10c8e87..daf4f86737d8 100644
--- a/vllm/model_executor/models/diffusion_gemma.py
+++ b/vllm/model_executor/models/diffusion_gemma.py
@@ -372,6 +372,29 @@ def _compute_num_rejected(
     return torch.where(is_denoise, query_lens, num_rejected)
 
 
+def _concat_logprob_stashes(
+    parts: list[LogprobsTensors], cu_num_generated_tokens: list[int]
+) -> LogprobsTensors:
+    """Join the logprobs stashed for the requests committing this step.
+
+    Each stash is as wide as the widest logprobs request in the batch at the
+    step that request converged, so stashes from different steps can differ
+    in width. Pad the narrow ones the way compute_topk_scores pads a mixed
+    batch: token id 0 at -inf, which the output processor never reports.
+    """
+    width = max(p.logprob_token_ids.shape[1] for p in parts)
+
+    def pad(t: torch.Tensor, value: float) -> torch.Tensor:
+        return F.pad(t, (0, width - t.shape[1]), value=value)
+
+    return LogprobsTensors(
+        logprob_token_ids=torch.cat([pad(p.logprob_token_ids, 0) for p in parts]),
+        logprobs=torch.cat([pad(p.logprobs, float("-inf")) for p in parts]),
+        selected_token_ranks=torch.cat([p.selected_token_ranks for p in parts]),
+        cu_num_generated_tokens=cu_num_generated_tokens,
+    )
+
+
 # Tiles specialize the graph on batch size 1, on a tile as wide as the canvas,
 # on compute_sc and on sizes that coincide with the state buffers. That set
 # is small but passes Dynamo's default of 8, after which every step would run
@@ -1519,7 +1542,7 @@ def __call__(
             committing_slots = (
                 set(decode_slots_np[is_committing.cpu().numpy()].tolist()) | emit_now
             )
-            parts_ids, parts_lp, parts_ranks = [], [], []
+            parts: list[LogprobsTensors] = []
             cu_gen: list[int] = []
             flat_offset = 0
             for i in range(num_reqs):
@@ -1527,17 +1550,10 @@ def __call__(
                 slot = int(slots_np[i])
                 if slot in committing_slots and slot in self._pending_logprobs:
                     lp = self._pending_logprobs.pop(slot)
-                    parts_ids.append(lp.logprob_token_ids)
-                    parts_lp.append(lp.logprobs)
-                    parts_ranks.append(lp.selected_token_ranks)
+                    parts.append(lp)
                     flat_offset += lp.logprobs.shape[0]
-            if parts_ids:
-                logprobs_tensors = LogprobsTensors(
-                    logprob_token_ids=torch.cat(parts_ids),
-                    logprobs=torch.cat(parts_lp),
-                    selected_token_ranks=torch.cat(parts_ranks),
-                    cu_num_generated_tokens=cu_gen,
-                )
+            if parts:
+                logprobs_tensors = _concat_logprob_stashes(parts, cu_gen)
 
         return self._build_output(
             input_batch,

From 58a1993127b1e8d73ec47674840c2430a3034978 Mon Sep 17 00:00:00 2001
From: Matt Mastracci 
Date: Mon, 21 Sep 2026 13:12:07 -0600
Subject: [PATCH 22/25] [Docs] Bound the example server's fan-out

The fixed sample count was capped at 32 but "auto_max" was not, and the
question groups of one request each got a thread. A request could ask for
any number of questions and reads and open that many threads. Cap the
questions per request, the reads per question under both policies, and the
groups read at once.

Co-Authored-By: Claude Fable 5.1 
Signed-off-by: Matt Mastracci 
---
 .../structured_diffusion/structured_server.py         | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/examples/features/structured_diffusion/structured_server.py b/examples/features/structured_diffusion/structured_server.py
index 5e43d2b4c5ec..a8f61da94e41 100644
--- a/examples/features/structured_diffusion/structured_server.py
+++ b/examples/features/structured_diffusion/structured_server.py
@@ -106,6 +106,9 @@
 TURN_CLOSE = 106
 PAD = 0
 TOPK = 20
+MAX_QUESTIONS = 64  # per request
+MAX_SAMPLES = 32  # reads per question, fixed or auto
+MAX_PARALLEL = 16  # question groups read at once
 # the empty thought block the chat template leaves to the model
 SCAFFOLD_TEXT = "<|channel>thought\n"
 SCAFFOLD = None
@@ -129,6 +132,8 @@ def parse_schema(value):
         or not value["questions"]
     ):
         raise SchemaError("schema: needs a non-empty questions array")
+    if len(value["questions"]) > MAX_QUESTIONS:
+        raise SchemaError(f"schema: at most {MAX_QUESTIONS} questions")
     qs = []
     seen = set()
     for q in value["questions"]:
@@ -212,11 +217,11 @@ def parse_schema(value):
     if samples == "auto":
         policy = {
             "mode": "auto",
-            "max": int(value.get("auto_max", 4)),
+            "max": max(1, min(int(value.get("auto_max", 4)), MAX_SAMPLES)),
             "threshold": float(value.get("auto_threshold", 0.1)),
         }
     elif isinstance(samples, int) and samples >= 1:
-        policy = {"mode": "fixed", "n": min(samples, 32)}
+        policy = {"mode": "fixed", "n": min(samples, MAX_SAMPLES)}
     else:
         raise SchemaError('schema: samples must be a positive count or "auto"')
     ask = value.get("ask")
@@ -828,7 +833,7 @@ def absorb(group, body, rows):
                 absorb(group, body, rows)
                 k += 1
         else:
-            with ThreadPoolExecutor(max_workers=len(groups)) as ex:
+            with ThreadPoolExecutor(max_workers=min(len(groups), MAX_PARALLEL)) as ex:
                 results = list(
                     ex.map(
                         lambda gk, conditioned=conditioned: run(

From dd8639536d745de03d526d6a2e7339a0c1f8b7cc Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson 
Date: Mon, 21 Sep 2026 21:05:03 +0000
Subject: [PATCH 23/25] Isolate diffusion request validation and simplify read
 emission

Co-authored-by: OpenAI Codex 

Signed-off-by: Lucas Wilkinson 
---
 tests/test_sampling_params.py                 |  25 ++--
 tests/v1/sample/test_diffusion_gemma_reads.py |  36 ++++-
 vllm/model_executor/models/diffusion_gemma.py |  62 +++------
 vllm/sampling_params.py                       | 124 +-----------------
 vllm/utils/diffusion.py                       | 101 ++++++++++++++
 vllm/v1/engine/input_processor.py             |  22 ++--
 6 files changed, 168 insertions(+), 202 deletions(-)
 create mode 100644 vllm/utils/diffusion.py

diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py
index 874237694ba0..635842e3a121 100644
--- a/tests/test_sampling_params.py
+++ b/tests/test_sampling_params.py
@@ -6,7 +6,9 @@
 import pytest
 
 from vllm import SamplingParams
+from vllm.config.diffusion import DiffusionConfig
 from vllm.exceptions import VLLMValidationError
+from vllm.utils.diffusion import validate_diffusion_sampling_params
 from vllm.v1.engine.input_processor import InputProcessor
 
 
@@ -62,21 +64,14 @@ def test_verify_leaves_logits_processors_to_admission():
     )
 
 
-@dataclass
-class MockDiffusionConfig:
-    canvas_length: int = 8
-
-
 def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None):
-    diffusion_config = (
-        MockDiffusionConfig(canvas_length) if canvas_length is not None else None
-    )
-    params.verify(
-        MockModelConfig(is_diffusion=True),
-        None,
-        None,
-        None,
-        diffusion_config=diffusion_config,
+    model_config = MockModelConfig(is_diffusion=True)
+    params.verify(model_config, None, None, None)
+    validate_diffusion_sampling_params(
+        params,
+        canvas_length=canvas_length,
+        vocab_size=model_config.get_vocab_size(),
+        async_scheduling=True,
     )
 
 
@@ -100,7 +95,7 @@ def test_narrow_diffusion_canvas_requires_async_scheduling(
         ),
         speculative_config=None,
         structured_outputs_config=None,
-        diffusion_config=MockDiffusionConfig(canvas_length=8),
+        diffusion_config=DiffusionConfig(canvas_length=8),
         tokenizer=None,
         validate_logits_processors_params=lambda params: None,
     )
diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py
index 50fb75ed149b..369af010f4c7 100644
--- a/tests/v1/sample/test_diffusion_gemma_reads.py
+++ b/tests/v1/sample/test_diffusion_gemma_reads.py
@@ -119,7 +119,7 @@ def _denoise_once(
     width: int = CL,
     embed_weight: torch.Tensor | None = None,
     embed_dtype: torch.dtype = torch.float32,
-) -> None:
+) -> tuple[torch.Tensor, torch.Tensor]:
     """One compiled denoise step over ``slots`` with flat logits, so nothing
     converges by stability or confidence and only the step cap can end it.
     ``width`` below CL runs the step on [:, :width] views, as the sampler
@@ -128,6 +128,8 @@ def _denoise_once(
     device = states.device
     decode_slots = torch.tensor(slots, dtype=torch.int64, device=device)
     decode_idx = torch.arange(n, dtype=torch.int64, device=device)
+    sampled = torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width]
+    num_sampled = torch.zeros(n, dtype=torch.int32, device=device)
     _compiled_sample_step(
         torch.zeros(n * width, VOCAB, device=device),
         decode_slots,
@@ -151,8 +153,9 @@ def _denoise_once(
         states.max_steps,
         states.pin_mask[:, :width],
         states.seed_canvas[:, :width],
-        torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width],
-        torch.zeros(n, dtype=torch.int32, device=device),
+        states.read_only,
+        sampled,
+        num_sampled,
         torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device),
         max_denoising_steps=float(MAX_STEPS),
         t_min=0.5,
@@ -168,6 +171,7 @@ def _denoise_once(
         tp_group_name="",
         compute_sc=compute_sc,
     )
+    return sampled, num_sampled
 
 
 def test_pinned_positions_hold_their_seed_through_a_step():
@@ -194,12 +198,10 @@ def test_add_request_clears_pins():
     states.add_request(0)
     states.set_seed_canvas(0, [1] * CL)
     states.set_pins(0, [2, 3])
-    assert states.pinned_slots == {0}
     assert states.pin_mask[0].tolist() == [False, False, True, True] + [False] * 4
 
     states.add_request(0)
 
-    assert not states.pinned_slots
     assert not states.pin_mask[0].any()
 
 
@@ -291,3 +293,27 @@ def test_stashes_of_different_widths_join():
     assert not out.logprob_token_ids[:3, 3:].any()
     assert torch.isneginf(out.logprobs[:3, 3:]).all()
     assert torch.equal(out.logprobs[3:], wide.logprobs)
+
+
+@pytest.mark.parametrize("width", [4, CL])
+@pytest.mark.parametrize("steps", [1, 2])
+def test_read_emits_at_convergence_while_generation_waits_for_commit(width, steps):
+    states = _states()
+    for slot in (2, 0):
+        states.add_request(slot)
+        states.is_encoder_phase[slot] = False
+        states.max_steps[slot] = steps
+    states.set_read_only(2)
+
+    for step in range(steps):
+        sampled, counts = _denoise_once(states, [2, 0], width=width)
+        assert counts.tolist() == ([width, 0] if step == steps - 1 else [0, 0])
+
+    assert torch.equal(sampled[0], states.argmax_canvas[2, :width].int())
+    assert not states.is_encoder_phase[2]
+    assert states.is_encoder_phase[0]
+    assert not states.self_conditioning_embeds[2].any()
+
+    sampled, counts = _denoise_once(states, [0], width=width)
+    assert counts.tolist() == [width]
+    assert torch.equal(sampled[0], states.argmax_canvas[0, :width].int())
diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py
index daf4f86737d8..a1b4eb3a0d73 100644
--- a/vllm/model_executor/models/diffusion_gemma.py
+++ b/vllm/model_executor/models/diffusion_gemma.py
@@ -423,6 +423,7 @@ def _compiled_sample_step(
     max_steps_tensor: torch.Tensor,  # [max_num_reqs] int32, per-slot step cap
     pin_mask: torch.Tensor,  # [max_num_reqs, CL] bool, positions held at the seed
     seed_canvas: torch.Tensor,  # [max_num_reqs, CL]
+    read_only: torch.Tensor,  # [max_num_reqs] bool, emit without a commit forward
     # Output tensors (modified in-place)
     sampled: torch.Tensor,  # [num_reqs, CL]
     num_sampled: torch.Tensor,  # [num_reqs]
@@ -540,16 +541,6 @@ def _compiled_sample_step(
     new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode))
     history_len_tensor[decode_slots] = new_hist_len
 
-    # Sampled output: commit → emit argmax_canvas, denoise → 0 (pre-zeroed)
-    sampled[decode_idx] = argmax_canvas[decode_slots].to(
-        sampled.dtype
-    ) * is_commit.unsqueeze(1).to(sampled.dtype)
-    # Commit only the real canvas length (== CL except for a canvas truncated
-    # near max_model_len); the padded tail positions are never emitted.
-    num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * valid_canvas_len.to(
-        num_sampled.dtype
-    )
-
     # ---- Phase 6: Stability + convergence ----
     ref = history[decode_slots, 0]
     mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32)
@@ -566,6 +557,10 @@ def _compiled_sample_step(
         is_commit, is_commit.new_zeros(num_decode), converged
     )
 
+    emit = is_commit | (is_denoise & converged & read_only[decode_slots])
+    sampled[decode_idx] = argmax_canvas[decode_slots].to(sampled.dtype) * emit[:, None]
+    num_sampled[decode_idx] = (emit * valid_canvas_len).to(num_sampled.dtype)
+
     # SC soft embedding: store ``probs @ embed_weight`` (the value the next step's
     # self-conditioning MLP consumes) only for slots that will denoise next — i.e.
     # this step denoised AND it isn't about to commit (is_encoder_phase now False).
@@ -607,6 +602,7 @@ def _compiled_sample_step(
     canvas[decode_slots] = torch.where(
         newly_converged, argmax_canvas[decode_slots], canvas[decode_slots]
     )
+    is_encoder_phase[decode_slots] &= ~read_only[decode_slots]
 
     # ---- Phase 7: Copy canvas → draft_tokens for all slots ----
     draft_tokens[all_slots, :CL] = canvas[all_slots]
@@ -696,7 +692,6 @@ def __init__(
         self.pin_mask = torch.zeros(
             max_num_reqs, canvas_length, dtype=torch.bool, device=device
         )
-        self.pinned_slots: set[int] = set()
         # Read-only slots emit on their converging step and skip the commit
         # forward.
         self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
@@ -740,7 +735,6 @@ def add_request(self, slot_idx: int) -> None:
         self.has_seed[slot_idx].fill_(False)
         self.seeded_slots.discard(slot_idx)
         self.pin_mask[slot_idx].fill_(False)
-        self.pinned_slots.discard(slot_idx)
         self.read_only[slot_idx].fill_(False)
         self.read_only_slots.discard(slot_idx)
         self.single_step_slots.discard(slot_idx)
@@ -753,7 +747,6 @@ def remove_request(self, slot_idx: int) -> None:
         self.accepted_canvas_history_len[slot_idx].fill_(0)
         self.self_conditioning_embeds[slot_idx] = 0
         self.seeded_slots.discard(slot_idx)
-        self.pinned_slots.discard(slot_idx)
         self.read_only_slots.discard(slot_idx)
         self.single_step_slots.discard(slot_idx)
 
@@ -773,7 +766,6 @@ def set_pins(self, slot_idx: int, positions: list[int]) -> None:
         self.pin_mask[
             slot_idx, async_tensor_h2d(positions, dtype=torch.int64, device=self.device)
         ] = True
-        self.pinned_slots.add(slot_idx)
 
     def set_read_only(self, slot_idx: int) -> None:
         self.read_only[slot_idx].fill_(True)
@@ -1446,6 +1438,7 @@ def __call__(
                     states.max_steps,
                     states.pin_mask[:, :W],
                     states.seed_canvas[:, :W],
+                    states.read_only,
                     # Output
                     sampled[:, :W],
                     num_sampled,
@@ -1466,10 +1459,11 @@ def __call__(
                     compute_sc=compute_sc,
                 )
 
-                # Logprobs for denoise steps that just converged (is_encoder_phase
-                # flipped False→True), stashed per tile so `scaled` is freed each tile.
+                # Stash newly converged logprobs, including reads that emit now.
                 if want_logprobs:
-                    converged_mask = states.is_encoder_phase[tile_slots]
+                    converged_mask = states.is_encoder_phase[tile_slots] | (
+                        num_sampled[decode_idx[sel]] > 0
+                    )
                     just_converged = converged_mask & ~is_committing[sel]
                     if just_converged.any():
                         flat_logits = scaled.reshape(-1, scaled.shape[-1])
@@ -1512,43 +1506,17 @@ def __call__(
                             )
             run_start = run_end
 
-        # Read-only slots that converged this step emit their argmax canvas
-        # now, skip the encoder phase, and hand out their logprobs below.
-        emit_now: set[int] = set()
-        if states.read_only_slots and not states.read_only_slots.isdisjoint(
-            decode_slots_np.tolist()
-        ):
-            # Read-only slots never enter the encoder phase, so the flag here
-            # means converged this step.
-            ro_mask = (
-                states.read_only[decode_slots] & states.is_encoder_phase[decode_slots]
-            )
-            if bool(ro_mask.any()):
-                ro_idx = decode_idx[ro_mask]
-                ro_slots = decode_slots[ro_mask]
-                sampled[ro_idx] = states.argmax_canvas[ro_slots].to(sampled.dtype)
-                num_sampled[ro_idx] = valid_canvas_len[ro_mask].to(num_sampled.dtype)
-                states.is_encoder_phase[ro_slots] = False
-                emit_now = set(ro_slots.tolist())
-
-        # Commit steps: is_committing was True at entry. Reassemble previously
-        # stashed logprobs and attach to SamplerOutput.
+        # Only emitting requests consume their stashed logprobs.
         logprobs_tensors = None
-        if (
-            want_logprobs
-            and (emit_now or bool(is_committing.any()))
-            and self._pending_logprobs
-        ):
-            committing_slots = (
-                set(decode_slots_np[is_committing.cpu().numpy()].tolist()) | emit_now
-            )
+        if want_logprobs and self._pending_logprobs:
+            emitting_slots = set(slots_np[num_sampled.cpu().numpy() > 0].tolist())
             parts: list[LogprobsTensors] = []
             cu_gen: list[int] = []
             flat_offset = 0
             for i in range(num_reqs):
                 cu_gen.append(flat_offset)
                 slot = int(slots_np[i])
-                if slot in committing_slots and slot in self._pending_logprobs:
+                if slot in emitting_slots and slot in self._pending_logprobs:
                     lp = self._pending_logprobs.pop(slot)
                     parts.append(lp)
                     flat_offset += lp.logprobs.shape[0]
diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py
index ceb35e8b1b8d..120dbaddfd80 100644
--- a/vllm/sampling_params.py
+++ b/vllm/sampling_params.py
@@ -15,12 +15,7 @@
 from pydantic.dataclasses import dataclass
 
 import vllm.envs as envs
-from vllm.config import (
-    DiffusionConfig,
-    ModelConfig,
-    SpeculativeConfig,
-    StructuredOutputsConfig,
-)
+from vllm.config import ModelConfig, SpeculativeConfig, StructuredOutputsConfig
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
 from vllm.tokenizers import TokenizerLike
@@ -823,7 +818,6 @@ def verify(
         speculative_config: SpeculativeConfig | None,
         structured_outputs_config: StructuredOutputsConfig | None,
         tokenizer: TokenizerLike | None,
-        diffusion_config: DiffusionConfig | None = None,
     ) -> None:
         self._validate_logprobs(model_config)
         self._validate_logit_bias(model_config)
@@ -831,7 +825,7 @@ def verify(
         self._validate_stop_token_ids(model_config)
         self._validate_allowed_token_ids(model_config)
         self._validate_spec_decode(speculative_config)
-        self._validate_diffusion(model_config, diffusion_config)
+        self._validate_diffusion(model_config)
         self._validate_structured_outputs(
             model_config, structured_outputs_config, tokenizer
         )
@@ -1043,11 +1037,7 @@ def _validate_spec_decode(
                 "are not yet supported with speculative decoding."
             )
 
-    def _validate_diffusion(
-        self,
-        model_config: ModelConfig,
-        diffusion_config: DiffusionConfig | None = None,
-    ) -> None:
+    def _validate_diffusion(self, model_config: ModelConfig) -> None:
         if not model_config.is_diffusion:
             return
 
@@ -1068,114 +1058,6 @@ def _validate_diffusion(
                 "bad_words, and allowed_token_ids sampling parameters "
                 "are not yet supported with diffusion models."
             )
-        self._validate_diffusion_extra_args(model_config, diffusion_config)
-
-    def _validate_diffusion_extra_args(
-        self,
-        model_config: ModelConfig,
-        diffusion_config: DiffusionConfig | None,
-    ) -> None:
-        """Reject bad diffusion fields in ``extra_args`` before the worker.
-
-        On the GPU a seed id outside the vocabulary is a device-side assert.
-        ``diffusion_seed_canvas``: token ids, one per canvas position, that
-        replace the random canvas after prefill. ``diffusion_pinned``: canvas
-        positions held at their seed value on every denoise step.
-        ``diffusion_max_steps``: denoise steps per canvas.
-        ``diffusion_read_only``: end the request on the first canvas that
-        converges.
-        """
-        extra = self.extra_args
-        if not extra:
-            return
-
-        canvas_length = (
-            None if diffusion_config is None else diffusion_config.canvas_length
-        )
-        width = extra.get("diffusion_canvas_length")
-        if width is not None and (
-            not isinstance(width, int)
-            or isinstance(width, bool)
-            or width < 1
-            or (canvas_length is not None and width > canvas_length)
-        ):
-            raise VLLMValidationError(
-                "diffusion_canvas_length must be a positive integer no larger "
-                "than the served canvas.",
-                parameter="extra_args",
-            )
-        # The request's canvas: its own width, else the served one.
-        expected_len = width or canvas_length
-
-        seed = extra.get("diffusion_seed_canvas")
-        if seed is not None:
-            if not isinstance(seed, (list, tuple)) or not all(
-                isinstance(t, int) and not isinstance(t, bool) for t in seed
-            ):
-                raise VLLMValidationError(
-                    "diffusion_seed_canvas must be a list of token ids.",
-                    parameter="extra_args",
-                )
-            vocab_size = model_config.get_vocab_size()
-            if any(t < 0 or t >= vocab_size for t in seed):
-                raise VLLMValidationError(
-                    f"diffusion_seed_canvas ids must be in [0, {vocab_size}).",
-                    parameter="extra_args",
-                )
-            if expected_len is not None and len(seed) != expected_len:
-                raise VLLMValidationError(
-                    "diffusion_seed_canvas must hold exactly "
-                    f"{expected_len} ids, got {len(seed)}.",
-                    parameter="extra_args",
-                )
-
-        pins = extra.get("diffusion_pinned")
-        if pins is not None:
-            if not isinstance(pins, (list, tuple)) or not all(
-                isinstance(p, int) and not isinstance(p, bool) for p in pins
-            ):
-                raise VLLMValidationError(
-                    "diffusion_pinned must be a list of canvas positions.",
-                    parameter="extra_args",
-                )
-            if seed is None:
-                raise VLLMValidationError(
-                    "diffusion_pinned needs a diffusion_seed_canvas to hold.",
-                    parameter="extra_args",
-                )
-            if any(
-                p < 0 or (expected_len is not None and p >= expected_len) for p in pins
-            ):
-                raise VLLMValidationError(
-                    "diffusion_pinned positions must lie inside the canvas.",
-                    parameter="extra_args",
-                )
-
-        cap = extra.get("diffusion_max_steps")
-        if cap is not None and (
-            not isinstance(cap, int) or isinstance(cap, bool) or cap < 1
-        ):
-            raise VLLMValidationError(
-                "diffusion_max_steps must be a positive integer.",
-                parameter="extra_args",
-            )
-
-        # The OpenAI server's vllm_xargs narrows JSON booleans to 0/1.
-        read_only = extra.get("diffusion_read_only")
-        if read_only is not None and (
-            not isinstance(read_only, (bool, int)) or read_only not in (0, 1)
-        ):
-            raise VLLMValidationError(
-                "diffusion_read_only must be a boolean (or 0/1).",
-                parameter="extra_args",
-            )
-        if read_only:
-            # One canvas is the whole output. Cap max_tokens so the scheduler
-            # ends the request there, and ignore EOS so an end-of-turn token
-            # drawn into a noise slot does not end it early.
-            if expected_len is not None:
-                self.max_tokens = min(self.max_tokens or expected_len, expected_len)
-            self.ignore_eos = True
 
     def _validate_structured_outputs(
         self,
diff --git a/vllm/utils/diffusion.py b/vllm/utils/diffusion.py
new file mode 100644
index 000000000000..ebf513fa8926
--- /dev/null
+++ b/vllm/utils/diffusion.py
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Request validation for discrete diffusion models."""
+
+from typing import TYPE_CHECKING
+
+from vllm.exceptions import VLLMValidationError
+
+if TYPE_CHECKING:
+    from vllm.sampling_params import SamplingParams
+
+
+def validate_diffusion_sampling_params(
+    params: "SamplingParams",
+    *,
+    canvas_length: int | None,
+    vocab_size: int,
+    async_scheduling: bool | None,
+) -> None:
+    """Validate request options and normalize one-canvas reads at admission."""
+    extra = params.extra_args
+    if not extra:
+        return
+
+    def check(valid: bool, message: str) -> None:
+        if not valid:
+            raise VLLMValidationError(message, parameter="extra_args")
+
+    width = extra.get("diffusion_canvas_length")
+    if width is not None:
+        check(
+            isinstance(width, int)
+            and not isinstance(width, bool)
+            and width > 0
+            and (canvas_length is None or width <= canvas_length),
+            "diffusion_canvas_length must be a positive integer no larger "
+            "than the served canvas.",
+        )
+    # The request's canvas: its own width, else the served one.
+    expected_len = width or canvas_length
+
+    seed = extra.get("diffusion_seed_canvas")
+    if seed is not None:
+        check(
+            isinstance(seed, (list, tuple))
+            and all(isinstance(t, int) and not isinstance(t, bool) for t in seed),
+            "diffusion_seed_canvas must be a list of token ids.",
+        )
+        check(
+            all(0 <= t < vocab_size for t in seed),
+            f"diffusion_seed_canvas ids must be in [0, {vocab_size}).",
+        )
+        check(
+            expected_len is None or len(seed) == expected_len,
+            f"diffusion_seed_canvas must hold exactly {expected_len} ids, "
+            f"got {len(seed)}.",
+        )
+
+    pins = extra.get("diffusion_pinned")
+    if pins is not None:
+        check(
+            isinstance(pins, (list, tuple))
+            and all(isinstance(p, int) and not isinstance(p, bool) for p in pins),
+            "diffusion_pinned must be a list of canvas positions.",
+        )
+        check(
+            seed is not None,
+            "diffusion_pinned needs a diffusion_seed_canvas to hold.",
+        )
+        check(
+            all(p >= 0 and (expected_len is None or p < expected_len) for p in pins),
+            "diffusion_pinned positions must lie inside the canvas.",
+        )
+
+    cap = extra.get("diffusion_max_steps")
+    if cap is not None:
+        check(
+            isinstance(cap, int) and not isinstance(cap, bool) and cap > 0,
+            "diffusion_max_steps must be a positive integer.",
+        )
+
+    # The OpenAI server's vllm_xargs narrows JSON booleans to 0/1.
+    read_only = extra.get("diffusion_read_only")
+    if read_only is not None:
+        check(
+            isinstance(read_only, (bool, int)) and read_only in (0, 1),
+            "diffusion_read_only must be a boolean (or 0/1).",
+        )
+    if read_only:
+        # One canvas is the whole output. Cap max_tokens so the scheduler
+        # ends the request there, and ignore EOS so an end-of-turn token
+        # drawn into a noise slot does not end it early.
+        if expected_len is not None:
+            params.max_tokens = min(params.max_tokens or expected_len, expected_len)
+        params.ignore_eos = True
+
+    check(
+        bool(async_scheduling) or width is None or width == canvas_length,
+        "A diffusion_canvas_length smaller than the served canvas "
+        "requires --async-scheduling.",
+    )
diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py
index db71de57afa6..6b66b5c484e1 100644
--- a/vllm/v1/engine/input_processor.py
+++ b/vllm/v1/engine/input_processor.py
@@ -30,6 +30,7 @@
 from vllm.tokenizers import TokenizerLike
 from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
 from vllm.utils.async_utils import make_async
+from vllm.utils.diffusion import validate_diffusion_sampling_params
 from vllm.utils.jsontree import json_iter_leaves
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.kv_hints import KvHintsEnvelope
@@ -126,24 +127,17 @@ def _validate_params(
                 self.speculative_config,
                 self.structured_outputs_config,
                 self.tokenizer,
-                diffusion_config=self.diffusion_config,
             )
 
             self.validate_logits_processors_params(params)
 
-            if (
-                self.diffusion_config is not None
-                and self.model_config.is_diffusion
-                and not self.vllm_config.scheduler_config.async_scheduling
-                and params.extra_args
-            ):
-                width = params.extra_args.get("diffusion_canvas_length")
-                if width is not None and width != self.diffusion_config.canvas_length:
-                    raise VLLMValidationError(
-                        "A diffusion_canvas_length smaller than the served canvas "
-                        "requires --async-scheduling.",
-                        parameter="extra_args",
-                    )
+            if self.diffusion_config is not None and self.model_config.is_diffusion:
+                validate_diffusion_sampling_params(
+                    params,
+                    canvas_length=self.diffusion_config.canvas_length,
+                    vocab_size=self.model_config.get_vocab_size(),
+                    async_scheduling=self.vllm_config.scheduler_config.async_scheduling,
+                )
 
             if self.model_config.return_sampling_mask:
                 if params.temperature <= 0:

From 7ba66756844edf9686cd193f9e06f4516046069d Mon Sep 17 00:00:00 2001
From: Matt Mastracci 
Date: Mon, 21 Sep 2026 18:06:58 -0600
Subject: [PATCH 24/25] [Core] Validate diffusion extra_args whenever the model
 is a diffusion model

The validator moved to admission and ran only with a diffusion config,
which the engine leaves unset without --diffusion-config. A diffusion
model served that way took seed canvases unchecked and read-only requests
unnormalised. Run it for any diffusion model, with the served canvas
unknown when there is no config; only the narrow-canvas check needs the
served width.

Co-Authored-By: Claude Fable 5.1 
Signed-off-by: Matt Mastracci 
---
 tests/test_sampling_params.py                     | 15 +++++++++++++++
 .../engine/test_input_processor_trace_replay.py   |  1 +
 vllm/utils/diffusion.py                           |  5 ++++-
 vllm/v1/engine/input_processor.py                 | 10 ++++++++--
 4 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py
index 635842e3a121..84ce45caa958 100644
--- a/tests/test_sampling_params.py
+++ b/tests/test_sampling_params.py
@@ -75,6 +75,21 @@ def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None):
     )
 
 
+def test_diffusion_extra_args_are_validated_without_a_served_canvas():
+    # No --diffusion-config: the canvas is unknown, the ids are still checked
+    # and a read-only request is still normalised.
+    params = SamplingParams(
+        max_tokens=64,
+        extra_args={"diffusion_seed_canvas": [0, 1], "diffusion_read_only": True},
+    )
+    _verify_diffusion(params, canvas_length=None)
+    assert params.ignore_eos is True
+
+    bad = SamplingParams(extra_args={"diffusion_seed_canvas": [0, 10**9]})
+    with pytest.raises(VLLMValidationError, match="ids must be in"):
+        _verify_diffusion(bad, canvas_length=None)
+
+
 @pytest.mark.parametrize("async_scheduling", [False, True])
 @pytest.mark.parametrize(
     "extra_args",
diff --git a/tests/v1/engine/test_input_processor_trace_replay.py b/tests/v1/engine/test_input_processor_trace_replay.py
index 5fc1118167e5..957f70b71858 100644
--- a/tests/v1/engine/test_input_processor_trace_replay.py
+++ b/tests/v1/engine/test_input_processor_trace_replay.py
@@ -80,6 +80,7 @@ def _validate(enable_trace_replay: bool) -> None:
         model_config=SimpleNamespace(
             return_sampling_mask=False,
             enable_trace_replay=enable_trace_replay,
+            is_diffusion=False,
         ),
         vllm_config=SimpleNamespace(reasoning_config=None),
         speculative_config=None,
diff --git a/vllm/utils/diffusion.py b/vllm/utils/diffusion.py
index ebf513fa8926..278c17d83823 100644
--- a/vllm/utils/diffusion.py
+++ b/vllm/utils/diffusion.py
@@ -95,7 +95,10 @@ def check(valid: bool, message: str) -> None:
         params.ignore_eos = True
 
     check(
-        bool(async_scheduling) or width is None or width == canvas_length,
+        bool(async_scheduling)
+        or width is None
+        or canvas_length is None
+        or width == canvas_length,
         "A diffusion_canvas_length smaller than the served canvas "
         "requires --async-scheduling.",
     )
diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py
index 6b66b5c484e1..199f8898531f 100644
--- a/vllm/v1/engine/input_processor.py
+++ b/vllm/v1/engine/input_processor.py
@@ -131,10 +131,16 @@ def _validate_params(
 
             self.validate_logits_processors_params(params)
 
-            if self.diffusion_config is not None and self.model_config.is_diffusion:
+            if self.model_config.is_diffusion:
+                # Without --diffusion-config the served canvas is unknown here;
+                # the ids and the read-only normalisation are still checked.
                 validate_diffusion_sampling_params(
                     params,
-                    canvas_length=self.diffusion_config.canvas_length,
+                    canvas_length=(
+                        self.diffusion_config.canvas_length
+                        if self.diffusion_config is not None
+                        else None
+                    ),
                     vocab_size=self.model_config.get_vocab_size(),
                     async_scheduling=self.vllm_config.scheduler_config.async_scheduling,
                 )

From a7c23ac96d7806e7c7e7d862eadbce5a33529b94 Mon Sep 17 00:00:00 2001
From: Matt Mastracci 
Date: Tue, 22 Sep 2026 12:03:07 -0600
Subject: [PATCH 25/25] [Bugfix] DiffusionGemma: register the top-level model
 type with the Gemma4 arch convertor

Checkpoints carry model_type "diffusion_gemma" at the top level and
"diffusion_gemma_text" inside text_config. The convertor table only knew
the inner name, so the base convertor ran and reported head_dim 256 for
every layer. Gemma4Config's heterogeneous head-dim handling then returned
early: no FA4 forced onto every layer, no TRITON_ATTN fallback.

At backend selection the full-attention layers still declare head_dim
512. FlashAttention drops out of that KV-cache group wherever FA4 is not
built (L40S, 5090, GB10 nightly) and FlashInfer is next. Its metadata
builder does `causal or self.use_xqa` on the per-request causal tensor,
so the engine dies at CUDA-graph capture with "Boolean value of Tensor
with more than one value is ambiguous". Three users hit this on the PR.
Our own launches always passed --attention-backend TRITON_ATTN.

Registering "diffusion_gemma" makes the Gemma4 logic run: FA4 on every
layer when available, else TRITON_ATTN. Verified on a GB10 with the
nightly image and no backend flag. Before: FLASHINFER chosen out of
['FLASHINFER', 'TRITON_ATTN'], then the capture crash. After: "FA4 not
available, forcing TRITON_ATTN", engine up, read batteries unchanged
(open extraction 33/33, grounded region 30/33, coupled pins 23/24,
question slots 8/8).

The use_non_causal guard's comment and log now say what the flag does:
only the sm_100 default order reads it. The README says how the backend
is picked.

Co-Authored-By: Claude Fable 5.1 
Signed-off-by: Matt Mastracci 
---
 .../features/structured_diffusion/README.md   |  5 ++++
 tests/config/test_model_arch_config.py        | 25 +++++++++++++++++++
 vllm/model_executor/models/config.py          |  9 ++++---
 .../model_arch_config_convertor.py            |  1 +
 4 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/examples/features/structured_diffusion/README.md b/examples/features/structured_diffusion/README.md
index 317fa4d4bc04..9840f9ea3f01 100644
--- a/examples/features/structured_diffusion/README.md
+++ b/examples/features/structured_diffusion/README.md
@@ -29,6 +29,11 @@ curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -
   ]}'
 ```
 
+The attention backend is picked as for Gemma 4: FlashAttention 4 on every
+layer when available, otherwise Triton. FlashInfer cannot serve this model (a
+batch mixes causal prefill with bidirectional denoising), and
+`--attention-backend FLASHINFER` is rejected.
+
 Per-request canvas widths smaller than the served canvas require async
 scheduling. The diffusion async scheduler is selected automatically; no
 `--scheduler-cls` argument is needed. Synchronous execution supports full-width
diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py
index b68820c0d51c..90dbbf278bd3 100644
--- a/tests/config/test_model_arch_config.py
+++ b/tests/config/test_model_arch_config.py
@@ -15,8 +15,10 @@
 
 from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
 from vllm.config.model_arch import ModelArchitectureConfig
+from vllm.transformers_utils.configs.diffusion_gemma import DiffusionGemmaConfig
 from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
 from vllm.transformers_utils.model_arch_config_convertor import (
+    MODEL_ARCH_CONFIG_CONVERTORS,
     Gemma4ModelArchConfigConvertor,
     ModelArchConfigConvertorBase,
 )
@@ -496,3 +498,26 @@ def test_deepseek_v4_convertor_splits_vision_architecture():
         conv = DeepseekV4ModelArchConfigConvertor(draft_cfg, draft_cfg)
         assert conv.get_architectures() == [draft_arch]
         assert draft_cfg.architectures == [draft_arch]
+
+
+def test_diffusion_gemma_head_dims_vary_by_layer_type():
+    """Checkpoints carry model_type `diffusion_gemma` at the top level. Under
+    the base convertor every layer reports the sliding head dim, Gemma4Config
+    skips its head-dim handling, and the full-attention KV-cache group falls to
+    FlashInfer, which cannot take a per-request causal tensor."""
+    text_config = dict(
+        num_hidden_layers=6,
+        hidden_size=64,
+        num_attention_heads=8,
+        num_key_value_heads=4,
+        head_dim=16,
+        global_head_dim=32,
+        layer_types=["sliding_attention"] * 5 + ["full_attention"],
+    )
+    hf_config = DiffusionGemmaConfig(text_config=text_config)
+
+    convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS[hf_config.model_type]
+    assert convertor_cls is Gemma4ModelArchConfigConvertor
+    arch = convertor_cls(hf_config, hf_config.text_config).convert()
+
+    assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32]
diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py
index 8288eab236f3..3a0392c52322 100644
--- a/vllm/model_executor/models/config.py
+++ b/vllm/model_executor/models/config.py
@@ -280,11 +280,14 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
                 "FLASH_ATTN or TRITON_ATTN instead."
             )
         if attention_config.backend is None and not attention_config.use_non_causal:
+            # Only the sm_100 default backend order reads this flag (it puts
+            # FlashInfer first for causal models). Elsewhere Gemma4Config above
+            # keeps FlashInfer out: FA4 on every layer, else TRITON_ATTN.
             attention_config.use_non_causal = True
             logger.info(
-                "DiffusionGemma uses mixed causal/bidirectional attention "
-                "within a batch; setting use_non_causal=True to exclude "
-                "FlashInfer from auto-selection."
+                "DiffusionGemma mixes causal and bidirectional attention in "
+                "one batch; setting use_non_causal=True so the default backend "
+                "order prefers FlashAttention over FlashInfer."
             )
 
         # Auto-create DiffusionConfig from HF config if not provided.
diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py
index 72f2920480a1..f0126ca34581 100644
--- a/vllm/transformers_utils/model_arch_config_convertor.py
+++ b/vllm/transformers_utils/model_arch_config_convertor.py
@@ -813,6 +813,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]:
     "dbrx": DbrxModelArchConfigConvertor,
     "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,
     "deepseek_v4": DeepseekV4ModelArchConfigConvertor,
+    "diffusion_gemma": Gemma4ModelArchConfigConvertor,
     "diffusion_gemma_text": Gemma4ModelArchConfigConvertor,
     "ernie_mtp": ErnieMTPModelArchConfigConvertor,
     "falcon": FalconModelArchConfigConvertor,