[Diffusion] Delete sgl-kernel outdated time_embedding kernel (#17278)

.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
 

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(

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",
 ]
 

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"

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, &timestep_embedding);
 }
 
 REGISTER_EXTENSION(common_ops)

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, &timestep_embedding);
 
   /*
    * From csrc/memory

sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu

@@ -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 <ATen/cuda/CUDAContext.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <math.h>
-#include <torch/all.h>
-
-#include <cassert>
-#include <cmath>
-
-#include "utils.h"
-
-template <bool flip_sin_to_cos = false, typename T_IN>
-__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<float4*>(output_batch_base_ptr + thread_offset * 4);
-      top_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
-    } else {
-      top_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
-      bottom_half = reinterpret_cast<float4*>(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<int>(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<float>(max_period)) * (-1.0f) / (static_cast<float>(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<true><<<grid, block, 0, stream>>>(
-              reinterpret_cast<scalar_t*>(t.data_ptr()),
-              reinterpret_cast<float*>(output.data_ptr()),
-              static_cast<int>(dim),
-              static_cast<float>(neg_log_max_period),
-              static_cast<float>(scale),
-              static_cast<int>(batch_size));
-        } else {
-          timestep_embedding_kernel<false><<<grid, block, 0, stream>>>(
-              reinterpret_cast<scalar_t*>(t.data_ptr()),
-              reinterpret_cast<float*>(output.data_ptr()),
-              static_cast<int>(dim),
-              static_cast<float>(neg_log_max_period),
-              static_cast<float>(scale),
-              static_cast<int>(batch_size));
-        }
-      });
-
-  return output;
-}

sgl-kernel/include/sgl_kernel_ops.h

@@ -1006,15 +1006,3 @@ std::vector<at::Tensor> fwd_kvcache_mla_fp8(
 
 std::vector<at::Tensor> 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);

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,

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
-    )

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"]

sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py

@@ -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__])