[Bugfix] DiffusionGemma: hand out stashed logprobs only on the committing step (#57414)
Matt Mastracci committed
Sep 17, 2026 at 18:24 UTC
2c88fb131c7ae0be01907cd8c276911db5e7aad4
1 file changed
+4
-2
vllm/model_executor/models/diffusion_gemma.py
+4
-2
@@ -1312,7 +1312,6 @@ class DiffusionSampler:
1312
is_committing = states.is_encoder_phase[decode_slots].clone()
1313
1314
slots_np = input_batch.idx_mapping_np[:num_reqs]
1315
- is_decode_np = per_req_nlogits_np > 0
1315
max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)
1316
1317
# Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in
@@ -1396,13 +1395,16 @@ class DiffusionSampler:
1395
# stashed logprobs and attach to SamplerOutput.
1396
logprobs_tensors = None
1397
if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs:
1398
+ committing_slots = set(
1399
+ decode_slots_np[is_committing.cpu().numpy()].tolist()
1400
+ )
1401
parts_ids, parts_lp, parts_ranks = [], [], []
1402
cu_gen: list[int] = []
1403
flat_offset = 0
1404
for i in range(num_reqs):
1405
cu_gen.append(flat_offset)
1406
slot = int(slots_np[i])
1405
- if is_decode_np[i] and slot in self._pending_logprobs:
1407
+ if slot in committing_slots and slot in self._pending_logprobs:
1408
lp = self._pending_logprobs.pop(slot)
1409
parts_ids.append(lp.logprob_token_ids)
1410
parts_lp.append(lp.logprobs)