[NPUW] Fixed chunked prefill for new LFM2 IR (#36933)

### Details:
- *This is a cherry-pick of
#36929 to master*
 - *Fixed chunked prefill with new IR version of LFM2*:
- *New IR version wires `attention_mask` with Conv states by slicing
that `attention_mask` from the left on the size of current prompt.*
- *NPUW places new tokens in the right end of the `input_ids` and
`attention_mask` in static shape scenario, so mask should be sliced from
the right end instead.*

### Tickets:
 - *CVS-190129*
 - *CVS-190528*

### AI Assistance:
 - *AI assistance used: yes*
- *Assistant is used for parallel debugging, small issues correction in
the fix and tests*

Author: AsyaPronina

src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp

@@ -20,6 +20,7 @@
 #include "npuw_transformations/replace_deepstack_scatter_with_add.hpp"
 #include "npuw_transformations/reshape_sliced_head_to_static.hpp"
 #include "npuw_transformations/reshape_to_static.hpp"
+#include "npuw_transformations/right_align_mask_slice_for_conv.hpp"
 #include "npuw_transformations/slice_out_embeds.hpp"
 #include "npuw_transformations/split_kvcache_into_blocks.hpp"
 #include "openvino/op/convert.hpp"
@@ -821,6 +822,8 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& m
     } else {
         LOG_DEBUG("Adding position_ids input in case it doesn't exist in model: LFM-2 case.");
         ov::npuw::AddPositionIdsParam().run_on_model(kvcache_model);
+        LOG_DEBUG("Right-align attention_mask slice for Conv operations: LFM-2 case.");
+        ov::npuw::RightAlignMaskSliceForConv().run_on_model(kvcache_model);
         LOG_DEBUG("Transform kvcache model from stateful to stateless.");
         ov::pass::StatefulToStateless().run_on_model(kvcache_model);
     }

src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/right_align_mask_slice_for_conv.cpp

@@ -0,0 +1,107 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "right_align_mask_slice_for_conv.hpp"
+
+#include "openvino/op/ops.hpp"
+#include "openvino/openvino.hpp"
+#include "openvino/pass/manager.hpp"
+#include "openvino/pass/matcher_pass.hpp"
+#include "openvino/pass/pass.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "openvino/pass/validate.hpp"
+
+namespace opp = ov::pass::pattern;
+
+namespace {
+
+// diagnostics warnings on OPENVINO_MATCHER_PASS_RTTI() definition: visibility hidden
+#ifdef __GNUC__
+#    pragma GCC diagnostic push
+#    pragma GCC diagnostic ignored "-Wattributes"
+#endif
+
+class RightAlignMaskSliceForConvImpl : public ov::pass::MatcherPass {
+public:
+    // In LFM2-like models, where `attention_mask` is also consumed by the Conv subgraph,
+    // the Slice operation is inserted to extract the current tokens from that `attention_mask`
+    // and pass to the Conv subgraph.
+    // This is done to allow skipping processing of some tokens in the prefill stage.
+    // Slice is done from zero till the number of current tokens, assuming that in most cases
+    // tokens from the right are not the useful ones.
+    // However, in NPUW static-shape mode, we place useful current tokens exactly to the right end
+    // in the prefill stage.
+    // Therefore, Slice operation will extract the left part of mask, that in case of chunked
+    // prefill will be padded either with zeroes (in the first iteration) or with ones
+    // (in subsequent iterations) but for past tokens.
+    // This pass re-align the Slice operation to extract the right part of the mask,
+    // that contains useful current tokens.
+    OPENVINO_MATCHER_PASS_RTTI("ov::npuw::RightAlignMaskSliceForConvImpl");
+    explicit RightAlignMaskSliceForConvImpl() {
+        auto attention_mask = opp::wrap_type<ov::op::v0::Parameter>();
+        auto attention_mask_slice = opp::wrap_type<ov::op::v8::Slice>(
+            {attention_mask, opp::any_input(), opp::any_input(), opp::any_input(), opp::any_input()});
+        auto unsqueeze = opp::wrap_type<ov::op::v0::Unsqueeze>({attention_mask_slice, opp::any_input()});
+        auto convert = opp::wrap_type<ov::op::v0::Convert>({unsqueeze});
+        auto multiply = opp::wrap_type<ov::op::v1::Multiply>({convert, opp::any_input()});
+        auto add = opp::wrap_type<ov::op::v1::Add>({multiply, opp::any_input()});
+        auto multiply_with_embedd = opp::wrap_type<ov::op::v1::Multiply>({opp::any_input(), add});
+        auto matmul = opp::wrap_type<ov::op::v0::MatMul>({multiply_with_embedd, opp::any_input()});
+        auto transpose = opp::wrap_type<ov::op::v1::Transpose>({matmul, opp::any_input()});
+        // VariadicSplit into B,C,~h in LFM2.
+        auto variadic_split =
+            opp::wrap_type<ov::op::v1::VariadicSplit>({transpose, opp::any_input(), opp::any_input()});
+
+        auto callback = [=](opp::Matcher& m) {
+            auto& node_to_output = m.get_pattern_value_map();
+
+            auto matched_attention_mask = node_to_output.at(attention_mask).get_node_shared_ptr();
+            const auto names = matched_attention_mask->output(0).get_names();
+            if (names.count(ov::npuw::attention_mask_name) == 0 &&
+                matched_attention_mask->get_friendly_name() != ov::npuw::attention_mask_name) {
+                return false;
+            }
+            auto matched_attention_mask_slice = node_to_output.at(attention_mask_slice).get_node_shared_ptr();
+            const auto slice_start_const = ov::as_type_ptr<ov::op::v0::Constant>(
+                matched_attention_mask_slice->input_value(1).get_node_shared_ptr());
+            if (!slice_start_const || slice_start_const->cast_vector<int64_t>() != std::vector<int64_t>{0}) {
+                return false;
+            }
+            auto matched_attention_mask_unsqueeze = node_to_output.at(unsqueeze).get_node_shared_ptr();
+
+            auto const_zero = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 0);
+            auto const_one = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 1);
+            auto mask_shape_of = std::make_shared<ov::op::v3::ShapeOf>(matched_attention_mask);
+            auto mask_len = std::make_shared<ov::op::v8::Gather>(mask_shape_of, const_one, const_zero);
+            auto mask_len_rank_one = std::make_shared<ov::op::v0::Unsqueeze>(mask_len, const_zero);
+            // 2nd argument to Slice is stop point:
+            auto current_len = matched_attention_mask_slice->input(2).get_source_output().get_node_shared_ptr();
+            auto current_offset = std::make_shared<ov::op::v1::Subtract>(mask_len, current_len);
+
+            auto const_one_rank_one = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{1}, 1);
+            auto new_slice = std::make_shared<ov::op::v8::Slice>(matched_attention_mask,
+                                                                 current_offset,
+                                                                 mask_len_rank_one,
+                                                                 const_one_rank_one,
+                                                                 const_one_rank_one);
+            matched_attention_mask_unsqueeze->input(0).replace_source_output(new_slice);
+            return true;
+        };
+        register_matcher(std::make_shared<opp::Matcher>(variadic_split, "RightAlignMaskSliceForConvImpl"),
+                         std::move(callback));
+    }
+};
+
+#ifdef __GNUC__
+#    pragma GCC diagnostic pop
+#endif
+}  // anonymous namespace
+
+bool ov::npuw::RightAlignMaskSliceForConv::run_on_model(const std::shared_ptr<ov::Model>& model) {
+    ov::pass::Manager manager("right-align-mask-slice-for-conv");
+    manager.set_per_pass_validation(true);
+    manager.register_pass<RightAlignMaskSliceForConvImpl>();
+    manager.run_passes(model);
+    return true;
+}

src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/right_align_mask_slice_for_conv.hpp

@@ -0,0 +1,18 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/openvino.hpp"
+
+namespace ov::npuw {
+
+constexpr const char* attention_mask_name = "attention_mask";
+
+class RightAlignMaskSliceForConv : public ov::pass::ModelPass {
+public:
+    OPENVINO_MODEL_PASS_RTTI("ov::npuw::RightAlignMaskSliceForConv");
+    bool run_on_model(const std::shared_ptr<ov::Model>& model) override;
+};
+}  // namespace ov::npuw

src/plugins/intel_npu/tests/unit/CMakeLists.txt

@@ -109,6 +109,7 @@ ov_add_test_target(
             ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/reshape_sliced_head_to_static.cpp
             ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/reshape_to_static.cpp
             ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/replace_deepstack_scatter_with_add.cpp
+            ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/right_align_mask_slice_for_conv.cpp
             ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/slice_out_embeds.cpp
             # kokoro pipeline
             ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/pipelines/kokoro/kokoro_compiled_model.cpp
@@ -161,6 +162,7 @@ target_sources(${TARGET_NAME} PRIVATE
         ${CMAKE_CURRENT_SOURCE_DIR}/npuw/subgraph_behavior_infer_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/npuw/preserve_const_matmul_pattern_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/npuw/pipeline_passes/add_position_ids_param_test.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/npuw/pipeline_passes/right_align_mask_slice_for_conv_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/npuw/pipeline_passes/collapse_unqdq_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/npuw/pipeline_passes/conv_to_matmul_test.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/npuw/pipeline_passes/patch_sliding_mask_test.cpp