[AMD][Diffusion] support timestep embedding kernel for AMD GPUs (#16766)

2026-01-12 22:17:07 -08:00
parent ff3ddb9d9b
commit 8716589826
9 changed files with 40 additions and 19 deletions
@@ -219,6 +219,18 @@ 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);
 }

 REGISTER_EXTENSION(common_ops)