Skip to content

Commit b42fd69

Browse files
committed
[experimental][kleidi] Allow weight zeros to be a nullptr
1 parent 24524a4 commit b42fd69

File tree

1 file changed

+8
-1
lines changed

1 file changed

+8
-1
lines changed

torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp_qsi4c32p.h

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,8 +99,15 @@ void prepare_weight_data(
9999
// TODO SIMDify this
100100
size_t n_groups = n * k / group_size;
101101
auto weight_scales_bf16 = std::vector<uint16_t>(n_groups, 0);
102+
103+
// We don't support weight zeros yet
104+
if (weight_zeros != nullptr) {
105+
for (size_t i = 0; i < n_groups; i++) {
106+
assert(weight_zeros[i] == 0);
107+
}
108+
}
109+
102110
for (size_t i = 0; i < n_groups; i++) {
103-
assert(weight_zeros[i] == 0);
104111
weight_scales_bf16[i] = get_bf16_from_float(weight_scales[i]);
105112
}
106113

0 commit comments

Comments
 (0)