[experimental][kleidi] Allow weight zeros to be a nullptr

pytorch · Oct 10, 2024 · b42fd69 · b42fd69
1 parent 24524a4
commit b42fd69
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp_qsi4c32p.h b/torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp_qsi4c32p.h
@@ -99,8 +99,15 @@ void prepare_weight_data(
   // TODO SIMDify this
   size_t n_groups = n * k / group_size;
   auto weight_scales_bf16 = std::vector<uint16_t>(n_groups, 0);
+
+  // We don't support weight zeros yet
+  if (weight_zeros != nullptr) {
+    for (size_t i = 0; i < n_groups; i++) {
+      assert(weight_zeros[i] == 0);
+    }
+  }
+
   for (size_t i = 0; i < n_groups; i++) {
-    assert(weight_zeros[i] == 0);
     weight_scales_bf16[i] = get_bf16_from_float(weight_scales[i]);
   }