[NPU][ModelSlim] adapt w2 quant layer for Minimax2.5 (#20905)

python/sglang/srt/layers/quantization/modelslim/modelslim.py

@@ -214,34 +214,26 @@ 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),
+        ]
+
+        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)
 
-        prefix_in_quant_config = prefix + ".0.gate_proj.weight"
-        is_moe_w4a4_dynamic = (
-            self.quant_description.get(prefix_in_quant_config, "STATIC")
-            == "W4A4_DYNAMIC"
-        )
-        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"{quant_schemes} in layer: {prefix}"
-                f"in layer: {prefix}"
         )
         return None
 

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(