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 6, 09:30 UTC
[Feat] Enable CompressedTensorW4A8Int for XPU (#37207)
Signed-off-by: Li, Tianmu <tianmu.li@intel.com>
Tianmu Li committed
Mar 19, 2026 at 19:34 UTC
47b7af0d87705f2e086ea0bc9d915fc7510e8e2f
4 files changed
+172
vllm/_xpu_ops.py
+54
@@ -37,6 +37,26 @@ if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"):
37
return torch.empty((M, N), dtype=input.dtype, device=input.device)
38
39
40
+if hasattr(torch.ops._xpu_C, "int4_gemm_w4a8"):
41
+
42
+ @register_fake("_xpu_C::int4_gemm_w4a8")
43
+ def _int4_gemm_w4a8_fake(
44
+ input: torch.Tensor,
45
+ input_scales: torch.Tensor,
46
+ input_zero_points: torch.Tensor,
47
+ q_weight: torch.Tensor,
48
+ weight_scale: torch.Tensor,
49
+ weight_zp: torch.Tensor,
50
+ group_size: int,
51
+ g_idx: torch.Tensor | None = None,
52
+ bias: torch.Tensor | None = None,
53
+ ) -> torch.Tensor:
54
+ input_2d = input.view(-1, input.shape[-1])
55
+ M = input_2d.size(0)
56
+ N = q_weight.size(1)
57
+ return torch.empty((M, N), dtype=torch.float16, device=input.device)
58
+
59
+
60
if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
61
62
@register_fake("_xpu_C::int4_gemm_w4a16")
@@ -87,6 +107,40 @@ _OPS_REGISTERED = False
107
108
109
class xpu_ops:
110
+ @staticmethod
111
+ @torch.compile
112
+ def dynamic_per_token_int8_quant_ref(
113
+ input: torch.Tensor, use_sym_quant: bool, bits: int
114
+ ):
115
+ original_sizes = input.size()
116
+ # view is not safe in torch.compile if input is not contiguous
117
+ input = input.reshape(
118
+ -1, original_sizes[-1]
119
+ ) # Flatten except for the last dimension
120
+ qmin = -(2 ** (bits - 1)) if use_sym_quant else 0
121
+ qmax = 2 ** (bits - 1) - 1 if use_sym_quant else 2**bits - 1
122
+ min_val = torch.min(input, dim=-1)[0].to(dtype=torch.float32).unsqueeze(-1)
123
+ max_val = torch.max(input, dim=-1)[0].to(dtype=torch.float32).unsqueeze(-1)
124
+ if use_sym_quant:
125
+ scale = (
126
+ torch.maximum(torch.abs(min_val), torch.abs(max_val)) / qmax
127
+ ).clamp(min=1e-5)
128
+ zero_point = torch.zeros_like(scale).to(dtype=torch.int32)
129
+ else:
130
+ scale = ((max_val - min_val) / qmax).clamp(min=1e-5)
131
+ zero_point = -1 * torch.round(min_val / scale).to(dtype=torch.int32)
132
+ scale = scale.to(dtype=input.dtype)
133
+ quantized = torch.clamp(
134
+ torch.round(input / scale.to(dtype=torch.float32) + zero_point),
135
+ qmin,
136
+ qmax,
137
+ ).to(dtype=torch.int8 if use_sym_quant else torch.uint8)
138
+ return (
139
+ quantized.view(original_sizes),
140
+ scale.view(original_sizes[:-1] + (1,)),
141
+ zero_point.view(original_sizes[:-1] + (1,)),
142
+ )
143
+
144
@staticmethod
145
def flash_attn_varlen_func(
146
q: torch.Tensor,
vllm/model_executor/kernels/linear/__init__.py
+3
@@ -48,6 +48,7 @@ from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
48
MarlinLinearKernel,
49
)
50
from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
51
+ XPUW4A8IntLinearKernel,
52
XPUwNa16LinearKernel,
53
)
54
from vllm.model_executor.kernels.linear.scaled_mm import (
@@ -138,6 +139,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
139
ExllamaLinearKernel,
140
],
141
PlatformEnum.XPU: [
142
+ XPUW4A8IntLinearKernel,
143
XPUwNa16LinearKernel,
144
],
145
PlatformEnum.CPU: [
@@ -391,5 +393,6 @@ __all__ = [
393
"ExllamaLinearKernel",
394
"MacheteLinearKernel",
395
"MarlinLinearKernel",
396
+ "XPUW4A8IntLinearKernel",
397
"XPUwNa16LinearKernel",
398
]
vllm/model_executor/kernels/linear/mixed_precision/__init__.py
+2
@@ -30,6 +30,7 @@ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
30
MPLinearLayerConfig,
31
)
32
from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
33
+ XPUW4A8IntLinearKernel,
34
XPUwNa16LinearKernel,
35
)
36
@@ -44,5 +45,6 @@ __all__ = [
45
"ExllamaLinearKernel",
46
"MacheteLinearKernel",
47
"MarlinLinearKernel",
48
+ "XPUW4A8IntLinearKernel",
49
"XPUwNa16LinearKernel",
50
]
vllm/model_executor/kernels/linear/mixed_precision/xpu.py
+113
@@ -5,6 +5,8 @@
5
import torch
6
from torch.nn.parameter import Parameter
7
8
+from vllm.logger import init_logger
9
+from vllm.model_executor.layers.quantization.utils import replace_parameter
10
from vllm.platforms import current_platform
11
from vllm.scalar_type import scalar_types
12
@@ -12,6 +14,8 @@ from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
14
15
_XPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)
16
17
+logger = init_logger(__name__)
18
+
19
20
class XPUwNa16LinearKernel(MPLinearKernel):
21
@classmethod
@@ -86,3 +90,112 @@ class XPUwNa16LinearKernel(MPLinearKernel):
90
layer.g_idx,
91
)
92
return out
93
+
94
+
95
+class XPUW4A8IntLinearKernel(MPLinearKernel):
96
+ """XPU kernel for W4A8 integer quantization using oneDNN int4_gemm_w4a8.
97
+
98
+ Weights are symmetric group-quantized int4 packed as uint4.
99
+ Activations are dynamically quantized per-token to symmetric int8.
100
+ """
101
+
102
+ @classmethod
103
+ def get_min_capability(cls) -> int:
104
+ return -1
105
+
106
+ @classmethod
107
+ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
108
+ if not current_platform.is_xpu():
109
+ return False, "XPUW4A8Int only supported on XPU"
110
+ if c.act_type not in (torch.bfloat16, torch.float16):
111
+ return False, "XPUW4A8Int requires BF16/FP16 activations"
112
+ if c.weight_type != scalar_types.int4:
113
+ return (
114
+ False,
115
+ f"XPUW4A8Int requires int4 weights, got {c.weight_type}",
116
+ )
117
+ if c.zero_points:
118
+ return False, "XPUW4A8Int only supports symmetric weight quantization"
119
+ if c.group_size != -1 and c.group_size % 32 != 0:
120
+ return (
121
+ False,
122
+ f"Group size ({c.group_size}) not supported by XPUW4A8Int, "
123
+ "must be a multiple of 32",
124
+ )
125
+ in_size, out_size = c.partition_weight_shape
126
+ if in_size % 8 != 0 or out_size % 8 != 0:
127
+ return (
128
+ False,
129
+ f"in/out sizes ({in_size}, {out_size}) must be multiples of 8",
130
+ )
131
+
132
+ if c.act_type != torch.float16:
133
+ logger.warning_once(
134
+ "XPUW4A8IntLinearKernel is running with model dtype %s, "
135
+ "but int4_gemm_w4a8 produces float16 output. Recommend "
136
+ "setting --dtype float16 for best performance.",
137
+ c.act_type,
138
+ )
139
+
140
+ return True, None
141
+
142
+ def _pack_int4_weight(self, w: torch.Tensor) -> torch.Tensor:
143
+ # w is [N, K] int8 with values in [-8, 7]
144
+ w_u4 = w.to(torch.int32) + 8 # shift to [0, 15]
145
+ w_u4 = w_u4.reshape(w.shape[0], w.shape[1] // 8, 8) # [N, K/8, 8]
146
+ shifts = torch.arange(0, 32, 4, dtype=torch.int32, device=w.device)
147
+ packed = ((w_u4 & 0xF) << shifts[None, None, :]).sum(dim=2).to(torch.int32)
148
+ return packed
149
+
150
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
151
+ layer.weight_scale.data = layer.weight_scale.data.t().contiguous()
152
+
153
+ device = layer.weight_packed.device
154
+ # TODO: support asymmetric quantization
155
+ weight_zero_point = torch.tensor([8], dtype=torch.int8, device=device)
156
+ layer.weight_zero_point = Parameter(weight_zero_point, requires_grad=False)
157
+
158
+ # weight_packed is [out, in] int8, signed int4 values in [-8, 7]
159
+ w = layer.weight_packed.data # [out, in]
160
+
161
+ # TODO: implement asym case
162
+ packed = self._pack_int4_weight(w) # [out, in/8] packed uint4
163
+
164
+ replace_parameter(
165
+ layer,
166
+ self.w_q_name,
167
+ torch.nn.Parameter(packed, requires_grad=False),
168
+ )
169
+
170
+ # Free the original unpacked int8 weight (still registered as "weight")
171
+ # to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory.
172
+ layer.register_parameter("weight", None)
173
+
174
+ def apply_weights(
175
+ self,
176
+ layer: torch.nn.Module,
177
+ x: torch.Tensor,
178
+ bias: torch.Tensor | None = None,
179
+ ) -> torch.Tensor:
180
+ reshaped_x = x.reshape(-1, x.shape[-1]) # [M, K]
181
+ from vllm._xpu_ops import xpu_ops as ops
182
+
183
+ # TODO: static and asymmetric quantization case
184
+ # Common code for CompressedTensorsW4A8Int does not read act symmetry data
185
+ quant_x, x_scale, x_zero = ops.dynamic_per_token_int8_quant_ref(
186
+ reshaped_x, True, 8
187
+ )
188
+
189
+ out = torch.ops._xpu_C.int4_gemm_w4a8(
190
+ quant_x,
191
+ x_scale,
192
+ x_zero,
193
+ layer.weight_packed.t(),
194
+ layer.weight_scale,
195
+ layer.weight_zero_point,
196
+ self.config.group_size,
197
+ None, # g_idx not currently supported
198
+ bias,
199
+ )
200
+
201
+ return out.to(x.dtype)