From 2cdd4370bc9217e958fa340f1f1754ccf169c372 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sat, 17 Jan 2026 12:21:22 +0800 Subject: [PATCH] [Diffusion] Move diffusion time embedding to jit kernel (#16879) --- .../csrc/diffusion/timestep_embedding.cuh | 173 ++++++++++++++++++ .../tests}/test_timestep_embedding.py | 12 +- .../sglang/jit_kernel/timestep_embedding.py | 44 +++++ .../runtime/layers/visual_embedding.py | 9 +- sgl-kernel/CMakeLists.txt | 1 - sgl-kernel/csrc/common_extension.cc | 13 -- .../elementwise/timestep_embedding.cu | 137 -------------- sgl-kernel/include/sgl_kernel_ops.h | 12 -- sgl-kernel/python/sgl_kernel/__init__.py | 1 - sgl-kernel/python/sgl_kernel/elementwise.py | 32 ---- 10 files changed, 228 insertions(+), 206 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh rename {sgl-kernel/tests/sgl_diffusion => python/sglang/jit_kernel/tests}/test_timestep_embedding.py (92%) create mode 100644 python/sglang/jit_kernel/timestep_embedding.py delete mode 100644 sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu diff --git a/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh b/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh new file mode 100644 index 000000000..ac969c587 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh @@ -0,0 +1,173 @@ +#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/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py b/python/sglang/jit_kernel/tests/test_timestep_embedding.py similarity index 92% rename from sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py rename to python/sglang/jit_kernel/tests/test_timestep_embedding.py index 53784444b..11834b355 100644 --- a/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py +++ b/python/sglang/jit_kernel/tests/test_timestep_embedding.py @@ -3,8 +3,10 @@ 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 @@ -12,9 +14,7 @@ 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.int32, torch.int64, torch.bfloat16, torch.float16] -) +@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) @@ -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.int32, torch.bfloat16]) +@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]) @@ -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.int32 + torch.float32 ) time_torch = perf_kernel_fn(timestep_embedding, t, dim) time_cuda = perf_kernel_fn(timestep_embedding_cuda, t, dim) diff --git a/python/sglang/jit_kernel/timestep_embedding.py b/python/sglang/jit_kernel/timestep_embedding.py new file mode 100644 index 000000000..f32f4c8f1 --- /dev/null +++ b/python/sglang/jit_kernel/timestep_embedding.py @@ -0,0 +1,44 @@ +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 71a88e788..6d3c51ec2 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py @@ -19,10 +19,12 @@ from diffusers.models.embeddings import ( ) try: - from sgl_kernel.elementwise import timestep_embedding as timestep_embedding_cuda + from sglang.jit_kernel.timestep_embedding 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. + # even if `jit_kernel` is not available. timestep_embedding_cuda = _get_timestep_embedding from sglang.multimodal_gen.runtime.layers.activation import get_act_fn @@ -86,14 +88,13 @@ class PatchEmbed(nn.Module): class Timesteps(_Timesteps): def forward(self, timesteps: torch.Tensor) -> torch.Tensor: - t_emb = timestep_embedding_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, ) - return t_emb class CombinedTimestepGuidanceTextProjEmbeddings( diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 81a020d43..e41cf6f0a 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -282,7 +282,6 @@ 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 b38eab218..d0b6fcf80 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -609,19 +609,6 @@ 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 deleted file mode 100644 index d241619e1..000000000 --- a/sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu +++ /dev/null @@ -1,137 +0,0 @@ -/* 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 5e3cf24f9..2eb0856aa 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -1006,15 +1006,3 @@ 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 1b97ef94f..8e8994e04 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -32,7 +32,6 @@ 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 68dc221d1..c7a6a0ed1 100644 --- a/sgl-kernel/python/sgl_kernel/elementwise.py +++ b/sgl-kernel/python/sgl_kernel/elementwise.py @@ -404,35 +404,3 @@ 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 - )