Humming support for 2/3/5/6/7-bit pack-quantized weight-only inference (#46389)
Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
HDCharles committed
Jun 24, 2026 at 13:53 UTC
56ca5997eac9a6c69c7af482c18ced4c184f35dc
4 files changed
+40
-22
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+3
-13
@@ -37,7 +37,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
37
CompressedTensorsMoEMethod,
38
)
39
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
40
- WNA16_SUPPORTED_BITS,
40
CompressedTensorsScheme,
41
CompressedTensorsW4A4Fp4,
42
CompressedTensorsW4A4Mxfp4,
@@ -680,14 +679,7 @@ class CompressedTensorsConfig(QuantizationConfig):
679
and output_quant.num_bits == 8
680
and not output_quant.dynamic
681
)
683
- # Static int8-activation layers, plus sub-byte weight-only layers (e.g.
684
- # 2-bit lm_head) that marlin-backed WNA16 cannot serve. Standard 4/8-bit
685
- # weight-only (no activations) falls through to WNA16.
686
- is_subbyte_weight_only = weight_quant.num_bits not in WNA16_SUPPORTED_BITS
687
- needs_wNa8o8 = is_intN_weight and (
688
- (is_static_int8_in and is_static_int8_out) or is_subbyte_weight_only
689
- )
690
- return needs_wNa8o8
682
+ return is_intN_weight and (is_static_int8_in or is_static_int8_out)
683
684
def _get_scheme_from_parts(
685
self,
@@ -740,10 +732,8 @@ class CompressedTensorsConfig(QuantizationConfig):
732
quant_format=format,
733
)
734
743
- if (
744
- self._is_wNa16_group_channel(weight_quant, input_quant)
745
- and (format == CompressionFormat.pack_quantized.value)
746
- and (weight_quant.num_bits in WNA16_SUPPORTED_BITS)
735
+ if self._is_wNa16_group_channel(weight_quant, input_quant) and (
736
+ format == CompressionFormat.pack_quantized.value
737
):
738
return CompressedTensorsWNA16(
739
num_bits=weight_quant.num_bits,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
+1
-2
@@ -11,7 +11,7 @@ from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
11
from .compressed_tensors_w8a8_mxfp8 import CompressedTensorsW8A8Mxfp8
12
from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8
13
from .compressed_tensors_wNa8o8 import CompressedTensorsWNA8O8Int
14
-from .compressed_tensors_wNa16 import WNA16_SUPPORTED_BITS, CompressedTensorsWNA16
14
+from .compressed_tensors_wNa16 import CompressedTensorsWNA16
15
16
__all__ = [
17
"CompressedTensorsScheme",
@@ -20,7 +20,6 @@ __all__ = [
20
"CompressedTensorsW8A16Fp8",
21
"CompressedTensorsW8A8Int8",
22
"CompressedTensorsW8A8Fp8",
23
- "WNA16_SUPPORTED_BITS",
23
"CompressedTensorsW4A4Mxfp4",
24
"CompressedTensorsW4A4Fp4",
25
"CompressedTensorsW4A8Int",
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
+33
-7
@@ -1,7 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
from collections.abc import Callable
6
+from fractions import Fraction
7
8
import torch
9
from compressed_tensors.quantization import ActivationOrdering
@@ -32,7 +34,15 @@ from vllm.scalar_type import scalar_types
34
logger = init_logger(__name__)
35
36
__all__ = ["CompressedTensorsWNA16"]
35
-WNA16_SUPPORTED_TYPES_MAP = {4: scalar_types.uint4b8, 8: scalar_types.uint8b128}
37
+WNA16_SUPPORTED_TYPES_MAP = {
38
+ 2: scalar_types.uint2b2,
39
+ 3: scalar_types.uint3b4,
40
+ 4: scalar_types.uint4b8,
41
+ 5: scalar_types.uint5b16,
42
+ 6: scalar_types.uint6b32,
43
+ 7: scalar_types.uint7b64,
44
+ 8: scalar_types.uint8b128,
45
+}
46
WNA16_ZP_SUPPORTED_TYPES_MAP = {4: scalar_types.uint4, 8: scalar_types.uint8}
47
WNA16_SUPPORTED_BITS = list(WNA16_SUPPORTED_TYPES_MAP.keys())
48
@@ -49,7 +59,8 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
59
actorder: ActivationOrdering | None = None,
60
layer_name: str | None = None,
61
):
52
- self.pack_factor = 32 // num_bits
62
+ self.num_bits = num_bits
63
+ self.pack_factor = Fraction(32, num_bits)
64
self.strategy = strategy
65
self.symmetric = symmetric
66
self.group_size = -1 if group_size is None else group_size
@@ -58,15 +69,22 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
69
70
if self.group_size == -1 and self.strategy != "channel":
71
raise ValueError(
61
- "Marlin kernels require group quantization or "
62
- "channelwise quantization, but found no group "
72
+ "Pack-quantized format requires group quantization "
73
+ "or channelwise quantization, but found no group "
74
"size and strategy is not channelwise."
75
)
76
77
if num_bits not in WNA16_SUPPORTED_TYPES_MAP:
78
raise ValueError(
79
f"Unsupported num_bits = {num_bits}. "
69
- f"Supported num_bits = {WNA16_SUPPORTED_TYPES_MAP.keys()}"
80
+ f"Supported num_bits = {list(WNA16_SUPPORTED_TYPES_MAP)}"
81
+ )
82
+
83
+ if not self.symmetric and num_bits not in WNA16_ZP_SUPPORTED_TYPES_MAP:
84
+ raise ValueError(
85
+ f"Asymmetric quantization not supported for "
86
+ f"num_bits = {num_bits}. Supported: "
87
+ f"{list(WNA16_ZP_SUPPORTED_TYPES_MAP)}"
88
)
89
90
self.quant_type = (
@@ -92,6 +110,12 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
110
**kwargs,
111
):
112
output_size_per_partition = sum(output_partition_sizes)
113
+ layer.input_size_per_partition = input_size_per_partition
114
+ layer.output_size_per_partition = output_size_per_partition
115
+ layer.output_partition_sizes = output_partition_sizes
116
+ layer.params_dtype = params_dtype
117
+ if not hasattr(layer, "has_bias"):
118
+ layer.has_bias = False
119
120
mp_linear_kernel_config = MPLinearLayerConfig(
121
full_weight_shape=(input_size, output_size),
@@ -130,6 +154,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
154
assert input_size_per_partition % group_size == 0
155
scales_and_zp_size = input_size_per_partition // group_size
156
157
+ packed_input_dim = math.ceil(input_size_per_partition * self.num_bits / 32)
158
weight = PackedvLLMParameter(
159
input_dim=1,
160
output_dim=0,
@@ -138,7 +163,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
163
packed_dim=1,
164
data=torch.empty(
165
output_size_per_partition,
141
- input_size_per_partition // self.pack_factor,
166
+ packed_input_dim,
167
dtype=torch.int32,
168
),
169
)
@@ -152,10 +177,11 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
177
),
178
}
179
180
+ packed_output_dim = math.ceil(output_size_per_partition * self.num_bits / 32)
181
zeros_args = {
182
"weight_loader": weight_loader,
183
"data": torch.zeros(
158
- output_size_per_partition // self.pack_factor,
184
+ packed_output_dim,
185
scales_and_zp_size,
186
dtype=torch.int32,
187
),
vllm/scalar_type.py
+3
@@ -348,6 +348,9 @@ class scalar_types:
348
uint2b2 = ScalarType.uint(2, 2)
349
uint3b4 = ScalarType.uint(3, 4)
350
uint4b8 = ScalarType.uint(4, 8)
351
+ uint5b16 = ScalarType.uint(5, 16)
352
+ uint6b32 = ScalarType.uint(6, 32)
353
+ uint7b64 = ScalarType.uint(7, 64)
354
uint8b128 = ScalarType.uint(8, 128)
355
356
# colloquial names