Add moe topk softmax templated from vllm (#4302)

2025-03-14 12:03:33 -07:00
parent 660305c38a
commit 61e4433caf
9 changed files with 716 additions and 6 deletions
--- a/sgl-kernel/python/sgl_kernel/init.py
+++ b/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,