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, 07:15 UTC
[Quant] add CompressedTensorsW8A8Mxfp8 for linear and MoE layers (#38815)
Signed-off-by: EdalatiAli <aliedalati@cohere.com>
EdalatiAli committed
Apr 11, 2026 at 19:21 UTC
df1e30e74be91bd48fd015c417e7699310a64119
6 files changed
+371
tests/quantization/test_compressed_tensors.py
+36
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
28
CompressedTensorsW4A16Fp4,
29
CompressedTensorsW8A8Fp8,
30
CompressedTensorsW8A8Int8,
31
+ CompressedTensorsW8A8Mxfp8,
32
CompressedTensorsW8A16Fp8,
33
CompressedTensorsWNA16,
34
)
@@ -632,3 +633,38 @@ def test_get_quant_method_returns_none_for_unmatched_parallel_lm_head():
633
assert method is None, (
634
f"Expected None for unmatched ParallelLMHead, got {type(method).__name__}"
635
)
636
+
637
+
638
+@pytest.mark.skipif(
639
+ not current_platform.is_cuda() or not current_platform.has_device_capability(75),
640
+ reason="MXFP8 requires Turing (sm_75+) or newer.",
641
+)
642
+def test_compressed_tensors_mxfp8_moe_setup(vllm_runner):
643
+ """Verify MXFP8 scheme, dtypes, and generation for a MoE model."""
644
+ model_path = "AliEdalati97/Qwen3-30B-A3B-MXFP8"
645
+ with vllm_runner(
646
+ model_path,
647
+ enforce_eager=True,
648
+ load_format="dummy",
649
+ hf_overrides={"num_hidden_layers": 4},
650
+ ) as llm:
651
+
652
+ def check_model(model):
653
+ from vllm.model_executor.layers.fused_moe import FusedMoE
654
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_mxfp8 import ( # noqa: E501
655
+ CompressedTensorsW8A8Mxfp8MoEMethod,
656
+ )
657
+
658
+ layer = model.model.layers[0]
659
+
660
+ qkv = layer.self_attn.qkv_proj
661
+ assert isinstance(qkv.quant_method, CompressedTensorsLinearMethod)
662
+ assert isinstance(qkv.scheme, CompressedTensorsW8A8Mxfp8)
663
+
664
+ experts = layer.mlp.experts
665
+ assert isinstance(experts, FusedMoE)
666
+ assert isinstance(experts.quant_method, CompressedTensorsW8A8Mxfp8MoEMethod)
667
+
668
+ llm.apply_model(check_model)
669
+ output = llm.generate_greedy("Hello my name is", max_tokens=4)
670
+ assert output
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+25
@@ -49,6 +49,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
49
CompressedTensorsW4A16Mxfp4,
50
CompressedTensorsW8A8Fp8,
51
CompressedTensorsW8A8Int8,
52
+ CompressedTensorsW8A8Mxfp8,
53
CompressedTensorsW8A16Fp8,
54
CompressedTensorsWNA16,
55
)
@@ -403,6 +404,27 @@ class CompressedTensorsConfig(QuantizationConfig):
404
and is_symmetric
405
)
406
407
+ @staticmethod
408
+ def _is_mxfp8(quant_args: QuantizationArgs) -> bool:
409
+ if quant_args is None:
410
+ return False
411
+
412
+ is_group_quant = quant_args.strategy == QuantizationStrategy.GROUP.value
413
+ is_symmetric = quant_args.symmetric
414
+ is_group_size_32 = quant_args.group_size == 32
415
+ is_float_type = quant_args.type == QuantizationType.FLOAT
416
+ is_8_bits = quant_args.num_bits == 8
417
+ is_mxfp8_scale_dtype = quant_args.scale_dtype == torch.uint8
418
+
419
+ return (
420
+ is_group_quant
421
+ and is_float_type
422
+ and is_8_bits
423
+ and is_group_size_32
424
+ and is_symmetric
425
+ and is_mxfp8_scale_dtype
426
+ )
427
+
428
@staticmethod
429
def _is_static_tensor_w8a8(
430
weight_quant: QuantizationArgs, input_quant: QuantizationArgs
@@ -606,6 +628,9 @@ class CompressedTensorsConfig(QuantizationConfig):
628
if self._is_mxfp4(weight_quant):
629
return CompressedTensorsW4A16Mxfp4()
630
631
+ if self._is_mxfp8(weight_quant):
632
+ return CompressedTensorsW8A8Mxfp8()
633
+
634
if self._is_fp8_w4a8_sm90(weight_quant, input_quant):
635
return CompressedTensorsW4A8Fp8(
636
num_bits=weight_quant.num_bits,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+7
@@ -68,6 +68,13 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
68
69
return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config)
70
71
+ if quant_config._is_mxfp8(weight_quant):
72
+ from .compressed_tensors_moe_w8a8_mxfp8 import (
73
+ CompressedTensorsW8A8Mxfp8MoEMethod,
74
+ )
75
+
76
+ return CompressedTensorsW8A8Mxfp8MoEMethod(layer.moe_config)
77
+
78
if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
79
# group_size=None means channelwise
80
group_size = weight_quant.group_size or -1
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py
new
+209
@@ -0,0 +1,209 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
+from vllm.model_executor.layers.fused_moe import (
8
+ FusedMoE,
9
+ FusedMoeWeightScaleSupported,
10
+)
11
+from vllm.model_executor.layers.fused_moe.config import (
12
+ FusedMoEConfig,
13
+ FusedMoEQuantConfig,
14
+)
15
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
16
+ convert_to_fp8_moe_kernel_format,
17
+ make_fp8_moe_kernel,
18
+ make_fp8_moe_quant_config,
19
+)
20
+from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
21
+ select_mxfp8_moe_backend,
22
+)
23
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe import ( # noqa: E501
24
+ CompressedTensorsMoEMethod,
25
+)
26
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
27
+ MXFP8_BLOCK_SIZE,
28
+ MXFP8_SCALE_DTYPE,
29
+ MXFP8_VALUE_DTYPE,
30
+)
31
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
32
+
33
+
34
+class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
35
+ """Compressed-tensors MoE method for pre-quantized MXFP8 (W8A8) checkpoints.
36
+
37
+ Loads FP8 (E4M3) weights with E8M0 uint8 per-group scales (group_size=32)
38
+ from checkpoint. Activations are dynamically quantized to MXFP8 at runtime.
39
+ Supports FlashInfer TRT-LLM and Marlin backends (auto-selected).
40
+ """
41
+
42
+ def __init__(self, moe: FusedMoEConfig):
43
+ super().__init__(moe)
44
+ self.weight_block_size = [1, MXFP8_BLOCK_SIZE]
45
+ self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)
46
+
47
+ def create_weights(
48
+ self,
49
+ layer: torch.nn.Module,
50
+ num_experts: int,
51
+ hidden_size: int,
52
+ intermediate_size_per_partition: int,
53
+ params_dtype: torch.dtype,
54
+ **extra_weight_attrs,
55
+ ):
56
+ layer.num_experts = num_experts
57
+ layer.params_dtype = params_dtype
58
+ w13_num_shards = 2 if self.moe.is_act_and_mul else 1
59
+
60
+ w13_weight = torch.nn.Parameter(
61
+ torch.empty(
62
+ num_experts,
63
+ w13_num_shards * intermediate_size_per_partition,
64
+ hidden_size,
65
+ dtype=MXFP8_VALUE_DTYPE,
66
+ ),
67
+ requires_grad=False,
68
+ )
69
+ layer.register_parameter("w13_weight", w13_weight)
70
+ set_weight_attrs(w13_weight, extra_weight_attrs)
71
+
72
+ w2_weight = torch.nn.Parameter(
73
+ torch.empty(
74
+ num_experts,
75
+ hidden_size,
76
+ intermediate_size_per_partition,
77
+ dtype=MXFP8_VALUE_DTYPE,
78
+ ),
79
+ requires_grad=False,
80
+ )
81
+ layer.register_parameter("w2_weight", w2_weight)
82
+ set_weight_attrs(w2_weight, extra_weight_attrs)
83
+
84
+ w13_weight_scale = torch.nn.Parameter(
85
+ torch.empty(
86
+ num_experts,
87
+ w13_num_shards * intermediate_size_per_partition,
88
+ hidden_size // MXFP8_BLOCK_SIZE,
89
+ dtype=MXFP8_SCALE_DTYPE,
90
+ ),
91
+ requires_grad=False,
92
+ )
93
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
94
+ extra_weight_attrs.update(
95
+ {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
96
+ )
97
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
98
+
99
+ w2_weight_scale = torch.nn.Parameter(
100
+ torch.empty(
101
+ num_experts,
102
+ hidden_size,
103
+ intermediate_size_per_partition // MXFP8_BLOCK_SIZE,
104
+ dtype=MXFP8_SCALE_DTYPE,
105
+ ),
106
+ requires_grad=False,
107
+ )
108
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
109
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
110
+
111
+ layer.w13_input_scale = None
112
+ layer.w2_input_scale = None
113
+
114
+ def process_weights_after_loading(self, layer: FusedMoE) -> None:
115
+ layer.weight_block_size = self.weight_block_size
116
+
117
+ w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
118
+ fp8_backend=self.fp8_backend,
119
+ layer=layer,
120
+ w13=layer.w13_weight,
121
+ w2=layer.w2_weight,
122
+ w13_scale=layer.w13_weight_scale,
123
+ w2_scale=layer.w2_weight_scale,
124
+ w13_input_scale=layer.w13_input_scale,
125
+ w2_input_scale=layer.w2_input_scale,
126
+ )
127
+
128
+ replace_parameter(layer, "w13_weight", w13)
129
+ replace_parameter(layer, "w2_weight", w2)
130
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
131
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
132
+
133
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
134
+ if self.moe_quant_config is not None:
135
+ assert self.experts_cls is not None
136
+ self.moe_kernel = make_fp8_moe_kernel(
137
+ moe_quant_config=self.moe_quant_config,
138
+ moe_config=self.moe,
139
+ fp8_backend=self.fp8_backend,
140
+ experts_cls=self.experts_cls,
141
+ routing_tables=layer._maybe_init_expert_routing_tables(),
142
+ shared_experts=layer.shared_experts,
143
+ )
144
+
145
+ def get_fused_moe_quant_config(
146
+ self, layer: torch.nn.Module
147
+ ) -> FusedMoEQuantConfig | None:
148
+ return make_fp8_moe_quant_config(
149
+ fp8_backend=self.fp8_backend,
150
+ w1_scale=layer.w13_weight_scale,
151
+ w2_scale=layer.w2_weight_scale,
152
+ a1_scale=layer.w13_input_scale,
153
+ a2_scale=layer.w2_input_scale,
154
+ block_shape=self.weight_block_size,
155
+ )
156
+
157
+ def maybe_make_prepare_finalize(
158
+ self,
159
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
160
+ ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
161
+ raise ValueError(
162
+ f"{self.__class__.__name__} uses the new modular kernel "
163
+ "initialization logic. This function should not be called."
164
+ )
165
+
166
+ def apply_monolithic(
167
+ self,
168
+ layer: FusedMoE,
169
+ x: torch.Tensor,
170
+ router_logits: torch.Tensor,
171
+ ) -> torch.Tensor:
172
+ assert self.moe_kernel is not None
173
+ return self.moe_kernel.apply_monolithic(
174
+ x,
175
+ layer.w13_weight,
176
+ layer.w2_weight,
177
+ router_logits,
178
+ activation=layer.activation,
179
+ global_num_experts=layer.global_num_experts,
180
+ expert_map=layer.expert_map,
181
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
182
+ num_expert_group=layer.num_expert_group,
183
+ topk_group=layer.topk_group,
184
+ e_score_correction_bias=layer.e_score_correction_bias,
185
+ routed_scaling_factor=layer.routed_scaling_factor,
186
+ )
187
+
188
+ def apply(
189
+ self,
190
+ layer: FusedMoE,
191
+ x: torch.Tensor,
192
+ topk_weights: torch.Tensor,
193
+ topk_ids: torch.Tensor,
194
+ shared_experts_input: torch.Tensor | None,
195
+ ) -> torch.Tensor:
196
+ assert not self.is_monolithic
197
+ assert self.moe_kernel is not None
198
+ return self.moe_kernel.apply(
199
+ x,
200
+ layer.w13_weight,
201
+ layer.w2_weight,
202
+ topk_weights,
203
+ topk_ids,
204
+ activation=layer.activation,
205
+ global_num_experts=layer.global_num_experts,
206
+ expert_map=layer.expert_map,
207
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
208
+ shared_experts_input=shared_experts_input,
209
+ )
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
+2
@@ -9,6 +9,7 @@ from .compressed_tensors_w4a16_mxfp4 import CompressedTensorsW4A16Mxfp4
9
from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4
10
from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
11
from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
12
+from .compressed_tensors_w8a8_mxfp8 import CompressedTensorsW8A8Mxfp8
13
from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8
14
from .compressed_tensors_wNa16 import WNA16_SUPPORTED_BITS, CompressedTensorsWNA16
15
@@ -28,4 +29,5 @@ __all__ = [
29
"CompressedTensorsW4A4Fp4",
30
"CompressedTensorsW4A8Int",
31
"CompressedTensorsW4A8Fp8",
32
+ "CompressedTensorsW8A8Mxfp8",
33
]
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py
new
+92
@@ -0,0 +1,92 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import torch
6
+
7
+from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel
8
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
9
+ CompressedTensorsScheme,
10
+)
11
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
12
+ MXFP8_BLOCK_SIZE,
13
+ MXFP8_SCALE_DTYPE,
14
+ MXFP8_VALUE_DTYPE,
15
+)
16
+from vllm.model_executor.parameter import (
17
+ GroupQuantScaleParameter,
18
+ ModelWeightParameter,
19
+)
20
+
21
+__all__ = ["CompressedTensorsW8A8Mxfp8"]
22
+
23
+
24
+class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme):
25
+ """
26
+ Compressed tensors scheme for MXFP8 quantization (W8A8).
27
+
28
+ Loads pre-quantized MXFP8 weights from compressed-tensors checkpoints.
29
+ Activations are dynamically quantized to MXFP8 at runtime.
30
+
31
+ MXFP8 format:
32
+ - 8-bit float weights (E4M3) stored as float8_e4m3fn
33
+ - Per-group E8M0 scales (uint8) with group_size=32
34
+ - Activations dynamically quantized to MXFP8 during inference
35
+ """
36
+
37
+ def __init__(self):
38
+ self.kernel = init_mxfp8_linear_kernel()
39
+
40
+ @classmethod
41
+ def get_min_capability(cls) -> int:
42
+ return 75
43
+
44
+ def create_weights(
45
+ self,
46
+ layer: torch.nn.Module,
47
+ output_partition_sizes: list[int],
48
+ input_size_per_partition: int,
49
+ params_dtype: torch.dtype,
50
+ weight_loader: Callable,
51
+ **kwargs,
52
+ ):
53
+ output_size_per_partition = sum(output_partition_sizes)
54
+ layer.logical_widths = output_partition_sizes
55
+ layer.input_size_per_partition = input_size_per_partition
56
+ layer.output_size_per_partition = output_size_per_partition
57
+ layer.params_dtype = params_dtype
58
+
59
+ weight = ModelWeightParameter(
60
+ data=torch.empty(
61
+ output_size_per_partition,
62
+ input_size_per_partition,
63
+ dtype=MXFP8_VALUE_DTYPE,
64
+ ),
65
+ input_dim=1,
66
+ output_dim=0,
67
+ weight_loader=weight_loader,
68
+ )
69
+ layer.register_parameter("weight", weight)
70
+
71
+ weight_scale = GroupQuantScaleParameter(
72
+ data=torch.empty(
73
+ output_size_per_partition,
74
+ input_size_per_partition // MXFP8_BLOCK_SIZE,
75
+ dtype=MXFP8_SCALE_DTYPE,
76
+ ),
77
+ input_dim=1,
78
+ output_dim=0,
79
+ weight_loader=weight_loader,
80
+ )
81
+ layer.register_parameter("weight_scale", weight_scale)
82
+
83
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
84
+ self.kernel.process_weights_after_loading(layer)
85
+
86
+ def apply_weights(
87
+ self,
88
+ layer: torch.nn.Module,
89
+ x: torch.Tensor,
90
+ bias: torch.Tensor | None = None,
91
+ ) -> torch.Tensor:
92
+ return self.kernel.apply_weights(layer, x, bias)