[Quantization] Support online quantization with partially pre-quantized checkpoints (#51392)

Signed-off-by: Felix Marty <Felix.Marty@amd.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mgoin <mgoin64@gmail.com>

fxmarty-amd committed Sep 4, 2026 at 20:46 UTC 5690b02c03832a4ac3af231d3ecebe649c188095
25 files changed +847 -176
docs/features/quantization/online.md
+18
@@ -99,6 +99,24 @@ vllm serve openai/gpt-oss-20b --quantization-config.moe.activation mxfp8
99
100 Combine with `--moe-backend` to pin a specific kernel family.
101
102 +### Online quantization on unquantized layers from partially-quantized checkpoints
103 +
104 +Online quantization can be used on already quantized checkpoints independently of their original `quant_method` (`modelopt`, `compressed-tensors`, `quark`, etc.), for layers that are left unquantized in the original checkpoint.
105 +
106 +The checkpoint `quant_method` remains responsible for its quantized layers, while the selected unquantized layers use the requested online method.
107 +
108 +For example:
109 +
110 +```bash
111 +vllm serve amd/Qwen3.5-35B-A3B-MXFP4 \
112 + --quantization-config.linear mxfp8
113 +```
114 +
115 +adds MXFP8 quantization to the dense linear layers of a Quark checkpoint where only MOE experts are quantized.
116 +
117 +!!! info
118 + `quantization_config.ignore` is an online-only exclusion: the original `quant_method` relies solely on its own ignore implementation and on the ignored layers specified in `config.json`.
119 +
120 ### Separate Schemes for Dense and MoE Layers
121
122 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.
tests/model_executor/test_eagle_quantization.py
+1
@@ -83,6 +83,7 @@ def test_fc_layer_quant_config_usage(default_vllm_config, dist_init, device) ->
83 assert fc_no_quant.output_size == output_size
84
85 mock_quant_config = Mock()
86 + mock_quant_config.online_quantization_config = None
87 fc_with_quant = ReplicatedLinear(
88 input_size=input_size,
89 output_size=output_size,
tests/quantization/test_config_utils.py
+1 -1
@@ -58,7 +58,7 @@ def test_online_ignore_supports_fnmatch_patterns():
58 )
59
60 with pytest.raises(ValueError, match="matches both quantization_config.ignore"):
61 - config._dispatch_target(layer_name, Mock())
61 + config.resolve_quant_method_cls(Mock(), layer_name)
62
63
64 @pytest.mark.parametrize(
tests/quantization/test_online.py
+514 -54
@@ -2,8 +2,10 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests online quantization."""
4
5 +import json
6 +from pathlib import Path
7 from types import SimpleNamespace
6 -from typing import cast
8 +from typing import Any, cast
9 from unittest.mock import Mock
10
11 import pytest
@@ -18,18 +20,39 @@ from tests.quantization.utils import (
20 from vllm import _custom_ops as ops
21 from vllm._aiter_ops import rocm_aiter_ops
22 from vllm._custom_ops import scaled_fp4_quant
21 -from vllm.config import ModelConfig
23 +from vllm.config.load import LoadConfig
24 +from vllm.config.model import ModelConfig
25 from vllm.config.quantization import (
26 QuantizationConfigArgs,
27 resolve_quantization_config,
28 )
29 +from vllm.config.vllm import VllmConfig
30 from vllm.forward_context import set_forward_context
31 +from vllm.model_executor.kernels.linear.mxfp8.emulation import (
32 + EmulationMxfp8LinearKernel,
33 +)
34 +from vllm.model_executor.kernels.linear.mxfp8.marlin import (
35 + MarlinMxfp8LinearKernel,
36 +)
37 from vllm.model_executor.layers.attention import Attention
38 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
39 from vllm.model_executor.layers.linear import (
40 ColumnParallelLinear,
41 LinearBase,
42 UnquantizedLinearMethod,
43 )
44 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
45 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
46 + CompressedTensorsConfig,
47 + CompressedTensorsLinearMethod,
48 +)
49 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
50 + CompressedTensorsMoEMethod,
51 +)
52 +from vllm.model_executor.layers.quantization.modelopt import (
53 + ModelOptFp8Config,
54 + ModelOptMxFp8Config,
55 +)
56 from vllm.model_executor.layers.quantization.online.base import (
57 OnlineQuantizationConfig,
58 )
@@ -50,19 +73,31 @@ from vllm.model_executor.layers.quantization.online.mxfp4 import (
73 Mxfp4OnlineLinearMethod,
74 Mxfp4OnlineMoEMethod,
75 )
76 +from vllm.model_executor.layers.quantization.online.mxfp8 import (
77 + Mxfp8OnlineLinearMethod,
78 +)
79 from vllm.model_executor.layers.quantization.online.nvfp4 import (
80 Nvfp4OnlineMoEMethod,
81 _quantize_moe_weight_to_nvfp4,
82 )
83 +from vllm.model_executor.layers.quantization.quark.quark import (
84 + QuarkConfig,
85 + QuarkLinearMethod,
86 +)
87 from vllm.model_executor.layers.quantization.utils import quant_utils
88 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
89 + MXFP8_VALUE_DTYPE,
90 +)
91 from vllm.model_executor.layers.quantization.utils.quant_utils import (
92 amax_for_moe_weight_quant,
93 amax_for_tp_weight_quant,
94 + kMxfp8Dynamic,
95 weight_amax,
96 )
97 from vllm.model_executor.layers.vocab_parallel_embedding import (
98 VocabParallelEmbedding,
99 )
100 +from vllm.model_executor.model_loader import weight_utils
101 from vllm.model_executor.model_loader.base_loader import log_online_quantization
102 from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
103 from vllm.model_executor.models.granitemoe import (
@@ -83,6 +118,10 @@ else:
118
119
120 DEVICE = current_platform.device_type
121 +GRANITE_MODEL_NAME = "ibm-granite/granite-3.0-1b-a400m-base"
122 +PARTIALLY_PREQUANTIZED_MODEL_NAME = (
123 + "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable"
124 +)
125
126
127 def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
@@ -132,15 +171,349 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
171 assert process_weights.call_count == 2
172
173
174 +def _fully_quantized_quark_config() -> QuarkConfig:
175 + return QuarkConfig(
176 + {
177 + "exclude": [],
178 + "global_quant_config": {
179 + "weight": {
180 + "dtype": "int8",
181 + "qscheme": "per_tensor",
182 + "is_dynamic": False,
183 + "symmetric": True,
184 + },
185 + "input_tensors": {
186 + "dtype": "int8",
187 + "qscheme": "per_tensor",
188 + "is_dynamic": False,
189 + "symmetric": True,
190 + },
191 + },
192 + "layer_quant_config": {},
193 + "layer_type_quant_config": {},
194 + }
195 + )
196 +
197 +
198 +def _fully_quantized_modelopt_config() -> ModelOptFp8Config:
199 + return ModelOptFp8Config(
200 + quant_method="FP8",
201 + is_checkpoint_fp8_serialized=True,
202 + kv_cache_quant_method=None,
203 + exclude_modules=[],
204 + )
205 +
206 +
207 +def _moe_only_compressed_tensors_config() -> CompressedTensorsConfig:
208 + return CompressedTensorsConfig(
209 + target_scheme_map={"RoutedExperts": {}},
210 + ignore=[],
211 + quant_format="pack-quantized",
212 + )
213 +
214 +
215 +def _write_minimal_llama_config(
216 + model_path: Path, quantization_config: dict[str, Any] | None = None
217 +) -> None:
218 + config: dict[str, Any] = {
219 + "architectures": ["LlamaForCausalLM"],
220 + "model_type": "llama",
221 + "hidden_size": 32,
222 + "intermediate_size": 64,
223 + "num_attention_heads": 4,
224 + "num_hidden_layers": 1,
225 + "vocab_size": 32,
226 + "max_position_embeddings": 32,
227 + }
228 + if quantization_config is not None:
229 + config["quantization_config"] = quantization_config
230 + (model_path / "config.json").write_text(json.dumps(config))
231 +
232 +
233 +@pytest.mark.parametrize(
234 + "checkpoint_config_factory,raises_conflict",
235 + [
236 + pytest.param(_fully_quantized_quark_config, True, id="quark"),
237 + pytest.param(_fully_quantized_modelopt_config, True, id="modelopt"),
238 + pytest.param(
239 + _moe_only_compressed_tensors_config,
240 + False,
241 + id="compressed_tensors",
242 + ),
243 + ],
244 +)
245 +def test_online_prequantized_compatibility(
246 + checkpoint_config_factory,
247 + raises_conflict: bool,
248 + default_vllm_config,
249 + dist_init,
250 +) -> None:
251 + """Online weights replace only layers left unquantized by a checkpoint."""
252 + default_vllm_config.model_config = ModelConfig()
253 + checkpoint_config = checkpoint_config_factory()
254 +
255 + checkpoint_config.online_quantization_config = OnlineQuantizationConfig(
256 + QuantizationConfigArgs(linear="mxfp8")
257 + )
258 + config = checkpoint_config
259 +
260 + layer_kwargs = {
261 + "input_size": 32,
262 + "output_size": 32,
263 + "bias": False,
264 + "params_dtype": torch.bfloat16,
265 + "quant_config": config,
266 + "prefix": "model.layers.0.self_attn.o_proj",
267 + "disable_tp": True,
268 + }
269 +
270 + if raises_conflict:
271 + with pytest.raises(ValueError, match="pre-quantized layer"):
272 + ColumnParallelLinear(**layer_kwargs)
273 + else:
274 + layer = ColumnParallelLinear(**layer_kwargs)
275 + assert isinstance(layer.quant_method, Mxfp8OnlineLinearMethod)
276 +
277 +
278 +def test_online_target_rejects_prequantized_layer(
279 + default_vllm_config, dist_init
280 +) -> None:
281 + """A targets match participates in checkpoint compatibility checks."""
282 + default_vllm_config.model_config = ModelConfig()
283 + prefix = "model.layers.0.self_attn.o_proj"
284 + quant_config = _fully_quantized_quark_config()
285 + quant_config.online_quantization_config = OnlineQuantizationConfig(
286 + QuantizationConfigArgs(targets={prefix: "mxfp8"})
287 + )
288 +
289 + with pytest.raises(ValueError, match="pre-quantized layer"):
290 + ColumnParallelLinear(
291 + input_size=32,
292 + output_size=32,
293 + bias=False,
294 + params_dtype=torch.bfloat16,
295 + quant_config=quant_config,
296 + prefix=prefix,
297 + disable_tp=True,
298 + )
299 +
300 +
301 +def test_online_ignore_keeps_checkpoint_quantization_linear(
302 + default_vllm_config, dist_init, monkeypatch
303 +):
304 + """Ignoring online quantization does not replace a checkpoint method (linear)."""
305 + default_vllm_config.model_config = ModelConfig()
306 + quant_config = _fully_quantized_quark_config()
307 + prefix = "model.layers.0.self_attn.o_proj"
308 + quant_config.online_quantization_config = OnlineQuantizationConfig(
309 + QuantizationConfigArgs(linear="mxfp8", ignore=[prefix])
310 + )
311 + monkeypatch.setattr(
312 + quant_config.online_quantization_config,
313 + "get_quant_method",
314 + lambda *args: pytest.fail("online method should not be constructed"),
315 + )
316 +
317 + layer = ColumnParallelLinear(
318 + input_size=32,
319 + output_size=32,
320 + bias=False,
321 + params_dtype=torch.bfloat16,
322 + quant_config=quant_config,
323 + prefix=prefix,
324 + disable_tp=True,
325 + )
326 +
327 + assert isinstance(layer.quant_method, QuarkLinearMethod)
328 +
329 +
330 +def test_online_quantization_rejects_prequantized_moe(
331 + default_vllm_config, dist_init
332 +) -> None:
333 + """Online linear and MoE quantization reject a pre-quantized MoE layer."""
334 + default_vllm_config.model_config = ModelConfig()
335 + prefix = "model.layers.0.mlp.experts"
336 + quant_config = _fully_quantized_quark_config()
337 + quant_config.online_quantization_config = OnlineQuantizationConfig(
338 + QuantizationConfigArgs(linear="mxfp4", moe="mxfp4")
339 + )
340 +
341 + with pytest.raises(ValueError, match="pre-quantized layer"):
342 + FusedMoEFactory(
343 + num_experts=4,
344 + top_k=2,
345 + hidden_size=32,
346 + intermediate_size=64,
347 + params_dtype=torch.bfloat16,
348 + quant_config=quant_config,
349 + prefix=prefix,
350 + )
351 +
352 +
353 +def test_activation_only_override_applies_to_checkpoint_method(
354 + default_vllm_config, tmp_path
355 +) -> None:
356 + """Activation-only overrides remain available to the checkpoint method."""
357 + _write_minimal_llama_config(tmp_path)
358 + model_config = ModelConfig(
359 + model=str(tmp_path),
360 + quantization="compressed-tensors",
361 + quantization_config=QuantizationConfigArgs(moe={"activation": "mxfp8"}),
362 + hf_overrides={
363 + "quantization_config": {
364 + "quant_method": "compressed-tensors",
365 + "format": "pack-quantized",
366 + "config_groups": {},
367 + }
368 + },
369 + )
370 +
371 + result = weight_utils.get_quant_config(model_config, LoadConfig())
372 +
373 + assert isinstance(result, CompressedTensorsConfig)
374 + assert isinstance(result.online_quantization_config, OnlineQuantizationConfig)
375 +
376 + default_vllm_config.model_config = model_config
377 + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
378 + _resolve_activation_key,
379 + )
380 +
381 + assert _resolve_activation_key(None) == kMxfp8Dynamic
382 +
383 +
384 +def test_online_overlay_requires_normalized_quantization_config(tmp_path) -> None:
385 + """The loader rejects a raw config dict instead of silently ignoring it."""
386 + _write_minimal_llama_config(tmp_path)
387 + model_config = ModelConfig(
388 + model=str(tmp_path),
389 + quantization="compressed-tensors",
390 + quantization_config={"linear": "mxfp8"},
391 + hf_overrides={
392 + "quantization_config": {
393 + "quant_method": "compressed-tensors",
394 + "format": "pack-quantized",
395 + "config_groups": {},
396 + }
397 + },
398 + )
399 +
400 + # assert isinstance(online_args, QuantizationConfigArgs)
401 + # in maybe_compose_online_quantization
402 + with pytest.raises(AssertionError):
403 + weight_utils.get_quant_config(model_config, LoadConfig())
404 +
405 +
406 +def test_online_overlay_loads_checkpoint_config_file(tmp_path) -> None:
407 + """An online overlay must not bypass checkpoint config file loading."""
408 + _write_minimal_llama_config(tmp_path)
409 + (tmp_path / "quantize_config.json").write_text(
410 + '{"w_bit": 4, "q_group_size": 128, "zero_point": true}'
411 + )
412 + model_config = ModelConfig(
413 + model=str(tmp_path),
414 + quantization="awq",
415 + quantization_config=QuantizationConfigArgs(linear="mxfp8"),
416 + )
417 +
418 + result = weight_utils.get_quant_config(model_config, LoadConfig())
419 +
420 + assert isinstance(result, AutoAWQConfig)
421 + assert result.weight_bits == 4
422 + assert result.group_size == 128
423 + assert result.zero_point
424 + assert isinstance(result.online_quantization_config, OnlineQuantizationConfig)
425 + assert result.online_quantization_config.args is model_config.quantization_config
426 +
427 +
428 +@pytest.mark.parametrize(
429 + "quantization,checkpoint_quantization_config,is_config_file",
430 + [
431 + pytest.param("fp8_per_block", None, False, id="fp8-per-block-online"),
432 + pytest.param("mxfp4", None, False, id="mxfp4-online"),
433 + pytest.param("mxfp8", None, False, id="mxfp8-online"),
434 + pytest.param(
435 + "mxfp8",
436 + {"quant_method": "mxfp8", "ignored_layers": []},
437 + False,
438 + id="mxfp8-inline-checkpoint-config",
439 + ),
440 + pytest.param(
441 + "mxfp8",
442 + {"quant_method": "mxfp8", "ignored_layers": []},
443 + True,
444 + id="mxfp8-checkpoint-config-file",
445 + ),
446 + ],
447 +)
448 +def test_online_shorthand_selects_checkpoint_or_online_config(
449 + tmp_path, quantization, checkpoint_quantization_config, is_config_file
450 +) -> None:
451 + """Online shorthands defer to checkpoint metadata when available."""
452 +
453 + if is_config_file:
454 + (tmp_path / "hf_quant_config.json").write_text(
455 + '{"quant_method": "mxfp8", "ignored_layers": []}'
456 + )
457 +
458 + _write_minimal_llama_config(
459 + tmp_path,
460 + None if is_config_file else checkpoint_quantization_config,
461 + )
462 + model_config = ModelConfig(model=str(tmp_path), quantization=quantization)
463 +
464 + result = weight_utils.get_quant_config(model_config, LoadConfig())
465 +
466 + if checkpoint_quantization_config is not None:
467 + assert isinstance(result, ModelOptMxFp8Config)
468 + assert result.online_quantization_config is None
469 + else:
470 + assert isinstance(result, OnlineQuantizationConfig)
471 + assert result.online_quantization_config is None
472 +
473 +
474 +def test_log_online_quantization_for_composable_config(monkeypatch) -> None:
475 + """Composable configs log their nested online quantization results."""
476 + online_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8"))
477 + online_config.quantized_layers = {
478 + "model.layers.0.self_attn.o_proj": ("linear", "mxfp8", None),
479 + "model.layers.1.self_attn.o_proj": ("linear", "mxfp8", None),
480 + }
481 + vllm_config = VllmConfig(quant_config=online_config)
482 + log_args = []
483 + monkeypatch.setattr(
484 + "vllm.model_executor.model_loader.base_loader.logger.info",
485 + lambda *args: log_args.append(args),
486 + )
487 +
488 + log_online_quantization(vllm_config)
489 +
490 + assert log_args == [
491 + (
492 + "Quantized %d layers of types: %s",
493 + 2,
494 + "; ".join(online_config.quantized_layer_summaries),
495 + )
496 + ]
497 +
498 +
499 +def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None:
500 + """A checkpoint quant method cannot be replaced by an online shorthand."""
501 + _write_minimal_llama_config(tmp_path, {"quant_method": "mxfp4"})
502 +
503 + with pytest.raises(ValueError, match="does not match the quantization"):
504 + ModelConfig(model=str(tmp_path), quantization="fp8_per_channel")
505 +
506 +
507 @pytest.mark.skipif(
508 not is_quant_method_supported("fp8"),
509 reason="FP8 is not supported on this GPU type.",
510 )
511 @pytest.mark.parametrize(
140 - "quant_scheme,online_quant_args,expected_linear_cls,expected_moe_cls",
512 + ("model_name,quant_scheme,online_quant_args,expected_linear_cls,expected_moe_cls"),
513 [
514 # simple case - quantization='fp8_per_tensor'
515 (
516 + GRANITE_MODEL_NAME,
517 "fp8_per_tensor",
518 None,
519 Fp8PerTensorOnlineLinearMethod,
@@ -148,12 +521,14 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
521 ),
522 # simple case - quantization='fp8_per_block'
523 (
524 + GRANITE_MODEL_NAME,
525 "fp8_per_block",
526 None,
527 Fp8PerBlockOnlineLinearMethod,
528 Fp8PerBlockOnlineMoEMethod,
529 ),
530 (
531 + GRANITE_MODEL_NAME,
532 "fp8_per_channel",
533 None,
534 Fp8PtpcOnlineLinearMethod,
@@ -161,6 +536,7 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
536 ),
537 # quantization='online' with per-layer-kind overrides
538 (
539 + GRANITE_MODEL_NAME,
540 "online",
541 {
542 "linear": "fp8_per_block",
@@ -171,6 +547,7 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
547 ),
548 # quantization='online' with per-layer target patterns
549 (
550 + GRANITE_MODEL_NAME,
551 "online",
552 {
553 "targets": {
@@ -183,6 +560,7 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
560 ),
561 # ignore with direct layer name
562 (
563 + GRANITE_MODEL_NAME,
564 "fp8_per_tensor",
565 # qkv_proj is fused from q_proj/k_proj/v_proj. The shard regex
566 # remains supported alongside direct fused-name regexes.
@@ -191,11 +569,20 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
569 Fp8PerTensorOnlineMoEMethod,
570 ),
571 (
572 + GRANITE_MODEL_NAME,
573 "mxfp4",
574 None,
575 Mxfp4OnlineLinearMethod,
576 Mxfp4OnlineMoEMethod,
577 ),
578 + pytest.param(
579 + PARTIALLY_PREQUANTIZED_MODEL_NAME,
580 + None,
581 + {"targets": {"model.layers.1.self_attn.o_proj": "mxfp8"}},
582 + Mxfp8OnlineLinearMethod,
583 + CompressedTensorsMoEMethod,
584 + id="partially_prequantized_checkpoint",
585 + ),
586 ],
587 ids=[
588 "fp8_per_tensor",
@@ -205,13 +592,15 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
592 "targets",
593 "ignore",
594 "mxfp4",
595 + "partially_prequantized_checkpoint",
596 ],
597 )
598 @pytest.mark.parametrize(
599 "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
600 )
601 def test_online_quantization(
214 - quant_scheme: str,
602 + model_name: str,
603 + quant_scheme: str | None,
604 online_quant_args: dict | None,
605 expected_linear_cls,
606 expected_moe_cls,
@@ -243,38 +632,74 @@ def test_online_quantization(
632 ):
633 pytest.skip(f"Skip test for online {quant_scheme} on XPU platform.")
634
246 - model_name = "ibm-granite/granite-3.0-1b-a400m-base"
247 - model, vllm_config = load_model_without_vllm_runner(
248 - model_name,
249 - dtype="bfloat16",
250 - quantization=quant_scheme,
251 - model_config_kwargs={
252 - "quantization_config": resolve_quantization_config(
253 - quant_scheme, online_quant_args
254 - ),
255 - "hf_overrides": {
256 - "num_hidden_layers": 3,
257 - "vocab_size": 256,
258 - "hidden_size": 256,
259 - "intermediate_size": 512,
260 - "num_attention_heads": 4,
261 - "num_key_value_heads": 4,
262 - "max_position_embeddings": 64,
263 - "num_local_experts": 4,
264 - "num_experts_per_tok": 2,
635 + if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME:
636 + model, vllm_config = load_model_without_vllm_runner(
637 + model_name,
638 + model_config_kwargs={
639 + "quantization_config": resolve_quantization_config(
640 + None, online_quant_args
641 + )
642 },
266 - },
267 - model_loader_cls=DummyModelLoader,
268 - )
643 + model_loader_cls=DummyModelLoader,
644 + )
645 + assert isinstance(
646 + model.model.layers[1].self_attn.o_proj.quant_method,
647 + expected_linear_cls,
648 + )
649 + assert isinstance(
650 + model.model.layers[0].mlp.experts._quant_method,
651 + expected_moe_cls,
652 + )
653 + assert isinstance(
654 + model.model.layers[0].self_attn.o_proj.quant_method,
655 + CompressedTensorsLinearMethod,
656 + )
657 + else:
658 + model, vllm_config = load_model_without_vllm_runner(
659 + model_name,
660 + dtype="bfloat16",
661 + quantization=quant_scheme,
662 + model_config_kwargs={
663 + "quantization_config": resolve_quantization_config(
664 + quant_scheme, online_quant_args
665 + ),
666 + "hf_overrides": {
667 + "num_hidden_layers": 3,
668 + "vocab_size": 256,
669 + "hidden_size": 256,
670 + "intermediate_size": 512,
671 + "num_attention_heads": 4,
672 + "num_key_value_heads": 4,
673 + "max_position_embeddings": 64,
674 + "num_local_experts": 4,
675 + "num_experts_per_tok": 2,
676 + },
677 + },
678 + model_loader_cls=DummyModelLoader,
679 + )
680
681 monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
682
272 - o_proj = model.model.layers[0].self_attn.o_proj
273 - moe = model.model.layers[0].block_sparse_moe.experts
683 + if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME:
684 + o_proj = model.model.layers[1].self_attn.o_proj
685 + moe = model.model.layers[0].mlp.experts
686 + else:
687 + o_proj = model.model.layers[0].self_attn.o_proj
688 + moe = model.model.layers[0].block_sparse_moe.experts
689 assert isinstance(o_proj.quant_method, expected_linear_cls)
690 assert isinstance(moe._quant_method, expected_moe_cls)
691
277 - if quant_scheme == "mxfp4":
692 + if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME and isinstance(
693 + o_proj.quant_method.kernel, MarlinMxfp8LinearKernel
694 + ):
695 + assert o_proj.weight.dtype == torch.int32
696 + elif model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME and isinstance(
697 + o_proj.quant_method.kernel, EmulationMxfp8LinearKernel
698 + ):
699 + assert o_proj.weight.dtype == torch.bfloat16
700 + elif model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME:
701 + assert o_proj.weight.dtype == MXFP8_VALUE_DTYPE
702 + elif quant_scheme == "mxfp4":
703 assert o_proj.weight.dtype == torch.uint8
704 elif current_platform.is_cuda() or current_platform.is_xpu():
705 assert o_proj.weight.dtype == torch.float8_e4m3fn
@@ -301,7 +726,11 @@ def test_online_quantization(
726 layer.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod
727 )
728
304 - if isinstance(online_quant_args, dict) and "targets" in online_quant_args:
729 + if (
730 + model_name == GRANITE_MODEL_NAME
731 + and isinstance(online_quant_args, dict)
732 + and "targets" in online_quant_args
733 + ):
734 for layer in model.model.layers:
735 assert isinstance(
736 layer.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod
@@ -319,8 +748,13 @@ def test_online_quantization(
748 not is_quant_method_supported("fp8"),
749 reason="FP8 is not supported on this GPU type.",
750 )
322 -def test_online_quantization_loads_real_weights(vllm_runner, monkeypatch) -> None:
323 - """Verify online quantization loads a Granite-MoE checkpoint end to end."""
751 +@pytest.mark.parametrize(
752 + "model_name", [GRANITE_MODEL_NAME, PARTIALLY_PREQUANTIZED_MODEL_NAME]
753 +)
754 +def test_online_quantization_loads_real_weights(
755 + model_name, vllm_runner, monkeypatch
756 +) -> None:
757 + """Verify online quantization loads real checkpoints end to end."""
758 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
759 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
760
@@ -336,26 +770,52 @@ def test_online_quantization_loads_real_weights(vllm_runner, monkeypatch) -> Non
770
771 monkeypatch.setattr(GraniteMoeModel, "load_weights", load_weights)
772
339 - with vllm_runner(
340 - "ibm-granite/granite-3.0-1b-a400m-base",
341 - quantization="fp8_per_tensor",
342 - dtype="bfloat16",
343 - enforce_eager=True,
344 - hf_overrides={"num_hidden_layers": 3},
345 - max_model_len=16,
346 - max_num_seqs=1,
347 - ) as llm:
348 -
349 - def check_model(model):
350 - layer = model.model.layers[0]
351 - assert isinstance(
352 - layer.self_attn.o_proj.quant_method,
353 - Fp8PerTensorOnlineLinearMethod,
354 - )
355 - assert isinstance(
356 - layer.block_sparse_moe.experts._quant_method,
357 - Fp8PerTensorOnlineMoEMethod,
358 - )
773 + if model_name == GRANITE_MODEL_NAME:
774 + runner_kwargs = {
775 + "quantization": "fp8_per_tensor",
776 + "dtype": "bfloat16",
777 + "enforce_eager": True,
778 + "hf_overrides": {"num_hidden_layers": 3},
779 + "max_model_len": 16,
780 + "max_num_seqs": 1,
781 + }
782 + else:
783 + runner_kwargs = {
784 + "quantization_config": {
785 + "targets": {"model.layers.1.self_attn.o_proj": "mxfp8"}
786 + },
787 + "enforce_eager": True,
788 + }
789 +
790 + with vllm_runner(model_name, **runner_kwargs) as llm:
791 + if model_name == GRANITE_MODEL_NAME:
792 +
793 + def check_model(model):
794 + layer = model.model.layers[0]
795 + assert isinstance(
796 + layer.self_attn.o_proj.quant_method,
797 + Fp8PerTensorOnlineLinearMethod,
798 + )
799 + assert isinstance(
800 + layer.block_sparse_moe.experts._quant_method,
801 + Fp8PerTensorOnlineMoEMethod,
802 + )
803 +
804 + else:
805 +
806 + def check_model(model):
807 + assert isinstance(
808 + model.model.layers[1].self_attn.o_proj.quant_method,
809 + Mxfp8OnlineLinearMethod,
810 + )
811 + assert isinstance(
812 + model.model.layers[0].self_attn.o_proj.quant_method,
813 + CompressedTensorsLinearMethod,
814 + )
815 + assert isinstance(
816 + model.model.layers[0].mlp.experts._quant_method,
817 + CompressedTensorsMoEMethod,
818 + )
819
820 llm.apply_model(check_model)
821 outputs = llm.generate_greedy(["Hello my name is"], max_tokens=1)
@@ -457,8 +917,8 @@ def test_online_quantization_targets_ignore_collision() -> None:
917 )
918 )
919 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)
920 + config.resolve_quant_method_cls(
921 + Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj"
922 )
923
924
tests/quantization/test_online_mxfp4.py
+1 -1
@@ -354,7 +354,7 @@ def test_online_mxfp4_moe_matches_quark(
354 weight_quant_key=kMxfp4Static,
355 activation_quant_key=kMxfp4Dynamic,
356 )
357 - online_method = Mxfp4OnlineMoEMethod(layer=online_layer)
357 + online_method = Mxfp4OnlineMoEMethod(moe=online_layer.moe_config)
358
359 # `RoutedExperts.__init__` applies this round-up in production; these
360 # layers are built without a quant config, so it is applied explicitly.
tests/quantization/test_quantization_config_args.py
+11 -8
@@ -95,6 +95,12 @@ def test_resolve_shorthand_only_populates_both_slots():
95 assert args.moe == QuantSpec(weight=kFp8Static128BlockSym)
96
97
98 +@pytest.mark.parametrize("quantization", ["mxfp4", "mxfp8"])
99 +def test_resolve_colliding_shorthand_is_deferred(quantization: str):
100 + """Checkpoint metadata determines whether an MXFP shorthand is online."""
101 + assert resolve_quantization_config(quantization, None) is None
102 +
103 +
104 def test_resolve_int8_shorthand_leaves_linear_unset():
105 # int8_per_channel_weight_only is MoE-only; linear stays None so that
106 # OnlineQuantizationConfig leaves Linear layers in full precision.
@@ -122,12 +128,9 @@ def test_resolve_merges_explicit_over_shorthand():
128 assert args.moe == QuantSpec(weight=kFp8StaticTensorSym)
129
130
125 -def test_resolve_rejects_quantization_config_with_non_shorthand_quant():
126 - # If --quantization names something other than an online shorthand,
127 - # quantization_config is not allowed via this path (checkpoint quant
128 - # paths read it directly off ModelConfig instead).
129 - with pytest.raises(ValueError, match="quantization_config is only supported"):
130 - resolve_quantization_config("gptq", {"linear": "fp8_per_block"})
131 +def test_resolve_quantization_config_with_checkpoint_quantization():
132 + args = resolve_quantization_config("gptq", {"linear": "fp8_per_block"})
133 + assert args == QuantizationConfigArgs(linear="fp8_per_block")
134
135
136 # ---- QUANT_KEY_NAMES coverage --------------------------------------------
@@ -206,6 +209,6 @@ def test_targets_reject_moe_only_shorthand_for_linear_layer():
209 )
210
211 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)
212 + config.resolve_quant_method_cls(
213 + Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj"
214 )
vllm/config/quantization.py
+19 -7
@@ -225,6 +225,10 @@ ONLINE_QUANT_SHORTHAND_NAMES: tuple[str, ...] = (
225 "online",
226 )
227
228 +# These names are also checkpoint quantization methods. Their online configs
229 +# are resolved only when checkpoint quantization metadata is absent.
230 +_DEFERRED_ONLINE_SHORTHANDS = frozenset(("mxfp4", "mxfp8"))
231 +
232
233 def resolve_quantization_config(
234 quantization: str | None,
@@ -239,17 +243,25 @@ def resolve_quantization_config(
243 take precedence over the shorthand.
244 """
245 if quantization is not None and quantization not in ONLINE_QUANT_SHORTHAND_NAMES:
242 - if quantization_config is not None:
243 - raise ValueError(
244 - f"quantization_config is only supported when quantization is "
245 - f"one of {sorted(ONLINE_QUANT_SHORTHAND_NAMES)}, "
246 - f"got quantization={quantization!r}"
247 - )
248 - return None
246 + # Pre-quantized checkpoints can be composed with online quantization
247 + # for layers that the base quant_method leaves unquantized. The
248 + # checkpoint quant_method remains the primary quantization method; composition
249 + # is performed after its config has been loaded.
250 + if quantization_config is None:
251 + return None
252 +
253 + # `quantization_config` may hold both:
254 + # 1. Base quantization method activation key override,
255 + # 2. online quantization config to apply on top of the base quant_method.
256 + if isinstance(quantization_config, dict):
257 + return QuantizationConfigArgs(**quantization_config)
258 + return quantization_config
259
260 base = _ONLINE_SHORTHANDS.get(quantization) if quantization else None
261
262 if quantization_config is None:
263 + if quantization in _DEFERRED_ONLINE_SHORTHANDS:
264 + return None
265 return base
266
267 if isinstance(quantization_config, dict):
vllm/model_executor/layers/attention/attention.py
+8 -3
@@ -19,7 +19,10 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
19 from vllm.model_executor.layers.linear import (
20 UnquantizedLinearMethod,
21 )
22 -from vllm.model_executor.layers.quantization import QuantizationConfig
22 +from vllm.model_executor.layers.quantization import (
23 + QuantizationConfig,
24 + resolve_quant_method,
25 +)
26 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
27 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
28 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
@@ -183,7 +186,9 @@ def _init_kv_cache_quant(
186 layer._o_scale_float = None
187
188 quant_method = (
186 - quant_config.get_quant_method(layer, prefix=prefix) if quant_config else None
189 + resolve_quant_method(quant_config, layer, prefix=prefix)
190 + if quant_config
191 + else None
192 )
193
194 # See [Note: Register q/k/v/prob scales in state dict]
@@ -584,7 +589,7 @@ class Attention(nn.Module, AttentionLayerBase):
589 # as the default value. See [Note: Register q/k/v/prob scales in state dict]
590 # for more details.
591 quant_method = (
587 - self.quant_config.get_quant_method(self, prefix=self.layer_name)
592 + resolve_quant_method(self.quant_config, self, prefix=self.layer_name)
593 if self.quant_config
594 else None
595 )
vllm/model_executor/layers/attention/mla_attention.py
+5 -2
@@ -246,7 +246,10 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
246 from vllm.model_executor.layers.linear import (
247 ColumnParallelLinear,
248 )
249 -from vllm.model_executor.layers.quantization import QuantizationConfig
249 +from vllm.model_executor.layers.quantization import (
250 + QuantizationConfig,
251 + resolve_quant_method,
252 +)
253 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
254 from vllm.model_executor.layers.quantization.utils.quant_utils import (
255 GroupShape,
@@ -1255,7 +1258,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1258 # as the default value. See [Note: Register q/k/v/prob scales in state dict]
1259 # for more details.
1260 quant_method = (
1258 - self.quant_config.get_quant_method(self, prefix=self.layer_name)
1261 + resolve_quant_method(self.quant_config, self, prefix=self.layer_name)
1262 if self.quant_config
1263 else None
1264 )
vllm/model_executor/layers/fused_moe/routed_experts.py
+2 -1
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
25 )
26 from vllm.model_executor.layers.quantization.base_config import (
27 QuantizationConfig,
28 + resolve_quant_method,
29 )
30
31 if TYPE_CHECKING:
@@ -199,7 +200,7 @@ class RoutedExperts(PluggableLayer):
200 """
201 quant_method = None
202 if quant_config is not None:
202 - quant_method = quant_config.get_quant_method(self, prefix)
203 + quant_method = resolve_quant_method(quant_config, self, prefix)
204 if quant_method is None:
205 quant_method = UnquantizedFusedMoEMethod(moe_config)
206 assert isinstance(quant_method, FusedMoEMethodBase)
vllm/model_executor/layers/linear.py
+2 -1
@@ -24,6 +24,7 @@ from vllm.model_executor.custom_op import PluggableLayer
24 from vllm.model_executor.determinism.batch_invariant import (
25 linear_batch_invariant,
26 )
27 +from vllm.model_executor.layers.quantization import resolve_quant_method
28 from vllm.model_executor.layers.quantization.base_config import (
29 QuantizationConfig,
30 QuantizeMethodBase,
@@ -282,7 +283,7 @@ class LinearBase(PluggableLayer):
283 self.quant_method: QuantizeMethodBase
284 if quant_config is None:
285 self.quant_method = UnquantizedLinearMethod()
285 - elif quant_method := quant_config.get_quant_method(self, prefix=prefix):
286 + elif quant_method := resolve_quant_method(quant_config, self, prefix=prefix):
287 self.quant_method = quant_method
288 else:
289 raise ValueError("All linear layers should support quant method.")
vllm/model_executor/layers/quantization/__init__.py
+5 -1
@@ -4,7 +4,10 @@
4 from typing import Literal, get_args
5
6 from vllm.logger import init_logger
7 -from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
7 +from vllm.model_executor.layers.quantization.base_config import (
8 + QuantizationConfig,
9 + resolve_quant_method,
10 +)
11 from vllm.platforms import current_platform
12
13 logger = init_logger(__name__)
@@ -182,6 +185,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
185
186 __all__ = [
187 "QuantizationConfig",
188 + "resolve_quant_method",
189 "QuantizationMethods",
190 "get_quantization_config",
191 "register_quantization_config",
vllm/model_executor/layers/quantization/base_config.py
+49 -1
@@ -12,6 +12,9 @@ from transformers import PretrainedConfig
12
13 if TYPE_CHECKING:
14 from vllm.model_executor.layers.quantization import QuantizationMethods
15 + from vllm.model_executor.layers.quantization.online.base import (
16 + OnlineQuantizationConfig,
17 + )
18 from vllm.model_executor.models.utils import WeightsMapper
19 else:
20 QuantizationMethods = str
@@ -105,6 +108,7 @@ class QuantizationConfig(ABC):
108 """Suffixes of quantization parameters that may be present in the checkpoint but
109 not in the model, and should be ignored if unexpected during loading. These are used
110 after remapping, so should be in vLLM format (e.g. .q_scale, not .q.scale)."""
111 + online_quantization_config: "OnlineQuantizationConfig | None" = None
112
113 def __init__(self):
114 super().__init__()
@@ -187,7 +191,8 @@ class QuantizationConfig(ABC):
191 def get_quant_method(
192 self, layer: torch.nn.Module, prefix: str
193 ) -> QuantizeMethodBase | None:
190 - """Get the quantize method to use for the quantized layer.
194 + """Get the quantize method to use for the quantized layer, from the
195 + pre-quantized checkpoint quant_method.
196
197 Args:
198 layer: The layer for the quant method.
@@ -265,3 +270,46 @@ class QuantizationConfig(ABC):
270 # TODO: revision is never passed currently in vllm.py,
271 # but is used in subclasses, should we remove this parameter?
272 pass
273 +
274 +
275 +def resolve_quant_method(
276 + quant_config: QuantizationConfig, layer: torch.nn.Module, prefix: str
277 +) -> QuantizeMethodBase | None:
278 + """Return the checkpoint method with configured online quantization."""
279 + from vllm.model_executor.layers.fused_moe import RoutedExperts
280 + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
281 + UnquantizedFusedMoEMethod,
282 + )
283 + from vllm.model_executor.layers.linear import (
284 + LinearBase,
285 + UnquantizedLinearMethod,
286 + )
287 +
288 + base_quant_method = quant_config.get_quant_method(layer, prefix)
289 + if quant_config.online_quantization_config is None:
290 + return base_quant_method
291 + # Online quantization currently supports only LinearBase and RoutedExperts.
292 + # Embeddings and ParallelLMHead retain their checkpoint quantization method.
293 + if not isinstance(layer, (LinearBase, RoutedExperts)):
294 + return base_quant_method
295 +
296 + quant_config.online_quantization_config.packed_modules_mapping = (
297 + quant_config.packed_modules_mapping
298 + )
299 + checkpoint_is_quantized = base_quant_method is not None and not isinstance(
300 + base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod)
301 + )
302 + online_target = quant_config.online_quantization_config.resolve_quant_method_cls(
303 + layer, prefix
304 + )
305 + if checkpoint_is_quantized:
306 + if online_target is not None:
307 + raise ValueError(
308 + f"Cannot apply requested online quantization {online_target[3]} to "
309 + f"pre-quantized layer {prefix}: {base_quant_method} was already "
310 + "selected by the checkpoint quantization config."
311 + )
312 + return base_quant_method
313 + if online_target is None:
314 + return base_quant_method
315 + return quant_config.online_quantization_config.get_quant_method(layer, prefix)
vllm/model_executor/layers/quantization/experts_int8.py
+1 -1
@@ -56,5 +56,5 @@ class ExpertsInt8Config(QuantizationConfig):
56 if isinstance(layer, LinearBase):
57 return UnquantizedLinearMethod()
58 elif isinstance(layer, RoutedExperts):
59 - return Int8OnlineMoEMethod(layer=layer)
59 + return Int8OnlineMoEMethod(moe=layer.moe_config)
60 return None
vllm/model_executor/layers/quantization/fp8.py
+1 -1
@@ -219,7 +219,7 @@ class Fp8Config(QuantizationConfig):
219 Fp8PerTensorOnlineMoEMethod,
220 )
221
222 - return Fp8PerTensorOnlineMoEMethod(layer=layer)
222 + return Fp8PerTensorOnlineMoEMethod(moe=layer.moe_config)
223 elif isinstance(layer, Attention):
224 return Fp8KVCacheMethod(self)
225 return None
vllm/model_executor/layers/quantization/online/base.py
+116 -50
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Mapping
5 +from enum import Enum
6 from types import MappingProxyType
7 from typing import Any
8
@@ -14,6 +15,7 @@ from vllm.config.quantization import (
15 )
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.fused_moe import (
18 + FusedMoEMethodBase,
19 RoutedExperts,
20 )
21 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
@@ -38,6 +40,7 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
40 Fp8PerTensorOnlineMoEMethod,
41 Fp8PtpcOnlineLinearMethod,
42 Fp8PtpcOnlineMoEMethod,
43 + OnlineLinearBase,
44 )
45 from vllm.model_executor.layers.quantization.online.int8 import (
46 Int8OnlineMoEMethod,
@@ -71,6 +74,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
74 logger = init_logger(__name__)
75
76
77 +class OnlineQuantizationSource(str, Enum):
78 + """Supported online quantization configuration sources."""
79 +
80 + linear = "linear" # LinearBase
81 + moe = "moe" # RoutedExperts
82 + targets = "targets"
83 +
84 +
85 # Online dispatch tables, keyed by the QuantSpec.weight QuantKey. The
86 # corresponding method class handles the activation choice via its
87 # `supported_activation_quant` set.
@@ -200,12 +211,23 @@ class OnlineQuantizationConfig(QuantizationConfig):
211 "quantization='fp8_per_tensor'/'fp8_per_block' instead."
212 )
213
203 - def _dispatch(
214 + def _get_method_cls(
215 self,
216 spec: QuantSpec | None,
217 table: dict[QuantKey, type],
218 layer: torch.nn.Module,
208 - ) -> "QuantizeMethodBase | None":
219 + ) -> type | None:
220 + """Resolve the online method class for a layer's quantization spec.
221 +
222 + Args:
223 + spec: Quantization specification to resolve.
224 + table: Mapping from weight quantization keys to method classes.
225 + layer: Layer that will use the resolved method.
226 +
227 + Returns:
228 + The matching method class, or None when ``spec`` has no weight
229 + quantization.
230 + """
231 if spec is None or spec.weight is None:
232 return None
233 cls = table.get(spec.weight)
@@ -223,13 +245,75 @@ class OnlineQuantizationConfig(QuantizationConfig):
245 f"activation override (activation={spec.activation}) is not "
246 f"yet supported for online {cls.__name__}"
247 )
226 - if isinstance(layer, RoutedExperts):
227 - return cls(layer=layer)
228 - return cls()
248 + return cls
249 +
250 + def resolve_quant_method_cls(
251 + self, layer: torch.nn.Module, prefix: str
252 + ) -> tuple[OnlineQuantizationSource, str, str | None, QuantSpec, type] | None:
253 + """Resolve quantization metadata and method class without instantiating it.
254 +
255 + Args:
256 + layer: Layer for which to resolve online quantization.
257 + prefix: Fully qualified layer name.
258 +
259 + Returns:
260 + A tuple of source, quantization key string, target pattern, spec,
261 + and method class. Returns None when online quantization does not
262 + apply to the layer.
263 + """
264 + quant_spec: QuantSpec | None
265 + if self.args.targets is not None:
266 + resolved_pattern = self._resolve_targets_quant_method_metadata(
267 + prefix, layer
268 + )
269 + if resolved_pattern is None:
270 + return None
271 + source, quant_key_str, target_pattern, quant_spec, table = resolved_pattern
272 + else:
273 + if isinstance(layer, LinearBase):
274 + source = OnlineQuantizationSource.linear
275 + quant_spec = self.args.linear
276 + table = _ONLINE_LINEAR_METHODS
277 + elif isinstance(layer, RoutedExperts):
278 + source = OnlineQuantizationSource.moe
279 + quant_spec = self.args.moe
280 + table = _ONLINE_MOE_METHODS
281 + else:
282 + return None
283
230 - def _dispatch_target(
284 + if should_ignore_layer(
285 + prefix,
286 + ignore=self.ignored_layers,
287 + fused_mapping=self.packed_modules_mapping,
288 + use_fnmatch=True,
289 + ):
290 + return None
291 + quant_key_str = str(quant_spec)
292 + target_pattern = None
293 +
294 + quant_method_cls = self._get_method_cls(quant_spec, table, layer)
295 + if quant_method_cls is None:
296 + return None
297 + assert quant_spec is not None
298 + return source, quant_key_str, target_pattern, quant_spec, quant_method_cls
299 +
300 + def _resolve_targets_quant_method_metadata(
301 self, prefix: str, layer: torch.nn.Module
232 - ) -> "QuantizeMethodBase | None":
302 + ) -> (
303 + tuple[OnlineQuantizationSource, str, str, QuantSpec, dict[QuantKey, type]]
304 + | None
305 + ):
306 + """Resolve target-pattern quantization metadata and dispatch table.
307 +
308 + Args:
309 + prefix: Fully qualified layer name.
310 + layer: Layer matched against configured target patterns.
311 +
312 + Returns:
313 + A tuple of source, quantization key string, target pattern, spec,
314 + and dispatch table. Returns None when no pattern applies or the
315 + layer is ignored.
316 + """
317 assert self.args.targets is not None
318 ignored = should_ignore_layer(
319 prefix,
@@ -254,7 +338,8 @@ class OnlineQuantizationConfig(QuantizationConfig):
338 f"targets patterns: {matches}. Each layer may match at most "
339 f"one target."
340 )
257 - quant_key_str = self.args.targets[matches[0]]
341 + target_pattern = matches[0]
342 + quant_key_str = self.args.targets[target_pattern]
343 shorthand = _ONLINE_SHORTHANDS[quant_key_str]
344 if isinstance(layer, LinearBase):
345 quant_spec = shorthand.linear
@@ -265,63 +350,44 @@ class OnlineQuantizationConfig(QuantizationConfig):
350 else:
351 raise ValueError(
352 f"Layer {prefix} was matched by quantization_config.targets "
268 - f"({matches[0]}), but online quantization is not supported for "
353 + f"({target_pattern}), but online quantization is not supported for "
354 f"{type(layer).__name__}."
355 )
356 if quant_spec is None:
357 raise ValueError(
273 - f"targets pattern {matches[0]} = {quant_key_str} does "
358 + f"targets pattern {target_pattern} = {quant_key_str} does "
359 f"not define a QuantSpec for {type(layer).__name__} layers "
360 f"(matched at {prefix})."
361 )
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
362 + return (
363 + OnlineQuantizationSource.targets,
364 + quant_key_str,
365 + target_pattern,
366 + quant_spec,
367 + table,
368 + )
369
370 def get_quant_method(
371 self, layer: torch.nn.Module, prefix: str
372 ) -> "QuantizeMethodBase | None":
373 # `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()
374 + resolved = self.resolve_quant_method_cls(layer, prefix)
375 + if resolved is not None:
376 + source, quant_key_str, target_pattern, _, quant_method_cls = resolved
377 + self.quantized_layers[prefix] = (
378 + source.value,
379 + quant_key_str,
380 + target_pattern,
381 + )
382 if isinstance(layer, RoutedExperts):
294 - return UnquantizedFusedMoEMethod(layer.moe_config)
295 - return None
383 + assert issubclass(quant_method_cls, FusedMoEMethodBase)
384 + return quant_method_cls(moe=layer.moe_config)
385 +
386 + assert issubclass(quant_method_cls, OnlineLinearBase)
387 + return quant_method_cls()
388
389 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)
306 - if method is not None:
307 - self.quantized_layers[prefix] = (
308 - "linear",
309 - str(self.args.linear),
310 - None,
311 - )
312 - return method
390 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)
323 - if method is not None:
324 - self.quantized_layers[prefix] = ("moe", str(self.args.moe), None)
325 - return method
391 + if isinstance(layer, RoutedExperts):
392 return UnquantizedFusedMoEMethod(layer.moe_config)
393 return None
vllm/model_executor/layers/quantization/online/fp8.py
+13 -12
@@ -23,6 +23,7 @@ from vllm.model_executor.kernels.linear.scaled_mm import (
23 MarlinFP8ScaledMMLinearKernel,
24 )
25 from vllm.model_executor.layers.fused_moe import RoutedExperts
26 +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
27 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
28 select_fp8_moe_backend,
29 )
@@ -111,7 +112,7 @@ def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool:
112 return is_row_parallel or (reduces_output_dim and is_column_parallel)
113
114
114 -class _Fp8OnlineLinearBase(LinearMethodBase):
115 +class OnlineLinearBase(LinearMethodBase):
116 """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
117 weights onto meta device and materializes them just-in-time."""
118
@@ -155,7 +156,7 @@ class _Fp8OnlineLinearBase(LinearMethodBase):
156 initialize_online_processing(layer)
157
158
158 -class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
159 +class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase):
160 """Online tensorwise FP8 linear quantization.
161 Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
162
@@ -256,7 +257,7 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
257 return self.fp8_linear.apply_weights(layer, x, bias)
258
259
259 -class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
260 +class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase):
261 """Online blockwise FP8 linear quantization.
262 Loads fp16/bf16 weights and quantizes them per-block during loading."""
263
@@ -336,7 +337,7 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
337 )
338
339
339 -class Fp8PtpcOnlineLinearMethod(_Fp8OnlineLinearBase):
340 +class Fp8PtpcOnlineLinearMethod(OnlineLinearBase):
341 """Online PTPC FP8 linear quantization.
342
343 Per-output-channel weight scale + dynamic per-token activation scale. The
@@ -443,12 +444,12 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
444 self,
445 *,
446 weight_block_size: list[int] | None,
446 - layer: torch.nn.Module,
447 + moe: FusedMoEConfig,
448 weight_key: "QuantKey | None" = None,
449 activation_key: "QuantKey | None" = None,
450 allow_vllm_cutlass: bool = False,
451 ):
451 - super().__init__(layer.moe_config)
452 + super().__init__(moe)
453 self.weight_block_size = weight_block_size
454 self.block_quant: bool = self.weight_block_size is not None
455 self.weight_scale_name = (
@@ -555,11 +556,11 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
556 def __init__(
557 self,
558 *,
558 - layer: torch.nn.Module,
559 + moe: FusedMoEConfig,
560 ):
561 super().__init__(
562 weight_block_size=None,
562 - layer=layer,
563 + moe=moe,
564 )
565
566 def process_weights_after_loading(self, layer: Module) -> None:
@@ -612,11 +613,11 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
613 def __init__(
614 self,
615 *,
615 - layer: torch.nn.Module,
616 + moe: FusedMoEConfig,
617 ):
618 super().__init__(
619 weight_block_size=[128, 128],
619 - layer=layer,
620 + moe=moe,
621 )
622
623 def maybe_roundup_sizes(
@@ -716,13 +717,13 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
717 def __init__(
718 self,
719 *,
719 - layer: torch.nn.Module,
720 + moe: FusedMoEConfig,
721 ):
722 from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
723
724 super().__init__(
725 weight_block_size=None,
725 - layer=layer,
726 + moe=moe,
727 weight_key=kFp8StaticChannelSym,
728 activation_key=kFp8DynamicTokenSym,
729 allow_vllm_cutlass=True,
vllm/model_executor/layers/quantization/online/int8.py
+3 -2
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
12 )
13
14 from vllm.model_executor.layers.fused_moe import RoutedExperts
15 +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
16 from vllm.model_executor.layers.fused_moe.oracle.int8 import (
17 convert_to_int8_moe_kernel_format,
18 make_int8_moe_kernel,
@@ -38,9 +39,9 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
39 def __init__(
40 self,
41 *,
41 - layer: torch.nn.Module,
42 + moe: FusedMoEConfig,
43 ):
43 - super().__init__(layer.moe_config)
44 + super().__init__(moe)
45 self.int8_backend, self.experts_cls = select_int8_moe_backend(
46 config=self.moe,
47 weight_key=kInt8StaticChannelSym,
vllm/model_executor/layers/quantization/online/mxfp4.py
+5 -4
@@ -17,6 +17,7 @@ if TYPE_CHECKING:
17 )
18
19 from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
20 +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
21 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
22 TRITON_BACKENDS,
23 Mxfp4MoeBackend,
@@ -27,7 +28,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
28 select_mxfp4_moe_backend,
29 )
30 from vllm.model_executor.layers.quantization.online.fp8 import (
30 - _Fp8OnlineLinearBase,
31 + OnlineLinearBase,
32 )
33 from vllm.model_executor.layers.quantization.online.moe_base import (
34 OnlineMoEMethodBase,
@@ -73,7 +74,7 @@ def _quantize_mxfp4_moe_weight(
74 return w_quant, w_scales
75
76
76 -class Mxfp4OnlineLinearMethod(_Fp8OnlineLinearBase):
77 +class Mxfp4OnlineLinearMethod(OnlineLinearBase):
78 """Online MXFP4 linear method.
79 Loads bf16/fp16 checkpoints and quantizes weights to MXFP4 (microscaling
80 FP4 with block-32 scales) during weight loading.
@@ -144,8 +145,8 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase):
145 experts_cls: "type[mk.FusedMoEExperts] | None"
146 activation_quant_key = kMxfp4Dynamic
147
147 - def __init__(self, *, layer: torch.nn.Module):
148 - super().__init__(layer.moe_config)
148 + def __init__(self, *, moe: FusedMoEConfig):
149 + super().__init__(moe)
150 self.weight_block_size: list[int] = [1, MXFP4_BLOCK_SIZE]
151 self.weight_scale_name = "weight_scale"
152
vllm/model_executor/layers/quantization/online/mxfp8.py
+6 -4
@@ -8,6 +8,8 @@ from typing import TYPE_CHECKING
8 import torch
9 from torch.nn import Module
10
11 +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
12 +
13 if TYPE_CHECKING:
14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
15 from vllm.model_executor.layers.fused_moe import (
@@ -21,7 +23,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
23 select_mxfp8_moe_backend,
24 )
25 from vllm.model_executor.layers.quantization.online.fp8 import (
24 - _Fp8OnlineLinearBase,
26 + OnlineLinearBase,
27 )
28 from vllm.model_executor.layers.quantization.online.moe_base import (
29 OnlineMoEMethodBase,
@@ -34,7 +36,7 @@ from vllm.model_executor.utils import replace_parameter
36 from vllm.platforms import current_platform
37
38
37 -class Mxfp8OnlineLinearMethod(_Fp8OnlineLinearBase):
39 +class Mxfp8OnlineLinearMethod(OnlineLinearBase):
40 """Online MXFP8 linear method.
41 Loads bf16/fp16 checkpoints and quantizes weights to MXFP8 (microscaling
42 FP8 with block-32 scales) during weight loading.
@@ -100,8 +102,8 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
102 fp8_backend: "Fp8MoeBackend"
103 experts_cls: "type[mk.FusedMoEExperts] | None"
104
103 - def __init__(self, *, layer: torch.nn.Module):
104 - super().__init__(layer.moe_config)
105 + def __init__(self, *, moe: FusedMoEConfig):
106 + super().__init__(moe)
107 self.weight_block_size: list[int] = [1, MXFP8_BLOCK_SIZE]
108 self.weight_scale_name = "weight_scale"
109
vllm/model_executor/layers/quantization/online/nvfp4.py
+6 -3
@@ -6,7 +6,10 @@ from torch.nn import Module
6
7 from vllm._custom_ops import scaled_fp4_quant
8 from vllm.model_executor.layers.fused_moe import RoutedExperts
9 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9 +from vllm.model_executor.layers.fused_moe.config import (
10 + FusedMoEConfig,
11 + FusedMoEQuantConfig,
12 +)
13 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
14 convert_to_nvfp4_moe_kernel_format,
15 make_nvfp4_moe_kernel,
@@ -87,13 +90,13 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
90 def __init__(
91 self,
92 *,
90 - layer: torch.nn.Module,
93 + moe: FusedMoEConfig,
94 ):
95 if not current_platform.is_device_capability_family(100):
96 raise ValueError(
97 "nvfp4_per_token online quantization requires a Blackwell (SM100) GPU."
98 )
96 - super().__init__(layer.moe_config)
99 + super().__init__(moe)
100 self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
101 config=self.moe,
102 weight_key=kNvfp4Static,
vllm/model_executor/layers/vocab_parallel_embedding.py
+2 -1
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.base_config import (
24 QuantizationConfig,
25 QuantizeMethodBase,
26 method_has_implemented_embedding,
27 + resolve_quant_method,
28 )
29 from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
30 from vllm.model_executor.parameter import BasevLLMParameter
@@ -291,7 +292,7 @@ class VocabParallelEmbedding(PluggableLayer):
292 # Avoid overriding a preselected model-specific method with generic
293 # config-based dispatch.
294 if quant_method is None and quant_config is not None:
294 - quant_method = quant_config.get_quant_method(self, prefix=prefix)
295 + quant_method = resolve_quant_method(quant_config, self, prefix=prefix)
296 if quant_method is None:
297 quant_method = UnquantizedEmbeddingMethod()
298
vllm/model_executor/model_loader/base_loader.py
+5
@@ -100,6 +100,11 @@ def log_online_quantization(vllm_config: VllmConfig) -> None:
100 )
101
102 quant_config = vllm_config.quant_config
103 + online_quantization_config = getattr(
104 + quant_config, "online_quantization_config", None
105 + )
106 + if isinstance(online_quantization_config, OnlineQuantizationConfig):
107 + quant_config = online_quantization_config
108 if not isinstance(quant_config, OnlineQuantizationConfig):
109 return
110
vllm/model_executor/model_loader/weight_utils.py
+48 -16
@@ -243,6 +243,26 @@ def get_quant_config(
243 if model_config.quantization is None:
244 raise ValueError("Model quantization method is not specified in the config.")
245 quant_cls = get_quantization_config(model_config.quantization)
246 + from vllm.config.quantization import _ONLINE_SHORTHANDS, QuantizationConfigArgs
247 + from vllm.model_executor.layers.quantization.online.base import (
248 + OnlineQuantizationConfig,
249 + )
250 +
251 + # resolve_quantization_config does not resolve `quantization_config` from
252 + # ambiguous shorthands as `"mxfp4"`, `"mxfp8"`.
253 + online_args = model_config.quantization_config
254 +
255 + def maybe_compose_online_quantization(
256 + checkpoint_config: QuantizationConfig,
257 + ) -> QuantizationConfig:
258 + if online_args is None:
259 + return checkpoint_config
260 + assert isinstance(online_args, QuantizationConfigArgs)
261 +
262 + checkpoint_config.online_quantization_config = OnlineQuantizationConfig(
263 + online_args
264 + )
265 + return checkpoint_config
266
267 # Read the quantization config from the HF model config, if available.
268 hf_quant_config = getattr(model_config.hf_config, "quantization_config", None)
@@ -284,7 +304,9 @@ def get_quant_config(
304 ):
305 pass # fall through to file-based loading below
306 else:
287 - return quant_cls.from_config(hf_quant_config)
307 + return maybe_compose_online_quantization(
308 + quant_cls.from_config(hf_quant_config)
309 + )
310
311 # if hf_quant_config is None, we will try to get config from
312 # hf_overrides
@@ -302,7 +324,9 @@ def get_quant_config(
324 quantization_config_file = hf_overrides.get("quantization_config_file")
325 if quantization_config_file is not None:
326 if hasattr(quant_cls, "from_config_file"):
305 - return quant_cls.from_config_file(quantization_config_file)
327 + return maybe_compose_online_quantization(
328 + quant_cls.from_config_file(quantization_config_file)
329 + )
330 else:
331 raise NotImplementedError(
332 "from_config_file is specified in hf_override config, "
@@ -312,7 +336,9 @@ def get_quant_config(
336 quantization_config_json = hf_overrides.get("quantization_config_dict_json")
337 if quantization_config_json is not None:
338 if hasattr(quant_cls, "from_config_dict_json"):
315 - return quant_cls.from_config_dict_json(quantization_config_json)
339 + return maybe_compose_online_quantization(
340 + quant_cls.from_config_dict_json(quantization_config_json)
341 + )
342 else:
343 raise NotImplementedError(
344 "from_config_dict_json is specified in hf_override config, "
@@ -320,16 +346,14 @@ def get_quant_config(
346 f"{quant_cls}"
347 )
348
323 - # Online quantization doesn't read from checkpoint configs - it quantizes
324 - # fp16/bf16 weights on the fly during loading.
325 - if model_config.quantization_config is not None:
326 - from vllm.config.quantization import QuantizationConfigArgs
327 - from vllm.model_executor.layers.quantization.online.base import (
328 - OnlineQuantizationConfig,
329 - )
330 -
331 - assert isinstance(model_config.quantization_config, QuantizationConfigArgs)
332 - return OnlineQuantizationConfig(args=model_config.quantization_config)
349 + # Raw unambiguous online quantization doesn't read from checkpoint configs
350 + # We must continue to load/read the config below in two cases:
351 + # 1. composed online quantization, before `maybe_compose_online_quantization`,
352 + # 2. Ambiguous online shorthands as `"mxpf4"`, `"mxfp8"`, for which `online_args`
353 + # is not set yet.
354 + if quant_cls is OnlineQuantizationConfig and online_args is not None:
355 + assert isinstance(online_args, QuantizationConfigArgs)
356 + return OnlineQuantizationConfig(args=online_args)
357
358 model_name_or_path = (
359 maybe_download_from_modelscope(
@@ -359,7 +383,11 @@ def get_quant_config(
383
384 # If the quantization config is not found, use the default config.
385 if not possible_config_filenames:
362 - return quant_cls()
386 + if model_config.quantization in _ONLINE_SHORTHANDS:
387 + args = online_args or _ONLINE_SHORTHANDS[model_config.quantization]
388 + assert isinstance(args, QuantizationConfigArgs)
389 + return OnlineQuantizationConfig(args=args)
390 + return maybe_compose_online_quantization(quant_cls())
391
392 config_files = glob.glob(os.path.join(hf_folder, "*.json"))
393
@@ -367,6 +395,10 @@ def get_quant_config(
395 f for f in config_files if any(f.endswith(x) for x in possible_config_filenames)
396 ]
397 if len(quant_config_files) == 0:
398 + if model_config.quantization in _ONLINE_SHORTHANDS:
399 + args = online_args or _ONLINE_SHORTHANDS[model_config.quantization]
400 + assert isinstance(args, QuantizationConfigArgs)
401 + return OnlineQuantizationConfig(args=args)
402 raise ValueError(f"Cannot find the config file for {model_config.quantization}")
403 if len(quant_config_files) > 1:
404 raise ValueError(
@@ -380,14 +412,14 @@ def get_quant_config(
412
413 if model_config.quantization in ("modelopt", "modelopt_mixed"):
414 if config.get("producer", {}).get("name") == "modelopt":
383 - return quant_cls.from_config(config)
415 + return maybe_compose_online_quantization(quant_cls.from_config(config))
416 else:
417 raise ValueError(
418 f"Unsupported quantization config"
419 f" found for {model_config.quantization} in {f}."
420 )
421
390 - return quant_cls.from_config(config)
422 + return maybe_compose_online_quantization(quant_cls.from_config(config))
423
424
425 def get_sparse_attention_config(
vllm/models/kimi_k3/nvidia/mla.py
+5 -2
@@ -69,7 +69,10 @@ from vllm.model_executor.layers.linear import (
69 RowParallelLinear,
70 UnquantizedLinearMethod,
71 )
72 -from vllm.model_executor.layers.quantization import QuantizationConfig
72 +from vllm.model_executor.layers.quantization import (
73 + QuantizationConfig,
74 + resolve_quant_method,
75 +)
76 from vllm.model_executor.layers.quantization.utils.quant_utils import (
77 get_and_maybe_dequant_weights,
78 )
@@ -462,7 +465,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
465 replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)
466
467 quant_method = (
465 - self.quant_config.get_quant_method(self, prefix=self.layer_name)
468 + resolve_quant_method(self.quant_config, self, prefix=self.layer_name)
469 if self.quant_config
470 else None
471 )