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 5, 22:15 UTC
[CPU] Support CT W4A16 on CPU MP kernel (#38219)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Mar 27, 2026 at 14:15 UTC
becaed6ec885b8f599429650262c56a60b1cbf8f
2 files changed
+42
-20
tests/quantization/test_cpu_wna16.py
+1
@@ -11,6 +11,7 @@ MODELS = [
11
"TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ",
12
"TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx
13
"Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx
14
+ "RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp
15
]
16
DTYPE = ["bfloat16"]
17
vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+41
-20
@@ -58,28 +58,29 @@ class CPUWNA16LinearKernel(MPLinearKernel):
58
return True, None
59
60
# note assumes that
61
- # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
62
- # `weight_scale` is: {input_dim = 0, output_dim = 1}
63
- # `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1}
61
+ # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0} (marlin)
62
+ # or: {input_dim = 1, output_dim = 0, packed_dim = 1} (CT)
63
+ # `weight_scale` is: {input_dim = 0, output_dim = 1} (marlin)
64
+ # or: {input_dim = 1, output_dim = 0} (CT)
65
+ # `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1} (marlin)
66
+ # or: {input_dim = 1, output_dim = 0, packed_dim = 0} (CT)
67
def _process_gptq_weights(self, layer: torch.nn.Module):
65
- packed_weight = layer.qweight.data
68
+ packed_weight = getattr(layer, self.w_q_name)
69
+ assert packed_weight.input_dim == packed_weight.packed_dim
70
+ is_ct_format = packed_weight.input_dim == 1
71
+ if is_ct_format:
72
+ packed_weight = packed_weight.t()
73
bits = self.config.weight_type.mantissa
74
pack_factor = 32 // bits
68
- p_w_k, p_w_n = packed_weight.size()
75
+ p_w_k, _ = packed_weight.size()
76
input_size = p_w_k * pack_factor
70
- output_size = p_w_n
71
- isa_hint = _get_isa_hint(layer.scales.dtype)
77
+ isa_hint = _get_isa_hint(getattr(layer, self.w_s_name).dtype)
78
layer.isa_hint = isa_hint
79
74
- layer.qzeros = None
75
- if not self.config.has_g_idx:
76
- layer.g_idx = None
77
-
80
# convert input dim packed to output dim packed
81
weight = unpack_quantized_values_into_int32(
80
- packed_weight, self.config.weight_type, 1
81
- ).view(p_w_k, p_w_n, pack_factor)
82
- weight = weight.permute(0, 2, 1).reshape(input_size, output_size).contiguous()
82
+ packed_weight, self.config.weight_type, 0
83
+ )
84
weight = pack_quantized_values_into_int32(weight, self.config.weight_type, 1)
85
# make 16 output channel as a block and transpose to the make
86
# the block contiguous
@@ -89,10 +90,29 @@ class CPUWNA16LinearKernel(MPLinearKernel):
90
.reshape(-1, input_size * 16 // pack_factor)
91
.contiguous()
92
)
92
- layer.qweight.data = weight
93
+ getattr(layer, self.w_q_name).data = weight
94
+
95
+ # transpose scale, zp for CT format
96
+ if is_ct_format:
97
+ scales = getattr(layer, self.w_s_name)
98
+ scales.data = scales.t().contiguous()
99
+ if self.config.zero_points:
100
+ assert self.w_zp_name
101
+ zp = getattr(layer, self.w_zp_name)
102
+ zp.data = zp.t().contiguous()
103
104
def process_weights_after_loading(self, layer: torch.nn.Module):
95
- if not self.config.zero_points:
105
+ if (not self.config.zero_points) and (self.w_zp_name is not None):
106
+ setattr(layer, self.w_zp_name, None)
107
+
108
+ if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
109
+ setattr(layer, self.w_gidx_name, None)
110
+
111
+ w_input_dim = getattr(layer, self.w_q_name).input_dim
112
+ w_pack_dim = getattr(layer, self.w_q_name).packed_dim
113
+ quant_method = "gptq" if w_pack_dim == w_input_dim else "awq"
114
+
115
+ if quant_method == "gptq":
116
# GPTQ
117
self._process_gptq_weights(layer)
118
else:
@@ -105,12 +125,13 @@ class CPUWNA16LinearKernel(MPLinearKernel):
125
x: torch.Tensor,
126
bias: torch.Tensor | None = None,
127
) -> torch.Tensor:
128
+ w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
129
x = ops.cpu_gemm_wna16(
130
input=x,
110
- q_weight=layer.qweight,
111
- scales=layer.scales,
112
- zeros=layer.qzeros,
113
- g_idx=layer.g_idx,
131
+ q_weight=w_q,
132
+ scales=w_s,
133
+ zeros=w_zp,
134
+ g_idx=w_gidx,
135
bias=bias,
136
pack_factor=8, # 32 // 4
137
isa_hint=layer.isa_hint,