[XPU] Route AWQ linear through choose_mp_linear_kernel (#50840)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Co-authored-by: opencode <noreply@opencode.ai> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
zofia committed
Aug 6, 2026 at 13:50 UTC
b976e980efdfb95d0ade2ff6658ca736c737f1bf
2 files changed
+11
-71
.buildkite/intel_jobs/quantization.yaml
+2
-1
@@ -24,5 +24,6 @@ steps:
24
# - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
25
- >-
26
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
27
- 'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]"'
27
+ 'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]"
28
+ && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_auto_awq.py'
29
vllm/model_executor/layers/quantization/auto_awq.py
+9
-70
@@ -296,13 +296,12 @@ class AutoAWQConfig(QuantizationConfig):
296
):
297
return UnquantizedLinearMethod()
298
299
- # Check if XPU - use XPU-specific linear method
300
- if current_platform.is_xpu():
301
- return AutoAWQXPULinearMethod(self)
302
-
303
- # On CPU, use Marlin linear method which uses choose_mp_linear_kernel
304
- # to select the best available kernel (CPUWNA16LinearKernel on CPU)
305
- if current_platform.is_cpu():
299
+ # On CPU and XPU, use the Marlin linear method which uses
300
+ # choose_mp_linear_kernel to select the best available kernel
301
+ # (CPUWNA16LinearKernel on CPU, XPUwNa16LinearKernel on XPU).
302
+ # The AWQ checkpoint is converted to the standard GPTQ-like format
303
+ # in process_weights_after_loading before being handed to the kernel.
304
+ if current_platform.is_cpu() or current_platform.is_xpu():
305
return AutoAWQMarlinLinearMethod(self)
306
307
# Check if Marlin is supported and not using batch invariant mode
@@ -428,8 +427,9 @@ class AutoAWQMarlinLinearMethod(LinearMethodBase):
427
self.quant_type = scalar_types.uint4
428
self.input_dtype = None
429
431
- # Skip Marlin verification on CPU - it will use CPUWNA16LinearKernel
432
- if not current_platform.is_cpu():
430
+ # Skip Marlin verification on CPU/XPU - they use dedicated WNA16
431
+ # kernels (CPUWNA16LinearKernel / XPUwNa16LinearKernel) instead.
432
+ if not (current_platform.is_cpu() or current_platform.is_xpu()):
433
verify_marlin_supported(
434
quant_type=self.quant_config.quant_type,
435
group_size=self.quant_config.group_size,
@@ -956,64 +956,3 @@ class AutoAWQLinearMethod(BaseAWQLinearMethod):
956
if bias is not None:
957
out.add_(bias)
958
return out.reshape(out_shape)
959
-
960
-
961
-class AutoAWQXPULinearMethod(BaseAWQLinearMethod):
962
- """Linear method for AWQ on XPU using int4 GEMM kernel.
963
-
964
- Args:
965
- quant_config: The AWQ quantization config.
966
- """
967
-
968
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
969
- layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
970
- layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
971
- layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
972
-
973
- try:
974
- from vllm_xpu_kernels.quantization._quantize_convert import (
975
- AWQUtils,
976
- transpose_onednn_woq_format,
977
- )
978
- except ImportError as e:
979
- raise ImportError(
980
- "XPU AWQ requires vllm-xpu-kernels. "
981
- "Please install it with: pip install vllm-xpu-kernels"
982
- ) from e
983
-
984
- layer.xpu_output_size = layer.qweight.size(1) * self.quant_config.pack_factor
985
- qweight_new, qzeros_new = AWQUtils.repack(layer.qweight, layer.qzeros)
986
- if qweight_new.shape != layer.qweight.data.shape:
987
- layer.qweight.data = layer.qweight.data.view_as(qweight_new)
988
- if qzeros_new.shape != layer.qzeros.data.shape:
989
- layer.qzeros.data = layer.qzeros.data.view_as(qzeros_new)
990
- layer.qweight.data.copy_(qweight_new)
991
- layer.qzeros.data.copy_(qzeros_new)
992
- transpose_onednn_woq_format(layer, "awq", False)
993
-
994
- def _get_group_size(self, layer: torch.nn.Module) -> int:
995
- """Get the effective group size for kernel computation."""
996
- if self.quant_config.group_size != -1:
997
- return self.quant_config.group_size
998
- return layer.qweight.shape[0] # input_size_per_partition
999
-
1000
- def apply(
1001
- self,
1002
- layer: torch.nn.Module,
1003
- x: torch.Tensor,
1004
- bias: torch.Tensor | None = None,
1005
- ) -> torch.Tensor:
1006
- reshaped_x = x.reshape(-1, x.shape[-1])
1007
- group_size = self._get_group_size(layer)
1008
-
1009
- out = torch.ops._xpu_C.int4_gemm_w4a16(
1010
- reshaped_x,
1011
- layer.qweight,
1012
- bias,
1013
- layer.scales,
1014
- layer.qzeros,
1015
- group_size,
1016
- None,
1017
- )
1018
- out_shape = x.shape[:-1] + (layer.xpu_output_size,)
1019
- return out.reshape(out_shape)