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 8, 17:45 UTC
[XPU][5/N] add wna16 xpu kernel (#33973)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
zofia committed
Feb 6, 2026 at 23:59 UTC
2ce9fe4ad03f5097f4abe4f188166ede8c62df21
2 files changed
+58
-66
.buildkite/scripts/hardware_ci/run-xpu-test.sh
+1
@@ -39,6 +39,7 @@ docker run \
39
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
40
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
41
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
42
+ python3 examples/offline_inference/basic/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager
43
python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2
44
python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
45
cd tests
vllm/model_executor/layers/quantization/kernels/mixed_precision/xpu.py
+57
-66
@@ -3,88 +3,71 @@
3
4
5
import torch
6
+from torch.nn.parameter import Parameter
7
8
from vllm.platforms import current_platform
9
+from vllm.scalar_type import scalar_types
10
11
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
12
13
+_XPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)
14
+
15
16
class XPUwNa16LinearKernel(MPLinearKernel):
17
@classmethod
18
def get_min_capability(cls) -> int:
15
- return 0
19
+ return -1
20
21
@classmethod
22
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
23
if not current_platform.is_xpu():
20
- return False, "IPEX wNa16 only supported on XPU/CPU devices"
24
+ return False, "XPUwNa16 only supported on XPU"
25
+
26
+ if c.act_type != torch.bfloat16 and c.act_type != torch.float16:
27
+ return False, "XPUwNa16 only supports BF16/FP16 activations"
28
22
- # TODO: (yiliu30) relax these restrictions in later PRs
23
- if c.zero_points:
24
- return False, "Zero points not supported for Now"
29
+ if c.weight_type not in _XPUWNA16_SUPPORTED_QUANT_TYPES:
30
+ return (
31
+ False,
32
+ f"Quant type ({c.weight_type}) not supported by "
33
+ "XPUwNa16, supported types are: "
34
+ f"{_XPUWNA16_SUPPORTED_QUANT_TYPES}",
35
+ )
36
+ if c.group_size != -1 and c.group_size % 32 != 0:
37
+ return (
38
+ False,
39
+ f"Group size ({c.group_size}) not supported by "
40
+ "XPUwNa16, supported group sizes are multiples of 32",
41
+ )
42
+
43
+ if c.partition_weight_shape[0] % 32 != 0:
44
+ return (
45
+ False,
46
+ f"Input size ({c.partition_weight_shape[0]}) not supported by "
47
+ "XPUwNa16, supported sizes are multiples of 32",
48
+ )
49
+
50
+ if c.partition_weight_shape[1] % 32 != 0:
51
+ return (
52
+ False,
53
+ f"Output size ({c.partition_weight_shape[1]}) not supported by "
54
+ "XPUWNA16, supported sizes are multiples of 32",
55
+ )
56
57
return True, None
58
28
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
29
- from packaging import version
30
-
31
- MIN_IPEX_VERSION = "2.6.0"
32
- bias = layer.bias if not layer.skip_bias_add else None
33
-
34
- try:
35
- import intel_extension_for_pytorch as ipex
36
-
37
- if version.parse(ipex.__version__) < version.parse(MIN_IPEX_VERSION):
38
- raise ImportError(
39
- "intel_extension_for_pytorch version is "
40
- "wrong. Please install "
41
- f"intel_extension_for_pytorch>={MIN_IPEX_VERSION}."
42
- )
43
- except ImportError as err:
44
- raise ImportError(
45
- "Please install "
46
- f"intel_extension_for_pytorch>={MIN_IPEX_VERSION} via "
47
- f"`pip install intel_extension_for_pytorch>={MIN_IPEX_VERSION}`"
48
- " to use IPEX-AWQ linear method."
49
- ) from err
50
- # Using the compute dtype (lowp_mode) as INT8 to leverage instructions
51
- # with better performance.
52
- lowp_mode = ipex.quantization.WoqLowpMode.INT8
53
- # The weight will be de-packed from INT4 to INT8.
54
- weight_dtype = ipex.quantization.WoqWeightDtype.INT4
55
- # The float activation will be quantized (dynamic, per-token) to INT8.
56
- act_quant_mode = ipex.quantization.WoqActQuantMode.PER_BATCH
57
-
58
- qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
59
- weight_dtype=weight_dtype,
60
- lowp_mode=lowp_mode,
61
- act_quant_mode=act_quant_mode,
62
- group_size=self.config.group_size,
63
- weight_qscheme=ipex.quantization.WoqWeightQScheme.SYMMETRIC,
64
- )
65
- qweight = layer.weight_packed
66
- g_idx = layer.weight_g_idx if self.config.has_g_idx else None
67
- scales = layer.weight_scale
68
- qzeros = None
59
+ def process_weights_after_loading(self, layer: torch.nn.Module):
60
+ layer.weight_scale.data = layer.weight_scale.t().contiguous()
61
+
62
if self.config.zero_points:
70
- qzeros = layer.weight_zero_point.contiguous()
71
- qweight = qweight.t().contiguous()
72
- scales = scales.t().contiguous()
73
- layer.ipex_output_size = self.config.partition_weight_shape[1]
74
- layer.ipex_qlinear = (
75
- ipex.llm.quantization.woq_linear.IPEXWeightOnlyQuantizedLinear.from_weight(
76
- qweight,
77
- scales,
78
- qzeros,
79
- in_features=self.config.partition_weight_shape[0],
80
- out_features=self.config.partition_weight_shape[1],
81
- qconfig=qconfig,
82
- g_idx=g_idx,
83
- bias=bias,
84
- group_size=self.config.group_size,
85
- quant_method=0, # `0` stands for the IPEX GPTQ
86
- )
87
- )
63
+ layer.weight_zero_point.data = layer.weight_zero_point.t().contiguous()
64
+ else:
65
+ weight_zero_point = torch.Tensor([8]).to(torch.int8).to("xpu")
66
+ layer.weight_zero_point = Parameter(weight_zero_point, requires_grad=False)
67
+ if self.config.has_g_idx:
68
+ layer.g_idx.data = layer.g_idx.t().contiguous()
69
+ else:
70
+ layer.g_idx = None
71
72
def apply_weights(
73
self,
@@ -93,5 +76,13 @@ class XPUwNa16LinearKernel(MPLinearKernel):
76
bias: torch.Tensor | None = None,
77
) -> torch.Tensor:
78
reshaped_x = x.reshape(-1, x.shape[-1])
96
- out = layer.ipex_qlinear(reshaped_x)
97
- return out.reshape(x.shape[:-1] + (layer.ipex_output_size,))
79
+ out = torch.ops._xpu_C.int4_gemm_w4a16(
80
+ reshaped_x,
81
+ layer.weight_packed.t(),
82
+ bias,
83
+ layer.weight_scale,
84
+ layer.weight_zero_point,
85
+ self.config.group_size,
86
+ layer.g_idx,
87
+ )
88
+ return out