[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

@@ -94,6 +94,7 @@
 #include "transformations/op_conversions/convert_topk3.hpp"
 #include "transformations/op_conversions/convert_topk11_downgrade.hpp"
 #include "transformations/op_conversions/scaled_dot_product_attention_decomposition.hpp"
+#include "transformations/op_conversions/convert_convertlike.hpp"
 #include "transformations/opset_conversions/convert_opset2_to_opset1.hpp"
 #include "transformations/opset_conversions/convert_opset3_to_opset2.hpp"
 #include "transformations/smart_reshape/matmul_sr.hpp"
@@ -141,7 +142,6 @@
 #include "transformations/cpu_opset/common/pass/swap_convert_transpose.hpp"
 #include "transformations/cpu_opset/common/pass/causal_mask_preprocess_fusion.hpp"
 #include "transformations/cpu_opset/common/pass/stateful_sdpa_fusion.hpp"
-#include "transformations/cpu_opset/x64/pass/sdpa_fuse_transpose_reshape.hpp"
 
 // Snippets
 #include "snippets/pass/tokenization.hpp"
@@ -618,14 +618,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
 
@@ -659,6 +651,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);
@@ -694,9 +687,22 @@ void Transformations::PreLpt(const std::vector<ov::element::Type>& defaultPrecis
     CPU_REGISTER_PASS_COMMON(manager, ov::pass::EnableDecompressionConvertConstantFolding);
     CPU_REGISTER_PASS_COMMON(manager, ov::pass::KeepConstAndDecompression);
     CPU_REGISTER_PASS_COMMON(manager, ov::pass::ConstantFolding);
+    CPU_REGISTER_PASS_COMMON(manager, ov::intel_cpu::StatefulSDPAFusion);
     CPU_REGISTER_PASS_COMMON(manager, ov::pass::LoraSubgraphFusion);
 
     manager.run_passes(model);
+
+    // Run SDPA decomposition pass separately to let the Snippets handle the default SDPA ops, which don't
+    // fuse special patterns
+    ov::pass::Manager sdpa_manager("CPU:SDPA");
+    CPU_REGISTER_PASS_COMMON(sdpa_manager, ov::pass::ScaledDotProductAttentionDecomposition);
+    CPU_REGISTER_PASS_COMMON(sdpa_manager, ov::pass::ConvertConvertLike);
+    CPU_REGISTER_PASS_COMMON(sdpa_manager, ov::pass::ConstantFolding);
+    CPU_REGISTER_PASS_COMMON(sdpa_manager, ov::pass::TransposeEltwise);
+    CPU_REGISTER_PASS_COMMON(sdpa_manager, ov::pass::TransposeFuse);
+    CPU_REGISTER_PASS_COMMON(sdpa_manager, ov::pass::TransposeMatMul);
+
+    sdpa_manager.run_passes(model);
 }
 
 void Transformations::Lpt(const std::vector<ov::element::Type>& defaultPrecisions) {
@@ -882,8 +888,6 @@ 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_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(postLPTPassManager,

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++) {