From ab3de6edf2b34507e066d998c09f2361a5c1d8ce Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Thu, 24 Sep 2026 14:15:08 -0600 Subject: [PATCH] [Perf] DiffusionGemma: constrained reads over the request's logprob_token_ids (#58216) Signed-off-by: Matt Mastracci Co-authored-by: Claude Fable 5.1 Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Misha Goin --- .../features/structured_diffusion/README.md | 4 +- .../structured_diffusion/structured_server.py | 16 ++ tests/test_sampling_params.py | 23 +++ tests/v1/sample/test_diffusion_gemma_reads.py | 49 ++++++ vllm/model_executor/models/diffusion_gemma.py | 164 +++++++++++++++++- vllm/utils/diffusion.py | 12 ++ 6 files changed, 262 insertions(+), 6 deletions(-) diff --git a/examples/features/structured_diffusion/README.md b/examples/features/structured_diffusion/README.md index 9840f9ea3f..37e58dbcdd 100644 --- a/examples/features/structured_diffusion/README.md +++ b/examples/features/structured_diffusion/README.md @@ -11,8 +11,10 @@ left as noise, one denoise step gives a distribution over each slot. These | `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 | +| `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 | -`structured_server.py` turns a question schema into those fields. It serves +`structured_server.py` turns a question schema into those fields, with +`diffusion_constrained` on for every read (`--no-constrained` turns it off). 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. diff --git a/examples/features/structured_diffusion/structured_server.py b/examples/features/structured_diffusion/structured_server.py index a8f61da94e..5cb98e4117 100644 --- a/examples/features/structured_diffusion/structured_server.py +++ b/examples/features/structured_diffusion/structured_server.py @@ -399,6 +399,14 @@ def canvas_width(template): return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP) +def constrained_xargs(): + """Read over the labels only. The engine runs the unembedding, sampler and + self-conditioning over the request's logprob_token_ids instead of the whole + vocabulary. Same argmax, about a quarter less GPU time per read. Logprobs + come back normalized over the labels.""" + return {"diffusion_constrained": True} if ARGS.constrained else {} + + def pin_xargs(template, slots, steps): """Past one denoise step the template must be held, or accept/renoise rewrites it: pin every canvas position that is not an answer slot.""" @@ -552,6 +560,7 @@ def one_read( "diffusion_max_steps": schema["steps"], "diffusion_read_only": True, **pin_xargs(template, slots, schema["steps"]), + **constrained_xargs(), }, } d = upstream_chat(body) @@ -601,6 +610,7 @@ def one_read_continuation(schema, template, slots, prompt_ids, seed): "diffusion_max_steps": schema["steps"], "diffusion_read_only": True, **pin_xargs(template, slots, schema["steps"]), + **constrained_xargs(), }, } d = upstream_completions(body) @@ -1496,6 +1506,12 @@ def main(): default=16, help="request widths round up to a multiple of this", ) + p.add_argument( + "--no-constrained", + dest="constrained", + action="store_false", + help="read over the whole vocabulary instead of the labels", + ) p.add_argument("--host", default="0.0.0.0") p.add_argument("--port", type=int, default=8011) p.add_argument( diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index 84ce45caa9..98ef59a089 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -145,6 +145,29 @@ def test_diffusion_rejects_bad_extra_args(extra_args: dict, match: str): _verify_diffusion(SamplingParams(extra_args=extra_args)) +@pytest.mark.parametrize("flag", [True, 1]) +def test_diffusion_constrained_needs_logprob_token_ids(flag): + bad = SamplingParams(extra_args={"diffusion_constrained": flag}) + with pytest.raises(VLLMValidationError, match="needs logprob_token_ids"): + _verify_diffusion(bad) + + ok = SamplingParams( + logprob_token_ids=[3, 5], extra_args={"diffusion_constrained": flag} + ) + _verify_diffusion(ok) + + # An unset or false flag needs no ids. + _verify_diffusion(SamplingParams(extra_args={"diffusion_constrained": False})) + _verify_diffusion(SamplingParams(extra_args={"diffusion_constrained": 0})) + + +@pytest.mark.parametrize("value", ["yes", 2, [1]]) +def test_diffusion_constrained_must_be_a_bool(value): + params = SamplingParams(extra_args={"diffusion_constrained": value}) + with pytest.raises(VLLMValidationError, match="must be a boolean"): + _verify_diffusion(params) + + 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"): diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py index 369af010f4..09b3a83223 100644 --- a/tests/v1/sample/test_diffusion_gemma_reads.py +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -8,9 +8,11 @@ import pytest import torch from vllm.model_executor.models.diffusion_gemma import ( + _MASKED_LOGIT, DiffusionGemmaRequestStates, _compiled_sample_step, _concat_logprob_stashes, + _mask_rows_to_allowed, ) from vllm.platforms import current_platform from vllm.v1.outputs import LogprobsTensors @@ -317,3 +319,50 @@ def test_read_emits_at_convergence_while_generation_waits_for_commit(width, step sampled, counts = _denoise_once(states, [0], width=width) assert counts.tolist() == [width] assert torch.equal(sampled[0], states.argmax_canvas[0, :width].int()) + + +def test_batch_allowed_needs_one_shared_set(): + states = _states() + states.constrained[0] = (1, 2, 3) + states.constrained[1] = (4, 5) + + # Mixed sets, or a slot with no set, fall back to per-row masking. + assert states.batch_allowed([0, 1]) is None + assert states.batch_allowed([0, 2]) is None + assert states.batch_allowed([]) is None + + shared = states.batch_allowed([0]) + assert shared.tolist() == [1, 2, 3] + assert shared.dtype == torch.int64 + assert states.batch_allowed([0, 0]) is shared + assert states.allowed_tensor((1, 2, 3)) is shared + + +def test_mask_rows_to_allowed_masks_only_constrained_rows(): + logits = torch.randn(5, VOCAB, device="cuda") + before = logits.clone() + first = torch.tensor([3, 7], device="cuda") + third = torch.tensor([0], device="cuda") + + out = _mask_rows_to_allowed(logits, [0, 2, 3], [2, 1, 2], [first, None, third]) + + assert out is not logits + assert torch.equal(logits, before) + # Masked columns hold the finite sentinel, so entropy stays finite. + kept = out > _MASKED_LOGIT + assert torch.isfinite(out).all() + assert kept[0:2].sum(dim=1).tolist() == [2, 2] + assert kept[0:2][:, first].all() + assert torch.equal(out[0:2][:, first], before[0:2][:, first]) + assert torch.equal(out[2], before[2]) + assert kept[3:5].sum(dim=1).tolist() == [1, 1] + assert torch.equal(out[3:5][:, third], before[3:5][:, third]) + # The masked row's softmax is the distribution renormalized over the set. + torch.testing.assert_close( + out[0].softmax(dim=-1)[first], before[0, first].softmax(dim=-1) + ) + + +def test_mask_rows_to_allowed_is_a_no_op_without_constrained_rows(): + logits = torch.randn(3, VOCAB, device="cuda") + assert _mask_rows_to_allowed(logits, [0, 1], [1, 2], [None, None]) is logits diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 84d52fb5ba..59fe7e22ba 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -16,7 +16,7 @@ via Gemma4MultimodalEmbedder. from __future__ import annotations -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from types import SimpleNamespace from typing import Any @@ -339,6 +339,17 @@ class DiffusionGemmaForConditionalGeneration( ) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + states = getattr(self, "diffusion_states", None) + allowed = states.step_allowed if states is not None else None + if allowed is not None: + # One column per allowed id: a [rows, K] GEMM over the K gathered + # rows of the tied embedding. + logits = torch.nn.functional.linear( + hidden_states, self.lm_head.weight[allowed] + ).float() + if self.final_logit_softcapping is not None: + logits = _softcap_logits(logits, self.final_logit_softcapping) + return logits logits = self.logits_processor(self.lm_head, hidden_states) if logits is not None and self.final_logit_softcapping is not None: logits = _softcap_logits(logits, self.final_logit_softcapping) @@ -395,6 +406,39 @@ def _concat_logprob_stashes( ) +_MASKED_LOGIT = -1e20 + + +def _mask_rows_to_allowed( + logits: torch.Tensor, + row_starts: Sequence[int] | np.ndarray, + row_lens: Sequence[int] | np.ndarray, + allowed_per_row: Sequence[torch.Tensor | None], +) -> torch.Tensor: + """Mask every column outside a request's allowed ids on that request's + rows. Request i owns rows [row_starts[i], +row_lens[i]); None leaves its + rows alone. Returns a copy when any row is masked, so the caller's tensor + (possibly the runner's) is never written. + + Softmax over a masked row equals the K-space distribution the shared + fast path computes, so both paths give the same reads. The mask value is + a large finite negative rather than -inf: the entropy is probs times + log-probs, and 0 * -inf is NaN, while 0 * -1e20 is 0. It stays finite + after the greedy temperature clamp (1e-10) scales it by 1e10. + """ + out: torch.Tensor | None = None + keep = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) + for start, n, allowed in zip(row_starts, row_lens, allowed_per_row): + if allowed is None: + continue + if out is None: + out = logits.clone() + keep.zero_() + keep[allowed] = True + out[int(start) : int(start) + int(n)].masked_fill_(~keep, _MASKED_LOGIT) + return logits if out is None else out + + # 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 @@ -446,6 +490,10 @@ def _compiled_sample_step( tp_size: int, tp_group_name: str, compute_sc: bool = True, + # Constrained reads: ``logits`` has one column per allowed id and + # ``embed_weight`` holds those ids' rows, so every argmax is an index + # into ``allowed``. ``vocab_size`` stays the real vocabulary for renoise. + allowed: torch.Tensor | None = None, ) -> torch.Tensor: """Compiled decode step: temperature → Gumbel sample → probs/confidence → accept/renoise → convergence, all as vectorized PyTorch ops. @@ -473,6 +521,9 @@ def _compiled_sample_step( argmax_tokens = ( scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL) ) + if allowed is not None: + new_tokens = allowed[new_tokens] + argmax_tokens = allowed[argmax_tokens] # ---- Phase 3: Probs, self-conditioning, confidence ---- log_probs = scaled.log_softmax(dim=-1) @@ -508,6 +559,10 @@ def _compiled_sample_step( step_tensor[decode_slots] = new_step_val # Random tokens for renoise / canvas reinit + # Renoise draws from the whole vocabulary even for a constrained read. + # The model was trained on random-vocabulary noise. Noise drawn from the + # allowed set looks like garbled text, and later steps denoise it into + # more garbled text. random_tokens = torch.randint( 0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype ) @@ -695,6 +750,12 @@ class DiffusionGemmaRequestStates: # 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) + # Constrained reads: slot -> the allowed ids (the request's + # logprob_token_ids). A step whose decode slots all share one set + # runs the unembedding, sampler and self-conditioning over that set. + self.constrained: dict[int, tuple[int, ...]] = {} + self.step_allowed: torch.Tensor | None = None + self._allowed_cache: dict[tuple[int, ...], torch.Tensor] = {} 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() @@ -738,6 +799,7 @@ class DiffusionGemmaRequestStates: self.read_only[slot_idx].fill_(False) self.read_only_slots.discard(slot_idx) self.single_step_slots.discard(slot_idx) + self.constrained.pop(slot_idx, None) self.canvas_width_np[slot_idx] = self.canvas_length def remove_request(self, slot_idx: int) -> None: @@ -749,6 +811,7 @@ class DiffusionGemmaRequestStates: self.seeded_slots.discard(slot_idx) self.read_only_slots.discard(slot_idx) self.single_step_slots.discard(slot_idx) + self.constrained.pop(slot_idx, None) def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None: """``ids`` covers the slot's canvas width; positions past it are never @@ -767,6 +830,30 @@ class DiffusionGemmaRequestStates: slot_idx, async_tensor_h2d(positions, dtype=torch.int64, device=self.device) ] = True + def allowed_tensor(self, ids: tuple[int, ...]) -> torch.Tensor: + """``ids`` as an int64 tensor on the device, built once per tuple.""" + t = self._allowed_cache.get(ids) + if t is None: + t = torch.tensor(ids, dtype=torch.int64, device=self.device) + if len(self._allowed_cache) > 64: + self._allowed_cache.clear() + self._allowed_cache[ids] = t + return t + + def batch_allowed(self, slots: list[int]) -> torch.Tensor | None: + """The allowed ids shared by every one of ``slots``, or None. + + None means the step runs over the full vocabulary. The sampler then + masks each constrained slot's logit rows to its own set, so a + constrained request reads the same way whoever shares its batch. + """ + if not slots or not self.constrained: + return None + first = self.constrained.get(slots[0]) + if first is None or any(self.constrained.get(s) != first for s in slots[1:]): + return None + return self.allowed_tensor(first) + def set_read_only(self, slot_idx: int) -> None: self.read_only[slot_idx].fill_(True) self.read_only_slots.add(slot_idx) @@ -828,6 +915,8 @@ class DiffusionGemmaModelState(ModelState): # the current step, so we add 1 to match the same behavior. stability_threshold=self.gen_config["stability_threshold"] + 1, ) + # compute_logits reads the step's shared allowed set from here. + self.model.diffusion_states = self.diffusion_states self._req_id_to_index: dict[str, int] = {} # Persistent buffer for per-request causal flags, updated in-place @@ -979,10 +1068,14 @@ class DiffusionGemmaModelState(ModelState): inputs_embeds[:num_tokens].copy_(self.model.embed_input_ids(input_ids)) # Apply self-conditioning ONLY for denoising decode requests. + states.step_allowed = None if input_batch.num_draft_tokens > 0 and self._req_id_to_index: slots_np = input_batch.idx_mapping_np[:num_reqs] num_logits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) is_decode_indices_np = np.where(num_logits_np > 0)[0] + states.step_allowed = states.batch_allowed( + slots_np[is_decode_indices_np].tolist() + ) self._apply_self_conditioning( slots_np[is_decode_indices_np], is_decode_indices_np, @@ -1177,6 +1270,16 @@ class DiffusionSampler: states.set_pins(req_idx, [int(p) for p in pins]) if extra.get("diffusion_read_only"): states.set_read_only(req_idx) + if extra.get("diffusion_constrained"): + ids = list(getattr(sampling_params, "logprob_token_ids", None) or []) + if not ids: + raise ValueError( + "diffusion_constrained needs logprob_token_ids: they are " + "the allowed set." + ) + if self.tp_size > 1: + raise ValueError("diffusion_constrained needs tensor parallel 1.") + states.constrained[req_idx] = tuple(int(t) for t in ids) def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() @@ -1361,6 +1464,30 @@ class DiffusionSampler: # since it mutates is_encoder_phase (commit→False, converge→True). is_committing = states.is_encoder_phase[decode_slots].clone() + # Constrained step: logits are [rows, K] over the shared allowed set, + # so the self-conditioning matmul only needs those K embedding rows. + allowed = states.step_allowed + embed_weight = self.embed_weight + vocab_size = self.vocab_size + sc_vocab_start, sc_vocab_end = self.sc_vocab_start, self.sc_vocab_end + if allowed is not None: + assert logits.shape[-1] == allowed.numel() + embed_weight = self.embed_weight[allowed] + sc_vocab_start, sc_vocab_end = 0, allowed.numel() + elif states.constrained and num_decode > 0: + # The decode slots do not share one set, so this step runs over + # the full vocabulary. Mask each constrained request's rows to + # its own set so it reads the same as on the shared path. The + # helper masks a copy, since the runner owns `logits`: one + # full-vocab copy per step, as for top_k/top_p above. + per_row = [ + None if ids is None else states.allowed_tensor(ids) + for ids in (states.constrained.get(s) for s in decode_slots_np.tolist()) + ] + logits = _mask_rows_to_allowed( + logits, row_starts_np, valid_canvas_len_np, per_row + ) + slots_np = input_batch.idx_mapping_np[:num_reqs] max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) # Requests may ask for specific token ids' logprobs instead of, or as @@ -1431,7 +1558,7 @@ class DiffusionSampler: states.is_encoder_phase, states.confident, states.self_conditioning_embeds[:, :W], - self.embed_weight, + embed_weight, self.normalizer, states.accepted_canvas_history[:, :, :W], states.accepted_canvas_history_len, @@ -1448,15 +1575,16 @@ class DiffusionSampler: t_min=self.t_min, t_max=self.t_max, confidence_threshold=self.confidence_threshold, - vocab_size=self.vocab_size, + vocab_size=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, + sc_vocab_start=sc_vocab_start, + sc_vocab_end=sc_vocab_end, tp_size=self.tp_size, tp_group_name=self.tp_group_name, compute_sc=compute_sc, + allowed=allowed, ) # Stash newly converged logprobs, including reads that emit now. @@ -1483,6 +1611,32 @@ class DiffusionSampler: if raw_flat is None: raw_flat = tile_logits.float() src = raw_flat + if allowed is not None: + # Column j of the K-space logits is allowed[j]: + # report every allowed id, renormalized over + # the set, with the argmax in column 0. + rows = src[pos : pos + k_i] + lp = rows.log_softmax(dim=-1) + am = argmax_tokens[local_idx][:k_i] + sel_lp = lp.gather(1, am.unsqueeze(1)) + # Same convention as _ranks_kernel: 1-based, + # ties count, so the rank is the number of + # columns at or above the selected one. Out + # of set ids have zero probability, so the + # rank within the set is the real rank. + ranks = (lp >= sel_lp).sum(dim=1) + self._pending_logprobs[slot] = LogprobsTensors( + logprob_token_ids=torch.cat( + ( + allowed[am].unsqueeze(1), + allowed.unsqueeze(0).expand(k_i, -1), + ), + dim=1, + ), + logprobs=torch.cat((sel_lp, lp), dim=1), + selected_token_ranks=ranks.to(torch.int64), + ) + continue per_req_ids = max_token_ids > 0 self._pending_logprobs[slot] = compute_topk_scores( src[pos : pos + k_i], diff --git a/vllm/utils/diffusion.py b/vllm/utils/diffusion.py index 278c17d838..4980f940d8 100644 --- a/vllm/utils/diffusion.py +++ b/vllm/utils/diffusion.py @@ -94,6 +94,18 @@ def validate_diffusion_sampling_params( params.max_tokens = min(params.max_tokens or expected_len, expected_len) params.ignore_eos = True + constrained = extra.get("diffusion_constrained") + if constrained is not None: + check( + isinstance(constrained, (bool, int)) and constrained in (0, 1), + "diffusion_constrained must be a boolean (or 0/1).", + ) + if constrained: + check( + bool(params.logprob_token_ids), + "diffusion_constrained needs logprob_token_ids: they are the allowed set.", + ) + check( bool(async_scheduling) or width is None