[DO NOT REVIEW][CPU] Change the order of the SDPA fusion transformations

src/common/transformations/src/transformations/op_conversions/scaled_dot_product_attention_decomposition.cpp

@@ -59,8 +59,7 @@ std::shared_ptr<ov::Node> ov::pass::ScaledDotProductAttentionDecomposition::deco
     auto query = node->input_value(0);
     auto key = node->input_value(1);
     auto value = node->input_value(2);
-    auto q_shape = register_new_node<v3::ShapeOf>(query, element::i32);
-    auto k_shape = register_new_node<v3::ShapeOf>(key, element::i32);
+
     auto minus_one = register_new_node(v0::Constant::create(element::i32, Shape{}, {-1}));
     auto minus_two = register_new_node(v0::Constant::create(element::i32, Shape{}, {-2}));
     auto zero_i = register_new_node(v0::Constant::create(element::i32, Shape{}, {0}));
@@ -70,7 +69,20 @@ std::shared_ptr<ov::Node> ov::pass::ScaledDotProductAttentionDecomposition::deco
 
     Output<Node> scale;
     if (node->get_input_size() < 5) {
-        scale = register_new_node<v8::Gather>(q_shape, minus_one, zero_i)->output(0);
+        auto&& query_shape = query.get_partial_shape();
+        // often the embeddings space size is known, so the dimension may be extracted into a constant
+        if (query_shape.rank().is_static()) {
+            auto&& last_dim = *(query_shape.rbegin());
+            if (last_dim.is_static()) {
+                scale = register_new_node(v0::Constant::create(element::i32, Shape{}, {last_dim.get_length()}));
+            }
+        }
+
+        if (!scale.get_node()) {
+            auto shape = register_new_node<v3::ShapeOf>(query, element::i32);
+            scale = register_new_node<v8::Gather>(shape, minus_one, zero_i);
+        }
+
         scale = register_new_node<v1::ConvertLike>(scale, query);
         auto sqrt_scale = register_new_node<v0::Sqrt>(scale);
         scale = register_new_node<v1::Divide>(one_f, sqrt_scale);
@@ -79,6 +91,7 @@ std::shared_ptr<ov::Node> ov::pass::ScaledDotProductAttentionDecomposition::deco
     }
 
     auto q_scaled = register_new_node<v1::Multiply>(query, scale);
+    auto k_shape = register_new_node<v3::ShapeOf>(key, element::i32);
     auto k_rank = register_new_node<v3::ShapeOf>(k_shape, element::i32)->output(0);
     auto k_last_dim = register_new_node<v1::Add>(k_rank, minus_one);
     auto k_next_dim = register_new_node<v1::Add>(k_rank, minus_two)->output(0);
@@ -112,6 +125,7 @@ std::shared_ptr<ov::Node> ov::pass::ScaledDotProductAttentionDecomposition::deco
                 atten_mask = mask;
             }
         } else {
+            auto q_shape = register_new_node<v3::ShapeOf>(query, element::i32);
             auto target_s_len = register_new_node<v8::Gather>(q_shape, minus_two, zero_i);
             auto source_s_len = register_new_node<v8::Gather>(k_shape, minus_two, zero_i);
             auto ssl = register_new_node<v0::Unsqueeze>(source_s_len, zero_i);

src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp

@@ -49,6 +49,7 @@
 #include "transformations/op_conversions/convert_batch_to_space.hpp"
 #include "transformations/op_conversions/convert_bitwise_to_logical_bool.hpp"
 #include "transformations/op_conversions/convert_broadcast_to_tiles.hpp"
+#include "transformations/op_conversions/convert_convertlike.hpp"
 #include "transformations/op_conversions/convert_depth_to_space.hpp"
 #include "transformations/op_conversions/convert_gather_downgrade.hpp"
 #include "transformations/op_conversions/convert_gather_to_compressed.hpp"
@@ -649,16 +650,6 @@ void Transformations::PreLpt(const std::vector<ov::element::Type>& defaultPrecis
     CPU_SET_CALLBACK_COMMON(manager, nmsCallback, ov::pass::ConvertNMS9ToNMSIEInternal);
     CPU_SET_CALLBACK_COMMON(manager, nmsCallback, ov::pass::ConvertMulticlassNmsToMulticlassNmsIE);
     CPU_SET_CALLBACK_COMMON(manager, nmsCallback, ov::pass::ConvertMatrixNmsToMatrixNmsIE);
-    CPU_SET_CALLBACK_COMMON(
-        manager,
-        [this](const_node_ptr& node) -> bool {
-            std::string errorMsg;
-            // Current SDPA impl is optimized only for LLM models, so we decompose it for others to avoid perf
-            // regression. Matching the pattern is a little complicated, so we just check if there is any state nodes.
-            return node::ScaledDotProductAttention::isSupportedOperation(node, errorMsg) &&
-                   model->get_variables().size() > 0;
-        },
-        ov::pass::ScaledDotProductAttentionDecomposition);
 
     // List of enabled/disabled transformations
 
@@ -692,6 +683,7 @@ void Transformations::PreLpt(const std::vector<ov::element::Type>& defaultPrecis
     CPU_DISABLE_PASS_COMMON(manager, ov::pass::MatMulConstTransposesExtraction);
     CPU_DISABLE_PASS_COMMON(manager, ov::pass::ConvertScatterNDUpdate15ToScatterNDUpdate3);
     CPU_DISABLE_PASS_COMMON(manager, ov::pass::ConvertSliceScatter);
+    CPU_DISABLE_PASS_COMMON(manager, ov::pass::ScaledDotProductAttentionDecomposition);
     CPU_DISABLE_PASS_X64(manager, ov::pass::HSigmoidDecomposition);
 
     CPU_DISABLE_PASS_X64(manager, ov::pass::ReduceL1Decomposition);
@@ -940,8 +932,36 @@ void Transformations::PostLpt() {
 #endif  // OPENVINO_ARCH_X86_64
 
     CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::transpose_sinking::TSShapeOfForward);
-    CPU_REGISTER_PASS_COMMON(postLPTPassManager, StatefulSDPAFusion);
+    CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::intel_cpu::StatefulSDPAFusion);
+
+    // If the SDPA patterns haven't been fused into the special CPU optimized SDPA nodes, we have to decompose
+    // these layers and run some auxilary transformation passes to let the snippets handle SDPA ops
+
+    if (!one_of(config.inferencePrecision, element::bf16, element::f16)) {  // So far Snippets don't support AMX MHA
+        CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::ScaledDotProductAttentionDecomposition);
+        CPU_SET_CALLBACK_COMMON(
+            postLPTPassManager,
+            [](const_node_ptr& node) {
+                // So far Snippets don't support AMX MHA
+                constexpr size_t QKV_inpt_number = 3ul;
+                for (size_t i = 0; i < QKV_inpt_number; ++i) {
+                    if (one_of(node->get_input_element_type(i), element::bf16, element::f16)) {
+                        return true;
+                    }
+                }
+                return false;
+            },
+            ov::pass::ScaledDotProductAttentionDecomposition);
+
+        CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::ConvertConvertLike);
+        CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::ConstantFolding);
+        CPU_REGISTER_PASS_COMMON(postLPTPassManager,
+                                 ov::pass::MoveEltwiseUpThroughDataMovScalar,
+                                 std::vector<DiscreteTypeInfo>{ov::op::v1::Transpose::get_type_info_static()});
+        CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::TransposeMatMul);
+    }
     CPU_REGISTER_PASS_X64(postLPTPassManager, ov::intel_cpu::SDPAFuseTransposeReshape);
+
     CPU_REGISTER_PASS_X64(postLPTPassManager, ov::pass::RMSFusion, false);
     CPU_REGISTER_PASS_X64(postLPTPassManager, ov::intel_cpu::DecomposeRMSNorm);
     CPU_SET_CALLBACK_X64(

src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/x64/fuse_reshape_transpose_to_sdpa.cpp

@@ -205,17 +205,10 @@ class FuseSDPAReshapeTransposeTest : virtual public ov::test::SubgraphBaseTest,
 
 TEST_P(FuseSDPAReshapeTransposeTest, CompareWithRefs) {
     SKIP_IF_CURRENT_TEST_IS_DISABLED();
-    bool reshape_transpose_fused = false;
     auto actualOutputs = run_test(function);
-    CheckNumberOfNodesWithType(compiledModel, "ScaledDotProductAttention", 1);
-    CheckNumberOfNodesWithType(compiledModel, "Reshape", 0);
+    CheckNumberOfNodesWithType(compiledModel, "Subgraph", 1);
+    CheckNumberOfNodesWithType(compiledModel, "Reshape", 4);
     CheckNumberOfNodesWithType(compiledModel, "Transpose", 0);
-    for (const auto& n : compiledModel.get_runtime_model()->get_ordered_ops()) {
-        if (n->get_friendly_name() == "mha/fused_reshape_transpose") {
-            reshape_transpose_fused = true;
-        }
-    }
-    ASSERT_TRUE(reshape_transpose_fused);
 
     auto expectedOutputs = run_test(functionRefs);
     for (size_t i = 0; i < actualOutputs.size(); i++) {