Support moe topk sigmoid kernel (#13049)

Co-authored-by: xuebi <xuebi@minimaxi.com>
This commit is contained in:
Roger Young
2025-11-20 00:24:37 +08:00
committed by GitHub
co-authored by xuebi
parent 196b940aed
commit e72cf13693
10 changed files with 992 additions and 0 deletions
+5
View File
@@ -100,6 +100,11 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
"moe_softcapping, Tensor? correction_bias) -> ()");
m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
m.def(
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, Tensor? "
"correction_bias) -> ()");
m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
/*
* From csrc/speculative
*/