[Core] structured generation mode for DiffusionGemma model (Jev-like) (#57250)

Signed-off-by: Matt Mastracci <matthew@mastracci.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Razorback16 <razorback16@protonmail.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Razorback16 <razorback16@protonmail.com>

examples/features/structured_diffusion/README.md

@@ -0,0 +1,74 @@
+# 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_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 |
+
+`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/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": [
+    {"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.\"}"}
+  ]}'
+```
+
+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
+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,
+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.
+
+```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.

tests/config/test_model_arch_config.py

@@ -15,8 +15,10 @@ from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
 
 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]

tests/test_sampling_params.py

@@ -1,11 +1,15 @@
 # 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.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
 
 
 @dataclass
@@ -13,6 +17,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
@@ -59,6 +64,183 @@ def test_verify_leaves_logits_processors_to_admission():
     )
 
 
+def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None):
+    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,
+    )
+
+
+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",
+    [
+        {},
+        {"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=DiffusionConfig(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",
+    [
+        ({"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"),
+        ({"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):
+    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)
+
+
+@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
+
+
+@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_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_pinned": [0, 1, 7],
+            "diffusion_max_steps": 4,
+            "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."""

tests/v1/core/test_scheduler.py

@@ -40,6 +40,10 @@ from vllm.utils.hashing import sha256
 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.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
@@ -6872,3 +6876,129 @@ def test_update_draft_token_ids_in_output_strips_padding():
         -1,
     ]
     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]
+    )
+    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
+
+
+@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 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
+
+
+@pytest.mark.parametrize("structured", [False, True])
+def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
+    """Padded worker drafts must be narrowed before scheduling or grammar validation."""
+    scheduler = _diffusion_scheduler()
+    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()])
+    output = scheduler.schedule()
+    scheduler.update_draft_token_ids_in_output(drafts, output)
+
+    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(
+        "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

tests/v1/core/utils.py

@@ -7,6 +7,7 @@ import vllm.envs as envs
 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,

tests/v1/engine/test_input_processor_trace_replay.py

@@ -80,10 +80,12 @@ 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,
         structured_outputs_config=None,
+        diffusion_config=None,
         tokenizer=None,
         validate_logits_processors_params=lambda params: None,
     )

tests/v1/sample/test_diffusion_gemma_reads.py

@@ -0,0 +1,319 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-request DiffusionGemma state behind structured reads: seed canvases,
+pinned positions, 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,
+    _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"
+)
+
+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_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)
+    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
+
+
+def _denoise_once(
+    states: DiffusionGemmaRequestStates,
+    slots: list[int],
+    compute_sc: bool = True,
+    width: int = CL,
+    embed_weight: torch.Tensor | None = None,
+    embed_dtype: torch.dtype = torch.float32,
+) -> 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
+    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)
+    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,
+        decode_idx,
+        decode_slots,
+        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[:, :width],
+        (
+            torch.zeros(VOCAB, 4, dtype=embed_dtype, device=device)
+            if embed_weight is None
+            else embed_weight
+        ),
+        torch.tensor(1.0, dtype=embed_dtype, device=device),
+        states.accepted_canvas_history[:, :, :width],
+        states.accepted_canvas_history_len,
+        states.max_steps,
+        states.pin_mask[:, :width],
+        states.seed_canvas[:, :width],
+        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,
+        t_max=1.0,
+        confidence_threshold=0.1,
+        vocab_size=VOCAB,
+        CL=width,
+        ST=states.stability_threshold,
+        entropy_bound=0.1,
+        sc_vocab_start=0,
+        sc_vocab_end=VOCAB,
+        tp_size=1,
+        tp_group_name="",
+        compute_sc=compute_sc,
+    )
+    return sampled, num_sampled
+
+
+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.pin_mask[0].tolist() == [False, False, True, True] + [False] * 4
+
+    states.add_request(0)
+
+    assert not states.pin_mask[0].any()
+
+
+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()
+
+
+@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)
+    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):
+        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]
+
+
+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)
+
+
+@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())

vllm/config/scheduler.py

@@ -232,13 +232,17 @@ class SchedulerConfig:
 
         # 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)

vllm/config/vllm.py

@@ -1894,6 +1894,16 @@ class VllmConfig:
             )
         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
         )

vllm/model_executor/models/config.py

@@ -280,11 +280,14 @@ class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig):
                 "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.

vllm/model_executor/models/diffusion_gemma.py

@@ -22,6 +22,7 @@ from typing import Any
 
 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,34 @@ 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
+# 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]
@@ -391,6 +420,10 @@ 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] 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]
@@ -412,6 +445,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.
@@ -478,8 +512,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(
@@ -504,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)
@@ -523,13 +550,17 @@ 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 >= max_steps_tensor[decode_slots]
     )
     # Commit done → denoise next (False); denoise converged → commit next (True)
     is_encoder_phase[decode_slots] = torch.where(
         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).
@@ -537,24 +568,41 @@ 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
+        # 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. 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:
+        # 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)
     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]
@@ -628,6 +676,32 @@ class DiffusionGemmaRequestStates:
         # Per-slot confidence flag, set by the sampler each step.
         self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
 
+        # 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
+        )
+        # 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()
+        # 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
+        )
+        # 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()
+        # 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
         # of the full [.., vocab] distribution shrinks this buffer by
@@ -657,11 +731,57 @@ class DiffusionGemmaRequestStates:
         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_(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.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
+        # 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
+        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:
+        """``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)
+        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
+
+    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):
@@ -1030,6 +1150,33 @@ class DiffusionSampler:
         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:
+            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)
+        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) != width:
+                raise ValueError(
+                    f"diffusion_seed_canvas must hold exactly {width} ids, "
+                    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)
 
     def apply_staged_writes(self) -> None:
         self.sampling_states.apply_staged_writes()
@@ -1071,6 +1218,7 @@ class DiffusionSampler:
             ps.astype(np.int64), device=states.is_encoder_phase.device
         )
         states.init_canvas(ps_gpu)
+        states.apply_seed_canvases(ps, ps_gpu)
         self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[
             ps_gpu
         ]
@@ -1151,7 +1299,6 @@ class DiffusionSampler:
         # --- 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])
 
@@ -1194,18 +1341,12 @@ class DiffusionSampler:
             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.
@@ -1228,125 +1369,159 @@ class DiffusionSampler:
         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))
+            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()
+                    )
+                )
 
-        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]
+                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,
+                    states.pin_mask[:, :W],
+                    states.seed_canvas[:, :W],
+                    states.read_only,
+                    # 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,
+                )
 
-            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,
-                # 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,
-            )
+                # Stash newly converged logprobs, including reads that emit now.
+                if want_logprobs:
+                    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])
+                        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
 
-            # 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)
-                    for local_idx in just_converged.nonzero(as_tuple=True)[0]:
-                        li = local_idx.item()
-                        slot = tile_slots[local_idx]
-                        # 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
-                        per_req_ids = max_token_ids > 0
-                        self._pending_logprobs[slot.item()] = compute_topk_scores(
-                            flat_logits[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.item(),
-                                    dtype=torch.int32,
-                                    device=flat_logits.device,
-                                )
-                                if per_req_ids
-                                else None
-                            ),
-                            max_per_req_token_ids=max_token_ids,
-                            logits_mode=self.logits_mode,
-                        )
-
-        # 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 is_committing.any() and self._pending_logprobs:
-            committing_slots = set(
-                decode_slots_np[is_committing.cpu().numpy()].tolist()
-            )
-            parts_ids, parts_lp, parts_ranks = [], [], []
+        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_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,

vllm/transformers_utils/model_arch_config_convertor.py

@@ -813,6 +813,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
     "dbrx": DbrxModelArchConfigConvertor,
     "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,
     "deepseek_v4": DeepseekV4ModelArchConfigConvertor,
+    "diffusion_gemma": Gemma4ModelArchConfigConvertor,
     "diffusion_gemma_text": Gemma4ModelArchConfigConvertor,
     "ernie_mtp": ErnieMTPModelArchConfigConvertor,
     "falcon": FalconModelArchConfigConvertor,

vllm/utils/diffusion.py

@@ -0,0 +1,104 @@
+# 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 canvas_length is None
+        or width == canvas_length,
+        "A diffusion_canvas_length smaller than the served canvas "
+        "requires --async-scheduling.",
+    )

vllm/v1/core/sched/diffusion_scheduler.py

@@ -0,0 +1,56 @@
+# 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.
+"""
+
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
+from vllm.v1.core.sched.output import SchedulerOutput
+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.
+
+    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:
+            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
+                # 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)

vllm/v1/engine/input_processor.py

@@ -30,6 +30,7 @@ from vllm.tasks import GENERATION_TASKS, POOLING_TASKS, SupportedTask
 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
@@ -53,6 +54,7 @@ class InputProcessor:
         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 = (
@@ -129,6 +131,20 @@ class InputProcessor:
 
             self.validate_logits_processors_params(params)
 
+            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
+                        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,
+                )
+
             if self.model_config.return_sampling_mask:
                 if params.temperature <= 0:
                     raise ValueError(