[FIX] Disable the scaled dot product attention fusion in ATSM platfor…

…m. (#2458)
intel · Oct 25, 2023 · 2ba07bb · 2ba07bb
1 parent a28e0c6
commit 2ba07bb
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 8 deletions.
diff --git a/itex/core/graph/remapper/mha_pattern.cc b/itex/core/graph/remapper/mha_pattern.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include "itex/core/graph/config_util.h"
 #include "itex/core/graph/remapper/constant_names.h"
 #include "itex/core/graph/remapper/fusion.h"
 #include "itex/core/graph/remapper/remapper.h"
@@ -85,6 +86,11 @@ class MHAFusionWithReshapeMatmul : public Fusion {
     MatchedProperties ret =
         FillProperties(&graph_view, graph_view.GetNode(node_index), pattern_);
 
+#ifndef INTEL_CPU_ONLY
+    if (!isxehpc_value) {
+      return ret.ToEmpty();
+    }
+#endif
     bool is_ok = !ret.Empty() && CheckShapes(ctx, ret);
 
     if (!is_ok) return ret.ToEmpty();
@@ -267,6 +273,11 @@ class MHAPatternWithMulAndAdd : public Fusion {
     MatchedProperties ret = FillProperties(
         &graph_view, graph_view.GetNode(node_index), pattern_, false);
 
+#ifndef INTEL_CPU_ONLY
+    if (!isxehpc_value) {
+      return ret.ToEmpty();
+    }
+#endif
     bool is_ok = !ret.Empty() && CheckShapes(ctx, ret);
 
     if (!is_ok) return ret.ToEmpty();

diff --git a/test/pattern/flash_attention.py b/test/pattern/flash_attention.py
@@ -14,21 +14,17 @@
 # ==============================================================================
 
 
+import os
+
 import numpy as np
 import tensorflow as tf
-from keras import backend as K
+
 from tensorflow import keras
 from tensorflow.python.framework import test_util
 from tensorflow.python.framework import config
-from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
 from tensorflow.core.protobuf import config_pb2
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import nn_ops
-import time
-import os
-import subprocess
-import sys
+from intel_extension_for_tensorflow.python.device import is_xehpc
 
 tf.compat.v1.disable_eager_execution()
 
@@ -81,6 +77,8 @@ def td_dot(self, a, b):
 class MHAFusionWithReshapeMatmulTest(test_util.TensorFlowTestCase):
 
     def testMHAFusionWithReshapeMatmul(self):
+        if config.list_logical_devices('XPU') and not is_xehpc():
+            self.skipTest("Only xehpc support xetla.")
         tf.random.set_seed(0)
         datatypes = [tf.float32, tf.bfloat16]
         if config.list_logical_devices('XPU'):
@@ -116,6 +114,8 @@ def testMHAFusionWithReshapeMatmul(self):
 class MHAPatternWithMulAndAddTest(test_util.TensorFlowTestCase):
 
     def testMHAPatternWithMulAndAdd(self):
+        if config.list_logical_devices('XPU') and not is_xehpc():
+            self.skipTest("Only xehpc support xetla.")
         tf.random.set_seed(0)
         datatypes = [tf.float32, tf.bfloat16]
         if config.list_logical_devices('XPU'):