rebase + update

Author: sbalandi

src/cpp/src/llm_pipeline_stateful.cpp

@@ -212,7 +212,6 @@ EncodedResults StatefulLLMPipeline::generate(
         reset_kv_state();
         m_model_runner.get_tensor("attention_mask").set_shape({1, 0});
         m_kv_cache_state.reset_state();
-        m_kv_history_trim_manager.reset();
     }
 
     auto start_time = std::chrono::steady_clock::now();

src/cpp/src/visual_language/inputs_embedder.cpp

@@ -39,16 +39,17 @@ void InputsEmbedder::IInputsEmbedder::start_chat(const std::string& system_messa
     m_history = {{{"role", "system"}, {"content", system_message}}};
 }
 
-void InputsEmbedder::IInputsEmbedder::update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status, size_t processed_tokens_amount) {
+void InputsEmbedder::IInputsEmbedder::update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status) {
     m_kv_cache_state.num_tokens_to_trim = 0;
     if (generation_finish_status == ov::genai::GenerationStatus::CANCEL) {
         // If chat generation process was cancelled by user, let's rollback to previous state of history
         m_history.pop_back();
 
         std::vector<int64_t>& state = m_kv_cache_state.get_state();
-        state.resize(state.size() - processed_tokens_amount);
+
+        m_kv_cache_state.num_tokens_to_trim = state.size() - m_prev_hist_length;
+        state.resize(m_prev_hist_length);
         m_kv_cache_state.reset_mem_state = state.empty();
-        m_kv_cache_state.num_tokens_to_trim = processed_tokens_amount;
     } else {
         // Tail of chat template is missing in KV cache.
         // Find the tail to concatenate it with the next input prompt.
@@ -142,6 +143,7 @@ ov::Tensor InputsEmbedder::IInputsEmbedder::update_history(const ov::Tensor& new
 ov::Tensor InputsEmbedder::IInputsEmbedder::get_encoded_input_ids(const std::string& prompt, ov::genai::VLMPerfMetrics& metrics) {
     const auto new_chat_tokens = apply_chat_template_tokenize(prompt, metrics);
     auto new_input_ids = update_history(new_chat_tokens);
+    m_prev_hist_length = m_kv_cache_state.get_state().size();
     m_kv_cache_state.add_inputs(new_input_ids);
 
     return new_input_ids;
@@ -244,8 +246,8 @@ void InputsEmbedder::start_chat(const std::string& system_message) {
     return m_impl->start_chat(system_message);
 }
 
-void InputsEmbedder::update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status, size_t processed_tokens_amount) {
-    return m_impl->update_chat_history(decoded_results, generation_finish_status, processed_tokens_amount);
+void InputsEmbedder::update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status) {
+    return m_impl->update_chat_history(decoded_results, generation_finish_status);
 }
 
 void InputsEmbedder::set_apply_chat_template_status(bool apply_chat_template) {

src/cpp/src/visual_language/inputs_embedder.hpp

@@ -51,7 +51,7 @@ class InputsEmbedder {
     void start_chat(const std::string& system_message);
 
     // adds currently generated text to chat history
-    void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status, size_t processed_tokens_amount);
+    void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status);
 
     // set the apply_chat_template flag, which determines whether chat template should be applied for non-chat scenarios
     void set_apply_chat_template_status(bool apply_chat_template);
@@ -83,6 +83,8 @@ class InputsEmbedder {
         ov::genai::GenerationStatus m_chat_generation_finish_status = ov::genai::GenerationStatus::RUNNING;
         // reflection of tokens contained in the kv cache
         utils::KVCacheState m_kv_cache_state;
+        // length of attention_mask/kv cache at the beginning of generation()
+        size_t m_prev_hist_length = 0;
     public:
         virtual ov::Tensor get_inputs_embeds(const std::string& prompt, const std::vector<ov::Tensor>& images, ov::genai::VLMPerfMetrics& metrics) = 0;
 
@@ -106,7 +108,7 @@ class InputsEmbedder {
 
         virtual void start_chat(const std::string& system_message);
 
-        virtual void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status, size_t processed_tokens_amount);
+        virtual void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status);
 
         virtual void finish_chat();
 

src/cpp/src/visual_language/minicpm/classes.cpp

@@ -667,8 +667,8 @@ ov::Tensor InputsEmbedderMiniCPM::get_inputs_embeds(const std::string& prompt, c
     return inputs_embeds;
 }
 
-void InputsEmbedderMiniCPM::update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status, size_t processed_tokens_amount) {
-    IInputsEmbedder::update_chat_history(decoded_results, generation_finish_status, processed_tokens_amount);
+void InputsEmbedderMiniCPM::update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status) {
+    IInputsEmbedder::update_chat_history(decoded_results, generation_finish_status);
     if (generation_finish_status == ov::genai::GenerationStatus::CANCEL)
         m_image_id = m_prev_image_id;
     else

src/cpp/src/visual_language/minicpm/classes.hpp

@@ -49,7 +49,7 @@ class InputsEmbedderMiniCPM : public InputsEmbedder::IInputsEmbedder {
 
     ov::Tensor get_inputs_embeds(const std::string& prompt, const std::vector<ov::Tensor>& images, ov::genai::VLMPerfMetrics& metrics) override;
 
-    void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status, size_t processed_tokens_amount) override;
+    void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status) override;
 
     void start_chat(const std::string& system_message) override;
 

src/cpp/src/visual_language/phi3_vision/classes.cpp

@@ -550,6 +550,7 @@ ov::Tensor InputsEmbedderPhi3V::get_inputs_embeds(const std::string& prompt, con
     }
     ov::Tensor new_merged_tokens = insert_image_placeholders(new_chat_tokens, m_tokens_per_images);
     ov::Tensor new_tokens = update_history(new_merged_tokens);
+    m_prev_hist_length = m_kv_cache_state.get_state().size();
     m_kv_cache_state.add_inputs(new_tokens);
 
     std::vector<ov::Tensor> tokens = drop_image_placeholders(new_tokens);
@@ -604,8 +605,8 @@ ov::Tensor InputsEmbedderPhi3V::get_inputs_embeds(const std::string& prompt, con
     return inputs_embeds;
 }
 
-void InputsEmbedderPhi3V::update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status, size_t processed_tokens_amount) {
-    IInputsEmbedder::update_chat_history(decoded_results, generation_finish_status, processed_tokens_amount);
+void InputsEmbedderPhi3V::update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status) {
+    IInputsEmbedder::update_chat_history(decoded_results, generation_finish_status);
     if (generation_finish_status == ov::genai::GenerationStatus::CANCEL)
         m_tokens_per_images = m_prev_tokens_per_images;
     else

src/cpp/src/visual_language/phi3_vision/classes.hpp

@@ -30,7 +30,7 @@ class InputsEmbedderPhi3V : public InputsEmbedder::IInputsEmbedder {
 
     ov::Tensor get_inputs_embeds(const std::string& prompt, const std::vector<ov::Tensor>& images, ov::genai::VLMPerfMetrics& metrics) override;
 
-    void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status, size_t processed_tokens_amount) override;
+    void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status) override;
 
     void start_chat(const std::string& system_message) override;
 

src/cpp/src/visual_language/pipeline.cpp

@@ -90,9 +90,6 @@ class VLMPipeline::VLMPipelineImpl : public VLMPipelineBase{
         ov::genai::utils::print_compiled_model_properties(compiled_language_model, "VLM language model");
 
         m_language = compiled_language_model.create_infer_request();
-
-        utils::KVCacheState& kv_cache_state = m_inputs_embedder->get_kv_cache_state();
-        kv_cache_state.seq_length_axis = kv_pos.seq_len;
         m_language.get_tensor("attention_mask").set_shape({1, 0});
 
         auto embedder_properties = device_propertes.empty()
@@ -102,6 +99,9 @@ class VLMPipeline::VLMPipelineImpl : public VLMPipelineBase{
         m_tokenizer = m_inputs_embedder->get_tokenizer();
         m_embedding = m_inputs_embedder->get_embedding_model();
 
+        utils::KVCacheState& kv_cache_state = m_inputs_embedder->get_kv_cache_state();
+        kv_cache_state.seq_length_axis = kv_pos.seq_len;
+
         // If eos_token_id was not provided, take value
         if (m_generation_config.eos_token_id == -1) {
             m_generation_config.set_eos_token_id(m_tokenizer.get_eos_token_id());
@@ -236,7 +236,7 @@ class VLMPipeline::VLMPipelineImpl : public VLMPipelineBase{
 
         std::string decoded_results = decoded.texts.at(0);
         if (m_is_chat_conversation)
-            m_inputs_embedder->update_chat_history(decoded_results, finish_info.streaming_finish_status, m_language.get_tensor("attention_mask").get_shape().at(1) - history_size);
+            m_inputs_embedder->update_chat_history(decoded_results, finish_info.streaming_finish_status);
         else
             kv_cache_state.reset_state();
 

tests/python_tests/test_vlm_pipeline.py

@@ -224,9 +224,11 @@ def test_vlm_with_scheduler_vs_default(config, cache):
 @pytest.mark.nightly
 @pytest.mark.parametrize("model_id", model_ids)
 @pytest.mark.parametrize("system_message", ["", "You are a helpful assistant."])
-@pytest.mark.parametrize("iteration_images", [[image_links_for_testing[0], image_links_for_testing[0]], [image_links_for_testing[0], image_links_for_testing[2], image_links_for_testing[0]],
-                                              [image_links_for_testing[1], image_links_for_testing[1]], [image_links_for_testing[1], image_links_for_testing[1], image_links_for_testing[1]],
-                                              [image_links_for_testing[2], image_links_for_testing[1]], [image_links_for_testing[2], image_links_for_testing[0], image_links_for_testing[1]]])
+@pytest.mark.parametrize("iteration_images", [[image_links_for_testing[0], image_links_for_testing[0]], # generation with text input only
+                                              [image_links_for_testing[0], image_links_for_testing[2], image_links_for_testing[0]], # combination of generations with text input and image input, empty string first
+                                              [image_links_for_testing[2], image_links_for_testing[1]], # text + image input
+                                              [image_links_for_testing[2], image_links_for_testing[0], image_links_for_testing[1]]] # combination of generations with text input and image input, image input first
+                         )
 def test_vlm_pipeline_chat(model_id, system_message, iteration_images, cache):
     def streamer(word: str) -> bool:
         nonlocal result_from_streamer