Implementation:Sgl project Sglang Fused MoE Marlin

Knowledge Sources	Sgl_project_Sglang
Domains	Kernel, Mixture of Experts, Quantization
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for the fused MoE Marlin WNA16 (weight-only N-bit, activation 16-bit) GEMM kernel.

Description

The fused_moe.py module wraps the C++ Marlin MoE kernel via moe_wna16_marlin_gemm. This function combines Marlin's optimized quantized GEMM with MoE token dispatch in a single fused operation. It takes quantized weights (b_q_weight), per-group or per-channel scales (b_scales), optional zero-points (b_zeros_or_none), optional group indices (g_idx_or_none), and MoE routing information including sorted_token_ids, expert_ids, num_tokens_post_padded, and topk_weights. Configuration flags control expert parallelism (is_ep), atomic add reduction (use_atomic_add), FP32 accumulation (use_fp32_reduce), and floating-point zero-points (is_zp_float). The function delegates to torch.ops.sgl_kernel.moe_wna16_marlin_gemm.

Usage

Use this function for high-performance quantized MoE inference where expert weights are stored in Marlin WNA16 format (e.g., 4-bit GPTQ or AWQ quantized models with MoE layers).

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/fused_moe.py
• Lines: 1-62

Signature

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_id: int,
    size_m: int,
    size_n: int,
    size_k: int,
    is_k_full: bool,
    use_atomic_add: bool,
    use_fp32_reduce: bool,
    is_zp_float: bool,
): ...

Import

from sgl_kernel.fused_moe import moe_wna16_marlin_gemm

I/O Contract

Inputs

Name	Type	Required	Description
a	torch.Tensor	Yes	Activation input tensor
c_or_none	Optional[torch.Tensor]	No	Pre-allocated output tensor, or None
b_q_weight	torch.Tensor	Yes	Quantized weight tensor in Marlin format
b_bias_or_none	Optional[torch.Tensor]	No	Bias tensor, or None
b_scales	torch.Tensor	Yes	Weight quantization scale factors
global_scale_or_none	Optional[torch.Tensor]	No	Global scale factor, or None
b_zeros_or_none	Optional[torch.Tensor]	No	Weight zero-point values, or None
g_idx_or_none	Optional[torch.Tensor]	No	Group index mapping for GPTQ, or None
perm_or_none	Optional[torch.Tensor]	No	Permutation tensor, or None
workspace	torch.Tensor	Yes	Workspace buffer for Marlin GEMM
sorted_token_ids	torch.Tensor	Yes	Token IDs sorted by expert assignment
expert_ids	torch.Tensor	Yes	Expert IDs for each token group
num_tokens_post_padded	torch.Tensor	Yes	Token count after padding per expert
topk_weights	torch.Tensor	Yes	Top-k routing weights from MoE gate
moe_block_size	int	Yes	Block size for MoE token grouping
top_k	int	Yes	Number of experts per token
mul_topk_weights	bool	Yes	Whether to multiply output by topk weights
is_ep	bool	Yes	Whether expert parallelism is used
b_q_type_id	int	Yes	Quantization type identifier
size_m	int	Yes	M dimension of the GEMM
size_n	int	Yes	N dimension (output features)
size_k	int	Yes	K dimension (input features)
is_k_full	bool	Yes	Whether K covers full input dimension
use_atomic_add	bool	Yes	Whether to use atomic add for output reduction
use_fp32_reduce	bool	Yes	Whether to accumulate in FP32
is_zp_float	bool	Yes	Whether zero-points are floating-point

Outputs

Name	Type	Description
result	torch.Tensor	GEMM output tensor with MoE-routed results

Usage Examples

from sgl_kernel.fused_moe import moe_wna16_marlin_gemm

result = moe_wna16_marlin_gemm(
    a=activations,
    c_or_none=None,
    b_q_weight=quantized_weights,
    b_bias_or_none=None,
    b_scales=weight_scales,
    global_scale_or_none=None,
    b_zeros_or_none=weight_zeros,
    g_idx_or_none=group_idx,
    perm_or_none=perm,
    workspace=workspace,
    sorted_token_ids=sorted_ids,
    expert_ids=expert_ids,
    num_tokens_post_padded=num_tokens,
    topk_weights=topk_w,
    moe_block_size=128,
    top_k=2,
    mul_topk_weights=True,
    is_ep=False,
    b_q_type_id=4,
    size_m=num_tokens_total,
    size_n=hidden_dim,
    size_k=input_dim,
    is_k_full=True,
    use_atomic_add=False,
    use_fp32_reduce=True,
    is_zp_float=False,
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime