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, 03:30 UTC
Enable ZP Support for Machete (#20268)
Signed-off-by: czhu-cohere <conway.zhu@cohere.com>
czhu-cohere committed
Jul 1, 2025 at 00:12 UTC
9909726d2a30d834d97efd7bf1c4fc0e52fa48b5
3 files changed
+19
-5
benchmarks/kernels/benchmark_machete.py
+2
@@ -234,8 +234,10 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
234
235
fn = lambda: ops.gptq_marlin_gemm(
236
a=bt.a,
237
+ c=None,
238
b_q_weight=w_q,
239
b_scales=w_s,
240
+ global_scale=None,
241
b_zeros=w_zp,
242
g_idx=g_idx,
243
perm=sort_indices,
tests/kernels/quantization/test_machete_mm.py
+1
-1
@@ -139,7 +139,7 @@ def maybe_convert_zeropoints(zps: Optional[torch.Tensor], s: torch.Tensor):
139
140
def group_size_valid(shape: tuple[int, int, int],
141
group_size: Optional[int]) -> bool:
142
- return group_size is None or group_size == -1 or group_size % shape[2] == 0
142
+ return group_size is None or group_size == -1 or shape[2] % group_size == 0
143
144
145
def machete_quantize_and_pack(atype: torch.dtype,
vllm/model_executor/layers/quantization/kernels/mixed_precision/machete.py
+16
-4
@@ -33,8 +33,6 @@ class MacheteLinearKernel(MPLinearKernel):
33
return False, "Act reordering currently not supported by Machete, "\
34
"when the input features are partitioned across "\
35
"devices"
36
- if c.zero_points:
37
- return False, "Zero points currently not supported by Machete"
36
37
if c.weight_type not in query_machete_supported_quant_types(
38
c.zero_points):
@@ -53,6 +51,7 @@ class MacheteLinearKernel(MPLinearKernel):
51
# note assumes that
52
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
53
# `weight_scale` is: {input_dim = 0, output_dim = 1}
54
+ # `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1}
55
def process_weights_after_loading(self, layer: torch.nn.Module):
56
c = self.config
57
@@ -90,16 +89,29 @@ class MacheteLinearKernel(MPLinearKernel):
89
x.data = x.data.contiguous()
90
return x
91
92
+ def transform_w_zp(x):
93
+ assert isinstance(x, BasevLLMParameter)
94
+ permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=1)
95
+ x_unpacked = unpack_quantized_values_into_int32(x.data,
96
+ c.weight_type,
97
+ packed_dim=1)
98
+ w_s = getattr(layer, self.w_s_name).data
99
+ # pre-apply scales to zero-points
100
+ x.data = (-1.0 * w_s * (x_unpacked.to(w_s.dtype))).contiguous()
101
+ return x
102
+
103
# Repack weights and scales for Machete
104
self._transform_param(layer, self.w_q_name, transform_w_q)
105
self._transform_param(layer, self.w_s_name, transform_w_s)
106
+ if c.zero_points:
107
+ self._transform_param(layer, self.w_zp_name, transform_w_zp)
108
109
def apply_weights(self,
110
layer: torch.nn.Module,
111
x: torch.Tensor,
112
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
113
c = self.config
102
- w_q, w_s, _, _ = self._get_weight_params(layer)
114
+ w_q, w_s, w_zp, _ = self._get_weight_params(layer)
115
116
x_2d = x.reshape(-1, x.shape[-1])
117
out_shape = x.shape[:-1] + (c.partition_weight_shape[1], )
@@ -110,7 +122,7 @@ class MacheteLinearKernel(MPLinearKernel):
122
output = ops.machete_mm(a=x_2d,
123
b_q=w_q,
124
b_type=c.weight_type,
113
- b_group_zeros=None,
125
+ b_group_zeros=w_zp,
126
b_group_scales=w_s,
127
b_group_size=c.group_size)
128