[Model] DiffusionGemma: honor logprob_token_ids on the converging step (#57417)
Signed-off-by: Matt Mastracci <matthew@mastracci.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Matt Mastracci committed
Sep 18, 2026 at 14:11 UTC
63d9ad0a3a435cdf3a44495028b10f390a38f960
1 file changed
+27
-3
vllm/model_executor/models/diffusion_gemma.py
+27
-3
@@ -967,6 +967,7 @@ class DiffusionSampler:
967
tp_group_name: str = "",
968
):
969
self.sampling_states = sampler.sampling_states
970
+ self.logprob_token_ids_state = sampler.logprob_token_ids_state
971
self.req_states = sampler.req_states
972
self.logits_mode = sampler.logprobs_mode in ("raw_logits", "processed_logits")
973
# Self-conditioning soft embed = probs @ embed_weight * normalizer,
@@ -1023,9 +1024,11 @@ class DiffusionSampler:
1024
# that was aborted between its converging denoise and commit steps.
1025
self._pending_logprobs.pop(req_idx, None)
1026
self.sampling_states.add_request(req_idx, sampling_params)
1027
+ self.logprob_token_ids_state.add_request(req_idx, sampling_params)
1028
1029
def apply_staged_writes(self) -> None:
1030
self.sampling_states.apply_staged_writes()
1031
+ self.logprob_token_ids_state.apply_staged_writes()
1032
1033
@property
1034
def penalties_state(self):
@@ -1214,6 +1217,11 @@ class DiffusionSampler:
1217
1218
slots_np = input_batch.idx_mapping_np[:num_reqs]
1219
max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)
1220
+ # Requests may ask for specific token ids' logprobs instead of, or as
1221
+ # well as, a top-k.
1222
+ max_token_ids = self.logprob_token_ids_state.max_num_token_ids(slots_np)
1223
+ want_logprobs = max_num_logprobs >= 0 or max_token_ids > 0
1224
+ num_logprobs = max(max_num_logprobs, 0)
1225
1226
# Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in
1227
# _compiled_sample_step keeps several live [group * CL, vocab] copies, so
@@ -1271,7 +1279,7 @@ class DiffusionSampler:
1279
1280
# Logprobs for denoise steps that just converged (is_encoder_phase
1281
# flipped False→True), stashed per tile so `scaled` is freed each tile.
1274
- if max_num_logprobs >= 0:
1282
+ if want_logprobs:
1283
converged_mask = states.is_encoder_phase[tile_slots]
1284
just_converged = converged_mask & ~is_committing[tile]
1285
if just_converged.any():
@@ -1285,17 +1293,33 @@ class DiffusionSampler:
1293
# positions are never emitted.
1294
k_i = int(valid_canvas_len_np[start_req + li])
1295
pos = li * CL
1296
+ per_req_ids = max_token_ids > 0
1297
self._pending_logprobs[slot.item()] = compute_topk_scores(
1298
flat_logits[pos : pos + k_i],
1290
- max_num_logprobs,
1299
+ num_logprobs,
1300
argmax_tokens[local_idx][:k_i],
1301
+ logprob_token_ids_state=(
1302
+ self.logprob_token_ids_state if per_req_ids else None
1303
+ ),
1304
+ # every row of this stash belongs to one slot
1305
+ expanded_idx_mapping=(
1306
+ torch.full(
1307
+ (k_i,),
1308
+ slot.item(),
1309
+ dtype=torch.int32,
1310
+ device=flat_logits.device,
1311
+ )
1312
+ if per_req_ids
1313
+ else None
1314
+ ),
1315
+ max_per_req_token_ids=max_token_ids,
1316
logits_mode=self.logits_mode,
1317
)
1318
1319
# Commit steps: is_committing was True at entry. Reassemble previously
1320
# stashed logprobs and attach to SamplerOutput.
1321
logprobs_tensors = None
1298
- if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs:
1322
+ if want_logprobs and is_committing.any() and self._pending_logprobs:
1323
committing_slots = set(
1324
decode_slots_np[is_committing.cpu().numpy()].tolist()
1325
)