chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-01-28 14:18:53 +08:00
committed by GitHub
parent 67fb492c9a
commit fb74e43707
12 changed files with 0 additions and 326 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.github/workflows/pr-test-amd.yml vendored
View File

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

1
3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt vendored
View File

@@ -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(

1
3rdparty/amd/sgl-kernel/rocm_hipify.py vendored
View File

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

1
sgl-kernel/CMakeLists.txt
View File

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

13
sgl-kernel/csrc/common_extension.cc
View File

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

12
sgl-kernel/csrc/common_extension_rocm.cc
View File

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

137
sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu
View File

@@ -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;
}

12
sgl-kernel/include/sgl_kernel_ops.h
View File

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

1
sgl-kernel/python/sgl_kernel/__init__.py
View File

@@ -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,

32
sgl-kernel/python/sgl_kernel/elementwise.py
View File

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

1
sgl-kernel/setup_rocm.py
View File

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

114
sgl-kernel/tests/sgl_diffusion/test_timestep_embedding.py
View File

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