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

src/common/snippets/include/snippets/op/vector_buffer.hpp

@@ -21,7 +21,7 @@ public:
 
     VectorBuffer(const ov::element::Type element_type = ov::element::f32);
 
-    bool visit_attributes(AttributeVisitor& visitor) override { return true;}
+    bool visit_attributes(AttributeVisitor& visitor) override;
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
     void validate_and_infer_types() override;
 

src/common/snippets/include/snippets/pass/tokenization.hpp

@@ -61,16 +61,17 @@ public:
      * @ingroup snippets
      */
     struct Config {
-        Config(size_t minimal_concurrency = 1, bool split_m_dimension = true, bool enable_transpose = true)
+        Config(size_t minimal_concurrency = 1, bool split_m_dimension = true, bool enable_transpose_on_output = true)
-            : minimal_concurrency(minimal_concurrency), split_m_dimension(split_m_dimension), mha_token_enable_transpose(enable_transpose) {}
+            : minimal_concurrency(minimal_concurrency), split_m_dimension(split_m_dimension),
+              mha_token_enable_transpose_on_output(enable_transpose_on_output) {}
 
         size_t minimal_concurrency = 1;
         // True if "SplitDimensionM" optimization is enabled. Otherwise, it's disabled.
         bool split_m_dimension = true;
-        // False if all Transposes aren't tokenized in MHA Tokenization.
+        // False if Transpose on output isn't tokenized in MHA Tokenization.
-        // Otherwise, they may be fused into Subgraph if possible
+        // Otherwise, it may be fused into Subgraph if possible
-        // TODO [106921]: Remove please when the ticket 106921 is implemented
+        // TODO [111813]: Remove please when the ticket 111813 is implemented
-        bool mha_token_enable_transpose = true;
+        bool mha_token_enable_transpose_on_output = true;
     };
 
     OPENVINO_RTTI("SnippetsTokenization", "0");

src/common/snippets/src/generator.cpp

@@ -81,7 +81,8 @@ Generator::opRegType Generator::get_op_reg_type(const std::shared_ptr<Node>& op)
              std::dynamic_pointer_cast<op::BroadcastMove>(op) ||
              std::dynamic_pointer_cast<op::Scalar>(op) ||
              std::dynamic_pointer_cast<op::HorizonMax>(op) ||
-             std::dynamic_pointer_cast<op::HorizonSum>(op))
+             std::dynamic_pointer_cast<op::HorizonSum>(op) ||
+             std::dynamic_pointer_cast<op::Fill>(op))
         return vec2vec;
     else
         return get_specific_op_reg_type(op);

src/common/snippets/src/lowered/pass/allocate_buffers.cpp

@@ -17,6 +17,7 @@ void AllocateBuffers::propagate_offset(const LinearIR& linear_ir, const Expressi
     // to correctly read and write data because all Buffers has the common data pointer on buffer scratchpad
 
     const auto buffer = ov::as_type_ptr<op::Buffer>(buffer_expr->get_node());
+    buffer->set_offset(static_cast<int64_t>(offset));
 
     // Propagate to up: in Store. Buffer can have only one Store
     {
@@ -57,26 +58,34 @@ void AllocateBuffers::propagate_offset(const LinearIR& linear_ir, const Expressi
 
 bool AllocateBuffers::run(LinearIR& linear_ir) {
     OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::AllocateBuffers");
+    // [113664] The pass contains two main logics: it defines which of buffers can be inplace (use the same memory) and
+    // allocate memory of needed size. It should be splitted into several passes and updated in bounds of the ticket 113664.
 
-    size_t offset = 0;
+    // [113664] At the moment New Memory Buffer is used only in BrgemmCPU for AMX case. This memory can be reused for each Brgemm.
+    // This plugin-specific condition will be removed in the near future after the task 113664 will be implemented
+    size_t offset = 0, new_memory_buffer_offset = 0;
+    size_t prev_data_size = 0, current_data_size = 0;
     std::set<ExpressionPtr> allocated_buffers;
+    bool new_memory_buffer_allocated = false;
 
     auto allocate = [&](const std::shared_ptr<op::Buffer>& buffer, const ExpressionPtr& expr, size_t buffer_size) {
         offset = m_buffer_scratchpad_size;
-        buffer->set_offset(static_cast<int64_t>(offset));
         propagate_offset(linear_ir, expr, offset);
         m_buffer_scratchpad_size += buffer_size;
         allocated_buffers.insert(expr);
+        prev_data_size = current_data_size;
     };
 
     for (auto expr_it = linear_ir.begin(); expr_it != linear_ir.end(); expr_it++) {
         const auto& expr = *expr_it;
         if (auto buffer = as_type_ptr<op::Buffer>(expr->get_node())) {
             const auto buffer_size = buffer->get_byte_size();
+            current_data_size = buffer->get_element_type().size();
             // If it's the first buffer, offsets are zero => nothing to propagate, can continue
             if (m_buffer_scratchpad_size == 0) {
                 m_buffer_scratchpad_size += buffer_size;
                 allocated_buffers.insert(expr);
+                prev_data_size = current_data_size;
                 continue;
             }
 
@@ -84,7 +93,7 @@ bool AllocateBuffers::run(LinearIR& linear_ir) {
                 const auto& parent_expr = expr->get_input_port_connector(0)->get_source().get_expr();
                 const auto& parent_node = parent_expr->get_node();
                 // Full MemoryAccess ops need new memory. Previous logic is to check for parent isn't Loop
-                // TODO: It should be unified in MemoryManager with memory reuse in the near future
+                // [113664] It should be unified in MemoryManager with memory reuse in the near future
                 const auto ma = ov::as_type_ptr<op::MemoryAccess>(parent_node);
                 if (ma && ma->is_full_memory_access_op()) {
                     allocate(buffer, *expr_it, buffer_size);
@@ -99,7 +108,7 @@ bool AllocateBuffers::run(LinearIR& linear_ir) {
                 // At the moment the pass support only sequentially implicit InPlace.
                 // If Buffer_0 is allocated firstly as Buffer after full memory access op,
                 // we cannot reuse this allocated memory for Buffer_1 - we must allocate new memory for it.
-                // TODO: It should be unified in MemoryManager with memory reuse in the near future
+                // [113664] It should be unified in MemoryManager with memory reuse in the near future
                 bool need_allocate = false;
                 const auto consumers = expr->get_output_port_connector(0)->get_consumers();
                 for (const auto& consumer : consumers) {
@@ -122,16 +131,26 @@ bool AllocateBuffers::run(LinearIR& linear_ir) {
                     continue;
                 }
 
+                // [113664] For more details and reason of the current solution, please, go to the ticket description
                 const auto current_allocated_memory_size = m_buffer_scratchpad_size - offset;
-                if (buffer_size > current_allocated_memory_size) {
+                if (((current_data_size == prev_data_size) && buffer_size > current_allocated_memory_size) ||
+                    ((current_data_size != prev_data_size) && buffer_size != current_allocated_memory_size)) {
                     allocate(buffer, expr, buffer_size);
                     continue;
                 }
                 propagate_offset(linear_ir, *expr_it, offset);
                 allocated_buffers.insert(expr);
+                prev_data_size = current_data_size;
             } else {
-                // Single Buffer without input should allocate new memory
+                if (!new_memory_buffer_allocated) {
-                allocate(buffer, *expr_it, buffer_size);
+                    allocate(buffer, *expr_it, buffer_size);
+                    new_memory_buffer_allocated = true;
+                    new_memory_buffer_offset = offset;
+                } else {
+                    propagate_offset(linear_ir, *expr_it, new_memory_buffer_offset);
+                    allocated_buffers.insert(expr);
+                    prev_data_size = current_data_size;
+                }
             }
         }
     }

src/common/snippets/src/lowered/pass/assign_registers.cpp

@@ -69,8 +69,12 @@ bool AssignRegisters::run(LinearIR& linear_ir) {
             const auto& input_expr = input_tensor->get_source().get_expr();
             const auto& input_expr_input_tensors = input_expr->get_input_port_connectors();
             for (const auto& tensor : input_expr_input_tensors) {
-                if (ov::is_type<op::VectorBuffer>(tensor->get_source().get_expr()->get_node())) {
+                const auto parent_expr = tensor->get_source().get_expr();
+                if (ov::is_type<op::Fill>(parent_expr->get_node())) {
                     manually_assigned_vecs[tensor] = static_cast<Reg>(accumulator_reg);
+                    if (ov::is_type<op::VectorBuffer>(parent_expr->get_input_port_connector(0)->get_source().get_expr()->get_node())) {
+                        manually_assigned_vecs[parent_expr->get_input_port_connector(0)] = static_cast<Reg>(accumulator_reg);
+                }
                 }
             }
             const auto& output_tensor = expr->get_output_port_connector(0);

src/common/snippets/src/lowered/pass/softmax_decomposition.cpp

@@ -41,6 +41,10 @@ bool SoftmaxDecomposition::run(LinearIR& linear_ir) {
             const auto tensor_out = softmax_expr->get_output_port_descriptor(0)->get_shape();
             const auto inner_work_amount = *(tensor_out.rbegin());
 
+            // Float constant values in byte representation
+            const auto float_min_constant = uint32_t(0xff7fffff);
+            const auto zero_constant = uint32_t(0x00000000);
+
             // We need an iterator to the inserted element
             auto push_node = [&linear_ir, &expr_it](const std::shared_ptr<Node>& n) {
                 const auto expr = linear_ir.insert(expr_it, n);
@@ -49,8 +53,10 @@ bool SoftmaxDecomposition::run(LinearIR& linear_ir) {
 
             // Note: VectorBuffer is a special case, since it should go before the initial Load. So we handle it separately
             const auto& vector_buffer_max = push_node(std::make_shared<op::VectorBuffer>());
+            // Init value of vector buffer for ReduceMax is -FLOAT_MIN.
+            const auto fill_max = push_node(std::make_shared<op::Fill>(vector_buffer_max.second, 0, float_min_constant));
             // ReduceMax loop
-            const auto& max = push_node(std::make_shared<ov::op::v1::Maximum>(softmax->get_input_source_output(0), vector_buffer_max.second));
+            const auto& max = push_node(std::make_shared<ov::op::v1::Maximum>(softmax->get_input_source_output(0), fill_max.second));
 
             const auto horizon_max = push_node(std::make_shared<op::HorizonMax>(max.second));
 
@@ -63,11 +69,13 @@ bool SoftmaxDecomposition::run(LinearIR& linear_ir) {
             const auto broadcast_horizon_max = push_node(
                     std::make_shared<op::BroadcastMove>(horizon_max.second, horizon_max.second->get_input_partial_shape(0)));
             const auto vector_buffer_sum = push_node(std::make_shared<op::VectorBuffer>());
+            // Init value of vector buffer for ReduceSum is zero.
+            const auto fill_sum = push_node(std::make_shared<op::Fill>(vector_buffer_sum.second, 0, zero_constant));
 
             // Sub + Exp + ReduceSum Loop
             const auto sub = push_node(std::make_shared<ov::op::v1::Subtract>(softmax->get_input_source_output(0), broadcast_horizon_max.second));
             const auto exp = push_node(std::make_shared<ov::op::v0::Exp>(sub.second));
-            const auto sum = push_node(std::make_shared<ov::op::v1::Add>(exp.second, vector_buffer_sum.second));
+            const auto sum = push_node(std::make_shared<ov::op::v1::Add>(exp.second, fill_sum.second));
 
             const auto horizon_sum = push_node(std::make_shared<op::HorizonSum>(sum.second));
 
@@ -114,8 +122,8 @@ bool SoftmaxDecomposition::run(LinearIR& linear_ir) {
             // For tail loop we should fill input of Max by float min and
             // input of Sum by zero to avoid math incorrect calculations
             // TODO [111383]: It should be covered via general pipeline (for example, via analyze in InsertTailLoop?)
-            max.second->input(0).get_rt_info()["set_fill"] = uint32_t(0xff7fffff);
+            max.second->input(0).get_rt_info()["set_fill"] = float_min_constant;
-            sum.second->input(0).get_rt_info()["set_fill"] = uint32_t(0x00000000);
+            sum.second->input(0).get_rt_info()["set_fill"] = zero_constant;
             modified = true;
         }
     }

src/common/snippets/src/op/subgraph.cpp

@@ -430,10 +430,26 @@ void snippets::op::Subgraph::align_element_types(const BlockedShapeVector& outpu
     for (size_t i = 0; i < outputShapes.size(); i++) {
         const auto needed_out_type = std::get<2>(outputShapes[i]);
         if (body_results[i]->get_input_element_type(0) != needed_out_type) {
-            const auto convert = std::make_shared<ov::snippets::op::ConvertSaturation>(
+            auto parent_output = body_results[i]->get_input_source_output(0);
-                body_results[i]->get_input_node_shared_ptr(0), needed_out_type);
+            std::shared_ptr<ov::Node> consumer = body_results[i];
-            body_results[i]->set_argument(0, convert);
+
-            body_results[i]->validate_and_infer_types();
+            // Snippets supports Transpose only after Parameter or before Result nodes
+            // So we have to insert Convert before Transpose (if there is) on Subgraph outputs
+            const auto transpose = ov::as_type_ptr<ov::op::v1::Transpose>(parent_output.get_node_shared_ptr());
+            if (transpose) {
+                OPENVINO_ASSERT(parent_output.get_target_inputs().size() == 1,
+                                "If Result has Transpose on input, this Result must be single consumer of the Transpose");
+                parent_output = transpose->get_input_source_output(0);
+                consumer = transpose;
+            }
+
+            const auto convert = std::make_shared<ov::snippets::op::ConvertSaturation>(parent_output, needed_out_type);
+            ov::copy_runtime_info(parent_output.get_node_shared_ptr(), convert);
+
+            consumer->set_argument(0, convert);
+            consumer->validate_and_infer_types();
+            if (consumer != body_results[i])
+                body_results[i]->validate_and_infer_types();
         }
     }
 
@@ -442,23 +458,37 @@ void snippets::op::Subgraph::align_element_types(const BlockedShapeVector& outpu
     for (size_t i = 0; i < inputShapes.size(); ++i) {
         const auto needed_in_type = std::get<2>(inputShapes[i]);
         const auto& parameter = parameters[i];
-        if (parameter->get_element_type() != needed_in_type) {
+        const auto original_type = parameter->get_element_type();
-            const auto parameter_output = parameter->output(0);
+        if (original_type != needed_in_type) {
-            const auto convert = std::make_shared<ov::snippets::op::ConvertSaturation>(
+            parameter->set_element_type(needed_in_type);
-                parameter_output,
+            parameter->validate_and_infer_types();
-                parameter_output.get_element_type());
-            ov::copy_runtime_info(parameter, convert);
 
-            for (const auto input : parameter_output.get_target_inputs()) {
+            auto parent_output = parameter->output(0);
+            auto consumer_inputs = parent_output.get_target_inputs();
+
+            // Snippets supports Transpose only after Parameter or before Result nodes
+            // So we have to insert Convert after Transpose (if there is) on Subgraph inputs
+            if (std::any_of(consumer_inputs.cbegin(), consumer_inputs.cend(),
+                [](const ov::Input<ov::Node>& input) { return ov::is_type<ov::op::v1::Transpose>(input.get_node()); })) {
+                OPENVINO_ASSERT(consumer_inputs.size() == 1,
+                                "If Parameter has Transpose on output, this Transpose must be single consumer of the Parameter");
+                const auto transpose = consumer_inputs.begin()->get_node()->shared_from_this();
+                transpose->validate_and_infer_types();
+
+                parent_output = transpose;
+                consumer_inputs = parent_output.get_target_inputs();
+            }
+
+            const auto convert = std::make_shared<ov::snippets::op::ConvertSaturation>(parent_output, original_type);
+            ov::copy_runtime_info(parent_output.get_node_shared_ptr(), convert);
+
+            for (const auto input : consumer_inputs) {
                 const auto& input_node = input.get_node();
                 if (input_node == convert.get()) {
                     continue;
                 }
                 input_node->set_argument(input.get_index(), convert->output(0));
             }
-
-            parameter->set_element_type(needed_in_type);
-            parameter->validate_and_infer_types();
         }
     }
 }

src/common/snippets/src/op/vector_buffer.cpp

@@ -10,7 +10,8 @@ namespace ov {
 namespace snippets {
 namespace op {
 
-VectorBuffer::VectorBuffer(const ov::element::Type element_type) : Op(), m_element_type(std::move(element_type)) {
+VectorBuffer::VectorBuffer(const ov::element::Type element_type)
+    : Op(), m_element_type(std::move(element_type)) {
     constructor_validate_and_infer_types();
 }
 
@@ -25,6 +26,12 @@ void VectorBuffer::validate_and_infer_types() {
     set_output_type(0, m_element_type, Shape{1lu});
 }
 
+bool VectorBuffer::visit_attributes(AttributeVisitor& visitor) {
+    INTERNAL_OP_SCOPE(VectorBuffer_visit_attributes);
+    visitor.on_attribute("element_type", m_element_type);
+    return true;
+}
+
 } // namespace op
 } // namespace snippets
 } // namespace ov

src/common/snippets/src/pass/collapse_subgraph.cpp

@@ -67,10 +67,12 @@ auto is_supported_op(const std::shared_ptr<const Node> &n) -> bool {
             const auto child = transpose->get_output_target_inputs(0).begin()->get_node()->shared_from_this();
             auto is_brgemm_case = ov::is_type<opset1::MatMul>(parent) || ov::is_type<opset1::MatMul>(child);
             // Check for Transpose parent is MatMul inside Subgraph
-            if (const auto subgraph = ov::as_type_ptr<op::Subgraph>(parent)) {
+            if (const auto subgraph = ov::as_type_ptr<const op::Subgraph>(parent)) {
-                const auto body = subgraph->body_ptr();
+                if (GetSnippetsSubgraphType(subgraph) != SnippetsSubgraphType::Completed) {
-                const auto subgraph_output = body->get_results()[transpose->input_value(0).get_index()]->get_input_node_shared_ptr(0);
+                    const auto body = subgraph->body_ptr();
-                is_brgemm_case = is_brgemm_case || ov::is_type<opset1::MatMul>(subgraph_output);
+                    const auto subgraph_output = body->get_results()[transpose->input_value(0).get_index()]->get_input_node_shared_ptr(0);
+                    is_brgemm_case = is_brgemm_case || ov::is_type<opset1::MatMul>(subgraph_output);
+                }
             }
 
             const auto& order = as_type_ptr<const opset1::Constant>(n->get_input_node_shared_ptr(1));

src/common/snippets/src/pass/mha_tokenization.cpp

@@ -191,7 +191,7 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken
     MATCHER_SCOPE(TokenizeMHASnippets);
 
     auto m_matmul0 = std::make_shared<ov::opset1::MatMul>(ov::pass::pattern::any_input(ov::pass::pattern::has_static_shape()),
-                                                              ov::pass::pattern::any_input(ov::pass::pattern::has_static_shape()));
+                                                          ov::pass::pattern::any_input(ov::pass::pattern::has_static_shape()));
 
     register_matcher(std::make_shared<ov::pass::pattern::Matcher>(m_matmul0, matcher_name),
         [=](ov::pass::pattern::Matcher &m) {
@@ -388,14 +388,9 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken
             }
         };
 
-        auto get_transpose = [config](const std::shared_ptr<ov::Node>& node) -> std::shared_ptr<ov::opset1::Transpose> {
+        const auto transpose1 = ov::as_type_ptr<ov::opset1::Transpose>(parent);
-            return config.mha_token_enable_transpose ? ov::as_type_ptr<ov::opset1::Transpose>(node)
+        const auto transpose0 = ov::as_type_ptr<ov::opset1::Transpose>(matmul0->get_input_node_shared_ptr(0));
-                                                     : nullptr;
+        const auto transpose2 = ov::as_type_ptr<ov::opset1::Transpose>(matmul1->get_input_node_shared_ptr(1));
-        };
-
-        const auto transpose1 = get_transpose(parent);
-        const auto transpose0 = get_transpose(matmul0->get_input_node_shared_ptr(0));
-        const auto transpose2 = get_transpose(matmul1->get_input_node_shared_ptr(1));
         tokenize_transpose(transpose1, is_transposed_b_0, {0, 2, 3, 1}, ordered_ops.begin());
         tokenize_transpose(transpose0, matmul0->get_transpose_a(), {0, 2, 1, 3}, ordered_ops.begin());
         tokenize_transpose(transpose2, matmul1->get_transpose_b(), {0, 2, 1, 3}, ordered_ops.end());
@@ -431,7 +426,7 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken
         //  <Supported ops>
         //    Transpose3
         if (!are_ops_after_matmul1) {
-            auto transpose3 = get_transpose(child);
+            auto transpose3 = config.mha_token_enable_transpose_on_output ? ov::as_type_ptr<ov::opset1::Transpose>(child) : nullptr;
             if (is_valid_transpose(transpose3, {0, 2, 1, 3}) &&
                 transpose3->get_input_element_type(0) == matmul1_out_type) {  // To avoid Convert between MatMul1 and Transpose3
                 ordered_ops.push_back(transpose3);

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
+    // [111813]: At the moment Snippets supports Transpose on output of MHA pattern only if it is an one node between MatMul and Result.
-    //      - ConvertSaturation[BF16->FP32] will be inserted after Parameters and before Transposes in canonicalization stage
+    // However there may be Convert [f32->bf16] before Result since:
-    //      - ConvertSaturation[FP32->BF16] will be inserted after Transposes and before Brgemm in precision propagation stage
+    //  - bf16 Brgemm has f32 output;
-    // Because of that Transposes won't be fused into Brgemm
+    //  - CPU Node Subgraph requires bf16 on output when inference precision is bf16.
-    // TODO [111813]: Need to update this pipeline to avoid Converts between Transposes and Brgemm on inputs
+    // To avoid sitations when Transpose is not alone node between MatMul and Result,
-    tokenization_config.mha_token_enable_transpose = (inferencePrecision == ov::element::f32);
+    // 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)
+            const auto in_type0 = matmul->get_input_element_type(0);
-                return !onlyFloatSupported && dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_vnni);
+            const auto in_type1 = matmul->get_input_element_type(1);
-            if (matmul->get_input_element_type(0) == ov::element::bf16 &&
+            if (in_type0 == ov::element::f32 && in_type1 == ov::element::f32 && inferencePrecision == ov::element::f32)
-                matmul->get_input_element_type(1) == ov::element::bf16)
+                return true;
-                return !onlyFloatSupported && dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core_bf16);
+            // [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 (funcInput.get_element_type() == ov::element::bf16)
-            if (patternType == 0 || expectedNode == "Subgraph")
+                tensor = ov::test::utils::create_and_fill_tensor(funcInput.get_element_type(), targetInputStaticShapes[i], 2, -1, 256);
-                tensor = ov::test::utils::create_and_fill_tensor_normal_distribution(funcInput.get_element_type(), targetInputStaticShapes[i], 1.0f, 0.5f);
             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;
+    ExpectedNodes expectedNodes;
-    std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetDevice) = this->GetParam();
+    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;
+        ExpectedNodes expectedNodes;
-        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=(";
@@ -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;
+        ExpectedNodes expectedNodes;
-        std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetDevice) = this->GetParam();
+        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;
+    ExpectedNodes expectedNodes;
-    std::tie(inputShapes, inputPrecisions, matMulIn0Precisions, patternType, expectedNode, targetDevice) = this->GetParam();
+    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 = {
+INSTANTIATE_TEST_SUITE_P(smoke_MHAQuant_Pattern0, MHAQuantTest,
-    0, 1, 2
-};
-
-INSTANTIATE_TEST_SUITE_P(smoke_MHAQuant, MHAQuantTest,
                         ::testing::Combine(
                                 ::testing::ValuesIn(static_shapes_to_test_representation(inputShapesQuant)),
                                 ::testing::ValuesIn(inputPrecisionsQuant),
                                 ::testing::ValuesIn(matMulIn0PrecisionsQuant),
-                                ::testing::ValuesIn(patternTypesQuant),
+                                ::testing::Values(0),
-                                ::testing::Values("MHA"),
+                                ::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);
 

src/tests/functional/plugin/shared/include/snippets/mha.hpp

@@ -66,6 +66,11 @@ protected:
     std::shared_ptr<SnippetsFunctionBase> get_subgraph() override;
 };
 
+class MHAQuantMatMul0 : public MHA {
+protected:
+    std::shared_ptr<SnippetsFunctionBase> get_subgraph() override;
+};
+
 class MHAFQAfterMatMul : public MHA {
 protected:
     std::shared_ptr<SnippetsFunctionBase> get_subgraph() override;

src/tests/functional/plugin/shared/src/snippets/mha.cpp

@@ -115,6 +115,10 @@ std::shared_ptr<SnippetsFunctionBase> MHAINT8MatMul::get_subgraph() {
     return std::make_shared<ov::test::snippets::MHAINT8MatMulFunction>(inputDynamicShapes);
 }
 
+std::shared_ptr<SnippetsFunctionBase> MHAQuantMatMul0::get_subgraph() {
+    return std::make_shared<ov::test::snippets::MHAQuantMatMul0Function>(inputDynamicShapes);
+}
+
 std::shared_ptr<SnippetsFunctionBase> MHAFQAfterMatMul::get_subgraph() {
     return std::make_shared<ov::test::snippets::MHAFQAfterMatMulFunction>(inputDynamicShapes);
 }
@@ -177,6 +181,12 @@ TEST_P(MHAINT8MatMul, CompareWithRefImpl) {
     validateNumSubgraphs();
 }
 
+TEST_P(MHAQuantMatMul0, CompareWithRefImpl) {
+    SKIP_IF_CURRENT_TEST_IS_DISABLED()
+    run();
+    validateNumSubgraphs();
+}
+
 TEST_P(MHAFQAfterMatMul, CompareWithRefImpl) {
     SKIP_IF_CURRENT_TEST_IS_DISABLED()
     run();

src/tests/ngraph_helpers/snippets_ngraph_functions/include/subgraph_mha.hpp

@@ -238,6 +238,33 @@ protected:
     std::shared_ptr<ov::Model> initOriginal() const override;
 };
 
+/* Graph:
+ *   FakeQuantize i8      Reshape1
+ *       Reshape0       Transpose1[0,2,3,1]
+ * Transpose0[0,2,1,3] FakeQuantize i8
+ *              \     /
+ *              MatMul0
+ *                 \   /
+ *                  Add        Reshape2
+ *                Softmax   Transpose2[0,2,1,3]
+ *                    \      /
+ *                     MatMul1
+ *                  FakeQuantize i8
+ *                  Transpose3[0,2,1,3]
+ *                    Reshape3
+ * Note: Reshapes are tosplit Tokenization between FQs and deq Mul and MHA since Snippets::Ignore_Callback may be enabled
+ */
+class MHAQuantMatMul0Function : public SnippetsFunctionBase {
+public:
+    explicit MHAQuantMatMul0Function(const std::vector<PartialShape>& inputShapes)
+            : SnippetsFunctionBase(inputShapes) {
+        NGRAPH_CHECK(input_shapes.size() == 4, "Got invalid number of input shapes");
+    }
+protected:
+    std::shared_ptr<ov::Model> initOriginal() const override;
+};
+
+
 /* Graph:
  *                                          Constant
  *   FakeQuantize u8      FakeQuantize u8   Convert

src/tests/ngraph_helpers/snippets_ngraph_functions/src/subgraph_mha.cpp

@@ -587,6 +587,55 @@ std::shared_ptr<ov::Model> MHAINT8MatMulFunction::initOriginal() const {
     ngraph::ResultVector results{std::make_shared<ngraph::opset1::Result>(transpose3)};
     return std::make_shared<ov::Model>(results, ngraphParam, "mha");
 }
+std::shared_ptr<ov::Model> MHAQuantMatMul0Function::initOriginal() const {
+    auto transpose0Param = std::make_shared<ngraph::opset1::Parameter>(precision, input_shapes[0]);
+    auto transpose1Param = std::make_shared<ngraph::opset1::Parameter>(precision, input_shapes[1]);
+    auto addParam = std::make_shared<ngraph::opset1::Parameter>(precision, input_shapes[2]);
+    auto transpose2Param = std::make_shared<ngraph::opset1::Parameter>(precision, input_shapes[3]);
+    ngraph::ParameterVector ngraphParam = {transpose0Param, transpose1Param, addParam, transpose2Param};
+
+    const auto channel = int64_t(12);
+    const auto last_dim = input_shapes[0].get_shape().back();
+    OPENVINO_ASSERT(last_dim % channel == 0, "Incorrect test configuration");
+    const auto new_shape = std::vector<int64_t>{0, 0, channel, static_cast<int64_t>(last_dim) / channel};
+
+    auto reshape0Const = ngraph::builder::makeConstant(ngraph::element::i64, {new_shape.size()}, new_shape);
+    auto reshape1Const = ngraph::builder::makeConstant(ngraph::element::i64, {new_shape.size()}, new_shape);
+    auto reshape2Const = ngraph::builder::makeConstant(ngraph::element::i64, {new_shape.size()}, new_shape);
+    auto reshape3Const = ngraph::builder::makeConstant(ngraph::element::i64, {input_shapes[0].size()}, std::vector<int64_t>{0, 0, -1});
+
+    auto transpose0Const = ngraph::builder::makeConstant(ngraph::element::i64, {4}, std::vector<int64_t>{0, 2, 1, 3});
+    auto transpose1Const = ngraph::builder::makeConstant(ngraph::element::i64, {4}, std::vector<int64_t>{0, 2, 3, 1});
+    auto transpose2Const = ngraph::builder::makeConstant(ngraph::element::i64, {4}, std::vector<int64_t>{0, 2, 1, 3});
+    auto transpose3Const = ngraph::builder::makeConstant(ngraph::element::i64, {4}, std::vector<int64_t>{0, 2, 1, 3});
+
+    const auto reshape1 = std::make_shared<ov::op::v1::Reshape>(transpose1Param, reshape1Const, true);
+    const auto reshape2 = std::make_shared<ov::op::v1::Reshape>(transpose2Param, reshape2Const, true);
+
+    const auto transpose1 = std::make_shared<ov::op::v1::Transpose>(reshape1, transpose1Const);
+    const auto transpose2 = std::make_shared<ov::op::v1::Transpose>(reshape2, transpose2Const);
+
+    auto fq0 = ngraph::builder::makeFakeQuantize(transpose0Param, ov::element::f32, 256, {1},
+                                                 {-12.5187311}, {12.4209289}, {-12.5187311}, {12.4209289});
+    auto fq1 = ngraph::builder::makeFakeQuantize(transpose1, ov::element::f32, 256, {1},
+                                                 {-1.43326699}, {1.42206954}, {-1.43326699}, {1.42206954});
+
+    const auto reshape0 = std::make_shared<ov::op::v1::Reshape>(fq0, reshape0Const, true);
+    const auto transpose0 = std::make_shared<ov::op::v1::Transpose>(reshape0, transpose0Const);
+
+    const auto matMul0 = std::make_shared<ngraph::opset3::MatMul>(transpose0, fq1);
+    const auto add = std::make_shared<ngraph::opset3::Add>(matMul0, addParam);
+    const auto softMax = std::make_shared<ngraph::opset8::Softmax>(add, -1);
+
+    const auto matMul1 = std::make_shared<ngraph::opset3::MatMul>(softMax, transpose2);
+    auto fq2 = ngraph::builder::makeFakeQuantize(matMul1, ov::element::f32, 256, {1},
+                                                 {-1.81826221}, {1.804057}, {-1.81826221}, {1.804057});
+    const auto transpose3 = std::make_shared<ov::op::v1::Transpose>(fq2, transpose3Const);
+    const auto reshape3 = std::make_shared<ov::op::v1::Reshape>(transpose3, reshape3Const, true);
+
+    ngraph::ResultVector results{std::make_shared<ngraph::opset1::Result>(reshape3)};
+    return std::make_shared<ov::Model>(results, ngraphParam, "mha");
+}
 std::shared_ptr<ov::Model> MHAFQFunction::initOriginal() const {
     auto transpose0Param = std::make_shared<ngraph::opset1::Parameter>(precision, input_shapes[0]);
     auto transpose1Param = std::make_shared<ngraph::opset1::Parameter>(precision, input_shapes[1]);