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, 14:45 UTC
[Misc] Update `compressed-tensors` WNA16 to support zero-points (#14211)
Dipika Sikka committed
Apr 15, 2025 at 09:33 UTC
54a66e5fee4a1ea62f1e4c79a078b20668e408c6
6 files changed
+85
-45
tests/quantization/test_compressed_tensors.py
+15
-6
@@ -261,16 +261,23 @@ def test_compressed_tensors_w8a8_dynamic_per_token(
261
262
@pytest.mark.parametrize(
263
"wNa16_args",
264
- [
265
- ("nm-testing/tinyllama-oneshot-w4a16-channel-v2", "channel", None, 8),
266
- ("nm-testing/tinyllama-oneshot-w4a16-group128-v2", "group", 128, 8),
267
- ("nm-testing/tinyllama-oneshot-w8a16-per-channel", "channel", None, 4),
268
- ],
264
+ [("nm-testing/tinyllama-oneshot-w4a16-channel-v2", "channel", None, 8,
265
+ True, False),
266
+ ("nm-testing/tinyllama-oneshot-w4a16-group128-v2", "group", 128, 8, True,
267
+ False),
268
+ ("nm-testing/tinyllama-oneshot-w8a16-per-channel", "channel", None, 4,
269
+ True, False),
270
+ ("nm-testing/TinyLlama-1.1B-Chat-v1.0-awq-group128-asym256", "group", 128,
271
+ 8, False, False),
272
+ ("nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-Channel",
273
+ "channel", None, 8, False, False),
274
+ ("nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder",
275
+ "group", 128, 8, False, True)],
276
)
277
@pytest.mark.skipif(not current_platform.is_cuda(),
278
reason="The tests are skipped on non-CUDA platform.")
279
def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
273
- model, strategy, group, pack_factor = wNa16_args
280
+ model, strategy, group, pack_factor, symmetric, has_g_idx = wNa16_args
281
with vllm_runner(model) as llm:
282
283
def check_model(model):
@@ -286,6 +293,8 @@ def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
293
if group is None else group)
294
295
assert qkv_proj.scheme.pack_factor == pack_factor
296
+ assert qkv_proj.scheme.symmetric == symmetric
297
+ assert qkv_proj.scheme.has_g_idx == has_g_idx
298
299
llm.apply_model(check_model)
300
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+3
-3
@@ -302,14 +302,12 @@ class CompressedTensorsConfig(QuantizationConfig):
302
def _is_wNa16_group_channel(self, weight_quant: BaseModel,
303
input_quant: BaseModel) -> bool:
304
input_quant_none = input_quant is None
305
- is_symmetric = weight_quant.symmetric
305
is_channel_group = (
306
weight_quant.strategy == QuantizationStrategy.CHANNEL.value
307
or weight_quant.strategy == QuantizationStrategy.GROUP.value)
308
is_static = not weight_quant.dynamic
309
311
- return (is_channel_group and input_quant_none and is_symmetric
312
- and is_static)
310
+ return (is_channel_group and input_quant_none and is_static)
311
312
def _get_scheme_from_parts(
313
self, weight_quant: BaseModel,
@@ -319,6 +317,7 @@ class CompressedTensorsConfig(QuantizationConfig):
317
if self._is_wNa16_group_channel(weight_quant, input_quant):
318
if (self.quant_format == CompressionFormat.marlin_24.value
319
and weight_quant.num_bits in W4A16SPARSE24_SUPPORTED_BITS):
320
+ assert weight_quant.symmetric
321
return CompressedTensorsW4A16Sparse24(
322
strategy=weight_quant.strategy,
323
num_bits=weight_quant.num_bits,
@@ -328,6 +327,7 @@ class CompressedTensorsConfig(QuantizationConfig):
327
return CompressedTensorsWNA16(
328
num_bits=weight_quant.num_bits,
329
strategy=weight_quant.strategy,
330
+ symmetric=weight_quant.symmetric,
331
group_size=weight_quant.group_size,
332
actorder=weight_quant.actorder)
333
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
+39
-3
@@ -12,11 +12,15 @@ from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
12
MPLinearLayerConfig, choose_mp_linear_kernel)
13
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
14
marlin_repeat_scales_on_all_ranks)
15
+# yapf conflicts with isort for this block
16
+# yapf: disable
17
from vllm.model_executor.parameter import (BasevLLMParameter,
18
ChannelQuantScaleParameter,
19
GroupQuantScaleParameter,
20
+ PackedColumnParameter,
21
PackedvLLMParameter,
22
RowvLLMParameter)
23
+# yapf: enable
24
from vllm.scalar_type import scalar_types
25
26
logger = init_logger(__name__)
@@ -26,6 +30,7 @@ WNA16_SUPPORTED_TYPES_MAP = {
30
4: scalar_types.uint4b8,
31
8: scalar_types.uint8b128
32
}
33
+WNA16_ZP_SUPPORTED_TYPES_MAP = {4: scalar_types.uint4, 8: scalar_types.uint8}
34
WNA16_SUPPORTED_BITS = list(WNA16_SUPPORTED_TYPES_MAP.keys())
35
36
@@ -36,10 +41,12 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
41
strategy: str,
42
num_bits: int,
43
group_size: Optional[int] = None,
44
+ symmetric: Optional[bool] = True,
45
actorder: Optional[ActivationOrdering] = None):
46
47
self.pack_factor = 32 // num_bits
48
self.strategy = strategy
49
+ self.symmetric = symmetric
50
self.group_size = -1 if group_size is None else group_size
51
self.has_g_idx = actorder == ActivationOrdering.GROUP
52
@@ -53,7 +60,9 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
60
f"Unsupported num_bits = {num_bits}. "
61
f"Supported num_bits = {WNA16_SUPPORTED_TYPES_MAP.keys()}")
62
56
- self.quant_type = WNA16_SUPPORTED_TYPES_MAP[num_bits]
63
+ self.quant_type = (WNA16_ZP_SUPPORTED_TYPES_MAP[num_bits]
64
+ if not self.symmetric else
65
+ WNA16_SUPPORTED_TYPES_MAP[num_bits])
66
67
@classmethod
68
def get_min_capability(cls) -> int:
@@ -75,7 +84,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
84
weight_type=self.quant_type,
85
act_type=params_dtype,
86
group_size=self.group_size,
78
- zero_points=False,
87
+ zero_points=not self.symmetric,
88
has_g_idx=self.has_g_idx
89
)
90
@@ -120,13 +129,37 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
129
dtype=params_dtype,
130
)
131
}
132
+
133
+ zeros_args = {
134
+ "weight_loader":
135
+ weight_loader,
136
+ "data":
137
+ torch.zeros(
138
+ output_size_per_partition // self.pack_factor,
139
+ scales_and_zp_size,
140
+ dtype=torch.int32,
141
+ )
142
+ }
143
+
144
if not partition_scales:
145
weight_scale = ChannelQuantScaleParameter(output_dim=0,
146
**weight_scale_args)
147
+
148
+ if not self.symmetric:
149
+ qzeros = PackedColumnParameter(output_dim=0,
150
+ packed_dim=0,
151
+ packed_factor=self.pack_factor,
152
+ **zeros_args)
153
else:
154
weight_scale = GroupQuantScaleParameter(output_dim=0,
155
input_dim=1,
156
**weight_scale_args)
157
+ if not self.symmetric:
158
+ qzeros = PackedvLLMParameter(input_dim=1,
159
+ output_dim=0,
160
+ packed_dim=0,
161
+ packed_factor=self.pack_factor,
162
+ **zeros_args)
163
164
# A 2D array defining the original shape of the weights
165
# before packing
@@ -138,6 +171,9 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
171
layer.register_parameter("weight_scale", weight_scale)
172
layer.register_parameter("weight_shape", weight_shape)
173
174
+ if not self.symmetric:
175
+ layer.register_parameter("weight_zero_point", qzeros)
176
+
177
# group index (for activation reordering)
178
if self.has_g_idx:
179
weight_g_idx = RowvLLMParameter(data=torch.empty(
@@ -151,7 +187,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
187
self.kernel = kernel_type(mp_linear_kernel_config,
188
w_q_param_name="weight_packed",
189
w_s_param_name="weight_scale",
154
- w_zp_param_name=None,
190
+ w_zp_param_name="weight_zero_point",
191
w_gidx_param_name="weight_g_idx")
192
193
# Checkpoints are serialized in compressed-tensors format, which is
vllm/model_executor/layers/quantization/kernels/mixed_precision/machete.py
+2
-5
@@ -26,17 +26,14 @@ class MacheteLinearKernel(MPLinearKernel):
26
@classmethod
27
def can_implement(cls,
28
c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:
29
+
30
if c.has_g_idx and\
31
c.partition_weight_shape[0] != c.full_weight_shape[0]:
32
return False, "Act reordering currently not supported by Machete, "\
33
"when the input features are partitioned across "\
34
"devices"
34
-
35
if c.zero_points:
36
- return False, "Zero points currently not supported by "\
37
- " Compressed Tensors + Machete. (Kernel supports it"\
38
- " but CompressedTensorsWNA16 does not so support has"\
39
- " not been added to MacheteWNA16Kernel yet"
36
+ return False, "Zero points currently not supported by Machete"
37
38
if c.weight_type not in query_machete_supported_quant_types(
39
c.zero_points):
vllm/model_executor/layers/quantization/kernels/mixed_precision/marlin.py
+24
-27
@@ -9,7 +9,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
9
MARLIN_SUPPORTED_GROUP_SIZES, apply_gptq_marlin_linear,
10
check_marlin_supports_shape, marlin_is_k_full, marlin_make_empty_g_idx,
11
marlin_make_workspace, marlin_permute_scales, marlin_sort_g_idx,
12
- query_marlin_supported_quant_types)
12
+ marlin_zero_points, query_marlin_supported_quant_types, unpack_cols)
13
from vllm.model_executor.parameter import (BasevLLMParameter,
14
permute_param_layout_)
15
@@ -25,10 +25,6 @@ class MarlinLinearKernel(MPLinearKernel):
25
@classmethod
26
def can_implement(cls,
27
c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:
28
- if c.zero_points:
29
- return False, "Zero points currently not supported by "\
30
- " MarlinLinearKernel. Will be added when AWQMarlin "\
31
- "is migrated over to using MPLinearKernel backend"
28
29
quant_types = query_marlin_supported_quant_types(c.zero_points)
30
if c.weight_type not in quant_types:
@@ -67,28 +63,6 @@ class MarlinLinearKernel(MPLinearKernel):
63
if self.w_zp_name is None:
64
self.w_zp_name = "w_zp"
65
70
- if c.has_g_idx:
71
- g_idx, g_idx_sort_indices = marlin_sort_g_idx(
72
- getattr(layer, self.w_gidx_name))
73
- self._transform_param(layer, self.w_gidx_name, lambda _: g_idx)
74
- layer.g_idx_sort_indices = g_idx_sort_indices
75
- else:
76
- setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device))
77
- layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
78
-
79
- if c.zero_points:
80
- pass
81
- # TODO (lucas): add the following when AWQMarlin is migrated over to
82
- # using MPLinearKernel backend
83
- # self._transform_param(layer, self.w_zp_name, lambda x: \
84
- # marlin_zero_points(
85
- # x,
86
- # size_k=c.partition_weight_shape[0],
87
- # size_n=c.partition_weight_shape[1],
88
- # num_bits=c.weight_type.size_bits))
89
- else:
90
- setattr(layer, self.w_zp_name, marlin_make_empty_g_idx(device))
91
-
66
def transform_w_q(x):
67
assert isinstance(x, BasevLLMParameter)
68
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
@@ -108,6 +82,28 @@ class MarlinLinearKernel(MPLinearKernel):
82
group_size=c.group_size)
83
return x
84
85
+ if c.has_g_idx:
86
+ g_idx, g_idx_sort_indices = marlin_sort_g_idx(
87
+ getattr(layer, self.w_gidx_name))
88
+ self._transform_param(layer, self.w_gidx_name, lambda _: g_idx)
89
+ layer.g_idx_sort_indices = g_idx_sort_indices
90
+ else:
91
+ setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device))
92
+ layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
93
+
94
+ if c.zero_points:
95
+ grouped_k = (c.partition_weight_shape[0] //
96
+ c.group_size if c.group_size != -1 else 1)
97
+ self._transform_param(layer, self.w_zp_name, lambda x: \
98
+ marlin_zero_points(
99
+ unpack_cols(x.t(), c.weight_type.size_bits,
100
+ grouped_k,
101
+ c.partition_weight_shape[1]),
102
+ size_k=grouped_k,
103
+ size_n=c.partition_weight_shape[1],
104
+ num_bits=c.weight_type.size_bits))
105
+ else:
106
+ setattr(layer, self.w_zp_name, marlin_make_empty_g_idx(device))
107
self._transform_param(layer, self.w_q_name, transform_w_q)
108
self._transform_param(layer, self.w_s_name, transform_w_s)
109
@@ -131,5 +127,6 @@ class MarlinLinearKernel(MPLinearKernel):
127
wtype=c.weight_type,
128
input_size_per_partition=c.partition_weight_shape[0],
129
output_size_per_partition=c.partition_weight_shape[1],
130
+ has_zp=self.config.zero_points,
131
is_k_full=self.is_k_full,
132
bias=bias)
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+2
-1
@@ -332,6 +332,7 @@ def apply_gptq_marlin_linear(
332
wtype: ScalarType,
333
output_size_per_partition: int,
334
input_size_per_partition: int,
335
+ has_zp: bool,
336
is_k_full: bool,
337
bias: Optional[torch.Tensor] = None,
338
use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT) -> torch.Tensor:
@@ -356,8 +357,8 @@ def apply_gptq_marlin_linear(
357
size_n=output_size_per_partition,
358
size_k=input_size_per_partition,
359
is_k_full=is_k_full,
359
- has_zp=False,
360
use_atomic_add=use_atomic_add,
361
+ has_zp=has_zp,
362
use_fp32_reduce=use_fp32_reduce,
363
is_zp_float=False)
364