Add eagle_tree_mask as input

Author: GuoliangShiIntel

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

@@ -1094,7 +1094,8 @@ void reshape_to_static(std::shared_ptr<ov::Model> model,
                        const uint32_t kvcache_size,
                        const KVAxesPosition& kv_axes_position,
                        const uint32_t lora_rank,
-                       const uint32_t lhs_seq_size = 0) {
+                       const uint32_t lhs_seq_size = 0,
+                       const bool is_prefill = false) {
     std::map<std::string, ov::PartialShape> new_shapes;
     for (const auto& input : model->inputs()) {
         const auto& input_name = input.get_any_name();
@@ -1129,8 +1130,9 @@ void reshape_to_static(std::shared_ptr<ov::Model> model,
             const auto& partial_shape = input.get_partial_shape();
             new_shape = partial_shape;
             new_shape[0] = 1;  // batch_dim
-        } else if (ov::npuw::matchEagle3HiddenStatesString(input_name)) {
-            new_shape = ov::npuw::Eagle3Extension::get_static_input(model, input, input_size);
+        } else if (ov::npuw::matchEagle3HiddenStatesString(input_name) ||
+                   ov::npuw::matchEagle3TreeMaskString(input_name)) {
+            new_shape = ov::npuw::Eagle3Extension::get_static_input(model, input, input_size, kvcache_size, is_prefill);
         } else if (ov::npuw::util::matchLoRAMatMulAString(input_name)) {
             new_shape = ov::PartialShape({lora_rank, input.get_partial_shape()[1]});
         } else if (ov::npuw::util::matchLoRAMatMulAlphaString(input_name)) {
@@ -2018,18 +2020,21 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& m
     LOG_DEBUG("Make prefill model with static shapes");
     m_max_lora_rank = m_cfg.get<::intel_npu::NPUW_LLM_MAX_LORA_RANK>();
     if (m_use_chunk_prefill) {
-        ReshapeToStatic(static_cast<uint32_t>(m_prefill_chunk_size),
-                        m_kvcache_desc.max_prompt_size,
-                        axes,
-                        m_max_lora_rank)
-            .run_on_model(prefill_model);
+        reshape_to_static(prefill_model,
+                          static_cast<uint32_t>(m_prefill_chunk_size),
+                          m_kvcache_desc.max_prompt_size,
+                          axes,
+                          m_max_lora_rank,
+                          0,
+                          true);
     } else {
-        ReshapeToStatic(m_kvcache_desc.max_prompt_size,
-                        m_kvcache_desc.max_prompt_size,
-                        axes,
-                        m_max_lora_rank,
-                        whisper_lhs_seq_size)
-            .run_on_model(prefill_model);
+        reshape_to_static(prefill_model,
+                          m_kvcache_desc.max_prompt_size,
+                          m_kvcache_desc.max_prompt_size,
+                          axes,
+                          m_max_lora_rank,
+                          whisper_lhs_seq_size,
+                          true);
     }
     LOG_DEBUG("Make kvcache model with static shapes");