[Model] Support top_k and top_p sampling for DiffusionGemma (#45429)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
Guan-Ming Chiu committed
Jul 26, 2026 at 16:39 UTC
21fd9e85a04f3a97da00eceddce91d8c6966b954
1 file changed
+18
-2
vllm/model_executor/models/diffusion_gemma.py
+18
-2
@@ -49,6 +49,7 @@ from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
49
from vllm.multimodal import MULTIMODAL_REGISTRY
50
from vllm.platforms import current_platform
51
from vllm.v1.outputs import LogprobsTensors
52
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
53
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
54
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu
55
from vllm.v1.worker.gpu.input_batch import InputBatch
@@ -1262,16 +1263,31 @@ class DiffusionSampler:
1263
valid_canvas_len_np.astype(np.int64), device=device
1264
)
1265
1266
+ # Per-request top_k/top_p, mirroring the AR sampler. Masked tokens
1267
+ # become -inf and survive the temperature scaling in the compiled
1268
+ # step, so Gumbel sampling, probs, and entropy all see the filtered
1269
+ # distribution. The committed argmax (always the top-1 token) is
1270
+ # unaffected; only the canvas exploration is constrained. Applied
1271
+ # before canvas padding so phantom positions stay uniform.
1272
+ if num_decode > 0:
1273
+ top_k, top_p = self.sampling_states.get_top_k_top_p(
1274
+ decode_slots.repeat_interleave(valid_canvas_len), decode_slots_np
1275
+ )
1276
+ if top_k is not None or top_p is not None:
1277
+ logits = apply_top_k_top_p(logits.float(), top_k, top_p)
1278
+
1279
# Pad any truncated canvas back to CL so the uniform-CL sampler math
1280
# holds. Phantom (padded) positions are zeroed → uniform logits → high
1281
# entropy (no premature convergence) and argmax 0 (stable); they are
1268
- # never committed (num_sampled == real length).
1282
+ # never committed (num_sampled == real length). masked_fill (not
1283
+ # multiply) so -inf entries from top_k/top_p filtering above don't
1284
+ # turn phantom rows into NaN.
1285
if num_decode > 0 and valid_canvas_len_np.min() < CL:
1286
ar = torch.arange(CL, device=device)
1287
starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req
1288
valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL]
1289
src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1)
1274
- logits = logits[src.reshape(-1)] * valid.reshape(-1, 1).to(logits.dtype)
1290
+ logits = logits[src.reshape(-1)].masked_fill_(~valid.reshape(-1, 1), 0)
1291
1292
# Clear once: the tiled loop below only scatters its own decode slots,
1293
# so it must not re-clear earlier tiles' writes.