[EP] Add cuda kernel for moe_ep_pre_reorder (#6699)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
@@ -240,6 +240,17 @@ void prepare_moe_input(
|
||||
const int64_t n,
|
||||
const int64_t k);
|
||||
|
||||
void ep_moe_pre_reorder(
|
||||
torch::Tensor input,
|
||||
torch::Tensor gateup_input,
|
||||
torch::Tensor src2dst,
|
||||
torch::Tensor topk_ids,
|
||||
torch::Tensor a1_scales,
|
||||
int64_t start_expert_id,
|
||||
int64_t end_expert_id,
|
||||
int64_t topk,
|
||||
bool use_per_token_if_dynamic);
|
||||
|
||||
/*
|
||||
* From csrc/speculative
|
||||
*/
|
||||
|
||||
Reference in New Issue
Block a user