disable MHA

intel · Mar 13, 2024 · f2e9171 · f2e9171
1 parent 863859b
commit f2e9171
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/neural_speed/core/layers/mha_dense.cpp b/neural_speed/core/layers/mha_dense.cpp
@@ -72,7 +72,9 @@ bool bestla_reordered_attn_fp32_support(const attn_shape_t* params) {
   // TODO(Yi): check K V's layout
   if (_cd->AMX_BF16()) return true;
 #endif
-  return !_cd->AVX512F() || _cd->AVX2();  // use avx2 and f16c on avx2 platforms
+  // use avx2 and f16c on avx2 platforms
+  // todo: check avx2 mha on sever
+  return !_cd->AVX512F() && _cd->AVX2();
 }
 // kv cache sizes in bytes per layer per batch per beam for;
 void bestla_reordered_attn_fp32_batch_kv_info(const kv_shape_t* params, kv_cache_info_t* out) {