[Bugfix][CPU][DiffusionGemma] Support narrower canvas w/ sync scheduling (#59107)
Signed-off-by: ojeda-e <estefania.b.ojeda@gmail.com>
Mia Ojeda committed
Sep 29, 2026 at 21:28 UTC
e006d761a526d0a69555a59f7c10c0c889babd98
7 files changed
+78
-41
examples/features/structured_diffusion/README.md
+3
-4
@@ -36,10 +36,9 @@ layer when available, otherwise Triton. FlashInfer cannot serve this model (a
36
batch mixes causal prefill with bidirectional denoising), and
37
`--attention-backend FLASHINFER` is rejected.
38
39
-Per-request canvas widths smaller than the served canvas require async
40
-scheduling. The diffusion async scheduler is selected automatically; no
41
-`--scheduler-cls` argument is needed. Synchronous execution supports full-width
42
-canvases only.
39
+Per-request canvas widths may be smaller than the served canvas with either
40
+synchronous or asynchronous scheduling. Omit `diffusion_canvas_length` to use
41
+the served canvas width.
42
43
Question types: `noul` (yes/no), `choice` with `options`, `score` with
44
ordered `levels`. Each label must be a single token in the answer template,
tests/test_sampling_params.py
+2
-7
@@ -71,7 +71,6 @@ def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None):
71
params,
72
canvas_length=canvas_length,
73
vocab_size=model_config.get_vocab_size(),
74
- async_scheduling=True,
74
)
75
76
@@ -100,7 +99,7 @@ def test_diffusion_extra_args_are_validated_without_a_served_canvas():
99
{"diffusion_canvas_length": 8},
100
],
101
)
103
-def test_narrow_diffusion_canvas_requires_async_scheduling(
102
+def test_diffusion_canvas_width_is_accepted_with_either_scheduler(
103
async_scheduling, extra_args
104
):
105
processor = SimpleNamespace(
@@ -115,11 +114,7 @@ def test_narrow_diffusion_canvas_requires_async_scheduling(
114
validate_logits_processors_params=lambda params: None,
115
)
116
params = SamplingParams(extra_args=extra_args)
118
- if not async_scheduling and extra_args.get("diffusion_canvas_length") == 4:
119
- with pytest.raises(VLLMValidationError, match="requires --async-scheduling"):
120
- InputProcessor._validate_params(processor, params, ("generate",))
121
- else:
122
- InputProcessor._validate_params(processor, params, ("generate",))
117
+ InputProcessor._validate_params(processor, params, ("generate",))
118
119
120
@pytest.mark.parametrize(
tests/v1/core/test_scheduler.py
+48
-9
@@ -42,6 +42,7 @@ from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
42
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
43
from vllm.v1.core.sched.diffusion_scheduler import (
44
DiffusionAsyncScheduler,
45
+ DiffusionScheduler,
46
diffusion_canvas_width,
47
)
48
from vllm.v1.core.sched.interface import PauseState
@@ -7109,14 +7110,15 @@ def diffusion_model_runner(monkeypatch):
7110
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
7111
7112
7112
-def _diffusion_scheduler(**kwargs) -> DiffusionAsyncScheduler:
7113
+def _diffusion_scheduler(async_scheduling: bool = True, **kwargs) -> DiffusionScheduler:
7114
+ scheduler_cls = DiffusionAsyncScheduler if async_scheduling else DiffusionScheduler
7115
scheduler = create_scheduler(
7114
- async_scheduling=True,
7116
+ async_scheduling=async_scheduling,
7117
diffusion_canvas_length=8,
7116
- scheduler_cls=DiffusionAsyncScheduler,
7118
+ scheduler_cls=scheduler_cls,
7119
**kwargs,
7120
)
7119
- assert isinstance(scheduler, DiffusionAsyncScheduler)
7121
+ assert isinstance(scheduler, scheduler_cls)
7122
return scheduler
7123
7124
@@ -7127,7 +7129,7 @@ def test_diffusion_scheduler_is_selected_by_default(async_scheduling):
7129
async_scheduling=async_scheduling, diffusion_canvas_length=8
7130
).vllm_config.scheduler_config
7131
assert config.get_scheduler_cls() is (
7130
- DiffusionAsyncScheduler if config.async_scheduling else Scheduler
7132
+ DiffusionAsyncScheduler if config.async_scheduling else DiffusionScheduler
7133
)
7134
7135
@@ -7149,10 +7151,13 @@ def test_diffusion_scheduler_narrows_the_canvas_per_request():
7151
7152
7153
@pytest.mark.parametrize("structured", [False, True])
7154
+@pytest.mark.parametrize("async_scheduling", [False, True])
7155
@pytest.mark.usefixtures("diffusion_model_runner")
7153
-def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
7156
+def test_diffusion_scheduler_trims_full_width_worker_drafts(
7157
+ structured, async_scheduling
7158
+):
7159
"""Padded worker drafts must be narrowed before scheduling or grammar validation."""
7155
- scheduler = _diffusion_scheduler()
7160
+ scheduler = _diffusion_scheduler(async_scheduling=async_scheduling)
7161
wide = _diffusion_request("wide", {})
7162
narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4})
7163
for request in (wide, narrow):
@@ -7167,8 +7172,12 @@ def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
7172
)
7173
tokens = list(range(8)) if structured else [-1] * 8
7174
drafts = DraftTokenIds(["wide", "narrow"], [tokens.copy(), tokens.copy()])
7170
- output = scheduler.schedule()
7171
- scheduler.update_draft_token_ids_in_output(drafts, output)
7175
+ if async_scheduling:
7176
+ output = scheduler.schedule()
7177
+ scheduler.update_draft_token_ids_in_output(drafts, output)
7178
+ else:
7179
+ scheduler.update_draft_token_ids(drafts)
7180
+ output = scheduler.schedule()
7181
7182
assert output.scheduled_spec_decode_tokens == {
7183
"wide": tokens,
@@ -7180,6 +7189,36 @@ def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
7189
assert narrow.structured_output_request.grammar.seen == [tokens[:4]]
7190
7191
7192
+@pytest.mark.parametrize(
7193
+ "sampled, expected_status",
7194
+ [
7195
+ pytest.param([], RequestStatus.RUNNING, id="denoising"),
7196
+ pytest.param([1] * 4, RequestStatus.FINISHED_LENGTH_CAPPED, id="final"),
7197
+ ],
7198
+)
7199
+@pytest.mark.usefixtures("diffusion_model_runner")
7200
+def test_sync_diffusion_step_commits_only_emitted_tokens(sampled, expected_status):
7201
+ """An empty denoising step reuses the canvas; a completed read emits it."""
7202
+ scheduler = _diffusion_scheduler(async_scheduling=False)
7203
+ request = _diffusion_request(
7204
+ "read", {"diffusion_canvas_length": 4, "diffusion_read_only": True}
7205
+ )
7206
+ request.max_tokens = request.sampling_params.max_tokens = 4
7207
+ scheduler.add_request(request)
7208
+ _model_output(scheduler, scheduler.schedule(), [[]])
7209
+
7210
+ scheduler.update_draft_token_ids(
7211
+ DraftTokenIds([request.request_id], [list(range(8))])
7212
+ )
7213
+ output = scheduler.schedule()
7214
+ assert output.num_scheduled_tokens == {request.request_id: 4}
7215
+ _model_output(scheduler, output, [sampled])
7216
+
7217
+ assert request.num_computed_tokens == request.num_prompt_tokens + len(sampled)
7218
+ assert list(request.output_token_ids) == sampled
7219
+ assert request.status == expected_status
7220
+
7221
+
7222
@pytest.mark.usefixtures("diffusion_model_runner")
7223
def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight():
7224
scheduler = _diffusion_scheduler()
vllm/config/vllm.py
+6
-2
@@ -2020,11 +2020,15 @@ class VllmConfig:
2020
# After the platform hook, which has the last word on async scheduling.
2021
if (
2022
self.diffusion_config is not None
2023
- and self.scheduler_config.async_scheduling
2023
and self.scheduler_config.scheduler_cls is None
2024
):
2025
+ scheduler_name = (
2026
+ "DiffusionAsyncScheduler"
2027
+ if self.scheduler_config.async_scheduling
2028
+ else "DiffusionScheduler"
2029
+ )
2030
self.scheduler_config.scheduler_cls = (
2027
- "vllm.v1.core.sched.diffusion_scheduler.DiffusionAsyncScheduler"
2031
+ f"vllm.v1.core.sched.diffusion_scheduler.{scheduler_name}"
2032
)
2033
2034
self._normalize_piecewise_cudagraph_mode(
vllm/utils/diffusion.py
-10
@@ -15,7 +15,6 @@ def validate_diffusion_sampling_params(
15
*,
16
canvas_length: int | None,
17
vocab_size: int,
18
- async_scheduling: bool | None,
18
) -> None:
19
"""Validate request options and normalize one-canvas reads at admission."""
20
extra = params.extra_args
@@ -105,12 +104,3 @@ def validate_diffusion_sampling_params(
104
bool(params.logprob_token_ids),
105
"diffusion_constrained needs logprob_token_ids: they are the allowed set.",
106
)
108
-
109
- check(
110
- bool(async_scheduling)
111
- or width is None
112
- or canvas_length is None
113
- or width == canvas_length,
114
- "A diffusion_canvas_length smaller than the served canvas "
115
- "requires --async-scheduling.",
116
- )
vllm/v1/core/sched/diffusion_scheduler.py
+19
-8
@@ -1,15 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Async scheduling for diffusion requests.
4
-
5
-These rules live here so Scheduler.schedule() and AsyncScheduler stay
6
-unchanged. VllmConfig selects this class for a diffusion model under async
7
-scheduling. A sync scheduler creates no output placeholders, which both rules
8
-read.
9
-"""
3
+"""Canvas-width handling and async read deferral for diffusion requests."""
4
5
from vllm.v1.core.sched.async_scheduler import AsyncScheduler
6
from vllm.v1.core.sched.output import SchedulerOutput
7
+from vllm.v1.core.sched.scheduler import Scheduler
8
+from vllm.v1.outputs import DraftTokenIds
9
from vllm.v1.request import Request
10
11
@@ -38,7 +34,22 @@ def _read_in_flight(request: Request, width: int) -> bool:
34
return request.num_output_placeholders >= steps * width
35
36
41
-class DiffusionAsyncScheduler(AsyncScheduler):
37
+class DiffusionScheduler(Scheduler):
38
+ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
39
+ for i, (req_id, token_ids) in enumerate(
40
+ zip(draft_token_ids.req_ids, draft_token_ids.draft_token_ids)
41
+ ):
42
+ request = self.requests.get(req_id)
43
+ if request is None:
44
+ continue
45
+ width = diffusion_canvas_width(request, self.num_spec_tokens)
46
+ # Remove padded positions before the base scheduler validates grammar.
47
+ if len(token_ids) > width:
48
+ draft_token_ids.draft_token_ids[i] = token_ids[:width]
49
+ super().update_draft_token_ids(draft_token_ids)
50
+
51
+
52
+class DiffusionAsyncScheduler(DiffusionScheduler, AsyncScheduler):
53
def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
54
for request in self.running:
55
width = diffusion_canvas_width(request, self.num_spec_tokens)
vllm/v1/engine/input_processor.py
-1
@@ -158,7 +158,6 @@ class InputProcessor:
158
else None
159
),
160
vocab_size=self.model_config.get_vocab_size(),
161
- async_scheduling=self.vllm_config.scheduler_config.async_scheduling,
161
)
162
163
if self.model_config.return_sampling_mask: