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 10, 09:00 UTC
[Quant] Support MXFP4 W4A16 for compressed-tensors dense models (#31926)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com>
Michael Goin committed
Jan 10, 2026 at 09:44 UTC
e6c6f2c79d2c1e0765d17b4dec83a4a8283342e4
3 files changed
+131
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+23
@@ -43,6 +43,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
43
CompressedTensorsW4A8Fp8,
44
CompressedTensorsW4A8Int,
45
CompressedTensorsW4A16Fp4,
46
+ CompressedTensorsW4A16Mxfp4,
47
CompressedTensorsW4A16Sparse24,
48
CompressedTensorsW8A8Fp8,
49
CompressedTensorsW8A8Int8,
@@ -350,6 +351,25 @@ class CompressedTensorsConfig(QuantizationConfig):
351
and is_symmetric
352
)
353
354
+ @staticmethod
355
+ def _is_mxfp4(quant_args: QuantizationArgs) -> bool:
356
+ if quant_args is None:
357
+ return False
358
+
359
+ is_group_quant = quant_args.strategy == QuantizationStrategy.GROUP.value
360
+ is_symmetric = quant_args.symmetric
361
+ is_group_size_32 = quant_args.group_size == 32
362
+ is_float_type = quant_args.type == QuantizationType.FLOAT
363
+ is_4_bits = quant_args.num_bits == 4
364
+
365
+ return (
366
+ is_group_quant
367
+ and is_float_type
368
+ and is_4_bits
369
+ and is_group_size_32
370
+ and is_symmetric
371
+ )
372
+
373
@staticmethod
374
def _is_static_tensor_w8a8(
375
weight_quant: QuantizationArgs, input_quant: QuantizationArgs
@@ -550,6 +570,9 @@ class CompressedTensorsConfig(QuantizationConfig):
570
if self._is_nvfp4_format(weight_quant) and input_quant is None:
571
return CompressedTensorsW4A16Fp4()
572
573
+ if self._is_mxfp4(weight_quant):
574
+ return CompressedTensorsW4A16Mxfp4()
575
+
576
if self._is_fp8_w4a8_sm90(weight_quant, input_quant):
577
return CompressedTensorsW4A8Fp8(
578
num_bits=weight_quant.num_bits,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
+2
@@ -9,6 +9,7 @@ from .compressed_tensors_w4a16_24 import (
9
W4A16SPARSE24_SUPPORTED_BITS,
10
CompressedTensorsW4A16Sparse24,
11
)
12
+from .compressed_tensors_w4a16_mxfp4 import CompressedTensorsW4A16Mxfp4
13
from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4
14
from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
15
from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
@@ -29,6 +30,7 @@ __all__ = [
30
"W4A16SPARSE24_SUPPORTED_BITS",
31
"CompressedTensors24",
32
"CompressedTensorsW4A16Fp4",
33
+ "CompressedTensorsW4A16Mxfp4",
34
"CompressedTensorsW4A4Fp4",
35
"CompressedTensorsW4A8Int",
36
"CompressedTensorsW4A8Fp8",
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_mxfp4.py
new
+106
@@ -0,0 +1,106 @@
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
+from torch.nn.parameter import Parameter
7
+
8
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
9
+ CompressedTensorsScheme,
10
+)
11
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
12
+ apply_fp4_marlin_linear,
13
+ prepare_fp4_layer_for_marlin,
14
+)
15
+from vllm.model_executor.parameter import (
16
+ GroupQuantScaleParameter,
17
+ ModelWeightParameter,
18
+)
19
+
20
+__all__ = ["CompressedTensorsW4A16Mxfp4"]
21
+
22
+
23
+class CompressedTensorsW4A16Mxfp4(CompressedTensorsScheme):
24
+ """
25
+ Compressed tensors scheme for MXFP4 weight-only quantization.
26
+
27
+ Supports models quantized with the compressed-tensors mxfp4-pack-quantized
28
+ format.
29
+
30
+ MXFP4 format:
31
+ - 4-bit float weights (E2M1) packed into uint8
32
+ - Per-group E8M0 scales with group_size=32
33
+ - No global scale (unlike NVFP4)
34
+ """
35
+
36
+ def __init__(self):
37
+ self.group_size = 32
38
+
39
+ @classmethod
40
+ def get_min_capability(cls) -> int:
41
+ return 80
42
+
43
+ def create_weights(
44
+ self,
45
+ layer: torch.nn.Module,
46
+ output_partition_sizes: list[int],
47
+ input_size_per_partition: int,
48
+ params_dtype: torch.dtype,
49
+ weight_loader: Callable,
50
+ **kwargs,
51
+ ):
52
+ output_size_per_partition = sum(output_partition_sizes)
53
+ layer.logical_widths = output_partition_sizes
54
+ layer.input_size_per_partition = input_size_per_partition
55
+ layer.output_size_per_partition = output_size_per_partition
56
+ layer.params_dtype = params_dtype
57
+
58
+ # Packed FP4 weights (2 values per byte)
59
+ weight = ModelWeightParameter(
60
+ data=torch.empty(
61
+ output_size_per_partition,
62
+ input_size_per_partition // 2,
63
+ dtype=torch.uint8,
64
+ ),
65
+ input_dim=1,
66
+ output_dim=0,
67
+ weight_loader=weight_loader,
68
+ )
69
+ layer.register_parameter("weight_packed", weight)
70
+
71
+ # Per-group E8M0 scales
72
+ weight_scale = GroupQuantScaleParameter(
73
+ data=torch.empty(
74
+ output_size_per_partition,
75
+ input_size_per_partition // self.group_size,
76
+ dtype=torch.uint8,
77
+ ),
78
+ input_dim=1,
79
+ output_dim=0,
80
+ weight_loader=weight_loader,
81
+ )
82
+ layer.register_parameter("weight_scale", weight_scale)
83
+
84
+ def process_weights_after_loading(self, layer) -> None:
85
+ # Rename weight_packed to weight that marlin expects
86
+ layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
87
+ del layer.weight_packed
88
+
89
+ prepare_fp4_layer_for_marlin(layer)
90
+
91
+ def apply_weights(
92
+ self,
93
+ layer: torch.nn.Module,
94
+ x: torch.Tensor,
95
+ bias: torch.Tensor | None = None,
96
+ ) -> torch.Tensor:
97
+ return apply_fp4_marlin_linear(
98
+ input=x,
99
+ weight=layer.weight,
100
+ weight_scale=layer.weight_scale,
101
+ weight_scale_2=None,
102
+ workspace=layer.workspace,
103
+ size_n=layer.output_size_per_partition,
104
+ size_k=layer.input_size_per_partition,
105
+ bias=bias,
106
+ )