[Diffusion] Move diffusion time embedding to jit kernel (#16879)

2026-01-17 12:21:22 +08:00
parent a7b5f75d88
commit 2cdd4370bc
10 changed files with 228 additions and 206 deletions
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/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);