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 1, 22:30 UTC
[Quant] Support compressed-tensors WNA8O8Int linears and WNInt embeddings (#44340)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jun 4, 2026 at 10:40 UTC
06ee2d8433831f69d5de3a6d9fa3d7d042dd394f
14 files changed
+744
-27
requirements/common.txt
+1
-1
@@ -38,7 +38,7 @@ pyyaml
38
six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
39
setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
40
einops # Required for Qwen2-VL.
41
-compressed-tensors == 0.15.0.1 # required for compressed-tensors
41
+compressed-tensors == 0.17.0 # required for compressed-tensors
42
depyf==0.20.0 # required for profiling and debugging with compilation config
43
cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
44
watchfiles # required for http server to monitor the updates of TLS files
requirements/cuda.txt
+1
-1
@@ -28,4 +28,4 @@ quack-kernels>=0.3.3
28
tokenspeed-mla==0.1.2
29
30
# Humming kernels for quantization gemm
31
-humming-kernels[cu13]==0.1.2
31
+humming-kernels[cu13]==0.1.4
requirements/test/rocm.txt
+1
-1
@@ -143,7 +143,7 @@ colorful==0.5.8
143
# via ray
144
colorlog==6.10.1
145
# via optuna
146
-compressed-tensors==0.15.0.1
146
+compressed-tensors==0.17.0
147
# via
148
# -c requirements/common.txt
149
# -r requirements/test/../common.txt
tests/kernels/quantization/test_quantized_embedding.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the Triton dequant-gather kernel used by
4
+``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup)."""
5
+
6
+import pytest
7
+import torch
8
+from compressed_tensors.compressors.pack_quantized.helpers import unpack_from_int32
9
+
10
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_embedding import ( # noqa: E501
11
+ _dequant_gather_triton,
12
+)
13
+from vllm.platforms import current_platform
14
+
15
+
16
+def _dequant_gather_torch(
17
+ ids: torch.Tensor,
18
+ weight_packed: torch.Tensor,
19
+ weight_scale: torch.Tensor,
20
+ hidden: int,
21
+ num_bits: int,
22
+) -> torch.Tensor:
23
+ """Reference: gather packed rows by id, unpack int32-packed INT, dequant."""
24
+ n = ids.shape[0]
25
+ int8 = unpack_from_int32(weight_packed[ids], num_bits, torch.Size([n, hidden]))
26
+ scale_rows = weight_scale[ids]
27
+ w = int8.to(scale_rows.dtype)
28
+ if scale_rows.shape[1] == 1:
29
+ return w * scale_rows
30
+ ng = scale_rows.shape[1]
31
+ return (w.view(n, ng, hidden // ng) * scale_rows.unsqueeze(-1)).view(n, hidden)
32
+
33
+
34
+@pytest.mark.skipif(
35
+ not current_platform.is_cuda(), reason="Triton dequant kernel requires CUDA"
36
+)
37
+@pytest.mark.parametrize("num_bits", [2, 4, 8])
38
+@pytest.mark.parametrize("group_size", [0, 256]) # 0 -> channel
39
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
40
+@pytest.mark.parametrize("num_ids", [1, 17, 4096])
41
+def test_dequant_gather(num_bits, group_size, dtype, num_ids):
42
+ torch.manual_seed(0)
43
+ device = "cuda"
44
+ vocab, hidden = 1000, 2048
45
+ pack_factor = 32 // num_bits
46
+
47
+ # Random full-range int32 packed weights (covers the sign bit -> exercises the
48
+ # arithmetic-shift + mask unpack path).
49
+ weight_packed = torch.randint(
50
+ -(2**31),
51
+ 2**31,
52
+ (vocab, hidden // pack_factor),
53
+ dtype=torch.int32,
54
+ device=device,
55
+ )
56
+
57
+ num_groups = 1 if group_size == 0 else hidden // group_size
58
+ weight_scale = torch.rand(vocab, num_groups, dtype=dtype, device=device) + 0.01
59
+
60
+ ids = torch.randint(0, vocab, (num_ids,), dtype=torch.long, device=device)
61
+
62
+ out = _dequant_gather_triton(ids, weight_packed, weight_scale, hidden, num_bits)
63
+ ref = _dequant_gather_torch(ids, weight_packed, weight_scale, hidden, num_bits)
64
+
65
+ assert out.shape == (num_ids, hidden)
66
+ assert out.dtype == dtype
67
+ torch.testing.assert_close(out, ref)
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -45,6 +45,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.dynamic_4bit import (
45
from vllm.model_executor.kernels.linear.mixed_precision.exllama import (
46
ExllamaLinearKernel,
47
)
48
+from vllm.model_executor.kernels.linear.mixed_precision.humming import (
49
+ HummingLinearKernel,
50
+)
51
from vllm.model_executor.kernels.linear.mixed_precision.machete import (
52
MacheteLinearKernel,
53
)
@@ -345,6 +348,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
348
MacheteLinearKernel,
349
AllSparkLinearKernel,
350
MarlinLinearKernel,
351
+ HummingLinearKernel,
352
ConchLinearKernel,
353
ExllamaLinearKernel,
354
TritonW4A16LinearKernel,
vllm/model_executor/kernels/linear/mixed_precision/humming.py
new
+61
@@ -0,0 +1,61 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Humming GEMM as a mixed-precision WNA16Int linear kernel."""
4
+
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.utils.import_utils import _has_module
9
+
10
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
11
+
12
+
13
+class HummingLinearKernel(MPLinearKernel):
14
+ @classmethod
15
+ def get_min_capability(cls) -> int:
16
+ return 75
17
+
18
+ @classmethod
19
+ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
20
+ if not current_platform.is_cuda():
21
+ return False, "Humming is only supported on CUDA"
22
+ if not _has_module("humming"):
23
+ return False, "Humming is not installed"
24
+ if c.has_g_idx:
25
+ return False, "Humming does not support act-order (g_idx)"
26
+ if c.zero_points:
27
+ return False, "Humming linear kernel only supports symmetric weights"
28
+ return True, None
29
+
30
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
31
+ from vllm.model_executor.layers.quantization.utils.humming_utils import (
32
+ convert_linear_layer_to_humming_standard,
33
+ prepare_humming_layer,
34
+ )
35
+
36
+ name_map = {"weight": self.w_q_name, "weight_scale": self.w_s_name}
37
+ group_size = self.config.group_size
38
+ quant_config = {
39
+ "quant_method": "humming",
40
+ "dtype": "int" + str(self.config.weight_type.size_bits),
41
+ "group_size": 0 if group_size == -1 else group_size,
42
+ }
43
+
44
+ convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
45
+ prepare_humming_layer(layer, quant_config)
46
+
47
+ def apply_weights(
48
+ self,
49
+ layer: torch.nn.Module,
50
+ x: torch.Tensor,
51
+ bias: torch.Tensor | None = None,
52
+ ) -> torch.Tensor:
53
+ from humming.layer import HummingMethod
54
+
55
+ flatten_inputs = x.view(-1, x.size(-1))
56
+ output = HummingMethod.forward_layer(
57
+ layer=layer,
58
+ inputs=flatten_inputs,
59
+ compute_config=layer.compute_config,
60
+ )
61
+ return output.view(*x.shape[:-1], output.size(-1))
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+101
-2
@@ -30,6 +30,9 @@ from vllm.model_executor.layers.quantization.base_config import (
30
QuantizationConfig,
31
QuantizeMethodBase,
32
)
33
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_embedding import ( # noqa: E501
34
+ CompressedTensorsEmbeddingWNA16Int,
35
+)
36
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
37
CompressedTensorsMoEMethod,
38
)
@@ -44,6 +47,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
47
CompressedTensorsW8A8Int8,
48
CompressedTensorsW8A8Mxfp8,
49
CompressedTensorsW8A16Fp8,
50
+ CompressedTensorsWNA8O8Int,
51
CompressedTensorsWNA16,
52
)
53
from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear import ( # noqa: E501
@@ -56,7 +60,10 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
60
should_ignore_layer,
61
)
62
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
59
-from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
63
+from vllm.model_executor.layers.vocab_parallel_embedding import (
64
+ ParallelLMHead,
65
+ VocabParallelEmbedding,
66
+)
67
from vllm.platforms import current_platform
68
69
if TYPE_CHECKING:
@@ -177,6 +184,24 @@ class CompressedTensorsConfig(QuantizationConfig):
184
layer.scheme = quant_scheme
185
return CompressedTensorsLinearMethod(self)
186
187
+ # ParallelLMHead subclasses VocabParallelEmbedding but is handled above as
188
+ # a linear; only true embedding lookups land here.
189
+ if isinstance(layer, VocabParallelEmbedding):
190
+ scheme_dict = self.get_scheme_dict(layer, layer_name=prefix)
191
+ weight_quant = scheme_dict.get("weights") if scheme_dict else None
192
+ if weight_quant is None:
193
+ return None # unquantized embedding
194
+ if not (
195
+ isinstance(weight_quant, QuantizationArgs)
196
+ and self._is_wNa16_group_channel(weight_quant, None)
197
+ and weight_quant.type == QuantizationType.INT
198
+ ):
199
+ raise ValueError(
200
+ "compressed-tensors embeddings only support weight-only INT "
201
+ f"group/channel (WNA16) quantization, got: {weight_quant}"
202
+ )
203
+ return CompressedTensorsEmbeddingWNA16Int(weight_quant)
204
+
205
if isinstance(layer, Attention):
206
return CompressedTensorsKVCacheMethod(self)
207
if isinstance(layer, RoutedExperts):
@@ -324,6 +349,15 @@ class CompressedTensorsConfig(QuantizationConfig):
349
quant_config.get("input_activations")
350
)
351
)
352
+
353
+ # Static output-activation quant is applied as a float fake-quant
354
+ # on the layer output; capture it when present.
355
+ target_scheme_map[target]["output_activations"] = None
356
+ output_activations = quant_config.get("output_activations")
357
+ if output_activations:
358
+ target_scheme_map[target]["output_activations"] = (
359
+ QuantizationArgs.model_validate(output_activations)
360
+ )
361
return target_scheme_map
362
363
@classmethod
@@ -604,10 +638,56 @@ class CompressedTensorsConfig(QuantizationConfig):
638
639
return is_channel_group and input_quant_none and is_static
640
641
+ @staticmethod
642
+ def _is_wNa8o8_int(
643
+ weight_quant: QuantizationArgs,
644
+ input_quant: QuantizationArgs | None,
645
+ output_quant: QuantizationArgs | None,
646
+ format: str | None,
647
+ ) -> bool:
648
+ """Weight N-bit INT (pack-quantized for sub-byte, int-quantized for 8-bit)
649
+ with static per-tensor INT8 input/output activation quant, applied as a float
650
+ fake-quant around a weight-only matmul."""
651
+ is_int_pack_format = format in (
652
+ CompressionFormat.pack_quantized.value,
653
+ CompressionFormat.int_quantized.value,
654
+ )
655
+ is_channel_group = weight_quant.strategy in (
656
+ QuantizationStrategy.CHANNEL.value,
657
+ QuantizationStrategy.GROUP.value,
658
+ )
659
+ is_static_int = (
660
+ weight_quant.type == QuantizationType.INT and not weight_quant.dynamic
661
+ )
662
+ is_intN_weight = is_static_int and is_channel_group and is_int_pack_format
663
+ is_static_int8_in = (
664
+ input_quant is not None
665
+ and input_quant.type == QuantizationType.INT
666
+ and input_quant.strategy == QuantizationStrategy.TENSOR.value
667
+ and input_quant.num_bits == 8
668
+ and not input_quant.dynamic
669
+ )
670
+ is_static_int8_out = (
671
+ output_quant is not None
672
+ and output_quant.type == QuantizationType.INT
673
+ and output_quant.strategy == QuantizationStrategy.TENSOR.value
674
+ and output_quant.num_bits == 8
675
+ and not output_quant.dynamic
676
+ )
677
+ # Static int8-activation layers, plus sub-byte weight-only layers (e.g.
678
+ # 2-bit lm_head) that marlin-backed WNA16 cannot serve. Standard 4/8-bit
679
+ # weight-only (no activations) falls through to WNA16.
680
+ is_subbyte_weight_only = weight_quant.num_bits not in WNA16_SUPPORTED_BITS
681
+ needs_wNa8o8 = is_intN_weight and (
682
+ (is_static_int8_in and is_static_int8_out) or is_subbyte_weight_only
683
+ )
684
+ return needs_wNa8o8
685
+
686
def _get_scheme_from_parts(
687
self,
688
weight_quant: QuantizationArgs,
689
input_quant: QuantizationArgs,
690
+ output_quant: QuantizationArgs | None = None,
691
format: str | None = None,
692
layer_name: str | None = None,
693
) -> "CompressedTensorsScheme":
@@ -641,6 +721,19 @@ class CompressedTensorsConfig(QuantizationConfig):
721
actorder=weight_quant.actorder,
722
)
723
724
+ # Must come before the WNA16 check; standard 4/8-bit weight-only (no
725
+ # output-activation scale) still falls through to WNA16.
726
+ if self._is_wNa8o8_int(weight_quant, input_quant, output_quant, format):
727
+ return CompressedTensorsWNA8O8Int(
728
+ num_bits=weight_quant.num_bits,
729
+ strategy=weight_quant.strategy,
730
+ group_size=weight_quant.group_size,
731
+ has_input_act=input_quant is not None,
732
+ has_output_act=output_quant is not None,
733
+ layer_name=layer_name,
734
+ quant_format=format,
735
+ )
736
+
737
if (
738
self._is_wNa16_group_channel(weight_quant, input_quant)
739
and (format == CompressionFormat.pack_quantized.value)
@@ -708,7 +801,10 @@ class CompressedTensorsConfig(QuantizationConfig):
801
input_symmetric=input_quant.symmetric,
802
)
803
711
- raise NotImplementedError("No compressed-tensors compatible scheme was found.")
804
+ raise NotImplementedError(
805
+ f"No compressed-tensors compatible scheme was found for {layer_name=}, "
806
+ f"{weight_quant=}, {input_quant=}, {output_quant=}, {format=}"
807
+ )
808
809
def get_scheme(
810
self, layer: torch.nn.Module, layer_name: str | None = None
@@ -731,10 +827,12 @@ class CompressedTensorsConfig(QuantizationConfig):
827
828
weight_quant = None
829
input_quant = None
830
+ output_quant = None
831
format = None
832
if scheme_dict:
833
weight_quant = scheme_dict.get("weights")
834
input_quant = scheme_dict.get("input_activations")
835
+ output_quant = scheme_dict.get("output_activations")
836
format = scheme_dict.get("format")
837
838
if weight_quant is None:
@@ -746,6 +844,7 @@ class CompressedTensorsConfig(QuantizationConfig):
844
scheme = self._get_scheme_from_parts( # type: ignore
845
weight_quant=weight_quant,
846
input_quant=input_quant,
847
+ output_quant=output_quant,
848
format=format,
849
layer_name=layer_name,
850
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_embedding.py
new
+170
@@ -0,0 +1,170 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Quantized embedding method for compressed-tensors.
4
+
5
+Adds dequant-on-lookup support for a pack-quantized ``VocabParallelEmbedding``
6
+(2-8 bit INT, channel- or group-quantized). Only the gathered token rows are
7
+unpacked and dequantized, so the packed weight is never densified.
8
+"""
9
+
10
+import torch
11
+from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
12
+
13
+from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
14
+from vllm.model_executor.parameter import (
15
+ BasevLLMParameter,
16
+ ChannelQuantScaleParameter,
17
+ GroupQuantScaleParameter,
18
+ PackedvLLMParameter,
19
+)
20
+from vllm.triton_utils import tl, triton
21
+
22
+__all__ = ["CompressedTensorsEmbeddingWNA16Int"]
23
+
24
+
25
+@triton.jit
26
+def _dequant_gather_kernel(
27
+ ids_ptr,
28
+ packed_ptr,
29
+ scale_ptr,
30
+ out_ptr,
31
+ hidden,
32
+ packed_cols,
33
+ num_groups,
34
+ NUM_BITS: tl.constexpr,
35
+ PACK_FACTOR: tl.constexpr,
36
+ GROUP_SIZE: tl.constexpr,
37
+ BLOCK: tl.constexpr,
38
+):
39
+ """Gather embedding rows by token id, unpack int32-packed INT weights, and
40
+ dequantize to ``out`` dtype in one pass (no int8 intermediate)."""
41
+ row = tl.program_id(0)
42
+ col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
43
+ col_mask = col < hidden
44
+ tid = tl.load(ids_ptr + row).to(tl.int64)
45
+
46
+ packed_idx = col // PACK_FACTOR
47
+ shift = (col % PACK_FACTOR) * NUM_BITS
48
+ packed = tl.load(
49
+ packed_ptr + tid * packed_cols + packed_idx, mask=col_mask, other=0
50
+ )
51
+ q = ((packed >> shift) & ((1 << NUM_BITS) - 1)) - (1 << (NUM_BITS - 1))
52
+
53
+ if GROUP_SIZE == 0: # channel: one scale per row
54
+ scale = tl.load(scale_ptr + tid)
55
+ else: # group: one scale per (row, group)
56
+ grp = col // GROUP_SIZE
57
+ scale = tl.load(scale_ptr + tid * num_groups + grp, mask=col_mask, other=0.0)
58
+
59
+ out = q.to(tl.float32) * scale.to(tl.float32)
60
+ tl.store(
61
+ out_ptr + row * hidden + col, out.to(out_ptr.dtype.element_ty), mask=col_mask
62
+ )
63
+
64
+
65
+def _dequant_gather_triton(
66
+ ids: torch.Tensor,
67
+ weight_packed: torch.Tensor,
68
+ weight_scale: torch.Tensor,
69
+ hidden: int,
70
+ num_bits: int,
71
+) -> torch.Tensor:
72
+ n = ids.numel()
73
+ out = torch.empty(n, hidden, dtype=weight_scale.dtype, device=weight_packed.device)
74
+ num_groups = weight_scale.shape[1]
75
+ group_size = 0 if num_groups == 1 else hidden // num_groups
76
+ block = min(triton.next_power_of_2(hidden), 1024)
77
+ grid = (n, triton.cdiv(hidden, block))
78
+ _dequant_gather_kernel[grid](
79
+ ids,
80
+ weight_packed,
81
+ weight_scale,
82
+ out,
83
+ hidden,
84
+ weight_packed.shape[1],
85
+ num_groups,
86
+ NUM_BITS=num_bits,
87
+ PACK_FACTOR=32 // num_bits,
88
+ GROUP_SIZE=group_size,
89
+ BLOCK=block,
90
+ )
91
+ return out
92
+
93
+
94
+class CompressedTensorsEmbeddingWNA16Int(QuantizeMethodBase):
95
+ def __init__(self, weight_quant: QuantizationArgs):
96
+ self.num_bits = weight_quant.num_bits
97
+ self.pack_factor = 32 // self.num_bits
98
+ self.strategy = weight_quant.strategy
99
+ self.group_size = weight_quant.group_size
100
+ self.is_group = (
101
+ self.strategy == QuantizationStrategy.GROUP.value
102
+ and self.group_size is not None
103
+ )
104
+
105
+ def create_weights(
106
+ self,
107
+ layer: torch.nn.Module,
108
+ input_size_per_partition: int,
109
+ output_partition_sizes: list[int],
110
+ input_size: int,
111
+ output_size: int,
112
+ params_dtype: torch.dtype,
113
+ **extra_weight_attrs,
114
+ ):
115
+ weight_loader = extra_weight_attrs["weight_loader"]
116
+ # Embedding weight is [num_embeddings(vocab), embedding_dim(hidden)];
117
+ # vocab is the output (partitioned) dim, hidden is the input dim.
118
+ vocab_pp = sum(output_partition_sizes)
119
+ hidden = input_size_per_partition
120
+ layer.hidden_size = hidden
121
+
122
+ weight_packed = PackedvLLMParameter(
123
+ input_dim=1,
124
+ output_dim=0,
125
+ packed_dim=1,
126
+ packed_factor=self.pack_factor,
127
+ weight_loader=weight_loader,
128
+ data=torch.empty(vocab_pp, hidden // self.pack_factor, dtype=torch.int32),
129
+ )
130
+
131
+ if self.is_group:
132
+ assert hidden % self.group_size == 0
133
+ weight_scale = GroupQuantScaleParameter(
134
+ output_dim=0,
135
+ input_dim=1,
136
+ weight_loader=weight_loader,
137
+ data=torch.empty(
138
+ vocab_pp, hidden // self.group_size, dtype=params_dtype
139
+ ),
140
+ )
141
+ else:
142
+ weight_scale = ChannelQuantScaleParameter(
143
+ output_dim=0,
144
+ weight_loader=weight_loader,
145
+ data=torch.empty(vocab_pp, 1, dtype=params_dtype),
146
+ )
147
+
148
+ weight_shape = BasevLLMParameter(
149
+ data=torch.empty(2, dtype=torch.int64), weight_loader=weight_loader
150
+ )
151
+
152
+ layer.register_parameter("weight_packed", weight_packed)
153
+ layer.register_parameter("weight_scale", weight_scale)
154
+ layer.register_parameter("weight_shape", weight_shape)
155
+
156
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
157
+ pass
158
+
159
+ def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor:
160
+ ids = input_.reshape(-1).contiguous()
161
+ hidden = layer.hidden_size
162
+ deq = _dequant_gather_triton(
163
+ ids, layer.weight_packed, layer.weight_scale, hidden, self.num_bits
164
+ )
165
+ return deq.reshape(*input_.shape, hidden)
166
+
167
+ def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
168
+ raise NotImplementedError(
169
+ "CompressedTensorsEmbeddingWNA16Int supports embedding lookup only"
170
+ )
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
+2
@@ -10,11 +10,13 @@ from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
10
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
15
16
__all__ = [
17
"CompressedTensorsScheme",
18
"CompressedTensorsWNA16",
19
+ "CompressedTensorsWNA8O8Int",
20
"CompressedTensorsW8A16Fp8",
21
"CompressedTensorsW8A8Int8",
22
"CompressedTensorsW8A8Fp8",
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py
new
+257
@@ -0,0 +1,257 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Weight N-bit INT scheme with static INT8 input/output activation quant.
4
+
5
+Handles compressed-tensors INT weight checkpoints that carry static per-tensor
6
+INT8 ``input_activations`` and/or ``output_activations``. The activation quant is
7
+reproduced as a float fake-quant on the layer input and output, around a
8
+weight-only matmul, rather than a fused int8 GEMM.
9
+"""
10
+
11
+from collections.abc import Callable
12
+
13
+import torch
14
+from compressed_tensors.compressors.pack_quantized.helpers import pack_to_int32
15
+
16
+from vllm.model_executor.kernels.linear import (
17
+ MPLinearLayerConfig,
18
+ choose_mp_linear_kernel,
19
+)
20
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
21
+ CompressedTensorsScheme,
22
+)
23
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
24
+ marlin_repeat_scales_on_all_ranks,
25
+)
26
+from vllm.model_executor.parameter import (
27
+ BasevLLMParameter,
28
+ ChannelQuantScaleParameter,
29
+ GroupQuantScaleParameter,
30
+ ModelWeightParameter,
31
+ PackedvLLMParameter,
32
+)
33
+from vllm.scalar_type import scalar_types
34
+
35
+__all__ = ["CompressedTensorsWNA8O8Int", "fake_quant_static_int8"]
36
+
37
+WNA8O8_SUPPORTED_TYPES_MAP = {
38
+ 2: scalar_types.uint2b2,
39
+ 4: scalar_types.uint4b8,
40
+ 8: scalar_types.uint8b128,
41
+}
42
+
43
+
44
+def fake_quant_static_int8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
45
+ """Static per-tensor symmetric INT8 quantize-dequantize, in x's dtype."""
46
+ scale = scale.to(x.dtype)
47
+ q = torch.clamp(torch.round(x / scale), -128.0, 127.0)
48
+ return q * scale
49
+
50
+
51
+class CompressedTensorsWNA8O8Int(CompressedTensorsScheme):
52
+ def __init__(
53
+ self,
54
+ num_bits: int,
55
+ strategy: str,
56
+ group_size: int | None = None,
57
+ has_input_act: bool = False,
58
+ has_output_act: bool = False,
59
+ layer_name: str | None = None,
60
+ quant_format: str = "pack-quantized",
61
+ ):
62
+ self.num_bits = num_bits
63
+ self.pack_factor = 32 // num_bits
64
+ self.strategy = strategy
65
+ self.group_size = -1 if group_size is None else group_size
66
+ self.has_input_act = has_input_act
67
+ self.has_output_act = has_output_act
68
+ self.layer_name = layer_name
69
+ # "pack-quantized" (sub-byte, int32-packed) or "int-quantized" (8-bit int8).
70
+ self.quant_format = quant_format
71
+ self.is_int_quantized = quant_format == "int-quantized"
72
+ if num_bits not in WNA8O8_SUPPORTED_TYPES_MAP:
73
+ raise ValueError(
74
+ f"Unsupported num_bits = {num_bits} for WNA8O8Int; "
75
+ f"supported = {sorted(WNA8O8_SUPPORTED_TYPES_MAP)}"
76
+ )
77
+ self.quant_type = WNA8O8_SUPPORTED_TYPES_MAP[num_bits]
78
+ self._input_scale: torch.Tensor | None = None
79
+ self._output_scale: torch.Tensor | None = None
80
+
81
+ @classmethod
82
+ def get_min_capability(cls) -> int:
83
+ return 70
84
+
85
+ def create_weights(
86
+ self,
87
+ layer: torch.nn.Module,
88
+ output_size: int,
89
+ input_size: int,
90
+ output_partition_sizes: list[int],
91
+ input_size_per_partition: int,
92
+ params_dtype: torch.dtype,
93
+ weight_loader: Callable,
94
+ **kwargs,
95
+ ):
96
+ output_size_per_partition = sum(output_partition_sizes)
97
+ layer.input_size_per_partition = input_size_per_partition
98
+ layer.output_size_per_partition = output_size_per_partition
99
+ # Set for kernels' weight prep; also covers ParallelLMHead, which does
100
+ # not set these in __init__.
101
+ layer.output_partition_sizes = output_partition_sizes
102
+ layer.params_dtype = params_dtype
103
+ if not hasattr(layer, "has_bias"):
104
+ layer.has_bias = False
105
+
106
+ mp_config = MPLinearLayerConfig(
107
+ full_weight_shape=(input_size, output_size),
108
+ partition_weight_shape=(
109
+ input_size_per_partition,
110
+ output_size_per_partition,
111
+ ),
112
+ weight_type=self.quant_type,
113
+ act_type=params_dtype, # activation quant applied externally (SRQ)
114
+ group_size=self.group_size,
115
+ zero_points=False,
116
+ has_g_idx=False,
117
+ )
118
+ self.kernel = choose_mp_linear_kernel(mp_config)(
119
+ mp_config,
120
+ w_q_param_name="weight_packed",
121
+ w_s_param_name="weight_scale",
122
+ )
123
+
124
+ self._register_weight(
125
+ layer, input_size, input_size_per_partition, params_dtype, weight_loader
126
+ )
127
+
128
+ def _register_weight(
129
+ self, layer, input_size, input_size_per_partition, params_dtype, weight_loader
130
+ ):
131
+ out = layer.output_size_per_partition
132
+ if self.is_int_quantized:
133
+ # Plain int8 weight; packed to the canonical int32 layout after load.
134
+ layer.register_parameter(
135
+ "weight",
136
+ ModelWeightParameter(
137
+ data=torch.empty(out, input_size_per_partition, dtype=torch.int8),
138
+ input_dim=1,
139
+ output_dim=0,
140
+ weight_loader=weight_loader,
141
+ ),
142
+ )
143
+ else:
144
+ layer.register_parameter(
145
+ "weight_packed",
146
+ PackedvLLMParameter(
147
+ input_dim=1,
148
+ output_dim=0,
149
+ packed_dim=1,
150
+ packed_factor=self.pack_factor,
151
+ weight_loader=weight_loader,
152
+ data=torch.empty(
153
+ out,
154
+ input_size_per_partition // self.pack_factor,
155
+ dtype=torch.int32,
156
+ ),
157
+ ),
158
+ )
159
+ layer.register_parameter(
160
+ "weight_shape",
161
+ BasevLLMParameter(
162
+ data=torch.empty(2, dtype=torch.int64), weight_loader=weight_loader
163
+ ),
164
+ )
165
+
166
+ # Scale: per-output-channel, or per group along the input dim under TP.
167
+ group_size = self.group_size if self.group_size != -1 else input_size
168
+ partitioned = not marlin_repeat_scales_on_all_ranks(
169
+ False, self.group_size, input_size != input_size_per_partition
170
+ )
171
+ scales = (input_size_per_partition if partitioned else input_size) // group_size
172
+ scale_data = torch.empty(out, scales, dtype=params_dtype)
173
+ if partitioned:
174
+ assert input_size_per_partition % group_size == 0
175
+ weight_scale = GroupQuantScaleParameter(
176
+ data=scale_data, output_dim=0, input_dim=1, weight_loader=weight_loader
177
+ )
178
+ else:
179
+ weight_scale = ChannelQuantScaleParameter(
180
+ data=scale_data, output_dim=0, weight_loader=weight_loader
181
+ )
182
+ layer.register_parameter("weight_scale", weight_scale)
183
+
184
+ for name, present in (
185
+ ("input_scale", self.has_input_act),
186
+ ("output_scale", self.has_output_act),
187
+ ):
188
+ if present:
189
+ layer.register_parameter(
190
+ name,
191
+ BasevLLMParameter(
192
+ data=torch.empty(1, dtype=torch.float32),
193
+ weight_loader=weight_loader,
194
+ ),
195
+ )
196
+
197
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
198
+ # Lift the static activation scales off the layer (applied externally) so
199
+ # the kernel only sees weight tensors. Drop uncalibrated (zero) scales.
200
+ self._input_scale = self._take_act_scale(layer, "input_scale")
201
+ self._output_scale = self._take_act_scale(layer, "output_scale")
202
+ self.has_input_act = self._input_scale is not None
203
+ self.has_output_act = self._output_scale is not None
204
+
205
+ if self.is_int_quantized:
206
+ self._pack_int_quantized_weight(layer)
207
+
208
+ self.kernel.process_weights_after_loading(layer)
209
+
210
+ def _pack_int_quantized_weight(self, layer: torch.nn.Module) -> None:
211
+ """Normalize an int-quantized (plain int8) weight to the canonical
212
+ ``weight_packed`` int32 + ``weight_shape`` layout the MP kernels expect."""
213
+ weight = layer.weight
214
+ out_features, in_features = weight.shape
215
+ packed = pack_to_int32(weight.data.contiguous(), self.num_bits)
216
+ delattr(layer, "weight")
217
+
218
+ def _noop_loader(*_, **__):
219
+ return None
220
+
221
+ layer.register_parameter(
222
+ "weight_packed",
223
+ PackedvLLMParameter(
224
+ data=packed.contiguous(),
225
+ input_dim=1,
226
+ output_dim=0,
227
+ packed_dim=1,
228
+ packed_factor=self.pack_factor,
229
+ weight_loader=_noop_loader,
230
+ ),
231
+ )
232
+ layer.register_parameter(
233
+ "weight_shape",
234
+ BasevLLMParameter(
235
+ data=torch.tensor([out_features, in_features], dtype=torch.int64),
236
+ weight_loader=_noop_loader,
237
+ ),
238
+ )
239
+
240
+ @staticmethod
241
+ def _take_act_scale(layer, name: str) -> torch.Tensor | None:
242
+ param = getattr(layer, name, None)
243
+ if param is None:
244
+ return None
245
+ scale = param.data.clone()
246
+ delattr(layer, name)
247
+ return None if float(scale.reshape(-1)[0]) == 0.0 else scale
248
+
249
+ def apply_weights(
250
+ self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
251
+ ) -> torch.Tensor:
252
+ if self.has_input_act:
253
+ x = fake_quant_static_int8(x, self._input_scale)
254
+ out = self.kernel.apply_weights(layer, x, bias)
255
+ if self.has_output_act:
256
+ out = fake_quant_static_int8(out, self._output_scale)
257
+ return out
vllm/model_executor/layers/quantization/utils/humming_utils.py
+61
-5
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import json
4
from typing import Any
5
6
import regex as re
@@ -42,16 +43,57 @@ def humming_is_layer_skipped(config: dict[str, Any], prefix: str):
43
return False
44
45
46
+def convert_linear_layer_to_humming_standard(
47
+ layer: LinearBase, name_map: dict[str, str]
48
+):
49
+ """Rename/reshape a linear layer's quantized params (the canonical MPLinear
50
+ layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names
51
+ and layout humming's weight schema expects (``weight`` / ``weight_scale``)."""
52
+ for name, checkpoint_name in name_map.items():
53
+ tensor = getattr(layer, checkpoint_name)
54
+ delattr(layer, checkpoint_name)
55
+
56
+ if name == "weight":
57
+ input_dim = getattr(tensor, "input_dim", 1)
58
+ output_dim = getattr(tensor, "output_dim", 0)
59
+
60
+ if input_dim == 0 and output_dim == 1:
61
+ tensor = tensor.transpose(1, 0).contiguous()
62
+ else:
63
+ assert output_dim == 0 and input_dim == 1
64
+
65
+ tensor = tensor.view(tensor.size(0), -1).view(torch.int32)
66
+ elif name in ["weight_scale", "zero_point"]:
67
+ if getattr(tensor, "output_dim", 0) == 1:
68
+ tensor = tensor.transpose(0, 1).contiguous()
69
+ if tensor.ndim == 1:
70
+ tensor = tensor.unsqueeze(1)
71
+
72
+ tensor = tensor.view(torch.int32) if name == "zero_point" else tensor
73
+
74
+ if isinstance(tensor, torch.nn.Parameter):
75
+ param = tensor
76
+ else:
77
+ param = torch.nn.Parameter(tensor, requires_grad=False)
78
+
79
+ setattr(layer, name, param)
80
+
81
+
82
def prepare_humming_layer(layer: LinearBase, quant_config: dict):
83
weight_schema = BaseWeightSchema.from_config(quant_config)
84
input_schema = HummingInputSchema()
85
49
- shape_k_stacks = [layer.input_size_per_partition]
86
+ # ReplicatedLinear has no TP partitioning and so does not set
87
+ # input_size_per_partition; for it that is just input_size.
88
+ input_size_per_partition = getattr(
89
+ layer, "input_size_per_partition", layer.input_size
90
+ )
91
+ shape_k_stacks = [input_size_per_partition]
92
shape_n_stacks = layer.output_partition_sizes
93
94
# Step 1: convert weight to humming standard format
95
weight_schema, tensors = weight_schema.convert_humming(
54
- tensors=layer.named_parameters(),
96
+ tensors=dict(layer.named_parameters()),
97
shape_n_stacks=shape_n_stacks,
98
shape_k_stacks=shape_k_stacks,
99
param_dtype=layer.params_dtype,
@@ -63,23 +105,37 @@ def prepare_humming_layer(layer: LinearBase, quant_config: dict):
105
delattr(layer, name)
106
107
for name, tensor in tensors.items():
108
+ if isinstance(tensor, torch.nn.Parameter):
109
+ tensor = tensor.data
110
param = torch.nn.Parameter(tensor, requires_grad=False)
111
setattr(layer, name, param)
112
113
# Step 2: transform weight (humming standard format) for forwarding
114
HummingMethod.prepare_layer_meta(
115
layer=layer,
72
- shape_n=layer.output_partition_sizes_sum,
73
- shape_k=layer.input_size_per_partition,
116
+ shape_n=sum(layer.output_partition_sizes),
117
+ shape_k=input_size_per_partition,
118
weight_schema=weight_schema,
119
input_schema=input_schema,
120
pad_n_to_multiple=256,
121
pad_k_to_multiple=128,
122
has_bias=layer.has_bias,
79
- torch_dtype=layer.param_dtype,
123
+ torch_dtype=layer.params_dtype,
124
)
125
126
HummingMethod.transform_humming_layer(layer)
127
+ if not hasattr(layer, "locks"):
128
+ device = layer.weight.device
129
+ locks = torch.zeros(1024, dtype=torch.int32, device=device)
130
+ layer.register_buffer("locks", locks)
131
+
132
+ compute_config = {
133
+ "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
134
+ "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
135
+ "gemm_type": "dense",
136
+ }
137
+
138
+ layer.compute_config = json.dumps(compute_config)
139
140
141
def prepare_humming_moe_layer(layer: RoutedExperts, quant_config: dict):
vllm/model_executor/models/gemma4.py
+3
-3
@@ -1057,7 +1057,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1057
"normalizer",
1058
torch.tensor(
1059
config.hidden_size**0.5,
1060
- dtype=self.embed_tokens.weight.dtype,
1060
+ dtype=vllm_config.model_config.dtype,
1061
),
1062
persistent=False,
1063
)
@@ -1111,7 +1111,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1111
)
1112
self.hidden_states = torch.zeros(
1113
(max_num_tokens, config.hidden_size),
1114
- dtype=self.embed_tokens.weight.dtype,
1114
+ dtype=vllm_config.model_config.dtype,
1115
device=device,
1116
)
1117
if (
@@ -1124,7 +1124,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1124
config.num_hidden_layers,
1125
self.hidden_size_per_layer_input,
1126
),
1127
- dtype=self.embed_tokens.weight.dtype,
1127
+ dtype=vllm_config.model_config.dtype,
1128
device=device,
1129
)
1130
else:
vllm/model_executor/models/gemma4_mm.py
+12
-12
@@ -1011,14 +1011,17 @@ class Gemma4ForConditionalGeneration(
1011
self.config = config
1012
self.quant_config = quant_config
1013
self.multimodal_config = multimodal_config
1014
+ self.model_dtype = vllm_config.model_config.dtype
1015
1016
# Only quantize towers when the quant method supports their
1017
# dimensions. BNB/torchao handle arbitrary sizes; other methods
1018
# (Marlin, FP8, …) require dimensions divisible by 64, which
1019
# the vision tower (intermediate_size=4304) does not satisfy.
1020
+ # TODO(mgoin): remove this by fixing kernel padding.
1021
if quant_config and quant_config.get_name() in [
1022
"bitsandbytes",
1023
"torchao",
1024
+ "compressed-tensors",
1025
]:
1026
tower_quant = quant_config
1027
else:
@@ -1081,12 +1084,13 @@ class Gemma4ForConditionalGeneration(
1084
# Some variants have hidden_size_per_layer_input=None (no PLE).
1085
ple_dim = config.text_config.hidden_size_per_layer_input
1086
if ple_dim is not None and ple_dim > 0:
1087
+ embed = self.language_model.model.embed_tokens
1088
self.per_layer_embeddings = torch.zeros(
1089
vllm_config.scheduler_config.max_num_batched_tokens,
1090
config.text_config.num_hidden_layers,
1091
ple_dim,
1088
- device=self.language_model.model.embed_tokens.weight.device,
1089
- dtype=self.language_model.model.embed_tokens.weight.dtype,
1092
+ device=next(embed.parameters()).device,
1093
+ dtype=vllm_config.model_config.dtype,
1094
)
1095
else:
1096
self.per_layer_embeddings = None
@@ -1246,7 +1250,6 @@ class Gemma4ForConditionalGeneration(
1250
vt = self.vision_tower
1251
vision_cfg = self.config.vision_config
1252
pooling_k2 = vision_cfg.pooling_kernel_size**2
1249
- target_dtype = self.language_model.model.embed_tokens.weight.dtype
1253
1254
# Concurrent requests with different image resolutions may
1255
# arrive as a list of per-image tensors, while same-resolution
@@ -1291,7 +1294,7 @@ class Gemma4ForConditionalGeneration(
1294
pv_tensor,
1295
pp_tensor,
1296
pad_tensor,
1294
- ).to(target_dtype)
1297
+ ).to(self.model_dtype)
1298
encoder_outputs = vt.encoder(
1299
inputs_embeds=inputs_embeds,
1300
attention_mask=~pad_tensor,
@@ -1328,12 +1331,8 @@ class Gemma4ForConditionalGeneration(
1331
all_valid_states[orig_idx] = valid_states
1332
valid_lens[orig_idx] = valid_states.shape[0]
1333
1331
- # Use embed_tokens dtype as compute dtype; embedding_projection.weight
1332
- # may be uint8 under BnB 4-bit, which would corrupt the cast.
1333
- target_dtype = self.language_model.model.embed_tokens.weight.dtype
1334
-
1334
# Project all images in a single batched call.
1336
- flat_valid_states = torch.cat(all_valid_states, dim=0).to(target_dtype)
1335
+ flat_valid_states = torch.cat(all_valid_states, dim=0).to(self.model_dtype)
1336
flat_proj_embs = self.embed_vision(
1337
inputs_embeds=flat_valid_states.unsqueeze(0)
1338
).squeeze(0)
@@ -1373,7 +1372,6 @@ class Gemma4ForConditionalGeneration(
1372
vt = self.vision_tower
1373
vision_cfg = self.config.vision_config
1374
pooling_k2 = vision_cfg.pooling_kernel_size**2
1376
- target_dtype = self.language_model.model.embed_tokens.weight.dtype
1375
1376
if isinstance(frame_counts, torch.Tensor):
1377
fc_list = frame_counts.tolist()
@@ -1405,7 +1403,7 @@ class Gemma4ForConditionalGeneration(
1403
pv_chunk,
1404
pp_chunk,
1405
pad_chunk,
1408
- ).to(target_dtype)
1406
+ ).to(self.model_dtype)
1407
encoder_outputs = vt.encoder(
1408
inputs_embeds=inputs_embeds,
1409
attention_mask=~pad_chunk,
@@ -1440,7 +1438,9 @@ class Gemma4ForConditionalGeneration(
1438
frame_valid_lens.append(valid_states.shape[0])
1439
1440
# Project all frames in a single batched call.
1443
- flat_valid_states = torch.cat(all_frame_valid_states, dim=0).to(target_dtype)
1441
+ flat_valid_states = torch.cat(all_frame_valid_states, dim=0).to(
1442
+ self.model_dtype
1443
+ )
1444
flat_proj_embs = self.embed_vision(
1445
inputs_embeds=flat_valid_states.unsqueeze(0)
1446
).squeeze(0)
vllm/model_executor/models/gemma4_unified.py
+3
-2
@@ -307,12 +307,13 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):
307
None,
308
)
309
if ple_dim is not None and ple_dim > 0:
310
+ embed = self.language_model.model.embed_tokens
311
self.per_layer_embeddings = torch.zeros(
312
vllm_config.scheduler_config.max_num_batched_tokens,
313
config.text_config.num_hidden_layers,
314
ple_dim,
314
- device=self.language_model.model.embed_tokens.weight.device,
315
- dtype=self.language_model.model.embed_tokens.weight.dtype,
315
+ device=next(embed.parameters()).device,
316
+ dtype=vllm_config.model_config.dtype,
317
)
318
else:
319
self.per_layer_embeddings = None