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

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

@@ -0,0 +1,150 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+namespace {
+
+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 = device::cast<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 * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 0));
+    vals.y = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 1));
+    vals.z = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 2));
+    vals.w = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 3));
+
+    float4 cos_vals;
+    cos_vals.x = device::math::cos(vals.x);
+    cos_vals.y = device::math::cos(vals.y);
+    cos_vals.z = device::math::cos(vals.z);
+    cos_vals.w = device::math::cos(vals.w);
+    *top_half = cos_vals;
+
+    float4 sin_vals;
+    sin_vals.x = device::math::sin(vals.x);
+    sin_vals.y = device::math::sin(vals.y);
+    sin_vals.z = device::math::sin(vals.z);
+    sin_vals.w = device::math::sin(vals.w);
+    *bottom_half = sin_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);
+  }
+}
+
+template <typename TIn>
+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})  // input
+      .with_strides({1})
+      .with_dtype<TIn>()
+      .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);
+
+  launch_timestep_embedding<TIn>(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period);
+}
+
+}  // namespace

python/sglang/jit_kernel/include/sgl_kernel/math.cuh

@@ -1,6 +1,8 @@
 #pragma once
 #include <sgl_kernel/type.cuh>
 
+#include <cmath>
+
 namespace device::math {
 
 inline constexpr float log2e = 1.44269504088896340736f;
@@ -33,4 +35,16 @@ SGL_DEVICE T rsqrt(T a) {
   return dtype_trait<T>::rsqrt(a);
 }
 
+SGL_DEVICE float exp(float a) {
+  return ::expf(a);
+}
+
+SGL_DEVICE float sin(float a) {
+  return ::sinf(a);
+}
+
+SGL_DEVICE float cos(float a) {
+  return ::cosf(a);
+}
+
 }  // namespace device::math

python/sglang/jit_kernel/include/sgl_kernel/type.cuh

@@ -37,6 +37,8 @@ struct dtype_trait {};
   static_assert(true)
 
 SGL_REGISTER_DTYPE_TRAIT(fp32_t, fp32x2_t, SGL_REGISTER_TYPE_END;  //
+                         SGL_REGISTER_FROM_FUNCTION(fp16_t, __half2float);
+                         SGL_REGISTER_FROM_FUNCTION(bf16_t, __bfloat162float);
                          SGL_REGISTER_UNARY_FUNCTION(abs, fabsf);
                          SGL_REGISTER_UNARY_FUNCTION(sqrt, sqrtf);
                          SGL_REGISTER_UNARY_FUNCTION(rsqrt, rsqrtf);

python/sglang/jit_kernel/tests/test_timestep_embedding.py

@@ -0,0 +1,160 @@
+import os
+
+import numpy as np
+import pytest
+import torch
+
+try:
+    import tabulate
+except Exception:
+    tabulate = None
+
+from sglang.jit_kernel.timestep_embedding import (
+    timestep_embedding as timestep_embedding_cuda,
+)
+
+
+def get_timestep_embedding_reference(
+    timesteps: torch.Tensor,
+    dim: int,
+    *,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+):
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    timesteps = timesteps.to(torch.float32)
+    half_dim = dim // 2
+    exponent = -torch.log(
+        torch.tensor(max_period, dtype=torch.float32, device=timesteps.device)
+    ) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    emb = scale * emb
+
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+    if dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+@pytest.mark.parametrize(
+    "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])
+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)
+    torch_output = get_timestep_embedding_reference(
+        t, dim, flip_sin_to_cos=True, downscale_freq_shift=0
+    )
+    cuda_output = timestep_embedding_cuda(
+        t, dim, flip_sin_to_cos=True, downscale_freq_shift=0
+    )
+    torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
+
+
+@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("flip_sin_to_cos", [False, True])
+@pytest.mark.parametrize("downscale_freq_shift", [0, 1])
+@pytest.mark.parametrize("scale", [1, 0.01])
+def test_timestep_embedding_correctness_with_diffusers(
+    batch_size, dim, flip_sin_to_cos, downscale_freq_shift, scale, dtype
+):
+    device = "cuda"
+    t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
+    torch_output = get_timestep_embedding_reference(
+        t,
+        dim,
+        flip_sin_to_cos=flip_sin_to_cos,
+        downscale_freq_shift=downscale_freq_shift,
+        scale=scale,
+        max_period=10000,
+    )
+    cuda_output = timestep_embedding_cuda(
+        t,
+        dim,
+        flip_sin_to_cos=flip_sin_to_cos,
+        downscale_freq_shift=downscale_freq_shift,
+        scale=scale,
+        max_period=10000,
+    )
+    torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
+
+
+def test_timestep_embedding_perf():
+    if os.environ.get("SGLANG_RUN_JIT_KERNEL_PERF_TESTS") != "1":
+        pytest.skip("Perf test disabled by default")
+    if tabulate is None:
+        pytest.skip("Optional dependency 'tabulate' is not installed")
+
+    NUM_BATCH = [1, 2, 8, 63, 256, 512, 613, 1024, 1536]
+    NUM_DIM = [32, 64, 128, 256, 512, 1024, 2048, 4096]
+
+    def perf_kernel_fn(kernel_fn: callable, *args, **kwargs):
+        warmup_times = 4
+        repeat_times = 20
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        for _ in range(warmup_times):
+            output_fn = kernel_fn(*args, **kwargs)
+        torch.cuda.synchronize()
+
+        start.record()
+        for _ in range(repeat_times):
+            output_fn = kernel_fn(*args, **kwargs)
+        end.record()
+        end.synchronize()
+        return start.elapsed_time(end) / repeat_times
+
+    device = "cuda"
+    results = []
+
+    cuda_speedups = []
+    for B in NUM_BATCH:
+        for dim in NUM_DIM:
+            t = torch.linspace(0, max(100000, B), steps=B, device=device).to(
+                torch.float32
+            )
+            time_torch = perf_kernel_fn(get_timestep_embedding_reference, t, dim)
+            time_cuda = perf_kernel_fn(timestep_embedding_cuda, t, dim)
+            speedup_cuda = time_torch / time_cuda
+
+            results.append(
+                {
+                    "Batch Size": B,
+                    "Dimension": dim,
+                    "Torch Time (ms)": time_torch,
+                    "CUDA Time (ms)": time_cuda,
+                    "Speedup (CUDA)": speedup_cuda,
+                }
+            )
+            cuda_speedups.append(speedup_cuda)
+
+    print("=== Timestep Embedding Benchmark Results ===")
+    print(
+        tabulate.tabulate(
+            results,
+            headers="keys",
+            tablefmt="fancy_grid",
+            floatfmt=(".0f", ".0f", ".6f", ".6f", ".5f"),
+        )
+    )
+    print(f"Average Speedup(cuda): {np.mean(cuda_speedups):.4f}")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

python/sglang/jit_kernel/timestep_embedding.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@functools.cache
+def _jit_timestep_embedding_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "timestep_embedding",
+        *args,
+        cuda_files=["diffusion/timestep_embedding.cuh"],
+        cuda_wrappers=[("timestep_embedding", f"timestep_embedding<{args}>")],
+    )
+
+
+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:
+    if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        t = t.to(dtype)
+    output = torch.empty((t.shape[0], dim), dtype=torch.float32, device=t.device)
+    module = _jit_timestep_embedding_module(t.dtype)
+    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

@@ -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(