| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Per-request DiffusionGemma state behind structured reads: seed canvases, |
| 4 | pinned positions, read-only slots and the per-slot step cap.""" |
| 5 | |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.model_executor.models.diffusion_gemma import ( |
| 11 | _MASKED_LOGIT, |
| 12 | DiffusionGemmaRequestStates, |
| 13 | _compiled_sample_step, |
| 14 | _concat_logprob_stashes, |
| 15 | _denoise_temperature, |
| 16 | _mask_rows_to_allowed, |
| 17 | sample_row_stats_reference, |
| 18 | ) |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.v1.outputs import LogprobsTensors |
| 21 | |
| 22 | pytestmark = pytest.mark.skipif( |
| 23 | not current_platform.is_cuda(), reason="the sampler state lives on the GPU" |
| 24 | ) |
| 25 | |
| 26 | CL = 8 |
| 27 | VOCAB = 64 |
| 28 | MAX_REQS = 4 |
| 29 | MAX_STEPS = 48 |
| 30 | |
| 31 | |
| 32 | def _states() -> DiffusionGemmaRequestStates: |
| 33 | return DiffusionGemmaRequestStates( |
| 34 | max_num_reqs=MAX_REQS, |
| 35 | canvas_length=CL, |
| 36 | vocab_size=VOCAB, |
| 37 | max_denoising_steps=MAX_STEPS, |
| 38 | device=torch.device("cuda"), |
| 39 | hidden_size=4, |
| 40 | stability_threshold=2, |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | def _slots(*idx: int) -> tuple[np.ndarray, torch.Tensor]: |
| 45 | slots = np.array(idx, dtype=np.int64) |
| 46 | return slots, torch.tensor(slots, device="cuda") |
| 47 | |
| 48 | |
| 49 | def test_seed_canvas_replaces_only_seeded_slots(): |
| 50 | states = _states() |
| 51 | for slot in range(3): |
| 52 | states.add_request(slot) |
| 53 | seed = list(range(CL)) |
| 54 | states.set_seed_canvas(1, seed) |
| 55 | slots, slots_gpu = _slots(0, 1, 2) |
| 56 | states.init_canvas(slots_gpu) |
| 57 | before = states.canvas[slots_gpu].clone() |
| 58 | |
| 59 | states.apply_seed_canvases(slots, slots_gpu) |
| 60 | |
| 61 | after = states.canvas[slots_gpu] |
| 62 | assert after[1].tolist() == seed |
| 63 | assert torch.equal(after[0], before[0]) |
| 64 | assert torch.equal(after[2], before[2]) |
| 65 | |
| 66 | |
| 67 | def test_apply_seed_canvases_leaves_unseeded_batches_alone(): |
| 68 | states = _states() |
| 69 | states.add_request(0) |
| 70 | slots, slots_gpu = _slots(0) |
| 71 | states.init_canvas(slots_gpu) |
| 72 | before = states.canvas[0].clone() |
| 73 | |
| 74 | states.apply_seed_canvases(slots, slots_gpu) |
| 75 | |
| 76 | assert torch.equal(states.canvas[0], before) |
| 77 | |
| 78 | |
| 79 | def test_add_request_clears_seed_and_read_only(): |
| 80 | states = _states() |
| 81 | states.add_request(0) |
| 82 | states.set_seed_canvas(0, [1] * CL) |
| 83 | states.set_read_only(0) |
| 84 | assert states.seeded_slots == {0} |
| 85 | assert states.read_only_slots == {0} |
| 86 | |
| 87 | states.add_request(0) |
| 88 | |
| 89 | assert not states.seeded_slots |
| 90 | assert not states.read_only_slots |
| 91 | assert not bool(states.has_seed[0]) |
| 92 | assert not bool(states.read_only[0]) |
| 93 | |
| 94 | |
| 95 | def test_canvas_width_resets_with_the_slot(): |
| 96 | states = _states() |
| 97 | states.add_request(0) |
| 98 | states.canvas_width_np[0] = 4 |
| 99 | states.set_seed_canvas(0, [7, 7, 7, 7]) |
| 100 | assert states.seed_canvas[0, :4].tolist() == [7, 7, 7, 7] |
| 101 | |
| 102 | states.add_request(0) |
| 103 | |
| 104 | assert states.canvas_width_np[0] == CL |
| 105 | |
| 106 | |
| 107 | def test_remove_request_forgets_the_slot(): |
| 108 | states = _states() |
| 109 | states.add_request(0) |
| 110 | states.set_seed_canvas(0, [1] * CL) |
| 111 | states.set_read_only(0) |
| 112 | |
| 113 | states.remove_request(0) |
| 114 | |
| 115 | assert not states.seeded_slots |
| 116 | assert not states.read_only_slots |
| 117 | |
| 118 | |
| 119 | def _denoise_once( |
| 120 | states: DiffusionGemmaRequestStates, |
| 121 | slots: list[int], |
| 122 | compute_sc: bool = True, |
| 123 | width: int = CL, |
| 124 | embed_weight: torch.Tensor | None = None, |
| 125 | embed_dtype: torch.dtype = torch.float32, |
| 126 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 127 | """One compiled denoise step over ``slots`` with flat logits, so nothing |
| 128 | converges by stability or confidence and only the step cap can end it. |
| 129 | ``width`` below CL runs the step on [:, :width] views, as the sampler |
| 130 | does for a narrow tile.""" |
| 131 | n = len(slots) |
| 132 | device = states.device |
| 133 | decode_slots = torch.tensor(slots, dtype=torch.int64, device=device) |
| 134 | decode_idx = torch.arange(n, dtype=torch.int64, device=device) |
| 135 | sampled = torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width] |
| 136 | num_sampled = torch.zeros(n, dtype=torch.int32, device=device) |
| 137 | temp = _denoise_temperature(states.step, decode_slots, float(MAX_STEPS), 0.5, 1.0) |
| 138 | argmax, sample, entropy, probs = sample_row_stats_reference( |
| 139 | torch.zeros(n * width, VOCAB, device=device), |
| 140 | temp, |
| 141 | width, |
| 142 | embed_dtype if compute_sc else None, |
| 143 | ) |
| 144 | _compiled_sample_step( |
| 145 | sample.view(n, width), |
| 146 | argmax.view(n, width), |
| 147 | entropy.view(n, width), |
| 148 | probs.view(n, width, -1) if probs is not None else None, |
| 149 | decode_slots, |
| 150 | decode_idx, |
| 151 | decode_slots, |
| 152 | torch.full((n,), width, dtype=torch.int64, device=device), |
| 153 | states.canvas[:, :width], |
| 154 | states.argmax_canvas[:, :width], |
| 155 | states.step, |
| 156 | states.is_encoder_phase, |
| 157 | states.confident, |
| 158 | states.self_conditioning_embeds[:, :width], |
| 159 | ( |
| 160 | torch.zeros(VOCAB, 4, dtype=embed_dtype, device=device) |
| 161 | if embed_weight is None |
| 162 | else embed_weight |
| 163 | ), |
| 164 | torch.tensor(1.0, dtype=embed_dtype, device=device), |
| 165 | states.accepted_canvas_history[:, :, :width], |
| 166 | states.accepted_canvas_history_len, |
| 167 | states.max_steps, |
| 168 | states.pin_mask[:, :width], |
| 169 | states.seed_canvas[:, :width], |
| 170 | states.read_only, |
| 171 | sampled, |
| 172 | num_sampled, |
| 173 | torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device), |
| 174 | confidence_threshold=0.1, |
| 175 | vocab_size=VOCAB, |
| 176 | CL=width, |
| 177 | ST=states.stability_threshold, |
| 178 | entropy_bound=0.1, |
| 179 | sc_vocab_start=0, |
| 180 | sc_vocab_end=VOCAB, |
| 181 | tp_size=1, |
| 182 | tp_group_name="", |
| 183 | compute_sc=compute_sc, |
| 184 | ) |
| 185 | return sampled, num_sampled |
| 186 | |
| 187 | |
| 188 | def test_pinned_positions_hold_their_seed_through_a_step(): |
| 189 | states = _states() |
| 190 | states.add_request(0) |
| 191 | states.is_encoder_phase[0] = False |
| 192 | seed = list(range(10, 10 + CL)) |
| 193 | states.set_seed_canvas(0, seed) |
| 194 | states.canvas[0] = torch.tensor(seed, device="cuda") |
| 195 | states.set_pins(0, [0, 1, 2, 3]) |
| 196 | |
| 197 | # Flat logits accept nothing, so every free position is renoised. |
| 198 | _denoise_once(states, [0], embed_weight=torch.ones(VOCAB, 4, device="cuda")) |
| 199 | |
| 200 | assert states.canvas[0, :4].tolist() == seed[:4] |
| 201 | # The soft embed is zero at pinned positions and non-zero elsewhere. |
| 202 | sc = states.self_conditioning_embeds[0] |
| 203 | assert not sc[:4].any() |
| 204 | assert sc[4:].any() |
| 205 | |
| 206 | |
| 207 | def test_add_request_clears_pins(): |
| 208 | states = _states() |
| 209 | states.add_request(0) |
| 210 | states.set_seed_canvas(0, [1] * CL) |
| 211 | states.set_pins(0, [2, 3]) |
| 212 | assert states.pin_mask[0].tolist() == [False, False, True, True] + [False] * 4 |
| 213 | |
| 214 | states.add_request(0) |
| 215 | |
| 216 | assert not states.pin_mask[0].any() |
| 217 | |
| 218 | |
| 219 | def test_single_step_tile_skips_self_conditioning(): |
| 220 | states = _states() |
| 221 | states.add_request(0) |
| 222 | states.is_encoder_phase[0] = False |
| 223 | states.self_conditioning_embeds[0] = 1.0 |
| 224 | |
| 225 | _denoise_once(states, [0], compute_sc=False) |
| 226 | |
| 227 | assert not states.self_conditioning_embeds[0].any() |
| 228 | |
| 229 | |
| 230 | @pytest.mark.parametrize("stance", ["default", "force_eager"]) |
| 231 | def test_self_conditioning_stores_a_bf16_model_in_the_fp32_buffer(stance): |
| 232 | # The model's embeddings are bf16 while the buffer is fp32. Compiled code |
| 233 | # casts on the store, but eager does not, and the step runs eager once |
| 234 | # torch.compile hits its recompile limit. |
| 235 | states = _states() |
| 236 | states.add_request(0) |
| 237 | states.is_encoder_phase[0] = False |
| 238 | |
| 239 | with torch.compiler.set_stance(stance): |
| 240 | _denoise_once( |
| 241 | states, |
| 242 | [0], |
| 243 | embed_weight=torch.ones(VOCAB, 4, dtype=torch.bfloat16, device="cuda"), |
| 244 | embed_dtype=torch.bfloat16, |
| 245 | ) |
| 246 | |
| 247 | assert states.self_conditioning_embeds.dtype == torch.float32 |
| 248 | # uniform probs @ an all-ones embedding |
| 249 | assert torch.allclose( |
| 250 | states.self_conditioning_embeds[0], torch.ones(CL, 4, device="cuda") |
| 251 | ) |
| 252 | |
| 253 | |
| 254 | def test_narrow_tile_leaves_the_rest_of_the_canvas_alone(): |
| 255 | states = _states() |
| 256 | states.add_request(0) |
| 257 | states.is_encoder_phase[0] = False |
| 258 | states.canvas[0] = 5 |
| 259 | states.argmax_canvas[0] = 5 |
| 260 | |
| 261 | _denoise_once(states, [0], width=4) |
| 262 | |
| 263 | # Columns past the width are untouched; the step counted. |
| 264 | assert states.canvas[0, 4:].tolist() == [5] * (CL - 4) |
| 265 | assert states.argmax_canvas[0, 4:].tolist() == [5] * (CL - 4) |
| 266 | assert states.step[0] == 1 |
| 267 | |
| 268 | |
| 269 | def test_step_cap_is_per_slot(): |
| 270 | states = _states() |
| 271 | for slot in (0, 1): |
| 272 | states.add_request(slot) |
| 273 | states.is_encoder_phase[slot] = False |
| 274 | states.max_steps[0] = 1 |
| 275 | |
| 276 | _denoise_once(states, [0, 1]) |
| 277 | |
| 278 | # Slot 0 hit its cap and moves to commit. Slot 1 keeps denoising. |
| 279 | assert states.is_encoder_phase[:2].tolist() == [True, False] |
| 280 | assert states.step[:2].tolist() == [1, 1] |
| 281 | |
| 282 | |
| 283 | def _stash(rows: int, width: int, first_id: int) -> LogprobsTensors: |
| 284 | ids = torch.arange(first_id, first_id + rows * width).reshape(rows, width) |
| 285 | return LogprobsTensors( |
| 286 | logprob_token_ids=ids, |
| 287 | logprobs=-ids.float(), |
| 288 | selected_token_ranks=torch.zeros(rows, dtype=torch.int64), |
| 289 | ) |
| 290 | |
| 291 | |
| 292 | def test_stashes_of_different_widths_join(): |
| 293 | # A read that asked for 10 label ids and a generation that asked for 2 |
| 294 | # logprobs converged in different steps, so their stashes differ in width. |
| 295 | wide, narrow = _stash(2, 11, 100), _stash(3, 3, 0) |
| 296 | |
| 297 | out = _concat_logprob_stashes([narrow, wide], [0, 3]) |
| 298 | |
| 299 | assert out.logprob_token_ids.shape == (5, 11) |
| 300 | assert out.logprobs.shape == (5, 11) |
| 301 | assert out.cu_num_generated_tokens == [0, 3] |
| 302 | # The narrow rows keep their values and pad with id 0 at -inf. |
| 303 | assert torch.equal(out.logprob_token_ids[:3, :3], narrow.logprob_token_ids) |
| 304 | assert not out.logprob_token_ids[:3, 3:].any() |
| 305 | assert torch.isneginf(out.logprobs[:3, 3:]).all() |
| 306 | assert torch.equal(out.logprobs[3:], wide.logprobs) |
| 307 | |
| 308 | |
| 309 | @pytest.mark.parametrize("width", [4, CL]) |
| 310 | @pytest.mark.parametrize("steps", [1, 2]) |
| 311 | def test_read_emits_at_convergence_while_generation_waits_for_commit(width, steps): |
| 312 | states = _states() |
| 313 | for slot in (2, 0): |
| 314 | states.add_request(slot) |
| 315 | states.is_encoder_phase[slot] = False |
| 316 | states.max_steps[slot] = steps |
| 317 | states.set_read_only(2) |
| 318 | |
| 319 | for step in range(steps): |
| 320 | sampled, counts = _denoise_once(states, [2, 0], width=width) |
| 321 | assert counts.tolist() == ([width, 0] if step == steps - 1 else [0, 0]) |
| 322 | |
| 323 | assert torch.equal(sampled[0], states.argmax_canvas[2, :width].int()) |
| 324 | assert not states.is_encoder_phase[2] |
| 325 | assert states.is_encoder_phase[0] |
| 326 | assert not states.self_conditioning_embeds[2].any() |
| 327 | |
| 328 | sampled, counts = _denoise_once(states, [0], width=width) |
| 329 | assert counts.tolist() == [width] |
| 330 | assert torch.equal(sampled[0], states.argmax_canvas[0, :width].int()) |
| 331 | |
| 332 | |
| 333 | def test_batch_allowed_needs_one_shared_set(): |
| 334 | states = _states() |
| 335 | states.constrained[0] = (1, 2, 3) |
| 336 | states.constrained[1] = (4, 5) |
| 337 | |
| 338 | # Mixed sets, or a slot with no set, fall back to per-row masking. |
| 339 | assert states.batch_allowed([0, 1]) is None |
| 340 | assert states.batch_allowed([0, 2]) is None |
| 341 | assert states.batch_allowed([]) is None |
| 342 | |
| 343 | shared = states.batch_allowed([0]) |
| 344 | assert shared is not None |
| 345 | assert shared.tolist() == [1, 2, 3] |
| 346 | assert shared.dtype == torch.int64 |
| 347 | assert states.batch_allowed([0, 0]) is shared |
| 348 | assert states.allowed_tensor((1, 2, 3)) is shared |
| 349 | |
| 350 | |
| 351 | def test_mask_rows_to_allowed_masks_only_constrained_rows(): |
| 352 | logits = torch.randn(5, VOCAB, device="cuda") |
| 353 | before = logits.clone() |
| 354 | first = torch.tensor([3, 7], device="cuda") |
| 355 | third = torch.tensor([0], device="cuda") |
| 356 | |
| 357 | out = _mask_rows_to_allowed(logits, [0, 2, 3], [2, 1, 2], [first, None, third]) |
| 358 | |
| 359 | assert out is not logits |
| 360 | assert torch.equal(logits, before) |
| 361 | # Masked columns hold the finite sentinel, so entropy stays finite. |
| 362 | kept = out > _MASKED_LOGIT |
| 363 | assert torch.isfinite(out).all() |
| 364 | assert kept[0:2].sum(dim=1).tolist() == [2, 2] |
| 365 | assert kept[0:2][:, first].all() |
| 366 | assert torch.equal(out[0:2][:, first], before[0:2][:, first]) |
| 367 | assert torch.equal(out[2], before[2]) |
| 368 | assert kept[3:5].sum(dim=1).tolist() == [1, 1] |
| 369 | assert torch.equal(out[3:5][:, third], before[3:5][:, third]) |
| 370 | # The masked row's softmax is the distribution renormalized over the set. |
| 371 | torch.testing.assert_close( |
| 372 | out[0].softmax(dim=-1)[first], before[0, first].softmax(dim=-1) |
| 373 | ) |
| 374 | |
| 375 | |
| 376 | def test_mask_rows_to_allowed_is_a_no_op_without_constrained_rows(): |
| 377 | logits = torch.randn(3, VOCAB, device="cuda") |
| 378 | assert _mask_rows_to_allowed(logits, [0, 1], [1, 2], [None, None]) is logits |