[XPU] add awq format for INCXPULinear (#43404)
Signed-off-by: Ma, Liangliang <liangliang.ma@intel.com>
Liangliang Ma committed
Jun 22, 2026 at 22:29 UTC
3da4a1b124a8839b1014e5d571784fc8d0953de7
2 files changed
+89
-16
tests/quantization/test_auto_round.py
+2
-2
@@ -43,8 +43,8 @@ MODELS = [
43
pytest.param(
44
"Intel/Qwen2-0.5B-Instruct-int4-sym-AutoRound",
45
marks=pytest.mark.skipif(
46
- not current_platform.is_cuda(),
47
- reason="AWQ AutoRound model only supports CUDA backend for now.",
46
+ not (current_platform.is_cuda() or current_platform.is_xpu()),
47
+ reason="AWQ AutoRound model only supports CUDA/XPU backend for now.",
48
),
49
id="auto_round:auto_awq",
50
),
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py
+87
-14
@@ -185,11 +185,17 @@ class INCWNA16LinearScheme(INCLinearScheme):
185
186
187
class INCXPULinearBase(INCLinearScheme):
188
+ # AWQ packs nibbles within each int32 in the order [0, 2, 4, 6, 1, 3, 5, 7];
189
+ # this permutation undoes that ordering so values can be repacked in
190
+ # standard sequential (GPTQ) order.
191
+ _REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
192
+
193
def __init__(self, layer_config: "INCLayerConfig") -> None:
194
self.weight_bits = layer_config.bits
195
self.group_size = layer_config.group_size
196
self.sym = layer_config.sym
197
self.pack_factor = 32 // self.weight_bits
198
+ self.is_awq_packed = layer_config.is_awq
199
200
@classmethod
201
def get_min_capability(cls) -> int:
@@ -206,18 +212,34 @@ class INCXPULinearBase(INCLinearScheme):
212
output_size_per_partition = sum(output_partition_sizes)
213
scales_and_zp_size = input_size_per_partition // self.group_size
214
209
- qweight = PackedvLLMParameter(
210
- data=torch.empty(
211
- input_size_per_partition // self.pack_factor,
212
- output_size_per_partition,
213
- dtype=torch.int32,
214
- ),
215
- input_dim=0,
216
- output_dim=1,
217
- packed_dim=0,
218
- packed_factor=self.pack_factor,
219
- weight_loader=weight_loader,
220
- )
215
+ if self.is_awq_packed:
216
+ # AWQ: qweight [in, out // pack_factor] packed along output dim
217
+ qweight = PackedvLLMParameter(
218
+ data=torch.empty(
219
+ input_size_per_partition,
220
+ output_size_per_partition // self.pack_factor,
221
+ dtype=torch.int32,
222
+ ),
223
+ input_dim=0,
224
+ output_dim=1,
225
+ packed_dim=1,
226
+ packed_factor=self.pack_factor,
227
+ weight_loader=weight_loader,
228
+ )
229
+ else:
230
+ # GPTQ: qweight [in // pack_factor, out] packed along input dim
231
+ qweight = PackedvLLMParameter(
232
+ data=torch.empty(
233
+ input_size_per_partition // self.pack_factor,
234
+ output_size_per_partition,
235
+ dtype=torch.int32,
236
+ ),
237
+ input_dim=0,
238
+ output_dim=1,
239
+ packed_dim=0,
240
+ packed_factor=self.pack_factor,
241
+ weight_loader=weight_loader,
242
+ )
243
scales = GroupQuantScaleParameter(
244
data=torch.empty(
245
scales_and_zp_size,
@@ -228,6 +250,8 @@ class INCXPULinearBase(INCLinearScheme):
250
output_dim=1,
251
weight_loader=weight_loader,
252
)
253
+ # Both AWQ and GPTQ checkpoints store qzeros with this shape; for
254
+ # symmetric quantization the values are ignored downstream.
255
qzeros = PackedvLLMParameter(
256
data=torch.empty(
257
scales_and_zp_size,
@@ -255,6 +279,37 @@ class INCXPULinearBase(INCLinearScheme):
279
)
280
layer.register_parameter("g_idx", g_idx)
281
282
+ def _convert_awq_qweight_to_gptq(self, qw: torch.Tensor) -> torch.Tensor:
283
+ """Convert AWQ qweight [K, N // pf] to GPTQ qweight [K // pf, N].
284
+
285
+ AWQ packs along the output dim with a non-standard nibble order; GPTQ
286
+ packs along the input dim with sequential nibble order. The conversion
287
+ is lossless — it only reshuffles bits.
288
+ """
289
+ size_bits = self.weight_bits
290
+ pack_factor = self.pack_factor
291
+ mask = (1 << size_bits) - 1
292
+ device = qw.device
293
+ reverse_order = torch.tensor(
294
+ self._REVERSE_AWQ_PACK_ORDER, dtype=torch.long, device=device
295
+ )
296
+ shifts = torch.arange(0, 32, size_bits, dtype=torch.int32, device=device)
297
+
298
+ K, N_packed = qw.shape
299
+ N = N_packed * pack_factor
300
+
301
+ # Unpack int32 → individual values, fix AWQ nibble ordering
302
+ unpacked = (qw.unsqueeze(-1) >> shifts) & mask # (K, N_packed, pf)
303
+ unpacked = unpacked[:, :, reverse_order]
304
+ unpacked = unpacked.reshape(K, N) # (K, N)
305
+
306
+ # Repack along input dim (dim 0) in sequential nibble order
307
+ unpacked = unpacked.reshape(K // pack_factor, pack_factor, N)
308
+ new_qw = (unpacked.to(torch.int32) << shifts[None, :, None]).sum(
309
+ dim=1, dtype=torch.int32
310
+ )
311
+ return new_qw.contiguous()
312
+
313
def create_weights(
314
self,
315
layer: torch.nn.Module,
@@ -276,10 +331,24 @@ class INCXPULinearBase(INCLinearScheme):
331
332
333
class INCXPULinearMethod(INCXPULinearBase):
334
+ """XPU linear method for INC w4a16 quantization (symmetric only).
335
+
336
+ Supports both GPTQ-packed (``auto_round:auto_gptq``) and AWQ-packed
337
+ (``auto_round:auto_awq``) AutoRound checkpoints. AWQ-packed qweights are
338
+ losslessly repacked into the GPTQ-style nibble layout during
339
+ ``process_weights_after_loading``, before the final oneDNN "NT" transpose
340
+ that ``torch.ops._xpu_C.int4_gemm_w4a16`` expects.
341
+ """
342
+
343
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
344
device = layer.qweight.data.device
345
282
- qweight_ct = layer.qweight.data.t().contiguous()
346
+ qweight_data = layer.qweight.data
347
+ if self.is_awq_packed:
348
+ # Lossless repack: AWQ [K, N // pf] → GPTQ [K // pf, N]
349
+ qweight_data = self._convert_awq_qweight_to_gptq(qweight_data)
350
+
351
+ qweight_ct = qweight_data.t().contiguous()
352
layer.qweight = Parameter(qweight_ct.t(), requires_grad=False)
353
layer.scales = Parameter(layer.scales.data, requires_grad=False)
354
layer.qzeros = Parameter(
@@ -370,7 +439,11 @@ class INCARKLinearMethod(INCXPULinearBase):
439
ark_linear.to(layer.qweight.device)
440
441
with torch.no_grad():
373
- ark_linear.qweight.copy_(layer.qweight.detach())
442
+ qweight_src = layer.qweight.detach()
443
+ if self.is_awq_packed:
444
+ # ARK consumes GPTQ-style packed nibbles; convert AWQ losslessly.
445
+ qweight_src = self._convert_awq_qweight_to_gptq(qweight_src)
446
+ ark_linear.qweight.copy_(qweight_src)
447
if hasattr(layer, "qzeros") and layer.qzeros is not None:
448
ark_linear.qzeros.copy_(layer.qzeros.detach())
449
else: