Avoid RoPE on NPU for Gemma Models (#32093)

Computing RoPE for Gemma Like Models in NPU has some accuracy issues.

This patch adds a matcher pass in avoid.cpp of npuw pipeline to create
an NPU avoid pattern for RoPE computation in Gemma like models.

This is added as an avoid pattern and to enable this the below flag :
"NPUW_ONLINE_AVOID": "P:GemmaRoPE/NPU" must be added to config. To
recognize this pattern check is added to "snapshot.cpp".

Signed-off-by: Ghosh, Tamoghna <[email protected]>

---------

Signed-off-by: Ghosh, Tamoghna <[email protected]>
Signed-off-by: tamoghna <[email protected]>
Co-authored-by: Alexey Smirnov <[email protected]>

Author: tamoghintel

src/plugins/intel_npu/src/plugin/npuw/partitioning/online/snapshot.cpp

@@ -451,18 +451,22 @@ void Snapshot::earlyAvoids() {
         }
         case PatternType::PATTERN: {
             // FIXME: refactor as more patterns are supported
-            if (avoid.pattern != "RMSNorm" && avoid.pattern != "SinCos") {
-                LOG_WARN(
-                    "OPENVINO_NPUW_AVOID only supports RMSNorm and SinCos as patterns (don't confuse with operations)."
-                    << " Avoid pattern " << avoid.pattern << " is skipped!");
+            if (avoid.pattern != "RMSNorm" && avoid.pattern != "SinCos" && avoid.pattern != "GemmaRoPE") {
+                LOG_WARN("OPENVINO_NPUW_AVOID only supports RMSNorm, SinCos and GemmaRoPE as patterns "
+                         "(don't confuse with operations). "
+                         "Avoid pattern "
+                         << avoid.pattern << " is skipped!");
                 break;
             }
             handle_patterns = true;
             if (avoid.pattern == "RMSNorm") {
                 rewr.add_matcher<ov::npuw::patterns::avoid::RMSNorm>(shared_from_this(), avoid.device);
             } else if (avoid.pattern == "SinCos") {
                 rewr.add_matcher<ov::npuw::patterns::avoid::SinCos>(shared_from_this(), avoid.device);
+            } else if (avoid.pattern == "GemmaRoPE") {
+                rewr.add_matcher<ov::npuw::patterns::avoid::GemmaRoPE>(shared_from_this(), avoid.device);
             }
+
             break;
         }
         }

src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/avoid.cpp

@@ -105,6 +105,35 @@ SinCos::SinCos(const std::shared_ptr<ov::npuw::online::Snapshot>& snapshot, cons
     };
     register_matcher(std::make_shared<opp::Matcher>(sin_cos, "TagSinCos"), std::move(callback));
 }
+GemmaRoPE::GemmaRoPE(const std::shared_ptr<ov::npuw::online::Snapshot>& snapshot, const std::string& avoid_device) {
+    auto power = opp::wrap_type<ov::op::v1::Power>({opp::any_input(), opp::any_input()});
+    auto unsqueeze1 = opp::wrap_type<ov::op::v0::Unsqueeze>({power, opp::wrap_type<ov::op::v0::Constant>()});
+    auto unsqueeze2 = opp::wrap_type<ov::op::v0::Unsqueeze>({unsqueeze1, opp::wrap_type<ov::op::v0::Constant>()});
+    auto divide = opp::wrap_type<ov::op::v1::Divide>({opp::wrap_type<ov::op::v0::Convert>(), unsqueeze2});
+    auto unsqueeze3 = opp::wrap_type<ov::op::v0::Unsqueeze>({divide, opp::wrap_type<ov::op::v0::Constant>()});
+    auto sin_cos = opp::wrap_type<ov::op::v0::Sin, ov::op::v0::Cos>({unsqueeze3});
+    auto node_to_gptr = snapshot->getNodeToGroupMap();
+
+    auto callback = [=](ov::pass::pattern::Matcher& m) {
+        auto& node_to_output = m.get_pattern_value_map();
+
+        auto matched_power = node_to_output.at(power).get_node_shared_ptr();
+        auto matched_unsqueeze1 = node_to_output.at(unsqueeze1).get_node_shared_ptr();
+        auto matched_unsqueeze2 = node_to_output.at(unsqueeze2).get_node_shared_ptr();
+        auto matched_divide = node_to_output.at(divide).get_node_shared_ptr();
+        auto matched_unsqueeze3 = node_to_output.at(unsqueeze3).get_node_shared_ptr();
+        auto matched_sin_cos = node_to_output.at(sin_cos).get_node_shared_ptr();
+
+        node_to_gptr->at(matched_power)->avoid(avoid_device);
+        node_to_gptr->at(matched_unsqueeze1)->avoid(avoid_device);
+        node_to_gptr->at(matched_unsqueeze2)->avoid(avoid_device);
+        node_to_gptr->at(matched_divide)->avoid(avoid_device);
+        node_to_gptr->at(matched_unsqueeze3)->avoid(avoid_device);
+        node_to_gptr->at(matched_sin_cos)->avoid(avoid_device);
+        return false;
+    };
+    register_matcher(std::make_shared<opp::Matcher>(sin_cos, "TagGemmaRoPE"), std::move(callback));
+}
 }  // namespace avoid
 }  // namespace patterns
 }  // namespace npuw

src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/avoid.hpp

@@ -34,6 +34,12 @@ class SinCos : public ov::pass::MatcherPass {
     SinCos(const std::shared_ptr<ov::npuw::online::Snapshot>& snapshot, const std::string& avoid_device);
 };
 
+class GemmaRoPE : public ov::pass::MatcherPass {
+public:
+    OPENVINO_MATCHER_PASS_RTTI("npuw::patterns::avoid::GemmaRoPE");
+    GemmaRoPE(const std::shared_ptr<ov::npuw::online::Snapshot>& snapshot, const std::string& avoid_device);
+};
+
 }  // namespace avoid
 }  // namespace patterns
 }  // namespace npuw