[Diffusion] Move diffusion time embedding to jit kernel (#16879)

python/sglang/multimodal_gen/runtime/layers/visual_embedding.py

@@ -19,10 +19,12 @@ from diffusers.models.embeddings import (
 )
 
 try:
-    from sgl_kernel.elementwise import timestep_embedding as timestep_embedding_cuda
+    from sglang.jit_kernel.timestep_embedding import (
+        timestep_embedding as timestep_embedding_cuda,
+    )
 except Exception as _e:
     # Fallback to diffusers implementation so downstream code can still run
-    # even if `sgl_kernel` is not installed/available.
+    # even if `jit_kernel` is not available.
     timestep_embedding_cuda = _get_timestep_embedding
 
 from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
@@ -86,14 +88,13 @@ class PatchEmbed(nn.Module):
 
 class Timesteps(_Timesteps):
     def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
-        t_emb = timestep_embedding_cuda(
+        return timestep_embedding_cuda(
             timesteps,
             self.num_channels,
             flip_sin_to_cos=self.flip_sin_to_cos,
             downscale_freq_shift=self.downscale_freq_shift,
             scale=self.scale,
         )
-        return t_emb
 
 
 class CombinedTimestepGuidanceTextProjEmbeddings(