diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 86a7bfbeb..3d0c9079a 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -13,7 +13,6 @@ from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, QuantizationConfig, ) -from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer from sglang.srt.layers.quantization.modelslim.schemes import ( ModelSlimW4A4Int4, ModelSlimW4A4Int4MoE, @@ -142,12 +141,7 @@ class ModelSlimConfig(QuantizationConfig): from sglang.srt.layers.moe.fused_moe_triton import FusedMoE if isinstance(layer, LinearBase): - if should_ignore_layer( - prefix, - ignore=self.ignore, - fused_mapping=self.packed_modules_mapping, - ): - return UnquantizedLinearMethod() + # TODO: we should remove this code and switch to the packed_modules_mapping declared inside the modeling files key = "model" if "vision_model" in prefix: key = "vision_model" @@ -163,57 +157,49 @@ class ModelSlimConfig(QuantizationConfig): prefix_in_quant_config = prefix.replace( proj_name, packed_modules_mapping_subset[proj_name][0] ) - - if self.is_layer_skipped(prefix, packed_modules_mapping_subset): + if self.is_layer_skipped( + prefix, packed_modules_mapping_subset + ) or self.is_layer_skipped(prefix, self.packed_modules_mapping): return UnquantizedLinearMethod() - scheme = self.get_linear_scheme( - layer=layer, layer_name=prefix_in_quant_config - ) - layer.scheme = scheme + layer.scheme = self.get_linear_scheme(layer, prefix_in_quant_config) return ModelSlimLinearMethod(self) elif isinstance(layer, FusedMoE): layer.scheme = self.get_moe_scheme(layer, prefix) return ModelSlimFusedMoEMethod(self) return None - def _get_scheme_from_parts( - self, - layer_name: str, - ) -> ModelSlimLinearScheme: - - quant_type = self.quant_description.get(layer_name + ".weight", "") - if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8": - return ModelSlimW8A8Int8( - quant_config=self.quant_description, prefix=layer_name - ) - elif quant_type == "W4A4_DYNAMIC": - return ModelSlimW4A4Int4( - quant_config=self.quant_description, prefix=layer_name - ) - raise NotImplementedError("No modelslim compatible scheme was found.") - def get_linear_scheme( - self, layer: torch.nn.Module, layer_name: Optional[str] = None + self, layer: torch.nn.Module, prefix: Optional[str] = None ) -> Optional[ModelSlimLinearScheme]: """ get_scheme method adjusted for modelslim, taken from python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py """ - scheme = self._get_scheme_from_parts( - layer_name=layer_name, - ) - # Ascend doesn't support device capability - logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name) - return scheme + linear_quant_schemes = [ + ("W4A4_DYNAMIC", ModelSlimW4A4Int4), + ("W8A8", ModelSlimW8A8Int8), + ("W8A8_DYNAMIC", ModelSlimW8A8Int8), + ] + + quant_schemes = [self.quant_description.get(prefix + ".weight", "")] + + for scheme_name, scheme_class in linear_quant_schemes: + if any(s == scheme_name for s in quant_schemes): + logger.info_once(f"Using {scheme_class.__name__}") + return scheme_class(quant_config=self.quant_description, prefix=prefix) + + logger.warning( + f"Unsupported Linear modelslim scheme: " + f"{quant_schemes} in layer: {prefix}" + ) + return None def get_moe_scheme( self, layer: torch.nn.Module, prefix: str, ) -> 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), @@ -222,7 +208,7 @@ class ModelSlimConfig(QuantizationConfig): moe_weight_suffixes = [".0.gate_proj.weight", ".0.w2.weight"] quant_schemes = [ - self.quant_description.get(prefix + suffix, "STATIC") + self.quant_description.get(prefix + suffix, "") for suffix in moe_weight_suffixes ] diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index b0884c681..e75e948b2 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -197,6 +197,7 @@ def _get_quantization_config( model_class, _ = get_model_architecture(model_config) packed_modules_mapping = getattr(model_class, "packed_modules_mapping", {}) remap_prefix = getattr(model_class, "remap_prefix", None) + # TODO: we should remove this code and switch to the packed_modules_mapping declared inside the modeling files if _is_npu: packed_modules_mapping.update( { diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2f430c2b9..53ed08a4c 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -881,7 +881,7 @@ ALL_DECODER_LAYER_TYPES = { class Qwen3_5ForCausalLM(nn.Module): """Qwen3.5 Model with support for dense variant.""" - if _is_gfx95: + if _is_gfx95 or _is_npu: packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], @@ -1310,7 +1310,7 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM): class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration): - if _is_gfx95: + if _is_gfx95 or _is_npu: packed_modules_mapping = Qwen3_5ForCausalLM.packed_modules_mapping hf_to_sglang_mapper = None @@ -1447,7 +1447,7 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration): class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): """Qwen3.5 MoE Vision-Language Model.""" - if _is_gfx95: + if _is_gfx95 or _is_npu: packed_modules_mapping = Qwen3_5ForCausalLM.packed_modules_mapping hf_to_sglang_mapper = None