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, 11:00 UTC
[Misc] Add compressed-tensors NVFP4A16 emulation support (#17914)
Signed-off-by: Dipika Sikka <dipikasikka1@gmail.com> Signed-off-by: Dipika <dipikasikka1@gmail.com>
Dipika Sikka committed
May 11, 2025 at 03:58 UTC
cd3edfc9082993d003a446d50624097bcf316365
5 files changed
+215
-7
tests/quantization/test_compressed_tensors.py
+23
-3
@@ -13,9 +13,9 @@ from compressed_tensors.quantization import QuantizationType
13
from tests.models.utils import check_logprobs_close
14
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
15
CompressedTensors24, CompressedTensorsLinearMethod,
16
- CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
17
- CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
18
- CompressedTensorsWNA16)
16
+ CompressedTensorsW4A16Fp4, CompressedTensorsW4A16Sparse24,
17
+ CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8,
18
+ CompressedTensorsW8A16Fp8, CompressedTensorsWNA16)
19
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20
sparse_cutlass_supported)
21
from vllm.platforms import current_platform
@@ -648,3 +648,23 @@ def test_compressed_tensors_2of4_sparse_compressed(vllm_runner, args_2of4):
648
output = llm.generate_greedy("Hello my name is", max_tokens=20)
649
print(output)
650
assert output
651
+
652
+
653
+def test_compressed_tensors_nvfp4a16(vllm_runner):
654
+ # run weight only example
655
+ model = "nm-testing/TinyLlama-1.1B-Chat-v1.0-FP4"
656
+ with vllm_runner(model, enforce_eager=True) as llm:
657
+
658
+ def check_model(model):
659
+ layer = model.model.layers[0]
660
+
661
+ qkv_proj = layer.self_attn.qkv_proj
662
+ assert isinstance(qkv_proj.quant_method,
663
+ CompressedTensorsLinearMethod)
664
+ assert isinstance(qkv_proj.scheme, CompressedTensorsW4A16Fp4)
665
+ assert qkv_proj.scheme.group_size == 16
666
+
667
+ llm.apply_model(check_model)
668
+ output = llm.generate_greedy("Hello my name is", max_tokens=20)
669
+ print(output)
670
+ assert output
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+22
-3
@@ -23,9 +23,10 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
23
CompressedTensorsMoEMethod)
24
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
25
W4A16SPARSE24_SUPPORTED_BITS, WNA16_SUPPORTED_BITS, CompressedTensors24,
26
- CompressedTensorsScheme, CompressedTensorsW4A16Sparse24,
27
- CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8,
28
- CompressedTensorsW8A16Fp8, CompressedTensorsWNA16)
26
+ CompressedTensorsScheme, CompressedTensorsW4A16Fp4,
27
+ CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
28
+ CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
29
+ CompressedTensorsWNA16)
30
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
31
find_matched_target, is_activation_quantization_format,
32
should_ignore_layer)
@@ -216,6 +217,21 @@ class CompressedTensorsConfig(QuantizationConfig):
217
else:
218
return False
219
220
+ def _is_fp4a16_nvfp4(self, weight_quant: BaseModel,
221
+ input_quant: BaseModel):
222
+
223
+ is_weight_only = weight_quant is not None and input_quant is None
224
+ is_group_quant = (
225
+ weight_quant.strategy == QuantizationStrategy.GROUP.value)
226
+ is_symmetric = weight_quant.symmetric
227
+
228
+ is_group_size_16 = weight_quant.group_size == 16
229
+ is_float_type = weight_quant.type == QuantizationType.FLOAT
230
+ is_4_bits = weight_quant.num_bits == 4
231
+
232
+ return (is_weight_only and is_group_quant and is_float_type
233
+ and is_4_bits and is_group_size_16 and is_symmetric)
234
+
235
def _is_static_tensor_w8a8(self, weight_quant: BaseModel,
236
input_quant: BaseModel) -> bool:
237
is_8_bits = weight_quant.num_bits == input_quant.num_bits == 8
@@ -315,6 +331,9 @@ class CompressedTensorsConfig(QuantizationConfig):
331
input_quant: BaseModel) -> "CompressedTensorsScheme":
332
333
# Detect If Mixed Precision
334
+ if self._is_fp4a16_nvfp4(weight_quant, input_quant):
335
+ return CompressedTensorsW4A16Fp4()
336
+
337
if self._is_wNa16_group_channel(weight_quant, input_quant):
338
if (self.quant_format == CompressionFormat.marlin_24.value
339
and weight_quant.num_bits in W4A16SPARSE24_SUPPORTED_BITS):
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
+2
-1
@@ -3,6 +3,7 @@
3
from .compressed_tensors_scheme import CompressedTensorsScheme
4
from .compressed_tensors_w4a16_24 import (W4A16SPARSE24_SUPPORTED_BITS,
5
CompressedTensorsW4A16Sparse24)
6
+from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4
7
from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
8
from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
9
from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8
@@ -16,5 +17,5 @@ __all__ = [
17
"CompressedTensorsW8A16Fp8", "CompressedTensorsW4A16Sparse24",
18
"CompressedTensorsW8A8Int8", "CompressedTensorsW8A8Fp8",
19
"WNA16_SUPPORTED_BITS", "W4A16SPARSE24_SUPPORTED_BITS",
19
- "CompressedTensors24"
20
+ "CompressedTensors24", "CompressedTensorsW4A16Fp4"
21
]
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_nvfp4.py
new
+107
@@ -0,0 +1,107 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Callable, List, Optional
3
+
4
+import torch
5
+import torch.nn.functional as F
6
+from torch.nn.parameter import Parameter
7
+
8
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
9
+ CompressedTensorsScheme)
10
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( # noqa: E501
11
+ dequantize_to_dtype)
12
+from vllm.model_executor.parameter import (GroupQuantScaleParameter,
13
+ ModelWeightParameter,
14
+ PerTensorScaleParameter)
15
+
16
+__all__ = ["CompressedTensorsW4A16Fp4"]
17
+
18
+
19
+class CompressedTensorsW4A16Fp4(CompressedTensorsScheme):
20
+
21
+ def __init__(self):
22
+ self.group_size = 16
23
+
24
+ @classmethod
25
+ def get_min_capability(cls) -> int:
26
+ # dont restrict as emulations
27
+ return 80
28
+
29
+ def create_weights(self, layer: torch.nn.Module,
30
+ output_partition_sizes: List[int],
31
+ input_size_per_partition: int,
32
+ params_dtype: torch.dtype, weight_loader: Callable,
33
+ **kwargs):
34
+
35
+ # Weight
36
+ weight = ModelWeightParameter(data=torch.empty(
37
+ sum(output_partition_sizes),
38
+ input_size_per_partition // 2,
39
+ dtype=torch.uint8),
40
+ input_dim=1,
41
+ output_dim=0,
42
+ weight_loader=weight_loader)
43
+ layer.register_parameter("weight_packed", weight)
44
+
45
+ # Global Weight Scale
46
+ weight_global_scale = PerTensorScaleParameter(
47
+ data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
48
+ weight_loader=weight_loader)
49
+ layer.register_parameter("weight_global_scale", weight_global_scale)
50
+
51
+ # Per Group Weight Scale
52
+ weight_scale = GroupQuantScaleParameter(data=torch.empty(
53
+ sum(output_partition_sizes),
54
+ input_size_per_partition // self.group_size,
55
+ dtype=torch.float8_e4m3fn,
56
+ ),
57
+ input_dim=1,
58
+ output_dim=0,
59
+ weight_loader=weight_loader)
60
+
61
+ layer.register_parameter("weight_scale", weight_scale)
62
+
63
+ def swizzle_blockscale(self, scale: torch.tensor):
64
+ assert (scale.dtype == torch.float8_e4m3fn)
65
+ # Pad and blockwise interleave weight_scale
66
+ scale_ndim = scale.ndim
67
+ if scale.ndim == 2:
68
+ scale = scale.unsqueeze(0)
69
+ assert scale.ndim == 3
70
+ B, M, K = scale.shape
71
+ round_up_multiple = lambda x, m: (x + m - 1) // m * m
72
+ M_padded = round_up_multiple(M, 128)
73
+ K_padded = round_up_multiple(K, 4)
74
+ padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
75
+ padded_scale[:B, :M, :K] = scale
76
+ batches, rows, cols = padded_scale.shape
77
+ assert rows % 128 == 0
78
+ assert cols % 4 == 0
79
+ padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
80
+ cols // 4, 4)
81
+ swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
82
+ swizzled_scale = swizzled_scale.contiguous().cuda()
83
+ return (swizzled_scale.reshape(M, K)
84
+ if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
85
+
86
+ def process_weights_after_loading(self, layer) -> None:
87
+ layer.weight_global_scale = Parameter(
88
+ layer.weight_global_scale.max().to(torch.float32),
89
+ requires_grad=False)
90
+ # Note: a post weight loading step but not required for the emulation
91
+ swizzled_weight_scale = self.swizzle_blockscale(layer.weight_scale)
92
+ layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
93
+ requires_grad=False)
94
+
95
+ def apply_weights(self,
96
+ layer: torch.nn.Module,
97
+ x: torch.Tensor,
98
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
99
+
100
+ w_fp4 = layer.weight_packed.data
101
+ w_global_scale = layer.weight_global_scale
102
+ w_blockscale = layer.weight_scale_swizzled.data
103
+ w_dq = dequantize_to_dtype(w_fp4, w_blockscale, w_global_scale,
104
+ x.dtype, x.device, self.group_size)
105
+ out = F.linear(x, w_dq)
106
+ del w_dq, w_fp4, w_global_scale, w_blockscale
107
+ return out
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
new
+61
@@ -0,0 +1,61 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import torch
3
+
4
+__all__ = [
5
+ "break_fp4_bytes",
6
+ "dequantize_to_dtype",
7
+]
8
+
9
+kE2M1ToFloat = torch.tensor([0., 0.5, 1., 1.5, 2., 3., 4., 6.],
10
+ dtype=torch.float32)
11
+
12
+
13
+def break_fp4_bytes(a, dtype):
14
+ assert a.dtype == torch.uint8
15
+ m, n = a.shape
16
+ # Vectorized nibble processing
17
+ a_flat = a.flatten()
18
+ high = (a_flat & 0xF0) >> 4 # Upper nibbles
19
+ low = a_flat & 0x0F # Lower nibbles
20
+ # Combine nibbles for batch processing
21
+ combined = torch.stack((low, high), dim=1).flatten()
22
+ # Vectorized sign and magnitude extraction
23
+ signs = (combined & 0x08).to(torch.bool) # Sign bits
24
+ abs_vals = (combined & 0x07).to(torch.long)
25
+ # Device-aware lookup and sign application
26
+ kE2M1 = kE2M1ToFloat.to(device=a.device)
27
+ values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)
28
+ # Reshape to final form
29
+ return values.reshape(m, n * 2).to(dtype=dtype)
30
+
31
+
32
+def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
33
+ m_tiles = (m + 128 - 1) // 128
34
+ f = block_size * 4
35
+ k_tiles = (k + f - 1) // f
36
+ tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
37
+ tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
38
+ out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
39
+ return out[0:m, 0:k]
40
+
41
+
42
+def dequantize_to_dtype(tensor_fp4,
43
+ tensor_sf,
44
+ global_scale,
45
+ dtype,
46
+ device,
47
+ block_size=16):
48
+ """Dequantize the fp4 tensor back to high precision."""
49
+ # Two fp4 values are packed into one uint8.
50
+ assert tensor_fp4.dtype == torch.uint8
51
+ m, packed_k = tensor_fp4.shape
52
+ k = packed_k * 2
53
+ tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32)
54
+ tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
55
+ tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
56
+ tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
57
+ tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
58
+
59
+ # scale the tensor
60
+ out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
61
+ return out.to(dtype)