[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
)