diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index e282eb936..84acecccc 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -214,36 +214,28 @@ class ModelSlimConfig(QuantizationConfig): ) -> Optional[ModelSlimMoEScheme]: # TODO: @dsikka: refactor this to use schemes as other kernels # are supported + check if the layer is being ignored. + moe_quant_schemes = [ + ("W4A4_DYNAMIC", ModelSlimW4A4Int4MoE), + ("W4A8_DYNAMIC", ModelSlimW4A8Int8MoE), + ("W8A8_DYNAMIC", ModelSlimW8A8Int8MoE), + ] - prefix_in_quant_config = prefix + ".0.gate_proj.weight" - is_moe_w4a4_dynamic = ( - self.quant_description.get(prefix_in_quant_config, "STATIC") - == "W4A4_DYNAMIC" + moe_weight_suffixes = [".0.gate_proj.weight", ".0.w2.weight"] + quant_schemes = [ + self.quant_description.get(prefix + suffix, "STATIC") + for suffix in moe_weight_suffixes + ] + + for scheme_name, scheme_class in moe_quant_schemes: + if any(s == scheme_name for s in quant_schemes): + logger.info_once(f"Using {scheme_class.__name__}") + return scheme_class(self) + + logger.warning( + f"Unsupported FusedMoe modelslim scheme: " + f"{quant_schemes} in layer: {prefix}" ) - is_moe_w4a8_dynamic = ( - self.quant_description.get(prefix_in_quant_config, "STATIC") - == "W4A8_DYNAMIC" - ) - is_moe_w8a8_dynamic = ( - self.quant_description.get(prefix_in_quant_config, "STATIC") - == "W8A8_DYNAMIC" - ) - if is_moe_w4a4_dynamic: - logger.info_once("Using ModelSlimW4A4Int4MoE") - return ModelSlimW4A4Int4MoE(self) - elif is_moe_w4a8_dynamic: - logger.info_once("Using ModelSlimW4A8Int8MoE") - return ModelSlimW4A8Int8MoE(self) - elif is_moe_w8a8_dynamic: - logger.info_once("Using ModelSlimW8A8Int8MoE") - return ModelSlimW8A8Int8MoE(self) - else: - logger.warning( - f"Unsupported FusedMoe modelslim scheme: " - f"{self.quant_description.get(prefix_in_quant_config.strip())} " - f"in layer: {prefix}" - ) - return None + return None def is_layer_skipped( self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({}) diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 11929f740..470b1e0f7 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -713,7 +713,7 @@ class MiniMaxM2DecoderLayer(nn.Module): config=config, layer_id=layer_id, quant_config=quant_config, - prefix=add_prefix("mlp", prefix), + prefix=add_prefix("block_sparse_moe", prefix), ) self.input_layernorm = RMSNorm(