update according to m_kv_cache_state

Author: sbalandi

src/cpp/src/llm_pipeline_stateful.cpp

@@ -60,7 +60,7 @@ StatefulLLMPipeline::StatefulLLMPipeline(
     }
 
     if (!m_use_full_chat_history)
-        m_kv_history_trim_manager.kv_cache_seq_length_axis = kv_pos.seq_len;
+        m_kv_cache_state.seq_length_axis = kv_pos.seq_len;
 
     auto filtered_properties = extract_adapters_from_properties(properties, &m_generation_config.adapters);
     if (m_generation_config.adapters) {
@@ -143,7 +143,7 @@ DecodedResults StatefulLLMPipeline::generate(
             if (m_use_full_chat_history) {
                 encoded_input = new_chat_tokens;
             } else {
-                ov::genai::align_kv_cache_and_history(m_kv_history_trim_manager, new_chat_tokens.input_ids, m_kv_cache_state);
+                ov::genai::align_kv_cache_and_history(new_chat_tokens.input_ids, m_kv_cache_state);
                 encoded_input = get_chat_encoded_input(new_chat_tokens.input_ids, m_kv_cache_state);
             }
             // TODO: Forbid LoRA config change if we are in the chat mode, because it requires regenerating the history with LoRA applied
@@ -232,7 +232,7 @@ EncodedResults StatefulLLMPipeline::generate(
     // Tail of previous output in chat mode is missing in KV cache.
     if (is_chat_conversation && m_chat_input_type == ov::genai::utils::GenerationChatInputsType::ENCODED_INPUTS) {
         ov::Tensor new_chat_tokens = ov::Tensor{ov::element::i64, {1, m_tokenized_chat_history.size()}, m_tokenized_chat_history.data()};
-        ov::genai::align_kv_cache_and_history(m_kv_history_trim_manager, new_chat_tokens, m_kv_cache_state);
+        ov::genai::align_kv_cache_and_history(new_chat_tokens, m_kv_cache_state);
 
         auto encoded_input = get_chat_encoded_input(new_chat_tokens, m_kv_cache_state);
         input_ids = encoded_input.input_ids;
@@ -278,8 +278,8 @@ EncodedResults StatefulLLMPipeline::generate(
         if (m_kv_cache_state.get_state().empty() || m_use_full_chat_history)
             reset_kv_state();
         else
-            ov::genai::utils::trim_kv_cache(m_model_runner, m_kv_history_trim_manager.num_tokens_to_trim,
-                                            m_kv_history_trim_manager.kv_cache_seq_length_axis, m_adapter_controller);
+            ov::genai::utils::trim_kv_cache(m_model_runner, m_kv_cache_state.num_tokens_to_trim,
+                                            m_kv_cache_state.seq_length_axis, m_adapter_controller);
     }
 
     size_t kv_cache_len = 0;
@@ -352,7 +352,7 @@ EncodedResults StatefulLLMPipeline::generate(
     m_chat_generation_finish_status = finish_info.streaming_finish_status;
 
     if (is_chat_conversation) {
-        m_kv_history_trim_manager.reset();
+        m_kv_cache_state.num_tokens_to_trim = 0;
 
         if (m_chat_input_type == ov::genai::utils::GenerationChatInputsType::ENCODED_INPUTS) {
             if (m_chat_generation_finish_status == ov::genai::GenerationStatus::CANCEL) {
@@ -361,7 +361,7 @@ EncodedResults StatefulLLMPipeline::generate(
                 std::copy(result.tokens[0].begin(), result.tokens[0].end(), std::back_inserter(m_tokenized_chat_history));
             }
         } else if (config.is_beam_search()) {
-            m_kv_history_trim_manager.num_tokens_to_trim = m_model_runner.get_tensor("attention_mask").get_shape()[1] - prev_attn_mask_size;
+            m_kv_cache_state.num_tokens_to_trim = m_model_runner.get_tensor("attention_mask").get_shape()[1] - prev_attn_mask_size;
         }
     } else {
         m_kv_cache_state.reset_state();
@@ -402,7 +402,6 @@ void StatefulLLMPipeline::reset_kv_state() {
 
 void StatefulLLMPipeline::finish_chat() {
     is_chat_conversation = false;
-    m_kv_history_trim_manager.reset();
     m_chat_input_type = ov::genai::utils::GenerationChatInputsType::UNDEF;
     bool have_state = 0 != m_model_runner.get_tensor("attention_mask").get_size();
     if (!m_kv_cache_state.get_state().empty() || have_state) {

src/cpp/src/llm_pipeline_stateful.hpp

@@ -20,17 +20,13 @@ class StatefulLLMPipeline final : public LLMPipelineImplBase {
     ChatHistory m_history;
     std::vector<int64_t> m_tokenized_chat_history;
     ov::genai::utils::GenerationChatInputsType m_chat_input_type = ov::genai::utils::GenerationChatInputsType::UNDEF;
-    // If sequence contains some symbols, which could be ambiguously encoded by tokenizer, we need to trim kv cache
-    // If we use beam search sampling with chat mode we need to remove last answer of the model from kv cache and add best answer to history 
-    // so, let's keep info about amount of tokens to trim from kv cache and amount of tokens to keep in history
-    ov::genai::KVCacheTrimManager m_kv_history_trim_manager = {0, 2};
     // Finish reason of last generation for chat scenario
     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
+    // include reflection of tokens contained in the kv cache and amount of tokens, which are needed to trim from kv cache on the next step of chat
     KVCacheState m_kv_cache_state;
 
     void reset_kv_state();

src/cpp/src/lm_encoding.cpp

@@ -325,7 +325,7 @@ TokenizedInputs get_chat_encoded_input(const ov::Tensor& new_chat_tokens, KVCach
 }
 
 
-void align_kv_cache_and_history(ov::genai::KVCacheTrimManager& kv_history_manager, const ov::Tensor& new_chat_tokens, KVCacheState& kv_cache_state) {
+void align_kv_cache_and_history(const ov::Tensor& new_chat_tokens, KVCacheState& kv_cache_state) {
     // KV cache in model already contains prompts and answers from previous iterations.
     // So only new prompt wrapped into chat template to be sent into model. Tokenizer always returns
     // token_ids = {<bos token>, ...<valuable tokens>}. So if tokenizer applies only to the new prompt,
@@ -343,7 +343,7 @@ void align_kv_cache_and_history(ov::genai::KVCacheTrimManager& kv_history_manage
     size_t first_diverse_tokens_idx = ov::genai::utils::get_first_history_difference(new_chat_tokens, state);
     // in the case of beam_search the longest answer is in the kv cache, but the best one is needed
     // so generated tokens were not added to KVCacheState and num_tokens_to_trim was set to the size of the generated serquence
-    kv_history_manager.num_tokens_to_trim = kv_history_manager.num_tokens_to_trim > 0 ? kv_history_manager.num_tokens_to_trim : (state.size() - first_diverse_tokens_idx);
+    kv_cache_state.num_tokens_to_trim = kv_cache_state.num_tokens_to_trim > 0 ? kv_cache_state.num_tokens_to_trim : (state.size() - first_diverse_tokens_idx);
     state.resize(first_diverse_tokens_idx);
 }
 

src/cpp/src/lm_encoding.hpp

@@ -11,6 +11,9 @@ namespace genai {
 class KVCacheState {
     std::vector<int64_t> state;
 public:
+    size_t num_tokens_to_trim = 0;
+    size_t seq_length_axis = 2;
+
     std::vector<int64_t>& get_state() {
         return state;
     }
@@ -20,18 +23,8 @@ class KVCacheState {
     }
 
     void reset_state() {
-        return state.clear();
-    }
-};
-
-
-struct KVCacheTrimManager
-{
-    size_t num_tokens_to_trim = 0;
-    size_t kv_cache_seq_length_axis = 2;
-
-    void reset() {
         num_tokens_to_trim = 0;
+        state.clear();
     }
 };
 
@@ -42,7 +35,7 @@ ov::genai::utils::GenerationFinishInfo get_lm_encoded_results(ov::InferRequest&
                                                               std::optional<int64_t> rope_delta = std::nullopt, const size_t max_kv_cache_size = std::numeric_limits<size_t>::max());
 
 
-void align_kv_cache_and_history(ov::genai::KVCacheTrimManager& kv_history_manager, const ov::Tensor& new_chat_tokens, KVCacheState& kv_cache_state);
+void align_kv_cache_and_history(const ov::Tensor& new_chat_tokens, KVCacheState& kv_cache_state);
 
 
 TokenizedInputs get_chat_encoded_input(const ov::Tensor& new_chat_tokens, KVCacheState& kv_cache_state);

src/cpp/src/visual_language/inputs_embedder.cpp

@@ -217,10 +217,6 @@ ov::Tensor InputsEmbedder::get_inputs_embeds(const std::string& prompt, const st
     return m_impl->get_inputs_embeds(prompt, images, metrics);
 }
 
-ov::Tensor InputsEmbedder::get_input_embeddings(const std::string& prompt, const std::vector<ov::Tensor>& images, ov::genai::VLMPerfMetrics& metrics) {
-    return m_impl->get_input_embeddings(prompt, images, metrics);
-}
-
 std::pair<ov::Tensor, std::optional<int64_t>> InputsEmbedder::get_position_ids(const size_t inputs_embeds_size, const size_t history_size) {
     return m_impl->get_position_ids(inputs_embeds_size, history_size);
 }
@@ -233,10 +229,6 @@ KVCacheState& InputsEmbedder::get_kv_cache_state() {
     return  m_impl->get_kv_cache_state();
 }
 
-bool InputsEmbedder::should_reset_kv_cache() const {
-    return m_impl->should_reset_kv_cache();
-}
-
 Tokenizer InputsEmbedder::get_tokenizer() const {
     return m_impl->get_tokenizer();
 }

src/cpp/src/visual_language/inputs_embedder.hpp

@@ -35,9 +35,6 @@ class InputsEmbedder {
     // compute input embedding for prompt and multiple images
     ov::Tensor get_inputs_embeds(const std::string& prompt, const std::vector<ov::Tensor>& images, ov::genai::VLMPerfMetrics& metrics);
 
-    // computes input embedding for prompt and multiple images and saves input_embeddings size
-    ov::Tensor get_input_embeddings(const std::string& prompt, const std::vector<ov::Tensor>& images, ov::genai::VLMPerfMetrics& metrics);
-
     // compute position ids for language model input
     std::pair<ov::Tensor, std::optional<int64_t>> get_position_ids(const size_t inputs_embeds_size, const size_t history_size);
 
@@ -50,9 +47,6 @@ class InputsEmbedder {
     // get reflection of tokens contained in the kv cache
     KVCacheState& get_kv_cache_state();
 
-    // 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);
 

src/cpp/src/visual_language/pipeline.cpp

@@ -42,8 +42,6 @@ class VLMPipeline::VLMPipelineImpl : public VLMPipelineBase{
     std::shared_ptr<InputsEmbedder> m_inputs_embedder;
     // Axis num in kv cache from m_language model, which contains information about history len
     size_t m_kv_cache_seq_length_axis = 2;
-    // Load pipeline time
-    float m_load_time_ms = 0;
     // Component for applying sampling to lm outputs
     Sampler m_sampler;
 public:
@@ -143,7 +141,7 @@ class VLMPipeline::VLMPipelineImpl : public VLMPipelineBase{
         m_inputs_embedder->set_apply_chat_template_status(generation_config.apply_chat_template);
 
         auto start_get_inputs_embeds = std::chrono::steady_clock::now();
-        ov::Tensor inputs_embeds = m_inputs_embedder->get_input_embeddings(prompt, rgbs, perf_metrics);
+        ov::Tensor inputs_embeds = m_inputs_embedder->get_inputs_embeds(prompt, rgbs, perf_metrics);
         auto end_get_inputs_embeds = std::chrono::steady_clock::now();
 
         KVCacheState& kv_cache_state = m_inputs_embedder->get_kv_cache_state();