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 4, 22:30 UTC
mxfp8 online quant move to new frontend (#40152)
Signed-off-by: Vasiliy Kuznetsov <vasiliy@meta.com>
Vasiliy Kuznetsov committed
Apr 20, 2026 at 09:26 UTC
38fa87cacadc73aec9b28fea52fa70a31070cec4
7 files changed
+108
-126
docs/features/quantization/online.md
+5
-2
@@ -17,6 +17,9 @@ llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor")
17
18
# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations)
19
llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
20
+
21
+# MXFP8 quantization for weights and activations
22
+llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp8")
23
```
24
25
Or with the CLI:
@@ -24,6 +27,7 @@ Or with the CLI:
27
```bash
28
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
29
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
30
+vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
31
```
32
33
## Supported Schemes
@@ -32,8 +36,7 @@ vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
36
| ------ | ------------- | ------------------ | ----- |
37
| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
38
| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
35
-
36
-Support for additional schemes will be added in future versions of vllm.
39
+| `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback |
40
41
## Advanced Configuration
42
vllm/config/quantization.py
+2
-1
@@ -23,7 +23,8 @@ class OnlineQuantScheme(Enum):
23
# Linear layers remain unquantized.
24
INT8_PER_CHANNEL_WEIGHT_ONLY = "int8_per_channel_weight_only"
25
26
- # TODO(future PRs): add more online quant schemes here: mxfp8, etc
26
+ # mxfp8, weights scaled in blocks of 1x32 elements (microscaling FP8)
27
+ MXFP8 = "mxfp8"
28
29
30
@config
vllm/model_executor/layers/quantization/__init__.py
+1
-3
@@ -31,7 +31,6 @@ QuantizationMethods = Literal[
31
"inc",
32
"mxfp4",
33
"gpt_oss_mxfp4",
34
- "mxfp8",
34
"cpu_awq",
35
"online",
36
# Below are values of the OnlineQuantScheme enum, specified as strings to
@@ -41,6 +40,7 @@ QuantizationMethods = Literal[
40
"fp8_per_tensor",
41
"fp8_per_block",
42
"int8_per_channel_weight_only",
43
+ "mxfp8",
44
]
45
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
46
@@ -135,7 +135,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
135
)
136
from .moe_wna16 import MoeWNA16Config
137
from .mxfp4 import GptOssMxfp4Config, Mxfp4Config
138
- from .mxfp8 import Mxfp8Config
138
from .online.base import OnlineQuantizationConfig
139
from .torchao import TorchAOConfig
140
@@ -162,7 +161,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
161
"inc": INCConfig,
162
"mxfp4": Mxfp4Config,
163
"gpt_oss_mxfp4": GptOssMxfp4Config,
165
- "mxfp8": Mxfp8Config,
164
"cpu_awq": CPUAWQConfig,
165
"online": OnlineQuantizationConfig,
166
}
vllm/model_executor/layers/quantization/fp8.py
-8
@@ -515,14 +515,6 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
515
516
initialize_online_processing(layer)
517
518
- # TODO: remove this check once the following RFC is resolved.
519
- # https://github.com/vllm-project/vllm/issues/33314
520
- # Subclasses (e.g. Mxfp8OnlineLinearMethod) only need the weight
521
- # registration above and manage their own kernel, so skip fp8_linear
522
- # kernel creation for them.
523
- if type(self) is not Fp8OnlineLinearMethod:
524
- return
525
-
518
self.fp8_linear = init_fp8_linear_kernel(
519
activation_quant_key=self.activation_quant_key,
520
weight_quant_key=self.weight_quant_key,
vllm/model_executor/layers/quantization/online/base.py
+8
@@ -37,6 +37,10 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
37
from vllm.model_executor.layers.quantization.online.int8 import (
38
Int8OnlineMoEMethod,
39
)
40
+from vllm.model_executor.layers.quantization.online.mxfp8 import (
41
+ Mxfp8OnlineLinearMethod,
42
+ Mxfp8OnlineMoEMethod,
43
+)
44
45
logger = init_logger(__name__)
46
@@ -110,6 +114,8 @@ class OnlineQuantizationConfig(QuantizationConfig):
114
return UnquantizedLinearMethod()
115
elif linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
116
return Fp8PerBlockOnlineLinearMethod()
117
+ elif linear_scheme == OnlineQuantScheme.MXFP8:
118
+ return Mxfp8OnlineLinearMethod()
119
else:
120
return Fp8PerTensorOnlineLinearMethod()
121
elif isinstance(layer, FusedMoE):
@@ -125,6 +131,8 @@ class OnlineQuantizationConfig(QuantizationConfig):
131
return Int8OnlineMoEMethod(layer=layer)
132
elif moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
133
return Fp8PerBlockOnlineMoEMethod(layer=layer)
134
+ elif moe_scheme == OnlineQuantScheme.MXFP8:
135
+ return Mxfp8OnlineMoEMethod(layer=layer)
136
else:
137
return Fp8PerTensorOnlineMoEMethod(layer=layer)
138
return None
vllm/model_executor/layers/quantization/online/mxfp8.py
renamed
+91
-111
@@ -1,131 +1,47 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-"""Online MXFP8 (microscaling FP8, block-32) quantization config and methods."""
4
+"""Online MXFP8 (microscaling FP8, block-32) quantization methods."""
5
6
-from typing import Any
6
+from typing import TYPE_CHECKING
7
8
import torch
9
from torch.nn import Module
10
11
-from vllm.logger import init_logger
11
+if TYPE_CHECKING:
12
+ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
+ from vllm.model_executor.layers.fused_moe import FusedMoE
14
+ from vllm.model_executor.layers.fused_moe.config import (
15
+ FusedMoEQuantConfig,
16
+ )
17
+ from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
18
+
19
from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel
13
-from vllm.model_executor.layers.attention import Attention
14
-from vllm.model_executor.layers.fused_moe import (
15
- FusedMoE,
16
- FusedMoEMethodBase,
17
-)
18
-from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
20
from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
21
select_mxfp8_moe_backend,
22
)
22
-from vllm.model_executor.layers.linear import (
23
- LinearBase,
24
- UnquantizedLinearMethod,
25
-)
26
-from vllm.model_executor.layers.quantization import QuantizationMethods
27
-from vllm.model_executor.layers.quantization.base_config import (
28
- QuantizeMethodBase,
23
+from vllm.model_executor.layers.quantization.online.fp8 import (
24
+ _Fp8OnlineLinearBase,
25
)
30
-from vllm.model_executor.layers.quantization.fp8 import (
31
- Fp8Config,
32
- Fp8KVCacheMethod,
33
- Fp8OnlineLinearMethod,
34
- Fp8OnlineMoEMethod,
26
+from vllm.model_executor.layers.quantization.online.moe_base import (
27
+ OnlineMoEMethodBase,
28
)
29
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
30
MXFP8_BLOCK_SIZE,
31
mxfp8_e4m3_quantize,
32
)
40
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
41
- is_layer_skipped,
42
-)
33
from vllm.model_executor.utils import replace_parameter
34
from vllm.platforms import current_platform
35
46
-logger = init_logger(__name__)
47
-
36
49
-class Mxfp8Config(Fp8Config):
50
- """Config class for online MXFP8 MoE quantization."""
51
-
52
- def __init__(
53
- self,
54
- activation_scheme: str = "dynamic",
55
- ignored_layers: list[str] | None = None,
56
- ) -> None:
57
- if activation_scheme != "dynamic":
58
- raise ValueError("mxfp8 only supports dynamic activation scheme.")
59
- super().__init__(
60
- is_checkpoint_fp8_serialized=False,
61
- activation_scheme=activation_scheme,
62
- ignored_layers=ignored_layers,
63
- weight_block_size=None,
64
- )
65
-
66
- @classmethod
67
- def get_name(cls) -> QuantizationMethods:
68
- return "mxfp8"
69
-
70
- @classmethod
71
- def get_min_capability(cls) -> int:
72
- # Marlin kernel supports MXFP8 on SM80+
73
- return 80
74
-
75
- @classmethod
76
- def from_config(cls, config: dict[str, Any]) -> "Mxfp8Config":
77
- activation_scheme = cls.get_from_keys_or(
78
- config, ["activation_scheme"], "dynamic"
79
- )
80
- ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)
81
- if not ignored_layers:
82
- ignored_layers = cls.get_from_keys_or(
83
- config, ["modules_to_not_convert"], None
84
- )
85
- return cls(
86
- activation_scheme=activation_scheme,
87
- ignored_layers=ignored_layers,
88
- )
89
-
90
- def get_quant_method(
91
- self, layer: torch.nn.Module, prefix: str
92
- ) -> "QuantizeMethodBase | None":
93
- if isinstance(layer, LinearBase):
94
- if is_layer_skipped(
95
- prefix=prefix,
96
- ignored_layers=self.ignored_layers,
97
- fused_mapping=self.packed_modules_mapping,
98
- skip_with_substr=True,
99
- ):
100
- return UnquantizedLinearMethod()
101
- return Mxfp8OnlineLinearMethod(self)
102
- elif isinstance(layer, FusedMoE):
103
- if is_layer_skipped(
104
- prefix=prefix,
105
- ignored_layers=self.ignored_layers,
106
- fused_mapping=self.packed_modules_mapping,
107
- skip_with_substr=True,
108
- ):
109
- return UnquantizedFusedMoEMethod(layer.moe_config)
110
- return Mxfp8OnlineMoEMethod(self, layer)
111
- elif isinstance(layer, Attention):
112
- return Fp8KVCacheMethod(self)
113
- return None
114
-
115
-
116
-class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
37
+class Mxfp8OnlineLinearMethod(_Fp8OnlineLinearBase):
38
"""Online MXFP8 linear method.
39
Loads bf16/fp16 checkpoints and quantizes weights to MXFP8 (microscaling
40
FP8 with block-32 scales) during weight loading.
120
-
121
- Args:
122
- quant_config: The MXFP8 quantization config.
41
"""
42
125
- uses_meta_device: bool = True
126
-
127
- def __init__(self, quant_config: "Mxfp8Config"):
128
- self.quant_config = quant_config
43
+ def __init__(self):
44
+ super().__init__()
45
self.kernel = init_mxfp8_linear_kernel()
46
47
def create_weights(
@@ -178,19 +94,15 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
94
return self.kernel.apply_weights(layer, x, bias)
95
96
181
-class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod):
97
+class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
98
"""MoE method for online MXFP8 (block) quantization."""
99
184
- uses_meta_device: bool = True
185
-
186
- def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
187
- FusedMoEMethodBase.__init__(self, layer.moe_config)
188
- self.quant_config = quant_config
189
- assert not quant_config.is_checkpoint_fp8_serialized
190
- assert quant_config.activation_scheme == "dynamic"
100
+ fp8_backend: "Fp8MoeBackend"
101
+ experts_cls: "type[mk.FusedMoEExperts] | None"
102
192
- self.weight_block_size = [1, MXFP8_BLOCK_SIZE]
193
- self.block_quant = True
103
+ def __init__(self, *, layer: torch.nn.Module):
104
+ super().__init__(layer.moe_config)
105
+ self.weight_block_size: list[int] = [1, MXFP8_BLOCK_SIZE]
106
self.weight_scale_name = "weight_scale"
107
108
self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)
@@ -247,6 +159,74 @@ class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod):
159
160
return w_quant, w_scales
161
162
+ def _setup_kernel(
163
+ self,
164
+ layer: "FusedMoE",
165
+ w13: torch.Tensor,
166
+ w2: torch.Tensor,
167
+ w13_scale: torch.Tensor,
168
+ w2_scale: torch.Tensor,
169
+ w13_input_scale: torch.Tensor | None,
170
+ w2_input_scale: torch.Tensor | None,
171
+ ) -> None:
172
+ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
173
+ convert_to_fp8_moe_kernel_format,
174
+ make_fp8_moe_kernel,
175
+ )
176
+
177
+ # Shuffle weights to runtime format.
178
+ w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
179
+ fp8_backend=self.fp8_backend,
180
+ layer=layer,
181
+ w13=w13,
182
+ w2=w2,
183
+ w13_scale=w13_scale,
184
+ w2_scale=w2_scale,
185
+ w13_input_scale=w13_input_scale,
186
+ w2_input_scale=w2_input_scale,
187
+ )
188
+
189
+ replace_parameter(layer, "w13_weight", w13)
190
+ replace_parameter(layer, "w2_weight", w2)
191
+ replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
192
+ replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
193
+
194
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
195
+ if self.moe_quant_config:
196
+ assert self.experts_cls is not None
197
+ self.moe_kernel = make_fp8_moe_kernel(
198
+ moe_quant_config=self.moe_quant_config,
199
+ moe_config=self.moe,
200
+ fp8_backend=self.fp8_backend,
201
+ experts_cls=self.experts_cls,
202
+ routing_tables=layer._maybe_init_expert_routing_tables(),
203
+ shared_experts=layer.shared_experts,
204
+ )
205
+
206
+ def get_fused_moe_quant_config(
207
+ self, layer: torch.nn.Module
208
+ ) -> "FusedMoEQuantConfig":
209
+ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
210
+ make_fp8_moe_quant_config,
211
+ )
212
+
213
+ w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
214
+ w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
215
+ a1_scale = layer.w13_input_scale
216
+ a2_scale = layer.w2_input_scale
217
+
218
+ quant_config = make_fp8_moe_quant_config(
219
+ fp8_backend=self.fp8_backend,
220
+ w1_scale=w1_scale,
221
+ w2_scale=w2_scale,
222
+ a1_scale=a1_scale,
223
+ a2_scale=a2_scale,
224
+ block_shape=self.weight_block_size,
225
+ )
226
+
227
+ self._maybe_inject_biases(quant_config, layer)
228
+ return quant_config
229
+
230
def process_weights_after_loading(self, layer: Module) -> None:
231
if getattr(layer, "_already_called_process_weights_after_loading", False):
232
return
vllm/model_executor/warmup/deep_gemm_warmup.py
+1
-1
@@ -19,7 +19,7 @@ from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
19
)
20
from vllm.model_executor.layers.linear import LinearBase
21
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
22
-from vllm.model_executor.layers.quantization.mxfp8 import Mxfp8OnlineLinearMethod
22
+from vllm.model_executor.layers.quantization.online.mxfp8 import Mxfp8OnlineLinearMethod
23
from vllm.tracing import instrument
24
from vllm.utils.deep_gemm import (
25
fp8_gemm_nt,