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

python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh

@@ -1,173 +0,0 @@
-#include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.h>
-
-#include <sgl_kernel/utils.cuh>
-
-#include <dlpack/dlpack.h>
-#include <tvm/ffi/container/tensor.h>
-
-#include <algorithm>
-#include <cmath>
-#include <cstdint>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <type_traits>
-
-namespace {
-
-template <typename T>
-__device__ __forceinline__ float cast_to_float(T v) {
-  if constexpr (std::is_same_v<T, nv_bfloat16>) {
-    return __bfloat162float(v);
-  } else if constexpr (std::is_same_v<T, half>) {
-    return __half2float(v);
-  } else {
-    return static_cast<float>(v);
-  }
-}
-
-template <bool kFlipSinToCos, typename TIn>
-__global__ void timestep_embedding_kernel(
-    const TIn* __restrict__ t_ptr,
-    float* __restrict__ output_ptr,
-    int dim,
-    float neg_log_max_period,
-    float scale,
-    int batch_size) {
-  int row_idx = static_cast<int>(blockIdx.x * blockDim.y + threadIdx.y);
-  if (row_idx >= batch_size) {
-    return;
-  }
-
-  float t_val = cast_to_float(t_ptr[row_idx]);
-  float* output_batch_base_ptr = output_ptr + row_idx * dim;
-
-  int half_dim = dim / 2;
-  int thread_offset = static_cast<int>(threadIdx.x);
-  while (thread_offset * 4 < half_dim) {
-    float4* top_half;
-    float4* bottom_half;
-    if constexpr (!kFlipSinToCos) {
-      bottom_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
-      top_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
-    } else {
-      top_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
-      bottom_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
-    }
-
-    float4 vals;
-    vals.x = scale * t_val * expf(neg_log_max_period * __int2float_rn(thread_offset * 4 + 0));
-    vals.y = scale * t_val * expf(neg_log_max_period * __int2float_rn(thread_offset * 4 + 1));
-    vals.z = scale * t_val * expf(neg_log_max_period * __int2float_rn(thread_offset * 4 + 2));
-    vals.w = scale * t_val * expf(neg_log_max_period * __int2float_rn(thread_offset * 4 + 3));
-
-    float4 sin_vals;
-    sin_vals.x = cosf(vals.x);
-    sin_vals.y = cosf(vals.y);
-    sin_vals.z = cosf(vals.z);
-    sin_vals.w = cosf(vals.w);
-    *top_half = sin_vals;
-
-    float4 cos_vals;
-    cos_vals.x = sinf(vals.x);
-    cos_vals.y = sinf(vals.y);
-    cos_vals.z = sinf(vals.z);
-    cos_vals.w = sinf(vals.w);
-    *bottom_half = cos_vals;
-
-    thread_offset += static_cast<int>(blockDim.x);
-  }
-}
-
-template <typename TIn>
-inline void launch_timestep_embedding(
-    const tvm::ffi::TensorView t,
-    const tvm::ffi::TensorView output,
-    int dim,
-    bool flip_sin_to_cos,
-    float downscale_freq_shift,
-    float scale,
-    int max_period) {
-  using namespace host;
-
-  const int batch_size = static_cast<int>(t.shape()[0]);
-  const int half_dim = dim / 2;
-
-  constexpr int kMaxThreadsPerBlock = 1024;
-  constexpr int kMinThreadsPerBlock = 128;
-
-  const int num_threads_per_row = std::min(kMaxThreadsPerBlock, half_dim / 4);
-  const int num_rows = (kMinThreadsPerBlock + num_threads_per_row - 1) / num_threads_per_row;
-
-  dim3 grid((batch_size + num_rows - 1) / num_rows);
-  dim3 block(num_threads_per_row, num_rows);
-
-  const float neg_log_max_period =
-      std::log(static_cast<float>(max_period)) * (-1.0f) / (static_cast<float>(half_dim) - downscale_freq_shift);
-
-  const DLDevice device = output.device();
-
-  if (flip_sin_to_cos) {
-    LaunchKernel(grid, block, device)(
-        timestep_embedding_kernel<true, TIn>,
-        static_cast<const TIn*>(t.data_ptr()),
-        static_cast<float*>(output.data_ptr()),
-        dim,
-        neg_log_max_period,
-        scale,
-        batch_size);
-  } else {
-    LaunchKernel(grid, block, device)(
-        timestep_embedding_kernel<false, TIn>,
-        static_cast<const TIn*>(t.data_ptr()),
-        static_cast<float*>(output.data_ptr()),
-        dim,
-        neg_log_max_period,
-        scale,
-        batch_size);
-  }
-}
-
-void timestep_embedding(
-    tvm::ffi::TensorView input,
-    tvm::ffi::TensorView output,
-    int dim,
-    bool flip_sin_to_cos,
-    float downscale_freq_shift,
-    float scale,
-    int max_period) {
-  using namespace host;
-
-  auto B = SymbolicSize{"batch_size"};
-  auto D = SymbolicSize{"dim"};
-  auto device = SymbolicDevice{};
-
-  TensorMatcher({B}).with_strides({1}).template with_device<kDLCUDA>(device).verify(input);
-
-  TensorMatcher({B, D}).with_strides({D, 1}).with_dtype<float>().template with_device<kDLCUDA>(device).verify(output);
-
-  RuntimeCheck(D.unwrap() == dim, "Output dim mismatch: ", D.unwrap(), " vs ", dim);
-  RuntimeCheck(dim % 8 == 0, "dim must align to 8, got ", dim);
-
-  const DLDataType in_dtype = input.dtype();
-
-  const bool input_dtype_supported = (in_dtype.code == kDLFloat && in_dtype.bits == 16) ||
-                                     (in_dtype.code == kDLBfloat && in_dtype.bits == 16) ||
-                                     (in_dtype.code == kDLFloat && in_dtype.bits == 32);
-  RuntimeCheck(input_dtype_supported, "input dtype must be fp16/bf16/fp32, but got ", in_dtype);
-
-  auto launch = [&]<typename TIn>() {
-    launch_timestep_embedding<TIn>(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period);
-  };
-
-  if (in_dtype.code == kDLFloat && in_dtype.bits == 32) {
-    launch.template operator()<float>();
-  } else if (in_dtype.code == kDLBfloat && in_dtype.bits == 16) {
-    launch.template operator()<nv_bfloat16>();
-  } else if (in_dtype.code == kDLFloat && in_dtype.bits == 16) {
-    launch.template operator()<half>();
-  }
-}
-
-}  // namespace

python/sglang/jit_kernel/timestep_embedding.py

@@ -1,44 +0,0 @@
-from __future__ import annotations
-
-import functools
-from typing import TYPE_CHECKING
-
-import torch
-
-from sglang.jit_kernel.utils import load_jit
-
-if TYPE_CHECKING:
-    from tvm_ffi.module import Module
-
-
-@functools.cache
-def _jit_timestep_embedding_module() -> Module:
-    return load_jit(
-        "timestep_embedding",
-        cuda_files=["diffusion/timestep_embedding.cuh"],
-        cuda_wrappers=[("timestep_embedding", "timestep_embedding")],
-    )
-
-
-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,
-) -> torch.Tensor:
-    dtype = torch.float32
-    output = torch.empty((t.shape[0], dim), dtype=dtype, device=t.device)
-    module = _jit_timestep_embedding_module()
-    module.timestep_embedding(
-        t,
-        output,
-        dim,
-        flip_sin_to_cos,
-        float(downscale_freq_shift),
-        float(scale),
-        int(max_period),
-    )
-    return output

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

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

sgl-kernel/CMakeLists.txt

@@ -282,6 +282,7 @@ set(SOURCES
     "csrc/elementwise/rope.cu"
     "csrc/elementwise/pos_enc.cu"
     "csrc/elementwise/topk.cu"
+    "csrc/sgl_diffusion/elementwise/timestep_embedding.cu"
     "csrc/expert_specialization/es_fp8_blockwise.cu"
     "csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu"
     "csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu"

sgl-kernel/csrc/common_extension.cc

@@ -609,6 +609,19 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
 
   m.def("fast_hadamard_transform_40N(Tensor x, float scale) -> Tensor");
   m.impl("fast_hadamard_transform_40N", torch::kCUDA, &fast_hadamard_transform_40N);
+
+  /*
+   * From csrc/sgl_diffusion/elementwise
+   */
+  m.def(
+      "timestep_embedding(Tensor input,"
+      "Tensor output,"
+      "int dim,"
+      "bool flip_sin_to_cos,"
+      "float downscale_freq_shift,"
+      "float scale,"
+      "int max_period) -> Tensor");
+  m.impl("timestep_embedding", torch::kCUDA, &timestep_embedding);
 }
 
 REGISTER_EXTENSION(common_ops)

sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu

@@ -0,0 +1,137 @@
+/* Copyright 2025 SGLang Team. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <math.h>
+#include <torch/all.h>
+
+#include <cassert>
+#include <cmath>
+
+#include "utils.h"
+
+template <bool flip_sin_to_cos = false, typename T_IN>
+__global__ void timestep_embedding_kernel(
+    T_IN* t_ptr, float* output_ptr, int dim, float neg_log_max_period, float scale, int batch_size) {
+  // Get the timestep for this batch
+  int row_idx = blockIdx.x * blockDim.y + threadIdx.y;
+  if (row_idx >= batch_size) {
+    return;
+  }
+  // Use the portable LDG helper (maps to __ldg on CUDA, plain load on ROCm/HIP).
+  float t_val = castToFloat(SGLANG_LDG(&t_ptr[row_idx]));
+  float* output_batch_base_ptr = output_ptr + row_idx * dim;
+
+  // Calculate half dimension
+  int half_dim = dim / 2;
+  int thread_offset = threadIdx.x % blockDim.x;
+  while (thread_offset * 4 < half_dim) {
+    float4* top_half;
+    float4* bottom_half;
+    if constexpr (flip_sin_to_cos == false) {
+      bottom_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
+      top_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
+    } else {
+      top_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
+      bottom_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
+    }
+
+    float4 vals;
+    vals.x = scale * t_val * expf(neg_log_max_period * __int2float_rn(thread_offset * 4 + 0));
+    vals.y = scale * t_val * expf(neg_log_max_period * __int2float_rn(thread_offset * 4 + 1));
+    vals.z = scale * t_val * expf(neg_log_max_period * __int2float_rn(thread_offset * 4 + 2));
+    vals.w = scale * t_val * expf(neg_log_max_period * __int2float_rn(thread_offset * 4 + 3));
+
+    float4 sin_vals;
+    sin_vals.x = cosf(vals.x);
+    sin_vals.y = cosf(vals.y);
+    sin_vals.z = cosf(vals.z);
+    sin_vals.w = cosf(vals.w);
+    *top_half = sin_vals;  // STG.128
+
+    float4 cos_vals;
+    cos_vals.x = sinf(vals.x);
+    cos_vals.y = sinf(vals.y);
+    cos_vals.z = sinf(vals.z);
+    cos_vals.w = sinf(vals.w);
+    *bottom_half = cos_vals;  // STG.128
+
+    thread_offset += blockDim.x;
+  }
+}
+
+torch::Tensor timestep_embedding(
+    const torch::Tensor& t,
+    torch::Tensor& output,
+    int64_t dim,
+    bool flip_sin_to_cos,
+    double downscale_freq_shift,
+    double scale,
+    int64_t max_period) {
+  TORCH_CHECK(t.dim() == 1 and t.stride(0) == 1, "t should be 1D");
+  TORCH_CHECK(output.dim() == 2 and output.is_contiguous(), "output should be a contiguous 2D tensor.");
+
+  const int batch_size = static_cast<int>(t.size(0));
+  TORCH_CHECK(output.size(0) == batch_size, "Output batch size doesn't match t");
+  TORCH_CHECK(output.size(1) == dim, "Output feature size doesn't match dim");
+
+  TORCH_CHECK(t.device().is_cuda(), "t must be a CUDA tensor");
+  TORCH_CHECK(output.device().is_cuda(), "output must be a CUDA tensor");
+  TORCH_CHECK(t.device() == output.device(), "t and output must be on the same device");
+
+  // To align with timestep_embedding python code.
+  TORCH_CHECK(output.scalar_type() == at::ScalarType::Float, "Output buffer should be float32.");
+
+  TORCH_CHECK(dim % 8 == 0, "dim should align to 8");
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  constexpr int MAX_THREADS_PER_BLOCK = 1024;
+  constexpr int MIN_THREADS_PER_BLOCK = 128;
+  int half_dim = dim / 2;
+  int num_threads_per_row = min(MAX_THREADS_PER_BLOCK, half_dim / 4);
+  int num_rows = (MIN_THREADS_PER_BLOCK + num_threads_per_row - 1) / num_threads_per_row;
+
+  dim3 grid((batch_size + num_rows - 1) / num_rows);
+  // assert float4 vectorize output
+  dim3 block(num_threads_per_row, num_rows);
+  float neg_log_max_period =
+      std::log(static_cast<float>(max_period)) * (-1.0f) / (static_cast<float>(half_dim) - downscale_freq_shift);
+
+  AT_DISPATCH_ALL_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, t.scalar_type(), "timestep_embedding_kernel", [&] {
+        if (flip_sin_to_cos == true) {
+          timestep_embedding_kernel<true><<<grid, block, 0, stream>>>(
+              reinterpret_cast<scalar_t*>(t.data_ptr()),
+              reinterpret_cast<float*>(output.data_ptr()),
+              static_cast<int>(dim),
+              static_cast<float>(neg_log_max_period),
+              static_cast<float>(scale),
+              static_cast<int>(batch_size));
+        } else {
+          timestep_embedding_kernel<false><<<grid, block, 0, stream>>>(
+              reinterpret_cast<scalar_t*>(t.data_ptr()),
+              reinterpret_cast<float*>(output.data_ptr()),
+              static_cast<int>(dim),
+              static_cast<float>(neg_log_max_period),
+              static_cast<float>(scale),
+              static_cast<int>(batch_size));
+        }
+      });
+
+  return output;
+}

sgl-kernel/include/sgl_kernel_ops.h

@@ -1006,3 +1006,15 @@ std::vector<at::Tensor> fwd_kvcache_mla_fp8(
 
 std::vector<at::Tensor> get_mla_decoding_metadata_dense_fp8(
     at::Tensor& seqlens_k, const int64_t num_heads_per_head_k, const int64_t num_heads_k);
+
+/*
+ * From csrc/sgl_diffusion/elementwise
+ */
+torch::Tensor timestep_embedding(
+    const torch::Tensor& t,
+    torch::Tensor& output,
+    int64_t dim,
+    bool flip_sin_to_cos,
+    double downscale_freq_shift,
+    double scale,
+    int64_t max_period);

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
+    )

sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py

@@ -3,10 +3,8 @@ import pytest
 import tabulate
 import torch
 from diffusers.models.embeddings import get_timestep_embedding
+from sgl_kernel.elementwise import timestep_embedding as timestep_embedding_cuda
 
-from sglang.jit_kernel.timestep_embedding import (
-    timestep_embedding as timestep_embedding_cuda,
-)
 from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
 
 
@@ -14,7 +12,9 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embed
     "batch_size", [1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384]
 )
 @pytest.mark.parametrize("dim", [32, 128, 256, 512, 1536, 2048, 4096, 8192])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize(
+    "dtype", [torch.int32, torch.int64, torch.bfloat16, torch.float16]
+)
 def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype):
     device = "cuda"
     t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
@@ -25,7 +25,7 @@ def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype):
 
 @pytest.mark.parametrize("batch_size", [1, 2, 8, 128, 256, 512, 1536, 2048, 16384])
 @pytest.mark.parametrize("dim", [32, 256, 512, 1536, 8192])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("dtype", [torch.int32, torch.bfloat16])
 @pytest.mark.parametrize("flip_sin_to_cos", [False, True])
 @pytest.mark.parametrize("downscale_freq_shift", [0, 1])
 @pytest.mark.parametrize("scale", [1, 0.01])
@@ -81,7 +81,7 @@ def test_timestep_embedding_perf():
     for B in NUM_BATCH:
         for dim in NUM_DIM:
             t = torch.linspace(0, max(100000, B), steps=B, device=device).to(
-                torch.float32
+                torch.int32
             )
             time_torch = perf_kernel_fn(timestep_embedding, t, dim)
             time_cuda = perf_kernel_fn(timestep_embedding_cuda, t, dim)