Implement CANCEL for streaming with VLM Pipeline

Author: sbalandi

src/cpp/src/visual_language/inputs_embedder.cpp

@@ -35,8 +35,10 @@ class InputsEmbedder::IInputsEmbedder {
     ChatHistory m_history;
     // Templated chat history
     std::string m_templated_chat_history;
-    // Tokenized chat history
+    // Tokenized history
     std::vector<int64_t> m_tokenized_history;
+    // Tokenized chat history on previous step
+    std::vector<int64_t> m_prev_tokenized_history;
     // Tail of previous output for LM in chat mode is missing in KV cache.
     std::optional<int64_t> m_last_disappeared_token = std::nullopt;
     // If sequence contains some symbols, which could be ambiguous encoded by tokenizer, we need to trim kv cache
@@ -72,21 +74,32 @@ class InputsEmbedder::IInputsEmbedder {
         return m_kv_history_manager.num_tokens_to_remove_from_kv_cache;
     }
 
+    bool should_reset_kv_cache() const {
+        return m_kv_history_manager.reset_kv_cache;
+    }
+
     void set_stop_token_ids(const std::set<int64_t>& stop_token_ids) {
         m_stop_token_ids = stop_token_ids;
     }
 
-    void update_tokenized_history(const std::vector<int64_t>& encoded_result, std::optional<int64_t> last_disappeared_token, bool is_beam_search, size_t last_answer_len) {
+    void update_tokenized_history(const ov::genai::utils::GenerationFinishInfo generation_finish_info, bool is_beam_search, size_t last_answer_len, size_t inputs_embeds_size) {
         if (is_beam_search) {
             m_kv_history_manager.trusted_history_length = m_tokenized_history.size();
             m_kv_history_manager.num_tokens_to_remove_from_kv_cache = last_answer_len;
         } else {
             m_kv_history_manager.reset();
         }
 
-        m_last_disappeared_token = last_disappeared_token;
-  
-        std::copy(encoded_result.begin(), encoded_result.end(), std::back_inserter(m_tokenized_history));
+        m_last_disappeared_token = generation_finish_info.probably_disappeared_token;
+
+        if (generation_finish_info.streaming_finish_status == ov::genai::GenerationStatus::CANCEL) {
+            // let's remove last answer and prompt
+            m_kv_history_manager.num_tokens_to_remove_from_kv_cache = inputs_embeds_size + last_answer_len;
+            m_tokenized_history = std::move(m_prev_tokenized_history);
+        } else {
+            auto encoded_result = generation_finish_info.results.tokens[0];
+            std::copy(encoded_result.begin(), encoded_result.end(), std::back_inserter(m_tokenized_history));
+        }
     }
 
     void set_apply_chat_template_status(bool apply_chat_template) {
@@ -100,6 +113,7 @@ class InputsEmbedder::IInputsEmbedder {
             m_history.clear();
             m_templated_chat_history.clear();
             m_tokenized_history.clear();
+            m_prev_tokenized_history.clear();
         }
         if (system_message.empty()) {
             return;
@@ -109,11 +123,16 @@ class InputsEmbedder::IInputsEmbedder {
         m_templated_chat_history = m_tokenizer.apply_chat_template(m_history, add_generation_prompt);
     }
 
-    void update_chat_history(const std::string& decoded_results) {
-        // Tail of chat template is missing in KV cache.
-        // Find the tail to concatenate it with the next input prompt.
-        m_templated_chat_history.append(decoded_results);
-        m_history.push_back({{"role", "assistant"}, {"content", decoded_results}});
+    void update_chat_history(const std::string& decoded_results, const ov::genai::GenerationStatus generation_finish_status) {
+        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();
+        } else {
+            // Tail of chat template is missing in KV cache.
+            // Find the tail to concatenate it with the next input prompt.
+            m_templated_chat_history.append(decoded_results);
+            m_history.push_back({{"role", "assistant"}, {"content", decoded_results}});
+        }
     }
 
     virtual void finish_chat() {
@@ -123,6 +142,7 @@ class InputsEmbedder::IInputsEmbedder {
         m_history.clear();
         m_templated_chat_history.clear();
         m_tokenized_history.clear();
+        m_prev_tokenized_history.clear();
     }
 
 protected:
@@ -213,6 +233,9 @@ class InputsEmbedder::IInputsEmbedder {
                 trusted_history_length = ov::genai::utils::get_first_history_difference(prev_chat_tokens, m_tokenized_history, m_stop_token_ids);
             }
 
+            m_prev_tokenized_history.clear();
+            std::copy_n(prev_chat_tokens.data<int64_t>(), prev_chat_tokens.get_size(), std::back_inserter(m_prev_tokenized_history));
+
             if (m_tokenized_history.empty()) {
                 encoded_input_ids = new_chat_tokens;
 
@@ -223,9 +246,14 @@ class InputsEmbedder::IInputsEmbedder {
                 if (m_kv_history_manager.does_history_cache_need_to_update()) {
                     trusted_history_length = m_kv_history_manager.trusted_history_length;
                 } else {
-                    m_kv_history_manager.num_tokens_to_remove_from_kv_cache = m_tokenized_history.size() - trusted_history_length;
+                    auto num_tokens_to_remove_from_kv_cache = m_tokenized_history.size() - trusted_history_length;
                     // last generated token is present in tokenized_history, but not included to attention mask, let's keep it in history
-                    m_kv_history_manager.num_tokens_to_remove_from_kv_cache -= 1;
+                    num_tokens_to_remove_from_kv_cache -= 1;
+
+                    // if streaming was used and cancelled on prev step, m_kv_history_manager.num_tokens_to_remove_from_kv_cache could be already set
+                    // and it would be bigger as it includes answer + prompt
+                    m_kv_history_manager.num_tokens_to_remove_from_kv_cache = m_kv_history_manager.num_tokens_to_remove_from_kv_cache > num_tokens_to_remove_from_kv_cache ?
+                                                                              m_kv_history_manager.num_tokens_to_remove_from_kv_cache : num_tokens_to_remove_from_kv_cache;
                 }
 
                 ov::Tensor new_tensor = ov::Tensor(new_chat_tokens.get_element_type(),
@@ -2040,14 +2068,18 @@ std::vector<int64_t> InputsEmbedder::get_tokenized_history() const {
     return m_impl->get_tokenized_history();
 }
 
-void InputsEmbedder::update_tokenized_history(const std::vector<int64_t>& encoded_result, std::optional<int64_t> last_disappeared_token, bool is_beam_search, size_t last_answer_len) {
-    return m_impl->update_tokenized_history(encoded_result, last_disappeared_token, is_beam_search, last_answer_len);
+void InputsEmbedder::update_tokenized_history(const ov::genai::utils::GenerationFinishInfo generation_finish_info, bool is_beam_search, size_t last_answer_len, size_t inputs_embeds_size) {
+    return m_impl->update_tokenized_history(generation_finish_info, is_beam_search, last_answer_len, inputs_embeds_size);
 }
 
 size_t InputsEmbedder::get_num_tokens_to_remove_from_hist() const {
     return m_impl->get_num_tokens_to_remove_from_hist();
 }
 
+bool InputsEmbedder::should_reset_kv_cache() const {
+    return m_impl->should_reset_kv_cache();
+}
+
 Tokenizer InputsEmbedder::get_tokenizer() const {
     return m_impl->get_tokenizer();
 }
@@ -2056,8 +2088,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) {
-    return m_impl->update_chat_history(decoded_results);
+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

@@ -7,6 +7,7 @@
 #include <vector>
 #include <filesystem>
 
+#include "utils.hpp"
 #include "openvino/genai/tokenizer.hpp"
 #include "openvino/genai/visual_language/pipeline.hpp"
 #include "openvino/runtime/tensor.hpp"
@@ -49,16 +50,19 @@ class InputsEmbedder {
     std::vector<int64_t> get_tokenized_history() const;
 
     // add new results to tokenized history
-    void update_tokenized_history(const std::vector<int64_t>& encoded_result, std::optional<int64_t> last_disappeared_token, bool is_beam_search, size_t last_answer_len);
+    void update_tokenized_history(const ov::genai::utils::GenerationFinishInfo generation_finish_info, bool is_beam_search, size_t last_answer_len, size_t inputs_embeds_size);
 
     // returns amount of elements, which need to remove from the end of the KV cache
     size_t get_num_tokens_to_remove_from_hist() const;
 
+    // returns true, if we need to remove full kv cache, in that case it's needed to reset it instead of manually updating
+    bool should_reset_kv_cache() const;
+
     // starts chat and adds optional system_message to chat history
     void start_chat(const std::string& system_message);
 
     // adds currently generated text to chat history
-    void update_chat_history(const std::string& decoded_results);
+    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);

src/cpp/src/visual_language/pipeline.cpp

@@ -169,7 +169,10 @@ class ov::genai::VLMPipeline::VLMPipelineImpl {
         if (generation_config.eos_token_id == -1)
             generation_config.set_eos_token_id(m_generation_config.eos_token_id);
         generation_config.validate();
-        
+
+        // keep it in case of generation will be canceled
+        auto prev_tokenized_history = m_inputs_embedder->get_tokenized_history();
+
         m_inputs_embedder->set_stop_token_ids(generation_config.stop_token_ids);
 
         m_inputs_embedder->set_apply_chat_template_status(generation_config.apply_chat_template);
@@ -179,7 +182,12 @@ class ov::genai::VLMPipeline::VLMPipelineImpl {
         auto end_get_inputs_embeds = std::chrono::steady_clock::now();
 
         auto to_remove_from_hist = m_inputs_embedder->get_num_tokens_to_remove_from_hist();
-        ov::genai::utils::trim_kv_cache(m_language, to_remove_from_hist, m_kv_cache_seq_length_axis, std::nullopt);
+        if (m_inputs_embedder->should_reset_kv_cache())
+            m_language.reset_state();
+        else
+            ov::genai::utils::trim_kv_cache(m_language, to_remove_from_hist, m_kv_cache_seq_length_axis, std::nullopt);
+
+        size_t attention_mask_size = m_language.get_tensor("attention_mask").get_shape().at(1);
 
         std::vector<SequenceGroup::Ptr> requests;
         size_t request_id = 0;
@@ -218,7 +226,6 @@ class ov::genai::VLMPipeline::VLMPipelineImpl {
                                                                                              position_ids, m_embedding, rope_delta);
         ov::genai::EncodedResults& encoded_result = finish_info.results;
 
-
         auto decode_start_time = std::chrono::steady_clock::now();
         VLMDecodedResults decoded;
         for (size_t idx = 0; idx < encoded_result.tokens.size(); ++idx) {
@@ -227,12 +234,11 @@ class ov::genai::VLMPipeline::VLMPipelineImpl {
         }
         auto decode_end_time = std::chrono::steady_clock::now();
 
-        m_inputs_embedder->update_tokenized_history(encoded_result.tokens[0], finish_info.probably_disappeared_token, generation_config.is_beam_search(),
-                                                    m_language.get_tensor("attention_mask").get_shape()[1] - (history_size + inputs_embeds_size));
+        m_inputs_embedder->update_tokenized_history(finish_info, generation_config.is_beam_search(), m_language.get_tensor("attention_mask").get_shape()[1] - (history_size + inputs_embeds_size), inputs_embeds_size);
 
         std::string decoded_results = decoded.texts.at(0);
         if (m_is_chat_conversation)
-            m_inputs_embedder->update_chat_history(decoded_results);
+            m_inputs_embedder->update_chat_history(decoded_results, finish_info.streaming_finish_status);
 
         auto generate_end_time = std::chrono::steady_clock::now();
         decoded.perf_metrics = encoded_result.perf_metrics;

tests/python_tests/test_vlm_pipeline.py

@@ -6,7 +6,7 @@
 import pytest
 import transformers
 from optimum.intel.openvino import OVModelForVisualCausalLM
-from openvino_genai import VLMPipeline, GenerationConfig
+from openvino_genai import VLMPipeline, GenerationConfig, StreamingStatus
 
 from utils.generation_config import get_beam_search, get_multinomial_all_parameters
 from utils.constants import get_default_llm_properties
@@ -184,3 +184,38 @@ def test_perf_metrics(cache):
     mean_dur, std_dur = perf_metrics.get_prepare_embeddings_duration()
     assert np.allclose(mean_dur, np.mean(raw_dur))
     assert np.allclose(std_dur, np.std(raw_dur))
+
+
+@pytest.mark.precommit
+@pytest.mark.nightly
+@pytest.mark.parametrize("model_id", model_ids)
+def test_vlm_pipeline_chat(model_id, cache):
+    callback_questions = [
+        '1+1=',
+        'Why is the Sun yellow?',
+        'What is the previous answer?'
+    ]
+
+    current_iter = 0
+    num_iters = 3
+    def streamer(subword):
+        nonlocal current_iter
+        current_iter += 1
+        return StreamingStatus.CANCEL if current_iter == num_iters else StreamingStatus.RUNNING
+    
+
+    models_path = get_ov_model(model_id, cache)
+    ov_pipe = VLMPipeline(models_path, "CPU")
+    generation_config = ov_pipe.get_generation_config()
+    generation_config.max_new_tokens = 30
+    generation_config.set_eos_token_id(ov_pipe.get_tokenizer().get_eos_token_id())
+
+    images = []
+    for link in image_links_for_testing[1]:
+        images.append(get_image_by_link(link))
+
+    ov_pipe.start_chat()
+    ov_pipe.generate(callback_questions[0], images=images, generation_config=generation_config)
+    ov_pipe.generate(callback_questions[1], generation_config=generation_config, streamer=streamer)
+    ov_pipe.generate(callback_questions[2], generation_config=generation_config)
+    ov_pipe.finish_chat()