MatmulNBits prepacking scales fix

onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc

@@ -359,9 +359,9 @@ Status MatMulNBits<MLFloat16>::PrePack(const Tensor& tensor, int input_idx, /*ou
     OpKernel::Info().TryGetConstantInput(InputIndex::scales, &scales);
     if (scales && MlasQNBitGemmScalesPacked(K_, nbits_, block_size_, compute_type_, has_zp_input_)) {
       auto sptr = scales->Data<MLFloat16>();
-      auto tensor_size = static_cast<size_t>(tensor.Shape().Size());
-      auto ptr = IAllocator::MakeUniquePtr<float>(alloc, tensor_size, true);
-      MlasConvertHalfToFloatBuffer(sptr, ptr.get(), tensor_size);
+      auto scales_size = static_cast<size_t>(scales->Shape().Size());
+      auto ptr = IAllocator::MakeUniquePtr<float>(alloc, scales_size, true);
+      MlasConvertHalfToFloatBuffer(sptr, ptr.get(), scales_size);
       scales_fp32_ = std::move(ptr);
     }
 

onnxruntime/test/shared_lib/test_inference.cc

@@ -341,6 +341,7 @@ static constexpr PATH_TYPE MODEL_URI = TSTR("testdata/mul_1.onnx");
 static constexpr PATH_TYPE CUDA_GRAPH_ANNOTATION_MODEL_URI = TSTR("testdata/mul_1_dynamic.onnx");
 #endif
 static constexpr PATH_TYPE MATMUL_MODEL_URI = TSTR("testdata/matmul_1.onnx");
+static constexpr PATH_TYPE PR_27412_FP16A_INT4B_MODEL_URI = TSTR("testdata/pr_27412_fp16A_int4B_model.onnx");
 #ifndef ORT_NO_RTTI
 static constexpr PATH_TYPE SEQUENCE_MODEL_URI = TSTR("testdata/sequence_length.onnx");
 #endif
@@ -624,6 +625,13 @@ TEST(CApiTest, TestDanglingInput) {
   ASSERT_TRUE(null_present);
 }
 
+// This test tests successful load of a fp16 A and Int4 B quantized model
+// which caused an issue before PR #27412 - there was a crash on Mac M4 while trying to pack scales
+TEST(CApiTest, LoadPr27412Fp16AInt4BModel) {
+  Ort::SessionOptions session_options;
+  EXPECT_NO_THROW(Ort::Session session(*ort_env, PR_27412_FP16A_INT4B_MODEL_URI, session_options));
+}
+
 #if !defined(DISABLE_SPARSE_TENSORS)
 TEST(CApiTest, SparseOutputModel) {
   std::vector<int64_t> dense_shape{3, 3};