[Snippets][CPU] Enabled MHA tokenization for quant and bf16 cases (#18403)

src/plugins/intel_cpu/src/emitters/x64/cpu_generator.cpp

@@ -55,7 +55,7 @@ ov::intel_cpu::CPUTargetMachine::CPUTargetMachine(dnnl::impl::cpu::x64::cpu_isa_
     jitters[ov::op::v0::Parameter::get_type_info_static()] = CREATE_EMITTER(NopEmitter);
     jitters[ov::op::v0::Result::get_type_info_static()] = CREATE_EMITTER(NopEmitter);
     jitters[snippets::op::Buffer::get_type_info_static()] = CREATE_EMITTER(NopEmitter);
-    jitters[snippets::op::VectorBuffer::get_type_info_static()] = CREATE_EMITTER(VectorBufferEmitter);
+    jitters[snippets::op::VectorBuffer::get_type_info_static()] = CREATE_EMITTER(NopEmitter);
     // jitters[ov::op::v1::Constant::get_type_info_static()] = CREATE_EMITTER(); // Not supported
 
     jitters[snippets::op::Load::get_type_info_static()] = CREATE_EMITTER(LoadEmitter);

src/plugins/intel_cpu/src/emitters/x64/jit_snippets_emitters.cpp

@@ -1510,31 +1510,6 @@ void HorizonEmitter::perform_op(const Vmm &vmm1, const Vmm &vmm2, const Vmm &vmm
     }
 }
 
-VectorBufferEmitter::VectorBufferEmitter(dnnl::impl::cpu::x64::jit_generator* h, dnnl::impl::cpu::x64::cpu_isa_t isa, const std::shared_ptr<ov::Node>& n) :
-    jit_emitter(h, isa, n, Precision::FP32, emitter_in_out_map::vec_to_vec) {}
-
-void VectorBufferEmitter::emit_impl(const std::vector<size_t>& in,
-                                    const std::vector<size_t>& out) const {
-    if (host_isa_ == dnnl::impl::cpu::x64::sse41) {
-        emit_isa<dnnl::impl::cpu::x64::sse41>(in, out);
-    } else if (host_isa_ == dnnl::impl::cpu::x64::avx2) {
-        emit_isa<dnnl::impl::cpu::x64::avx2>(in, out);
-    } else if (host_isa_ == dnnl::impl::cpu::x64::avx512_core) {
-        emit_isa<dnnl::impl::cpu::x64::avx512_core>(in, out);
-    } else {
-        IE_THROW() << "Zero emitter doesn't support " << host_isa_;
-    }
-}
-
-template <dnnl::impl::cpu::x64::cpu_isa_t isa>
-void VectorBufferEmitter::emit_isa(const std::vector<size_t> &in, const std::vector<size_t> &out) const {
-    using Vmm = typename dnnl::impl::utils::conditional3<isa == dnnl::impl::cpu::x64::sse41,
-            Xmm, isa == dnnl::impl::cpu::x64::avx2, Ymm, Zmm>::type;
-
-    Vmm vmm = Vmm(out[0]);
-    h->uni_vpxor(vmm, vmm, vmm);
-}
-
 FillEmitter::FillEmitter(dnnl::impl::cpu::x64::jit_generator* h, dnnl::impl::cpu::x64::cpu_isa_t isa, const std::shared_ptr<ov::Node>& n) :
     jit_emitter(h, isa, n, Precision::FP32, emitter_in_out_map::vec_to_vec) {
     const auto fill = ov::as_type_ptr<snippets::op::Fill>(n);
@@ -1544,10 +1519,18 @@ FillEmitter::FillEmitter(dnnl::impl::cpu::x64::jit_generator* h, dnnl::impl::cpu
 
     offset = fill->get_offset();
     fill_value = fill->get_fill_value();
+    if (!is_optimized())
+        push_arg_entry_of("value", fill_value, true);
     prepare_table();
 }
 
 size_t FillEmitter::aux_gprs_count() const {
+    // Optimized version (fill full vector by zero) doesn't need additional register
+    if (is_optimized())
+        return 0;
+    // + 1 reg for table value in full vector case
+    if (is_full_reg())
+        return 1;
     // + 1 reg for temp reg for mask in avx512
     return one_of(host_isa_, dnnl::impl::cpu::x64::avx512_core) ? 2 : 1;
 }
@@ -1573,6 +1556,25 @@ void FillEmitter::emit_isa(const std::vector<size_t> &in, const std::vector<size
     Vmm src_vmm = Vmm(in[0]);
     Vmm dst_vmm = Vmm(out[0]);
 
+    if (is_full_reg())
+        fill_full<Vmm>(dst_vmm);
+    else
+        fill_tail<Vmm>(src_vmm, dst_vmm);
+}
+
+template <typename Vmm>
+void FillEmitter::fill_full(const Vmm& dst_vmm) const {
+    // Optimized impl for zero
+    if (is_optimized()) {
+        h->uni_vpxor(dst_vmm, dst_vmm, dst_vmm);
+        return;
+    }
+
+    h->uni_vbroadcastss(dst_vmm, table_val("value"));
+}
+
+template <typename Vmm>
+void FillEmitter::fill_tail(const Vmm& src_vmm, const Vmm& dst_vmm) const {
     if (one_of(host_isa_, dnnl::impl::cpu::x64::avx512_core)) {
         uint64_t tail_mask = 1;
         tail_mask = ~((tail_mask << offset) - tail_mask);
@@ -1584,15 +1586,12 @@ void FillEmitter::emit_isa(const std::vector<size_t> &in, const std::vector<size
         imm = ~((imm << offset) - imm);  // shift load_num bit
         if (host_isa_ == dnnl::impl::cpu::x64::sse41 && src_vmm.getIdx() != dst_vmm.getIdx()) {
             h->uni_vmovups(dst_vmm, src_vmm);
-            src_vmm = Vmm(dst_vmm.getIdx());
+            h->uni_vblendps(dst_vmm, dst_vmm, table_val("value"), imm);
+        } else {
+            h->uni_vblendps(dst_vmm, src_vmm, table_val("value"), imm);
         }
-        h->uni_vblendps(dst_vmm, src_vmm, table_val("value"), imm);
     }
 }
 
-void FillEmitter::register_table_entries() {
-    push_arg_entry_of("value", fill_value, true);
-}
-
 }   // namespace intel_cpu
 }   // namespace ov

src/plugins/intel_cpu/src/emitters/x64/jit_snippets_emitters.hpp

@@ -455,21 +455,6 @@ private:
     enum class OpType { max, sum };
     OpType m_op_type = OpType::max;
 };
-
-class VectorBufferEmitter : public jit_emitter {
-public:
-    VectorBufferEmitter(dnnl::impl::cpu::x64::jit_generator* h, dnnl::impl::cpu::x64::cpu_isa_t isa, const std::shared_ptr<ov::Node>& n);
-
-    size_t get_inputs_num() const override {return 0;}
-
-private:
-    void emit_impl(const std::vector<size_t>& in,
-                   const std::vector<size_t>& out) const override;
-
-    template <dnnl::impl::cpu::x64::cpu_isa_t isa>
-    void emit_isa(const std::vector<size_t> &in, const std::vector<size_t> &out) const;
-};
-
 class FillEmitter : public jit_emitter {
 public:
     FillEmitter(dnnl::impl::cpu::x64::jit_generator* h, dnnl::impl::cpu::x64::cpu_isa_t isa, const std::shared_ptr<ov::Node>& n);
@@ -485,8 +470,13 @@ private:
 
     template <dnnl::impl::cpu::x64::cpu_isa_t isa>
     void emit_isa(const std::vector<size_t> &in, const std::vector<size_t> &out) const;
+    template <typename Vmm>
+    void fill_full(const Vmm& vmm_dst) const;
+    template <typename Vmm>
+    void fill_tail(const Vmm& vmm_src, const Vmm& vmm_dst) const;
 
-    void register_table_entries() override;
+    bool is_full_reg() const { return offset == 0; }
+    bool is_optimized() const { return is_full_reg() && fill_value == uint32_t(0x0); }
 
     size_t offset = 0;
     uint32_t fill_value = 0x0;

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

@@ -96,7 +96,6 @@
 // CPU specific transformations
 #include "transformations/cpu_opset/convert_to_cpu_specific_opset.hpp"
 #include "transformations/snippets/x64/pass/snippets_mark_skipped.hpp"
-#include "transformations/cpu_opset/x64/pass/mha_fusion.hpp"
 #include "transformations/cpu_opset/x64/pass/convert_to_interaction.hpp"
 #include "transformations/cpu_opset/arm/pass/convert_group_conv.hpp"
 #include "transformations/cpu_opset/arm/pass/convert_group_conv1d.hpp"
@@ -560,36 +559,8 @@ void Transformations::PostLpt() {
 
     CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::ConstantFolding);
 
-    // Snippets may brake MHA patterns so the fusion has to performed before
-    CPU_REGISTER_PASS_X64(postLPTPassManager, MHAFusion);
     CPU_REGISTER_PASS_X64(postLPTPassManager, FuseFQtoInteraction);
 
-    CPU_SET_CALLBACK_X64(postLPTPassManager,
-        ([this](const std::shared_ptr<const ov::Node>& n) -> bool {
-            std::string errorMessage;
-
-            if (!node::MHA::isSupportedOperation(n, errorMessage))
-                return true;
-
-            // Implementation calls AMX BF16 brgemm only for tensors with K and N aligned on 2, otherwise fallbacks on vector impl
-            // Vector madd BF16 instruction on SPR has reduced performance on HW level, which results in overall perf degradation
-            size_t bf16Factor = 2;
-            if (dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_amx) &&
-                (n->get_input_element_type(0) == element::bf16 || (n->get_input_element_type(0) == element::f32 && inferencePrecision == ov::element::bf16)) &&
-                (n->get_input_shape(0)[3] % bf16Factor != 0 || n->get_input_shape(1)[1] % bf16Factor != 0 || n->get_input_shape(3)[3] % bf16Factor != 0)) {
-                return true;
-            }
-
-            return false;
-        }),
-        MHAFloatFusion, MHAFloatFusion2, MHAQuantFusion, MHAQuantFusion2);
-
-    // Float MHA is supported by snippets now
-    if (inferencePrecision == ov::element::f32) {
-        CPU_DISABLE_PASS_X64(postLPTPassManager, MHAFloatFusion);
-        CPU_DISABLE_PASS_X64(postLPTPassManager, MHAFloatFusion2);
-    }
-
     // Execute before snippets. Otherwise FQ will be converted to Subgraph
     CPU_REGISTER_PASS_X64(postLPTPassManager, ConvertFqRnnToQuantizedRnn);
     postLPTPassManager.run_passes(model);
@@ -601,12 +572,13 @@ void Transformations::MainSnippets(void) {
         return;
 
     ov::snippets::pass::SnippetsTokenization::Config tokenization_config;
-    // At the moment Snippets supports Transposes in MHA pattern only in FP32 case since
-    //      - ConvertSaturation[BF16->FP32] will be inserted after Parameters and before Transposes in canonicalization stage
-    //      - ConvertSaturation[FP32->BF16] will be inserted after Transposes and before Brgemm in precision propagation stage
-    // Because of that Transposes won't be fused into Brgemm
-    // TODO [111813]: Need to update this pipeline to avoid Converts between Transposes and Brgemm on inputs
-    tokenization_config.mha_token_enable_transpose = (inferencePrecision == ov::element::f32);
+    // [111813]: At the moment Snippets supports Transpose on output of MHA pattern only if it is an one node between MatMul and Result.
+    // However there may be Convert [f32->bf16] before Result since:
+    //  - bf16 Brgemm has f32 output;
+    //  - CPU Node Subgraph requires bf16 on output when inference precision is bf16.
+    // To avoid sitations when Transpose is not alone node between MatMul and Result,
+    // Plugin disables Transpose tokenization on output
+    tokenization_config.mha_token_enable_transpose_on_output = (inferencePrecision == ov::element::f32);
     tokenization_config.minimal_concurrency = parallel_get_num_threads();
     // The optimization "SplitDimensionM" depends on target machine (thread count).
     // To avoid uncontrolled behavior in tests, we disabled the optimization when there is Config::SnippetsMode::IgnoreCallback
@@ -618,11 +590,7 @@ void Transformations::MainSnippets(void) {
         CPU_REGISTER_PASS_X64(snippetsManager, SnippetsMarkSkipped, inferencePrecision != ov::element::f32);
     CPU_REGISTER_PASS_X64(snippetsManager, snippets::pass::SnippetsTokenization, tokenization_config);
 
-    // Tokenize MHA in quantized model or with BF16 only in tests.
-    // TODO [106921]: Please enable the tokenization when the ticket 106921 with blocking support for BRGEMM will be implemented
-    const bool onlyFloatSupported = snippetsMode != Config::SnippetsMode::IgnoreCallback;
     const bool isMHASupported =
-            IMPLICATION(inferencePrecision != ov::element::f32, !onlyFloatSupported) &&
             dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core);  // MHA has BRGEMM that is supported only on AVX512 platforms
     if (!isMHASupported) {
         CPU_DISABLE_PASS_X64(snippetsManager, snippets::pass::TokenizeMHASnippets);
@@ -631,15 +599,38 @@ void Transformations::MainSnippets(void) {
 
     if (snippetsMode != Config::SnippetsMode::IgnoreCallback) {
 #if defined(OPENVINO_ARCH_X86_64)
-        auto is_supported_matmul = [onlyFloatSupported](const std::shared_ptr<const ov::Node>& n) {
+        auto is_supported_matmul = [this](const std::shared_ptr<const ov::Node>& n) {
             const auto matmul = ov::as_type_ptr<const ov::op::v0::MatMul>(n);
             if (!matmul)
                 return false;
-            if (matmul->get_input_element_type(1) == ov::element::i8)
-                return !onlyFloatSupported && dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_vnni);
-            if (matmul->get_input_element_type(0) == ov::element::bf16 &&
-                matmul->get_input_element_type(1) == ov::element::bf16)
-                return !onlyFloatSupported && dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_bf16);
+            const auto in_type0 = matmul->get_input_element_type(0);
+            const auto in_type1 = matmul->get_input_element_type(1);
+            if (in_type0 == ov::element::f32 && in_type1 == ov::element::f32 && inferencePrecision == ov::element::f32)
+                return true;
+            // [114487] brgemm kernel in oneDNN requires brgemm_copy_b kernel if MatMul node has transposed_b=True
+            // The current solution with ExtractExplicitMatMulTranspose pass is slower for non-f32 cases than using of brgemm_copy_b kernel
+            if (matmul->get_transpose_a() || matmul->get_transpose_b())
+                return false;
+            // [115165] At the moment Quantized and BF16 Brgemm doesn't support blocking by K and N.
+            // Big shapes may lead to perf degradation
+            const auto K = *(matmul->get_input_partial_shape(0).rbegin());
+            const auto N = *(matmul->get_input_partial_shape(1).rbegin());
+            if ((K.is_static() && K.get_length() > 512) || // heuristic values
+                (N.is_static() && N.get_length() > 256))
+                return false;
+            if (in_type0 == ov::element::i8)
+                return dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_vnni);
+            if ((in_type0 == ov::element::bf16 && in_type1 == ov::element::bf16) ||
+                ((in_type0 == element::f32 && in_type1 == ov::element::f32 && inferencePrecision == ov::element::bf16))) {
+                // Implementation calls AMX BF16 brgemm only for tensors with K and N aligned on 2, otherwise fallbacks on vector impl
+                // Vector madd BF16 instruction on SPR has reduced performance on HW level, which results in overall perf degradation
+                size_t bf16Factor = 2;
+                if (dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_amx)) {
+                    return K.is_static() && (K.get_length() % bf16Factor == 0) &&
+                           N.is_static() && (N.get_length() % bf16Factor == 0);
+                }
+                return dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_bf16);
+            }
             return true;
         };
         auto is_unsupported_parallel_work_amount = [&](const std::shared_ptr<const ov::Node>& n, const ov::Shape& shape) {

src/plugins/intel_cpu/tests/functional/shared_tests_instances/skip_tests_config.cpp

@@ -264,6 +264,7 @@ std::vector<std::string> disabledTestPatterns() {
         retVector.emplace_back(R"(.*Snippets.*MatMul.*Quantized.*)");
         retVector.emplace_back(R"(.*Snippets.*MHAFQ.*)");
         retVector.emplace_back(R"(.*Snippets.*MHAINT8.*)");
+        retVector.emplace_back(R"(.*Snippets.*MHAQuant.*)");
     }
     if (!InferenceEngine::with_cpu_x86_avx512_core_amx_int8())
         //TODO: Issue 92895

src/plugins/intel_cpu/tests/functional/shared_tests_instances/snippets/mha.cpp

@@ -203,6 +203,18 @@ INSTANTIATE_TEST_SUITE_P(smoke_Snippets_MHAINT8MatMul, MHAINT8MatMul,
                                  ::testing::Values(CPUTestUtils::cpuEmptyPluginConfig)),
                          MHA::getTestCaseName);
 
+INSTANTIATE_TEST_SUITE_P(smoke_Snippets_MHAQuantMatMul0, MHAQuantMatMul0,
+                         ::testing::Combine(
+                                 ::testing::Values(std::vector<ov::PartialShape>{{1, 128, 768}, {1, 128, 768}, {1, 1, 1, 128}, {1, 128, 768}}),
+                                 ::testing::Values(std::vector<element::Type>{}),
+                                 ::testing::Values(ov::element::f32),
+                                 ::testing::Values(false), // The graph doesn't contain Multiply
+                                 ::testing::Values(8),     // FQ on input + MHA + Transpose on output + 4 Reshapes + Deq Mul
+                                 ::testing::Values(3),     // FQ on input + MHA + Deq Mul
+                                 ::testing::Values(CommonTestUtils::DEVICE_CPU),
+                                 ::testing::Values(CPUTestUtils::cpuEmptyPluginConfig)),
+                         MHA::getTestCaseName);
+
 INSTANTIATE_TEST_SUITE_P(smoke_Snippets_MHAFQAfterMatMul, MHAFQAfterMatMul,
                          ::testing::Combine(
                                  ::testing::ValuesIn(inputShapes),

src/plugins/intel_cpu/tests/functional/subgraph_tests/src/mha.cpp

@@ -21,12 +21,14 @@ using namespace ngraph::helpers;
 
 namespace CPUSubgraphTestsDefinitions {
 
+using ExpectedNodes = std::vector<std::pair<std::string, size_t>>;
+
 typedef std::tuple<
         std::vector<InputShape>,   // Input shapes
         std::vector<ElementType>,  // Input precisions
         std::vector<ElementType>,  // MatMul input #0 precisions
         size_t,                    // pattern type #
-        std::string,               // Expected node
+        ExpectedNodes,             // Expected node -> count
         std::string                // Device name
 > MHATuple;
 
@@ -157,9 +159,9 @@ public:
         std::vector<ElementType> inputPrecisions;
         std::vector<ElementType> matMulIn0Precisions;
         size_t patternType;
-        std::string expectedNode;
+        ExpectedNodes expectedNodes;
         std::string targetName;
-        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetName) = obj.param;
+        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNodes, targetName) = obj.param;
         std::ostringstream results;
 
         results << "IS=(";
@@ -176,7 +178,10 @@ public:
             results << "InPRC" << std::to_string(i) << "=" << inputPrecisions[i] << "_";
         }
         results << "patternType=" << patternType;
-        results << "expect=" << expectedNode;
+        results << "expect=";
+        for (const auto& node : expectedNodes) {
+            results << node.first << "[" << node.second << "]" << "_";
+        }
         results << "targetDevice=" << targetName;
 
         return results.str();
@@ -188,24 +193,23 @@ public:
         for (size_t i = 0; i < funcInputs.size(); ++i) {
             const auto& funcInput = funcInputs[i];
             ov::Tensor tensor;
-            // TODO: after snippets fixed should remove 2nd condition, ticket: 105339
-            if (patternType == 0 || expectedNode == "Subgraph")
-                tensor = ov::test::utils::create_and_fill_tensor_normal_distribution(funcInput.get_element_type(), targetInputStaticShapes[i], 1.0f, 0.5f);
+            if (funcInput.get_element_type() == ov::element::bf16)
+                tensor = ov::test::utils::create_and_fill_tensor(funcInput.get_element_type(), targetInputStaticShapes[i], 2, -1, 256);
             else
-                // generate all negative inputs
-                tensor = ov::test::utils::create_and_fill_tensor_unique_sequence(funcInput.get_element_type(), targetInputStaticShapes[i], -1, -5);
+                tensor = ov::test::utils::create_and_fill_tensor_unique_sequence(funcInput.get_element_type(), targetInputStaticShapes[i], -1, 5);
+            inputs.insert({funcInput.get_node_shared_ptr(), tensor});
             inputs.insert({funcInput.get_node_shared_ptr(), tensor});
         }
     }
 
 protected:
     size_t patternType;
-    std::string expectedNode;
+    ExpectedNodes expectedNodes;
     void SetUp() override {
         std::vector<InputShape> inputShapes;
         std::vector<ElementType> inputPrecisions;
         std::vector<ElementType> matMulIn0Precisions;
-        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetDevice) = this->GetParam();
+        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNodes, targetDevice) = this->GetParam();
 
         init_input_shapes(inputShapes);
 
@@ -240,8 +244,8 @@ TEST_P(MHATest, CompareWithRefs) {
     std::vector<ElementType> inputPrecisions;
     std::vector<ElementType> matMulIn0Precisions;
     size_t patternType;
-    std::string expectedNode;
-    std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetDevice) = this->GetParam();
+    ExpectedNodes expectedNodes;
+    std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNodes, targetDevice) = this->GetParam();
 
     if (inputPrecisions[0] == ElementType::bf16 && !InferenceEngine::with_cpu_x86_bfloat16())
         GTEST_SKIP();
@@ -250,7 +254,10 @@ TEST_P(MHATest, CompareWithRefs) {
         GTEST_SKIP();
 
     run();
-    CheckNumberOfNodesWithType(compiledModel, expectedNode, 1);
+
+    for (const auto& node : expectedNodes) {
+        CheckNumberOfNodesWithType(compiledModel, node.first, node.second);
+    }
 }
 
 namespace {
@@ -273,23 +280,24 @@ std::vector<size_t> patternTypes = {
     0, 1
 };
 
-INSTANTIATE_TEST_SUITE_P(smoke_Snippets_MHA, MHATest,
+INSTANTIATE_TEST_SUITE_P(smoke_MHA, MHATest,
                         ::testing::Combine(
                                 ::testing::ValuesIn(static_shapes_to_test_representation(inputShapes)),
                                 ::testing::Values(std::vector<ElementType>{ ElementType::f32, ElementType::f32, ElementType::f32, ElementType::f32 }),
                                 ::testing::ValuesIn(matMulIn0Precisions),
                                 ::testing::ValuesIn(patternTypes),
-                                ::testing::Values("Subgraph"),
+                                ::testing::Values(ExpectedNodes{{"Subgraph", 1}}),
                                 ::testing::Values(ov::test::utils::DEVICE_CPU)),
                         MHATest::getTestCaseName);
 
-INSTANTIATE_TEST_SUITE_P(smoke_MHA, MHATest,
+INSTANTIATE_TEST_SUITE_P(smoke_MHA_BF16, MHATest,
                          ::testing::Combine(
                                  ::testing::ValuesIn(static_shapes_to_test_representation(inputShapes)),
                                  ::testing::Values(std::vector<ElementType>{ ElementType::bf16, ElementType::bf16, ElementType::bf16, ElementType::bf16 }),
                                  ::testing::ValuesIn(matMulIn0Precisions),
                                  ::testing::ValuesIn(patternTypes),
-                                 ::testing::Values("MHA"),  // Snippets don't support BF16 MHA pattern yet
+                                 ::testing::Values(ExpectedNodes{{"Subgraph", 1},
+                                                                 {"Transpose", 1}}),  // Plugin disables tokenization of Transpose on output
                                  ::testing::Values(ov::test::utils::DEVICE_CPU)),
                          MHATest::getTestCaseName);
 
@@ -454,8 +462,8 @@ public:
         std::vector<ElementType> matMulIn0Precisions;
         size_t patternType;
         std::string targetName;
-        std::string expectedNode;
-        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetName) = obj.param;
+        ExpectedNodes expectedNodes;
+        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNodes, targetName) = obj.param;
         std::ostringstream results;
 
         results << "IS=(";
@@ -475,7 +483,10 @@ public:
             results << "MatMulIn0PRC" << std::to_string(i) << "=" << matMulIn0Precisions[i] << "_";
         }
         results << "patternType=" << patternType;
-        results << "expect=" << expectedNode;
+        results << "expect=";
+        for (const auto& node : expectedNodes) {
+            results << node.first << "[" << node.second << "]" << "_";
+        }
         results << "targetDevice=" << targetName;
 
         return results.str();
@@ -505,8 +516,8 @@ protected:
         std::vector<ElementType> inputPrecisions;
         std::vector<ElementType> matMulIn0Precisions;
         size_t patternType;
-        std::string expectedNode;
-        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetDevice) = this->GetParam();
+        ExpectedNodes expectedNodes;
+        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNodes, targetDevice) = this->GetParam();
 
         init_input_shapes(inputShapes);
 
@@ -534,8 +545,8 @@ TEST_P(MHAQuantTest, CompareWithRefs) {
     std::vector<ElementType> inputPrecisions;
     std::vector<ElementType> matMulIn0Precisions;
     size_t patternType;
-    std::string expectedNode;
-    std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetDevice) = this->GetParam();
+    ExpectedNodes expectedNodes;
+    std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNodes, targetDevice) = this->GetParam();
 
     if (inputPrecisions[0] == ElementType::bf16 && !InferenceEngine::with_cpu_x86_bfloat16())
         GTEST_SKIP();
@@ -544,7 +555,10 @@ TEST_P(MHAQuantTest, CompareWithRefs) {
         GTEST_SKIP();
 
     run();
-    CheckNumberOfNodesWithType(compiledModel, expectedNode, 1);
+
+    for (const auto& node : expectedNodes) {
+        CheckNumberOfNodesWithType(compiledModel, node.first, node.second);
+    }
 }
 
 namespace {
@@ -570,17 +584,37 @@ std::vector<std::vector<ElementType>> matMulIn0PrecisionsQuant = {
     { ElementType::i8, ElementType::u8 },
 };
 
-std::vector<size_t> patternTypesQuant = {
-    0, 1, 2
-};
-
-INSTANTIATE_TEST_SUITE_P(smoke_MHAQuant, MHAQuantTest,
+INSTANTIATE_TEST_SUITE_P(smoke_MHAQuant_Pattern0, MHAQuantTest,
                         ::testing::Combine(
                                 ::testing::ValuesIn(static_shapes_to_test_representation(inputShapesQuant)),
                                 ::testing::ValuesIn(inputPrecisionsQuant),
                                 ::testing::ValuesIn(matMulIn0PrecisionsQuant),
-                                ::testing::ValuesIn(patternTypesQuant),
-                                ::testing::Values("MHA"),
+                                ::testing::Values(0),
+                                ::testing::Values(ExpectedNodes{{"Subgraph", 5},  // FQs on inputs x 3 + MHA + Deq Mul
+                                                                {"Transpose", 1}}),  // Transpose between MHA and Deq Mul
+                                ::testing::Values(CommonTestUtils::DEVICE_CPU)),
+                        MHAQuantTest::getTestCaseName);
+
+
+INSTANTIATE_TEST_SUITE_P(smoke_MHAQuant_Pattern1, MHAQuantTest,
+                        ::testing::Combine(
+                                ::testing::ValuesIn(static_shapes_to_test_representation(inputShapesQuant)),
+                                ::testing::ValuesIn(inputPrecisionsQuant),
+                                ::testing::ValuesIn(matMulIn0PrecisionsQuant),
+                                ::testing::Values(1),
+                                ::testing::Values(ExpectedNodes{{"Subgraph", 3},  // FQ on input + MHA + Deq Mul
+                                                                {"Transpose", 1}}),  // Transpose between MHA and Deq Mul
+                                ::testing::Values(CommonTestUtils::DEVICE_CPU)),
+                        MHAQuantTest::getTestCaseName);
+
+INSTANTIATE_TEST_SUITE_P(smoke_MHAQuant_Pattern2, MHAQuantTest,
+                        ::testing::Combine(
+                                ::testing::ValuesIn(static_shapes_to_test_representation(inputShapesQuant)),
+                                ::testing::ValuesIn(inputPrecisionsQuant),
+                                ::testing::ValuesIn(matMulIn0PrecisionsQuant),
+                                ::testing::Values(2),
+                                ::testing::Values(ExpectedNodes{{"Subgraph", 2},  // MHA + Deq Mul
+                                                                {"Transpose", 0}}), // Transpose is fused
                                 ::testing::Values(ov::test::utils::DEVICE_CPU)),
                         MHAQuantTest::getTestCaseName);