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 9, 06:15 UTC
[CPU][Feat] Enable KleidiAI accelerated int4 dynamic quant with BF16 activations on Arm CPUs (#33122)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Jan 31, 2026 at 07:16 UTC
1618e2549281324100690db3bd8c5f497fa4fc3f
1 file changed
+18
-1
vllm/model_executor/layers/quantization/kernels/mixed_precision/dynamic_4bit.py
+18
-1
@@ -11,6 +11,18 @@ from vllm.scalar_type import scalar_types
11
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
12
13
14
+# This implementation is for the KleidiAI-accelerated w4a8int quantization
15
+# scheme on Arm CPUs:
16
+# torch.ops.aten._dyn_quant_matmul_4bit performs dynamic quantized matmul
17
+# it takes:
18
+# - int4 weights packed along with bias/scales by
19
+# torch.ops.aten._dyn_quant_pack_4bit_weight
20
+# - float32/bfloat16 activations
21
+# then it leverages KleidiAI ukernels that:
22
+# - dynamically quantize the activations to int8
23
+# - unpack the int4 weights to int8
24
+# - perform int8 x int8 -> int32 matmul
25
+# - dequantize the int32 output to float32/bfloat16 outputs
26
class Dynamic4bitLinearKernel(MPLinearKernel):
27
SUPPORTED_QUANT_TYPES = [scalar_types.int4]
28
@@ -29,9 +41,14 @@ class Dynamic4bitLinearKernel(MPLinearKernel):
41
and c.act_type
42
not in [
43
torch.float32,
44
+ torch.bfloat16,
45
]
46
):
34
- return False, "Dynamic4bitLinearKernel on Arm requires Float32 activations"
47
+ return (
48
+ False,
49
+ "Dynamic4bitLinearKernel on Arm requires Float32 or"
50
+ " BFloat16 activations",
51
+ )
52
if c.full_weight_shape[0] % c.group_size != 0:
53
return (
54
False,