From 46be74b4b4ef0fe2c738ba8b8c606afe16f2d136 Mon Sep 17 00:00:00 2001 From: 66RING Date: Fri, 19 Dec 2025 20:59:50 +0800 Subject: [PATCH] [diffusion] kernel: timestep embedding kernel implementation (#12995) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 戚余航 Co-authored-by: Qi Yuhang <45795032+HydraQYH@users.noreply.github.com> --- .../runtime/layers/visual_embedding.py | 60 ++++++++ sgl-kernel/CMakeLists.txt | 1 + sgl-kernel/csrc/common_extension.cc | 13 ++ .../elementwise/timestep_embedding.cu | 136 ++++++++++++++++++ sgl-kernel/include/sgl_kernel_ops.h | 12 ++ sgl-kernel/python/sgl_kernel/__init__.py | 1 + sgl-kernel/python/sgl_kernel/elementwise.py | 32 +++++ .../sgl_diffusion/test_timestep_embedding.py | 114 +++++++++++++++ 8 files changed, 369 insertions(+) create mode 100644 sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu create mode 100644 sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py diff --git a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py index d556ab584..39fd32593 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py @@ -6,6 +6,15 @@ import math import torch import torch.nn as nn +from diffusers.models.embeddings import ( + CombinedTimestepGuidanceTextProjEmbeddings as _CombinedTimestepGuidanceTextProjEmbeddings, +) +from diffusers.models.embeddings import ( + CombinedTimestepTextProjEmbeddings as _CombinedTimestepTextProjEmbeddings, +) +from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding +from diffusers.models.embeddings import Timesteps as _Timesteps +from sgl_kernel.elementwise 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 ReplicatedLinear @@ -66,6 +75,57 @@ class PatchEmbed(nn.Module): return x +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 + + +class CombinedTimestepGuidanceTextProjEmbeddings( + _CombinedTimestepGuidanceTextProjEmbeddings +): + def __init__(self, embedding_dim, pooled_projection_dim): + nn.Module.__init__(self) + + # use sgld op + self.time_proj = Timesteps( + num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + # use diffusers op + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim + ) + self.guidance_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim + ) + self.text_embedder = PixArtAlphaTextProjection( + pooled_projection_dim, embedding_dim, act_fn="silu" + ) + + +class CombinedTimestepTextProjEmbeddings(_CombinedTimestepTextProjEmbeddings): + def __init__(self, embedding_dim, pooled_projection_dim): + nn.Module.__init__(self) + + # use sgld op + self.time_proj = Timesteps( + num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + # use diffusers op + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim + ) + self.text_embedder = PixArtAlphaTextProjection( + pooled_projection_dim, embedding_dim, act_fn="silu" + ) + + class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 557b82510..1862e8707 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 8c36e9a8f..463b924fa 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..e239569f5 --- /dev/null +++ b/sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu @@ -0,0 +1,136 @@ +/* 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; + } + float t_val = castToFloat(__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/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py b/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py new file mode 100644 index 000000000..53784444b --- /dev/null +++ b/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py @@ -0,0 +1,114 @@ +import numpy as np +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.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding + + +@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.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) + torch_output = timestep_embedding(t, dim) + cuda_output = timestep_embedding_cuda(t, dim, flip_sin_to_cos=True) + 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.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]) +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( + 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(): + 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.int32 + ) + time_torch = perf_kernel_fn(timestep_embedding, 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__])