From b4bed8c3981178d57a0581e32f97484530fc787a Mon Sep 17 00:00:00 2001
From: Yihao Wang <42559837+AgainstEntropy@users.noreply.github.com>
Date: Wed, 10 Jun 2026 19:29:21 -0700
Subject: [PATCH] [diffusion] fix: cast to float32 (from float64) in triton
 kernel to unblock torch.compile (#27709)

---
 .../sglang/jit_kernel/diffusion/triton/sana_wm_gdn_chunkwise.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn_chunkwise.py b/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn_chunkwise.py
index 70e5c0b8f..2014d2138 100644
--- a/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn_chunkwise.py
+++ b/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn_chunkwise.py
@@ -485,7 +485,7 @@ def _phase_a_z_kernel(
             k_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0
         ).to(tl.float32)
         K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :]
-        K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale
+        K = (tl.where(K_normed > 0, K_normed, 0.0) * k_scale).to(tl.float32)
 
         beta_K = beta_t[:, None] * K