@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ab3de6edf2

[Perf] DiffusionGemma: constrained reads over the request's logprob_token_ids (#58216)

Signed-off-by: Matt Mastracci <matthew@mastracci.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Misha Goin <mgoin64@gmail.com>

Matt Mastracci committed Sep 24, 2026 at 14:15 UTC ab3de6edf2b34507e066d998c09f2361a5c1d8ce

6 files changed +262 -6

examples/features/structured_diffusion/README.md

+3 -1

@@ -11,8 +11,10 @@ left as noise, one denoise step gives a distribution over each slot. These

11 | `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 |

12 | `diffusion_max_steps` | `int` | denoise steps before the canvas is emitted |

13 | `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 |

14 +| `diffusion_constrained` | `bool` | run the unembedding, sampler and self-conditioning over the request's `logprob_token_ids` only. Logprobs are normalized over that set. A step uses this only when every read in it has the same set |

15 -`structured_server.py` turns a question schema into those fields. It serves

16 +`structured_server.py` turns a question schema into those fields, with

17 +`diffusion_constrained` on for every read (`--no-constrained` turns it off). It serves

18 `/v1/chat/completions`: the system message is the schema, the user message

19 is the state JSON, and the reply content is one distribution per question

20 with a standard error over a few noise draws.

examples/features/structured_diffusion/structured_server.py

+16

@@ -399,6 +399,14 @@ def canvas_width(template):

399 return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP)

400

401

402 +def constrained_xargs():

403 + """Read over the labels only. The engine runs the unembedding, sampler and

404 + self-conditioning over the request's logprob_token_ids instead of the whole

405 + vocabulary. Same argmax, about a quarter less GPU time per read. Logprobs

406 + come back normalized over the labels."""

407 + return {"diffusion_constrained": True} if ARGS.constrained else {}

408 +

409 +

410 def pin_xargs(template, slots, steps):

411 """Past one denoise step the template must be held, or accept/renoise

412 rewrites it: pin every canvas position that is not an answer slot."""

@@ -552,6 +560,7 @@ def one_read(

560 "diffusion_max_steps": schema["steps"],

561 "diffusion_read_only": True,

562 **pin_xargs(template, slots, schema["steps"]),

563 + **constrained_xargs(),

564 },

565 }

566 d = upstream_chat(body)

@@ -601,6 +610,7 @@ def one_read_continuation(schema, template, slots, prompt_ids, seed):

610 "diffusion_max_steps": schema["steps"],

611 "diffusion_read_only": True,

612 **pin_xargs(template, slots, schema["steps"]),

613 + **constrained_xargs(),

614 },

615 }

616 d = upstream_completions(body)

@@ -1496,6 +1506,12 @@ def main():

1506 default=16,

1507 help="request widths round up to a multiple of this",

1508 )

1509 + p.add_argument(

1510 + "--no-constrained",

1511 + dest="constrained",

1512 + action="store_false",

1513 + help="read over the whole vocabulary instead of the labels",

1514 + )

1515 p.add_argument("--host", default="0.0.0.0")

1516 p.add_argument("--port", type=int, default=8011)

1517 p.add_argument(

tests/test_sampling_params.py

+23

@@ -145,6 +145,29 @@ def test_diffusion_rejects_bad_extra_args(extra_args: dict, match: str):

145 _verify_diffusion(SamplingParams(extra_args=extra_args))

146

147

148 +@pytest.mark.parametrize("flag", [True, 1])

149 +def test_diffusion_constrained_needs_logprob_token_ids(flag):

150 + bad = SamplingParams(extra_args={"diffusion_constrained": flag})

151 + with pytest.raises(VLLMValidationError, match="needs logprob_token_ids"):

152 + _verify_diffusion(bad)

153 +

154 + ok = SamplingParams(

155 + logprob_token_ids=[3, 5], extra_args={"diffusion_constrained": flag}

156 + )

157 + _verify_diffusion(ok)

158 +

159 + # An unset or false flag needs no ids.

160 + _verify_diffusion(SamplingParams(extra_args={"diffusion_constrained": False}))

161 + _verify_diffusion(SamplingParams(extra_args={"diffusion_constrained": 0}))

162 +

163 +

164 +@pytest.mark.parametrize("value", ["yes", 2, [1]])

165 +def test_diffusion_constrained_must_be_a_bool(value):

166 + params = SamplingParams(extra_args={"diffusion_constrained": value})

167 + with pytest.raises(VLLMValidationError, match="must be a boolean"):

168 + _verify_diffusion(params)

169 +

170 +

171 def test_diffusion_seed_canvas_must_fill_the_canvas():

172 params = SamplingParams(extra_args={"diffusion_seed_canvas": [0] * 7})

173 with pytest.raises(VLLMValidationError, match="exactly 8 ids, got 7"):

tests/v1/sample/test_diffusion_gemma_reads.py

+49

@@ -8,9 +8,11 @@ import pytest

8 import torch

10 from vllm.model_executor.models.diffusion_gemma import (

11 + _MASKED_LOGIT,

12 DiffusionGemmaRequestStates,

13 _compiled_sample_step,

14 _concat_logprob_stashes,

15 + _mask_rows_to_allowed,

16 )

17 from vllm.platforms import current_platform

18 from vllm.v1.outputs import LogprobsTensors

@@ -317,3 +319,50 @@ def test_read_emits_at_convergence_while_generation_waits_for_commit(width, step

319 sampled, counts = _denoise_once(states, [0], width=width)

320 assert counts.tolist() == [width]

321 assert torch.equal(sampled[0], states.argmax_canvas[0, :width].int())

322 +

323 +

324 +def test_batch_allowed_needs_one_shared_set():

325 + states = _states()

326 + states.constrained[0] = (1, 2, 3)

327 + states.constrained[1] = (4, 5)

328 +

329 + # Mixed sets, or a slot with no set, fall back to per-row masking.

330 + assert states.batch_allowed([0, 1]) is None

331 + assert states.batch_allowed([0, 2]) is None

332 + assert states.batch_allowed([]) is None

333 +

334 + shared = states.batch_allowed([0])

335 + assert shared.tolist() == [1, 2, 3]

336 + assert shared.dtype == torch.int64

337 + assert states.batch_allowed([0, 0]) is shared

338 + assert states.allowed_tensor((1, 2, 3)) is shared

339 +

340 +

341 +def test_mask_rows_to_allowed_masks_only_constrained_rows():

342 + logits = torch.randn(5, VOCAB, device="cuda")

343 + before = logits.clone()

344 + first = torch.tensor([3, 7], device="cuda")

345 + third = torch.tensor([0], device="cuda")

346 +

347 + out = _mask_rows_to_allowed(logits, [0, 2, 3], [2, 1, 2], [first, None, third])

348 +

349 + assert out is not logits

350 + assert torch.equal(logits, before)

351 + # Masked columns hold the finite sentinel, so entropy stays finite.

352 + kept = out > _MASKED_LOGIT

353 + assert torch.isfinite(out).all()

354 + assert kept[0:2].sum(dim=1).tolist() == [2, 2]

355 + assert kept[0:2][:, first].all()

356 + assert torch.equal(out[0:2][:, first], before[0:2][:, first])

357 + assert torch.equal(out[2], before[2])

358 + assert kept[3:5].sum(dim=1).tolist() == [1, 1]

359 + assert torch.equal(out[3:5][:, third], before[3:5][:, third])

360 + # The masked row's softmax is the distribution renormalized over the set.

361 + torch.testing.assert_close(

362 + out[0].softmax(dim=-1)[first], before[0, first].softmax(dim=-1)

363 + )

364 +

365 +

366 +def test_mask_rows_to_allowed_is_a_no_op_without_constrained_rows():

367 + logits = torch.randn(3, VOCAB, device="cuda")

368 + assert _mask_rows_to_allowed(logits, [0, 1], [1, 2], [None, None]) is logits

vllm/model_executor/models/diffusion_gemma.py

+159 -5

@@ -16,7 +16,7 @@ via Gemma4MultimodalEmbedder.

17 from __future__ import annotations

19 -from collections.abc import Iterable, Mapping

19 +from collections.abc import Iterable, Mapping, Sequence

20 from types import SimpleNamespace

21 from typing import Any

@@ -339,6 +339,17 @@ class DiffusionGemmaForConditionalGeneration(

339 )

340

341 def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

342 + states = getattr(self, "diffusion_states", None)

343 + allowed = states.step_allowed if states is not None else None

344 + if allowed is not None:

345 + # One column per allowed id: a [rows, K] GEMM over the K gathered

346 + # rows of the tied embedding.

347 + logits = torch.nn.functional.linear(

348 + hidden_states, self.lm_head.weight[allowed]

349 + ).float()

350 + if self.final_logit_softcapping is not None:

351 + logits = _softcap_logits(logits, self.final_logit_softcapping)

352 + return logits

353 logits = self.logits_processor(self.lm_head, hidden_states)

354 if logits is not None and self.final_logit_softcapping is not None:

355 logits = _softcap_logits(logits, self.final_logit_softcapping)

@@ -395,6 +406,39 @@ def _concat_logprob_stashes(

406 )

407

408

409 +_MASKED_LOGIT = -1e20

410 +

411 +

412 +def _mask_rows_to_allowed(

413 + logits: torch.Tensor,

414 + row_starts: Sequence[int] | np.ndarray,

415 + row_lens: Sequence[int] | np.ndarray,

416 + allowed_per_row: Sequence[torch.Tensor | None],

417 +) -> torch.Tensor:

418 + """Mask every column outside a request's allowed ids on that request's

419 + rows. Request i owns rows [row_starts[i], +row_lens[i]); None leaves its

420 + rows alone. Returns a copy when any row is masked, so the caller's tensor

421 + (possibly the runner's) is never written.

422 +

423 + Softmax over a masked row equals the K-space distribution the shared

424 + fast path computes, so both paths give the same reads. The mask value is

425 + a large finite negative rather than -inf: the entropy is probs times

426 + log-probs, and 0 * -inf is NaN, while 0 * -1e20 is 0. It stays finite

427 + after the greedy temperature clamp (1e-10) scales it by 1e10.

428 + """

429 + out: torch.Tensor | None = None

430 + keep = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device)

431 + for start, n, allowed in zip(row_starts, row_lens, allowed_per_row):

432 + if allowed is None:

433 + continue

434 + if out is None:

435 + out = logits.clone()

436 + keep.zero_()

437 + keep[allowed] = True

438 + out[int(start) : int(start) + int(n)].masked_fill_(~keep, _MASKED_LOGIT)

439 + return logits if out is None else out

440 +

441 +

442 # Tiles specialize the graph on batch size 1, on a tile as wide as the canvas,

443 # on compute_sc and on sizes that coincide with the state buffers. That set

444 # is small but passes Dynamo's default of 8, after which every step would run

@@ -446,6 +490,10 @@ def _compiled_sample_step(

490 tp_size: int,

491 tp_group_name: str,

492 compute_sc: bool = True,

493 + # Constrained reads: ``logits`` has one column per allowed id and

494 + # ``embed_weight`` holds those ids' rows, so every argmax is an index

495 + # into ``allowed``. ``vocab_size`` stays the real vocabulary for renoise.

496 + allowed: torch.Tensor | None = None,

497 ) -> torch.Tensor:

498 """Compiled decode step: temperature → Gumbel sample → probs/confidence →

499 accept/renoise → convergence, all as vectorized PyTorch ops.

@@ -473,6 +521,9 @@ def _compiled_sample_step(

521 argmax_tokens = (

522 scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL)

523 )

524 + if allowed is not None:

525 + new_tokens = allowed[new_tokens]

526 + argmax_tokens = allowed[argmax_tokens]

527

528 # ---- Phase 3: Probs, self-conditioning, confidence ----

529 log_probs = scaled.log_softmax(dim=-1)

@@ -508,6 +559,10 @@ def _compiled_sample_step(

559 step_tensor[decode_slots] = new_step_val

560

561 # Random tokens for renoise / canvas reinit

562 + # Renoise draws from the whole vocabulary even for a constrained read.

563 + # The model was trained on random-vocabulary noise. Noise drawn from the

564 + # allowed set looks like garbled text, and later steps denoise it into

565 + # more garbled text.

566 random_tokens = torch.randint(

567 0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype

568 )

@@ -695,6 +750,12 @@ class DiffusionGemmaRequestStates:

750 # Read-only slots emit on their converging step and skip the commit

751 # forward.

752 self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)

753 + # Constrained reads: slot -> the allowed ids (the request's

754 + # logprob_token_ids). A step whose decode slots all share one set

755 + # runs the unembedding, sampler and self-conditioning over that set.

756 + self.constrained: dict[int, tuple[int, ...]] = {}

757 + self.step_allowed: torch.Tensor | None = None

758 + self._allowed_cache: dict[tuple[int, ...], torch.Tensor] = {}

759 self.read_only_slots: set[int] = set()

760 # Slots capped at one denoise step never consume a soft embed.

761 self.single_step_slots: set[int] = set()

@@ -738,6 +799,7 @@ class DiffusionGemmaRequestStates:

799 self.read_only[slot_idx].fill_(False)

800 self.read_only_slots.discard(slot_idx)

801 self.single_step_slots.discard(slot_idx)

802 + self.constrained.pop(slot_idx, None)

803 self.canvas_width_np[slot_idx] = self.canvas_length

804

805 def remove_request(self, slot_idx: int) -> None:

@@ -749,6 +811,7 @@ class DiffusionGemmaRequestStates:

811 self.seeded_slots.discard(slot_idx)

812 self.read_only_slots.discard(slot_idx)

813 self.single_step_slots.discard(slot_idx)

814 + self.constrained.pop(slot_idx, None)

815

816 def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None:

817 """``ids`` covers the slot's canvas width; positions past it are never

@@ -767,6 +830,30 @@ class DiffusionGemmaRequestStates:

830 slot_idx, async_tensor_h2d(positions, dtype=torch.int64, device=self.device)

831 ] = True

832

833 + def allowed_tensor(self, ids: tuple[int, ...]) -> torch.Tensor:

834 + """``ids`` as an int64 tensor on the device, built once per tuple."""

835 + t = self._allowed_cache.get(ids)

836 + if t is None:

837 + t = torch.tensor(ids, dtype=torch.int64, device=self.device)

838 + if len(self._allowed_cache) > 64:

839 + self._allowed_cache.clear()

840 + self._allowed_cache[ids] = t

841 + return t

842 +

843 + def batch_allowed(self, slots: list[int]) -> torch.Tensor | None:

844 + """The allowed ids shared by every one of ``slots``, or None.

845 +

846 + None means the step runs over the full vocabulary. The sampler then

847 + masks each constrained slot's logit rows to its own set, so a

848 + constrained request reads the same way whoever shares its batch.

849 + """

850 + if not slots or not self.constrained:

851 + return None

852 + first = self.constrained.get(slots[0])

853 + if first is None or any(self.constrained.get(s) != first for s in slots[1:]):

854 + return None

855 + return self.allowed_tensor(first)

856 +

857 def set_read_only(self, slot_idx: int) -> None:

858 self.read_only[slot_idx].fill_(True)

859 self.read_only_slots.add(slot_idx)

@@ -828,6 +915,8 @@ class DiffusionGemmaModelState(ModelState):

915 # the current step, so we add 1 to match the same behavior.

916 stability_threshold=self.gen_config["stability_threshold"] + 1,

917 )

918 + # compute_logits reads the step's shared allowed set from here.

919 + self.model.diffusion_states = self.diffusion_states

920 self._req_id_to_index: dict[str, int] = {}

921

922 # Persistent buffer for per-request causal flags, updated in-place

@@ -979,10 +1068,14 @@ class DiffusionGemmaModelState(ModelState):

1068 inputs_embeds[:num_tokens].copy_(self.model.embed_input_ids(input_ids))

1069

1070 # Apply self-conditioning ONLY for denoising decode requests.

1071 + states.step_allowed = None

1072 if input_batch.num_draft_tokens > 0 and self._req_id_to_index:

1073 slots_np = input_batch.idx_mapping_np[:num_reqs]

1074 num_logits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1])

1075 is_decode_indices_np = np.where(num_logits_np > 0)[0]

1076 + states.step_allowed = states.batch_allowed(

1077 + slots_np[is_decode_indices_np].tolist()

1078 + )

1079 self._apply_self_conditioning(

1080 slots_np[is_decode_indices_np],

1081 is_decode_indices_np,

@@ -1177,6 +1270,16 @@ class DiffusionSampler:

1270 states.set_pins(req_idx, [int(p) for p in pins])

1271 if extra.get("diffusion_read_only"):

1272 states.set_read_only(req_idx)

1273 + if extra.get("diffusion_constrained"):

1274 + ids = list(getattr(sampling_params, "logprob_token_ids", None) or [])

1275 + if not ids:

1276 + raise ValueError(

1277 + "diffusion_constrained needs logprob_token_ids: they are "

1278 + "the allowed set."

1279 + )

1280 + if self.tp_size > 1:

1281 + raise ValueError("diffusion_constrained needs tensor parallel 1.")

1282 + states.constrained[req_idx] = tuple(int(t) for t in ids)

1283

1284 def apply_staged_writes(self) -> None:

1285 self.sampling_states.apply_staged_writes()

@@ -1361,6 +1464,30 @@ class DiffusionSampler:

1464 # since it mutates is_encoder_phase (commit→False, converge→True).

1465 is_committing = states.is_encoder_phase[decode_slots].clone()

1466

1467 + # Constrained step: logits are [rows, K] over the shared allowed set,

1468 + # so the self-conditioning matmul only needs those K embedding rows.

1469 + allowed = states.step_allowed

1470 + embed_weight = self.embed_weight

1471 + vocab_size = self.vocab_size

1472 + sc_vocab_start, sc_vocab_end = self.sc_vocab_start, self.sc_vocab_end

1473 + if allowed is not None:

1474 + assert logits.shape[-1] == allowed.numel()

1475 + embed_weight = self.embed_weight[allowed]

1476 + sc_vocab_start, sc_vocab_end = 0, allowed.numel()

1477 + elif states.constrained and num_decode > 0:

1478 + # The decode slots do not share one set, so this step runs over

1479 + # the full vocabulary. Mask each constrained request's rows to

1480 + # its own set so it reads the same as on the shared path. The

1481 + # helper masks a copy, since the runner owns `logits`: one

1482 + # full-vocab copy per step, as for top_k/top_p above.

1483 + per_row = [

1484 + None if ids is None else states.allowed_tensor(ids)

1485 + for ids in (states.constrained.get(s) for s in decode_slots_np.tolist())

1486 + ]

1487 + logits = _mask_rows_to_allowed(

1488 + logits, row_starts_np, valid_canvas_len_np, per_row

1489 + )

1490 +

1491 slots_np = input_batch.idx_mapping_np[:num_reqs]

1492 max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)

1493 # Requests may ask for specific token ids' logprobs instead of, or as

@@ -1431,7 +1558,7 @@ class DiffusionSampler:

1558 states.is_encoder_phase,

1559 states.confident,

1560 states.self_conditioning_embeds[:, :W],

1434 - self.embed_weight,

1561 + embed_weight,

1562 self.normalizer,

1563 states.accepted_canvas_history[:, :, :W],

1564 states.accepted_canvas_history_len,

@@ -1448,15 +1575,16 @@ class DiffusionSampler:

1575 t_min=self.t_min,

1576 t_max=self.t_max,

1577 confidence_threshold=self.confidence_threshold,

1451 - vocab_size=self.vocab_size,

1578 + vocab_size=vocab_size,

1579 CL=W,

1580 ST=states.stability_threshold,

1581 entropy_bound=self.entropy_bound,

1455 - sc_vocab_start=self.sc_vocab_start,

1456 - sc_vocab_end=self.sc_vocab_end,

1582 + sc_vocab_start=sc_vocab_start,

1583 + sc_vocab_end=sc_vocab_end,

1584 tp_size=self.tp_size,

1585 tp_group_name=self.tp_group_name,

1586 compute_sc=compute_sc,

1587 + allowed=allowed,

1588 )

1589

1590 # Stash newly converged logprobs, including reads that emit now.

@@ -1483,6 +1611,32 @@ class DiffusionSampler:

1611 if raw_flat is None:

1612 raw_flat = tile_logits.float()

1613 src = raw_flat

1614 + if allowed is not None:

1615 + # Column j of the K-space logits is allowed[j]:

1616 + # report every allowed id, renormalized over

1617 + # the set, with the argmax in column 0.

1618 + rows = src[pos : pos + k_i]

1619 + lp = rows.log_softmax(dim=-1)

1620 + am = argmax_tokens[local_idx][:k_i]

1621 + sel_lp = lp.gather(1, am.unsqueeze(1))

1622 + # Same convention as _ranks_kernel: 1-based,

1623 + # ties count, so the rank is the number of

1624 + # columns at or above the selected one. Out

1625 + # of set ids have zero probability, so the

1626 + # rank within the set is the real rank.

1627 + ranks = (lp >= sel_lp).sum(dim=1)

1628 + self._pending_logprobs[slot] = LogprobsTensors(

1629 + logprob_token_ids=torch.cat(

1630 + (

1631 + allowed[am].unsqueeze(1),

1632 + allowed.unsqueeze(0).expand(k_i, -1),

1633 + ),

1634 + dim=1,

1635 + ),

1636 + logprobs=torch.cat((sel_lp, lp), dim=1),

1637 + selected_token_ranks=ranks.to(torch.int64),

1638 + )

1639 + continue

1640 per_req_ids = max_token_ids > 0

1641 self._pending_logprobs[slot] = compute_topk_scores(

1642 src[pos : pos + k_i],

vllm/utils/diffusion.py

+12

@@ -94,6 +94,18 @@ def validate_diffusion_sampling_params(

94 params.max_tokens = min(params.max_tokens or expected_len, expected_len)

95 params.ignore_eos = True

97 + constrained = extra.get("diffusion_constrained")

98 + if constrained is not None:

99 + check(

100 + isinstance(constrained, (bool, int)) and constrained in (0, 1),

101 + "diffusion_constrained must be a boolean (or 0/1).",

102 + )

103 + if constrained:

104 + check(

105 + bool(params.logprob_token_ids),

106 + "diffusion_constrained needs logprob_token_ids: they are the allowed set.",

107 + )

108 +

109 check(

110 bool(async_scheduling)

111 or width is None