[NPU] [Bugfix] Wan quantization fix (#24540)

Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Артем Савкин
2026-05-11 08:32:57 +03:00
committed by GitHub
co-authored by ronnie_zheng
parent 5207f074a4
commit 9ec2880eca
5 changed files with 206 additions and 188 deletions
@@ -21,11 +21,12 @@ from sglang.srt.layers.quantization.modelslim.schemes import (
) )
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
from sglang.srt.layers.quantization.modelslim.schemes import ( from sglang.srt.layers.quantization.modelslim.schemes import (
ModelSlimLinearScheme, ModelSlimLinearScheme,
) )
from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -40,7 +41,11 @@ class ModelSlimConfig(QuantizationConfig):
- W8A8 dynamic linear - W8A8 dynamic linear
""" """
def __init__(self, quant_config: Dict[str, Any] = {}): def __init__(
self,
quant_config: Dict[str, Any] = {},
reverse_param_names_mapping: dict = None,
):
super().__init__() super().__init__()
self.quant_description = quant_config self.quant_description = quant_config
ignore = cast(List[str], quant_config.get("ignore", [])) ignore = cast(List[str], quant_config.get("ignore", []))
@@ -49,6 +54,11 @@ class ModelSlimConfig(QuantizationConfig):
self.packed_modules_mapping = ( self.packed_modules_mapping = (
packed_modules_mapping if packed_modules_mapping is not None else {} packed_modules_mapping if packed_modules_mapping is not None else {}
) )
self._name_mapper = (
get_param_names_mapping(reverse_param_names_mapping)
if reverse_param_names_mapping is not None
else None
)
def get_linear_method(self) -> ModelSlimLinearMethod: def get_linear_method(self) -> ModelSlimLinearMethod:
return ModelSlimLinearMethod(self) return ModelSlimLinearMethod(self)
@@ -71,8 +81,10 @@ class ModelSlimConfig(QuantizationConfig):
return filenames return filenames
@classmethod @classmethod
def from_config(cls, config: Dict[str, Any]) -> ModelSlimConfig: def from_config(
return cls(config) cls, config: Dict[str, Any], reverse_param_names_mapping: dict = None
) -> ModelSlimConfig:
return cls(config, reverse_param_names_mapping)
def get_quant_method( def get_quant_method(
self, self,
@@ -109,16 +121,18 @@ class ModelSlimConfig(QuantizationConfig):
self, self,
layer_name: str, layer_name: str,
) -> ModelSlimLinearScheme: ) -> ModelSlimLinearScheme:
full_weight_name = layer_name + ".weight"
if self._name_mapper is not None:
mapped_name, _, _ = self._name_mapper(full_weight_name)
else:
mapped_name = full_weight_name
quant_type = self.quant_description.get(layer_name + ".weight", "") quant_type = self.quant_description.get(mapped_name, "")
prefix = mapped_name.removesuffix(".weight")
if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8": if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8":
return ModelSlimW8A8Int8( return ModelSlimW8A8Int8(quant_config=self.quant_description, prefix=prefix)
quant_config=self.quant_description, prefix=layer_name
)
elif quant_type == "W4A4_DYNAMIC": elif quant_type == "W4A4_DYNAMIC":
return ModelSlimW4A4Int4( return ModelSlimW4A4Int4(quant_config=self.quant_description, prefix=prefix)
quant_config=self.quant_description, prefix=layer_name
)
elif quant_type == "W8A8_MXFP8": elif quant_type == "W8A8_MXFP8":
from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp8_scheme import ( from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp8_scheme import (
ModelSlimMXFP8Scheme, ModelSlimMXFP8Scheme,
@@ -493,8 +493,11 @@ def _resolve_quant_config(
reverse_param_names_mapping_dict = getattr( reverse_param_names_mapping_dict = getattr(
arch_config, "reverse_param_names_mapping", None arch_config, "reverse_param_names_mapping", None
) )
quant_config = get_quant_config(
quant_config = get_quant_config(hf_config, component_model_path) hf_config,
component_model_path,
reverse_param_names_mapping=reverse_param_names_mapping_dict,
)
quant_config_name = _get_quant_config_name(quant_config) quant_config_name = _get_quant_config_name(quant_config)
inferred_nvfp4_config = None inferred_nvfp4_config = None
if quant_config is None or quant_config_name == "modelopt_fp4": if quant_config is None or quant_config_name == "modelopt_fp4":
@@ -127,12 +127,13 @@ def get_quant_config(
model_config, model_config,
component_model_path: str, component_model_path: str,
packed_modules_mapping: Dict[str, List[str]] = {}, packed_modules_mapping: Dict[str, List[str]] = {},
reverse_param_names_mapping: Dict[str, List[str]] = {},
remap_prefix: Dict[str, str] | None = None, remap_prefix: Dict[str, str] | None = None,
) -> QuantizationConfig: ) -> QuantizationConfig:
quant_cfg = find_quant_modelslim_config(model_config, component_model_path) quant_cfg = find_quant_modelslim_config(model_config, component_model_path)
if quant_cfg is not None: if quant_cfg is not None:
quant_cls = _load_quant_cls(quant_cfg) quant_cls = _load_quant_cls(quant_cfg)
return quant_cls.from_config(quant_cfg) return quant_cls.from_config(quant_cfg, reverse_param_names_mapping)
if "quantization_config" not in model_config: if "quantization_config" not in model_config:
return None return None
@@ -7,68 +7,68 @@
"scenarios": { "scenarios": {
"flux_image_t2i_npu": { "flux_image_t2i_npu": {
"stages_ms": { "stages_ms": {
"InputValidationStage": 0.07, "InputValidationStage": 0.06,
"TextEncodingStage": 154.51, "TextEncodingStage": 468.6,
"TimestepPreparationStage": 53.52, "TimestepPreparationStage": 41.44,
"LatentPreparationStage": 0.39, "LatentPreparationStage": 0.28,
"DenoisingStage": 19423.39, "DenoisingStage": 19973.94,
"DecodingStage": 196.62 "DecodingStage": 8.58
}, },
"denoise_step_ms": { "denoise_step_ms": {
"0": 123.16, "0": 100.15,
"1": 91.7, "1": 95.5,
"2": 265.62, "2": 334.93,
"3": 402.68, "3": 413.4,
"4": 402.86, "4": 413.43,
"5": 402.78, "5": 413.33,
"6": 402.99, "6": 413.6,
"7": 402.77, "7": 413.2,
"8": 402.59, "8": 413.4,
"9": 402.93, "9": 413.15,
"10": 402.05, "10": 413.24,
"11": 402.99, "11": 413.37,
"12": 402.29, "12": 413.63,
"13": 403.07, "13": 413.39,
"14": 402.62, "14": 413.49,
"15": 402.99, "15": 413.69,
"16": 402.68, "16": 413.61,
"17": 403.0, "17": 413.38,
"18": 402.74, "18": 413.32,
"19": 402.85, "19": 413.3,
"20": 402.83, "20": 413.45,
"21": 403.03, "21": 413.49,
"22": 402.56, "22": 413.31,
"23": 402.84, "23": 413.14,
"24": 402.79, "24": 413.46,
"25": 402.95, "25": 413.31,
"26": 402.65, "26": 413.61,
"27": 403.01, "27": 413.58,
"28": 402.66, "28": 413.52,
"29": 402.92, "29": 413.37,
"30": 402.75, "30": 413.46,
"31": 403.0, "31": 413.83,
"32": 402.9, "32": 413.23,
"33": 402.48, "33": 413.49,
"34": 402.85, "34": 413.45,
"35": 402.03, "35": 413.41,
"36": 402.93, "36": 413.4,
"37": 402.3, "37": 413.46,
"38": 403.12, "38": 413.38,
"39": 402.83, "39": 413.43,
"40": 402.84, "40": 413.52,
"41": 402.75, "41": 413.56,
"42": 402.97, "42": 413.63,
"43": 402.62, "43": 413.46,
"44": 402.91, "44": 413.47,
"45": 402.81, "45": 413.38,
"46": 402.97, "46": 413.57,
"47": 402.57, "47": 413.37,
"48": 403.0, "48": 413.19,
"49": 402.75 "49": 413.38
}, },
"expected_e2e_ms": 23819.1, "expected_e2e_ms": 20670.75,
"expected_avg_denoise_ms": 388.22, "expected_avg_denoise_ms": 399.24,
"expected_median_denoise_ms": 402.82 "expected_median_denoise_ms": 413.41
}, },
"flux_2_image_t2i_2npu": { "flux_2_image_t2i_2npu": {
"stages_ms": { "stages_ms": {
@@ -138,125 +138,125 @@
}, },
"wan2_1_t2v_1.3b_1_npu": { "wan2_1_t2v_1.3b_1_npu": {
"stages_ms": { "stages_ms": {
"InputValidationStage": 0.07, "InputValidationStage": 0.06,
"TextEncodingStage": 876.11, "TextEncodingStage": 2386.55,
"LatentPreparationStage": 0.25,
"TimestepPreparationStage": 2.9,
"DenoisingStage": 26188.0,
"DecodingStage": 650.1,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 153.0,
"1": 329.59,
"2": 545.23,
"3": 537.0,
"4": 536.27,
"5": 536.29,
"6": 536.33,
"7": 536.0,
"8": 536.17,
"9": 536.28,
"10": 535.53,
"11": 536.04,
"12": 536.42,
"13": 536.09,
"14": 536.32,
"15": 536.25,
"16": 536.36,
"17": 536.21,
"18": 536.29,
"19": 536.15,
"20": 536.28,
"21": 536.5,
"22": 536.46,
"23": 536.06,
"24": 536.45,
"25": 536.24,
"26": 536.14,
"27": 536.13,
"28": 536.22,
"29": 536.15,
"30": 535.94,
"31": 536.1,
"32": 536.13,
"33": 536.2,
"34": 536.24,
"35": 536.34,
"36": 536.54,
"37": 536.42,
"38": 536.41,
"39": 536.42,
"40": 536.13,
"41": 536.32,
"42": 536.23,
"43": 536.16,
"44": 536.05,
"45": 536.18,
"46": 536.08,
"47": 536.34,
"48": 536.26,
"49": 535.41
},
"expected_e2e_ms": 38738.17,
"expected_avg_denoise_ms": 523.62,
"expected_median_denoise_ms": 536.23
},
"wan2_2_t2v_14b_w8a8_8npu": {
"stages_ms": {
"InputValidationStage": 0.07,
"TextEncodingStage": 1200.21,
"LatentPreparationStage": 0.2, "LatentPreparationStage": 0.2,
"TimestepPreparationStage": 2.68, "TimestepPreparationStage": 3.03,
"DenoisingStage": 83661.46, "DenoisingStage": 26240.89,
"DecodingStage": 1080.05, "DecodingStage": 720.28,
"per_frame_generation": null "per_frame_generation": null
}, },
"denoise_step_ms": { "denoise_step_ms": {
"0": 1919.92, "0": 101.91,
"1": 2099.45, "1": 303.37,
"2": 2092.11, "2": 548.94,
"3": 2090.84, "3": 542.02,
"4": 2089.89, "4": 537.7,
"5": 2090.6, "5": 537.26,
"6": 2090.77, "6": 537.78,
"7": 2091.43, "7": 537.27,
"8": 2091.24, "8": 537.72,
"9": 2067.83, "9": 537.65,
"10": 2078.02, "10": 537.49,
"11": 2090.75, "11": 537.41,
"12": 2108.36, "12": 537.31,
"13": 2096.16, "13": 537.43,
"14": 2091.74, "14": 537.41,
"15": 2091.47, "15": 537.51,
"16": 2091.6, "16": 537.54,
"17": 2091.94, "17": 537.56,
"18": 2091.39, "18": 538.09,
"19": 2090.69, "19": 537.09,
"20": 2090.27, "20": 537.72,
"21": 2090.77, "21": 537.56,
"22": 2090.24, "22": 537.85,
"23": 2091.65, "23": 537.54,
"24": 2091.21, "24": 537.67,
"25": 2126.82, "25": 537.66,
"26": 2338.39, "26": 537.25,
"27": 2085.18, "27": 537.71,
"28": 2084.68, "28": 537.77,
"29": 2084.71, "29": 537.16,
"30": 2051.48, "30": 537.6,
"31": 2104.3, "31": 537.58,
"32": 2084.58, "32": 537.6,
"33": 2085.04, "33": 537.51,
"34": 2085.03, "34": 537.78,
"35": 2084.58, "35": 537.43,
"36": 2084.41, "36": 537.52,
"37": 2085.16, "37": 537.59,
"38": 2084.88, "38": 537.46,
"39": 2083.54 "39": 537.77,
"40": 537.41,
"41": 538.11,
"42": 537.41,
"43": 537.18,
"44": 537.64,
"45": 537.32,
"46": 537.52,
"47": 537.52,
"48": 537.29,
"49": 547.88
}, },
"expected_e2e_ms": 91733.92, "expected_e2e_ms": 29360.56,
"expected_avg_denoise_ms": 2091.33, "expected_avg_denoise_ms": 524.67,
"expected_median_denoise_ms": 2090.72 "expected_median_denoise_ms": 537.54
},
"wan2_2_t2v_14b_w8a8_8npu": {
"stages_ms": {
"InputValidationStage": 0.14,
"TextEncodingStage": 3020.73,
"LatentPreparationStage": 0.19,
"TimestepPreparationStage": 5.01,
"DenoisingStage": 82744.33,
"DecodingStage": 932.41,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 1232.32,
"1": 2091.77,
"2": 2097.62,
"3": 2087.53,
"4": 2088.54,
"5": 2087.96,
"6": 2088.28,
"7": 2089.77,
"8": 2101.9,
"9": 2088.73,
"10": 2088.04,
"11": 2087.53,
"12": 2088.89,
"13": 2087.09,
"14": 2088.25,
"15": 2087.96,
"16": 2088.24,
"17": 2088.45,
"18": 2104.7,
"19": 2088.44,
"20": 2087.19,
"21": 2088.19,
"22": 2088.37,
"23": 2087.6,
"24": 2088.13,
"25": 2088.06,
"26": 2126.23,
"27": 2089.92,
"28": 2087.37,
"29": 2089.21,
"30": 2088.29,
"31": 2087.89,
"32": 2073.1,
"33": 2086.71,
"34": 2087.88,
"35": 2088.64,
"36": 2088.1,
"37": 2089.14,
"38": 2087.5,
"39": 2087.86
},
"expected_e2e_ms": 86719.57,
"expected_avg_denoise_ms": 2068.43,
"expected_median_denoise_ms": 2088.21
}, },
"qwen_image_t2i_2npu": { "qwen_image_t2i_2npu": {
"stages_ms": { "stages_ms": {
@@ -154,7 +154,7 @@ def npu_format_cast(
logger.warning_once( logger.warning_once(
"Warning: The conversion from 'ND' to 'NZ' does not work on the CPU. " "Warning: The conversion from 'ND' to 'NZ' does not work on the CPU. "
"Please disable offloading, otherwise the performance will be " "Please disable offloading, otherwise the performance will be "
"significantly reduced." "significantly reduced. --dit-cpu-offload false"
) )
return tensor return tensor