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