Support moe topk sigmoid kernel (#13049)

Co-authored-by: xuebi <xuebi@minimaxi.com>
2025-11-20 00:24:37 +08:00
parent 196b940aed
commit e72cf13693
10 changed files with 992 additions and 0 deletions
@@ -230,6 +230,11 @@ TORCH_LIBRARY_FRAGMENT(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);
+
  m.def("moe_sum_reduce(Tensor input, Tensor output, float routed_scaling_factor) -> ()");
  m.impl("moe_sum_reduce", torch::kCUDA, &moe_sum_reduce);