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 14, 05:30 UTC
Quantization: support FP4 quantized models on AMD CDNA2/CDNA3 GPUs (#22527)
Signed-off-by: feng <fengli1702@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Daifeng Li committed
Aug 23, 2025 at 10:53 UTC
fa78de9dc373c3bcc365ccc0f0631010df22ebe0
8 files changed
+451
-5
requirements/rocm.txt
+1
-1
@@ -17,4 +17,4 @@ setuptools>=77.0.3,<80.0.0
17
setuptools-scm>=8
18
runai-model-streamer==0.11.0
19
runai-model-streamer-s3==0.11.0
20
-conch-triton-kernels==1.2.1
20
+conch-triton-kernels==1.2.1
\ No newline at end of file
setup.py
+2
@@ -695,6 +695,8 @@ setup(
695
"video": [], # Kept for backwards compatibility
696
# FlashInfer should be updated together with the Dockerfile
697
"flashinfer": ["flashinfer-python==0.2.12"],
698
+ # Optional deps for AMD FP4 quantization support
699
+ "petit-kernel": ["petit-kernel"],
700
},
701
cmdclass=cmdclass,
702
package_data=package_data,
vllm/config/__init__.py
+15
-3
@@ -1119,9 +1119,20 @@ class ModelConfig:
1119
def _verify_quantization(self) -> None:
1120
supported_quantization = me_quant.QUANTIZATION_METHODS
1121
optimized_quantization_methods = [
1122
- "fp8", "modelopt", "gptq_marlin_24", "gptq_marlin", "awq_marlin",
1123
- "fbgemm_fp8", "compressed-tensors", "experts_int8", "quark",
1124
- "modelopt_fp4", "bitblas", "gptq_bitblas", "inc"
1122
+ "fp8",
1123
+ "modelopt",
1124
+ "gptq_marlin_24",
1125
+ "gptq_marlin",
1126
+ "awq_marlin",
1127
+ "fbgemm_fp8",
1128
+ "compressed-tensors",
1129
+ "experts_int8",
1130
+ "quark",
1131
+ "modelopt_fp4",
1132
+ "bitblas",
1133
+ "gptq_bitblas",
1134
+ "inc",
1135
+ "petit_nvfp4",
1136
]
1137
if self.quantization is not None:
1138
self.quantization = cast(me_quant.QuantizationMethods,
@@ -1153,6 +1164,7 @@ class ModelConfig:
1164
"moe_wna16",
1165
"modelopt",
1166
"modelopt_fp4",
1167
+ "petit_nvfp4",
1168
]
1169
quantization_methods = [
1170
q for q in supported_quantization if q not in overrides
vllm/model_executor/layers/linear.py
+1
@@ -52,6 +52,7 @@ WEIGHT_LOADER_V2_SUPPORTED = [
52
"HQQMarlinMethod",
53
"QuarkLinearMethod",
54
"ModelOptNvFp4LinearMethod",
55
+ "PetitNvFp4LinearMethod",
56
]
57
58
vllm/model_executor/layers/quantization/__init__.py
+3
@@ -35,6 +35,7 @@ QuantizationMethods = Literal[
35
"rtn",
36
"inc",
37
"mxfp4",
38
+ "petit_nvfp4",
39
]
40
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
41
@@ -108,6 +109,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
109
from .moe_wna16 import MoeWNA16Config
110
from .mxfp4 import Mxfp4Config
111
from .neuron_quant import NeuronQuantConfig
112
+ from .petit import PetitNvFp4Config
113
from .ptpc_fp8 import PTPCFp8Config
114
from .rtn import RTNConfig
115
from .torchao import TorchAOConfig
@@ -142,6 +144,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
144
"rtn": RTNConfig,
145
"inc": INCConfig,
146
"mxfp4": Mxfp4Config,
147
+ "petit_nvfp4": PetitNvFp4Config,
148
}
149
# Update the `method_to_config` with customized quantization methods.
150
method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
vllm/model_executor/layers/quantization/petit.py
new
+306
@@ -0,0 +1,306 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py
4
+
5
+from typing import Any, Optional
6
+
7
+import regex as re
8
+import torch
9
+from torch.nn.parameter import Parameter
10
+
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
13
+ UnquantizedLinearMethod)
14
+from vllm.model_executor.layers.quantization import QuantizationMethods
15
+from vllm.model_executor.layers.quantization.base_config import (
16
+ QuantizationConfig, QuantizeMethodBase)
17
+from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
18
+from vllm.model_executor.layers.quantization.utils.petit_utils import (
19
+ apply_petit_nvfp4_linear, prepare_nvfp4_layer_for_petit,
20
+ verify_petit_nvfp4_supported)
21
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
+ is_layer_skipped)
23
+from vllm.model_executor.parameter import (ModelWeightParameter,
24
+ PerTensorScaleParameter)
25
+from vllm.platforms import current_platform
26
+
27
+# Initialize logger for the module
28
+logger = init_logger(__name__)
29
+
30
+
31
+# Configuration class to support the NVFP4 quantized model
32
+# generated by the ModelOpt quantization tool
33
+class PetitNvFp4Config(QuantizationConfig):
34
+ """Config class for Petit FP4."""
35
+
36
+ def __init__(
37
+ self,
38
+ is_checkpoint_nvfp4_serialized: bool = False,
39
+ kv_cache_quant_algo: Optional[str] = None,
40
+ group_size: Optional[int] = None,
41
+ exclude_modules: Optional[list[str]] = None,
42
+ ) -> None:
43
+ self._check_hardware_support()
44
+ self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
45
+ if is_checkpoint_nvfp4_serialized:
46
+ logger.warning("Detected nvfp4 checkpoint. Please note that the "
47
+ "format is experimental and subject to change.")
48
+ self.group_size = group_size
49
+ self.kv_cache_quant_algo = kv_cache_quant_algo
50
+ self.exclude_modules = exclude_modules
51
+
52
+ def _check_hardware_support(self) -> None:
53
+ """
54
+ Verifies that the current hardware is supported by the Petit backend.
55
+ This backend is specifically designed for AMD GPUs and is not
56
+ supported on the CUDA platform.
57
+ """
58
+ # This check ensures the code is NOT running on an NVIDIA GPU.
59
+ if current_platform.is_cuda():
60
+ raise ValueError(
61
+ "The 'petit' quantization backend is designed for AMD GPUs "
62
+ "and is not supported on the CUDA platform. For NVIDIA GPUs, "
63
+ "please use a different quantization method such as FP8, AWQ, "
64
+ "or GPTQ.")
65
+
66
+ @classmethod
67
+ def get_name(cls) -> QuantizationMethods:
68
+ return "petit_nvfp4"
69
+
70
+ @classmethod
71
+ def get_supported_act_dtypes(cls) -> list[torch.dtype]:
72
+ return [torch.bfloat16, torch.half]
73
+
74
+ @classmethod
75
+ def get_min_capability(cls) -> int:
76
+ # Petit supports the gfx90a and gfx942 GPUs
77
+ return 90
78
+
79
+ @classmethod
80
+ def get_config_filenames(cls) -> list[str]:
81
+ return ["hf_quant_config.json"]
82
+
83
+ @classmethod
84
+ def from_config(cls, config: dict[str, Any]) -> "PetitNvFp4Config":
85
+ qc = cls.get_from_keys(config, ["quantization"])
86
+
87
+ quant_method_raw = qc.get("quant_algo")
88
+ if not isinstance(quant_method_raw, str) or not quant_method_raw:
89
+ raise ValueError(
90
+ "Missing or invalid 'quant_algo' in quantization config.")
91
+ quant_method = quant_method_raw.upper()
92
+
93
+ group_size_raw = qc.get("group_size")
94
+ if not isinstance(group_size_raw, int):
95
+ raise ValueError(
96
+ "Missing or invalid 'group_size' (int) in hf_quant_config.json."
97
+ )
98
+ group_size = group_size_raw
99
+
100
+ verify_petit_nvfp4_supported(quant_method, group_size)
101
+
102
+ kv_cache_quant_algo_raw = qc.get("kv_cache_quant_algo") or "auto"
103
+ if not isinstance(kv_cache_quant_algo_raw, str):
104
+ raise ValueError(
105
+ "'kv_cache_quant_algo' must be a string if provided.")
106
+ kv_cache_quant_algo = kv_cache_quant_algo_raw
107
+
108
+ exclude_raw = qc.get("exclude_modules", [])
109
+ if exclude_raw is None:
110
+ exclude_modules: list[str] = []
111
+ elif isinstance(exclude_raw, list) and all(
112
+ isinstance(x, str) for x in exclude_raw):
113
+ exclude_modules = exclude_raw
114
+ else:
115
+ raise ValueError(
116
+ "'exclude_modules' must be a list[str] (or omitted).")
117
+
118
+ is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method
119
+
120
+ return cls(
121
+ is_checkpoint_nvfp4_serialized=is_checkpoint_nvfp4_serialized,
122
+ kv_cache_quant_algo=kv_cache_quant_algo,
123
+ group_size=group_size,
124
+ exclude_modules=exclude_modules,
125
+ )
126
+
127
+ @classmethod
128
+ def override_quantization_method(
129
+ cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:
130
+ if not current_platform.is_rocm():
131
+ return None
132
+
133
+ qc = hf_quant_cfg.get("quantization", hf_quant_cfg)
134
+ algo = (qc.get("quant_algo") or qc.get("quant_method") or "").upper()
135
+ if algo in ("NVFP4", "MODELOPT_FP4", "MODELOPT"):
136
+ return cls.get_name() # "petit_nvfp4"
137
+ return None
138
+
139
+ @classmethod
140
+ def is_petit_nvfp4_compatible(cls, quant_config: dict[str, Any]) -> bool:
141
+ qc = quant_config.get("quantization", quant_config)
142
+ algo = (qc.get("quant_algo") or qc.get("quant_method") or "").upper()
143
+ return algo == "NVFP4"
144
+
145
+ def is_layer_excluded(self, prefix: str,
146
+ exclude_modules: list[str]) -> bool:
147
+ for pattern in exclude_modules:
148
+ regex_str = pattern.replace(".", r"\.").replace("*", r".*")
149
+ if re.fullmatch(regex_str, prefix):
150
+ return True
151
+ return False
152
+
153
+ def get_quant_method(self, layer: torch.nn.Module,
154
+ prefix: str) -> Optional["QuantizeMethodBase"]:
155
+ from vllm.attention.layer import Attention # Avoid circular import
156
+
157
+ exclude = self.require_exclude_modules()
158
+
159
+ if isinstance(layer, LinearBase):
160
+ if is_layer_skipped(prefix, exclude) or self.is_layer_excluded(
161
+ prefix, exclude):
162
+ return UnquantizedLinearMethod()
163
+ return PetitNvFp4LinearMethod(self)
164
+ elif isinstance(layer, Attention):
165
+ return PetitFp8KVCacheMethod(self)
166
+ return None
167
+
168
+ def get_scaled_act_names(self) -> list[str]:
169
+ return []
170
+
171
+ def require_group_size(self) -> int:
172
+ if self.group_size is None:
173
+ logger.warning("group_size not set; defaulting to 16 for NVFP4.")
174
+ return 16
175
+ return self.group_size
176
+
177
+ def require_kv_cache_quant_algo(self) -> str:
178
+ return self.kv_cache_quant_algo or "auto"
179
+
180
+ def require_exclude_modules(self) -> list[str]:
181
+ return list(self.exclude_modules or [])
182
+
183
+
184
+class PetitFp8KVCacheMethod(BaseKVCacheMethod):
185
+ """
186
+ Supports loading kv-cache scaling factors from FP8 checkpoints.
187
+ """
188
+
189
+ def __init__(self, quant_config: PetitNvFp4Config):
190
+ super().__init__(quant_config)
191
+
192
+
193
+class PetitNvFp4LinearMethod(LinearMethodBase):
194
+ """Linear method for NVFP4.
195
+ Supports loading NVFP4 checkpoints with the following structure:
196
+
197
+ |Tensor Name | datatype | shape |
198
+ |----------------------------------------------------|
199
+ |input_scale | torch.float32 | scalar |
200
+ |weight | NVFP4(SE2M1) | [1, X, y/2] |
201
+ |weight_scale | FP8-E4M3 | [X, Y] |
202
+ |weight_scale_2 | torch.float32 | scalar |
203
+
204
+ The weights are quantized per block of 16 elements.
205
+ Args: quant_config: The ModelOpt quantization config.
206
+ """
207
+
208
+ def __init__(self, quant_config: PetitNvFp4Config):
209
+ self.quant_config = quant_config
210
+
211
+ def create_weights(
212
+ self,
213
+ layer: torch.nn.Module,
214
+ input_size_per_partition: int,
215
+ output_partition_sizes: list[int],
216
+ input_size: int,
217
+ output_size: int,
218
+ params_dtype: torch.dtype,
219
+ **extra_weight_attrs,
220
+ ):
221
+ del input_size, output_size
222
+ if not self.quant_config.is_checkpoint_nvfp4_serialized:
223
+ raise ValueError("NVFP4 quantization was selected, "
224
+ " dynamic quantization is not supported.")
225
+
226
+ output_size_per_partition = sum(output_partition_sizes)
227
+ weight_loader = extra_weight_attrs.get("weight_loader")
228
+
229
+ layer.logical_widths = output_partition_sizes
230
+
231
+ layer.input_size_per_partition = input_size_per_partition
232
+ layer.output_size_per_partition = output_size_per_partition
233
+ if input_size_per_partition % 16 != 0:
234
+ raise ValueError("Unsupported model when in features size is "
235
+ "not multiple of 16")
236
+
237
+ weight_dtype = (torch.float8_e4m3fn
238
+ if self.quant_config.is_checkpoint_nvfp4_serialized
239
+ else params_dtype)
240
+
241
+ weight = ModelWeightParameter(
242
+ data=torch.empty(
243
+ # 2 fp4 data is packed in one uint8 in the input dimension
244
+ output_size_per_partition,
245
+ input_size_per_partition // 2,
246
+ dtype=torch.uint8,
247
+ ),
248
+ input_dim=1,
249
+ output_dim=0,
250
+ weight_loader=weight_loader,
251
+ )
252
+ layer.register_parameter("weight", weight)
253
+
254
+ input_scale = PerTensorScaleParameter(
255
+ data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
256
+ weight_loader=weight_loader,
257
+ )
258
+
259
+ layer.register_parameter("input_scale", input_scale)
260
+
261
+ weight_scale_2 = PerTensorScaleParameter(
262
+ data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
263
+ weight_loader=weight_loader,
264
+ )
265
+ layer.register_parameter("weight_scale_2", weight_scale_2)
266
+
267
+ group_size = self.quant_config.require_group_size()
268
+ weight_scale = ModelWeightParameter(
269
+ data=torch.empty(
270
+ output_size_per_partition,
271
+ input_size_per_partition // group_size,
272
+ dtype=weight_dtype,
273
+ ),
274
+ input_dim=1,
275
+ output_dim=0,
276
+ weight_loader=weight_loader,
277
+ )
278
+
279
+ layer.register_parameter("weight_scale", weight_scale)
280
+
281
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
282
+ input_scale_2 = layer.input_scale.max().to(torch.float32)
283
+ weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
284
+ layer.input_scale = Parameter(input_scale_2, requires_grad=False)
285
+ layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False)
286
+ layer.alpha = Parameter(layer.input_scale * layer.weight_scale_2,
287
+ requires_grad=False)
288
+
289
+ prepare_nvfp4_layer_for_petit(layer)
290
+ del layer.input_scale
291
+
292
+ def apply(
293
+ self,
294
+ layer: torch.nn.Module,
295
+ x: torch.Tensor,
296
+ bias: Optional[torch.Tensor] = None,
297
+ ) -> torch.Tensor:
298
+ return apply_petit_nvfp4_linear(
299
+ input=x,
300
+ weight=layer.weight,
301
+ weight_scale=layer.weight_scale,
302
+ weight_scale_2=layer.weight_scale_2,
303
+ size_n=layer.output_size_per_partition,
304
+ size_k=layer.input_size_per_partition,
305
+ bias=bias,
306
+ )
vllm/model_executor/layers/quantization/utils/petit_utils.py
new
+122
@@ -0,0 +1,122 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import TYPE_CHECKING, Optional
4
+
5
+import torch
6
+
7
+# TYPE_CHECKING is used for static type analysis to prevent circular imports.
8
+if TYPE_CHECKING:
9
+ from types import ModuleType
10
+
11
+# 1. Create a global variable as a placeholder for the module
12
+_petit_kernel: Optional["ModuleType"] = None
13
+
14
+_PETIT_INSTALL_MSG = ("Petit is not installed. Please install it with "
15
+ "`pip install petit-kernel`.")
16
+
17
+
18
+def _import_petit_kernel() -> "ModuleType":
19
+ """
20
+ A helper function to handle the lazy import.
21
+ The first time this function is called, it will import the petit_kernel
22
+ library and store it in the global _petit_kernel variable.
23
+ Subsequent calls will return the already-loaded module directly.
24
+ """
25
+ global _petit_kernel
26
+ if _petit_kernel is not None:
27
+ return _petit_kernel
28
+
29
+ try:
30
+ import petit_kernel
31
+ _petit_kernel = petit_kernel
32
+ return _petit_kernel
33
+ except ImportError:
34
+ # The 'from None' syntax prevents chaining the original ImportError,
35
+ # making the traceback cleaner.
36
+ raise ImportError(_PETIT_INSTALL_MSG) from None
37
+
38
+
39
+# The _require_petit function can now be a simple alias for consistency.
40
+_require_petit = _import_petit_kernel
41
+
42
+
43
+def _check_petit_nvfp4_supported(
44
+ quant_method: str,
45
+ group_size: Optional[int]) -> tuple[bool, Optional[str]]:
46
+ if quant_method != "NVFP4":
47
+ return (
48
+ False,
49
+ ("Petit currently only supports: NVFP4 quantizations in sglang. "
50
+ "Please check the `hf_quant_config.json` file for your model's "
51
+ "quant configuration."),
52
+ )
53
+ if group_size is not None and group_size != 16:
54
+ return (
55
+ False,
56
+ "Petit currently only supports: group_size=16 quantizations.",
57
+ )
58
+ return (True, None)
59
+
60
+
61
+def verify_petit_nvfp4_supported(quant_method: str,
62
+ group_size: Optional[int]) -> None:
63
+ supported, error_msg = _check_petit_nvfp4_supported(
64
+ quant_method, group_size)
65
+ if not supported:
66
+ assert error_msg is not None
67
+ raise ValueError(error_msg)
68
+
69
+
70
+def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None:
71
+ # 2. Call _import_petit_kernel() to trigger (or get) the import.
72
+ petit_kernel = _import_petit_kernel()
73
+
74
+ # Repack weights to petit format
75
+ part_size_n = layer.output_size_per_partition
76
+ part_size_k = layer.input_size_per_partition
77
+ qweight = layer.weight.view(torch.int32).contiguous()
78
+
79
+ # 3. Call functions through the imported module variable.
80
+ petit_qweight = petit_kernel.repack_nvfp4(qweight,
81
+ size_n=part_size_n,
82
+ size_k=part_size_k)
83
+ layer.weight = torch.nn.Parameter(petit_qweight, requires_grad=False)
84
+
85
+ # Permute scales
86
+ weight_scale = petit_kernel.process_nvfp4_scales(scales=layer.weight_scale,
87
+ size_k=part_size_k,
88
+ size_n=part_size_n)
89
+ layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
90
+
91
+
92
+def apply_petit_nvfp4_linear(
93
+ input: torch.Tensor,
94
+ weight: torch.Tensor,
95
+ weight_scale: torch.Tensor,
96
+ weight_scale_2: torch.Tensor,
97
+ size_n: int,
98
+ size_k: int,
99
+ bias: Optional[torch.Tensor] = None,
100
+) -> torch.Tensor:
101
+ # Trigger (or get) the import here as well.
102
+ petit_kernel = _import_petit_kernel()
103
+
104
+ reshaped_x = input.reshape(-1, input.shape[-1])
105
+ out_shape = input.shape[:-1] + (size_n, )
106
+
107
+ # TODO: Use auto-tuning to find the performant solution_id
108
+ # Call the function via the module variable.
109
+ output = petit_kernel.mul_nvfp4_a16(
110
+ a=reshaped_x,
111
+ b=weight,
112
+ s=weight_scale,
113
+ global_scale=weight_scale_2,
114
+ size_m=reshaped_x.size(0),
115
+ size_n=size_n,
116
+ size_k=size_k,
117
+ solution_id=-1,
118
+ )
119
+ if bias is not None:
120
+ output.add_(bias) # In-place add
121
+
122
+ return output.reshape(out_shape)
vllm/platforms/rocm.py
+1
-1
@@ -171,7 +171,7 @@ class RocmPlatform(Platform):
171
172
supported_quantization: list[str] = [
173
"awq", "gptq", "fp8", "compressed-tensors", "fbgemm_fp8", "gguf",
174
- "quark", "ptpc_fp8", "mxfp4"
174
+ "quark", "ptpc_fp8", "mxfp4", "petit_nvfp4"
175
]
176
177
@classmethod