[Kernel Slimming] Migrate marlin moe kernel to JIT (#19181)

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py

@@ -0,0 +1,251 @@
+import os
+
+import torch
+import triton
+import triton.testing
+from sgl_kernel.scalar_type import scalar_types
+
+from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm as jit_fn
+from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
+from sglang.test.test_marlin_utils import marlin_quantize
+
+try:
+    from sgl_kernel import moe_wna16_marlin_gemm as _aot_import  # noqa: F401
+
+    AOT_AVAILABLE = True
+except (ImportError, AttributeError):
+    AOT_AVAILABLE = False
+
+IS_CI = (
+    os.getenv("CI", "false").lower() == "true"
+    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
+)
+
+
+def stack_and_dev(tensors):
+    dev = tensors[0].device
+    return torch.stack(tensors, dim=0).to(dev)
+
+
+# Fixed problem dimensions
+E = 8
+SIZE_K = 4096
+SIZE_N = 4096
+GROUP_SIZE = 128
+TOPK = 2
+QUANT_TYPE = scalar_types.uint4b8
+DTYPE = torch.float16
+BLOCK_SIZE_M = 64
+
+# Quantize weights once (per-expert)
+torch.manual_seed(0)
+_qweight_l, _scales_l, _w_ref_l = [], [], []
+for i in range(E):
+    _w = torch.randn((SIZE_N, SIZE_K), dtype=DTYPE, device="cuda") / 20
+    _perm = torch.randperm(SIZE_K)
+    _w_ref, _qw, _s, _, _, _ = marlin_quantize(_w, QUANT_TYPE, GROUP_SIZE, False, _perm)
+    _w_ref_l.append(_w_ref.T)
+    _qweight_l.append(_qw)
+    _scales_l.append(_s)
+
+_qweight = stack_and_dev(_qweight_l).contiguous()
+_scales = stack_and_dev(_scales_l)
+
+_sms = torch.cuda.get_device_properties("cuda").multi_processor_count
+
+
+def _make_inputs(size_m):
+    a = torch.randn((size_m, SIZE_K), dtype=DTYPE, device="cuda") / 10
+    score = torch.randn((size_m, E), dtype=DTYPE, device="cuda")
+    score_softmax = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(score_softmax, TOPK)
+
+    sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+        topk_ids, BLOCK_SIZE_M, E
+    )
+
+    max_workspace_size = (SIZE_N // 64) * (sorted_token_ids.size(0) // BLOCK_SIZE_M)
+    max_workspace_size = min(max_workspace_size, _sms * 4)
+    workspace = torch.zeros(max_workspace_size, dtype=torch.int, device="cuda")
+
+    c = torch.empty((size_m * TOPK, SIZE_N), dtype=DTYPE, device="cuda")
+
+    return (
+        a,
+        c,
+        topk_weights,
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        workspace,
+    )
+
+
+def _run_jit(
+    a,
+    c,
+    topk_weights,
+    sorted_token_ids,
+    expert_ids,
+    num_tokens_post_padded,
+    workspace,
+    size_m,
+):
+    return jit_fn(
+        a,
+        c,
+        _qweight,
+        None,
+        _scales,
+        None,
+        None,
+        None,
+        None,
+        workspace,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        topk_weights,
+        moe_block_size=BLOCK_SIZE_M,
+        top_k=TOPK,
+        mul_topk_weights=False,
+        is_ep=False,
+        b_q_type=QUANT_TYPE,
+        size_m=size_m,
+        size_n=SIZE_N,
+        size_k=SIZE_K,
+        is_k_full=True,
+        use_atomic_add=True,
+        use_fp32_reduce=True,
+        is_zp_float=False,
+    )
+
+
+def _run_aot(
+    a,
+    c,
+    topk_weights,
+    sorted_token_ids,
+    expert_ids,
+    num_tokens_post_padded,
+    workspace,
+    size_m,
+):
+    return torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default(
+        a,
+        c,
+        _qweight,
+        None,
+        _scales,
+        None,
+        None,
+        None,
+        None,
+        workspace,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        topk_weights,
+        moe_block_size=BLOCK_SIZE_M,
+        top_k=TOPK,
+        mul_topk_weights=False,
+        is_ep=False,
+        b_q_type_id=QUANT_TYPE.id,
+        size_m=size_m,
+        size_n=SIZE_N,
+        size_k=SIZE_K,
+        is_k_full=True,
+        use_atomic_add=True,
+        use_fp32_reduce=True,
+        is_zp_float=False,
+    )
+
+
+def check_correctness():
+    if not AOT_AVAILABLE:
+        print("sgl_kernel AOT not available, skipping correctness check")
+        return
+    size_m = 16
+    a, c, topk_weights, topk_ids, sorted_token_ids, expert_ids, ntp, workspace = (
+        _make_inputs(size_m)
+    )
+    c_jit = c.clone()
+    c_aot = c.clone()
+    _run_jit(
+        a, c_jit, topk_weights, sorted_token_ids, expert_ids, ntp, workspace, size_m
+    )
+    _run_aot(
+        a, c_aot, topk_weights, sorted_token_ids, expert_ids, ntp, workspace, size_m
+    )
+    torch.testing.assert_close(c_jit, c_aot, rtol=1e-3, atol=1e-3)
+    print("Correctness check passed (JIT vs AOT)")
+
+
+if IS_CI:
+    m_range = [1, 16, 128]
+else:
+    m_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
+
+if AOT_AVAILABLE:
+    line_vals = ["jit", "aot"]
+    line_names = ["JIT Kernel", "AOT Kernel"]
+    styles = [("blue", "-"), ("green", "-")]
+else:
+    line_vals = ["jit"]
+    line_names = ["JIT Kernel"]
+    styles = [("blue", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["size_m"],
+        x_vals=m_range,
+        line_arg="provider",
+        line_vals=line_vals,
+        line_names=line_names,
+        styles=styles,
+        ylabel="us",
+        plot_name="moe-wna16-marlin-gemm-performance",
+        args={},
+    )
+)
+def benchmark(size_m, provider):
+    a, c, topk_weights, topk_ids, sorted_token_ids, expert_ids, ntp, workspace = (
+        _make_inputs(size_m)
+    )
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "jit":
+        fn = lambda: _run_jit(
+            a,
+            c.clone(),
+            topk_weights,
+            sorted_token_ids,
+            expert_ids,
+            ntp,
+            workspace,
+            size_m,
+        )
+    elif provider == "aot":
+        fn = lambda: _run_aot(
+            a,
+            c.clone(),
+            topk_weights,
+            sorted_token_ids,
+            expert_ids,
+            ntp,
+            workspace,
+            size_m,
+        )
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    check_correctness()
+    benchmark.run(print_data=True)

python/sglang/jit_kernel/csrc/gemm/marlin_moe/kernel.h

@@ -0,0 +1,37 @@
+
+#include <sgl_kernel/scalar_type.hpp>
+
+#include "../marlin/marlin.cuh"
+#include "../marlin/marlin_dtypes.cuh"
+
+#define MARLIN_KERNEL_PARAMS                                                                                         \
+  const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp,            \
+      const int4 *__restrict__ b_bias_ptr, const int4 *__restrict__ scales_ptr,                                      \
+      const uint16_t *__restrict__ scale2_ptr, const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx,       \
+      const int32_t *__restrict__ sorted_token_ids_ptr, const int32_t *__restrict__ expert_ids_ptr,                  \
+      const int32_t *__restrict__ num_tokens_past_padded_ptr, const float *__restrict__ topk_weights_ptr, int top_k, \
+      bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, int prob_n, int prob_k, int *locks,             \
+      bool has_bias, bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem
+
+namespace device::marlin_moe {
+template <
+    typename scalar_t,                   // compute dtype, half or nv_float16
+    const host::ScalarTypeId w_type_id,  // weight ScalarType id
+    const host::ScalarTypeId s_type_id,  // weight scale ScalarType id
+    const int threads,                   // number of threads in a threadblock
+    const int thread_m_blocks,           // number of 16x16 blocks in the m
+                                         // dimension (batchsize) of the
+                                         // threadblock
+    const int thread_n_blocks,           // same for n dimension (output)
+    const int thread_k_blocks,           // same for k dimension (reduction)
+    const bool m_block_size_8,           // whether m_block_size == 8
+                                         // only works when thread_m_blocks == 1
+    const int stages,                    // number of stages for the async global->shared
+                                         // fetch pipeline
+    const int group_blocks,              // number of consecutive 16x16 blocks
+                                         // with a separate quantization scale
+    const bool is_zp_float               // is zero point of float16 type?
+    >
+__global__ void Marlin(MARLIN_KERNEL_PARAMS);
+
+}  // namespace device::marlin_moe

python/sglang/jit_kernel/moe_wna16_marlin.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from sgl_kernel.scalar_type import ScalarType
+    from tvm_ffi.module import Module
+
+# Constants matching device::marlin_moe:: in marlin.cuh
+_MAX_THREAD_N = 256
+
+
+@cache_once
+def _jit_moe_wna16_marlin_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "moe_wna16_marlin",
+        *args,
+        cuda_files=["gemm/marlin_moe/moe_wna16_marlin.cuh"],
+        cuda_wrappers=[
+            (
+                "moe_wna16_marlin_gemm",
+                f"moe_wna16_marlin_gemm<{args}>",
+            )
+        ],
+    )
+
+
+def _or_empty(
+    t: Optional[torch.Tensor], device: torch.device, dtype: torch.dtype
+) -> torch.Tensor:
+    return t if t is not None else torch.empty(0, device=device, dtype=dtype)
+
+
+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],
+    workspace: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    topk_weights: torch.Tensor,
+    moe_block_size: int,
+    top_k: int,
+    mul_topk_weights: bool,
+    is_ep: bool,
+    b_q_type: ScalarType,
+    size_m: int,
+    size_n: int,
+    size_k: int,
+    is_k_full: bool = True,
+    use_atomic_add: bool = False,
+    use_fp32_reduce: bool = False,
+    is_zp_float: bool = False,
+) -> torch.Tensor:
+    device = a.device
+
+    # Allocate output if not provided
+    if c_or_none is not None:
+        c = c_or_none
+    else:
+        c = torch.empty((size_m * top_k, size_n), dtype=a.dtype, device=device)
+
+    # Early return for zero-size M
+    if size_m == 0:
+        return c
+
+    # Determine activation ordering
+    has_act_order = (
+        g_idx_or_none is not None
+        and perm_or_none is not None
+        and g_idx_or_none.numel() > 0
+        and perm_or_none.numel() > 0
+        and g_idx_or_none.size(-1) > 0
+        and perm_or_none.size(-1) > 0
+    )
+
+    # Determine has_zp
+    has_zp = b_zeros_or_none is not None and b_zeros_or_none.numel() > 0
+
+    # Determine has_bias
+    has_bias = b_bias_or_none is not None
+
+    # Derive num_groups and group_size from b_scales
+    num_groups = b_scales.size(1)
+    if has_act_order:
+        if is_k_full:
+            group_size = size_k // num_groups
+        else:
+            group_size = 0
+    else:
+        if num_groups > 1:
+            group_size = size_k // num_groups
+        else:
+            group_size = -1
+
+    # Allocate a_tmp for act_order column permutation
+    if has_act_order:
+        a_tmp = torch.empty((size_m * top_k, size_k), dtype=a.dtype, device=device)
+    else:
+        a_tmp = torch.empty(0, dtype=a.dtype, device=device)
+
+    # Allocate c_tmp for fp32 reduce
+    if use_fp32_reduce and not use_atomic_add:
+        sms = torch.cuda.get_device_properties(device).multi_processor_count
+        # max num of threadblocks is sms * 4
+        max_c_tmp_size = min(
+            size_n * sorted_token_ids.size(0),
+            sms * 4 * moe_block_size * _MAX_THREAD_N,
+        )
+        if moe_block_size == 8:
+            max_c_tmp_size *= 2
+        c_tmp = torch.empty(max_c_tmp_size, dtype=torch.float32, device=device)
+    else:
+        c_tmp = torch.empty(0, dtype=torch.float32, device=device)
+
+    # Convert Optional tensors to empty tensors
+    g_idx_t = _or_empty(g_idx_or_none, device, torch.int32)
+    perm_t = _or_empty(perm_or_none, device, torch.int32)
+    b_zeros_t = _or_empty(b_zeros_or_none, device, a.dtype)
+    b_bias_t = _or_empty(b_bias_or_none, device, a.dtype)
+    global_scale_t = _or_empty(global_scale_or_none, device, a.dtype)
+
+    module = _jit_moe_wna16_marlin_module(a.dtype)
+    module.moe_wna16_marlin_gemm(
+        a,
+        c,
+        b_q_weight,
+        b_bias_t,
+        b_scales,
+        global_scale_t,
+        b_zeros_t,
+        g_idx_t,
+        perm_t,
+        workspace,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        topk_weights,
+        a_tmp,
+        c_tmp,
+        moe_block_size,
+        top_k,
+        mul_topk_weights,
+        is_ep,
+        b_q_type.id,
+        size_m,
+        size_n,
+        size_k,
+        has_act_order,
+        has_bias,
+        is_k_full,
+        has_zp,
+        num_groups,
+        group_size,
+        use_atomic_add,
+        use_fp32_reduce,
+        is_zp_float,
+    )
+
+    return c

python/sglang/jit_kernel/tests/test_moe_wna16_marlin.py

@@ -0,0 +1,329 @@
+import itertools
+
+import pytest
+import torch
+from sgl_kernel import moe_wna16_marlin_gemm as aot_moe_wna16_marlin_gemm
+from sgl_kernel.scalar_type import scalar_types
+
+from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
+from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
+from sglang.test.test_marlin_utils import awq_marlin_quantize, marlin_quantize
+
+
+def stack_and_dev(tensors: list[torch.Tensor]):
+    dev = tensors[0].device
+    return torch.stack(tensors, dim=0).to(dev)
+
+
+def _get_scalar_type(num_bits: int, has_zp: bool):
+    if has_zp:
+        assert num_bits == 4
+        return scalar_types.uint4
+    else:
+        return scalar_types.uint4b8 if num_bits == 4 else scalar_types.uint8b128
+
+
+def _setup_moe_weights(e, n, k, quant_type, group_size, act_order, dtype):
+    """Set up quantized MoE weights for a single gate (e experts, output n, input k)."""
+    has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
+
+    w = torch.randn((e, n, k), device="cuda", dtype=dtype) / 20
+
+    w_ref_l = []
+    qweight_l = []
+    scales_l = []
+    zeros_l = []
+    g_idx_l = []
+    sort_indices_l = []
+
+    for i in range(e):
+        if has_zp:
+            w_ref, qweight, scales, zeros = awq_marlin_quantize(
+                w[i].transpose(1, 0), quant_type, group_size
+            )
+            w_ref_l.append(w_ref.T)
+            qweight_l.append(qweight)
+            scales_l.append(scales)
+            zeros_l.append(zeros)
+        else:
+            test_perm = torch.randperm(k)
+            w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
+                w[i].transpose(1, 0), quant_type, group_size, act_order, test_perm
+            )
+            w_ref_l.append(w_ref.T)
+            qweight_l.append(qweight)
+            scales_l.append(scales)
+            g_idx_l.append(g_idx)
+            sort_indices_l.append(sort_indices)
+
+    w_ref = stack_and_dev(w_ref_l)
+    qweight = stack_and_dev(qweight_l).contiguous()
+    scales = stack_and_dev(scales_l)
+    g_idx = stack_and_dev(g_idx_l) if g_idx_l else None
+    sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None
+    zeros = stack_and_dev(zeros_l) if zeros_l else None
+
+    return w_ref, qweight, scales, zeros, g_idx, sort_indices
+
+
+def _run_single_gemm(
+    fn,
+    a,
+    c,
+    qweight,
+    scales,
+    zeros,
+    g_idx,
+    sort_indices,
+    workspace,
+    sorted_token_ids,
+    expert_ids,
+    num_tokens_post_padded,
+    topk_weights,
+    quant_type,
+    block_size_m,
+    topk,
+    size_m,
+    size_n,
+    size_k,
+    mul_topk_weights,
+    is_k_full,
+    use_atomic_add,
+):
+    return fn(
+        a,
+        c,
+        qweight,
+        None,  # b_bias
+        scales,
+        None,  # global_scale
+        zeros,
+        g_idx,
+        sort_indices,
+        workspace,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        topk_weights,
+        moe_block_size=block_size_m,
+        top_k=topk,
+        mul_topk_weights=mul_topk_weights,
+        is_ep=False,
+        b_q_type=quant_type,
+        size_m=size_m,
+        size_n=size_n,
+        size_k=size_k,
+        is_k_full=is_k_full,
+        use_atomic_add=use_atomic_add,
+        use_fp32_reduce=True,
+        is_zp_float=False,
+    )
+
+
+def _run_single_gemm_aot(
+    a,
+    c,
+    qweight,
+    scales,
+    zeros,
+    g_idx,
+    sort_indices,
+    workspace,
+    sorted_token_ids,
+    expert_ids,
+    num_tokens_post_padded,
+    topk_weights,
+    quant_type,
+    block_size_m,
+    topk,
+    size_m,
+    size_n,
+    size_k,
+    mul_topk_weights,
+    is_k_full,
+    use_atomic_add,
+):
+    return aot_moe_wna16_marlin_gemm(
+        a,
+        c,
+        qweight,
+        None,  # b_bias
+        scales,
+        None,  # global_scale
+        zeros,
+        g_idx,
+        sort_indices,
+        workspace,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        topk_weights,
+        moe_block_size=block_size_m,
+        top_k=topk,
+        mul_topk_weights=mul_topk_weights,
+        is_ep=False,
+        b_q_type_id=quant_type.id,
+        size_m=size_m,
+        size_n=size_n,
+        size_k=size_k,
+        is_k_full=is_k_full,
+        use_atomic_add=use_atomic_add,
+        use_fp32_reduce=True,
+        is_zp_float=False,
+    )
+
+
+def generate_test_cases():
+    m_list = [1, 123]
+    n_list = [128, 1024]
+    k_list = [256]
+    e_list = [4]
+    topk_list = [2]
+    dtype_list = [torch.float16, torch.bfloat16]
+    group_size_list = [128]
+    act_order_list = [False, True]
+    quant_type_list = [scalar_types.uint4, scalar_types.uint4b8]
+
+    all_combinations = itertools.product(
+        m_list,
+        n_list,
+        k_list,
+        e_list,
+        topk_list,
+        dtype_list,
+        group_size_list,
+        act_order_list,
+        quant_type_list,
+    )
+
+    def is_valid(m, n, k, e, topk, dtype, group_size, act_order, quant_type):
+        has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
+        if act_order:
+            if group_size == -1 or group_size == k:
+                return False
+            if has_zp:
+                return False
+        if group_size > 0 and k % group_size != 0:
+            return False
+        return True
+
+    return [case for case in all_combinations if is_valid(*case)]
+
+
+TEST_CASES = generate_test_cases()
+
+
+@pytest.mark.parametrize(
+    "m,n,k,e,topk,dtype,group_size,act_order,quant_type",
+    TEST_CASES,
+    ids=[
+        f"m{c[0]}_n{c[1]}_k{c[2]}_e{c[3]}_t{c[4]}_{c[5].__name__ if hasattr(c[5], '__name__') else str(c[5]).split('.')[-1]}_g{c[6]}_act{c[7]}_{c[8]}"
+        for c in TEST_CASES
+    ],
+)
+def test_moe_wna16_marlin_gemm(
+    m, n, k, e, topk, dtype, group_size, act_order, quant_type
+):
+    torch.manual_seed(0)
+
+    has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
+
+    a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
+
+    # Set up quantized weights for first gemm (gate_up: output 2*n, input k)
+    w_ref1, qweight1, scales1, zeros1, g_idx1, sort_indices1 = _setup_moe_weights(
+        e, 2 * n, k, quant_type, group_size, act_order, dtype
+    )
+
+    # Compute block_size_m
+    for block_size_m in [8, 16, 32, 48, 64]:
+        if m * topk / e / block_size_m < 0.9:
+            break
+
+    # Align tokens
+    score = torch.randn((m, e), device="cuda", dtype=dtype)
+    score_softmax = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(score_softmax, topk)
+
+    sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+        topk_ids, block_size_m, e
+    )
+
+    # Workspace
+    sms = torch.cuda.get_device_properties("cuda").multi_processor_count
+    max_workspace_size = (max(2 * n, k) // 64) * (
+        sorted_token_ids.size(0) // block_size_m
+    )
+    max_workspace_size = min(max_workspace_size, sms * 4)
+    workspace = torch.zeros(
+        max_workspace_size, dtype=torch.int, device="cuda", requires_grad=False
+    )
+
+    use_atomic_add = (
+        dtype == torch.half or torch.cuda.get_device_capability("cuda")[0] >= 9
+    )
+
+    scalar_type = _get_scalar_type(4, has_zp)
+
+    # --- Run JIT kernel ---
+    c_jit = torch.empty((m * topk, 2 * n), dtype=dtype, device="cuda")
+    c_jit = _run_single_gemm(
+        moe_wna16_marlin_gemm,
+        a,
+        c_jit,
+        qweight1,
+        scales1,
+        zeros1,
+        g_idx1,
+        sort_indices1,
+        workspace,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        topk_weights,
+        scalar_type,
+        block_size_m,
+        topk,
+        m,
+        2 * n,
+        k,
+        False,
+        True,
+        use_atomic_add,
+    )
+
+    torch.cuda.synchronize()
+
+    # --- Check bitwise equality with AOT kernel ---
+    c_aot = torch.empty((m * topk, 2 * n), dtype=dtype, device="cuda")
+    c_aot = _run_single_gemm_aot(
+        a,
+        c_aot,
+        qweight1,
+        scales1,
+        zeros1,
+        g_idx1,
+        sort_indices1,
+        workspace,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        topk_weights,
+        scalar_type,
+        block_size_m,
+        topk,
+        m,
+        2 * n,
+        k,
+        False,
+        True,
+        use_atomic_add,
+    )
+    torch.cuda.synchronize()
+    torch.testing.assert_close(c_jit, c_aot, rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    import subprocess
+
+    subprocess.call(["pytest", "--tb=short", "-v", str(__file__)])

python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py

@@ -10,6 +10,8 @@ _is_cuda = is_cuda()
 if _is_cuda:
     from sgl_kernel import moe_sum_reduce, silu_and_mul
 
+    from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
+
 
 def get_scalar_type(num_bits: int, has_zp: bool):
     from sgl_kernel.scalar_type import scalar_types
@@ -142,7 +144,7 @@ def fused_marlin_moe(
         or torch.cuda.get_device_capability(hidden_states.device)[0] >= 9
     )
 
-    intermediate_cache1 = torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default(
+    intermediate_cache1 = moe_wna16_marlin_gemm(
         hidden_states,
         intermediate_cache1,
         w1,
@@ -161,7 +163,7 @@ def fused_marlin_moe(
         top_k=topk,
         mul_topk_weights=False,
         is_ep=expert_map is not None,
-        b_q_type_id=scalar_type1.id,
+        b_q_type=scalar_type1,
         size_m=M,
         size_n=2 * N,
         size_k=K,
@@ -176,7 +178,7 @@ def fused_marlin_moe(
     if expert_map is not None:
         intermediate_cache3.zero_()
 
-    intermediate_cache3 = torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default(
+    intermediate_cache3 = moe_wna16_marlin_gemm(
         intermediate_cache2,
         intermediate_cache3,
         w2,
@@ -195,7 +197,7 @@ def fused_marlin_moe(
         top_k=1,
         mul_topk_weights=True,
         is_ep=expert_map is not None,
-        b_q_type_id=scalar_type2.id,
+        b_q_type=scalar_type2,
         size_m=M * topk,
         size_n=K,
         size_k=N,