Add moe topk softmax templated from vllm (#4302)

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
+    )