[diffusion] fix: cast to float32 (from float64) in triton kernel to unblock torch.compile (#27709)

2026-06-11 10:29:21 +08:00
parent 588d1f7bc9
commit b4bed8c398
1 changed files with 1 additions and 1 deletions
@@ -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