[diffusion] kernel: timestep embedding kernel implementation (#12995)

Co-authored-by: 戚余航 <qiyuhang@bytedance.com> Co-authored-by: Qi Yuhang <45795032+HydraQYH@users.noreply.github.com>
2025-12-19 20:59:50 +08:00
parent 1c65802648
commit 46be74b4b4
8 changed files with 369 additions and 0 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)