Revert "[Diffusion] Move diffusion time embedding to jit kernel" (#17257)

2026-01-17 21:29:22 +08:00
parent 9c2530642c
commit 53609e5e5b
10 changed files with 206 additions and 228 deletions
@@ -609,6 +609,19 @@ 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)