A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 17:45 UTC
[Quantization] Rework quantization_config to use QuantKey and allow for activation override (#41566)
Michael Goin committed
May 13, 2026 at 16:58 UTC
8efd508204e5a3f55a27407f39869ccd8b1c1af9
16 files changed
+626
-327
docs/features/quantization/online.md
+45
-3
@@ -42,18 +42,59 @@ vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
42
43
For fine-grained control, use a `quantization_config` dictionary.
44
45
+### Schema
46
+
47
+```yaml
48
+quantization_config:
49
+ linear:
50
+ weight: <name> # see QUANT_KEY_NAMES in vllm/config/quantization.py
51
+ activation: <name>
52
+ moe:
53
+ weight: <name>
54
+ activation: <name>
55
+ ignore: [<layer-name-or-regex>, ...]
56
+```
57
+
58
+`linear` and `moe` accept a full `{weight, activation}` dict, or a bare
59
+string. A string resolves first against the `--quantization` shorthands
60
+(taking the matching layer-kind slot), then against `QUANT_KEY_NAMES` as a
61
+weight name. Unset fields fall back to the `--quantization` shorthand's
62
+defaults, or for already-quantized checkpoints to whatever the checkpoint
63
+declares.
64
+
65
+The CLI accepts the same shape as JSON or as dotted keys:
66
+
67
+```bash
68
+vllm serve <model> --quantization-config '{"moe":{"activation":"mxfp8"}}'
69
+vllm serve <model> --quantization-config.moe.activation mxfp8
70
+```
71
+
72
+### Activation overrides on already-quantized checkpoints
73
+
74
+For checkpoint-quantized models, `quantization_config` lets you pick an
75
+activation format independently of the baked-in weights. The supported
76
+overrides are checkpoint-specific; today this is wired up for MXFP4 MoE
77
+checkpoints (gpt-oss) where you can opt into FP8 activations:
78
+
79
+```bash
80
+vllm serve openai/gpt-oss-20b --quantization-config.moe.activation mxfp8
81
+```
82
+
83
+Combine with `--moe-backend` to pin a specific kernel family.
84
+
85
### Separate Schemes for Dense and MoE Layers
86
47
-You can apply different quantization schemes to dense linear layers and MoE expert layers:
87
+You can apply different quantization schemes to dense linear layers and MoE expert layers via the `linear` and `moe` fields. Each accepts either a full spec dict, or a bare string naming an online shorthand (e.g. `"fp8_per_block"`) or weight format (e.g. `"fp8_per_block_static"`); fields not set fall back to the shorthand defaults.
88
89
```python
90
from vllm import LLM
91
92
+# Linear: per-block FP8; MoE: per-tensor FP8 (inherited from the shorthand)
93
llm = LLM(
94
"ibm-granite/granite-3.0-1b-a400m-base",
95
quantization="fp8_per_tensor",
96
quantization_config={
56
- "linear_scheme_override": "fp8_per_block",
97
+ "linear": "fp8_per_block",
98
},
99
)
100
```
@@ -63,11 +104,12 @@ Or,
104
```python
105
from vllm import LLM
106
107
+# Linear: per-tensor FP8 (inherited); MoE: per-block FP8
108
llm = LLM(
109
"ibm-granite/granite-3.0-1b-a400m-base",
110
quantization="fp8_per_tensor",
111
quantization_config={
70
- "moe_scheme_override": "fp8_per_block",
112
+ "moe": "fp8_per_block",
113
},
114
)
115
```
tests/compile/fusions_e2e/conftest.py
-7
@@ -103,13 +103,6 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
103
if not supports_trtllm_attention():
104
matches = matches._replace(attn_quant_fusion=0)
105
106
- # TODO: remove this after finishing migration from envs to model kwargs
107
- if model_name == "openai/gpt-oss-20b":
108
- from .common import is_blackwell
109
-
110
- if is_blackwell():
111
- monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
112
-
106
# Disable, compile cache to make sure custom passes run.
107
# Otherwise, we can't verify fusion happened through the logs.
108
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
tests/compile/fusions_e2e/models.py
+5
@@ -214,4 +214,9 @@ gpt_oss_20b = ModelFusionInfo(
214
sequence_parallel=n_layers * 2 + 1,
215
async_tp=n_layers * 2,
216
),
217
+ model_kwargs=(
218
+ {"quantization_config": {"moe": {"activation": "mxfp8"}}}
219
+ if is_blackwell()
220
+ else {}
221
+ ),
222
)
tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml
+1
-3
@@ -3,6 +3,4 @@
3
model_name: "openai/gpt-oss-20b"
4
metric_threshold: 0.568
5
reasoning_effort: "low"
6
-server_args: "--tensor-parallel-size 2"
7
-env:
8
- VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: "1"
6
+server_args: "--tensor-parallel-size 2 --moe-backend flashinfer_cutlass --quantization-config.moe.activation mxfp8"
tests/evals/gpt_oss/configs/gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml
+1
-3
@@ -3,6 +3,4 @@
3
model_name: "openai/gpt-oss-20b"
4
metric_threshold: 0.568
5
reasoning_effort: "low"
6
-server_args: "--tensor-parallel-size 2"
7
-env:
8
- VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: "1"
6
+server_args: "--tensor-parallel-size 2 --quantization-config.moe.activation mxfp8"
tests/quantization/test_blackwell_moe.py
+18
-4
@@ -190,13 +190,27 @@ def test_gptoss_mxfp4bf16_moe_flashinfer(monkeypatch: pytest.MonkeyPatch):
190
191
192
def test_gptoss_mxfp4mxfp8_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
193
- monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS", "1")
194
- can_initialize("openai/gpt-oss-20b", hf_overrides=HF_OVERRIDE_TEXT)
193
+ can_initialize(
194
+ "openai/gpt-oss-20b",
195
+ hf_overrides=HF_OVERRIDE_TEXT,
196
+ extra_args=[
197
+ "--moe-backend",
198
+ "flashinfer_cutlass",
199
+ "--quantization-config.moe.activation",
200
+ "mxfp8",
201
+ ],
202
+ )
203
204
205
def test_gptoss_mxfp4mxfp8_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
198
- monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
199
- can_initialize("openai/gpt-oss-20b", hf_overrides=HF_OVERRIDE_TEXT)
206
+ can_initialize(
207
+ "openai/gpt-oss-20b",
208
+ hf_overrides=HF_OVERRIDE_TEXT,
209
+ extra_args=[
210
+ "--quantization-config.moe.activation",
211
+ "mxfp8",
212
+ ],
213
+ )
214
215
216
def test_gptoss_eager(monkeypatch: pytest.MonkeyPatch):
tests/quantization/test_online.py
+3
-4
@@ -40,13 +40,12 @@ from vllm.platforms import current_platform
40
Fp8PerBlockOnlineLinearMethod,
41
Fp8PerBlockOnlineMoEMethod,
42
),
43
- # quantization='online with linear_scheme_override and
44
- # moe_scheme_override
43
+ # quantization='online' with per-layer-kind overrides
44
(
45
"online",
46
{
48
- "linear_scheme_override": "fp8_per_block",
49
- "moe_scheme_override": "fp8_per_tensor",
47
+ "linear": "fp8_per_block",
48
+ "moe": "fp8_per_tensor",
49
},
50
Fp8PerBlockOnlineLinearMethod,
51
Fp8PerTensorOnlineMoEMethod,
tests/quantization/test_quantization_config_args.py
new
+134
@@ -0,0 +1,134 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for QuantizationConfigArgs parsing."""
4
+
5
+import pytest
6
+
7
+from vllm.config.quantization import (
8
+ QUANT_KEY_NAMES,
9
+ QuantizationConfigArgs,
10
+ QuantSpec,
11
+ resolve_quantization_config,
12
+)
13
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
+ kFp8Dynamic128Sym,
15
+ kFp8DynamicTokenSym,
16
+ kFp8Static128BlockSym,
17
+ kFp8StaticTensorSym,
18
+ kInt8StaticChannelSym,
19
+ kMxfp8Dynamic,
20
+)
21
+
22
+# ---- QuantSpec ------------------------------------------------------------
23
+
24
+
25
+def test_quant_spec_resolves_string_to_quant_key():
26
+ spec = QuantSpec(weight="mxfp8", activation="fp8_per_token")
27
+ assert spec.weight == kMxfp8Dynamic
28
+ assert spec.activation == kFp8DynamicTokenSym
29
+
30
+
31
+def test_quant_spec_accepts_quant_key_directly():
32
+ spec = QuantSpec(weight=kFp8StaticTensorSym)
33
+ assert spec.weight is kFp8StaticTensorSym
34
+ assert spec.activation is None
35
+
36
+
37
+def test_quant_spec_rejects_unknown_name():
38
+ with pytest.raises(ValueError, match="unknown quantization name"):
39
+ QuantSpec(weight="not_a_real_format")
40
+
41
+
42
+# ---- QuantizationConfigArgs string shorthand on linear/moe ----------------
43
+
44
+
45
+def test_args_linear_string_resolves_via_quant_key_names():
46
+ # A bare QUANT_KEY_NAMES entry desugars to QuantSpec(weight=<key>).
47
+ args = QuantizationConfigArgs(linear="fp8_per_block_static")
48
+ assert args.linear == QuantSpec(weight=kFp8Static128BlockSym)
49
+ assert args.moe is None
50
+
51
+
52
+def test_args_moe_string_resolves_via_online_shorthand():
53
+ # An online-shorthand name pulls the matching slot from _ONLINE_SHORTHANDS
54
+ # (so `linear: "fp8_per_block"` and `moe: "fp8_per_block"` produce the
55
+ # same per-layer-kind spec the `--quantization fp8_per_block` shorthand
56
+ # would).
57
+ args = QuantizationConfigArgs(moe="fp8_per_block")
58
+ assert args.moe == QuantSpec(weight=kFp8Static128BlockSym)
59
+
60
+
61
+def test_args_string_shorthand_missing_slot_raises():
62
+ # int8_per_channel_weight_only sets only `moe`; using it on `linear`
63
+ # has no defined spec and should raise rather than silently no-op.
64
+ with pytest.raises(ValueError, match="does not define a linear spec"):
65
+ QuantizationConfigArgs(linear="int8_per_channel_weight_only")
66
+
67
+
68
+def test_args_accepts_dict_form():
69
+ args = QuantizationConfigArgs(moe={"activation": "mxfp8"})
70
+ assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic)
71
+
72
+
73
+# ---- resolve_quantization_config -----------------------------------------
74
+
75
+
76
+def test_resolve_shorthand_only_populates_both_slots():
77
+ args = resolve_quantization_config("fp8_per_block", None)
78
+ assert args.linear == QuantSpec(weight=kFp8Static128BlockSym)
79
+ assert args.moe == QuantSpec(weight=kFp8Static128BlockSym)
80
+
81
+
82
+def test_resolve_int8_shorthand_leaves_linear_unset():
83
+ # int8_per_channel_weight_only is MoE-only; linear stays None so that
84
+ # OnlineQuantizationConfig leaves Linear layers in full precision.
85
+ args = resolve_quantization_config("int8_per_channel_weight_only", None)
86
+ assert args.linear is None
87
+ assert args.moe == QuantSpec(weight=kInt8StaticChannelSym)
88
+
89
+
90
+def test_resolve_quantization_config_only():
91
+ # When only `quantization_config` is given (e.g. for an already-quantized
92
+ # checkpoint that needs an activation override), it's returned as-is.
93
+ args = resolve_quantization_config(None, {"moe": {"activation": "mxfp8"}})
94
+ assert args.linear is None
95
+ assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic)
96
+
97
+
98
+def test_resolve_merges_explicit_over_shorthand():
99
+ # Explicit linear in quantization_config wins; moe falls back to the
100
+ # shorthand's slot.
101
+ args = resolve_quantization_config(
102
+ "fp8_per_tensor",
103
+ {"linear": "fp8_per_block"},
104
+ )
105
+ assert args.linear == QuantSpec(weight=kFp8Static128BlockSym)
106
+ assert args.moe == QuantSpec(weight=kFp8StaticTensorSym)
107
+
108
+
109
+def test_resolve_rejects_quantization_config_with_non_shorthand_quant():
110
+ # If --quantization names something other than an online shorthand,
111
+ # quantization_config is not allowed via this path (checkpoint quant
112
+ # paths read it directly off ModelConfig instead).
113
+ with pytest.raises(ValueError, match="quantization_config is only supported"):
114
+ resolve_quantization_config("gptq", {"linear": "fp8_per_block"})
115
+
116
+
117
+# ---- QUANT_KEY_NAMES coverage --------------------------------------------
118
+
119
+
120
+def test_quant_key_names_round_trip():
121
+ # Every advertised name should round-trip through QuantSpec without error
122
+ # and produce the same QuantKey it maps to.
123
+ for name, expected in QUANT_KEY_NAMES.items():
124
+ assert QuantSpec(weight=name).weight == expected, name
125
+ assert QuantSpec(activation=name).activation == expected, name
126
+
127
+
128
+def test_static_block_weight_paired_with_dynamic_block_activation():
129
+ # The block-FP8 shorthand pair: 128x128 static weights + 1x128 dynamic
130
+ # activations. Pinning this so renames in QUANT_KEY_NAMES don't quietly
131
+ # rewire the kernel dispatch.
132
+ spec = QuantSpec(weight="fp8_per_block_static", activation="fp8_per_block_dynamic")
133
+ assert spec.weight == kFp8Static128BlockSym
134
+ assert spec.activation == kFp8Dynamic128Sym
vllm/config/model.py
+6
-5
@@ -21,7 +21,7 @@ from vllm.config.multimodal import (
21
MultiModalConfig,
22
)
23
from vllm.config.pooler import PoolerConfig
24
-from vllm.config.quantization import OnlineQuantizationConfigArgs
24
+from vllm.config.quantization import QuantizationConfigArgs
25
from vllm.config.scheduler import RunnerType
26
from vllm.config.utils import config, getattr_iter
27
from vllm.logger import init_logger
@@ -206,10 +206,11 @@ class ModelConfig:
206
`quantization_config` attribute in the model config file. If that is
207
`None`, we assume the model weights are not quantized and use `dtype` to
208
determine the data type of the weights."""
209
- quantization_config: dict[str, Any] | OnlineQuantizationConfigArgs | None = None
210
- """Arguments for online quantization.
211
- Auto-created when `quantization` equals to one of the string values of
212
- the `OnlineQuantScheme` enum."""
209
+ quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None
210
+ """User-facing quantization configuration. Carries per-layer-kind specs
211
+ (linear, moe) and ignore patterns; see :class:`QuantizationConfigArgs`.
212
+ Auto-populated from the matching online shorthand when `quantization` is
213
+ one of the values in `ONLINE_QUANT_SHORTHAND_NAMES`."""
214
allow_deprecated_quantization: bool = False
215
"""Whether to allow deprecated quantization methods."""
216
enforce_eager: bool = False
vllm/config/quantization.py
+143
-94
@@ -1,126 +1,175 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from enum import Enum
5
-from typing import Any
4
+from typing import Annotated, Any
5
7
-from pydantic import Field, field_validator
6
+from pydantic import Field, GetPydanticSchema, ValidationInfo, field_validator
7
+from pydantic_core import core_schema
8
9
from vllm.config.utils import config
10
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
11
+ QuantKey,
12
+ kFp8Dynamic128Sym,
13
+ kFp8DynamicTensorSym,
14
+ kFp8DynamicTokenSym,
15
+ kFp8Static128BlockSym,
16
+ kFp8StaticTensorSym,
17
+ kInt8StaticChannelSym,
18
+ kMxfp4Dynamic,
19
+ kMxfp8Dynamic,
20
+)
21
+
22
+# User-facing names addressable from quantization_config.
23
+QUANT_KEY_NAMES: dict[str, QuantKey] = {
24
+ "fp8_per_tensor_static": kFp8StaticTensorSym,
25
+ "fp8_per_tensor_dynamic": kFp8DynamicTensorSym,
26
+ "fp8_per_token": kFp8DynamicTokenSym,
27
+ "fp8_per_block_static": kFp8Static128BlockSym,
28
+ "fp8_per_block_dynamic": kFp8Dynamic128Sym,
29
+ "mxfp8": kMxfp8Dynamic,
30
+ "mxfp4": kMxfp4Dynamic,
31
+ "int8_per_channel_static": kInt8StaticChannelSym,
32
+}
33
+
34
+
35
+def _coerce_quant_key(v: Any) -> QuantKey | None:
36
+ if v is None or isinstance(v, QuantKey):
37
+ return v
38
+ if not isinstance(v, str):
39
+ raise TypeError(f"expected str or QuantKey, got {type(v).__name__}")
40
+ try:
41
+ return QUANT_KEY_NAMES[v]
42
+ except KeyError:
43
+ raise ValueError(
44
+ f"unknown quantization name {v!r}; "
45
+ f"expected one of {sorted(QUANT_KEY_NAMES)}"
46
+ ) from None
47
+
48
+
49
+# Stop pydantic from introspecting QuantKey: it transitively contains a
50
+# NamedTuple with `ClassVar[GroupShape]` declarations that pydantic refuses.
51
+QuantKeyField = Annotated[
52
+ QuantKey | None,
53
+ GetPydanticSchema(
54
+ lambda _src, _handler: core_schema.no_info_plain_validator_function(
55
+ _coerce_quant_key
56
+ )
57
+ ),
58
+]
59
60
12
-class OnlineQuantScheme(Enum):
13
- """Supported online quantization schemes."""
14
-
15
- # fp8, weights and activations scaled per-tensor
16
- FP8_PER_TENSOR = "fp8_per_tensor"
61
+@config
62
+class QuantSpec:
63
+ """Quantization spec for one layer kind (linear or MoE).
64
18
- # fp8, activations scaled in blocks of 1x128 elements, weights scaled in
19
- # blocks of 128x128 elements (popularized by DeepSeek)
20
- FP8_PER_BLOCK = "fp8_per_block"
65
+ `None` on either side means the method class falls back to its own default
66
+ (typically inherited from the checkpoint, or unquantized for online).
67
+ """
68
22
- # int8, weight-only per-channel quantization for MoE expert weights.
23
- # Linear layers remain unquantized.
24
- INT8_PER_CHANNEL_WEIGHT_ONLY = "int8_per_channel_weight_only"
69
+ weight: QuantKeyField = None
70
+ """Weight quantization key, or a name from QUANT_KEY_NAMES."""
71
26
- # mxfp8, weights scaled in blocks of 1x32 elements (microscaling FP8)
27
- MXFP8 = "mxfp8"
72
+ activation: QuantKeyField = None
73
+ """Activation quantization key, or a name from QUANT_KEY_NAMES."""
74
75
76
@config
31
-class OnlineQuantizationConfigArgs:
32
- """Configuration for online quantization.
77
+class QuantizationConfigArgs:
78
+ """User-facing quantization configuration.
79
34
- Controls how ``OnlineQuantizationConfig`` is applied to a model.
35
- At least one of ``global_scheme``, ``linear_scheme_override``, or
36
- ``moe_scheme_override`` must be set.
80
+ See `docs/features/quantization/online.md` for the schema and shorthand
81
+ string forms accepted on `linear` and `moe`.
82
"""
83
39
- global_scheme: OnlineQuantScheme | None = None
40
- """Quantization scheme applied to every supported layer."""
41
-
42
- linear_scheme_override: OnlineQuantScheme | None = None
43
- """Quantization scheme override for ``LinearBase`` layers."""
84
+ linear: QuantSpec | None = None
85
+ """Spec applied to ``LinearBase`` layers."""
86
45
- moe_scheme_override: OnlineQuantScheme | None = None
46
- """Quantization scheme override for ``FusedMoE`` layers."""
87
+ moe: QuantSpec | None = None
88
+ """Spec applied to ``FusedMoE`` layers."""
89
90
ignore: list[str] = Field(default_factory=list)
49
- """Layers to skip quantization for. Supports exact names and regex
50
- patterns with ``re:`` prefix (e.g. ``re:.*attn.*``), consistent with
51
- compressed_tensors layer skipping."""
91
+ """Layers to skip quantization for."""
92
53
- @field_validator(
54
- "global_scheme", "linear_scheme_override", "moe_scheme_override", mode="before"
55
- )
93
+ @field_validator("linear", "moe", mode="before")
94
@classmethod
57
- def _coerce_scheme(
58
- cls, v: str | OnlineQuantScheme | None
59
- ) -> OnlineQuantScheme | None:
60
- if isinstance(v, str):
61
- return OnlineQuantScheme(v)
62
- return v
63
-
64
-
65
-def resolve_online_quant_config(
95
+ def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any:
96
+ if not isinstance(v, str):
97
+ return v
98
+ field_name = info.field_name
99
+ assert field_name is not None
100
+ if v in _ONLINE_SHORTHANDS:
101
+ spec = getattr(_ONLINE_SHORTHANDS[v], field_name)
102
+ if spec is None:
103
+ raise ValueError(
104
+ f"online shorthand {v!r} does not define a {field_name} spec"
105
+ )
106
+ return spec
107
+ return QuantSpec(weight=_coerce_quant_key(v))
108
+
109
+
110
+# CLI shorthands accepted by `--quantization`. Each desugars to a full
111
+# QuantizationConfigArgs; activation overrides go through quantization_config.
112
+_ONLINE_SHORTHANDS: dict[str, QuantizationConfigArgs] = {
113
+ "fp8_per_tensor": QuantizationConfigArgs(
114
+ linear=QuantSpec(weight=kFp8StaticTensorSym),
115
+ moe=QuantSpec(weight=kFp8StaticTensorSym),
116
+ ),
117
+ "fp8_per_block": QuantizationConfigArgs(
118
+ linear=QuantSpec(weight=kFp8Static128BlockSym),
119
+ moe=QuantSpec(weight=kFp8Static128BlockSym),
120
+ ),
121
+ "mxfp8": QuantizationConfigArgs(
122
+ linear=QuantSpec(weight=kMxfp8Dynamic),
123
+ moe=QuantSpec(weight=kMxfp8Dynamic),
124
+ ),
125
+ # INT8 weight-only on MoE; linear stays unquantized (no `linear` field).
126
+ "int8_per_channel_weight_only": QuantizationConfigArgs(
127
+ moe=QuantSpec(weight=kInt8StaticChannelSym),
128
+ ),
129
+}
130
+
131
+
132
+# Names accepted by `--quantization`; "online" means "use quantization_config".
133
+ONLINE_QUANT_SHORTHAND_NAMES: tuple[str, ...] = (
134
+ *_ONLINE_SHORTHANDS.keys(),
135
+ "online",
136
+)
137
+
138
+
139
+def resolve_quantization_config(
140
quantization: str | None,
67
- quantization_config: dict[str, Any] | OnlineQuantizationConfigArgs | None,
68
-) -> OnlineQuantizationConfigArgs | None:
69
- """Resolve online quant scheme shorthand into a quantization config.
70
-
71
- If ``quantization`` is an online quant scheme (e.g. ``'fp8_per_tensor'``),
72
- ensures ``quantization_config`` has a matching ``global_scheme`` and casts
73
- it to :class:`OnlineQuantizationConfigArgs` if needed.
141
+ quantization_config: dict[str, Any] | QuantizationConfigArgs | None,
142
+) -> QuantizationConfigArgs | None:
143
+ """Resolve `--quantization` shorthand and `--quantization-config` into a
144
+ QuantizationConfigArgs.
145
+
146
+ `quantization` is a CLI shorthand that desugars into a base config via
147
+ `_ONLINE_SHORTHANDS`. `quantization_config` is a dict or pre-built args
148
+ object. When both are given, fields explicitly set in `quantization_config`
149
+ take precedence over the shorthand.
150
"""
75
- online_quant_values = {s.value for s in OnlineQuantScheme}
76
- valid_quantization_values = online_quant_values | {"online"}
77
- if quantization not in valid_quantization_values:
151
+ if quantization is not None and quantization not in ONLINE_QUANT_SHORTHAND_NAMES:
152
if quantization_config is not None:
153
raise ValueError(
80
- f"quantization_config is only supported when quantization "
81
- f"is one of {sorted(valid_quantization_values)}, "
154
+ f"quantization_config is only supported when quantization is "
155
+ f"one of {sorted(ONLINE_QUANT_SHORTHAND_NAMES)}, "
156
f"got quantization={quantization!r}"
157
)
158
return None
159
86
- if quantization in online_quant_values:
87
- scheme = OnlineQuantScheme(quantization)
88
-
89
- if quantization_config is None:
90
- quantization_config = {
91
- "global_scheme": scheme.value,
92
- }
93
- elif isinstance(quantization_config, OnlineQuantizationConfigArgs):
94
- if quantization_config.global_scheme is None:
95
- quantization_config.global_scheme = scheme
96
- elif quantization_config.global_scheme != scheme:
97
- raise ValueError(
98
- f"quantization={quantization!r} conflicts with "
99
- f"quantization_config.global_scheme="
100
- f"{quantization_config.global_scheme.value!r}. "
101
- f"These must match when both are specified."
102
- )
103
- elif isinstance(quantization_config, dict):
104
- existing = quantization_config.get("global_scheme")
105
- if existing is None:
106
- quantization_config["global_scheme"] = scheme.value
107
- else:
108
- # Coerce to enum for comparison
109
- existing_scheme = (
110
- OnlineQuantScheme(existing)
111
- if isinstance(existing, str)
112
- else existing
113
- )
114
- if existing_scheme != scheme:
115
- raise ValueError(
116
- f"quantization={quantization!r} conflicts "
117
- f"with quantization_config"
118
- f"['global_scheme']={existing!r}. "
119
- f"These must match when both are specified."
120
- )
121
-
122
- # Cast dict to OnlineQuantizationConfigArgs
160
+ base = _ONLINE_SHORTHANDS.get(quantization) if quantization else None
161
+
162
+ if quantization_config is None:
163
+ return base
164
+
165
if isinstance(quantization_config, dict):
124
- quantization_config = OnlineQuantizationConfigArgs(**quantization_config)
166
+ quantization_config = QuantizationConfigArgs(**quantization_config)
167
126
- return quantization_config
168
+ if base is None:
169
+ return quantization_config
170
+
171
+ return QuantizationConfigArgs(
172
+ linear=quantization_config.linear or base.linear,
173
+ moe=quantization_config.moe or base.moe,
174
+ ignore=quantization_config.ignore or base.ignore,
175
+ )
vllm/engine/arg_utils.py
+12
-4
@@ -118,7 +118,7 @@ from vllm.v1.sample.logits_processor import LogitsProcessor
118
from vllm.version import __version__ as VLLM_VERSION
119
120
if TYPE_CHECKING:
121
- from vllm.config.quantization import OnlineQuantizationConfigArgs
121
+ from vllm.config.quantization import QuantizationConfigArgs
122
from vllm.model_executor.layers.quantization import QuantizationMethods
123
from vllm.model_executor.model_loader import LoadFormats
124
from vllm.usage.usage_lib import UsageContext
@@ -527,7 +527,12 @@ class EngineArgs:
527
hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides")
528
tokenizer_revision: str | None = ModelConfig.tokenizer_revision
529
quantization: QuantizationMethods | str | None = ModelConfig.quantization
530
- quantization_config: "dict[str, Any] | OnlineQuantizationConfigArgs | None" = None
530
+ quantization_config: "dict[str, Any] | QuantizationConfigArgs | None" = None
531
+ """User-facing quantization configuration. Carries per-layer-kind
532
+ QuantSpecs (linear, moe) and ignore patterns; see
533
+ :class:`QuantizationConfigArgs`. Auto-populated from the matching online
534
+ shorthand when `quantization` is one of the values in
535
+ `ONLINE_QUANT_SHORTHAND_NAMES`."""
536
allow_deprecated_quantization: bool = ModelConfig.allow_deprecated_quantization
537
enforce_eager: bool = ModelConfig.enforce_eager
538
disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce
@@ -719,9 +724,9 @@ class EngineArgs:
724
if isinstance(self.ir_op_priority, dict):
725
self.ir_op_priority = IrOpPriorityConfig(**self.ir_op_priority)
726
722
- from vllm.config.quantization import resolve_online_quant_config
727
+ from vllm.config.quantization import resolve_quantization_config
728
724
- self.quantization_config = resolve_online_quant_config(
729
+ self.quantization_config = resolve_quantization_config(
730
self.quantization, self.quantization_config
731
)
732
@@ -785,6 +790,9 @@ class EngineArgs:
790
)
791
model_group.add_argument("--max-model-len", **model_kwargs["max_model_len"])
792
model_group.add_argument("--quantization", "-q", **model_kwargs["quantization"])
793
+ model_group.add_argument(
794
+ "--quantization-config", **model_kwargs["quantization_config"]
795
+ )
796
model_group.add_argument(
797
"--allow-deprecated-quantization",
798
**model_kwargs["allow_deprecated_quantization"],
vllm/entrypoints/llm.py
+2
-4
@@ -35,7 +35,7 @@ from vllm.config.model import (
35
TokenizerMode,
36
)
37
from vllm.config.quantization import (
38
- OnlineQuantizationConfigArgs,
38
+ QuantizationConfigArgs,
39
)
40
from vllm.distributed.weight_transfer.base import (
41
WeightTransferInitRequest,
@@ -248,9 +248,7 @@ class LLM:
248
attention_config: dict[str, Any] | AttentionConfig | None = None,
249
kv_cache_memory_bytes: int | None = None,
250
compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
251
- quantization_config: dict[str, Any]
252
- | OnlineQuantizationConfigArgs
253
- | None = None,
251
+ quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None,
252
logits_processors: list[str | type[LogitsProcessor]] | None = None,
253
**kwargs: Any,
254
) -> None:
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+173
-122
@@ -1,13 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from enum import Enum
4
-from typing import TYPE_CHECKING, Union
4
+from typing import TYPE_CHECKING, Literal, Union
5
6
import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm import envs
10
+from vllm.config import get_current_vllm_config
11
from vllm.config.kernel import MoEBackend
12
+from vllm.config.quantization import QuantizationConfigArgs
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.fused_moe import (
15
FusedMoEConfig,
@@ -207,25 +209,38 @@ def backend_to_kernel_cls(
209
raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}")
210
211
210
-def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend:
211
- """Map user's moe_backend string to Mxfp4MoeBackend."""
212
- mapping: dict[str, Mxfp4MoeBackend] = {
213
- "deep_gemm": Mxfp4MoeBackend.DEEPGEMM_MXFP4,
214
- "flashinfer_trtllm": Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
215
- "flashinfer_trtllm_afp8": Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
216
- "flashinfer_cutlass": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
217
- "flashinfer_cutlass_afp8": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
218
- "triton": Mxfp4MoeBackend.TRITON,
219
- "triton_unfused": Mxfp4MoeBackend.TRITON_UNFUSED,
220
- "humming": Mxfp4MoeBackend.HUMMING,
221
- "marlin": Mxfp4MoeBackend.MARLIN,
222
- "aiter": Mxfp4MoeBackend.AITER_MXFP4_BF16,
223
- "aiter_mxfp4_fp8": Mxfp4MoeBackend.AITER_MXFP4_FP8,
224
- "xpu": Mxfp4MoeBackend.XPU,
225
- "emulation": Mxfp4MoeBackend.EMULATION,
212
+def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:
213
+ """Map a moe_backend string to its candidate Mxfp4MoeBackends.
214
+
215
+ Vendor families return all activation variants; the caller picks one
216
+ via ``activation_key`` and ``is_supported_config``.
217
+ """
218
+ mapping: dict[str, list[Mxfp4MoeBackend]] = {
219
+ "deep_gemm": [Mxfp4MoeBackend.DEEPGEMM_MXFP4],
220
+ "flashinfer_trtllm": [
221
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
222
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
223
+ ],
224
+ "flashinfer_trtllm_afp8": [Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8],
225
+ "flashinfer_cutlass": [
226
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
227
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
228
+ ],
229
+ "flashinfer_cutlass_afp8": [Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8],
230
+ "triton": [Mxfp4MoeBackend.TRITON],
231
+ "triton_unfused": [Mxfp4MoeBackend.TRITON_UNFUSED],
232
+ "humming": [Mxfp4MoeBackend.HUMMING],
233
+ "marlin": [Mxfp4MoeBackend.MARLIN],
234
+ "aiter": [
235
+ Mxfp4MoeBackend.AITER_MXFP4_BF16,
236
+ Mxfp4MoeBackend.AITER_MXFP4_FP8,
237
+ ],
238
+ "aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8],
239
+ "xpu": [Mxfp4MoeBackend.XPU],
240
+ "emulation": [Mxfp4MoeBackend.EMULATION],
241
}
227
- if backend := mapping.get(runner_backend):
228
- return backend
242
+ if backends := mapping.get(runner_backend):
243
+ return backends
244
raise ValueError(
245
f"moe_backend='{runner_backend}' is not supported for MXFP4 MoE. "
246
f"Expected one of {list(mapping.keys())}."
@@ -233,16 +248,16 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend:
248
249
250
def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:
236
- """
237
- Get available backends in priority order based on platform and config.
238
- Only includes BF16 backends. MXFP8 backends are selected via env vars.
239
- """
251
+ """Available backends in priority order, BF16-act variant before
252
+ activation-quantized variant within each vendor family."""
253
_AVAILABLE_BACKENDS = [
254
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
255
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
256
Mxfp4MoeBackend.AITER_MXFP4_BF16,
257
Mxfp4MoeBackend.AITER_MXFP4_FP8,
258
Mxfp4MoeBackend.TRITON,
259
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
260
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
261
# TRITON_UNFUSED has bug with MTP support
262
# TODO re-enable after kernel is fixed
263
# TRITON_UNFUSED
@@ -288,6 +303,79 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:
303
return None # BF16 activation
304
305
306
+def _user_moe_activation_override() -> QuantKey | None:
307
+ """User's MoE activation override from quantization_config, or None."""
308
+ args = get_current_vllm_config().model_config.quantization_config
309
+ if not isinstance(args, QuantizationConfigArgs) or args.moe is None:
310
+ return None
311
+ return args.moe.activation
312
+
313
+
314
+def _resolve_activation_key(
315
+ model_activation_key: QuantKey | None,
316
+) -> QuantKey | None:
317
+ """Combine the model-supplied activation key with the user override.
318
+ Raises on conflict (both set and disagreeing)."""
319
+ user_override = _user_moe_activation_override()
320
+ if user_override is None:
321
+ return model_activation_key
322
+ if model_activation_key is None or model_activation_key == user_override:
323
+ return user_override
324
+ raise ValueError(
325
+ f"checkpoint declares MoE activation={model_activation_key} but "
326
+ f"quantization_config.moe.activation={user_override}; remove the "
327
+ f"override or align it with the checkpoint."
328
+ )
329
+
330
+
331
+def _make_log_backend(backend: Mxfp4MoeBackend) -> str:
332
+ return f"Using '{backend.value}' Mxfp4 MoE backend."
333
+
334
+
335
+def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str:
336
+ base = (
337
+ f"Mxfp4 MoE backend '{backend.value}' does not support the "
338
+ f"deployment configuration"
339
+ )
340
+ return f"{base} since {reason}." if reason else f"{base}."
341
+
342
+
343
+def _return_or_raise(
344
+ backend: Mxfp4MoeBackend,
345
+ config: FusedMoEConfig,
346
+ weight_key: QuantKey | None,
347
+ activation_key: QuantKey | None,
348
+ activation_format: mk.FusedMoEActivationFormat,
349
+ scope: Literal["process", "global", "local"] = "local",
350
+) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts]]:
351
+ reason: str | None = None
352
+ for k_cls in backend_to_kernel_cls(backend):
353
+ supported, reason = k_cls.is_supported_config(
354
+ k_cls, config, weight_key, activation_key, activation_format
355
+ )
356
+ if supported:
357
+ logger.info_once(_make_log_backend(backend), scope=scope)
358
+ return backend, k_cls
359
+ raise ValueError(_make_log_unsupported(backend, reason))
360
+
361
+
362
+def _filter_by_activation(
363
+ backends: list[Mxfp4MoeBackend],
364
+ requested_activation_key: QuantKey | None,
365
+) -> list[Mxfp4MoeBackend]:
366
+ """Pick variants matching ``requested_activation_key``; without one,
367
+ prefer BF16 if the list has any, else keep the list as-is so explicit
368
+ non-BF16 picks (e.g. the ``_afp8`` aliases) still land."""
369
+ if requested_activation_key is not None:
370
+ return [
371
+ b
372
+ for b in backends
373
+ if _backend_activation_key(b) == requested_activation_key
374
+ ]
375
+ bf16 = [b for b in backends if _backend_activation_key(b) is None]
376
+ return bf16 if bf16 else backends
377
+
378
+
379
def select_mxfp4_moe_backend(
380
config: FusedMoEConfig,
381
activation_key: QuantKey | None = None,
@@ -303,8 +391,7 @@ def select_mxfp4_moe_backend(
391
392
Note: Shape-specific fallbacks may still occur at runtime.
393
"""
306
- # If activation_key is explicitly provided (e.g., W4A8), use it
307
- requested_activation_key = activation_key
394
+ requested_activation_key = _resolve_activation_key(activation_key)
395
396
activation_format = (
397
mk.FusedMoEActivationFormat.BatchedExperts
@@ -312,67 +399,55 @@ def select_mxfp4_moe_backend(
399
else mk.FusedMoEActivationFormat.Standard
400
)
401
315
- def _make_log_backend(backend: Mxfp4MoeBackend):
316
- return f"Using '{backend.value}' Mxfp4 MoE backend."
317
-
318
- def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str:
319
- if reason:
320
- return (
321
- f"Mxfp4 MoE backend '{backend.value}' does not support the "
322
- f"deployment configuration since {reason}."
323
- )
324
- return (
325
- f"Mxfp4 MoE backend '{backend.value}' does not support the "
326
- "deployment configuration."
327
- )
328
-
329
- def _return_or_raise(
330
- backend: Mxfp4MoeBackend,
331
- config: FusedMoEConfig,
332
- weight_key: QuantKey | None,
333
- activation_key: QuantKey | None,
334
- activation_format: mk.FusedMoEActivationFormat,
335
- ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts]]:
336
- reason: str | None = None
337
- for k_cls in backend_to_kernel_cls(backend):
338
- supported, reason = k_cls.is_supported_config(
339
- k_cls, config, weight_key, activation_key, activation_format
340
- )
341
- if supported:
342
- logger.info_once(_make_log_backend(backend))
343
- return backend, k_cls
344
- raise ValueError(_make_log_unsupported(backend, reason))
345
-
402
runner_backend = config.moe_backend
403
if runner_backend != "auto":
348
- requested_backend = map_mxfp4_backend(runner_backend)
349
- if (
350
- activation_format == mk.FusedMoEActivationFormat.BatchedExperts
351
- and requested_backend == Mxfp4MoeBackend.MARLIN
352
- ):
353
- requested_backend = Mxfp4MoeBackend.BATCHED_MARLIN
354
- # Use requested_activation_key if provided, otherwise use backend default
355
- act_key = (
356
- requested_activation_key
357
- if requested_activation_key is not None
358
- else _backend_activation_key(requested_backend)
359
- )
360
- return _return_or_raise(
361
- requested_backend,
362
- config,
363
- kMxfp4Static,
364
- act_key,
365
- activation_format,
366
- )
404
+ requested_backends = map_mxfp4_backend(runner_backend)
405
+ if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
406
+ requested_backends = [
407
+ Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b
408
+ for b in requested_backends
409
+ ]
410
+ candidates = _filter_by_activation(requested_backends, requested_activation_key)
411
+ if not candidates:
412
+ raise ValueError(
413
+ f"moe_backend={runner_backend!r} does not support "
414
+ f"activation={requested_activation_key}; supported variants: "
415
+ f"{[b.name for b in requested_backends]}"
416
+ )
417
+ last_error: Exception | None = None
418
+ for requested_backend in candidates:
419
+ act_key = (
420
+ requested_activation_key
421
+ if requested_activation_key is not None
422
+ else _backend_activation_key(requested_backend)
423
+ )
424
+ try:
425
+ return _return_or_raise(
426
+ requested_backend,
427
+ config,
428
+ kMxfp4Static,
429
+ act_key,
430
+ activation_format,
431
+ )
432
+ except ValueError as e:
433
+ last_error = e
434
+ assert last_error is not None
435
+ raise last_error
436
437
# Select kernels in order of backend.
369
- AVAILABLE_BACKENDS = _get_priority_backends_for_gpt_oss()
438
+ AVAILABLE_BACKENDS = _filter_by_activation(
439
+ _get_priority_backends_for_gpt_oss(), requested_activation_key
440
+ )
441
442
# Handle explicit FlashInfer MXFP4 BF16 configuration.
443
if envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16"):
444
if not envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16:
374
- AVAILABLE_BACKENDS.remove(Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16)
375
- AVAILABLE_BACKENDS.remove(Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16)
445
+ for _b in (
446
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
447
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
448
+ ):
449
+ if _b in AVAILABLE_BACKENDS:
450
+ AVAILABLE_BACKENDS.remove(_b)
451
else:
452
if current_platform.is_device_capability(90):
453
return _return_or_raise(
@@ -481,54 +556,30 @@ def select_deepseek_v4_mxfp4_moe_backend(
556
else mk.FusedMoEActivationFormat.Standard
557
)
558
484
- def _make_log_backend(backend: Mxfp4MoeBackend):
485
- return f"Using '{backend.value}' Mxfp4 MoE backend."
486
-
487
- def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str:
488
- if reason:
489
- return (
490
- f"Mxfp4 MoE backend '{backend.value}' does not support the "
491
- f"deployment configuration since {reason}."
492
- )
493
- return (
494
- f"Mxfp4 MoE backend '{backend.value}' does not support the "
495
- "deployment configuration."
496
- )
497
-
498
- def _return_or_raise(
499
- backend: Mxfp4MoeBackend,
500
- config: FusedMoEConfig,
501
- weight_key: QuantKey | None,
502
- activation_key: QuantKey | None,
503
- activation_format: mk.FusedMoEActivationFormat,
504
- ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts]]:
505
- reason: str | None = None
506
- for k_cls in backend_to_kernel_cls(backend):
507
- supported, reason = k_cls.is_supported_config(
508
- k_cls, config, weight_key, activation_key, activation_format
509
- )
510
- if supported:
511
- logger.info_once(_make_log_backend(backend), scope="local")
512
- return backend, k_cls
513
- raise ValueError(_make_log_unsupported(backend, reason))
514
-
559
# Honor explicit moe_backend (e.g. "marlin", "triton_unfused") before
560
# falling back to the auto priority list.
561
runner_backend = config.moe_backend
562
if runner_backend != "auto":
519
- requested_backend = map_mxfp4_backend(runner_backend)
520
- if (
521
- activation_format == mk.FusedMoEActivationFormat.BatchedExperts
522
- and requested_backend == Mxfp4MoeBackend.MARLIN
523
- ):
524
- requested_backend = Mxfp4MoeBackend.BATCHED_MARLIN
525
- return _return_or_raise(
526
- requested_backend,
527
- config,
528
- kMxfp4Static,
529
- _backend_activation_key(requested_backend),
530
- activation_format,
531
- )
563
+ requested_backends = map_mxfp4_backend(runner_backend)
564
+ if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
565
+ requested_backends = [
566
+ Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b
567
+ for b in requested_backends
568
+ ]
569
+ last_error: Exception | None = None
570
+ for requested_backend in requested_backends:
571
+ try:
572
+ return _return_or_raise(
573
+ requested_backend,
574
+ config,
575
+ kMxfp4Static,
576
+ _backend_activation_key(requested_backend),
577
+ activation_format,
578
+ )
579
+ except ValueError as e:
580
+ last_error = e
581
+ assert last_error is not None
582
+ raise last_error
583
584
# DeepSeek-V4 on ROCm is more accurate with the unfused Triton MXFP4 path
585
# than the default AITER path. Prefer Triton-unfused for this routing mode,
vllm/model_executor/layers/quantization/__init__.py
+11
-13
@@ -35,10 +35,9 @@ QuantizationMethods = Literal[
35
"deepseek_v4_fp8",
36
"cpu_awq",
37
"online",
38
- # Below are values of the OnlineQuantScheme enum, specified as strings to
39
- # avoid circular import issues. This is here to provide a shortcut where
40
- # the user can specify "LLM(..., quantization='fp8_per_tensor')" as
41
- # shorthand for creating a more complicated online quant config object
38
+ # Below are online quant shorthand names (see vllm.config.quantization).
39
+ # Listed here as strings to avoid a circular import; kept in sync with
40
+ # _ONLINE_SHORTHANDS by the assertion in get_quantization_config().
41
"fp8_per_tensor",
42
"fp8_per_block",
43
"int8_per_channel_weight_only",
@@ -111,7 +110,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
110
raise ValueError(f"Invalid quantization method: {quantization}")
111
112
# lazy import to avoid triggering `torch.compile` too early
114
- from vllm.config.quantization import OnlineQuantScheme
113
+ from vllm.config.quantization import _ONLINE_SHORTHANDS
114
from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
115
from vllm.model_executor.models.deepseek_v4 import DeepseekV4FP8Config
116
@@ -171,16 +170,15 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
170
"online": OnlineQuantizationConfig,
171
}
172
174
- # Below are values of the OnlineQuantScheme enum. This is here to provide
175
- # a shortcut where the user can specify
176
- # "LLM(..., quantization='fp8_per_tensor')" as shorthand for creating a
177
- # more complicated online quant config object
178
- for scheme in OnlineQuantScheme:
179
- assert scheme.value not in method_to_config, (
180
- f"Online quant scheme {scheme.value!r} conflicts with an "
173
+ # Register online shorthands as quantization methods so the user can
174
+ # specify "LLM(..., quantization='fp8_per_tensor')" as shorthand for
175
+ # creating a more complicated online quant config object.
176
+ for shorthand in _ONLINE_SHORTHANDS:
177
+ assert shorthand not in method_to_config, (
178
+ f"Online quant shorthand {shorthand!r} conflicts with an "
179
f"existing quantization method"
180
)
183
- method_to_config[scheme.value] = OnlineQuantizationConfig
181
+ method_to_config[shorthand] = OnlineQuantizationConfig
182
183
# Update the `method_to_config` with customized quantization methods.
184
method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
vllm/model_executor/layers/quantization/online/base.py
+65
-38
@@ -5,10 +5,7 @@ from typing import Any
5
6
import torch
7
8
-from vllm.config.quantization import (
9
- OnlineQuantizationConfigArgs,
10
- OnlineQuantScheme,
11
-)
8
+from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.fused_moe import (
11
RoutedExperts,
@@ -41,31 +38,50 @@ from vllm.model_executor.layers.quantization.online.mxfp8 import (
38
Mxfp8OnlineLinearMethod,
39
Mxfp8OnlineMoEMethod,
40
)
41
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
42
+ QuantKey,
43
+ kFp8Static128BlockSym,
44
+ kFp8StaticTensorSym,
45
+ kInt8StaticChannelSym,
46
+ kMxfp8Dynamic,
47
+)
48
49
logger = init_logger(__name__)
50
51
52
+# Online dispatch tables, keyed by the QuantSpec.weight QuantKey. The
53
+# corresponding method class handles the activation choice via its
54
+# `supported_activation_quant` set.
55
+_ONLINE_LINEAR_METHODS: dict[QuantKey, type] = {
56
+ kFp8StaticTensorSym: Fp8PerTensorOnlineLinearMethod,
57
+ kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod,
58
+ kMxfp8Dynamic: Mxfp8OnlineLinearMethod,
59
+}
60
+
61
+_ONLINE_MOE_METHODS: dict[QuantKey, type] = {
62
+ kFp8StaticTensorSym: Fp8PerTensorOnlineMoEMethod,
63
+ kFp8Static128BlockSym: Fp8PerBlockOnlineMoEMethod,
64
+ kMxfp8Dynamic: Mxfp8OnlineMoEMethod,
65
+ kInt8StaticChannelSym: Int8OnlineMoEMethod,
66
+}
67
+
68
+
69
class OnlineQuantizationConfig(QuantizationConfig):
49
- """Model-level config class for online quantization (quantize fp16/bf16 weights
70
+ """Model-level config for online quantization (quantize fp16/bf16 weights
71
during model loading, without requiring a pre-quantized checkpoint)."""
72
73
def __init__(
74
self,
54
- args: OnlineQuantizationConfigArgs,
75
+ args: QuantizationConfigArgs,
76
) -> None:
77
super().__init__()
57
- if (
58
- args.global_scheme is None
59
- and args.linear_scheme_override is None
60
- and args.moe_scheme_override is None
61
- ):
78
+ if args.linear is None and args.moe is None:
79
raise ValueError(
80
"OnlineQuantizationConfig requires at least one of "
64
- "global_scheme, linear_scheme_override, or "
65
- "moe_scheme_override to be set."
81
+ "quantization_config.linear or quantization_config.moe "
82
+ "to be set."
83
)
84
self.args = args
68
- self.quant_scheme = args.global_scheme
85
self.ignored_layers: list[str] = args.ignore
86
87
@classmethod
@@ -94,6 +110,33 @@ class OnlineQuantizationConfig(QuantizationConfig):
110
"quantization='fp8_per_tensor'/'fp8_per_block' instead."
111
)
112
113
+ def _dispatch(
114
+ self,
115
+ spec: QuantSpec | None,
116
+ table: dict[QuantKey, type],
117
+ layer: torch.nn.Module,
118
+ ) -> "QuantizeMethodBase | None":
119
+ if spec is None or spec.weight is None:
120
+ return None
121
+ cls = table.get(spec.weight)
122
+ if cls is None:
123
+ raise ValueError(
124
+ f"online quantization for {type(layer).__name__} with "
125
+ f"weight={spec.weight} is not supported; supported weight "
126
+ f"keys: {sorted(str(k) for k in table)}"
127
+ )
128
+ # Online method classes pick their own activation format internally.
129
+ # Per-class activation overrides are not yet wired through; reject
130
+ # explicit overrides until the relevant method class opts in.
131
+ if spec.activation is not None:
132
+ raise ValueError(
133
+ f"activation override (activation={spec.activation}) is not "
134
+ f"yet supported for online {cls.__name__}"
135
+ )
136
+ if isinstance(layer, RoutedExperts):
137
+ return cls(layer=layer)
138
+ return cls()
139
+
140
def get_quant_method(
141
self, layer: torch.nn.Module, prefix: str
142
) -> "QuantizeMethodBase | None":
@@ -104,20 +147,8 @@ class OnlineQuantizationConfig(QuantizationConfig):
147
fused_mapping=self.packed_modules_mapping,
148
):
149
return UnquantizedLinearMethod()
107
-
108
- linear_scheme = self.args.linear_scheme_override or self.args.global_scheme
109
- if linear_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
110
- logger.warning_once(
111
- "INT8 online quantization only quantizes MoE expert "
112
- "weights. linear layers remain in full precision."
113
- )
114
- return UnquantizedLinearMethod()
115
- elif linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
116
- return Fp8PerBlockOnlineLinearMethod()
117
- elif linear_scheme == OnlineQuantScheme.MXFP8:
118
- return Mxfp8OnlineLinearMethod()
119
- else:
120
- return Fp8PerTensorOnlineLinearMethod()
150
+ method = self._dispatch(self.args.linear, _ONLINE_LINEAR_METHODS, layer)
151
+ return method if method is not None else UnquantizedLinearMethod()
152
elif isinstance(layer, RoutedExperts):
153
if should_ignore_layer(
154
prefix,
@@ -125,14 +156,10 @@ class OnlineQuantizationConfig(QuantizationConfig):
156
fused_mapping=self.packed_modules_mapping,
157
):
158
return UnquantizedFusedMoEMethod(layer.moe_config)
128
-
129
- moe_scheme = self.args.moe_scheme_override or self.args.global_scheme
130
- if moe_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
131
- return Int8OnlineMoEMethod(layer=layer)
132
- elif moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
133
- return Fp8PerBlockOnlineMoEMethod(layer=layer)
134
- elif moe_scheme == OnlineQuantScheme.MXFP8:
135
- return Mxfp8OnlineMoEMethod(layer=layer)
136
- else:
137
- return Fp8PerTensorOnlineMoEMethod(layer=layer)
159
+ method = self._dispatch(self.args.moe, _ONLINE_MOE_METHODS, layer)
160
+ return (
161
+ method
162
+ if method is not None
163
+ else UnquantizedFusedMoEMethod(layer.moe_config)
164
+ )
165
return None
vllm/model_executor/model_loader/weight_utils.py
+7
-23
@@ -300,12 +300,10 @@ def get_quant_config(
300
)
301
302
if hf_quant_config is not None:
303
- if model_config.quantization_config is not None:
304
- raise ValueError(
305
- "Setting `quantization_config` for online "
306
- "quantization when the model checkpoint already "
307
- "has a `quantization_config` is not supported"
308
- )
303
+ # `model_config.quantization_config` may be set alongside a checkpoint
304
+ # quant config: the checkpoint determines `quant_cls`, and the user's
305
+ # QuantizationConfigArgs is consulted by individual quant methods
306
+ # (e.g. for activation overrides via the MXFP4 oracle).
307
308
# For modelopt_mixed, config.json's quantization_config may or may
309
# not contain the per-layer quantized_layers map. Newer checkpoints
@@ -330,12 +328,6 @@ def get_quant_config(
328
quantization_config_file = hf_overrides.get("quantization_config_file", None)
329
if quantization_config_file is not None:
330
if hasattr(quant_cls, "from_config_file"):
333
- if model_config.quantization_config is not None:
334
- raise ValueError(
335
- "Setting `quantization_config` for online "
336
- "quantization when the model checkpoint already "
337
- "has a `quantization_config` is not supported"
338
- )
331
return quant_cls.from_config_file(quantization_config_file)
332
else:
333
raise NotImplementedError(
@@ -346,12 +338,6 @@ def get_quant_config(
338
quantization_config_json = hf_overrides.get("quantization_config_dict_json", None)
339
if quantization_config_json is not None:
340
if hasattr(quant_cls, "from_config_dict_json"):
349
- if model_config.quantization_config is not None:
350
- raise ValueError(
351
- "Setting `quantization_config` for online "
352
- "quantization when the model checkpoint already "
353
- "has a `quantization_config` is not supported"
354
- )
341
return quant_cls.from_config_dict_json(quantization_config_json)
342
else:
343
raise NotImplementedError(
@@ -360,17 +346,15 @@ def get_quant_config(
346
f"{quant_cls}"
347
)
348
363
- # Online quantization doesn't read from checkpoint configs — it quantizes
349
+ # Online quantization doesn't read from checkpoint configs - it quantizes
350
# fp16/bf16 weights on the fly during loading.
351
if model_config.quantization_config is not None:
366
- from vllm.config.quantization import OnlineQuantizationConfigArgs
352
+ from vllm.config.quantization import QuantizationConfigArgs
353
from vllm.model_executor.layers.quantization.online.base import (
354
OnlineQuantizationConfig,
355
)
356
371
- assert isinstance(
372
- model_config.quantization_config, OnlineQuantizationConfigArgs
373
- )
357
+ assert isinstance(model_config.quantization_config, QuantizationConfigArgs)
358
return OnlineQuantizationConfig(args=model_config.quantization_config)
359
360
# Inflight BNB quantization