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..2d29da50b --- /dev/null +++ b/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh @@ -0,0 +1,150 @@ +#include +#include + +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include + +namespace { + +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 = device::cast(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 * 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(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); + } +} + +template +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() + .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); + + launch_timestep_embedding(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period); +} + +} // namespace diff --git a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh b/python/sglang/jit_kernel/include/sgl_kernel/math.cuh index 3d4aa7473..97e6cf637 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/math.cuh @@ -1,6 +1,8 @@ #pragma once #include +#include + namespace device::math { inline constexpr float log2e = 1.44269504088896340736f; @@ -33,4 +35,16 @@ SGL_DEVICE T rsqrt(T a) { return dtype_trait::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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh b/python/sglang/jit_kernel/include/sgl_kernel/type.cuh index 047e5c24b..ff31c5ff3 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh +++ b/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); diff --git a/python/sglang/jit_kernel/tests/test_timestep_embedding.py b/python/sglang/jit_kernel/tests/test_timestep_embedding.py new file mode 100644 index 000000000..2ed242429 --- /dev/null +++ b/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__]) diff --git a/python/sglang/jit_kernel/timestep_embedding.py b/python/sglang/jit_kernel/timestep_embedding.py new file mode 100644 index 000000000..75d4bfe60 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py index 71a88e788..535bf3311 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/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(