@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7b1a7423be
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 14:45 UTC

[Frontend] new online quantization frontend (#38138)

Signed-off-by: Vasiliy Kuznetsov <vasiliy@meta.com>

Vasiliy Kuznetsov committed Apr 3, 2026 at 11:58 UTC 7b1a7423bea1705bd51d838e34bef99e8a01cbbd
13 files changed +1205
tests/quantization/test_online.py new
+179
@@ -0,0 +1,179 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests online quantization."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from tests.quantization.utils import (
9 + _test_online_quant_peak_mem_impl,
10 + is_quant_method_supported,
11 +)
12 +from vllm.model_executor.layers.linear import UnquantizedLinearMethod
13 +from vllm.model_executor.layers.quantization.online.fp8 import (
14 + Fp8PerBlockOnlineLinearMethod,
15 + Fp8PerBlockOnlineMoEMethod,
16 + Fp8PerTensorOnlineLinearMethod,
17 + Fp8PerTensorOnlineMoEMethod,
18 +)
19 +from vllm.platforms import current_platform
20 +
21 +
22 +@pytest.mark.skipif(
23 + not is_quant_method_supported("fp8"),
24 + reason="FP8 is not supported on this GPU type.",
25 +)
26 +@pytest.mark.parametrize(
27 + "quant_scheme,online_quant_args,expected_linear_cls,expected_moe_cls",
28 + [
29 + # simple case - quantization='fp8_per_tensor'
30 + (
31 + "fp8_per_tensor",
32 + None,
33 + Fp8PerTensorOnlineLinearMethod,
34 + Fp8PerTensorOnlineMoEMethod,
35 + ),
36 + # simple case - quantization='fp8_per_block'
37 + (
38 + "fp8_per_block",
39 + None,
40 + Fp8PerBlockOnlineLinearMethod,
41 + Fp8PerBlockOnlineMoEMethod,
42 + ),
43 + # quantization='online with linear_scheme_override and
44 + # moe_scheme_override
45 + (
46 + "online",
47 + {
48 + "linear_scheme_override": "fp8_per_block",
49 + "moe_scheme_override": "fp8_per_tensor",
50 + },
51 + Fp8PerBlockOnlineLinearMethod,
52 + Fp8PerTensorOnlineMoEMethod,
53 + ),
54 + # ignore with direct layer name
55 + (
56 + "fp8_per_tensor",
57 + # qkv_proj is fused from q_proj/k_proj/v_proj, so currently the
58 + # ignore regex must match the unfused shard names
59 + # TODO(future PR): also make 're:.*qkv_proj.*' work
60 + {"ignore": ["model.layers.1.self_attn.o_proj", "re:.*[qkv]_proj"]},
61 + Fp8PerTensorOnlineLinearMethod,
62 + Fp8PerTensorOnlineMoEMethod,
63 + ),
64 + ],
65 +)
66 +@pytest.mark.parametrize(
67 + "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
68 +)
69 +def test_online_quantization(
70 + vllm_runner,
71 + quant_scheme: str,
72 + online_quant_args: dict | None,
73 + expected_linear_cls,
74 + expected_moe_cls,
75 + use_rocm_aiter: bool,
76 + monkeypatch,
77 +) -> None:
78 + """
79 + Tests that online quantization frontend configuration works -
80 + selecting quant schemes, overriding quant schemes by type, ignoring
81 + layers.
82 +
83 + Does not test performance, peak memory usage, etc.
84 + """
85 +
86 + if use_rocm_aiter:
87 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
88 +
89 + # `LLM.apply_model` requires pickling a function.
90 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
91 +
92 + # a tiny model with both dense and MoE layers
93 + model_name = "ibm-granite/granite-3.0-1b-a400m-base"
94 +
95 + runner_kwargs = dict(
96 + quantization=quant_scheme,
97 + enforce_eager=True,
98 + )
99 + if online_quant_args is not None:
100 + runner_kwargs["quantization_config"] = online_quant_args
101 +
102 + with vllm_runner(
103 + model_name,
104 + **runner_kwargs,
105 + ) as llm:
106 +
107 + def check_model(model):
108 + # checks further down in the test case are hardcoded for this
109 + # model
110 + assert model_name == "ibm-granite/granite-3.0-1b-a400m-base"
111 +
112 + o_proj = model.model.layers[0].self_attn.o_proj
113 + moe = model.model.layers[0].block_sparse_moe.experts
114 +
115 + # o_proj and moe in layer 0 are always quantized (never ignored)
116 + # because of how we craft the test case inputs
117 + assert isinstance(o_proj.quant_method, expected_linear_cls)
118 + if moe is not None:
119 + assert isinstance(moe.quant_method, expected_moe_cls)
120 +
121 + if current_platform.is_cuda():
122 + assert o_proj.weight.dtype == torch.float8_e4m3fn
123 + elif current_platform.is_rocm():
124 + assert o_proj.weight.dtype == current_platform.fp8_dtype()
125 + else:
126 + pytest.skip("Only runs on CUDA and ROCm.")
127 +
128 + # Verify ignored layers are unquantized.
129 + if isinstance(online_quant_args, dict) and "ignore" in online_quant_args:
130 + # only .*1.self_attn_o_proj is skipped
131 + for layer_idx in range(len(model.model.layers)):
132 + o_proj = model.model.layers[layer_idx].self_attn.o_proj
133 + if layer_idx == 1:
134 + assert isinstance(o_proj.quant_method, UnquantizedLinearMethod)
135 + else:
136 + assert isinstance(o_proj.quant_method, expected_linear_cls)
137 +
138 + # every .*self_attn.qkv_proj is skipped
139 + for layer_idx in range(len(model.model.layers)):
140 + qkv_proj = model.model.layers[layer_idx].self_attn.qkv_proj
141 + assert isinstance(qkv_proj.quant_method, UnquantizedLinearMethod)
142 +
143 + llm.apply_model(check_model)
144 +
145 + outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
146 + print(outputs[0][1])
147 +
148 +
149 +@pytest.mark.skipif(
150 + not is_quant_method_supported("fp8"),
151 + reason="FP8 is not supported on this GPU type.",
152 +)
153 +def test_online_quant_peak_mem(
154 + vllm_runner,
155 + caplog_mp_spawn,
156 + monkeypatch,
157 +) -> None:
158 + _test_online_quant_peak_mem_impl(
159 + "fp8_per_tensor", vllm_runner, caplog_mp_spawn, monkeypatch
160 + )
161 +
162 +
163 +@pytest.mark.skipif(
164 + not is_quant_method_supported("fp8"),
165 + reason="FP8 is not supported on this GPU type.",
166 +)
167 +def test_online_quant_load_format_dummy(
168 + vllm_runner,
169 + monkeypatch,
170 + caplog,
171 +) -> None:
172 + with vllm_runner(
173 + "ibm-granite/granite-3.0-1b-a400m-base",
174 + quantization="fp8_per_tensor",
175 + enforce_eager=True,
176 + load_format="dummy",
177 + ) as llm:
178 + outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
179 + print(outputs[0][1])
tests/quantization/utils.py
+75
@@ -1,6 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import logging
5 +
6 +import regex as re
7 +
8 from vllm.model_executor.layers.quantization import get_quantization_config
9 from vllm.platforms import current_platform
10
@@ -21,3 +25,74 @@ def is_quant_method_supported(quant_method: str) -> bool:
25 min_capability = get_quantization_config(quant_method).get_min_capability()
26
27 return capability.to_int() >= min_capability
28 +
29 +
30 +def _test_online_quant_peak_mem_impl(
31 + quantization_arg_value,
32 + vllm_runner,
33 + caplog_mp_spawn,
34 + monkeypatch,
35 +) -> None:
36 + # Note: `allenai/OLMoE-1B-7B-0125-Instruct` was selected because:
37 + # 1. it covers both Linear and MoE paths
38 + # 2. it is already used by other tests in CI, so adding it here
39 + # does not increase disk space for CI runners
40 + # I really wanted to use `ibm-granite/granite-3.0-1b-a400m-base`
41 + # which I think is the smallest MoE model in vLLM (2.5 GiB bf16,
42 + # 1.3 GiB fp8), but could not as adding one more model makes CI
43 + # run out of disk space.
44 + model_name = "allenai/OLMoE-1B-7B-0125-Instruct"
45 +
46 + # Force spawn to ensure caplog_mp_spawn works consistently
47 + # (it relies on VLLM_LOGGING_CONFIG_PATH which spawn reads but fork ignores)
48 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
49 +
50 + with (
51 + caplog_mp_spawn(logging.DEBUG) as log_holder,
52 + vllm_runner(
53 + model_name,
54 + quantization=quantization_arg_value,
55 + enforce_eager=True,
56 + ) as llm,
57 + ):
58 + outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
59 + print(outputs[0][1])
60 +
61 + log_text = log_holder.text
62 +
63 + # Parse memory usage from captured logs
64 + model_memory_gib = None
65 + peak_memory_gib = None
66 + for line in log_text.splitlines():
67 + if model_memory_gib is None:
68 + match = re.search(r"Model loading took ([\d.]+) GiB memory", line)
69 + if match:
70 + model_memory_gib = float(match.group(1))
71 + if peak_memory_gib is None:
72 + match = re.search(
73 + r"Peak GPU memory after loading weights: ([\d.]+) GiB", line
74 + )
75 + if match:
76 + peak_memory_gib = float(match.group(1))
77 +
78 + assert model_memory_gib is not None, "Could not find model loading memory log"
79 + assert peak_memory_gib is not None, "Could not find peak memory log"
80 + print(f"GPU memory used after loading weights: {model_memory_gib} GiB")
81 + print(f"Peak GPU memory usage while loading weights: {peak_memory_gib} GiB")
82 +
83 + # model specific, allenai/OLMoE-1B-7B-0125-Instruct fp8 online quant
84 + # uses 6.65 GiB for weight loading (bf16 checkpoint is ~12.89 GiB)
85 + expected_model_memory_gib = 6.7
86 +
87 + # for allenai/OLMoE-1B-7B-0125-Instruct the number we see today is 9.06
88 + # GiB, which is 1.36x above model_memory_gib. A slightly higher number is
89 + # expected as when we load and quantize weights in a streaming fashion we
90 + # need to have individual weights in bf16 + fp8 alive at the same time.
91 + expected_peak_memory_gib = expected_model_memory_gib * 1.4
92 +
93 + assert model_memory_gib < expected_model_memory_gib, (
94 + f"{model_memory_gib=} higher than {expected_model_memory_gib}"
95 + )
96 + assert peak_memory_gib < expected_peak_memory_gib, (
97 + f"{peak_memory_gib=} higher than {expected_peak_memory_gib}"
98 + )
vllm/config/model.py
+5
@@ -21,6 +21,7 @@ from vllm.config.multimodal import (
21 MultiModalConfig,
22 )
23 from vllm.config.pooler import PoolerConfig
24 +from vllm.config.quantization import OnlineQuantizationConfigArgs
25 from vllm.config.scheduler import RunnerType
26 from vllm.config.utils import config, getattr_iter
27 from vllm.logger import init_logger
@@ -199,6 +200,10 @@ class ModelConfig:
200 `quantization_config` attribute in the model config file. If that is
201 `None`, we assume the model weights are not quantized and use `dtype` to
202 determine the data type of the weights."""
203 + quantization_config: dict[str, Any] | OnlineQuantizationConfigArgs | None = None
204 + """Arguments for online quantization.
205 + Auto-created when `quantization` equals to one of the string values of
206 + the `OnlineQuantScheme` enum."""
207 allow_deprecated_quantization: bool = False
208 """Whether to allow deprecated quantization methods."""
209 enforce_eager: bool = False
vllm/config/quantization.py new
+121
@@ -0,0 +1,121 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from enum import Enum
5 +from typing import Any
6 +
7 +from pydantic import Field, field_validator
8 +
9 +from vllm.config.utils import config
10 +
11 +
12 +class OnlineQuantScheme(Enum):
13 + """Supported online quantization schemes."""
14 +
15 + # fp8, weights and activations scaled per-tensor
16 + FP8_PER_TENSOR = "fp8_per_tensor"
17 +
18 + # fp8, activations scaled in blocks of 1x128 elements, weights scaled in
19 + # blocks of 128x128 elements (popularized by DeepSeek)
20 + FP8_PER_BLOCK = "fp8_per_block"
21 +
22 + # TODO(future PRs): add more online quant schemes here: mxfp8, etc
23 +
24 +
25 +@config
26 +class OnlineQuantizationConfigArgs:
27 + """Configuration for online quantization.
28 +
29 + Controls how ``OnlineQuantizationConfig`` is applied to a model.
30 + At least one of ``global_scheme``, ``linear_scheme_override``, or
31 + ``moe_scheme_override`` must be set.
32 + """
33 +
34 + global_scheme: OnlineQuantScheme | None = None
35 + """Quantization scheme applied to every supported layer."""
36 +
37 + linear_scheme_override: OnlineQuantScheme | None = None
38 + """Quantization scheme override for ``LinearBase`` layers."""
39 +
40 + moe_scheme_override: OnlineQuantScheme | None = None
41 + """Quantization scheme override for ``FusedMoE`` layers."""
42 +
43 + ignore: list[str] = Field(default_factory=list)
44 + """Layers to skip quantization for. Supports exact names and regex
45 + patterns with ``re:`` prefix (e.g. ``re:.*attn.*``), consistent with
46 + compressed_tensors layer skipping."""
47 +
48 + @field_validator(
49 + "global_scheme", "linear_scheme_override", "moe_scheme_override", mode="before"
50 + )
51 + @classmethod
52 + def _coerce_scheme(
53 + cls, v: str | OnlineQuantScheme | None
54 + ) -> OnlineQuantScheme | None:
55 + if isinstance(v, str):
56 + return OnlineQuantScheme(v)
57 + return v
58 +
59 +
60 +def resolve_online_quant_config(
61 + quantization: str | None,
62 + quantization_config: dict[str, Any] | OnlineQuantizationConfigArgs | None,
63 +) -> OnlineQuantizationConfigArgs | None:
64 + """Resolve online quant scheme shorthand into a quantization config.
65 +
66 + If ``quantization`` is an online quant scheme (e.g. ``'fp8_per_tensor'``),
67 + ensures ``quantization_config`` has a matching ``global_scheme`` and casts
68 + it to :class:`OnlineQuantizationConfigArgs` if needed.
69 + """
70 + online_quant_values = {s.value for s in OnlineQuantScheme}
71 + valid_quantization_values = online_quant_values | {"online"}
72 + if quantization not in valid_quantization_values:
73 + if quantization_config is not None:
74 + raise ValueError(
75 + f"quantization_config is only supported when quantization "
76 + f"is one of {sorted(valid_quantization_values)}, "
77 + f"got quantization={quantization!r}"
78 + )
79 + return None
80 +
81 + if quantization in online_quant_values:
82 + scheme = OnlineQuantScheme(quantization)
83 +
84 + if quantization_config is None:
85 + quantization_config = {
86 + "global_scheme": scheme.value,
87 + }
88 + elif isinstance(quantization_config, OnlineQuantizationConfigArgs):
89 + if quantization_config.global_scheme is None:
90 + quantization_config.global_scheme = scheme
91 + elif quantization_config.global_scheme != scheme:
92 + raise ValueError(
93 + f"quantization={quantization!r} conflicts with "
94 + f"quantization_config.global_scheme="
95 + f"{quantization_config.global_scheme.value!r}. "
96 + f"These must match when both are specified."
97 + )
98 + elif isinstance(quantization_config, dict):
99 + existing = quantization_config.get("global_scheme")
100 + if existing is None:
101 + quantization_config["global_scheme"] = scheme.value
102 + else:
103 + # Coerce to enum for comparison
104 + existing_scheme = (
105 + OnlineQuantScheme(existing)
106 + if isinstance(existing, str)
107 + else existing
108 + )
109 + if existing_scheme != scheme:
110 + raise ValueError(
111 + f"quantization={quantization!r} conflicts "
112 + f"with quantization_config"
113 + f"['global_scheme']={existing!r}. "
114 + f"These must match when both are specified."
115 + )
116 +
117 + # Cast dict to OnlineQuantizationConfigArgs
118 + if isinstance(quantization_config, dict):
119 + quantization_config = OnlineQuantizationConfigArgs(**quantization_config)
120 +
121 + return quantization_config
vllm/config/vllm.py
+1
@@ -1713,6 +1713,7 @@ class VllmConfig:
1713 f"dcp_comm_backend={self.parallel_config.dcp_comm_backend}, " # noqa
1714 f"disable_custom_all_reduce={self.parallel_config.disable_custom_all_reduce}, " # noqa
1715 f"quantization={self.model_config.quantization}, "
1716 + f"quantization_config={self.model_config.quantization_config}, " # noqa
1717 f"enforce_eager={self.model_config.enforce_eager}, "
1718 f"enable_return_routed_experts={self.model_config.enable_return_routed_experts}, " # noqa
1719 f"kv_cache_dtype={self.cache_config.cache_dtype}, "
vllm/engine/arg_utils.py
+9
@@ -112,6 +112,7 @@ from vllm.v1.sample.logits_processor import LogitsProcessor
112 from vllm.version import __version__ as VLLM_VERSION
113
114 if TYPE_CHECKING:
115 + from vllm.config.quantization import OnlineQuantizationConfigArgs
116 from vllm.model_executor.layers.quantization import QuantizationMethods
117 from vllm.model_executor.model_loader import LoadFormats
118 from vllm.usage.usage_lib import UsageContext
@@ -483,6 +484,7 @@ class EngineArgs:
484 hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides")
485 tokenizer_revision: str | None = ModelConfig.tokenizer_revision
486 quantization: QuantizationMethods | str | None = ModelConfig.quantization
487 + quantization_config: "dict[str, Any] | OnlineQuantizationConfigArgs | None" = None
488 allow_deprecated_quantization: bool = ModelConfig.allow_deprecated_quantization
489 enforce_eager: bool = ModelConfig.enforce_eager
490 disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce
@@ -661,6 +663,12 @@ class EngineArgs:
663 if isinstance(self.ir_op_priority, dict):
664 self.ir_op_priority = IrOpPriorityConfig(**self.ir_op_priority)
665
666 + from vllm.config.quantization import resolve_online_quant_config
667 +
668 + self.quantization_config = resolve_online_quant_config(
669 + self.quantization, self.quantization_config
670 + )
671 +
672 # Setup plugins
673 from vllm.plugins import load_general_plugins
674
@@ -1431,6 +1439,7 @@ class EngineArgs:
1439 tokenizer_revision=self.tokenizer_revision,
1440 max_model_len=self.max_model_len,
1441 quantization=self.quantization,
1442 + quantization_config=self.quantization_config,
1443 allow_deprecated_quantization=self.allow_deprecated_quantization,
1444 enforce_eager=self.enforce_eager,
1445 enable_return_routed_experts=self.enable_return_routed_experts,
vllm/entrypoints/llm.py
+7
@@ -34,6 +34,9 @@ from vllm.config.model import (
34 RunnerOption,
35 TokenizerMode,
36 )
37 +from vllm.config.quantization import (
38 + OnlineQuantizationConfigArgs,
39 +)
40 from vllm.distributed.weight_transfer.base import (
41 WeightTransferInitRequest,
42 WeightTransferUpdateRequest,
@@ -247,6 +250,9 @@ class LLM:
250 attention_config: dict[str, Any] | AttentionConfig | None = None,
251 kv_cache_memory_bytes: int | None = None,
252 compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
253 + quantization_config: dict[str, Any]
254 + | OnlineQuantizationConfigArgs
255 + | None = None,
256 logits_processors: list[str | type[LogitsProcessor]] | None = None,
257 **kwargs: Any,
258 ) -> None:
@@ -367,6 +373,7 @@ class LLM:
373 profiler_config=profiler_config_instance,
374 attention_config=attention_config_instance,
375 compilation_config=compilation_config_instance,
376 + quantization_config=quantization_config,
377 logits_processors=logits_processors,
378 **kwargs,
379 )
vllm/model_executor/layers/quantization/__init__.py
+22
@@ -33,6 +33,13 @@ QuantizationMethods = Literal[
33 "mxfp8",
34 "petit_nvfp4",
35 "cpu_awq",
36 + "online",
37 + # Below are values of the OnlineQuantScheme enum, specified as strings to
38 + # avoid circular import issues. This is here to provide a shortcut where
39 + # the user can specify "LLM(..., quantization='fp8_per_tensor')" as
40 + # shorthand for creating a more complicated online quant config object
41 + "fp8_per_tensor",
42 + "fp8_per_block",
43 ]
44 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
45
@@ -103,6 +110,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
110 raise ValueError(f"Invalid quantization method: {quantization}")
111
112 # lazy import to avoid triggering `torch.compile` too early
113 + from vllm.config.quantization import OnlineQuantScheme
114 from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
115
116 from .awq import AWQConfig
@@ -129,6 +137,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
137 from .moe_wna16 import MoeWNA16Config
138 from .mxfp4 import Mxfp4Config
139 from .mxfp8 import Mxfp8Config
140 + from .online.base import OnlineQuantizationConfig
141 from .petit import PetitNvFp4Config
142 from .torchao import TorchAOConfig
143
@@ -157,7 +166,20 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
166 "mxfp8": Mxfp8Config,
167 "petit_nvfp4": PetitNvFp4Config,
168 "cpu_awq": CPUAWQConfig,
169 + "online": OnlineQuantizationConfig,
170 }
171 +
172 + # Below are values of the OnlineQuantScheme enum. This is here to provide
173 + # a shortcut where the user can specify
174 + # "LLM(..., quantization='fp8_per_tensor')" as shorthand for creating a
175 + # more complicated online quant config object
176 + for scheme in OnlineQuantScheme:
177 + assert scheme.value not in method_to_config, (
178 + f"Online quant scheme {scheme.value!r} conflicts with an "
179 + f"existing quantization method"
180 + )
181 + method_to_config[scheme.value] = OnlineQuantizationConfig
182 +
183 # Update the `method_to_config` with customized quantization methods.
184 method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
185
vllm/model_executor/layers/quantization/fp8.py
+4
@@ -497,6 +497,8 @@ class Fp8LinearMethod(LinearMethodBase):
497 return self.fp8_linear.apply_weights(layer, x, bias)
498
499
500 +# TODO(future PR): remove this class in favor of
501 +# online/fp8.py::Fp8PerTensorOnlineLinearMethod
502 class Fp8OnlineLinearMethod(Fp8LinearMethod):
503 """Online version of Fp8LinearMethod which loads a full precision checkpoint
504 and quantizes weights during loading."""
@@ -919,6 +921,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
921 )
922
923
924 +# TODO(future PR): remove this class in favor of
925 +# online/fp8.py::Fp8PerTensorOnlineMoEMethod
926 class Fp8OnlineMoEMethod(Fp8MoEMethod):
927 """MoE method for online FP8 quantization.
928 Supports loading quantized FP16/BF16 model checkpoints with dynamic
vllm/model_executor/layers/quantization/online/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/layers/quantization/online/base.py new
+116
@@ -0,0 +1,116 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any
5 +
6 +import torch
7 +
8 +from vllm.config.quantization import (
9 + OnlineQuantizationConfigArgs,
10 + OnlineQuantScheme,
11 +)
12 +from vllm.model_executor.layers.fused_moe import (
13 + FusedMoE,
14 +)
15 +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
16 + UnquantizedFusedMoEMethod,
17 +)
18 +from vllm.model_executor.layers.linear import (
19 + LinearBase,
20 + UnquantizedLinearMethod,
21 +)
22 +from vllm.model_executor.layers.quantization import QuantizationMethods
23 +from vllm.model_executor.layers.quantization.base_config import (
24 + QuantizationConfig,
25 + QuantizeMethodBase,
26 +)
27 +from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
28 + should_ignore_layer,
29 +)
30 +from vllm.model_executor.layers.quantization.online.fp8 import (
31 + Fp8PerBlockOnlineLinearMethod,
32 + Fp8PerBlockOnlineMoEMethod,
33 + Fp8PerTensorOnlineLinearMethod,
34 + Fp8PerTensorOnlineMoEMethod,
35 +)
36 +
37 +
38 +class OnlineQuantizationConfig(QuantizationConfig):
39 + """Model-level config class for online quantization (quantize fp16/bf16 weights
40 + during model loading, without requiring a pre-quantized checkpoint)."""
41 +
42 + def __init__(
43 + self,
44 + args: OnlineQuantizationConfigArgs,
45 + ) -> None:
46 + super().__init__()
47 + if (
48 + args.global_scheme is None
49 + and args.linear_scheme_override is None
50 + and args.moe_scheme_override is None
51 + ):
52 + raise ValueError(
53 + "OnlineQuantizationConfig requires at least one of "
54 + "global_scheme, linear_scheme_override, or "
55 + "moe_scheme_override to be set."
56 + )
57 + self.args = args
58 + self.quant_scheme = args.global_scheme
59 + self.ignored_layers: list[str] = args.ignore
60 +
61 + @classmethod
62 + def get_name(cls) -> QuantizationMethods:
63 + return "online"
64 +
65 + @classmethod
66 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:
67 + return [torch.bfloat16, torch.half]
68 +
69 + @classmethod
70 + def get_min_capability(cls) -> int:
71 + # Note: as more online quant schemes will be added, this
72 + # value will become the minimum across all supported schemes.
73 + return 75
74 +
75 + @classmethod
76 + def get_config_filenames(cls) -> list[str]:
77 + return []
78 +
79 + @classmethod
80 + def from_config(cls, config: dict[str, Any]) -> "OnlineQuantizationConfig":
81 + raise NotImplementedError(
82 + "OnlineQuantizationConfig does not support loading from a "
83 + "checkpoint config. Use quantization_config or "
84 + "quantization='fp8_per_tensor'/'fp8_per_block' instead."
85 + )
86 +
87 + def get_quant_method(
88 + self, layer: torch.nn.Module, prefix: str
89 + ) -> "QuantizeMethodBase | None":
90 + if isinstance(layer, LinearBase):
91 + if should_ignore_layer(
92 + prefix,
93 + ignore=self.ignored_layers,
94 + fused_mapping=self.packed_modules_mapping,
95 + ):
96 + return UnquantizedLinearMethod()
97 +
98 + linear_scheme = self.args.linear_scheme_override or self.args.global_scheme
99 + if linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
100 + return Fp8PerBlockOnlineLinearMethod()
101 + else:
102 + return Fp8PerTensorOnlineLinearMethod()
103 + elif isinstance(layer, FusedMoE):
104 + if should_ignore_layer(
105 + prefix,
106 + ignore=self.ignored_layers,
107 + fused_mapping=self.packed_modules_mapping,
108 + ):
109 + return UnquantizedFusedMoEMethod(layer.moe_config)
110 +
111 + moe_scheme = self.args.moe_scheme_override or self.args.global_scheme
112 + if moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
113 + return Fp8PerBlockOnlineMoEMethod(layer=layer)
114 + else:
115 + return Fp8PerTensorOnlineMoEMethod(layer=layer)
116 + return None
vllm/model_executor/layers/quantization/online/fp8.py new
+632
@@ -0,0 +1,632 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import TYPE_CHECKING
5 +
6 +import torch
7 +from torch.nn import Module
8 +
9 +if TYPE_CHECKING:
10 + import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 + from vllm.model_executor.layers.fused_moe import FusedMoE
12 + from vllm.model_executor.layers.fused_moe.config import (
13 + FusedMoEConfig,
14 + FusedMoEQuantConfig,
15 + )
16 + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
17 +
18 +import vllm.envs as envs
19 +from vllm import _custom_ops as ops
20 +from vllm._aiter_ops import rocm_aiter_ops
21 +from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
22 +from vllm.model_executor.layers.fused_moe import (
23 + FusedMoEMethodBase,
24 +)
25 +from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
26 + select_fp8_moe_backend,
27 +)
28 +from vllm.model_executor.layers.linear import (
29 + LinearMethodBase,
30 +)
31 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
32 + W8A8BlockFp8LinearOp,
33 + maybe_post_process_fp8_weight_block,
34 + process_fp8_weight_block_strategy,
35 +)
36 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
37 + GroupShape,
38 + kFp8Dynamic128Sym,
39 + kFp8DynamicTensorSym,
40 + kFp8DynamicTokenSym,
41 + kFp8Static128BlockSym,
42 + kFp8StaticTensorSym,
43 +)
44 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
45 + cutlass_block_fp8_supported,
46 + cutlass_fp8_supported,
47 +)
48 +from vllm.model_executor.model_loader.reload.layerwise import (
49 + initialize_online_processing,
50 +)
51 +from vllm.model_executor.parameter import ModelWeightParameter
52 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs
53 +from vllm.platforms import current_platform
54 +from vllm.utils.deep_gemm import is_deep_gemm_supported, per_block_cast_to_fp8
55 +
56 +# ---------------------------------------------------------------------------
57 +# Online FP8 Linear Methods
58 +# ---------------------------------------------------------------------------
59 +
60 +
61 +class _Fp8OnlineLinearBase(LinearMethodBase):
62 + """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
63 + weights onto meta device and materializes them just-in-time."""
64 +
65 + uses_meta_device: bool = True
66 +
67 + def create_weights(
68 + self,
69 + layer: torch.nn.Module,
70 + input_size_per_partition: int,
71 + output_partition_sizes: list[int],
72 + input_size: int,
73 + output_size: int,
74 + params_dtype: torch.dtype,
75 + **extra_weight_attrs,
76 + ):
77 + output_size_per_partition = sum(output_partition_sizes)
78 + weight_loader = extra_weight_attrs.get("weight_loader")
79 + layer.logical_widths = output_partition_sizes
80 + layer.input_size_per_partition = input_size_per_partition
81 + layer.output_size_per_partition = output_size_per_partition
82 + layer.orig_dtype = params_dtype
83 + layer.weight_block_size = None
84 +
85 + weight = ModelWeightParameter(
86 + data=torch.empty(
87 + output_size_per_partition,
88 + input_size_per_partition,
89 + device="meta", # materialized and processed during loading
90 + dtype=params_dtype,
91 + ),
92 + input_dim=1,
93 + output_dim=0,
94 + weight_loader=weight_loader,
95 + )
96 + layer.register_parameter("weight", weight)
97 +
98 + initialize_online_processing(layer)
99 +
100 +
101 +class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
102 + """Online tensorwise FP8 linear quantization.
103 + Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
104 +
105 + def __init__(self):
106 + self.out_dtype = torch.get_default_dtype()
107 +
108 + # Use per-token quantization for better perf if dynamic and cutlass
109 + if cutlass_fp8_supported():
110 + activation_quant_key = kFp8DynamicTokenSym
111 + else:
112 + activation_quant_key = kFp8DynamicTensorSym
113 +
114 + self.fp8_linear = init_fp8_linear_kernel(
115 + activation_quant_key=activation_quant_key,
116 + weight_quant_key=kFp8StaticTensorSym,
117 + out_dtype=torch.get_default_dtype(),
118 + module_name=self.__class__.__name__,
119 + )
120 +
121 + def process_weights_after_loading(self, layer: Module) -> None:
122 + if getattr(layer, "_already_called_process_weights_after_loading", False):
123 + return
124 +
125 + layer.input_scale = None
126 + qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
127 +
128 + # Update layer with new values.
129 + replace_parameter(layer, "weight", qweight.t().data)
130 + replace_parameter(layer, "weight_scale", weight_scale.data)
131 +
132 + # Prevent duplicate processing (e.g., during weight reload)
133 + layer._already_called_process_weights_after_loading = True
134 +
135 + def apply(
136 + self,
137 + layer: torch.nn.Module,
138 + x: torch.Tensor,
139 + bias: torch.Tensor | None = None,
140 + ) -> torch.Tensor:
141 + # if batch invariant mode is enabled, use BF16 dequant
142 + if envs.VLLM_BATCH_INVARIANT:
143 + weight_fp8 = layer.weight.to(torch.bfloat16)
144 + weight_scale = layer.weight_scale.to(torch.bfloat16)
145 + if weight_scale.numel() == 1:
146 + # Per-tensor: simple scalar multiplication
147 + weight_bf16 = weight_fp8 * weight_scale
148 + else:
149 + # Multiple scales (fused modules like QKV)
150 + if (
151 + weight_scale.dim() == 1
152 + and weight_scale.shape[0] == weight_fp8.shape[0]
153 + ):
154 + # Per-row scaling
155 + weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1)
156 + else:
157 + # Fallback
158 + weight_bf16 = weight_fp8 * weight_scale
159 + return torch.nn.functional.linear(x, weight_bf16.t(), bias)
160 +
161 + return self.fp8_linear.apply_weights(layer, x, bias)
162 +
163 +
164 +class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
165 + """Online blockwise FP8 linear quantization.
166 + Loads fp16/bf16 weights and quantizes them per-block during loading."""
167 +
168 + def __init__(self):
169 + self.out_dtype = torch.get_default_dtype()
170 + self.weight_block_size = [128, 128]
171 +
172 + self.use_deep_gemm = is_deep_gemm_supported()
173 + self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
174 + self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
175 +
176 + self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
177 + weight_group_shape=GroupShape(*self.weight_block_size),
178 + act_quant_group_shape=GroupShape(1, self.weight_block_size[0]),
179 + cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
180 + use_aiter_and_is_supported=self.use_aiter_and_is_supported,
181 + use_deep_gemm=self.use_deep_gemm,
182 + )
183 +
184 + def create_weights(
185 + self,
186 + layer: torch.nn.Module,
187 + input_size_per_partition: int,
188 + output_partition_sizes: list[int],
189 + input_size: int,
190 + output_size: int,
191 + params_dtype: torch.dtype,
192 + **extra_weight_attrs,
193 + ):
194 + super().create_weights(
195 + layer,
196 + input_size_per_partition,
197 + output_partition_sizes,
198 + input_size,
199 + output_size,
200 + params_dtype,
201 + **extra_weight_attrs,
202 + )
203 + layer.weight_block_size = self.weight_block_size
204 +
205 + def process_weights_after_loading(self, layer: Module) -> None:
206 + if getattr(layer, "_already_called_process_weights_after_loading", False):
207 + return
208 +
209 + layer.input_scale = None
210 + block_size = self.weight_block_size
211 +
212 + qweight, weight_scale_inv = per_block_cast_to_fp8(
213 + layer.weight, block_size=block_size, use_ue8m0=False
214 + )
215 +
216 + qweight, weight_scale_inv = process_fp8_weight_block_strategy(
217 + qweight, weight_scale_inv
218 + )
219 +
220 + replace_parameter(layer, "weight", qweight.data)
221 + replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
222 +
223 + maybe_post_process_fp8_weight_block(layer)
224 +
225 + # Prevent duplicate processing (e.g., during weight reload)
226 + layer._already_called_process_weights_after_loading = True
227 +
228 + def apply(
229 + self,
230 + layer: torch.nn.Module,
231 + x: torch.Tensor,
232 + bias: torch.Tensor | None = None,
233 + ) -> torch.Tensor:
234 + assert self.weight_block_size is not None
235 +
236 + # Note: batch invariance already handled in the function below
237 + return self.w8a8_block_fp8_linear.apply(
238 + input=x,
239 + weight=layer.weight,
240 + weight_scale=layer.weight_scale_inv,
241 + input_scale=layer.input_scale,
242 + bias=bias,
243 + )
244 +
245 +
246 +# ---------------------------------------------------------------------------
247 +# Online FP8 MoE Methods
248 +# ---------------------------------------------------------------------------
249 +
250 +
251 +class _Fp8OnlineMoEBase(FusedMoEMethodBase):
252 + """Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint
253 + weights onto meta device and materializes them just-in-time."""
254 +
255 + uses_meta_device: bool = True
256 +
257 + # Declared here for mypy; actual values are set in __init__.
258 + fp8_backend: "Fp8MoeBackend"
259 + experts_cls: "type[mk.FusedMoEExperts] | None"
260 + weight_scale_name: str
261 + weight_block_size: list[int] | None
262 + moe: "FusedMoEConfig"
263 + is_monolithic: bool
264 + moe_quant_config: "FusedMoEQuantConfig | None"
265 + moe_kernel: "mk.FusedMoEKernel | None"
266 +
267 + def __init__(
268 + self,
269 + *,
270 + weight_block_size: list[int] | None,
271 + layer: torch.nn.Module,
272 + ):
273 + super().__init__(layer.moe_config)
274 + self.weight_block_size = weight_block_size
275 + self.block_quant: bool = self.weight_block_size is not None
276 + self.weight_scale_name = (
277 + "weight_scale_inv" if self.block_quant else "weight_scale"
278 + )
279 +
280 + # Set weight key and activation key for kernel compatibility
281 + if self.block_quant:
282 + weight_key = kFp8Static128BlockSym
283 + activation_key = kFp8Dynamic128Sym
284 + else:
285 + weight_key = kFp8StaticTensorSym
286 + activation_key = kFp8DynamicTensorSym
287 +
288 + # Select Fp8 MoE backend
289 + self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
290 + config=self.moe,
291 + weight_key=weight_key,
292 + activation_key=activation_key,
293 + allow_vllm_cutlass=False,
294 + )
295 +
296 + def create_weights(
297 + self,
298 + layer: Module,
299 + num_experts: int,
300 + hidden_size: int,
301 + intermediate_size_per_partition: int,
302 + params_dtype: torch.dtype,
303 + **extra_weight_attrs,
304 + ):
305 + layer.num_experts = num_experts
306 + layer.orig_dtype = params_dtype
307 + layer.weight_block_size = None
308 +
309 + # WEIGHTS
310 + w13_weight = torch.nn.Parameter(
311 + torch.empty(
312 + num_experts,
313 + 2 * intermediate_size_per_partition,
314 + hidden_size,
315 + device="meta",
316 + dtype=params_dtype,
317 + ),
318 + requires_grad=False,
319 + )
320 + layer.register_parameter("w13_weight", w13_weight)
321 + set_weight_attrs(w13_weight, extra_weight_attrs)
322 +
323 + w2_weight = torch.nn.Parameter(
324 + torch.empty(
325 + num_experts,
326 + hidden_size,
327 + intermediate_size_per_partition,
328 + device="meta", # materialized and processed during loading
329 + dtype=params_dtype,
330 + ),
331 + requires_grad=False,
332 + )
333 + layer.register_parameter("w2_weight", w2_weight)
334 + set_weight_attrs(w2_weight, extra_weight_attrs)
335 +
336 + # BIASES (for models like GPT-OSS that have biased MoE)
337 + if self.moe.has_bias:
338 + w13_bias = torch.nn.Parameter(
339 + torch.zeros(
340 + num_experts,
341 + 2 * intermediate_size_per_partition,
342 + device="meta", # materialized and processed during loading
343 + dtype=layer.orig_dtype,
344 + ),
345 + requires_grad=False,
346 + )
347 + layer.register_parameter("w13_bias", w13_bias)
348 + set_weight_attrs(w13_bias, extra_weight_attrs)
349 +
350 + w2_bias = torch.nn.Parameter(
351 + torch.zeros(
352 + num_experts,
353 + hidden_size,
354 + device="meta", # materialized and processed during loading
355 + dtype=layer.orig_dtype,
356 + ),
357 + requires_grad=False,
358 + )
359 + layer.register_parameter("w2_bias", w2_bias)
360 + set_weight_attrs(w2_bias, extra_weight_attrs)
361 +
362 + layer.w13_input_scale = None
363 + layer.w2_input_scale = None
364 +
365 + initialize_online_processing(layer)
366 +
367 + def _setup_kernel(
368 + self,
369 + layer: "FusedMoE",
370 + w13: torch.Tensor,
371 + w2: torch.Tensor,
372 + w13_scale: torch.Tensor,
373 + w2_scale: torch.Tensor,
374 + w13_input_scale: torch.Tensor | None,
375 + w2_input_scale: torch.Tensor | None,
376 + ) -> None:
377 + from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
378 + convert_to_fp8_moe_kernel_format,
379 + make_fp8_moe_kernel,
380 + )
381 +
382 + # Shuffle weights to runtime format.
383 + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
384 + fp8_backend=self.fp8_backend,
385 + layer=layer,
386 + w13=w13,
387 + w2=w2,
388 + w13_scale=w13_scale,
389 + w2_scale=w2_scale,
390 + w13_input_scale=w13_input_scale,
391 + w2_input_scale=w2_input_scale,
392 + )
393 +
394 + # Replace parameters with updated versions. Note that this helper
395 + # function ensures the replacement is compatible with RL weight reloads.
396 + replace_parameter(layer, "w13_weight", w13)
397 + replace_parameter(layer, "w2_weight", w2)
398 + replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
399 + replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
400 +
401 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
402 + if self.moe_quant_config:
403 + assert self.experts_cls is not None
404 + self.moe_kernel = make_fp8_moe_kernel(
405 + moe_quant_config=self.moe_quant_config,
406 + moe_config=self.moe,
407 + fp8_backend=self.fp8_backend,
408 + experts_cls=self.experts_cls,
409 + routing_tables=layer._maybe_init_expert_routing_tables(),
410 + shared_experts=layer.shared_experts,
411 + )
412 +
413 + def maybe_make_prepare_finalize(
414 + self,
415 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
416 + ) -> "mk.FusedMoEPrepareAndFinalizeModular | None":
417 + raise ValueError(
418 + f"{self.__class__.__name__} uses the new modular kernel "
419 + "initialization logic. This function should not be called."
420 + )
421 +
422 + def get_fused_moe_quant_config(
423 + self, layer: torch.nn.Module
424 + ) -> "FusedMoEQuantConfig":
425 + from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
426 + make_fp8_moe_quant_config,
427 + )
428 +
429 + w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
430 + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
431 + a1_scale = layer.w13_input_scale
432 + a2_scale = layer.w2_input_scale
433 +
434 + quant_config = make_fp8_moe_quant_config(
435 + fp8_backend=self.fp8_backend,
436 + w1_scale=w1_scale,
437 + w2_scale=w2_scale,
438 + a1_scale=a1_scale,
439 + a2_scale=a2_scale,
440 + block_shape=self.weight_block_size,
441 + )
442 +
443 + # Inject biases into the quant config if the model has them
444 + # (e.g. GPT-OSS biased MoE)
445 + if quant_config is not None and self.moe.has_bias:
446 + w13_bias = getattr(layer, "w13_bias", None)
447 + w2_bias = getattr(layer, "w2_bias", None)
448 + if w13_bias is not None:
449 + quant_config._w1.bias = w13_bias
450 + if w2_bias is not None:
451 + quant_config._w2.bias = w2_bias
452 +
453 + return quant_config
454 +
455 + @property
456 + def supports_eplb(self) -> bool:
457 + return True
458 +
459 + def apply_monolithic(
460 + self,
461 + layer: "FusedMoE",
462 + x: torch.Tensor,
463 + router_logits: torch.Tensor,
464 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
465 + assert self.is_monolithic
466 + assert self.moe_kernel is not None
467 + return self.moe_kernel.apply_monolithic(
468 + x,
469 + layer.w13_weight,
470 + layer.w2_weight,
471 + router_logits,
472 + activation=layer.activation,
473 + global_num_experts=layer.global_num_experts,
474 + expert_map=layer.expert_map,
475 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
476 + num_expert_group=layer.num_expert_group,
477 + topk_group=layer.topk_group,
478 + e_score_correction_bias=layer.e_score_correction_bias,
479 + routed_scaling_factor=layer.routed_scaling_factor,
480 + )
481 +
482 + def apply(
483 + self,
484 + layer: "FusedMoE",
485 + x: torch.Tensor,
486 + topk_weights: torch.Tensor,
487 + topk_ids: torch.Tensor,
488 + shared_experts_input: torch.Tensor | None,
489 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
490 + assert not self.is_monolithic
491 + assert self.moe_kernel is not None
492 + return self.moe_kernel.apply(
493 + x,
494 + layer.w13_weight,
495 + layer.w2_weight,
496 + topk_weights,
497 + topk_ids,
498 + activation=layer.activation,
499 + global_num_experts=layer.global_num_experts,
500 + expert_map=layer.expert_map,
501 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
502 + shared_experts_input=shared_experts_input,
503 + )
504 +
505 +
506 +class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
507 + """Online tensorwise FP8 MoE quantization.
508 + Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
509 +
510 + def __init__(
511 + self,
512 + *,
513 + layer: torch.nn.Module,
514 + ):
515 + super().__init__(
516 + weight_block_size=None,
517 + layer=layer,
518 + )
519 +
520 + def process_weights_after_loading(self, layer: Module) -> None:
521 + # TODO(@ksayers): inplace fp8 quant kernel, initialize scales with ones
522 + if getattr(layer, "_already_called_process_weights_after_loading", False):
523 + return
524 +
525 + # If checkpoint is fp16, quantize in place.
526 + fp8_dtype = current_platform.fp8_dtype()
527 + w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
528 + w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
529 + w13_scale = torch.ones(
530 + layer.num_experts, device=w13.device, dtype=torch.float32
531 + )
532 + w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32)
533 + layer.w13_input_scale = None
534 + layer.w2_input_scale = None
535 +
536 + for expert in range(layer.local_num_experts):
537 + w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
538 + layer.w13_weight[expert, :, :]
539 + )
540 + w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
541 + layer.w2_weight[expert, :, :]
542 + )
543 +
544 + # Shuffle weights to runtime format and setup kernel.
545 + self._setup_kernel(
546 + layer,
547 + w13,
548 + w2,
549 + w13_scale,
550 + w2_scale,
551 + w13_input_scale=layer.w13_input_scale,
552 + w2_input_scale=layer.w2_input_scale,
553 + )
554 +
555 + # Prevent duplicate processing (e.g., during weight reload)
556 + layer._already_called_process_weights_after_loading = True
557 +
558 +
559 +class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
560 + """Online blockwise FP8 MoE quantization.
561 + Loads fp16/bf16 weights and quantizes them per-block during loading."""
562 +
563 + def __init__(
564 + self,
565 + *,
566 + layer: torch.nn.Module,
567 + ):
568 + super().__init__(
569 + weight_block_size=[128, 128],
570 + layer=layer,
571 + )
572 +
573 + def process_weights_after_loading(self, layer: Module) -> None:
574 + if getattr(layer, "_already_called_process_weights_after_loading", False):
575 + return
576 +
577 + fp8_dtype = current_platform.fp8_dtype()
578 + w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
579 + w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
580 +
581 + block_size = self.weight_block_size
582 + assert block_size is not None
583 + block_n, block_k = block_size
584 +
585 + # Create block-shaped scales (computed here rather than in
586 + # create_weights because online quant doesn't need them until now).
587 + num_experts = layer.local_num_experts
588 + _, w13_out, w13_in = layer.w13_weight.shape
589 + _, w2_out, w2_in = layer.w2_weight.shape
590 +
591 + w13_scale = torch.ones(
592 + num_experts,
593 + (w13_out + block_n - 1) // block_n,
594 + (w13_in + block_k - 1) // block_k,
595 + dtype=torch.float32,
596 + device=w13.device,
597 + )
598 + w2_scale = torch.ones(
599 + num_experts,
600 + (w2_out + block_n - 1) // block_n,
601 + (w2_in + block_k - 1) // block_k,
602 + dtype=torch.float32,
603 + device=w2.device,
604 + )
605 +
606 + for expert in range(num_experts):
607 + w13[expert], w13_scale[expert] = per_block_cast_to_fp8(
608 + layer.w13_weight[expert],
609 + block_size=block_size,
610 + use_ue8m0=False,
611 + )
612 + w2[expert], w2_scale[expert] = per_block_cast_to_fp8(
613 + layer.w2_weight[expert],
614 + block_size=block_size,
615 + use_ue8m0=False,
616 + )
617 +
618 + layer.weight_block_size = block_size
619 +
620 + # Shuffle weights to runtime format and setup kernel.
621 + self._setup_kernel(
622 + layer,
623 + w13,
624 + w2,
625 + w13_scale,
626 + w2_scale,
627 + layer.w13_input_scale,
628 + layer.w2_input_scale,
629 + )
630 +
631 + # Prevent duplicate processing (e.g., during weight reload)
632 + layer._already_called_process_weights_after_loading = True
vllm/model_executor/model_loader/weight_utils.py
+32
@@ -296,6 +296,13 @@ def get_quant_config(
296 )
297
298 if hf_quant_config is not None:
299 + if model_config.quantization_config is not None:
300 + raise ValueError(
301 + "Setting `quantization_config` for online "
302 + "quantization when the model checkpoint already "
303 + "has a `quantization_config` is not supported"
304 + )
305 +
306 # For modelopt_mixed, config.json's quantization_config may or may
307 # not contain the per-layer quantized_layers map. Newer checkpoints
308 # embed it directly; older ones keep it only in hf_quant_config.json.
@@ -319,6 +326,12 @@ def get_quant_config(
326 quantization_config_file = hf_overrides.get("quantization_config_file", None)
327 if quantization_config_file is not None:
328 if hasattr(quant_cls, "from_config_file"):
329 + if model_config.quantization_config is not None:
330 + raise ValueError(
331 + "Setting `quantization_config` for online "
332 + "quantization when the model checkpoint already "
333 + "has a `quantization_config` is not supported"
334 + )
335 return quant_cls.from_config_file(quantization_config_file)
336 else:
337 raise NotImplementedError(
@@ -329,6 +342,12 @@ def get_quant_config(
342 quantization_config_json = hf_overrides.get("quantization_config_dict_json", None)
343 if quantization_config_json is not None:
344 if hasattr(quant_cls, "from_config_dict_json"):
345 + if model_config.quantization_config is not None:
346 + raise ValueError(
347 + "Setting `quantization_config` for online "
348 + "quantization when the model checkpoint already "
349 + "has a `quantization_config` is not supported"
350 + )
351 return quant_cls.from_config_dict_json(quantization_config_json)
352 else:
353 raise NotImplementedError(
@@ -337,6 +356,19 @@ def get_quant_config(
356 f"{quant_cls}"
357 )
358
359 + # Online quantization doesn't read from checkpoint configs — it quantizes
360 + # fp16/bf16 weights on the fly during loading.
361 + if model_config.quantization_config is not None:
362 + from vllm.config.quantization import OnlineQuantizationConfigArgs
363 + from vllm.model_executor.layers.quantization.online.base import (
364 + OnlineQuantizationConfig,
365 + )
366 +
367 + assert isinstance(
368 + model_config.quantization_config, OnlineQuantizationConfigArgs
369 + )
370 + return OnlineQuantizationConfig(args=model_config.quantization_config)
371 +
372 # Inflight BNB quantization
373 if model_config.quantization == "bitsandbytes":
374 return quant_cls.from_config({})