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
@@ -1006,3 +1006,15 @@ 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);