Add moe topk softmax templated from vllm (#4302)

sgl-kernel/python/sgl_kernel/__init__.py

@@ -33,7 +33,7 @@ from sgl_kernel.gemm import (
     sgl_per_token_group_quant_fp8,
     sgl_per_token_quant_fp8,
 )
-from sgl_kernel.moe import moe_align_block_size
+from sgl_kernel.moe import moe_align_block_size, topk_softmax
 from sgl_kernel.sampling import (
     min_p_sampling_from_probs,
     top_k_renorm_prob,

sgl-kernel/python/sgl_kernel/moe.py

@@ -21,3 +21,14 @@ def moe_align_block_size(
         token_cnts_buffer,
         cumsum_buffer,
     )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indices: torch.Tensor,
+    gating_output: float,
+) -> None:
+    torch.ops.sgl_kernel.topk_softmax(
+        topk_weights, topk_ids, token_expert_indices, gating_output
+    )