| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 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 tests.quantization.utils import quant_config_args, quant_spec |
| 10 | from vllm.config.quantization import ( |
| 11 | QUANT_KEY_NAMES, |
| 12 | QuantizationConfigArgs, |
| 13 | QuantSpec, |
| 14 | resolve_quantization_config, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.linear import LinearBase |
| 17 | from vllm.model_executor.layers.quantization.online.base import ( |
| 18 | OnlineQuantizationConfig, |
| 19 | _find_matching_targets, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 22 | kFp8Dynamic128Sym, |
| 23 | kFp8DynamicTokenSym, |
| 24 | kFp8Static128BlockSym, |
| 25 | kFp8StaticTensorSym, |
| 26 | kInt8StaticChannelSym, |
| 27 | kMxfp8Dynamic, |
| 28 | ) |
| 29 | |
| 30 | # ---- QuantSpec ------------------------------------------------------------ |
| 31 | |
| 32 | |
| 33 | def test_quant_spec_resolves_string_to_quant_key(): |
| 34 | spec = quant_spec(weight="mxfp8", activation="fp8_per_token") |
| 35 | assert spec.weight == kMxfp8Dynamic |
| 36 | assert spec.activation == kFp8DynamicTokenSym |
| 37 | |
| 38 | |
| 39 | def test_quant_spec_accepts_quant_key_directly(): |
| 40 | spec = QuantSpec(weight=kFp8StaticTensorSym) |
| 41 | assert spec.weight is kFp8StaticTensorSym |
| 42 | assert spec.activation is None |
| 43 | |
| 44 | |
| 45 | def test_quant_spec_string_representation_uses_quantization_name(): |
| 46 | assert str(quant_spec(weight="mxfp4")) == "mxfp4" |
| 47 | |
| 48 | |
| 49 | def test_quant_spec_rejects_unknown_name(): |
| 50 | with pytest.raises(ValueError, match="unknown quantization name"): |
| 51 | quant_spec(weight="not_a_real_format") |
| 52 | |
| 53 | |
| 54 | # ---- QuantizationConfigArgs string shorthand on linear/moe ---------------- |
| 55 | |
| 56 | |
| 57 | def test_args_linear_string_resolves_via_quant_key_names(): |
| 58 | # A bare QUANT_KEY_NAMES entry desugars to QuantSpec(weight=<key>). |
| 59 | args = quant_config_args(linear="fp8_per_block_static") |
| 60 | assert args.linear == QuantSpec(weight=kFp8Static128BlockSym) |
| 61 | assert args.moe is None |
| 62 | |
| 63 | |
| 64 | def test_args_moe_string_resolves_via_online_shorthand(): |
| 65 | # An online-shorthand name pulls the matching slot from _ONLINE_SHORTHANDS |
| 66 | # (so `linear: "fp8_per_block"` and `moe: "fp8_per_block"` produce the |
| 67 | # same per-layer-kind spec the `--quantization fp8_per_block` shorthand |
| 68 | # would). |
| 69 | args = quant_config_args(moe="fp8_per_block") |
| 70 | assert args.moe == QuantSpec(weight=kFp8Static128BlockSym) |
| 71 | |
| 72 | |
| 73 | def test_args_string_shorthand_missing_slot_raises(): |
| 74 | # int8_per_channel_weight_only sets only `moe`; using it on `linear` |
| 75 | # has no defined spec and should raise rather than silently no-op. |
| 76 | with pytest.raises(ValueError, match="does not define a linear spec"): |
| 77 | quant_config_args(linear="int8_per_channel_weight_only") |
| 78 | |
| 79 | |
| 80 | def test_args_accepts_dict_form(): |
| 81 | args = quant_config_args(moe={"activation": "mxfp8"}) |
| 82 | assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic) |
| 83 | |
| 84 | |
| 85 | def test_targets_reject_non_string_keys(): |
| 86 | with pytest.raises(ValueError, match="targets keys must be strings"): |
| 87 | QuantizationConfigArgs._validate_targets({123: "mxfp8"}) |
| 88 | |
| 89 | |
| 90 | # ---- resolve_quantization_config ----------------------------------------- |
| 91 | |
| 92 | |
| 93 | def test_resolve_shorthand_only_populates_both_slots(): |
| 94 | args = resolve_quantization_config("fp8_per_block", None) |
| 95 | assert args is not None |
| 96 | assert args.linear == QuantSpec(weight=kFp8Static128BlockSym) |
| 97 | assert args.moe == QuantSpec(weight=kFp8Static128BlockSym) |
| 98 | |
| 99 | |
| 100 | @pytest.mark.parametrize("quantization", ["mxfp4", "mxfp8"]) |
| 101 | def test_resolve_colliding_shorthand_is_deferred(quantization: str): |
| 102 | """Checkpoint metadata determines whether an MXFP shorthand is online.""" |
| 103 | assert resolve_quantization_config(quantization, None) is None |
| 104 | |
| 105 | |
| 106 | def test_resolve_int8_shorthand_leaves_linear_unset(): |
| 107 | # int8_per_channel_weight_only is MoE-only; linear stays None so that |
| 108 | # OnlineQuantizationConfig leaves Linear layers in full precision. |
| 109 | args = resolve_quantization_config("int8_per_channel_weight_only", None) |
| 110 | assert args is not None |
| 111 | assert args.linear is None |
| 112 | assert args.moe == QuantSpec(weight=kInt8StaticChannelSym) |
| 113 | |
| 114 | |
| 115 | def test_resolve_quantization_config_only(): |
| 116 | # When only `quantization_config` is given (e.g. for an already-quantized |
| 117 | # checkpoint that needs an activation override), it's returned as-is. |
| 118 | args = resolve_quantization_config(None, {"moe": {"activation": "mxfp8"}}) |
| 119 | assert args is not None |
| 120 | assert args.linear is None |
| 121 | assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic) |
| 122 | |
| 123 | |
| 124 | def test_resolve_merges_explicit_over_shorthand(): |
| 125 | # Explicit linear in quantization_config wins; moe falls back to the |
| 126 | # shorthand's slot. |
| 127 | args = resolve_quantization_config( |
| 128 | "fp8_per_tensor", |
| 129 | {"linear": "fp8_per_block"}, |
| 130 | ) |
| 131 | assert args is not None |
| 132 | assert args.linear == QuantSpec(weight=kFp8Static128BlockSym) |
| 133 | assert args.moe == QuantSpec(weight=kFp8StaticTensorSym) |
| 134 | |
| 135 | |
| 136 | def test_resolve_quantization_config_with_checkpoint_quantization(): |
| 137 | args = resolve_quantization_config("gptq", {"linear": "fp8_per_block"}) |
| 138 | assert args == quant_config_args(linear="fp8_per_block") |
| 139 | |
| 140 | |
| 141 | # ---- QUANT_KEY_NAMES coverage -------------------------------------------- |
| 142 | |
| 143 | |
| 144 | def test_quant_key_names_round_trip(): |
| 145 | # Every advertised name should round-trip through QuantSpec without error |
| 146 | # and produce the same QuantKey it maps to. |
| 147 | for name, expected in QUANT_KEY_NAMES.items(): |
| 148 | assert quant_spec(weight=name).weight == expected, name |
| 149 | assert quant_spec(activation=name).activation == expected, name |
| 150 | |
| 151 | |
| 152 | def test_static_block_weight_paired_with_dynamic_block_activation(): |
| 153 | # The block-FP8 shorthand pair: 128x128 static weights + 1x128 dynamic |
| 154 | # activations. Pinning this so renames in QUANT_KEY_NAMES don't quietly |
| 155 | # rewire the kernel dispatch. |
| 156 | spec = quant_spec(weight="fp8_per_block_static", activation="fp8_per_block_dynamic") |
| 157 | assert spec.weight == kFp8Static128BlockSym |
| 158 | assert spec.activation == kFp8Dynamic128Sym |
| 159 | |
| 160 | |
| 161 | def test_targets_allow_distinct_patterns_with_the_same_shorthand(): |
| 162 | layer_name = "model.layers.0.self_attn.qkv_proj" |
| 163 | fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} |
| 164 | targets = { |
| 165 | r"re:.*\.q_proj$": "mxfp8", |
| 166 | r"re:.*\.k_proj$": "mxfp8", |
| 167 | r"re:.*\.v_proj$": "mxfp8", |
| 168 | } |
| 169 | |
| 170 | matches = _find_matching_targets(layer_name, targets, fused_mapping) |
| 171 | assert len(matches) == 1 |
| 172 | assert targets[matches[0]] == "mxfp8" |
| 173 | |
| 174 | |
| 175 | def test_targets_reject_overlapping_patterns(): |
| 176 | targets = { |
| 177 | r"re:.*o_proj": "fp8_per_tensor", |
| 178 | "model.layers.0.self_attn.o_proj": "fp8_per_block", |
| 179 | } |
| 180 | |
| 181 | with pytest.raises(ValueError, match="multiple quantization_config.targets"): |
| 182 | _find_matching_targets("model.layers.0.self_attn.o_proj", targets) |
| 183 | |
| 184 | |
| 185 | def test_targets_reject_partially_matched_fused_layer(): |
| 186 | targets = {r"re:.*q_proj": "fp8_per_tensor"} |
| 187 | fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} |
| 188 | |
| 189 | with pytest.raises(ValueError, match="unmatched shards"): |
| 190 | _find_matching_targets( |
| 191 | "model.layers.0.self_attn.qkv_proj", targets, fused_mapping |
| 192 | ) |
| 193 | |
| 194 | |
| 195 | def test_targets_reject_fused_shards_with_different_schemes(): |
| 196 | targets = { |
| 197 | r"re:.*\.q_proj$": "fp8_per_tensor", |
| 198 | r"re:.*\.k_proj$": "fp8_per_block", |
| 199 | r"re:.*\.v_proj$": "fp8_per_tensor", |
| 200 | } |
| 201 | fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} |
| 202 | |
| 203 | with pytest.raises(ValueError, match="different quantization_config.targets"): |
| 204 | _find_matching_targets( |
| 205 | "model.layers.0.self_attn.qkv_proj", targets, fused_mapping |
| 206 | ) |
| 207 | |
| 208 | |
| 209 | def test_targets_reject_moe_only_shorthand_for_linear_layer(): |
| 210 | config = OnlineQuantizationConfig( |
| 211 | QuantizationConfigArgs( |
| 212 | targets={"model.layers.0.self_attn.o_proj": "nvfp4_per_token"} |
| 213 | ) |
| 214 | ) |
| 215 | |
| 216 | with pytest.raises(ValueError, match="does not define a QuantSpec"): |
| 217 | config.resolve_quant_method_cls( |
| 218 | Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj" |
| 219 | ) |