Revert "[Diffusion] Move diffusion time embedding to jit kernel" (#17257)

sgl-kernel/python/sgl_kernel/__init__.py

@@ -32,6 +32,7 @@ from sgl_kernel.elementwise import (
     rmsnorm,
     rotary_embedding,
     silu_and_mul,
+    timestep_embedding,
 )
 from sgl_kernel.expert_specialization import (
     es_fp8_blockwise_scaled_grouped_mm,

sgl-kernel/python/sgl_kernel/elementwise.py

@@ -404,3 +404,35 @@ def concat_mla_absorb_q(
     )
     torch.ops.sgl_kernel.concat_mla_absorb_q(a, b, out)
     return out
+
+
+def timestep_embedding(
+    t: torch.Tensor,
+    dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 0.0,
+    scale: float = 1,
+    max_period: int = 10000,
+    dtype: torch.dtype = torch.float32,
+):
+    """
+    Create sinusoidal timestep embeddings.
+
+    # TODO: review, output dtype always be float32. According to python code:
+    #  sglang/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py
+
+    Args:
+        t: Tensor of shape [B] with timesteps
+        dim: Embedding dimension
+        max_period: Controls the minimum frequency of the embeddings
+
+    Returns:
+        Tensor of shape [B, dim] with embeddings
+    """
+    dtype = torch.float32
+
+    batch_size = t.shape[0]
+    output = torch.empty((batch_size, dim), dtype=dtype, device=t.device)
+    return torch.ops.sgl_kernel.timestep_embedding(
+        t, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period
+    )