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, 21:30 UTC
[CPU][BugFix] Allow w8a8 oneDNN quantized matmul to support 3D inputs (#33727)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Feb 5, 2026 at 06:26 UTC
fd03538bf97cd7f4fedd6da4584c89635878174f
1 file changed
+3
-1
vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
+3
-1
@@ -182,6 +182,8 @@ class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
182
x: torch.Tensor,
183
bias: torch.Tensor | None = None,
184
) -> torch.Tensor:
185
+ x_shape = x.shape
186
+ x = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x
187
w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
188
189
# ops.scaled_int8_quant supports both dynamic and static quant:
@@ -195,7 +197,7 @@ class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
197
n = self.dnnl_handler.n
198
out = torch.empty((m, n), dtype=x.dtype)
199
ops.onednn_scaled_mm(self.dnnl_handler, x_q, out, x_s, x_zp, azp_adj, bias)
198
-
200
+ out = out.reshape(x_shape[:-1] + (n,)) if len(x_shape) > 2 else out
201
return out
202
203
def _apply_weights_sgl(