| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Supports FP-Quant compression, see https://arxiv.org/abs/2509.23202 |
| 5 | |
| 6 | from typing import Any, Literal, cast |
| 7 | |
| 8 | import torch |
| 9 | from torch.nn.parameter import Parameter |
| 10 | |
| 11 | from vllm._custom_ops import ( |
| 12 | cutlass_scaled_fp4_mm, |
| 13 | fusedQuantizeMx, |
| 14 | fusedQuantizeNv, |
| 15 | matmul_mxf4_bf16_tn, |
| 16 | ) |
| 17 | from vllm.model_executor.layers.linear import ( |
| 18 | LinearBase, |
| 19 | LinearMethodBase, |
| 20 | UnquantizedLinearMethod, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 23 | from vllm.model_executor.layers.quantization.base_config import QuantizationConfig |
| 24 | from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked |
| 25 | from vllm.model_executor.utils import set_weight_attrs |
| 26 | from vllm.platforms import current_platform |
| 27 | from vllm.utils.torch_utils import direct_register_custom_op |
| 28 | |
| 29 | |
| 30 | class FPQuantConfig(QuantizationConfig): |
| 31 | """Config class for FPQuant.""" |
| 32 | |
| 33 | def __init__( |
| 34 | self, |
| 35 | hadamard_group_size: int = 32, |
| 36 | forward_dtype: str = "mxfp4", |
| 37 | forward_method: str = "abs_max", |
| 38 | modules_to_not_convert: list[str] | None = None, |
| 39 | ) -> None: |
| 40 | super().__init__() |
| 41 | self.hadamard_group_size = hadamard_group_size |
| 42 | self.forward_dtype = forward_dtype |
| 43 | self.forward_method = forward_method |
| 44 | self.modules_to_not_convert = modules_to_not_convert |
| 45 | |
| 46 | def __repr__(self) -> str: |
| 47 | return ( |
| 48 | f"FPQuantConfig(hadamard_group_size={self.hadamard_group_size}, " |
| 49 | f"forward_dtype={self.forward_dtype}, " |
| 50 | f"forward_method={self.forward_method}, " |
| 51 | f"modules_to_not_convert={self.modules_to_not_convert})" |
| 52 | ) |
| 53 | |
| 54 | @classmethod |
| 55 | def get_name(cls) -> QuantizationMethods: |
| 56 | return "fp_quant" |
| 57 | |
| 58 | @classmethod |
| 59 | def get_supported_act_dtypes(cls) -> list[torch.dtype]: |
| 60 | return [torch.bfloat16] |
| 61 | |
| 62 | @classmethod |
| 63 | def get_min_capability(cls) -> int: |
| 64 | return 100 |
| 65 | |
| 66 | @classmethod |
| 67 | def get_config_filenames(cls) -> list[str]: |
| 68 | return [] # no extra configs. |
| 69 | |
| 70 | @classmethod |
| 71 | def from_config(cls, config: dict[str, Any]) -> "FPQuantConfig": |
| 72 | hadamard_group_size = cls.get_from_keys(config, ["hadamard_group_size"]) |
| 73 | forward_dtype = cls.get_from_keys(config, ["forward_dtype"]) |
| 74 | forward_method = cls.get_from_keys(config, ["forward_method"]) |
| 75 | modules_to_not_convert = cls.get_from_keys(config, ["modules_to_not_convert"]) |
| 76 | return cls( |
| 77 | hadamard_group_size, |
| 78 | forward_dtype, |
| 79 | forward_method, |
| 80 | modules_to_not_convert, |
| 81 | ) |
| 82 | |
| 83 | def get_quant_method( |
| 84 | self, layer: torch.nn.Module, prefix: str |
| 85 | ) -> LinearMethodBase | None: |
| 86 | if self.modules_to_not_convert is not None and any( |
| 87 | prefix.endswith(module) for module in self.modules_to_not_convert |
| 88 | ): |
| 89 | return UnquantizedLinearMethod() |
| 90 | |
| 91 | if isinstance(layer, LinearBase): |
| 92 | return FPQuantLinearMethod(self) |
| 93 | return None |
| 94 | |
| 95 | |
| 96 | class FPQuantLinearMethod(LinearMethodBase): |
| 97 | """Linear method for FPQuant. |
| 98 | |
| 99 | Args: |
| 100 | quant_config: The FPQuant quantization config. |
| 101 | |
| 102 | """ |
| 103 | |
| 104 | def __init__(self, quant_config: FPQuantConfig): |
| 105 | self.quant_config = quant_config |
| 106 | |
| 107 | def create_weights( |
| 108 | self, |
| 109 | layer: torch.nn.Module, |
| 110 | input_size_per_partition: int, |
| 111 | output_partition_sizes: list[int], |
| 112 | input_size: int, |
| 113 | output_size: int, |
| 114 | params_dtype: torch.dtype, |
| 115 | **extra_weight_attrs, |
| 116 | ): |
| 117 | del output_size # Unused. |
| 118 | del input_size # Unused. |
| 119 | |
| 120 | if params_dtype != torch.bfloat16: |
| 121 | raise ValueError("Only bfloat16 is currently supported by FPQuant") |
| 122 | if input_size_per_partition % self.quant_config.hadamard_group_size != 0: # noqa: E501 |
| 123 | raise ValueError( |
| 124 | "The input size is not aligned with the quantized " |
| 125 | "weight shape. This can be caused by too large " |
| 126 | "tensor parallel size. Or other skill issues." |
| 127 | ) |
| 128 | |
| 129 | assert self.quant_config.forward_dtype in ["mxfp4", "nvfp4"], ( |
| 130 | "Only mxfp4 and nvfp4 are supported for now" |
| 131 | ) |
| 132 | if self.quant_config.forward_dtype == "mxfp4": |
| 133 | group_size = 32 |
| 134 | elif self.quant_config.forward_dtype == "nvfp4": |
| 135 | group_size = 16 |
| 136 | else: |
| 137 | raise ValueError( |
| 138 | f"Unsupported forward_dtype: {self.quant_config.forward_dtype}" |
| 139 | ) |
| 140 | |
| 141 | qweight = Parameter( |
| 142 | torch.empty( |
| 143 | sum(output_partition_sizes), |
| 144 | input_size_per_partition // 2, |
| 145 | dtype=torch.uint8, |
| 146 | ), |
| 147 | requires_grad=False, |
| 148 | ) |
| 149 | set_weight_attrs( |
| 150 | qweight, |
| 151 | { |
| 152 | "input_dim": 1, |
| 153 | "output_dim": 0, |
| 154 | "packed_dim": 1, |
| 155 | "pack_factor": 2, |
| 156 | } |
| 157 | | extra_weight_attrs, |
| 158 | ) |
| 159 | layer.register_parameter("qweight", qweight) |
| 160 | |
| 161 | scales = Parameter( |
| 162 | torch.empty( |
| 163 | sum(output_partition_sizes), |
| 164 | input_size_per_partition // group_size, |
| 165 | dtype=torch.uint8, |
| 166 | ), |
| 167 | requires_grad=False, |
| 168 | ) |
| 169 | set_weight_attrs( |
| 170 | scales, |
| 171 | { |
| 172 | "input_dim": 1, |
| 173 | "output_dim": 0, |
| 174 | "packed_dim": 1, |
| 175 | "pack_factor": group_size, |
| 176 | } |
| 177 | | extra_weight_attrs, |
| 178 | ) |
| 179 | layer.register_parameter("scales", scales) |
| 180 | |
| 181 | weight_global_scale = Parameter( |
| 182 | torch.empty(1, dtype=torch.float32), |
| 183 | requires_grad=False, |
| 184 | ) |
| 185 | set_weight_attrs( |
| 186 | weight_global_scale, {"ignore_warning": True} | extra_weight_attrs |
| 187 | ) |
| 188 | layer.register_parameter("weight_global_scale", weight_global_scale) |
| 189 | |
| 190 | act_global_scale = Parameter( |
| 191 | torch.empty(1, dtype=torch.float32), |
| 192 | requires_grad=False, |
| 193 | ) |
| 194 | set_weight_attrs( |
| 195 | act_global_scale, {"ignore_warning": True} | extra_weight_attrs |
| 196 | ) |
| 197 | layer.register_parameter("act_global_scale", act_global_scale) |
| 198 | |
| 199 | forward_hadamard_matrix = Parameter( |
| 200 | torch.empty( |
| 201 | self.quant_config.hadamard_group_size, |
| 202 | self.quant_config.hadamard_group_size, |
| 203 | dtype=params_dtype, |
| 204 | ), |
| 205 | requires_grad=False, |
| 206 | ) |
| 207 | set_weight_attrs( |
| 208 | forward_hadamard_matrix, {"ignore_warning": True} | extra_weight_attrs |
| 209 | ) |
| 210 | layer.register_parameter("forward_hadamard_matrix", forward_hadamard_matrix) |
| 211 | |
| 212 | def apply( |
| 213 | self, |
| 214 | layer: torch.nn.Module, |
| 215 | x: torch.Tensor, |
| 216 | bias: torch.Tensor | None = None, |
| 217 | ) -> torch.Tensor: |
| 218 | return quantized_forward( |
| 219 | x, |
| 220 | layer.qweight, |
| 221 | layer.scales, |
| 222 | layer.weight_global_scale, |
| 223 | layer.act_global_scale, |
| 224 | bias, |
| 225 | layer.forward_hadamard_matrix, |
| 226 | self.quant_config.forward_method, |
| 227 | self.quant_config.forward_dtype, |
| 228 | ) |
| 229 | |
| 230 | |
| 231 | def fused_quantize_mx( |
| 232 | x_flat: torch.Tensor, hadamard_matrix: torch.Tensor, forward_method: str |
| 233 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 234 | return fusedQuantizeMx( |
| 235 | x_flat, |
| 236 | hadamard_matrix, |
| 237 | method=cast(Literal["quest", "abs_max"], forward_method), |
| 238 | ) |
| 239 | |
| 240 | |
| 241 | def fused_quantize_mx_fake(x_flat, hadamard_matrix, forward_method): |
| 242 | rows, cols = x_flat.size(0), x_flat.size(1) // 32 |
| 243 | padded_rows = ((rows + 128 - 1) // 128) * 128 |
| 244 | padded_cols = ((cols + 4 - 1) // 4) * 4 |
| 245 | |
| 246 | xh_e2m1 = torch.empty( |
| 247 | x_flat.size(0), x_flat.size(1) // 2, dtype=torch.uint8, device=x_flat.device |
| 248 | ) |
| 249 | xh_e8m0 = torch.empty( |
| 250 | padded_rows, padded_cols, dtype=torch.float8_e8m0fnu, device=x_flat.device |
| 251 | ) |
| 252 | |
| 253 | return xh_e2m1, xh_e8m0 |
| 254 | |
| 255 | |
| 256 | direct_register_custom_op( |
| 257 | op_name="fused_quantize_mx", |
| 258 | op_func=fused_quantize_mx, |
| 259 | mutates_args=[], |
| 260 | fake_impl=fused_quantize_mx_fake, |
| 261 | dispatch_key=current_platform.dispatch_key, |
| 262 | ) |
| 263 | |
| 264 | |
| 265 | def matmul_mxf4_bf16( |
| 266 | x: torch.Tensor, |
| 267 | w: torch.Tensor, |
| 268 | xs: torch.Tensor, |
| 269 | ws: torch.Tensor, |
| 270 | alpha: torch.Tensor, |
| 271 | ) -> torch.Tensor: |
| 272 | return matmul_mxf4_bf16_tn( |
| 273 | x, |
| 274 | w, |
| 275 | to_blocked(xs, backend="triton").view(torch.float8_e8m0fnu), |
| 276 | to_blocked(ws, backend="triton").view(torch.float8_e8m0fnu), |
| 277 | alpha, |
| 278 | ) |
| 279 | |
| 280 | |
| 281 | def matmul_mxf4_bf16_fake(x, w, xs, ws, alpha): |
| 282 | return torch.empty(*x.shape[:-1], w.shape[0], dtype=torch.bfloat16, device=x.device) |
| 283 | |
| 284 | |
| 285 | direct_register_custom_op( |
| 286 | op_name="matmul_mxf4_bf16", |
| 287 | op_func=matmul_mxf4_bf16, |
| 288 | mutates_args=[], |
| 289 | fake_impl=matmul_mxf4_bf16_fake, |
| 290 | dispatch_key=current_platform.dispatch_key, |
| 291 | ) |
| 292 | |
| 293 | |
| 294 | def fused_quantize_nv( |
| 295 | x_flat: torch.Tensor, hadamard_matrix: torch.Tensor, global_scale: torch.Tensor |
| 296 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 297 | return fusedQuantizeNv(x_flat, hadamard_matrix, global_scale) |
| 298 | |
| 299 | |
| 300 | def fused_quantize_nv_fake(x_flat, hadamard_matrix, global_scale): |
| 301 | rows, cols = x_flat.size(0), x_flat.size(1) // 16 |
| 302 | padded_rows = ((rows + 128 - 1) // 128) * 128 |
| 303 | padded_cols = ((cols + 4 - 1) // 4) * 4 |
| 304 | |
| 305 | xh_e2m1 = torch.empty( |
| 306 | x_flat.size(0), x_flat.size(1) // 2, dtype=torch.uint8, device=x_flat.device |
| 307 | ) |
| 308 | xh_e8m0 = torch.empty( |
| 309 | padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=x_flat.device |
| 310 | ) |
| 311 | |
| 312 | return xh_e2m1, xh_e8m0 |
| 313 | |
| 314 | |
| 315 | direct_register_custom_op( |
| 316 | op_name="fused_quantize_nv", |
| 317 | op_func=fused_quantize_nv, |
| 318 | mutates_args=[], |
| 319 | fake_impl=fused_quantize_nv_fake, |
| 320 | dispatch_key=current_platform.dispatch_key, |
| 321 | ) |
| 322 | |
| 323 | |
| 324 | def matmul_nvf4_bf16( |
| 325 | x: torch.Tensor, |
| 326 | w: torch.Tensor, |
| 327 | xs: torch.Tensor, |
| 328 | ws: torch.Tensor, |
| 329 | alpha: torch.Tensor, |
| 330 | ) -> torch.Tensor: |
| 331 | return cutlass_scaled_fp4_mm( |
| 332 | x, |
| 333 | w, |
| 334 | to_blocked(xs, backend="triton") |
| 335 | .view(torch.float8_e4m3fn) |
| 336 | .view(-1, x.shape[1] // 8), # *2//16 |
| 337 | to_blocked(ws, backend="triton") |
| 338 | .view(torch.float8_e4m3fn) |
| 339 | .view(-1, x.shape[1] // 8), |
| 340 | alpha, |
| 341 | torch.bfloat16, |
| 342 | ) |
| 343 | |
| 344 | |
| 345 | def matmul_nvf4_bf16_fake(x, w, xs, ws, alpha): |
| 346 | return torch.empty(*x.shape[:-1], w.shape[0], dtype=torch.bfloat16, device=x.device) |
| 347 | |
| 348 | |
| 349 | direct_register_custom_op( |
| 350 | op_name="matmul_nvf4_bf16", |
| 351 | op_func=matmul_nvf4_bf16, |
| 352 | mutates_args=[], |
| 353 | fake_impl=matmul_nvf4_bf16_fake, |
| 354 | dispatch_key=current_platform.dispatch_key, |
| 355 | ) |
| 356 | |
| 357 | |
| 358 | def quantized_forward( |
| 359 | x: torch.Tensor, |
| 360 | qweight: torch.Tensor, |
| 361 | weight_scales: torch.Tensor, |
| 362 | weight_global_scale: torch.Tensor, |
| 363 | act_global_scale: torch.Tensor, |
| 364 | bias: torch.Tensor | None, |
| 365 | forward_hadamard_matrix: torch.Tensor, |
| 366 | forward_method: str, |
| 367 | forward_dtype: str, |
| 368 | ) -> torch.Tensor: |
| 369 | x_flat = x.contiguous().flatten(end_dim=-2) |
| 370 | |
| 371 | if forward_dtype == "mxfp4": |
| 372 | x_flat_q, x_flat_scales = torch.ops.vllm.fused_quantize_mx( |
| 373 | x_flat, forward_hadamard_matrix, forward_method |
| 374 | ) |
| 375 | y = torch.ops.vllm.matmul_mxf4_bf16( |
| 376 | x_flat_q, |
| 377 | qweight, |
| 378 | x_flat_scales, |
| 379 | weight_scales, |
| 380 | 1 / (weight_global_scale * act_global_scale), |
| 381 | ) |
| 382 | elif forward_dtype == "nvfp4": |
| 383 | x_flat_q, x_flat_scales = torch.ops.vllm.fused_quantize_nv( |
| 384 | x_flat, forward_hadamard_matrix, act_global_scale |
| 385 | ) |
| 386 | y = torch.ops.vllm.matmul_nvf4_bf16( |
| 387 | x_flat_q, |
| 388 | qweight, |
| 389 | x_flat_scales, |
| 390 | weight_scales, |
| 391 | 1 / (weight_global_scale * act_global_scale), |
| 392 | ) |
| 393 | else: |
| 394 | raise ValueError(f"Unsupported forward_dtype: {forward_dtype}") |
| 395 | |
| 396 | y = y.view(*x.shape[:-1], y.shape[-1]) |
| 397 | if bias is not None: |
| 398 | y += bias |
| 399 | |
| 400 | return y |