[2/2] Add python wrapper for CUTLASS FP8 Blockscale MoE Kernel. (#5694)

2025-05-16 13:14:07 -07:00
parent 839fb31e5f
commit 6fc9357503
12 changed files with 896 additions and 41 deletions
--- a/sgl-kernel/python/sgl_kernel/init.py
+++ b/sgl-kernel/python/sgl_kernel/init.py
@@ -47,6 +47,7 @@ from sgl_kernel.moe import (
    fp8_blockwise_scaled_grouped_mm,
    moe_align_block_size,
    moe_fused_gate,
+    prepare_moe_input,
    topk_softmax,
 )
 from sgl_kernel.sampling import (