[Bugfix] Enhance extra_config handling for layer name suffix matching (#48589)
Signed-off-by: Xin He <xin3.he@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Xin He committed
Jul 28, 2026 at 20:37 UTC
948107acf7ef8813b8ec94fff7c5ab62d4aba9ae
2 files changed
+57
tests/quantization/test_auto_round.py
+49
@@ -243,6 +243,27 @@ def test_inc_config_parser_parallel_lm_head_defaults_to_unquantized() -> None:
243
assert layer_config.bits == 16
244
245
246
+def test_inc_config_parser_suffix_match_for_lm_head() -> None:
247
+ """Short extra_config key should match fully-qualified lm_head layer name."""
248
+ layer = object.__new__(ParallelLMHead)
249
+ config = make_config(
250
+ extra_config={
251
+ "lm_head": {
252
+ "bits": 4,
253
+ "group_size": 128,
254
+ "sym": True,
255
+ }
256
+ }
257
+ )
258
+
259
+ layer_config = config.config_parser.resolve(layer, "model.language_model.lm_head")
260
+
261
+ assert layer_config.quantized is True
262
+ assert layer_config.bits == 4
263
+ assert layer_config.group_size == 128
264
+ assert layer_config.sym is True
265
+
266
+
267
def test_inc_config_parser_fused_moe_requires_consistent_configs() -> None:
268
config = make_config(
269
extra_config={
@@ -790,6 +811,34 @@ def test_inc_get_quant_method_linear_uses_resolved_scheme(monkeypatch) -> None:
811
assert method is sentinel
812
813
814
+def test_inc_get_quant_method_lm_head_uses_suffix_match(monkeypatch) -> None:
815
+ """lm_head extra_config should apply to fully-qualified prefix."""
816
+ config = make_config(
817
+ extra_config={
818
+ "lm_head": {
819
+ "bits": 4,
820
+ "group_size": 128,
821
+ "sym": True,
822
+ }
823
+ }
824
+ )
825
+ layer = object.__new__(ParallelLMHead)
826
+ sentinel = object()
827
+
828
+ class DummyScheme:
829
+ def get_linear_method(self, _config, _layer, _prefix, _layer_config):
830
+ return sentinel
831
+
832
+ monkeypatch.setattr(
833
+ "vllm.model_executor.layers.quantization.inc.schemes.factory.resolve_scheme",
834
+ lambda _layer_config: DummyScheme(),
835
+ )
836
+
837
+ method = config.get_quant_method(layer, "model.language_model.lm_head")
838
+
839
+ assert method is sentinel
840
+
841
+
842
def test_inc_get_quant_method_moe_uses_resolved_scheme(monkeypatch) -> None:
843
config = make_config()
844
layer = object.__new__(RoutedExperts)
vllm/model_executor/layers/quantization/inc/config_parser.py
+8
@@ -142,6 +142,14 @@ class INCConfigParser:
142
if self._config.extra_config and layer_name in self._config.extra_config:
143
return get_config(layer_name)
144
145
+ # Suffix match: handle cases where extra_config keys use short names
146
+ # (e.g. "lm_head") but the layer_name is fully qualified
147
+ # (e.g. "model.language_model.lm_head") due to model nesting.
148
+ if self._config.extra_config:
149
+ for cfg_key in self._config.extra_config:
150
+ if layer_name.endswith(f".{cfg_key}"):
151
+ return get_config(cfg_key)
152
+
153
quantized = not isinstance(layer, ParallelLMHead)
154
if self._config.block_name_to_quantize:
155
quantized = any(