Add flag to disable MFA GEMM.

liuliu · Nov 9, 2023 · a2f75ed · a2f75ed
1 parent 8a11643
commit a2f75ed
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/lib/nnc/ccv_nnc.h b/lib/nnc/ccv_nnc.h
@@ -32,6 +32,7 @@ enum {
 	CCV_NNC_DISABLE_MIXED_MPS_SOFTMAX = 0x2,
 	CCV_NNC_DISABLE_MMAP_MTL_BUFFER = 0x4,
 	CCV_NNC_DISABLE_METAL_FLASH_ATTENTION = 0x8,
+	CCV_NNC_DISABLE_MFA_GEMM = 0x16,
 };
 /**
  * Enable system-wide specific flag.

diff --git a/lib/nnc/cmd/blas/mps/ccv_nnc_gemm_mps.m b/lib/nnc/cmd/blas/mps/ccv_nnc_gemm_mps.m
@@ -176,7 +176,7 @@ static int _ccv_nnc_gemm_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint
 
 		ccv_nnc_mfa_context_t* context = ccv_nnc_default_mfa_context();
 		const int is_mfa_supported =
-			ccv_nnc_mfa_context_supported(context) && is_contiguous && is_same_dtype && is_supported_dtype && (!is_batched || is_mfa_compatible_batch) && !(ccv_nnc_flags() & CCV_NNC_DISABLE_METAL_FLASH_ATTENTION);
+			ccv_nnc_mfa_context_supported(context) && is_contiguous && is_same_dtype && is_supported_dtype && (!is_batched || is_mfa_compatible_batch) && !(ccv_nnc_flags() & CCV_NNC_DISABLE_METAL_FLASH_ATTENTION) && !(ccv_nnc_flags() & CCV_NNC_DISABLE_MFA_GEMM);
 
 		size_t a_data_size = 0;
 		if (CCV_GET_DATA_TYPE(a->info.datatype) == CCV_QX)