[Diffusion] Add diffusion time embedding to jit kernel (#17658)

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

@@ -15,19 +15,18 @@ from diffusers.models.embeddings import (
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding
 from diffusers.models.embeddings import Timesteps as _Timesteps
 from diffusers.models.embeddings import (
-    get_timestep_embedding as _get_timestep_embedding,
+    get_timestep_embedding as timestep_embedding_diffusers,
 )
 
-try:
-    from sgl_kernel.elementwise 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.
-    timestep_embedding_cuda = _get_timestep_embedding
-
+from sglang.jit_kernel.timestep_embedding import (
+    timestep_embedding as timestep_embedding_cuda,
+)
 from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
 from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_cuda = current_platform.is_cuda()
 
 
 class PatchEmbed(nn.Module):
@@ -86,14 +85,22 @@ class PatchEmbed(nn.Module):
 
 class Timesteps(_Timesteps):
     def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
-        t_emb = 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
+        if _is_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,
+            )
+        else:
+            return timestep_embedding_diffusers(
+                timesteps,
+                self.num_channels,
+                flip_sin_to_cos=self.flip_sin_to_cos,
+                downscale_freq_shift=self.downscale_freq_shift,
+                scale=self.scale,
+            )
 
 
 class CombinedTimestepGuidanceTextProjEmbeddings(