[EP] Add cuda kernel for moe_ep_pre_reorder (#6699)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2025-06-02 11:49:01 +08:00
committed by GitHub
parent 20fd53b8f6
commit 55444ed667
7 changed files with 230 additions and 0 deletions

View File

@@ -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
*/