[diffusion] kernel: timestep embedding kernel implementation (#12995)
Co-authored-by: 戚余航 <qiyuhang@bytedance.com> Co-authored-by: Qi Yuhang <45795032+HydraQYH@users.noreply.github.com>
This commit is contained in:
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user