diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 44fe6d45c..58f06c980 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -34,7 +34,7 @@ from sgl_kernel.elementwise import ( silu_and_mul, ) from sgl_kernel.expert_specialization import es_fp8_blockwise_scaled_grouped_mm -from sgl_kernel.fused_moe import fused_marlin_moe +from sgl_kernel.fused_moe import fused_marlin_moe, moe_wna16_marlin_gemm from sgl_kernel.gemm import ( awq_dequantize, bmm_fp8, diff --git a/sgl-kernel/python/sgl_kernel/fused_moe.py b/sgl-kernel/python/sgl_kernel/fused_moe.py index cc2e00fd1..8e0cea934 100644 --- a/sgl-kernel/python/sgl_kernel/fused_moe.py +++ b/sgl-kernel/python/sgl_kernel/fused_moe.py @@ -15,6 +15,60 @@ def get_scalar_type(num_bits: int, has_zp: bool): return scalar_types.uint4b8 if num_bits == 4 else scalar_types.uint8b128 +def moe_wna16_marlin_gemm( + a: torch.Tensor, + c_or_none: Optional[torch.Tensor], + b_q_weight: torch.Tensor, + b_scales: 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, +): + return torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default( + a, + c_or_none, + b_q_weight, + b_scales, + b_zeros_or_none, + g_idx_or_none, + perm_or_none, + workspace, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_weights, + moe_block_size=moe_block_size, + top_k=top_k, + mul_topk_weights=mul_topk_weights, + is_ep=is_ep, + b_q_type_id=b_q_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=use_fp32_reduce, + is_zp_float=is_zp_float, + ) + + def fused_marlin_moe( hidden_states: torch.Tensor, w1: torch.Tensor,