[CPU] Fix performance regressions in Llama & Gemma (openvinotoolkit#23396)

### Details:
- the sin/cos cache is removed from LlamaRotaryEmbedding, sin/cos table
used for ROPE is generated on the fly for each 2nd token and each layer,
this PR is match such subgraph and share the first match across all
macthes
-
[There](https://github.com/openvinotoolkit/openvino/blob/5b7d7c87a8a5847f97d869473fa22e8f1290a116/src/common/transformations/src/transformations/common_optimizations/shared_ops_optimization.cpp#L64)
is generic common transformation doing similar optimizations but failed
due to
`cpu___module.model.layers.1.self_attn.rotary_emb/aten::to/Convert` is
not constant (but will be const folded later, and after that it can pass
the check). another limitation is SharedOpOptimization is only for OP
sharing rather than a subgraph sharing.

- the causal mask maintenance logic in original Llama & Gemma model is
very inefficient, a big triangular causal mask tensor is updated on
every token generation:

```python
def _update_causal_mask():
   ... ...
    min_dtype = torch.finfo(dtype).min
    causal_mask = self.causal_mask[None, None, :, :].repeat(batch_size, 1, 1, 1).to(dtype) * min_dtype
    causal_mask = causal_mask.to(dtype=dtype, device=device)

    mask_length = attention_mask.shape[-1]
    padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
    causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype)
``` 
due to limitations in ngraph OP set, PT FE generated even more costly
subgraph : a Range node generating 8K-by-8K 1D indices to
ScatterElementsUpdate OP. but actually only very small slice of the
result mask is being used by SDPA as combined causal/attention mask.

```python
      if attention_mask is not None and cache_position is not None:
          causal_mask = causal_mask[:, :, cache_position, : key_states.shape[-2]]
```
I introduced a cpu custom ngraph & node `CausalMaskPreprocess` to fuse
above logic.

- gen_pattern code is also partially refactored to display better
verbose logs when pattern matching failed.

### Tickets:
 - *CVS-133878*
 - *CVS-134592*

---------

Co-authored-by: guozhong wang <guozhong.wang@intel.com>

Author: usstq

src/plugins/intel_cpu/src/cpu_types.cpp

@@ -219,6 +219,7 @@ static const TypeToNameMap& get_type_to_name_tbl() {
         {"ScaledDotProductAttentionWithKVCache", Type::ScaledDotProductAttention},
         {"PagedAttentionExtension", Type::ScaledDotProductAttention},
         {"RoPE", Type::RoPE},
+        {"CausalMaskPreprocess", Type::CausalMaskPreprocess},
     };
     return type_to_name_tbl;
 }
@@ -337,6 +338,7 @@ std::string NameFromType(const Type type) {
         CASE(Ngram);
         CASE(ScaledDotProductAttention);
         CASE(RoPE);
+        CASE(CausalMaskPreprocess);
         CASE(Unknown);
     }
 #undef CASE

src/plugins/intel_cpu/src/cpu_types.h

@@ -116,6 +116,7 @@ enum class Type {
     Ngram,
     ScaledDotProductAttention,
     RoPE,
+    CausalMaskPreprocess,
 };
 
 enum class Algorithm {

src/plugins/intel_cpu/src/extension.cpp

@@ -12,6 +12,7 @@
 #include "ov_ops/nms_static_shape_ie.hpp"
 #include "ov_ops/type_relaxed.hpp"
 #include "snippets/op/subgraph.hpp"
+#include "transformations/cpu_opset/common/op/causal_mask_preprocess.hpp"
 #include "transformations/cpu_opset/common/op/fully_connected.hpp"
 #include "transformations/cpu_opset/common/op/leaky_relu.hpp"
 #include "transformations/cpu_opset/common/op/ngram.hpp"
@@ -69,6 +70,7 @@ class TypeRelaxedExtension : public ov::OpExtension<ov::op::TypeRelaxed<Op>> {
     OP_EXTENSION(ov::intel_cpu::LeakyReluNode)                              \
     OP_EXTENSION(ov::intel_cpu::PowerStaticNode)                            \
     OP_EXTENSION(ov::intel_cpu::RoPENode)                                   \
+    OP_EXTENSION(ov::intel_cpu::CausalMaskPreprocessNode)                             \
     OP_EXTENSION(ov::intel_cpu::SwishNode)                                  \
     OP_EXTENSION(ov::intel_cpu::NgramNode)                                  \
     OP_EXTENSION(ov::op::internal::NonMaxSuppressionIEInternal)             \

src/plugins/intel_cpu/src/nodes/causal_mask_preprocess.cpp

@@ -0,0 +1,158 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "causal_mask_preprocess.h"
+
+#include "common/bfloat16.hpp"
+#include "common/cpu_memcpy.h"
+#include "cpu/x64/cpu_isa_traits.hpp"
+#include "shape_inference/shape_inference_internal_dyn.hpp"
+#include "utils/plain_tensor.hpp"
+
+#include <chrono>
+#include <string>
+#include <vector>
+
+namespace ov {
+namespace intel_cpu {
+namespace node {
+
+/*
+CausalMaskPreprocess:
+    inputs:
+        0: attention_mask            : i64[N, kv_len]
+                                    0 means mask-out, 1 means attends to
+        1: batch_size (size_Gather)  : i32[1]
+        2: cache_positions  i32[q_len];
+        3: kvLen            i32[1];
+    outputs
+        0: causal mask for SDPA : f32[batch_size, 1, q_len, kvLen]
+
+The functionality is equivalent to following python code:
+
+    ##### preprocess
+    min_dtype = torch.finfo(dtype).min
+    causal_mask = self.causal_mask[None, None, :, :].repeat(batch_size, 1, 1, 1).to(dtype) * min_dtype
+    causal_mask = causal_mask.to(dtype=dtype, device=device)
+
+    mask_length = attention_mask.shape[-1]
+    padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
+    causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype)
+
+    ##### when being used will be further sliced
+    causal_mask = attention_mask
+        if attention_mask is not None and cache_position is not None:
+            causal_mask = causal_mask[:, :, cache_position, : key_states.shape[-2]]
+*/
+template <typename T>
+struct CausalMaskPreprocess::ExecutorCausalMaskPreprocess : public CausalMaskPreprocess::Executor {
+    void execute(dnnl::stream strm,
+                 intel_cpu::Node * pnode,
+                 const intel_cpu::CausalMaskPreprocessNode::Config& config) override {
+        ov::intel_cpu::PlainTensor t_attention_mask(pnode->getSrcMemoryAtPort(0));
+        ov::intel_cpu::PlainTensor t_batch_size(pnode->getSrcMemoryAtPort(1));
+        ov::intel_cpu::PlainTensor t_cache_positions(pnode->getSrcMemoryAtPort(2));
+        ov::intel_cpu::PlainTensor t_kvLen(pnode->getSrcMemoryAtPort(3));
+
+        auto mask_length = t_attention_mask.size(-1);
+        auto batch_size = static_cast<size_t>(*t_batch_size.ptr<int32_t>(0));
+        auto kvLen = static_cast<size_t>(*t_kvLen.ptr<int32_t>(0));
+        auto qLen = t_cache_positions.size(0);
+
+        VectorDims newDims{batch_size, 1, qLen, kvLen};
+        pnode->redefineOutputMemory({newDims});
+        ov::intel_cpu::PlainTensor t_dst(pnode->getDstMemoryAtPort(0));
+
+        DEBUG_LOG("CausalMaskPreprocess::execute", config.type, "  batch_size=", batch_size, " qLen=", qLen, " kvLen=", kvLen);
+        DEBUG_LOG("CausalMaskPreprocess::execute  attention_mask=", t_attention_mask);
+        DEBUG_LOG("CausalMaskPreprocess::execute  cache_positions=", t_cache_positions);
+
+        // raw_causal_mask is already ensured to be triu by transformation
+        auto* prow = t_cache_positions.ptr<int32_t>(0);
+        T min_dtype = std::numeric_limits<T>::lowest();
+
+        parallel_for2d(batch_size, qLen, [&](size_t n, size_t i) {
+            auto* pamask = t_attention_mask.ptr<int32_t>(n, 0);
+            auto* pdst = t_dst.ptr<T>(n, 0, i);
+            auto row = static_cast<size_t>(prow[i]);
+            size_t j = 0;
+            for (; j < mask_length; j++) {
+                bool cmask_eq0 = (j <= row);
+                bool amask_eq0 = (pamask[j] == 0);
+                bool padding_mask = (cmask_eq0 && amask_eq0);
+                pdst[j] = (padding_mask | (!cmask_eq0))? min_dtype : T(0);
+            }
+            for (; j < kvLen; j++) {
+                bool cmask_eq0 = (j <= row);
+                pdst[j] = cmask_eq0 ? T(0) : min_dtype;
+            }
+        });
+        DEBUG_LOG("CausalMaskPreprocess::execute  dst=", t_dst);
+    }
+};
+
+CausalMaskPreprocess::CausalMaskPreprocess(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context)
+    : Node(op, context, InternalDynShapeInferFactory()) {
+    std::string errorMessage;
+    if (!isSupportedOperation(op, errorMessage)) {
+        OPENVINO_THROW("CPU: " + errorMessage);
+    }
+
+    const auto node = std::dynamic_pointer_cast<const intel_cpu::CausalMaskPreprocessNode>(op);
+    m_config = node->get_config();
+}
+
+bool CausalMaskPreprocess::isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept {
+    try {
+        const auto node = std::dynamic_pointer_cast<const intel_cpu::CausalMaskPreprocessNode>(op);
+        if (!node) {
+            errorMessage = "Only CausalMaskPreprocessNode operation is supported";
+            return false;
+        }
+    } catch (...) {
+        return false;
+    }
+    return true;
+}
+
+void CausalMaskPreprocess::initSupportedPrimitiveDescriptors() {
+    if (!supportedPrimitiveDescriptors.empty())
+        return;
+
+    std::vector<ov::element::Type> iprecs = getOriginalInputPrecisions();
+    std::vector<ov::element::Type> oprecs = getOriginalOutputPrecisions();
+
+    // precision preferences
+    if (m_config.type == "CausalMaskPreprocess") {
+        if (oprecs[0] == ov::element::bf16) {
+            m_executor = std::make_shared<ExecutorCausalMaskPreprocess<ov::bfloat16>>();
+        } else {
+            // fallback to default precision
+            m_executor = std::make_shared<ExecutorCausalMaskPreprocess<float>>();
+            oprecs[0] = ov::element::f32;
+        }
+        // all input precisions must be int32
+        for (auto& prec : iprecs) prec = ov::element::i32;
+    } else {
+        OPENVINO_THROW("CPU: CausalMaskPreprocess type not supported : " + m_config.type);
+    }
+
+    std::vector<PortConfigurator> inPortConfigs;
+    for (size_t i = 0; i < getOriginalInputsNumber(); i++)
+        inPortConfigs.emplace_back(LayoutType::ncsp, iprecs[i], getInputShapeAtPort(i), false, -1);
+
+    std::vector<PortConfigurator> outPortConfigs;
+    for (size_t i = 0; i < getOriginalOutputsNumber(); i++)
+        outPortConfigs.emplace_back(LayoutType::ncsp, oprecs[i], getOutputShapeAtPort(i), false, -1);
+
+    addSupportedPrimDesc(inPortConfigs, outPortConfigs, impl_desc_type::ref_any);
+}
+
+void CausalMaskPreprocess::execute(dnnl::stream strm) {
+    m_executor->execute(strm, this, m_config);
+}
+
+}  // namespace node
+}  // namespace intel_cpu
+}  // namespace ov

src/plugins/intel_cpu/src/nodes/causal_mask_preprocess.h

@@ -0,0 +1,46 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "node.h"
+#include "transformations/cpu_opset/common/op/causal_mask_preprocess.hpp"
+
+namespace ov {
+namespace intel_cpu {
+namespace node {
+
+class CausalMaskPreprocess : public Node {
+public:
+    CausalMaskPreprocess(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context);
+
+    void getSupportedDescriptors() override {}
+    bool created() const override {
+        return getType() == Type::CausalMaskPreprocess;
+    }
+    bool needPrepareParams() const override {
+        return false;
+    };
+    void executeDynamicImpl(dnnl::stream strm) override {
+        execute(strm);
+    }
+    void initSupportedPrimitiveDescriptors() override;
+    void execute(dnnl::stream strm) override;
+    static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;
+
+private:
+    struct Executor {
+        virtual void execute(dnnl::stream strm,
+                             intel_cpu::Node * pnode,
+                             const intel_cpu::CausalMaskPreprocessNode::Config& config) = 0;
+    };
+    template <typename T>
+    struct ExecutorCausalMaskPreprocess;
+    intel_cpu::CausalMaskPreprocessNode::Config m_config;
+    std::shared_ptr<Executor> m_executor;
+};
+
+}  // namespace node
+}  // namespace intel_cpu
+}  // namespace ov

src/plugins/intel_cpu/src/nodes_factory.cpp

@@ -96,6 +96,7 @@
 #include "nodes/topk.h"
 #include "nodes/transpose.h"
 #include "nodes/unique.hpp"
+#include "nodes/causal_mask_preprocess.h"
 
 namespace ov {
 namespace intel_cpu {
@@ -183,6 +184,7 @@ Node::NodesFactory::NodesFactory() : Factory("NodesFactory") {
     INTEL_CPU_NODE(Unique, Type::Unique);
     INTEL_CPU_NODE(Ngram, Type::Ngram);
     INTEL_CPU_NODE(RoPE, Type::RoPE);
+    INTEL_CPU_NODE(CausalMaskPreprocess, Type::CausalMaskPreprocess);
     INTEL_CPU_NODE(Interpolate, Type::Interpolate);
     INTEL_CPU_NODE(Inverse, Type::Inverse);
     INTEL_CPU_NODE(RandomUniform, Type::RandomUniform);

src/plugins/intel_cpu/src/transformations/cpu_opset/common/op/causal_mask_preprocess.cpp

@@ -0,0 +1,47 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+#include "causal_mask_preprocess.hpp"
+
+#include <algorithm>
+
+#include "transformations/itt.hpp"
+
+ov::intel_cpu::CausalMaskPreprocessNode::CausalMaskPreprocessNode(const OutputVector& args, const Config& cfg) : Op(args), m_config(cfg) {
+    constructor_validate_and_infer_types();
+}
+
+std::shared_ptr<ov::Node> ov::intel_cpu::CausalMaskPreprocessNode::clone_with_new_inputs(const ov::OutputVector& new_args) const {
+    INTERNAL_OP_SCOPE(CausalMaskPreprocessNode_with_new_inputs);
+    check_new_args_count(this, new_args);
+    return std::make_shared<ov::intel_cpu::CausalMaskPreprocessNode>(new_args, m_config);
+}
+
+void ov::intel_cpu::CausalMaskPreprocessNode::validate_and_infer_types() {
+    INTERNAL_OP_SCOPE(CausalMaskPreprocessNode_validate_and_infer_types);
+    if (m_config.type == "CausalMaskPreprocess") {
+        // inputs:
+        //   0: attention_mask            : i64[N, kv_len]
+        //                             0 means mask-out, 1 means attends to
+        //   1: batch_size (size_Gather)  : i32[1]
+        //   2: cache_positions  i32[q_len];
+        //   3: kvLen            i32[1];
+        // outputs
+        //   0: causal mask for SDPA : f32[batch_size, 1, q_len, kvLen]
+
+        auto batch_size = Dimension::dynamic();
+        auto q_len = get_input_partial_shape(2)[0];
+        auto kv_len = Dimension::dynamic();
+        set_output_type(0, ov::element::f32, {batch_size, 1, q_len, kv_len});
+        return;
+    }
+    NODE_VALIDATION_CHECK(this, false, "unsupported type : ", m_config.type);
+}
+
+bool ov::intel_cpu::CausalMaskPreprocessNode::visit_attributes(ov::AttributeVisitor& visitor) {
+    INTERNAL_OP_SCOPE(CausalMaskPreprocessNode_visit_attributes);
+    visitor.start_structure("config");
+    visitor.on_attribute("type", m_config.type);
+    visitor.finish_structure();
+    return true;
+}

src/plugins/intel_cpu/src/transformations/cpu_opset/common/op/causal_mask_preprocess.hpp

@@ -0,0 +1,43 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/op/op.hpp"
+
+namespace ov {
+namespace intel_cpu {
+
+class CausalMaskPreprocessNode : public ov::op::Op {
+public:
+    OPENVINO_OP("CausalMaskPreprocess", "cpu_plugin_opset");
+
+    CausalMaskPreprocessNode() = default;
+
+    struct Config {
+        std::string type;
+    };
+
+    CausalMaskPreprocessNode(const OutputVector& args, const Config& cfg);
+
+    bool visit_attributes(ov::AttributeVisitor& visitor) override;
+
+    void validate_and_infer_types() override;
+
+    std::shared_ptr<Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
+
+    const Config& get_config() const {
+        return m_config;
+    }
+
+    Config& get_config() {
+        return m_config;
+    }
+
+private:
+    Config m_config;
+};
+
+}  // namespace intel_cpu
+}  // namespace ov