[opt kimi k2 1 / n] Add kimi k2 moe fused gate (#13287)

2025-11-15 17:14:19 +08:00
parent 8e9f05ece1
commit 1d3d42bda0
8 changed files with 646 additions and 0 deletions
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -331,6 +331,14 @@ std::vector<at::Tensor> moe_fused_gate(
    double routed_scaling_factor,
    bool apply_routed_scaling_factor_on_output);

+std::vector<at::Tensor> kimi_k2_moe_fused_gate(
+    at::Tensor& input,
+    at::Tensor& bias,
+    int64_t topk,
+    bool renormalize,
+    double routed_scaling_factor,
+    bool apply_routed_scaling_factor_on_output);
+
 void fp8_blockwise_scaled_grouped_mm(
    torch::Tensor& output,
    torch::Tensor& a_ptrs,