iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[XPU] Route AWQ linear through choose_mp_linear_kernel (#50840)

Browse Source 
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>
This commit is contained in:
zofia
2026-08-06 13:50:11 +08:00
committed by GitHub
co-authored by opencode Kunshang Ji mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 2a0323b4e0
commit b976e980ef
2 changed files with 11 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/intel_jobs/quantization.yaml
+2 -1
View File
@@ -24,5 +24,6 @@ steps:
# - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'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]"'
'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]"
&& VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_auto_awq.py'
vllm/model_executor/layers/quantization/auto_awq.py
+9 -70
View File
@@ -296,13 +296,12 @@ class AutoAWQConfig(QuantizationConfig):
):
return UnquantizedLinearMethod()
# Check if XPU - use XPU-specific linear method
if current_platform.is_xpu():
return AutoAWQXPULinearMethod(self)
# On CPU, use Marlin linear method which uses choose_mp_linear_kernel
# to select the best available kernel (CPUWNA16LinearKernel on CPU)
if current_platform.is_cpu():
# On CPU and XPU, use the Marlin linear method which uses
# choose_mp_linear_kernel to select the best available kernel
# (CPUWNA16LinearKernel on CPU, XPUwNa16LinearKernel on XPU).
# The AWQ checkpoint is converted to the standard GPTQ-like format
# in process_weights_after_loading before being handed to the kernel.
if current_platform.is_cpu() or current_platform.is_xpu():
return AutoAWQMarlinLinearMethod(self)
# Check if Marlin is supported and not using batch invariant mode
@@ -428,8 +427,9 @@ class AutoAWQMarlinLinearMethod(LinearMethodBase):
self.quant_type = scalar_types.uint4
self.input_dtype = None
# Skip Marlin verification on CPU - it will use CPUWNA16LinearKernel
if not current_platform.is_cpu():
# Skip Marlin verification on CPU/XPU - they use dedicated WNA16
# kernels (CPUWNA16LinearKernel / XPUwNa16LinearKernel) instead.
if not (current_platform.is_cpu() or current_platform.is_xpu()):
verify_marlin_supported(
quant_type=self.quant_config.quant_type,
group_size=self.quant_config.group_size,
@@ -956,64 +956,3 @@ class AutoAWQLinearMethod(BaseAWQLinearMethod):
if bias is not None:
out.add_(bias)
return out.reshape(out_shape)
class AutoAWQXPULinearMethod(BaseAWQLinearMethod):
"""Linear method for AWQ on XPU using int4 GEMM kernel.
Args:
quant_config: The AWQ quantization config.
"""
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
try:
from vllm_xpu_kernels.quantization._quantize_convert import (
AWQUtils,
transpose_onednn_woq_format,
)
except ImportError as e:
raise ImportError(
"XPU AWQ requires vllm-xpu-kernels. "
"Please install it with: pip install vllm-xpu-kernels"
) from e
layer.xpu_output_size = layer.qweight.size(1) * self.quant_config.pack_factor
qweight_new, qzeros_new = AWQUtils.repack(layer.qweight, layer.qzeros)
if qweight_new.shape != layer.qweight.data.shape:
layer.qweight.data = layer.qweight.data.view_as(qweight_new)
if qzeros_new.shape != layer.qzeros.data.shape:
layer.qzeros.data = layer.qzeros.data.view_as(qzeros_new)
layer.qweight.data.copy_(qweight_new)
layer.qzeros.data.copy_(qzeros_new)
transpose_onednn_woq_format(layer, "awq", False)
def _get_group_size(self, layer: torch.nn.Module) -> int:
"""Get the effective group size for kernel computation."""
if self.quant_config.group_size != -1:
return self.quant_config.group_size
return layer.qweight.shape[0] # input_size_per_partition
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
reshaped_x = x.reshape(-1, x.shape[-1])
group_size = self._get_group_size(layer)
out = torch.ops._xpu_C.int4_gemm_w4a16(
reshaped_x,
layer.qweight,
bias,
layer.scales,
layer.qzeros,
group_size,
None,
)
out_shape = x.shape[:-1] + (layer.xpu_output_size,)
return out.reshape(out_shape)