Add new moe wna16 marlin gemm (#14122)

2025-12-01 23:07:53 +08:00
parent 2e8f54e61e
commit 3de09aadbc
15 changed files with 823 additions and 876 deletions
--- a/sgl-kernel/python/sgl_kernel/fused_moe.py
+++ b/sgl-kernel/python/sgl_kernel/fused_moe.py
@@ -7,7 +7,9 @@ def moe_wna16_marlin_gemm(
    a: torch.Tensor,
    c_or_none: Optional[torch.Tensor],
    b_q_weight: torch.Tensor,
+    b_bias_or_none: Optional[torch.Tensor],
    b_scales: torch.Tensor,
+    global_scale_or_none: Optional[torch.Tensor],
    b_zeros_or_none: Optional[torch.Tensor],
    g_idx_or_none: Optional[torch.Tensor],
    perm_or_none: Optional[torch.Tensor],
@@ -33,7 +35,9 @@ def moe_wna16_marlin_gemm(
        a,
        c_or_none,
        b_q_weight,
+        b_bias_or_none,
        b_scales,
+        global_scale_or_none,
        b_zeros_or_none,
        g_idx_or_none,
        perm_or_none,