@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7894394b07

[Online quantization] Add targeted online quantization configuration based on user patterns (#51285)

fxmarty-amd committed Sep 2, 2026 at 14:55 UTC 7894394b0759ba0f2d9c9cca249a6e49ecdd814f
11 files changed +892 -172
docs/features/quantization/online.md
+47 -3
@@ -67,7 +67,7 @@ quantization_config:
67 moe:
68 weight: <name>
69 activation: <name>
70 - ignore: [<layer-name-or-regex>, ...]
70 + ignore: [<layer-name-or-regex-or-fnmatch-pattern>, ...]
71 ```
72
73 `linear` and `moe` accept a full `{weight, activation}` dict, or a bare
@@ -133,7 +133,7 @@ llm = LLM(
133
134 ### Excluding Layers from Quantization
135
136 -Use the `ignore` parameter to skip specific layers. It accepts exact layer names and regex patterns (prefixed with `re:`):
136 +Use the `ignore` parameter to skip specific layers. It accepts exact layer names, regex patterns (prefixed with `re:`), and patterns understood by [`fnmatch.fnmatch`](https://docs.python.org/3/library/fnmatch.html#fnmatch.fnmatch):
137
138 ```python
139 from vllm import LLM
@@ -147,10 +147,54 @@ llm = LLM(
147 "model.layers.1.self_attn.o_proj",
148 # regex: skip all QKV projections
149 "re:.*[qkv]_proj",
150 + # fnmatch: skip all MoE experts
151 + "*mlp.experts*",
152 ],
153 },
154 )
155 ```
156
157 !!! note
156 - For fused layers (e.g., `qkv_proj` which fuses `q_proj`, `k_proj`, `v_proj`), the ignore pattern must match the **unfused** shard names (`q_proj`, `k_proj`, `v_proj`), not the fused name.
158 + For fused layers (e.g., `qkv_proj` which fuses `q_proj`, `k_proj`, `v_proj`), patterns may match the fused name directly or all of its unfused shard names.
159 +
160 +### Fine-Grained Per-Layer Quantization Schemes
161 +
162 +Use the `targets` parameter to apply different online shorthands to different layers, instead of one scheme applied everywhere via `linear`/`moe`. Keys are exact layer names, regex patterns (prefixed with `re:`), or patterns understood by [`fnmatch.fnmatch`](https://docs.python.org/3/library/fnmatch.html#fnmatch.fnmatch); values are shorthand names (`fp8_per_tensor`, `fp8_per_block`, `fp8_per_channel`, `mxfp8`, `int8_per_channel_weight_only`, `nvfp4_per_token`).
163 +
164 +Example:
165 +
166 +```python
167 +from vllm import LLM
168 +
169 +llm = LLM(
170 + "Qwen/Qwen3.5-35B-A3B",
171 + quantization="online",
172 + quantization_config={
173 + "targets": {
174 + # exact layer name
175 + "model.layers.0.self_attn.o_proj": "fp8_per_tensor",
176 + # regex: quantize all QKV projections
177 + r"re:.*self_attn\.qkv_proj.*": "mxfp8",
178 + # fnmatch: quantize all MoE experts
179 + "*mlp.experts*": "mxfp4",
180 + },
181 + },
182 +)
183 +```
184 +
185 +Or with the CLI:
186 +
187 +```bash
188 +vllm serve Qwen/Qwen3.5-35B-A3B \
189 + --quantization online \
190 + --quantization-config '{"targets":{"model.layers.0.self_attn.o_proj":"fp8_per_tensor","re:.*self_attn\\.qkv_proj.*":"mxfp8","*mlp.experts*":"mxfp4"}}'
191 +```
192 +
193 +!!! info
194 +
195 + - `targets` is mutually exclusive with online `linear` and `moe`: set one or the other, not both.
196 + - A layer that matches no `targets` pattern is left unchanged from its checkpoint dtype.
197 + - A layer name may not match both `targets` and `ignore`, raising an error.
198 + - A layer may not match more than one `targets` pattern, raising an error.
199 + - fnmatch-style patterns are supported only by online quantization; they are
200 + not applied to Quark or compressed-tensors configurations.
tests/quantization/test_config_utils.py new
+195
@@ -0,0 +1,195 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests quantization configuration matching utilities."""
4 +
5 +from unittest.mock import Mock
6 +
7 +import pytest
8 +
9 +from vllm.config.quantization import QuantizationConfigArgs
10 +from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
11 + should_ignore_layer,
12 +)
13 +from vllm.model_executor.layers.quantization.online.base import (
14 + OnlineQuantizationConfig,
15 + _find_matching_targets,
16 +)
17 +from vllm.model_executor.layers.quantization.quark.utils import (
18 + should_ignore_layer as quark_should_ignore_layer,
19 +)
20 +from vllm.model_executor.layers.quantization.utils.config_utils import (
21 + find_matching_patterns,
22 + get_layer_name_after_index,
23 + is_equal_or_regex_match,
24 +)
25 +
26 +
27 +def test_is_equal_or_regex_match():
28 + assert is_equal_or_regex_match(
29 + "model.layers.0.mlp.down_proj", "model.layers.0.mlp.down_proj"
30 + )
31 + assert is_equal_or_regex_match("model.layers.0.mlp.down_proj", r"re:.*down_proj")
32 + assert not is_equal_or_regex_match("model.layers.0.mlp.down_proj", "other")
33 +
34 +
35 +def test_fnmatch_is_opt_in():
36 + layer_name = "model.layers.0.mlp.experts"
37 + pattern = "*mlp.experts*"
38 +
39 + assert not is_equal_or_regex_match(layer_name, pattern)
40 + assert is_equal_or_regex_match(layer_name, pattern, use_fnmatch=True)
41 + assert not should_ignore_layer(layer_name, [pattern])
42 + assert not quark_should_ignore_layer(layer_name, [pattern])
43 + assert should_ignore_layer(layer_name, [pattern], use_fnmatch=True)
44 +
45 +
46 +def test_online_targets_support_fnmatch_patterns():
47 + layer_name = "model.layers.0.mlp.experts"
48 + targets = {"*mlp.experts*": "mxfp4"}
49 +
50 + assert _find_matching_targets(layer_name, targets) == ["*mlp.experts*"]
51 +
52 +
53 +def test_online_ignore_supports_fnmatch_patterns():
54 + layer_name = "model.layers.0.mlp.experts"
55 + ignore = ["*mlp.experts*"]
56 + config = OnlineQuantizationConfig(
57 + QuantizationConfigArgs(targets={"*mlp.experts*": "mxfp4"}, ignore=ignore)
58 + )
59 +
60 + with pytest.raises(ValueError, match="matches both quantization_config.ignore"):
61 + config._dispatch_target(layer_name, Mock())
62 +
63 +
64 +@pytest.mark.parametrize(
65 + "layer_name,expected",
66 + [
67 + ("model.layers.1.self_attn.qkv_proj", "self_attn.qkv_proj"),
68 + ("model.layers.2.mlp.down_proj", "mlp.down_proj"),
69 + ("lm_head", "lm_head"),
70 + ],
71 +)
72 +def test_get_layer_name_after_index(layer_name, expected):
73 + assert get_layer_name_after_index(layer_name) == expected
74 +
75 +
76 +@pytest.mark.parametrize(
77 + "patterns",
78 + [
79 + [r"re:.*qkv_proj.*"],
80 + [r"re:.*\.[qkv]_proj$"],
81 + ],
82 + ids=["direct_fused_regex", "fused_shard_regexes"],
83 +)
84 +def test_find_matching_patterns_for_fused_regexes(patterns):
85 + layer_name = "model.layers.0.self_attn.qkv_proj"
86 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
87 +
88 + matches = find_matching_patterns(layer_name, patterns, fused_mapping)
89 + assert all(len(shard_matches) == 1 for shard_matches in matches)
90 +
91 +
92 +@pytest.mark.parametrize(
93 + "patterns,expected",
94 + [
95 + (
96 + [r"re:.*\.qkv_proj$"],
97 + [{r"re:.*\.qkv_proj$"}],
98 + ),
99 + (
100 + [r"re:.*\.q_proj$", r"re:.*\.k_proj$", r"re:.*\.v_proj$"],
101 + [
102 + {r"re:.*\.q_proj$"},
103 + {r"re:.*\.k_proj$"},
104 + {r"re:.*\.v_proj$"},
105 + ],
106 + ),
107 + ],
108 + ids=["direct_fused_match", "individual_shard_matches"],
109 +)
110 +def test_find_matching_patterns_distinguishes_direct_and_shard_matches(
111 + patterns, expected
112 +):
113 + layer_name = "model.layers.0.self_attn.qkv_proj"
114 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
115 +
116 + assert find_matching_patterns(layer_name, patterns, fused_mapping) == expected
117 + assert should_ignore_layer(layer_name, patterns, fused_mapping)
118 +
119 +
120 +def test_should_ignore_layer_rejects_partially_matched_fused_layer():
121 + layer_name = "model.layers.0.self_attn.qkv_proj"
122 + patterns = [r"re:.*\.q_proj$"]
123 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
124 +
125 + assert find_matching_patterns(layer_name, patterns, fused_mapping) == [
126 + {r"re:.*\.q_proj$"},
127 + set(),
128 + set(),
129 + ]
130 + with pytest.raises(ValueError, match="different quantization schemes"):
131 + should_ignore_layer(layer_name, patterns, fused_mapping)
132 +
133 +
134 +def test_should_ignore_layer_returns_false_when_no_fused_pattern_matches():
135 + layer_name = "model.layers.0.self_attn.qkv_proj"
136 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
137 +
138 + assert find_matching_patterns(layer_name, [], fused_mapping) == [
139 + set(),
140 + set(),
141 + set(),
142 + ]
143 + assert not should_ignore_layer(layer_name, [], fused_mapping)
144 +
145 +
146 +@pytest.mark.parametrize(
147 + "patterns",
148 + [
149 + [r"re:.*qkv_proj.*"],
150 + [r"re:.*\.[qkv]_proj$"],
151 + ],
152 + ids=["direct_fused_regex", "fused_shard_regexes"],
153 +)
154 +def test_ignore_and_targets_match_fused_regexes_identically(patterns):
155 + layer_name = "model.layers.0.self_attn.qkv_proj"
156 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
157 + targets = {pattern: "fp8_per_block" for pattern in patterns}
158 +
159 + assert should_ignore_layer(layer_name, patterns, fused_mapping)
160 + matches = _find_matching_targets(layer_name, targets, fused_mapping)
161 + assert len(matches) == 1
162 + assert targets[matches[0]] == "fp8_per_block"
163 +
164 +
165 +def test_ignore_allows_individually_matched_fused_shards():
166 + layer_name = "model.layers.0.self_attn.qkv_proj"
167 + patterns = [
168 + "model.layers.0.self_attn.q_proj",
169 + "model.layers.0.self_attn.k_proj",
170 + "model.layers.0.self_attn.v_proj",
171 + ]
172 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
173 +
174 + assert should_ignore_layer(layer_name, patterns, fused_mapping)
175 +
176 +
177 +@pytest.mark.parametrize(
178 + "patterns",
179 + [
180 + ["model.layers.0.self_attn.qkv_proj"],
181 + [r"re:.*\.qkv_proj$"],
182 + [
183 + "model.layers.0.self_attn.q_proj",
184 + "model.layers.0.self_attn.k_proj",
185 + "model.layers.0.self_attn.v_proj",
186 + ],
187 + ],
188 + ids=["direct_fused_name", "direct_fused_regex", "individual_shards"],
189 +)
190 +def test_quark_and_compressed_tensors_ignore_fused_layers_identically(patterns):
191 + layer_name = "model.layers.0.self_attn.qkv_proj"
192 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
193 +
194 + assert should_ignore_layer(layer_name, patterns, fused_mapping)
195 + assert quark_should_ignore_layer(layer_name, patterns, fused_mapping)
tests/quantization/test_online.py
+198 -5
@@ -18,10 +18,21 @@ from tests.quantization.utils import (
18 from vllm import _custom_ops as ops
19 from vllm._aiter_ops import rocm_aiter_ops
20 from vllm._custom_ops import scaled_fp4_quant
21 -from vllm.config.quantization import resolve_quantization_config
21 +from vllm.config import ModelConfig
22 +from vllm.config.quantization import (
23 + QuantizationConfigArgs,
24 + resolve_quantization_config,
25 +)
26 from vllm.forward_context import set_forward_context
27 from vllm.model_executor.layers.attention import Attention
24 -from vllm.model_executor.layers.linear import UnquantizedLinearMethod
28 +from vllm.model_executor.layers.linear import (
29 + ColumnParallelLinear,
30 + LinearBase,
31 + UnquantizedLinearMethod,
32 +)
33 +from vllm.model_executor.layers.quantization.online.base import (
34 + OnlineQuantizationConfig,
35 +)
36 from vllm.model_executor.layers.quantization.online.fp8 import (
37 Fp8PerBlockOnlineLinearMethod,
38 Fp8PerBlockOnlineMoEMethod,
@@ -49,6 +60,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
60 amax_for_tp_weight_quant,
61 weight_amax,
62 )
63 +from vllm.model_executor.layers.vocab_parallel_embedding import (
64 + VocabParallelEmbedding,
65 +)
66 +from vllm.model_executor.model_loader.base_loader import log_online_quantization
67 from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
68 from vllm.model_executor.models.granitemoe import (
69 GraniteMoeModel,
@@ -154,12 +169,23 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
169 Fp8PerBlockOnlineLinearMethod,
170 Fp8PerTensorOnlineMoEMethod,
171 ),
172 + # quantization='online' with per-layer target patterns
173 + (
174 + "online",
175 + {
176 + "targets": {
177 + r"re:.*self_attn\.o_proj": "fp8_per_block",
178 + r"re:.*block_sparse_moe\.experts": "fp8_per_tensor",
179 + }
180 + },
181 + Fp8PerBlockOnlineLinearMethod,
182 + Fp8PerTensorOnlineMoEMethod,
183 + ),
184 # ignore with direct layer name
185 (
186 "fp8_per_tensor",
160 - # qkv_proj is fused from q_proj/k_proj/v_proj, so currently the
161 - # ignore regex must match the unfused shard names
162 - # TODO(future PR): also make 're:.*qkv_proj.*' work
187 + # qkv_proj is fused from q_proj/k_proj/v_proj. The shard regex
188 + # remains supported alongside direct fused-name regexes.
189 {"ignore": ["model.layers.1.self_attn.o_proj", "re:.*[qkv]_proj"]},
190 Fp8PerTensorOnlineLinearMethod,
191 Fp8PerTensorOnlineMoEMethod,
@@ -171,6 +197,15 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
197 Mxfp4OnlineMoEMethod,
198 ),
199 ],
200 + ids=[
201 + "fp8_per_tensor",
202 + "fp8_per_block",
203 + "fp8_per_channel",
204 + "per_layer_kind_overrides",
205 + "targets",
206 + "ignore",
207 + "mxfp4",
208 + ],
209 )
210 @pytest.mark.parametrize(
211 "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
@@ -266,6 +301,12 @@ def test_online_quantization(
301 layer.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod
302 )
303
304 + if isinstance(online_quant_args, dict) and "targets" in online_quant_args:
305 + for layer in model.model.layers:
306 + assert isinstance(
307 + layer.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod
308 + )
309 +
310 input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE)
311 positions = torch.arange(input_ids.numel(), device=DEVICE)
312 with set_forward_context(None, vllm_config, num_tokens=input_ids.numel()):
@@ -321,6 +362,158 @@ def test_online_quantization_loads_real_weights(vllm_runner, monkeypatch) -> Non
362 assert outputs
363
364
365 +@pytest.mark.skipif(
366 + not is_quant_method_supported("fp8"),
367 + reason="FP8 is not supported on this GPU type.",
368 +)
369 +@pytest.mark.parametrize(
370 + "targets,prefix,expected_method_cls,unmatched_prefix,expected_metadata",
371 + [
372 + (
373 + {r"re:.*self_attn\.o_proj": "fp8_per_block"},
374 + "model.layers.0.self_attn.o_proj",
375 + Fp8PerBlockOnlineLinearMethod,
376 + "model.layers.0.self_attn.qkv_proj",
377 + ("targets", "fp8_per_block", r"re:.*self_attn\.o_proj"),
378 + ),
379 + (
380 + {r"re:.*qkv_proj.*": "fp8_per_tensor"},
381 + "model.layers.0.self_attn.qkv_proj",
382 + Fp8PerTensorOnlineLinearMethod,
383 + "model.layers.0.self_attn.o_proj",
384 + ("targets", "fp8_per_tensor", r"re:.*qkv_proj.*"),
385 + ),
386 + (
387 + {r"re:.*[qkv]_proj": "fp8_per_tensor"},
388 + "model.layers.0.self_attn.qkv_proj",
389 + Fp8PerTensorOnlineLinearMethod,
390 + "model.layers.0.self_attn.o_proj",
391 + ("targets", "fp8_per_tensor", r"re:.*[qkv]_proj"),
392 + ),
393 + ],
394 + ids=["linear_regex", "direct_fused_regex", "legacy_fused_regex"],
395 +)
396 +def test_online_quantization_targets(
397 + default_vllm_config,
398 + dist_init,
399 + targets: dict[str, str],
400 + prefix: str,
401 + expected_method_cls,
402 + unmatched_prefix: str,
403 + expected_metadata: tuple[str, str, str],
404 +) -> None:
405 + """Target patterns select the real online linear methods."""
406 + default_vllm_config.model_config = ModelConfig()
407 + config = OnlineQuantizationConfig(QuantizationConfigArgs(targets=targets))
408 + config.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
409 +
410 + layer = ColumnParallelLinear(
411 + input_size=1,
412 + output_size=1,
413 + bias=False,
414 + disable_tp=True,
415 + )
416 +
417 + method = config.get_quant_method(layer, prefix)
418 + assert isinstance(method, expected_method_cls)
419 + assert config.quantized_layers == {prefix: expected_metadata}
420 +
421 + unmatched_method = config.get_quant_method(layer, unmatched_prefix)
422 + assert isinstance(unmatched_method, UnquantizedLinearMethod)
423 + assert config.quantized_layers == {prefix: expected_metadata}
424 +
425 +
426 +@pytest.mark.skipif(
427 + not is_quant_method_supported("fp8"),
428 + reason="FP8 is not supported on this GPU type.",
429 +)
430 +def test_online_quantization_records_global_config(
431 + default_vllm_config, dist_init
432 +) -> None:
433 + default_vllm_config.model_config = ModelConfig()
434 + config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_block"))
435 + prefix = "model.layers.0.self_attn.o_proj"
436 + layer = ColumnParallelLinear(
437 + input_size=1,
438 + output_size=1,
439 + bias=False,
440 + disable_tp=True,
441 + )
442 +
443 + method = config.get_quant_method(layer, prefix)
444 +
445 + assert isinstance(method, Fp8PerBlockOnlineLinearMethod)
446 + assert config.quantized_layers == {
447 + prefix: ("linear", str(config.args.linear), None)
448 + }
449 +
450 +
451 +def test_online_quantization_targets_ignore_collision() -> None:
452 + """A targets/ignore collision is reported when the layer is dispatched."""
453 + config = OnlineQuantizationConfig(
454 + QuantizationConfigArgs(
455 + targets={"model.layers.0.self_attn.o_proj": "fp8_per_tensor"},
456 + ignore=["model.layers.0.self_attn.o_proj"],
457 + )
458 + )
459 + with pytest.raises(ValueError, match="matches both quantization_config.ignore"):
460 + config._dispatch_target(
461 + "model.layers.0.self_attn.o_proj", Mock(spec=LinearBase)
462 + )
463 +
464 +
465 +def test_online_quantization_targets_reject_unsupported_layer() -> None:
466 + """A targets match on a non-linear, non-MoE layer is rejected."""
467 + config = OnlineQuantizationConfig(
468 + QuantizationConfigArgs(targets={"lm_head": "fp8_per_tensor"})
469 + )
470 + lm_head = VocabParallelEmbedding(
471 + num_embeddings=1,
472 + embedding_dim=1,
473 + disable_tp=True,
474 + )
475 +
476 + with pytest.raises(
477 + ValueError,
478 + match=(
479 + "Layer lm_head was matched by quantization_config.targets "
480 + r"\(lm_head\), but online quantization is not supported for "
481 + "VocabParallelEmbedding."
482 + ),
483 + ):
484 + config.get_quant_method(lm_head, "lm_head")
485 +
486 +
487 +def test_log_online_quantization(default_vllm_config, monkeypatch) -> None:
488 + config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_tensor"))
489 + config.quantized_layers = {
490 + "model.layers.0.mlp.down_proj": ("linear", "fp8_per_tensor", None),
491 + "model.layers.1.mlp.down_proj": ("linear", "fp8_per_tensor", None),
492 + "model.layers.0.self_attn.qkv_proj": (
493 + "targets",
494 + "mxfp4",
495 + r"re:.*qkv_proj.*",
496 + ),
497 + }
498 + default_vllm_config.quant_config = config
499 +
500 + logged_messages: list[str] = []
501 +
502 + def record_info(message: str, *args: object) -> None:
503 + logged_messages.append(message % args)
504 +
505 + monkeypatch.setattr(
506 + "vllm.model_executor.model_loader.base_loader.logger.info", record_info
507 + )
508 + log_online_quantization(default_vllm_config)
509 +
510 + assert logged_messages == [
511 + "Quantized 3 layers of types: mlp.down_proj: 2 (from linear: "
512 + "fp8_per_tensor); self_attn.qkv_proj: 1 (from targets: "
513 + "re:.*qkv_proj.*, mxfp4)"
514 + ]
515 +
516 +
517 @pytest.mark.skipif(
518 not (
519 current_platform.is_cuda()
tests/quantization/test_quantization_config_args.py
+77
@@ -2,6 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for QuantizationConfigArgs parsing."""
4
5 +from unittest.mock import Mock
6 +
7 import pytest
8
9 from vllm.config.quantization import (
@@ -10,6 +12,11 @@ from vllm.config.quantization import (
12 QuantSpec,
13 resolve_quantization_config,
14 )
15 +from vllm.model_executor.layers.linear import LinearBase
16 +from vllm.model_executor.layers.quantization.online.base import (
17 + OnlineQuantizationConfig,
18 + _find_matching_targets,
19 +)
20 from vllm.model_executor.layers.quantization.utils.quant_utils import (
21 kFp8Dynamic128Sym,
22 kFp8DynamicTokenSym,
@@ -34,6 +41,10 @@ def test_quant_spec_accepts_quant_key_directly():
41 assert spec.activation is None
42
43
44 +def test_quant_spec_string_representation_uses_quantization_name():
45 + assert str(QuantSpec(weight="mxfp4")) == "mxfp4"
46 +
47 +
48 def test_quant_spec_rejects_unknown_name():
49 with pytest.raises(ValueError, match="unknown quantization name"):
50 QuantSpec(weight="not_a_real_format")
@@ -70,6 +81,11 @@ def test_args_accepts_dict_form():
81 assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic)
82
83
84 +def test_targets_reject_non_string_keys():
85 + with pytest.raises(ValueError, match="targets keys must be strings"):
86 + QuantizationConfigArgs._validate_targets({123: "mxfp8"})
87 +
88 +
89 # ---- resolve_quantization_config -----------------------------------------
90
91
@@ -132,3 +148,64 @@ def test_static_block_weight_paired_with_dynamic_block_activation():
148 spec = QuantSpec(weight="fp8_per_block_static", activation="fp8_per_block_dynamic")
149 assert spec.weight == kFp8Static128BlockSym
150 assert spec.activation == kFp8Dynamic128Sym
151 +
152 +
153 +def test_targets_allow_distinct_patterns_with_the_same_shorthand():
154 + layer_name = "model.layers.0.self_attn.qkv_proj"
155 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
156 + targets = {
157 + r"re:.*\.q_proj$": "mxfp8",
158 + r"re:.*\.k_proj$": "mxfp8",
159 + r"re:.*\.v_proj$": "mxfp8",
160 + }
161 +
162 + matches = _find_matching_targets(layer_name, targets, fused_mapping)
163 + assert len(matches) == 1
164 + assert targets[matches[0]] == "mxfp8"
165 +
166 +
167 +def test_targets_reject_overlapping_patterns():
168 + targets = {
169 + r"re:.*o_proj": "fp8_per_tensor",
170 + "model.layers.0.self_attn.o_proj": "fp8_per_block",
171 + }
172 +
173 + with pytest.raises(ValueError, match="multiple quantization_config.targets"):
174 + _find_matching_targets("model.layers.0.self_attn.o_proj", targets)
175 +
176 +
177 +def test_targets_reject_partially_matched_fused_layer():
178 + targets = {r"re:.*q_proj": "fp8_per_tensor"}
179 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
180 +
181 + with pytest.raises(ValueError, match="unmatched shards"):
182 + _find_matching_targets(
183 + "model.layers.0.self_attn.qkv_proj", targets, fused_mapping
184 + )
185 +
186 +
187 +def test_targets_reject_fused_shards_with_different_schemes():
188 + targets = {
189 + r"re:.*\.q_proj$": "fp8_per_tensor",
190 + r"re:.*\.k_proj$": "fp8_per_block",
191 + r"re:.*\.v_proj$": "fp8_per_tensor",
192 + }
193 + fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
194 +
195 + with pytest.raises(ValueError, match="different quantization_config.targets"):
196 + _find_matching_targets(
197 + "model.layers.0.self_attn.qkv_proj", targets, fused_mapping
198 + )
199 +
200 +
201 +def test_targets_reject_moe_only_shorthand_for_linear_layer():
202 + config = OnlineQuantizationConfig(
203 + QuantizationConfigArgs(
204 + targets={"model.layers.0.self_attn.o_proj": "nvfp4_per_token"}
205 + )
206 + )
207 +
208 + with pytest.raises(ValueError, match="does not define a QuantSpec"):
209 + config._dispatch_target(
210 + "model.layers.0.self_attn.o_proj", Mock(spec=LinearBase)
211 + )
tests/quantization/test_quark.py
+21 -1
@@ -50,7 +50,10 @@ from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E
50 QuarkW8A8Int8MoEMethod,
51 )
52 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
53 -from vllm.model_executor.layers.quantization.quark.utils import QuarkQTensorHint
53 +from vllm.model_executor.layers.quantization.quark.utils import (
54 + QuarkQTensorHint,
55 + should_ignore_layer,
56 +)
57 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
58 quant_dequant_mxfp4,
59 )
@@ -1294,6 +1297,23 @@ FUSED_MAPPING = {
1297 }
1298
1299
1300 +def test_quark_should_ignore_layer_checks_children():
1301 + assert should_ignore_layer(
1302 + "model.layers.78.mlp.experts",
1303 + ["model.layers.78.mlp.experts.0.down_proj"],
1304 + check_children=True,
1305 + )
1306 +
1307 +
1308 +def test_quark_should_ignore_layer_rejects_partial_fused_matches():
1309 + with pytest.raises(ValueError, match="different quantization schemes"):
1310 + should_ignore_layer(
1311 + "model.layers.0.self_attn.qkv_proj",
1312 + ["model.layers.0.self_attn.q_proj"],
1313 + FUSED_MAPPING,
1314 + )
1315 +
1316 +
1317 def test_fused_name_listed_directly_is_skipped():
1318 # Regression for Step-3.5-Flash-FP8: the checkpoint lists the fused
1319 # name (``qkv_proj``) directly in ``modules_to_not_convert``. When a
vllm/config/quantization.py
+74 -2
@@ -3,7 +3,14 @@
3
4 from typing import Annotated, Any
5
6 -from pydantic import Field, GetPydanticSchema, ValidationInfo, field_validator
6 +import regex as re
7 +from pydantic import (
8 + Field,
9 + GetPydanticSchema,
10 + ValidationInfo,
11 + field_validator,
12 + model_validator,
13 +)
14 from pydantic_core import core_schema
15
16 from vllm.config.utils import config
@@ -76,6 +83,21 @@ class QuantSpec:
83 activation: QuantKeyField = None
84 """Activation quantization key, or a name from QUANT_KEY_NAMES."""
85
86 + def __str__(self) -> str:
87 + def quant_key_str(quant_key: QuantKey | None) -> str:
88 + if quant_key is None:
89 + return "None"
90 + return next(
91 + (
92 + name
93 + for name, known_quant_key in QUANT_KEY_NAMES.items()
94 + if known_quant_key == quant_key
95 + ),
96 + str(quant_key),
97 + )
98 +
99 + return quant_key_str(self.weight)
100 +
101
102 @config
103 class QuantizationConfigArgs:
@@ -92,7 +114,16 @@ class QuantizationConfigArgs:
114 """Spec applied to ``FusedMoEFactory`` layers."""
115
116 ignore: list[str] = Field(default_factory=list)
95 - """Layers to skip quantization for."""
117 + """Layers to skip quantization for. Online quantization also supports
118 + fnmatch-style patterns."""
119 +
120 + targets: dict[str, str] | None = None
121 + """Per-layer online quantization overrides, keyed by exact layer name or
122 + regex patterns with a `re:`, or fnmatch-style patterns for online
123 + quantization, mapping to an online shorthand name (see
124 + `_ONLINE_SHORTHANDS`). A layer that matches no pattern is left unquantized.
125 + Mutually exclusive with `linear` and `moe`.
126 + """
127
128 @field_validator("linear", "moe", mode="before")
129 @classmethod
@@ -110,6 +141,46 @@ class QuantizationConfigArgs:
141 return spec
142 return QuantSpec(weight=_coerce_quant_key(v))
143
144 + @field_validator("targets", mode="before")
145 + @classmethod
146 + def _validate_targets(cls, v: Any) -> Any:
147 + if v is None:
148 + return v
149 + if not isinstance(v, dict):
150 + raise TypeError(f"targets must be a dict, got {type(v).__name__}")
151 + for pattern, shorthand in v.items():
152 + if not isinstance(pattern, str):
153 + raise ValueError(
154 + f"targets keys must be strings, got {type(pattern).__name__}"
155 + )
156 + if not isinstance(shorthand, str) or shorthand not in _ONLINE_SHORTHANDS:
157 + raise ValueError(
158 + f"targets[{pattern}] = {shorthand} is not a valid "
159 + f"online shorthand name; expected one of "
160 + f"{sorted(_ONLINE_SHORTHANDS)}"
161 + )
162 + if pattern.startswith("re:"):
163 + try:
164 + re.compile(pattern[3:])
165 + except re.error as e:
166 + raise ValueError(
167 + f"targets key {pattern} is not a valid regex: {e}"
168 + ) from e
169 + return v
170 +
171 + @model_validator(mode="after")
172 + def _validate_targets_exclusivity(self) -> "QuantizationConfigArgs":
173 + if self.targets is None:
174 + return self
175 + if self.linear is not None or self.moe is not None:
176 + raise ValueError(
177 + "quantization_config.targets is mutually exclusive with "
178 + f"quantization_config.linear/moe, got "
179 + f"targets={self.targets}, linear={self.linear}, "
180 + f"moe={self.moe}."
181 + )
182 + return self
183 +
184
185 # CLI shorthands accepted by `--quantization`. Each desugars to a full
186 # QuantizationConfigArgs; activation overrides go through quantization_config.
@@ -191,4 +262,5 @@ def resolve_quantization_config(
262 linear=quantization_config.linear or base.linear,
263 moe=quantization_config.moe or base.moe,
264 ignore=quantization_config.ignore or base.ignore,
265 + targets=quantization_config.targets or base.targets,
266 )
vllm/model_executor/layers/quantization/compressed_tensors/utils.py
+16 -77
@@ -4,11 +4,14 @@
4 from collections.abc import Iterable, Mapping
5 from types import MappingProxyType
6
7 -import regex as re
7 from compressed_tensors import CompressionFormat
8 from compressed_tensors.quantization import QuantizationStrategy
9 from torch.nn import Module
10
11 +from vllm.model_executor.layers.quantization.utils.config_utils import (
12 + find_matching_patterns,
13 + is_equal_or_regex_match,
14 +)
15 from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 kFp8Static128BlockSym,
17 kFp8StaticChannelSym,
@@ -51,63 +54,20 @@ def should_ignore_layer(
54 layer_name: str | None,
55 ignore: Iterable[str] = tuple(),
56 fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
57 + use_fnmatch: bool = False,
58 ) -> bool:
59 if layer_name is None:
60 return False
57 -
58 - # layer_name = model.layers.0.self_attn.qkv_proj
59 - # proj_name = qkv_proj
60 - proj_name = layer_name.split(".")[-1]
61 -
62 - # Fused layers like gate_up_proj or qkv_proj will not be fused
63 - # in the safetensors checkpoint. So, we convert the name
64 - # from the fused version to unfused + check to make sure that
65 - # each shard of the fused layer has the same scheme.
66 - if proj_name in fused_mapping and layer_name not in ignore:
67 - shard_proj_names = fused_mapping[proj_name]
68 -
69 - # Convert fused_name --> [shard_names]
70 - shard_names = [
71 - layer_name.replace(proj_name, shard_proj_name)
72 - for shard_proj_name in shard_proj_names
73 - ]
74 -
75 - # Layer should be ignored if shards are ignored.
76 - should_ignore_layer = None
77 - for shard_name in shard_names:
78 - should_ignore_shard = check_equal_or_regex_match(
79 - layer_name=shard_name, targets=ignore
80 - )
81 -
82 - # If shard_idx=0, set layer ignore to match shard.
83 - if should_ignore_layer is None:
84 - should_ignore_layer = should_ignore_shard
85 -
86 - # If shard_idx=1+ confirm scheme matches prior shards.
87 - elif should_ignore_shard != should_ignore_layer:
88 - raise ValueError(
89 - f"Found a different quantization schemes for "
90 - f"{shard_proj_names} in {layer_name}. vLLM "
91 - "requires all to use the same scheme."
92 - )
93 -
94 - # Unfused layers like down_proj and o_proj will match
95 - # the safetensors checkpoint already.
96 - else:
97 - should_ignore_layer = check_equal_or_regex_match(
98 - layer_name=layer_name, targets=ignore
61 + per_shard_matches = find_matching_patterns(
62 + layer_name, ignore, fused_mapping, use_fnmatch=use_fnmatch
63 + )
64 + shards_ignored = [len(matches) > 0 for matches in per_shard_matches]
65 + if any(shards_ignored) and not all(shards_ignored):
66 + raise ValueError(
67 + f"Found different quantization schemes for the shards of "
68 + f"{layer_name}. vLLM requires all to use the same scheme."
69 )
100 -
101 - assert should_ignore_layer is not None
102 - return should_ignore_layer
103 -
104 -
105 -def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:
106 - """
107 - Checks whether a layer_name is exactly equal or a regex match for
108 - if target starts with 're:' to any target in list.
109 - """
110 - return any(_is_equal_or_regex_match(layer_name, target) for target in targets)
70 + return all(shards_ignored)
71
72
73 def find_matched_target(
@@ -167,32 +127,11 @@ def _find_first_match(
127 """
128
129 for target in targets:
170 - if _is_equal_or_regex_match(value, target, check_contains=check_contains):
130 + if is_equal_or_regex_match(value, target, check_contains=check_contains):
131 return target
132 return None
133
134
175 -def _is_equal_or_regex_match(
176 - value: str, target: str, check_contains: bool = False
177 -) -> bool:
178 - """
179 - Checks whether a value is exactly equal or a regex match for target
180 - if target starts with 're:'. If check_contains is set to True,
181 - additionally checks if the target string is contained within the value.
182 - """
183 -
184 - if target.startswith("re:"):
185 - pattern = target[3:]
186 - if re.match(pattern, value):
187 - return True
188 - elif check_contains:
189 - if target.lower() in value.lower():
190 - return True
191 - elif target == value:
192 - return True
193 - return False
194 -
195 -
135 def _match_fused_layer(
136 layer_name: str,
137 target_layers: Iterable[str],
@@ -230,7 +169,7 @@ def _match_fused_layer(
169 unfused_matches: list[str | None] = []
170 for unfused in unfused_paths:
171 for target in target_layers:
233 - if _is_equal_or_regex_match(unfused, target):
172 + if is_equal_or_regex_match(unfused, target):
173 unfused_matches.append(target)
174 break
175 else:
vllm/model_executor/layers/quantization/online/base.py
+155 -10
@@ -1,11 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from collections.abc import Mapping
5 +from types import MappingProxyType
6 from typing import Any
7
8 import torch
9
8 -from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
10 +from vllm.config.quantization import (
11 + _ONLINE_SHORTHANDS,
12 + QuantizationConfigArgs,
13 + QuantSpec,
14 +)
15 from vllm.logger import init_logger
16 from vllm.model_executor.layers.fused_moe import (
17 RoutedExperts,
@@ -47,6 +53,10 @@ from vllm.model_executor.layers.quantization.online.mxfp8 import (
53 from vllm.model_executor.layers.quantization.online.nvfp4 import (
54 Nvfp4OnlineMoEMethod,
55 )
56 +from vllm.model_executor.layers.quantization.utils.config_utils import (
57 + find_matching_patterns,
58 + get_layer_name_after_index,
59 +)
60 from vllm.model_executor.layers.quantization.utils.quant_utils import (
61 QuantKey,
62 kFp8Static128BlockSym,
@@ -83,6 +93,39 @@ _ONLINE_MOE_METHODS: dict[QuantKey, type] = {
93 }
94
95
96 +def _find_matching_targets(
97 + prefix: str,
98 + targets: Mapping[str, str],
99 + fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
100 +) -> list[str]:
101 + per_shard_matches = find_matching_patterns(
102 + prefix, targets, fused_mapping, use_fnmatch=True
103 + )
104 + if all(len(matches) == 0 for matches in per_shard_matches):
105 + return []
106 + if any(len(matches) == 0 for matches in per_shard_matches):
107 + raise ValueError(
108 + f"Found unmatched shards for {prefix}: {per_shard_matches}. vLLM "
109 + "requires all shards of a fused layer to match a target."
110 + )
111 + if any(len(matches) > 1 for matches in per_shard_matches):
112 + raise ValueError(
113 + f"Found multiple quantization_config.targets matches for the "
114 + f"shards of {prefix}: {per_shard_matches}. Each shard may match "
115 + "at most one target."
116 + )
117 +
118 + matched_patterns = [next(iter(matches)) for matches in per_shard_matches]
119 + quant_key_strs = {targets[pattern] for pattern in matched_patterns}
120 + if len(quant_key_strs) > 1:
121 + raise ValueError(
122 + f"Found different quantization_config.targets values for the "
123 + f"shards of {prefix}: {matched_patterns}. vLLM requires all "
124 + "shards of a fused layer to use the same target."
125 + )
126 + return [matched_patterns[0]]
127 +
128 +
129 class OnlineQuantizationConfig(QuantizationConfig):
130 """Model-level config for online quantization (quantize fp16/bf16 weights
131 during model loading, without requiring a pre-quantized checkpoint)."""
@@ -92,14 +135,44 @@ class OnlineQuantizationConfig(QuantizationConfig):
135 args: QuantizationConfigArgs,
136 ) -> None:
137 super().__init__()
95 - if args.linear is None and args.moe is None:
138 + if args.linear is None and args.moe is None and args.targets is None:
139 raise ValueError(
140 "OnlineQuantizationConfig requires at least one of "
98 - "quantization_config.linear or quantization_config.moe "
99 - "to be set."
141 + "quantization_config.linear, quantization_config.moe, or "
142 + "quantization_config.targets to be set."
143 )
144 self.args = args
145 self.ignored_layers: list[str] = args.ignore
146 + self.quantized_layers: dict[str, tuple[str, str, str | None]] = {}
147 +
148 + @property
149 + def quantized_layer_summaries(self) -> list[str]:
150 + counts: dict[tuple[str, str, str | None, str], int] = {}
151 + for layer_name, (
152 + source,
153 + quant_key_str,
154 + target_pattern,
155 + ) in self.quantized_layers.items():
156 + key = (
157 + get_layer_name_after_index(layer_name),
158 + source,
159 + target_pattern,
160 + quant_key_str,
161 + )
162 + counts[key] = counts.get(key, 0) + 1
163 +
164 + summaries = []
165 + # Build summary entries as
166 + # `self_attn.o_proj: 24 (from targets: re:.*self_attn\.o_proj, mxfp4`
167 + for (layer_type, source, target_pattern, quant_key_str), count in sorted(
168 + counts.items()
169 + ):
170 + pattern_prefix = f"{target_pattern}, " if target_pattern else ""
171 + summaries.append(
172 + f"{layer_type}: {count} "
173 + f"(from {source}: {pattern_prefix}{quant_key_str})"
174 + )
175 + return summaries
176
177 @classmethod
178 def get_name(cls) -> QuantizationMethods:
@@ -154,29 +227,101 @@ class OnlineQuantizationConfig(QuantizationConfig):
227 return cls(layer=layer)
228 return cls()
229
230 + def _dispatch_target(
231 + self, prefix: str, layer: torch.nn.Module
232 + ) -> "QuantizeMethodBase | None":
233 + assert self.args.targets is not None
234 + ignored = should_ignore_layer(
235 + prefix,
236 + ignore=self.ignored_layers,
237 + fused_mapping=self.packed_modules_mapping,
238 + use_fnmatch=True,
239 + )
240 + matches = _find_matching_targets(
241 + prefix, self.args.targets, fused_mapping=self.packed_modules_mapping
242 + )
243 + if ignored and matches:
244 + raise ValueError(
245 + f"Layer {prefix} matches both quantization_config.ignore "
246 + f"and quantization_config.targets ({matches}); a layer may "
247 + f"not be referenced by both."
248 + )
249 + if ignored or not matches:
250 + return None
251 + if len(matches) > 1:
252 + raise ValueError(
253 + f"Layer {prefix} matches multiple quantization_config."
254 + f"targets patterns: {matches}. Each layer may match at most "
255 + f"one target."
256 + )
257 + quant_key_str = self.args.targets[matches[0]]
258 + shorthand = _ONLINE_SHORTHANDS[quant_key_str]
259 + if isinstance(layer, LinearBase):
260 + quant_spec = shorthand.linear
261 + table = _ONLINE_LINEAR_METHODS
262 + elif isinstance(layer, RoutedExperts):
263 + quant_spec = shorthand.moe
264 + table = _ONLINE_MOE_METHODS
265 + else:
266 + raise ValueError(
267 + f"Layer {prefix} was matched by quantization_config.targets "
268 + f"({matches[0]}), but online quantization is not supported for "
269 + f"{type(layer).__name__}."
270 + )
271 + if quant_spec is None:
272 + raise ValueError(
273 + f"targets pattern {matches[0]} = {quant_key_str} does "
274 + f"not define a QuantSpec for {type(layer).__name__} layers "
275 + f"(matched at {prefix})."
276 + )
277 + method = self._dispatch(quant_spec, table, layer)
278 + if method is not None:
279 + self.quantized_layers[prefix] = ("targets", quant_key_str, matches[0])
280 + return method
281 +
282 def get_quant_method(
283 self, layer: torch.nn.Module, prefix: str
284 ) -> "QuantizeMethodBase | None":
285 + # `targets` takes precedence over `moe` and `linear` and is exclusive.
286 + if self.args.targets is not None:
287 + method = self._dispatch_target(prefix, layer)
288 + if method is not None:
289 + return method
290 +
291 + if isinstance(layer, LinearBase):
292 + return UnquantizedLinearMethod()
293 + if isinstance(layer, RoutedExperts):
294 + return UnquantizedFusedMoEMethod(layer.moe_config)
295 + return None
296 +
297 if isinstance(layer, LinearBase):
298 if should_ignore_layer(
299 prefix,
300 ignore=self.ignored_layers,
301 fused_mapping=self.packed_modules_mapping,
302 + use_fnmatch=True,
303 ):
304 return UnquantizedLinearMethod()
305 method = self._dispatch(self.args.linear, _ONLINE_LINEAR_METHODS, layer)
168 - return method if method is not None else UnquantizedLinearMethod()
306 + if method is not None:
307 + self.quantized_layers[prefix] = (
308 + "linear",
309 + str(self.args.linear),
310 + None,
311 + )
312 + return method
313 + return UnquantizedLinearMethod()
314 elif isinstance(layer, RoutedExperts):
315 if should_ignore_layer(
316 prefix,
317 ignore=self.ignored_layers,
318 fused_mapping=self.packed_modules_mapping,
319 + use_fnmatch=True,
320 ):
321 return UnquantizedFusedMoEMethod(layer.moe_config)
322 method = self._dispatch(self.args.moe, _ONLINE_MOE_METHODS, layer)
177 - return (
178 - method
179 - if method is not None
180 - else UnquantizedFusedMoEMethod(layer.moe_config)
181 - )
323 + if method is not None:
324 + self.quantized_layers[prefix] = ("moe", str(self.args.moe), None)
325 + return method
326 + return UnquantizedFusedMoEMethod(layer.moe_config)
327 return None
vllm/model_executor/layers/quantization/quark/utils.py
+15 -74
@@ -5,9 +5,12 @@ from collections.abc import Iterable, Mapping
5 from types import MappingProxyType
6 from typing import Any
7
8 -import regex as re
8 import torch
9
10 +from vllm.model_executor.layers.quantization.utils.config_utils import (
11 + find_matching_patterns,
12 +)
13 +
14 QuarkQTensorHint = dict[str, Any] | list[dict[str, Any]] | None
15
16
@@ -49,80 +52,18 @@ def should_ignore_layer(
52 ):
53 return True
54
52 - # layer_name = model.layers.0.self_attn.qkv_proj
53 - # proj_name = qkv_proj
54 - proj_name = layer_name.split(".")[-1]
55 -
56 - # Fused layers like gate_up_proj or qkv_proj will not be fused
57 - # in the safetensors checkpoint. So, we convert the name
58 - # from the fused version to unfused + check to make sure that
59 - # each shard of the fused layer has the same scheme.
60 - if proj_name in fused_mapping:
61 - shard_proj_names = fused_mapping[proj_name]
62 -
63 - # Convert fused_name --> [shard_names]
64 - shard_names = [
65 - layer_name.replace(proj_name, shard_proj_name)
66 - for shard_proj_name in shard_proj_names
67 - ]
68 -
69 - # Layer should be ignored if shards are ignored.
70 - should_ignore_layer = None
71 - for shard_name in shard_names:
72 - should_ignore_shard = check_equal_or_regex_match(
73 - layer_name=shard_name, targets=ignore
74 - )
75 -
76 - # If shard_idx=0, set layer ignore to match shard.
77 - if should_ignore_layer is None:
78 - should_ignore_layer = should_ignore_shard
79 -
80 - # If shard_idx=1+ confirm scheme matches prior shards.
81 - elif should_ignore_shard != should_ignore_layer:
82 - raise ValueError(
83 - f"Found a different quantization schemes for "
84 - f"{shard_proj_names} in {layer_name}. vLLM "
85 - "requires all to use the same scheme."
86 - )
87 -
88 - # Unfused layers like down_proj and o_proj will match
89 - # the safetensors checkpoint already.
90 - else:
91 - should_ignore_layer = check_equal_or_regex_match(
92 - layer_name=layer_name, targets=ignore
55 + # A direct fused-layer pattern takes precedence over expansion. For
56 + # model.layers.0.self_attn.qkv_proj,
57 + # ignore=["re:.*qkv_proj.*"] yields [{"re:.*qkv_proj.*"}]. In contrast,
58 + # ignore=["re:.*[qkv]_proj"] yields one matching set per expanded shard.
59 + per_shard_matches = find_matching_patterns(layer_name, ignore, fused_mapping)
60 + shards_ignored = [len(matches) > 0 for matches in per_shard_matches]
61 + if any(shards_ignored) and not all(shards_ignored):
62 + raise ValueError(
63 + f"Found different quantization schemes for the shards of "
64 + f"{layer_name}. vLLM requires all to use the same scheme."
65 )
94 -
95 - assert should_ignore_layer is not None
96 - return should_ignore_layer
97 -
98 -
99 -def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:
100 - """
101 - Checks whether a layer_name is exactly equal or a regex match for
102 - if target starts with 're:' to any target in list.
103 - """
104 - return any(_is_equal_or_regex_match(layer_name, target) for target in targets)
105 -
106 -
107 -def _is_equal_or_regex_match(
108 - value: str, target: str, check_contains: bool = False
109 -) -> bool:
110 - """
111 - Checks whether a value is exactly equal or a regex match for target
112 - if target starts with 're:'. If check_contains is set to True,
113 - additionally checks if the target string is contained within the value.
114 - """
115 -
116 - if target.startswith("re:"):
117 - pattern = target[3:]
118 - if re.match(pattern, value):
119 - return True
120 - elif check_contains:
121 - if target.lower() in value.lower():
122 - return True
123 - elif target == value:
124 - return True
125 - return False
66 + return all(shards_ignored)
67
68
69 # utility for tensor dims > 2 cases
vllm/model_executor/layers/quantization/utils/config_utils.py
+76
@@ -1,6 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from collections.abc import Iterable, Mapping
5 +from fnmatch import fnmatch
6 +from types import MappingProxyType
7 from typing import TYPE_CHECKING
8
9 import regex as re
@@ -244,3 +247,76 @@ def is_shared_expert_quant_fse_compatible(
247 "shared-expert FSE quantization compatibility is not implemented for "
248 f"{type(quant_config).__name__}",
249 )
250 +
251 +
252 +def find_matching_patterns(
253 + layer_name: str,
254 + patterns: Iterable[str],
255 + fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
256 + use_fnmatch: bool = False,
257 +) -> list[set[str]]:
258 + """Return matching patterns for a layer or each shard of a fused layer.
259 +
260 + A pattern matching the fused layer directly takes precedence. Otherwise,
261 + return one set of matching patterns for every shard.
262 + """
263 + patterns = list(patterns)
264 + matches = [
265 + pattern
266 + for pattern in patterns
267 + if is_equal_or_regex_match(layer_name, pattern, use_fnmatch=use_fnmatch)
268 + ]
269 + if matches:
270 + return [set(matches)]
271 +
272 + proj_name = layer_name.split(".")[-1]
273 + if proj_name not in fused_mapping:
274 + return [set()]
275 +
276 + shard_names = [
277 + layer_name.replace(proj_name, shard_proj_name)
278 + for shard_proj_name in fused_mapping[proj_name]
279 + ]
280 + per_shard_matches = [
281 + {
282 + pattern
283 + for pattern in patterns
284 + if is_equal_or_regex_match(shard_name, pattern, use_fnmatch=use_fnmatch)
285 + }
286 + for shard_name in shard_names
287 + ]
288 + return per_shard_matches
289 +
290 +
291 +def get_layer_name_after_index(layer_name: str) -> str:
292 + """Return the suffix following the final numeric component of a layer name."""
293 + parts = layer_name.split(".")
294 + for index in range(len(parts) - 1, -1, -1):
295 + if parts[index].isdigit():
296 + return ".".join(parts[index + 1 :])
297 + return layer_name
298 +
299 +
300 +def is_equal_or_regex_match(
301 + value: str,
302 + target: str,
303 + check_contains: bool = False,
304 + use_fnmatch: bool = False,
305 +) -> bool:
306 + """
307 + Checks whether a value is exactly equal or a regex match for target
308 + if target starts with 're:'. If check_contains is set to True,
309 + additionally checks if the target string is contained within the value.
310 + If use_fnmatch is set, supports shell-style patterns in target.
311 + """
312 +
313 + if target.startswith("re:"):
314 + pattern = target[3:]
315 + if re.match(pattern, value):
316 + return True
317 + elif check_contains:
318 + if target.lower() in value.lower():
319 + return True
320 + elif target == value or use_fnmatch and fnmatch(value, target):
321 + return True
322 + return False
vllm/model_executor/model_loader/base_loader.py
+18
@@ -58,6 +58,7 @@ class BaseModelLoader(ABC):
58 prefix=prefix,
59 )
60
61 + log_online_quantization(vllm_config)
62 log_model_inspection(model)
63
64 logger.debug("Loading weights on %s ...", load_device)
@@ -92,6 +93,23 @@ def log_model_inspection(model: nn.Module) -> None:
93 logger.info("vLLM model structure:\n%s", format_model_inspection(model))
94
95
96 +def log_online_quantization(vllm_config: VllmConfig) -> None:
97 + """Log the online-quantized layer count and types, when applicable."""
98 + from vllm.model_executor.layers.quantization.online.base import (
99 + OnlineQuantizationConfig,
100 + )
101 +
102 + quant_config = vllm_config.quant_config
103 + if not isinstance(quant_config, OnlineQuantizationConfig):
104 + return
105 +
106 + logger.info(
107 + "Quantized %d layers of types: %s",
108 + len(quant_config.quantized_layers),
109 + "; ".join(quant_config.quantized_layer_summaries),
110 + )
111 +
112 +
113 def _has_online_quant(model: nn.Module):
114 for module in model.modules():
115 quant_method = getattr(module, "quant_method", None)