More fixes for accuracy

Author: AsyaPronina

src/cpp/src/continuous_batching/pipeline.cpp

@@ -19,16 +19,6 @@
 using namespace ov::genai;
 
 namespace {
-ov::genai::ModelDesc
-extract_draft_model_from_config(ov::AnyMap& config) {
-    ov::genai::ModelDesc draft_model;
-    if (config.find(utils::DRAFT_MODEL_ARG_NAME) != config.end()) {
-        draft_model = config.at(utils::DRAFT_MODEL_ARG_NAME).as<ov::genai::ModelDesc>();
-        config.erase(utils::DRAFT_MODEL_ARG_NAME);
-    }
-    return draft_model;
-}
-
 bool
 extract_prompt_lookup_from_config(ov::AnyMap& config) {
     bool res = false;

src/cpp/src/llm/pipeline.cpp

@@ -80,12 +80,12 @@ ov::genai::LLMPipeline::LLMPipeline(
     auto start_time = std::chrono::steady_clock::now();
     auto [properties, attention_backend] = utils::extract_attention_backend(user_properties);
 
-    // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
-    if (utils::explicitly_requires_paged_attention(user_properties)) {
+    if (is_npu_requested(device, properties)) {
+        m_pimpl = std::make_unique<StatefulLLMPipelineNPU>(models_path, tokenizer, properties);
+    } else if (utils::explicitly_requires_paged_attention(user_properties)) {
+        // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
         auto [device_properties, scheduler_config] = utils::extract_scheduler_config(properties, utils::get_latency_oriented_scheduler_config());
         m_pimpl = std::make_unique<ContinuousBatchingAdapter>(models_path, tokenizer, scheduler_config, device, device_properties);
-    } else if (device == "NPU") {
-        m_pimpl = std::make_unique<StatefulLLMPipelineNPU>(models_path, tokenizer, properties);
     } else if (attention_backend == PA_BACKEND) {
         // try to call CB adapter one more time, but with safe guard to silent exception
         try {
@@ -115,13 +115,10 @@ ov::genai::LLMPipeline::LLMPipeline(
 
     auto [properties, attention_backend] = utils::extract_attention_backend(user_properties);
 
-    // First -> check draft model. for NPU leave it as is for the main model.
-    // if NPU
-    // if draft model is on NPU
-    // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
-    if (device == "NPU") {
+    if (is_npu_requested(device, properties)) {
         m_pimpl = std::make_unique<StatefulLLMPipelineNPU>(models_path, properties);
     } else if (utils::explicitly_requires_paged_attention(user_properties)) {
+        // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
         auto [device_properties, scheduler_config] = utils::extract_scheduler_config(properties, utils::get_latency_oriented_scheduler_config());
         m_pimpl = std::make_unique<ContinuousBatchingAdapter>(models_path, scheduler_config, device, device_properties);
 
@@ -157,17 +154,17 @@ ov::genai::LLMPipeline::LLMPipeline(
 
     auto [properties, attention_backend] = utils::extract_attention_backend(user_properties);
 
-    // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
-    if (utils::explicitly_requires_paged_attention(user_properties)) {
-        auto [device_properties, scheduler_config] = utils::extract_scheduler_config(properties, utils::get_latency_oriented_scheduler_config());
-        m_pimpl = std::make_unique<ContinuousBatchingAdapter>(model_str, weights_tensor,
-                                                              tokenizer, scheduler_config, device, device_properties, generation_config);
-    } else if (device == "NPU") {
+    if (is_npu_requested(device, properties)) {
         m_pimpl = std::make_unique<StatefulLLMPipelineNPU>(
                 utils::singleton_core().read_model(model_str, weights_tensor),
                 tokenizer,
                 properties,
                 generation_config);
+    } else if (utils::explicitly_requires_paged_attention(user_properties)) {
+        // If CB is invoked explicitly, create CB adapter as is and re-throw in case if internal issues
+        auto [device_properties, scheduler_config] = utils::extract_scheduler_config(properties, utils::get_latency_oriented_scheduler_config());
+        m_pimpl = std::make_unique<ContinuousBatchingAdapter>(model_str, weights_tensor,
+                                                              tokenizer, scheduler_config, device, device_properties, generation_config);
     } else if (attention_backend == PA_BACKEND) {
         // try to call CB adapter one more time, but with safe guard to silent exception
         try {

src/cpp/src/llm/pipeline_stateful_npu.cpp

@@ -14,18 +14,6 @@
 #include "openvino/core/parallel.hpp"
 #include "openvino/genai/text_streamer.hpp"
 
-namespace {
-    ov::genai::ModelDesc
-    extract_draft_model_from_config(ov::AnyMap& config) {
-        ov::genai::ModelDesc draft_model;
-        if (config.find(ov::genai::utils::DRAFT_MODEL_ARG_NAME) != config.end()) {
-            draft_model = config.at(ov::genai::utils::DRAFT_MODEL_ARG_NAME).as<ov::genai::ModelDesc>();
-            config.erase(ov::genai::utils::DRAFT_MODEL_ARG_NAME);
-        }
-        return draft_model;
-}
-} // anonymous namespace
-
 namespace ov::genai {
 
 // NB: No constructor for creation of pipeline from infer request, as pipeline from infer request

src/cpp/src/speculative_decoding/speculative_decoding_npu.cpp

@@ -509,6 +509,9 @@ EncodedResults SpeculativeLLMPipelineNPU::generate(
     */
     // Last generated token by draft model needs to be prepended before next run if it is accepted by the main model!
     // So it will get into context too.
+    // Remove debug lines.
+    // std::cout << std::endl << "Launching spec decode for " << config.get_max_new_tokens(prompt_len) << " max new tokens." << std::endl << std::endl;
+    // std::vector<std::pair<int,int>> accepted_tokens;
     int64_t draft_prefix_token = -1;
     while (m_main_request->can_infer() && (streaming_status == ov::genai::StreamingStatus::RUNNING)) {
         // Phase 1: Generation of candidates with the draft model:
@@ -534,7 +537,7 @@ EncodedResults SpeculativeLLMPipelineNPU::generate(
             candidate = m_draft_request->infer_next(candidate);
             candidates.push_back(candidate);
         }
-        
+
         // Phase 2. Main inference.
         // For the main network, candidates_size + 1 tokens will be fed at once in a single infer request:
         // last token from previous main inference + all candidates from the draft stage
@@ -548,22 +551,23 @@ EncodedResults SpeculativeLLMPipelineNPU::generate(
         // that is generated based on subsequence [first token,...,`t`]
         // of the input prompt.
         // TODO: Handle OOM exception for static model here.
-        auto ref_out_tokens = m_main_request->infer_next_return_all(input_for_main);
+        auto ref_tokens = m_main_request->infer_next_return_all(input_for_main);
 
         // Phase 3. Check if main model produced the same tokens as input candidates:
         size_t accepted_tokens_number = 0u;
         // Last token is a new token from the main model, skip it:
-        for (size_t i = 0; i < ref_out_tokens.size() - 1; ++i) {
-            if (ref_out_tokens[i] != candidates[i]) {
+        for (size_t i = 0; i < ref_tokens.size() - 1; ++i) {
+            if (ref_tokens[i] != candidates[i]) {
                 break;
             }
             accepted_tokens_number++;
         }
 
+        // FIXME: Remove debug line
+        // accepted_tokens.push_back({accepted_tokens_number, candidates.size()});
         auto mismatched_candidates = candidates.size() - accepted_tokens_number;
-        std::vector<int64_t> validated_tokens(candidates.begin(), candidates.end() - mismatched_candidates);
-        out_token = ref_out_tokens.back();
-        validated_tokens.push_back(out_token);
+        std::vector<int64_t> validated_tokens(ref_tokens.begin(), ref_tokens.end() - mismatched_candidates);
+        out_token = validated_tokens.back();
 
         // Phase 4: Update inference wrappers based on found matches and mismatches
         // This is the case when main model accepted all candidates from draft model
@@ -573,6 +577,7 @@ EncodedResults SpeculativeLLMPipelineNPU::generate(
         } else {
             m_draft_request->trimm_kv_cache(mismatched_candidates - 1);
             m_main_request->trimm_kv_cache(mismatched_candidates);
+            draft_prefix_token = -1;
         }
 
         m_speculative_config.update_candidate_strategy(accepted_tokens_number);
@@ -587,6 +592,14 @@ EncodedResults SpeculativeLLMPipelineNPU::generate(
         streamer_ptr->end();
     }
 
+    // Remove debug lines
+    // std::cout << std::endl << std::endl << "Acceptance ratios for each iteration from total of " << accepted_tokens.size() << "." << std::endl;
+    // std::cout << "Format: n/m per iteration, `n` accepted tokens from `m` candidates." << std::endl;
+    // for (int i = 0; i < accepted_tokens.size(); ++i) {
+    //     std::cout << accepted_tokens[i].first << "/" << accepted_tokens[i].second << ", ";
+    // }
+    m_speculative_config.num_pred_tokens = 5;
+
     m_draft_request->reset_state();
     m_main_request->reset_state();
 

src/cpp/src/utils.cpp

@@ -101,7 +101,6 @@ inline bool is_paged_attention_available() {
     return false;
 #endif
 }
-
 } // anonymous
 
 namespace ov {
@@ -202,6 +201,28 @@ ProcessorConfig from_any_map(
     return extracted_config;
 }
 
+ov::genai::ModelDesc extract_draft_model_from_config(ov::AnyMap& config) {
+    ov::genai::ModelDesc draft_model;
+    if (config.find(utils::DRAFT_MODEL_ARG_NAME) != config.end()) {
+        draft_model = config.at(utils::DRAFT_MODEL_ARG_NAME).as<ov::genai::ModelDesc>();
+        config.erase(utils::DRAFT_MODEL_ARG_NAME);
+    }
+    return draft_model;
+}
+
+bool is_npu_requested(const std::string& device, const ov::AnyMap& properties) {
+    if (device == "NPU") {
+        return true;
+    }
+
+    auto draft_model_descr = extract_draft_model_from_config(properties);
+    if (draft_model_descr.model != nullptr) {
+        return draft_model_descr.device == "NPU";
+    }
+
+    return false;
+}
+
 ov::genai::TokenizedInputs subtract_chat_tokenized_inputs(const ov::genai::TokenizedInputs& minuend, const ov::genai::TokenizedInputs& subtrahend) {
     auto minuend_size = minuend.input_ids.get_size();
     auto subtrahend_size = subtrahend.input_ids.get_size();

src/cpp/src/utils.hpp

@@ -118,6 +118,10 @@ ProcessorConfig from_any_map(
     const ProcessorConfig& initial
 );
 
+ov::genai::ModelDesc extract_draft_model_from_config(ov::AnyMap& config);
+
+bool is_npu_requested(const std::string& device, const ov::AnyMap& properties)
+
 ov::genai::TokenizedInputs subtract_chat_tokenized_inputs(const ov::genai::TokenizedInputs& minuend, const ov::genai::TokenizedInputs& subtrahend);
 
 void apply_slice_before_matmul_transformation(std::shared_ptr<ov::Model> model);