diff --git a/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh b/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh deleted file mode 100644 index ac969c587..000000000 --- a/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh +++ /dev/null @@ -1,173 +0,0 @@ -#include -#include - -#include - -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace { - -template -__device__ __forceinline__ float cast_to_float(T v) { - if constexpr (std::is_same_v) { - return __bfloat162float(v); - } else if constexpr (std::is_same_v) { - return __half2float(v); - } else { - return static_cast(v); - } -} - -template -__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(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(threadIdx.x); - while (thread_offset * 4 < half_dim) { - float4* top_half; - float4* bottom_half; - if constexpr (!kFlipSinToCos) { - bottom_half = reinterpret_cast(output_batch_base_ptr + thread_offset * 4); - top_half = reinterpret_cast(output_batch_base_ptr + half_dim + thread_offset * 4); - } else { - top_half = reinterpret_cast(output_batch_base_ptr + thread_offset * 4); - bottom_half = reinterpret_cast(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(blockDim.x); - } -} - -template -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(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(max_period)) * (-1.0f) / (static_cast(half_dim) - downscale_freq_shift); - - const DLDevice device = output.device(); - - if (flip_sin_to_cos) { - LaunchKernel(grid, block, device)( - timestep_embedding_kernel, - static_cast(t.data_ptr()), - static_cast(output.data_ptr()), - dim, - neg_log_max_period, - scale, - batch_size); - } else { - LaunchKernel(grid, block, device)( - timestep_embedding_kernel, - static_cast(t.data_ptr()), - static_cast(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(device).verify(input); - - TensorMatcher({B, D}).with_strides({D, 1}).with_dtype().template with_device(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 = [&]() { - launch_timestep_embedding(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()(); - } else if (in_dtype.code == kDLBfloat && in_dtype.bits == 16) { - launch.template operator()(); - } else if (in_dtype.code == kDLFloat && in_dtype.bits == 16) { - launch.template operator()(); - } -} - -} // namespace diff --git a/python/sglang/jit_kernel/timestep_embedding.py b/python/sglang/jit_kernel/timestep_embedding.py deleted file mode 100644 index f32f4c8f1..000000000 --- a/python/sglang/jit_kernel/timestep_embedding.py +++ /dev/null @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py index 6d3c51ec2..71a88e788 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/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 ( - timestep_embedding as timestep_embedding_cuda, - ) + 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 `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( diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index e41cf6f0a..81a020d43 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/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" diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index d0b6fcf80..b38eab218 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/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, ×tep_embedding); } REGISTER_EXTENSION(common_ops) diff --git a/sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu new file mode 100644 index 000000000..d241619e1 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +#include +#include + +#include "utils.h" + +template +__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(output_batch_base_ptr + thread_offset * 4); + top_half = reinterpret_cast(output_batch_base_ptr + half_dim + thread_offset * 4); + } else { + top_half = reinterpret_cast(output_batch_base_ptr + thread_offset * 4); + bottom_half = reinterpret_cast(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(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(max_period)) * (-1.0f) / (static_cast(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<<>>( + reinterpret_cast(t.data_ptr()), + reinterpret_cast(output.data_ptr()), + static_cast(dim), + static_cast(neg_log_max_period), + static_cast(scale), + static_cast(batch_size)); + } else { + timestep_embedding_kernel<<>>( + reinterpret_cast(t.data_ptr()), + reinterpret_cast(output.data_ptr()), + static_cast(dim), + static_cast(neg_log_max_period), + static_cast(scale), + static_cast(batch_size)); + } + }); + + return output; +} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 2eb0856aa..5e3cf24f9 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -1006,3 +1006,15 @@ std::vector fwd_kvcache_mla_fp8( std::vector 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); diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 8e8994e04..1b97ef94f 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/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, diff --git a/sgl-kernel/python/sgl_kernel/elementwise.py b/sgl-kernel/python/sgl_kernel/elementwise.py index c7a6a0ed1..68dc221d1 100644 --- a/sgl-kernel/python/sgl_kernel/elementwise.py +++ b/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 + ) diff --git a/python/sglang/jit_kernel/tests/test_timestep_embedding.py b/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py similarity index 92% rename from python/sglang/jit_kernel/tests/test_timestep_embedding.py rename to sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py index 11834b355..53784444b 100644 --- a/python/sglang/jit_kernel/tests/test_timestep_embedding.py +++ b/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)