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 15, 06:00 UTC
[Quantization] Add compressed-tensors NVFP4 support (#18312)
Dipika Sikka committed
Jun 8, 2025 at 06:05 UTC
c123bc33f9c139b293e906d9612e247c8e29dbbf
6 files changed
+267
-16
tests/quantization/test_compressed_tensors.py
+12
-7
@@ -14,9 +14,10 @@ from compressed_tensors.quantization import QuantizationType
14
from tests.models.utils import check_logprobs_close
15
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
16
CompressedTensors24, CompressedTensorsLinearMethod,
17
- CompressedTensorsW4A16Fp4, CompressedTensorsW4A16Sparse24,
18
- CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8,
19
- CompressedTensorsW8A16Fp8, CompressedTensorsWNA16)
17
+ CompressedTensorsW4A4Fp4, CompressedTensorsW4A16Fp4,
18
+ CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
19
+ CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
20
+ CompressedTensorsWNA16)
21
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
22
sparse_cutlass_supported)
23
from vllm.platforms import current_platform
@@ -651,9 +652,13 @@ def test_compressed_tensors_2of4_sparse_compressed(vllm_runner, args_2of4):
652
assert output
653
654
654
-def test_compressed_tensors_nvfp4a16(vllm_runner):
655
- # run weight only example
656
- model = "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16"
655
+@pytest.mark.parametrize(
656
+ "args",
657
+ [("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16",
658
+ CompressedTensorsW4A16Fp4),
659
+ ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4", CompressedTensorsW4A4Fp4)])
660
+def test_compressed_tensors_nvfp4(vllm_runner, args):
661
+ model, scheme = args
662
with vllm_runner(model, enforce_eager=True) as llm:
663
664
def check_model(model):
@@ -662,7 +667,7 @@ def test_compressed_tensors_nvfp4a16(vllm_runner):
667
qkv_proj = layer.self_attn.qkv_proj
668
assert isinstance(qkv_proj.quant_method,
669
CompressedTensorsLinearMethod)
665
- assert isinstance(qkv_proj.scheme, CompressedTensorsW4A16Fp4)
670
+ assert isinstance(qkv_proj.scheme, scheme)
671
assert qkv_proj.scheme.group_size == 16
672
673
llm.apply_model(check_model)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+27
-4
@@ -24,10 +24,10 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
24
CompressedTensorsMoEMethod)
25
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
26
W4A16SPARSE24_SUPPORTED_BITS, WNA16_SUPPORTED_BITS, CompressedTensors24,
27
- CompressedTensorsScheme, CompressedTensorsW4A16Fp4,
28
- CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
29
- CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
30
- CompressedTensorsWNA16)
27
+ CompressedTensorsScheme, CompressedTensorsW4A4Fp4,
28
+ CompressedTensorsW4A16Fp4, CompressedTensorsW4A16Sparse24,
29
+ CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8,
30
+ CompressedTensorsW8A16Fp8, CompressedTensorsWNA16)
31
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
32
find_matched_target, is_activation_quantization_format,
33
should_ignore_layer)
@@ -218,6 +218,26 @@ class CompressedTensorsConfig(QuantizationConfig):
218
else:
219
return False
220
221
+ def _is_fp4a4_nvfp4(self, weight_quant: BaseModel, input_quant: BaseModel):
222
+
223
+ if weight_quant is None or input_quant is None:
224
+ return False
225
+
226
+ is_tensor_group_quant = (weight_quant.strategy
227
+ == QuantizationStrategy.TENSOR_GROUP.value
228
+ and input_quant.strategy
229
+ == QuantizationStrategy.TENSOR_GROUP.value)
230
+ is_symmetric = weight_quant.symmetric and input_quant.symmetric
231
+
232
+ is_group_size_16 = (weight_quant.group_size == 16
233
+ and input_quant.group_size == 16)
234
+ is_float_type = (weight_quant.type == QuantizationType.FLOAT
235
+ and input_quant.type == QuantizationType.FLOAT.value)
236
+ is_4_bits = weight_quant.num_bits == 4 and input_quant.num_bits == 4
237
+
238
+ return (is_tensor_group_quant and is_float_type and is_4_bits
239
+ and is_group_size_16 and is_symmetric)
240
+
241
def _is_fp4a16_nvfp4(self, weight_quant: BaseModel,
242
input_quant: BaseModel):
243
@@ -353,6 +373,9 @@ class CompressedTensorsConfig(QuantizationConfig):
373
actorder=weight_quant.actorder)
374
375
if is_activation_quantization_format(self.quant_format):
376
+ if self._is_fp4a4_nvfp4(weight_quant, input_quant):
377
+ return CompressedTensorsW4A4Fp4()
378
+
379
if self._is_fp8_w8a8(weight_quant, input_quant):
380
is_fp8_w8a8_supported = self._check_scheme_supported(
381
CompressedTensorsW8A8Fp8.get_min_capability(), error=False)
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
+3
-1
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from .compressed_tensors_scheme import CompressedTensorsScheme
5
+from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
6
from .compressed_tensors_w4a16_24 import (W4A16SPARSE24_SUPPORTED_BITS,
7
CompressedTensorsW4A16Sparse24)
8
from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4
@@ -18,5 +19,6 @@ __all__ = [
19
"CompressedTensorsW8A16Fp8", "CompressedTensorsW4A16Sparse24",
20
"CompressedTensorsW8A8Int8", "CompressedTensorsW8A8Fp8",
21
"WNA16_SUPPORTED_BITS", "W4A16SPARSE24_SUPPORTED_BITS",
21
- "CompressedTensors24", "CompressedTensorsW4A16Fp4"
22
+ "CompressedTensors24", "CompressedTensorsW4A16Fp4",
23
+ "CompressedTensorsW4A4Fp4"
24
]
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
new
+178
@@ -0,0 +1,178 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Callable, Optional
3
+
4
+import torch
5
+from torch.nn.parameter import Parameter
6
+
7
+from vllm._custom_ops import (cutlass_scaled_fp4_mm,
8
+ cutlass_scaled_mm_supports_fp4, scaled_fp4_quant)
9
+from vllm.logger import init_logger
10
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11
+ CompressedTensorsScheme)
12
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( # noqa: E501
13
+ dequantize_to_dtype, ref_nvfp4_quant)
14
+from vllm.model_executor.parameter import (GroupQuantScaleParameter,
15
+ ModelWeightParameter,
16
+ PerTensorScaleParameter)
17
+from vllm.platforms import current_platform
18
+
19
+logger = init_logger(__name__)
20
+
21
+__all__ = ["CompressedTensorsW4A4Fp4"]
22
+
23
+
24
+def cutlass_fp4_supported() -> bool:
25
+ if not current_platform.is_cuda():
26
+ return False
27
+ capability_tuple = current_platform.get_device_capability()
28
+ capability = -1 if capability_tuple is None else capability_tuple.to_int()
29
+ return cutlass_scaled_mm_supports_fp4(capability)
30
+
31
+
32
+class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
33
+
34
+ def __init__(self):
35
+ self.group_size = 16
36
+ self.cutlass_nvfp4_supported = cutlass_fp4_supported()
37
+ if not self.cutlass_nvfp4_supported:
38
+ logger.warning("Current platform does not support cutlass NVFP4."
39
+ " Running emulations.")
40
+
41
+ @classmethod
42
+ def get_min_capability(cls) -> int:
43
+ # dont restrict as emulations
44
+ return 80
45
+
46
+ def run_nvfp4_emulations(self, x: torch.Tensor, layer):
47
+ x_m, x_k = x.shape
48
+ output_dtype = x.dtype
49
+
50
+ # quantize input to (FP4 and interleaved block scale)
51
+ x_fp4, x_blockscale = ref_nvfp4_quant(x, layer.input_global_scale,
52
+ self.group_size)
53
+
54
+ # dequantize input
55
+ x_fp4 = x_fp4.reshape(x_m, x_k // self.group_size, self.group_size)
56
+ x_blockscale = x_blockscale.unsqueeze(-1) / layer.input_global_scale
57
+ x_dq = (x_fp4 * x_blockscale).reshape(x_m, x_k).to(output_dtype)
58
+ del x_fp4, x_blockscale
59
+
60
+ # dequantize weight
61
+ w_fp4 = layer.weight.data.view(torch.uint8)
62
+ w_blockscale = layer.weight_scale_swizzled.data
63
+ w_global_scale = layer.weight_global_scale
64
+ w_dq = dequantize_to_dtype(w_fp4, w_blockscale, w_global_scale,
65
+ output_dtype, x.device, self.group_size)
66
+
67
+ # matmul
68
+ out = torch.matmul(x_dq, w_dq.t())
69
+ del w_dq, x_dq
70
+ return out
71
+
72
+ def create_weights(self, layer: torch.nn.Module,
73
+ output_partition_sizes: list[int],
74
+ input_size_per_partition: int,
75
+ params_dtype: torch.dtype, weight_loader: Callable,
76
+ **kwargs):
77
+ output_size_per_partition = sum(output_partition_sizes)
78
+ layer.logical_widths = output_partition_sizes
79
+ layer.input_size_per_partition = input_size_per_partition
80
+ layer.output_size_per_partition = output_size_per_partition
81
+
82
+ # Weight
83
+ weight = ModelWeightParameter(data=torch.empty(
84
+ sum(output_partition_sizes),
85
+ input_size_per_partition // 2,
86
+ dtype=torch.uint8),
87
+ input_dim=1,
88
+ output_dim=0,
89
+ weight_loader=weight_loader)
90
+ layer.register_parameter("weight_packed", weight)
91
+
92
+ # Global Weight Scale
93
+ weight_global_scale = PerTensorScaleParameter(
94
+ data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
95
+ weight_loader=weight_loader)
96
+ layer.register_parameter("weight_global_scale", weight_global_scale)
97
+
98
+ # Per Group Weight Scale
99
+ weight_scale = GroupQuantScaleParameter(data=torch.empty(
100
+ sum(output_partition_sizes),
101
+ input_size_per_partition // self.group_size,
102
+ dtype=torch.float8_e4m3fn,
103
+ ),
104
+ input_dim=1,
105
+ output_dim=0,
106
+ weight_loader=weight_loader)
107
+
108
+ layer.register_parameter("weight_scale", weight_scale)
109
+
110
+ input_global_scale = PerTensorScaleParameter(
111
+ data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
112
+ weight_loader=weight_loader)
113
+ layer.register_parameter("input_global_scale", input_global_scale)
114
+
115
+ def swizzle_blockscale(self, scale: torch.tensor):
116
+ assert (scale.dtype == torch.float8_e4m3fn)
117
+ # Pad and blockwise interleave weight_scale
118
+ scale_ndim = scale.ndim
119
+ if scale.ndim == 2:
120
+ scale = scale.unsqueeze(0)
121
+ assert scale.ndim == 3
122
+ B, M, K = scale.shape
123
+ round_up_multiple = lambda x, m: (x + m - 1) // m * m
124
+ M_padded = round_up_multiple(M, 128)
125
+ K_padded = round_up_multiple(K, 4)
126
+ padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
127
+ padded_scale[:B, :M, :K] = scale
128
+ batches, rows, cols = padded_scale.shape
129
+ assert rows % 128 == 0
130
+ assert cols % 4 == 0
131
+ padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
132
+ cols // 4, 4)
133
+ swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
134
+ swizzled_scale = swizzled_scale.contiguous().cuda()
135
+ return (swizzled_scale.reshape(M, K)
136
+ if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
137
+
138
+ def process_weights_after_loading(self, layer) -> None:
139
+
140
+ global_input_scale = layer.input_global_scale.max().to(torch.float32)
141
+ layer.input_global_scale = Parameter(global_input_scale,
142
+ requires_grad=False)
143
+
144
+ layer.weight_global_scale = Parameter(
145
+ layer.weight_global_scale.max().to(torch.float32),
146
+ requires_grad=False)
147
+
148
+ swizzled_weight_scale = self.swizzle_blockscale(layer.weight_scale)
149
+ layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
150
+ requires_grad=False)
151
+
152
+ # required by cutlass kernel; need Parameter, not ModelWeightParameter
153
+ layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
154
+
155
+ if self.cutlass_nvfp4_supported:
156
+ layer.alpha = Parameter(layer.input_global_scale *
157
+ layer.weight_global_scale,
158
+ requires_grad=False)
159
+
160
+ def apply_weights(self,
161
+ layer: torch.nn.Module,
162
+ x: torch.Tensor,
163
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
164
+
165
+ if self.cutlass_nvfp4_supported:
166
+ output_dtype = x.dtype
167
+ output_shape = [x.shape[0], layer.weight.shape[0]]
168
+
169
+ # quantize BF16 or FP16 to (FP4 and interleaved block scale)
170
+ x_fp4, x_blockscale = scaled_fp4_quant(x, layer.input_global_scale)
171
+
172
+ out = cutlass_scaled_fp4_mm(x_fp4, layer.weight, x_blockscale,
173
+ layer.weight_scale_swizzled,
174
+ 1 / layer.alpha, output_dtype)
175
+ if bias is not None:
176
+ out = out + bias
177
+ return out.view(*output_shape)
178
+ return self.run_nvfp4_emulations(x, layer)
vllm/model_executor/layers/quantization/compressed_tensors/utils.py
+1
@@ -15,6 +15,7 @@ def is_activation_quantization_format(format: str) -> bool:
15
CompressionFormat.naive_quantized.value,
16
CompressionFormat.int_quantized.value,
17
CompressionFormat.float_quantized.value,
18
+ CompressionFormat.nvfp4_pack_quantized.value
19
]
20
return format in _ACTIVATION_QUANTIZATION_FORMATS
21
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+46
-4
@@ -2,10 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
-__all__ = [
6
- "break_fp4_bytes",
7
- "dequantize_to_dtype",
8
-]
5
+from vllm.scalar_type import scalar_types
6
+
7
+__all__ = ["break_fp4_bytes", "dequantize_to_dtype", "ref_nvfp4_quant"]
8
+
9
+FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
10
11
kE2M1ToFloat = torch.tensor([0., 0.5, 1., 1.5, 2., 3., 4., 6.],
12
dtype=torch.float32)
@@ -60,3 +61,44 @@ def dequantize_to_dtype(tensor_fp4,
61
# scale the tensor
62
out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
63
return out.to(dtype)
64
+
65
+
66
+def get_reciprocal(x):
67
+ if isinstance(x, torch.Tensor):
68
+ return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x)
69
+ elif isinstance(x, (float, int)):
70
+ return 0.0 if x == 0 else 1.0 / x
71
+ else:
72
+ raise TypeError("Input must be a float, int, or a torch.Tensor.")
73
+
74
+
75
+def cast_to_fp4(x):
76
+ sign = torch.sign(x)
77
+ x = torch.abs(x)
78
+ x[(x >= 0.0) & (x <= 0.25)] = 0.0
79
+ x[(x > 0.25) & (x < 0.75)] = 0.5
80
+ x[(x >= 0.75) & (x <= 1.25)] = 1.0
81
+ x[(x > 1.25) & (x < 1.75)] = 1.5
82
+ x[(x >= 1.75) & (x <= 2.5)] = 2.0
83
+ x[(x > 2.5) & (x < 3.5)] = 3.0
84
+ x[(x >= 3.5) & (x <= 5.0)] = 4.0
85
+ x[x > 5.0] = 6.0
86
+ return x * sign
87
+
88
+
89
+def ref_nvfp4_quant(x, global_scale, block_size):
90
+ assert global_scale.dtype == torch.float32
91
+ assert x.ndim == 2
92
+ m, n = x.shape
93
+ x = torch.reshape(x, (m, n // block_size, block_size))
94
+ vec_max = torch.max(torch.abs(x), dim=-1,
95
+ keepdim=True)[0].to(torch.float32)
96
+ scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX))
97
+ scale = torch.clamp(scale, max=448, min=-448)
98
+ scale = scale.to(torch.float8_e4m3fn).to(torch.float32)
99
+ output_scale = get_reciprocal(scale * get_reciprocal(global_scale))
100
+
101
+ scaled_x = x.to(torch.float32) * output_scale
102
+ clipped_x = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n)
103
+ # both outputs are float32
104
+ return cast_to_fp4(clipped_x), scale.squeeze(-1)