[Bugfix] DiffusionGemma: cast the self-conditioning soft embed to the buffer dtype (#57462)
Signed-off-by: Matt Mastracci <matthew@mastracci.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Matt Mastracci committed
Sep 19, 2026 at 22:58 UTC
0c90804c0b8b139c252a86c79a3af26d181e11a0
1 file changed
+1
-1
vllm/model_executor/models/diffusion_gemma.py
+1
-1
@@ -543,7 +543,7 @@ def _compiled_sample_step(
543
if tp_size > 1:
544
soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name)
545
soft_embeds = soft_embeds * normalizer
546
- sc_embeds[decode_slots] = soft_embeds * sc_keep
546
+ sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype)
547
548
# Overwrite canvas with argmax for newly converged denoise requests
549
newly_converged = (converged & is_denoise).unsqueeze(1)