[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:
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
from sglang.srt.layers.quantization.modelslim.schemes import (
ModelSlimLinearScheme,
)
from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
logger = logging.getLogger(__name__)
@@ -40,7 +41,11 @@ class ModelSlimConfig(QuantizationConfig):
- 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__()
self.quant_description = quant_config
ignore = cast(List[str], quant_config.get("ignore", []))
@@ -49,6 +54,11 @@ class ModelSlimConfig(QuantizationConfig):
self.packed_modules_mapping = (
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:
return ModelSlimLinearMethod(self)
@@ -71,8 +81,10 @@ class ModelSlimConfig(QuantizationConfig):
return filenames
@classmethod
def from_config(cls, config: Dict[str, Any]) -> ModelSlimConfig:
return cls(config)
def from_config(
cls, config: Dict[str, Any], reverse_param_names_mapping: dict = None
) -> ModelSlimConfig:
return cls(config, reverse_param_names_mapping)
def get_quant_method(
self,
@@ -109,16 +121,18 @@ class ModelSlimConfig(QuantizationConfig):
self,
layer_name: str,
) -> 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":
return ModelSlimW8A8Int8(
quant_config=self.quant_description, prefix=layer_name
)
return ModelSlimW8A8Int8(quant_config=self.quant_description, prefix=prefix)
elif quant_type == "W4A4_DYNAMIC":
return ModelSlimW4A4Int4(
quant_config=self.quant_description, prefix=layer_name
)
return ModelSlimW4A4Int4(quant_config=self.quant_description, prefix=prefix)
elif quant_type == "W8A8_MXFP8":
from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp8_scheme import (
ModelSlimMXFP8Scheme,
@@ -493,8 +493,11 @@ def _resolve_quant_config(
reverse_param_names_mapping_dict = getattr(
arch_config, "reverse_param_names_mapping", None
)
quant_config = get_quant_config(hf_config, component_model_path)
quant_config = get_quant_config(
hf_config,
component_model_path,
reverse_param_names_mapping=reverse_param_names_mapping_dict,
)
quant_config_name = _get_quant_config_name(quant_config)
inferred_nvfp4_config = None
if quant_config is None or quant_config_name == "modelopt_fp4":
@@ -127,12 +127,13 @@ def get_quant_config(
model_config,
component_model_path: str,
packed_modules_mapping: Dict[str, List[str]] = {},
reverse_param_names_mapping: Dict[str, List[str]] = {},
remap_prefix: Dict[str, str] | None = None,
) -> QuantizationConfig:
quant_cfg = find_quant_modelslim_config(model_config, component_model_path)
if quant_cfg is not None:
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:
return None
@@ -7,68 +7,68 @@
"scenarios": {
"flux_image_t2i_npu": {
"stages_ms": {
"InputValidationStage": 0.07,
"TextEncodingStage": 154.51,
"TimestepPreparationStage": 53.52,
"LatentPreparationStage": 0.39,
"DenoisingStage": 19423.39,
"DecodingStage": 196.62
"InputValidationStage": 0.06,
"TextEncodingStage": 468.6,
"TimestepPreparationStage": 41.44,
"LatentPreparationStage": 0.28,
"DenoisingStage": 19973.94,
"DecodingStage": 8.58
},
"denoise_step_ms": {
"0": 123.16,
"1": 91.7,
"2": 265.62,
"3": 402.68,
"4": 402.86,
"5": 402.78,
"6": 402.99,
"7": 402.77,
"8": 402.59,
"9": 402.93,
"10": 402.05,
"11": 402.99,
"12": 402.29,
"13": 403.07,
"14": 402.62,
"15": 402.99,
"16": 402.68,
"17": 403.0,
"18": 402.74,
"19": 402.85,
"20": 402.83,
"21": 403.03,
"22": 402.56,
"23": 402.84,
"24": 402.79,
"25": 402.95,
"26": 402.65,
"27": 403.01,
"28": 402.66,
"29": 402.92,
"30": 402.75,
"31": 403.0,
"32": 402.9,
"33": 402.48,
"34": 402.85,
"35": 402.03,
"36": 402.93,
"37": 402.3,
"38": 403.12,
"39": 402.83,
"40": 402.84,
"41": 402.75,
"42": 402.97,
"43": 402.62,
"44": 402.91,
"45": 402.81,
"46": 402.97,
"47": 402.57,
"48": 403.0,
"49": 402.75
"0": 100.15,
"1": 95.5,
"2": 334.93,
"3": 413.4,
"4": 413.43,
"5": 413.33,
"6": 413.6,
"7": 413.2,
"8": 413.4,
"9": 413.15,
"10": 413.24,
"11": 413.37,
"12": 413.63,
"13": 413.39,
"14": 413.49,
"15": 413.69,
"16": 413.61,
"17": 413.38,
"18": 413.32,
"19": 413.3,
"20": 413.45,
"21": 413.49,
"22": 413.31,
"23": 413.14,
"24": 413.46,
"25": 413.31,
"26": 413.61,
"27": 413.58,
"28": 413.52,
"29": 413.37,
"30": 413.46,
"31": 413.83,
"32": 413.23,
"33": 413.49,
"34": 413.45,
"35": 413.41,
"36": 413.4,
"37": 413.46,
"38": 413.38,
"39": 413.43,
"40": 413.52,
"41": 413.56,
"42": 413.63,
"43": 413.46,
"44": 413.47,
"45": 413.38,
"46": 413.57,
"47": 413.37,
"48": 413.19,
"49": 413.38
},
"expected_e2e_ms": 23819.1,
"expected_avg_denoise_ms": 388.22,
"expected_median_denoise_ms": 402.82
"expected_e2e_ms": 20670.75,
"expected_avg_denoise_ms": 399.24,
"expected_median_denoise_ms": 413.41
},
"flux_2_image_t2i_2npu": {
"stages_ms": {
@@ -138,125 +138,125 @@
},
"wan2_1_t2v_1.3b_1_npu": {
"stages_ms": {
"InputValidationStage": 0.07,
"TextEncodingStage": 876.11,
"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,
"InputValidationStage": 0.06,
"TextEncodingStage": 2386.55,
"LatentPreparationStage": 0.2,
"TimestepPreparationStage": 2.68,
"DenoisingStage": 83661.46,
"DecodingStage": 1080.05,
"TimestepPreparationStage": 3.03,
"DenoisingStage": 26240.89,
"DecodingStage": 720.28,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 1919.92,
"1": 2099.45,
"2": 2092.11,
"3": 2090.84,
"4": 2089.89,
"5": 2090.6,
"6": 2090.77,
"7": 2091.43,
"8": 2091.24,
"9": 2067.83,
"10": 2078.02,
"11": 2090.75,
"12": 2108.36,
"13": 2096.16,
"14": 2091.74,
"15": 2091.47,
"16": 2091.6,
"17": 2091.94,
"18": 2091.39,
"19": 2090.69,
"20": 2090.27,
"21": 2090.77,
"22": 2090.24,
"23": 2091.65,
"24": 2091.21,
"25": 2126.82,
"26": 2338.39,
"27": 2085.18,
"28": 2084.68,
"29": 2084.71,
"30": 2051.48,
"31": 2104.3,
"32": 2084.58,
"33": 2085.04,
"34": 2085.03,
"35": 2084.58,
"36": 2084.41,
"37": 2085.16,
"38": 2084.88,
"39": 2083.54
"0": 101.91,
"1": 303.37,
"2": 548.94,
"3": 542.02,
"4": 537.7,
"5": 537.26,
"6": 537.78,
"7": 537.27,
"8": 537.72,
"9": 537.65,
"10": 537.49,
"11": 537.41,
"12": 537.31,
"13": 537.43,
"14": 537.41,
"15": 537.51,
"16": 537.54,
"17": 537.56,
"18": 538.09,
"19": 537.09,
"20": 537.72,
"21": 537.56,
"22": 537.85,
"23": 537.54,
"24": 537.67,
"25": 537.66,
"26": 537.25,
"27": 537.71,
"28": 537.77,
"29": 537.16,
"30": 537.6,
"31": 537.58,
"32": 537.6,
"33": 537.51,
"34": 537.78,
"35": 537.43,
"36": 537.52,
"37": 537.59,
"38": 537.46,
"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_avg_denoise_ms": 2091.33,
"expected_median_denoise_ms": 2090.72
"expected_e2e_ms": 29360.56,
"expected_avg_denoise_ms": 524.67,
"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": {
"stages_ms": {
@@ -154,7 +154,7 @@ def npu_format_cast(
logger.warning_once(
"Warning: The conversion from 'ND' to 'NZ' does not work on the CPU. "
"Please disable offloading, otherwise the performance will be "
"significantly reduced."
"significantly reduced. --dit-cpu-offload false"
)
return tensor