diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index c2211b9fc..a2e8e9988 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -150,7 +150,6 @@ jobs: docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_activation.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_topk.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_kvcacheio.py - docker exec -w /sglang-checkout/sgl-kernel/tests/sgl_diffusion ci_sglang python3 -m pytest test_timestep_embedding.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_topk_sigmoid.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_torch_defaults_reset.py diff --git a/3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt b/3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt index e4d29ae73..86ba249c1 100644 --- a/3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt +++ b/3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt @@ -109,7 +109,6 @@ ${PROJ_ROOT}/csrc/kvcacheio/transfer.hip ${PROJ_ROOT}/csrc/moe/moe_align_kernel.hip ${PROJ_ROOT}/csrc/moe/moe_topk_softmax_kernels.hip ${PROJ_ROOT}/csrc/moe/moe_topk_sigmoid_kernels.hip -${PROJ_ROOT}/csrc/sgl_diffusion/elementwise/timestep_embedding.hip ${PROJ_ROOT}/csrc/speculative/eagle_utils.hip ) set_source_files_properties( diff --git a/3rdparty/amd/sgl-kernel/rocm_hipify.py b/3rdparty/amd/sgl-kernel/rocm_hipify.py index c758fe0f7..8373f741d 100644 --- a/3rdparty/amd/sgl-kernel/rocm_hipify.py +++ b/3rdparty/amd/sgl-kernel/rocm_hipify.py @@ -24,7 +24,6 @@ sources = [ "csrc/moe/moe_align_kernel.cu", "csrc/moe/moe_topk_softmax_kernels.cu", "csrc/moe/moe_topk_sigmoid_kernels.cu", - "csrc/sgl_diffusion/elementwise/timestep_embedding.cu", "csrc/speculative/eagle_utils.cu", ] 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/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index add827787..495fb43a2 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/sgl-kernel/csrc/common_extension_rocm.cc @@ -219,18 +219,6 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { " Tensor!? key, int head_size," " Tensor cos_sin_cache, bool is_neox) -> ()"); m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding); - /* - * 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); /* * From csrc/memory 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 - ) diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index 16fdbbd2f..66713bf0a 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -55,7 +55,6 @@ sources = [ "csrc/kvcacheio/transfer.cu", "csrc/memory/weak_ref_tensor.cpp", "csrc/elementwise/pos_enc.cu", - "csrc/sgl_diffusion/elementwise/timestep_embedding.cu", ] cxx_flags = ["-O3"] diff --git a/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py b/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py deleted file mode 100644 index 53784444b..000000000 --- a/sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py +++ /dev/null @@ -1,114 +0,0 @@ -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__])