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
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -317,6 +317,13 @@ void topk_softmax(
    double moe_softcapping,
    const c10::optional<torch::Tensor>& correction_bias);

+void topk_sigmoid(
+    torch::Tensor& topk_weights,
+    torch::Tensor& topk_indices,
+    torch::Tensor& gating_output,
+    bool renormalize,
+    const c10::optional<torch::Tensor>& correction_bias);
+
 void moe_sum_reduce(at::Tensor& input, at::Tensor& output, double routed_scaling_factor);

 void moe_sum(torch::Tensor& input, torch::Tensor& output);