diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 59ca3f9cc..640e2b504 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -70,6 +70,7 @@ from sglang.srt.utils import ( get_bool_env_var, is_cpu, is_hip, + is_npu, print_info_once, round_up, ) @@ -78,6 +79,7 @@ from sglang.srt.utils.custom_op import register_custom_op _is_hip = is_hip() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() +_is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip @@ -778,6 +780,13 @@ class FusedMoE(torch.nn.Module): # expert weights into block layout. During weight update, we must restore # canonical load-time shapes before copying checkpoint tensors. if isinstance(method, UnquantizedFusedMoEMethod): + if _is_npu: + if weight_name.endswith(".experts.w2_weight"): + if param.data.shape[1] != loaded_weight.shape[0]: + param.data = param.data.transpose(1, 2).contiguous() + if weight_name.endswith(".experts.w13_weight"): + if param.data.shape[2] != loaded_weight.shape[1]: + param.data = param.data.transpose(1, 2).contiguous() method.maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load( layer=self, param=param,