Switch NPU LLM execution to ov::genai::StatefulLLMPipeline (#1677)

Co-authored-by: Ilya Lavrenov <ilya.lavrenov@intel.com>

Author: TolyaTalamanov

src/cpp/src/llm_pipeline.cpp

@@ -128,7 +128,9 @@ ov::genai::LLMPipeline::LLMPipeline(
     }
 
     if (m_pimpl == nullptr && device == "NPU") {
-        m_pimpl = static_llm::LLMPipelineFactory::create(models_path, tokenizer, device, properties);
+        m_pimpl = properties.count("STATIC_PIPELINE")
+            ? static_llm::LLMPipelineFactory::create(models_path, tokenizer, properties)
+            : std::make_unique<StatefulLLMPipeline>(models_path, tokenizer, device, properties);
     }
 
     // try to call CB adapter one more time, but with safe guard to silent exception
@@ -166,7 +168,9 @@ ov::genai::LLMPipeline::LLMPipeline(
     }
 
     if (m_pimpl == nullptr && device == "NPU") {
-        m_pimpl = static_llm::LLMPipelineFactory::create(models_path, device, properties);
+        m_pimpl = properties.count("STATIC_PIPELINE")
+            ? static_llm::LLMPipelineFactory::create(models_path, properties)
+            : std::make_unique<StatefulLLMPipeline>(models_path, device, properties);
     }
 
     // try to call CB adapter one more time, but with safe guard to silent exception
@@ -208,13 +212,18 @@ ov::genai::LLMPipeline::LLMPipeline(
     }
 
     if (m_pimpl == nullptr && device == "NPU") {
-        m_pimpl = static_llm::LLMPipelineFactory::create(
-            utils::singleton_core().read_model(model_str, weights_tensor),
-            tokenizer,
-            device,
-            properties,
-            generation_config
-        );
+        m_pimpl = properties.count("STATIC_PIPELINE")
+            ? static_llm::LLMPipelineFactory::create(
+                  utils::singleton_core().read_model(model_str, weights_tensor),
+                  tokenizer,
+                  properties,
+                  generation_config)
+            : std::make_unique<StatefulLLMPipeline>(
+                utils::singleton_core().read_model(model_str, weights_tensor),
+                tokenizer,
+                device,
+                properties,
+                generation_config);
     }
 
     // try to call CB adapter one more time, but with safe guard to silent exception

src/cpp/src/llm_pipeline_stateful.cpp

@@ -17,7 +17,17 @@ StatefulLLMPipeline::StatefulLLMPipeline(
     const ov::genai::Tokenizer& tokenizer,
     OptionalGenerationConfig generation_config)
     : LLMPipelineImplBase(tokenizer, generation_config.has_value() ? *generation_config : GenerationConfig()),
-    m_model_runner(request) {}
+    m_model_runner(request) {
+    auto compiled_model = m_model_runner.get_compiled_model();
+    auto execution_devices = compiled_model.get_property(ov::execution_devices);
+    if (execution_devices[0].find("NPU") != std::string::npos) {
+        OPENVINO_ASSERT(execution_devices.size() == 1u);
+        m_is_npu = true;
+        const auto max_prompt_len = compiled_model.get_property("NPUW_LLM_MAX_PROMPT_LEN").as<uint32_t>();
+        const auto min_response_len = compiled_model.get_property("NPUW_LLM_MIN_RESPONSE_LEN").as<uint32_t>();
+        m_max_kv_cache_size = max_prompt_len + min_response_len;
+    }
+}
 
 StatefulLLMPipeline::StatefulLLMPipeline(
     const std::filesystem::path& models_path,
@@ -29,30 +39,44 @@ StatefulLLMPipeline::StatefulLLMPipeline(
         tokenizer,
         device,
         properties,
-        utils::from_config_json_if_exists(models_path)
+        utils::from_config_json_if_exists(models_path),
+        models_path
     } {}
 
 StatefulLLMPipeline::StatefulLLMPipeline(
     const std::shared_ptr<ov::Model>& model,
     const ov::genai::Tokenizer& tokenizer,
     const std::string& device,
     const ov::AnyMap& properties,
-    const ov::genai::GenerationConfig& generation_config)
+    const ov::genai::GenerationConfig& generation_config,
+    const std::filesystem::path& models_path)
     : LLMPipelineImplBase(tokenizer, generation_config), m_sampler(m_tokenizer) {
     utils::apply_slice_before_matmul_transformation(model);
+    auto kv_pos = ov::genai::utils::get_kv_axes_pos(model);
 
-    if (device.find("NPU") != std::string::npos)
+    if (device.find("NPU") != std::string::npos) {
+        m_is_npu = true;
         m_use_full_chat_history = true;
+    }
 
     if (!m_use_full_chat_history)
-        m_kv_history_trim_manager.kv_cache_seq_length_axis = ov::genai::utils::get_kv_axes_pos(model).seq_len;
+        m_kv_history_trim_manager.kv_cache_seq_length_axis = kv_pos.seq_len;
 
     auto filtered_properties = extract_adapters_from_properties(properties, &m_generation_config.adapters);
     if (m_generation_config.adapters) {
         m_generation_config.adapters->set_tensor_name_prefix("base_model.model.");
         m_adapter_controller = AdapterController(model, *m_generation_config.adapters, device);   // TODO: Make the prefix name configurable
     }
-    ov::CompiledModel compiled_model = utils::singleton_core().compile_model(model, device, *filtered_properties);
+    ov::CompiledModel compiled_model;
+    if (m_is_npu) {
+        utils::KVDesc kv_desc;
+        std::tie(compiled_model, kv_desc) = utils::compile_decoder_for_npu(
+            model, *filtered_properties, kv_pos, models_path
+        );
+        m_max_kv_cache_size = kv_desc.max_prompt_len + kv_desc.min_response_len;
+    } else {
+       compiled_model = utils::singleton_core().compile_model(model, device, *filtered_properties);
+    }
     m_model_runner = compiled_model.create_infer_request();
     ov::genai::utils::print_compiled_model_properties(compiled_model, "Stateful LLM model");
 
@@ -225,12 +249,21 @@ EncodedResults StatefulLLMPipeline::generate(
         config.set_eos_token_id(m_generation_config.eos_token_id);
     config.validate();
 
+    auto batch_size = input_ids.get_shape().at(0);
+
+    if (m_is_npu) {
+        OPENVINO_ASSERT(batch_size == 1u, "Currently only batch size equal to 1 is supported for NPU device!");
+        OPENVINO_ASSERT(config.is_greedy_decoding() || config.is_multinomial(),
+            "Currently only greedy and multinomial decoding are supported for NPU device!");
+        OPENVINO_ASSERT(config.num_return_sequences == 1u,
+            "Currently only \"num_return_sequences\" equal to 1 is supported for NPU device!");
+    }
+
     // Stateful pipeline does not provide logprobs for prompt tokens
     OPENVINO_ASSERT(config.echo == false, "Echo is not supported in the stateful pipeline");
 
     std::shared_ptr<StreamerBase> streamer_ptr = ov::genai::utils::create_streamer(streamer, m_tokenizer);
 
-    auto batch_size = input_ids.get_shape().at(0);
     OPENVINO_ASSERT(streamer_ptr == nullptr || batch_size == 1 && config.num_return_sequences == 1 &&
         (config.is_greedy_decoding() || config.is_multinomial()),
         "Currently streaming is possible only with batch size=1 and only for greedy or multinomial decoding");
@@ -314,7 +347,7 @@ EncodedResults StatefulLLMPipeline::generate(
     }
 
     ov::genai::utils::GenerationFinishInfo finish_info = get_lm_encoded_results(m_model_runner, input_ids, concatenated_attention_mask, streamer_ptr, m_sampler,
-                                                                                requests, position_ids, m_kv_cache_state, std::nullopt, std::nullopt);
+                                                                                requests, position_ids, m_kv_cache_state, std::nullopt, std::nullopt, m_max_kv_cache_size);
     ov::genai::EncodedResults& result = finish_info.results;
     m_chat_generation_finish_status = finish_info.streaming_finish_status;
 

src/cpp/src/llm_pipeline_stateful.hpp

@@ -2,6 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 
+#include <limits>
+
 #include "llm_pipeline_base.hpp"
 #include "lm_encoding.hpp"
 #include "sampler.hpp"
@@ -26,6 +28,8 @@ class StatefulLLMPipeline final : public LLMPipelineImplBase {
     ov::genai::GenerationStatus m_chat_generation_finish_status = ov::genai::GenerationStatus::RUNNING;
     // if True, full history will be used as prompt on each chat generation
     bool m_use_full_chat_history = false;
+    size_t m_max_kv_cache_size = std::numeric_limits<size_t>::max();
+    bool m_is_npu = false;
     // reflection of tokens contained in the kv cache
     KVCacheState m_kv_cache_state;
 
@@ -50,7 +54,8 @@ class StatefulLLMPipeline final : public LLMPipelineImplBase {
         const ov::genai::Tokenizer& tokenizer,
         const std::string& device,
         const ov::AnyMap& config,
-        const ov::genai::GenerationConfig& generation_config
+        const ov::genai::GenerationConfig& generation_config,
+        const std::filesystem::path& models_path = {}
     );
 
     StatefulLLMPipeline(