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 7, 17:45 UTC
[CPU][Feat] Enable KleidiAI INT8_W4A8 for all input dtypes (#34890)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Fadi Arafeh committed
Feb 26, 2026 at 05:00 UTC
4171ff6dd9ce18f452c4e9267f5bf090c0989b04
1 file changed
+29
-3
vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py
+29
-3
@@ -42,12 +42,13 @@ class Dynamic4bitLinearKernel(MPLinearKernel):
42
not in [
43
torch.float32,
44
torch.bfloat16,
45
+ torch.float16,
46
]
47
):
48
return (
49
False,
50
"Dynamic4bitLinearKernel on Arm requires Float32 or"
50
- " BFloat16 activations",
51
+ " BFloat16 or Float16 activations",
52
)
53
if c.full_weight_shape[0] % c.group_size != 0:
54
return (
@@ -118,8 +119,30 @@ class Dynamic4bitLinearKernel(MPLinearKernel):
119
x: torch.Tensor,
120
bias: torch.Tensor | None = None,
121
) -> torch.Tensor:
122
+ # PyTorch / KleidiAI kernels natively support the following configs:
123
+ # - channelwise with bfloat16 / float32 activations
124
+ # - groupwise with float32 activations
125
+ # To support:
126
+ # - groupwise with bfloat16/float16 activations: we need to upcast
127
+ # activations to float32 before matmul and downcast back to bfloat16/float16
128
+ # - channelwise with float16 activations, we need to upcast activations to
129
+ # float32 before matmul and downcast back to float16
130
+ # Note: these activations will be dynamically quantized to int8 by the kernel.
131
+
132
c = self.config
133
+ is_groupwise = c.group_size != c.partition_weight_shape[0]
134
+ # dtype of activations before they get dynamically quantized to int8
135
+ original_pre_quant_act_dtype = x.dtype
136
+ pre_quant_act_dtype = original_pre_quant_act_dtype
137
+ if (
138
+ is_groupwise and pre_quant_act_dtype == torch.bfloat16
139
+ ) or pre_quant_act_dtype == torch.float16:
140
+ pre_quant_act_dtype = torch.float32
141
+
142
x_2d = x.reshape(-1, x.shape[-1])
143
+ if pre_quant_act_dtype != original_pre_quant_act_dtype:
144
+ x_2d = x_2d.to(pre_quant_act_dtype)
145
+
146
out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
147
148
w_q = getattr(layer, self.w_q_name)
@@ -129,5 +152,8 @@ class Dynamic4bitLinearKernel(MPLinearKernel):
152
c.group_size,
153
c.partition_weight_shape[0],
154
c.partition_weight_shape[1],
132
- )
133
- return output.reshape(out_shape)
155
+ ).reshape(out_shape)
156
+
157
+ if pre_quant_act_dtype != original_pre_quant_act_dtype:
158
+ output = output.to(original_pre_quant_act_dtype)
159
+ return output