[Speculative decoding] Fix draft_model tun in case of long prompt (openvinotoolkit#1114)

Ticket:
* CVS-156390

Details:
in case of big prompt draft_model request was marked as paused, so not
scheduled

Author: iefode

src/cpp/src/sampler.cpp

@@ -586,7 +586,7 @@ void register_new_token(const Token& sampled_token_id,
         running_sequence->append_token(sampled_token_id.m_index, sampled_token_id.m_log_prob);
     }
     if (!is_validation_mode_enabled &&
-        std::fabs(sampled_token_id.m_log_prob) < logit_processor.get_assistant_confidence_threshold()) { 
+        std::fabs(std::exp(sampled_token_id.m_log_prob)) < logit_processor.get_assistant_confidence_threshold()) {
         auto sequence_group = running_sequence->get_sequence_group_ptr();
         sequence_group->pause_generation(true);
     }

src/cpp/src/speculative_decoding/continuous_batching_for_speculative_decoding_impl.cpp

@@ -254,6 +254,10 @@ ContinuousBatchingPipeline::ContinuousBatchingForSpeculativeDecodingImpl::update
                      prompt_len = request->get_prompt_len(),
                      updated_context_len = min_candidate_len + prompt_len,
                      max_new_tokens = request->get_sampling_parameters().max_new_tokens;
+        // prompt phase
+        if (request->get_context_len() < request->get_prompt_len() && result.inserted_tokens_cnt == 0) {
+            return result;
+        }
         size_t generated_len = request->get_context_len() - request->get_prompt_len();
         if (num_processed_tokens > 0) {
             request->update_processed_tokens_num(num_processed_tokens - result.removed_tokens_cnt);

src/cpp/src/speculative_decoding/speculative_decoding_impl.cpp

@@ -48,7 +48,7 @@ ContinuousBatchingPipeline::SpeculativeDecodingImpl::SpeculativeDecodingImpl(
         auto k = static_cast<float>(draft_model_cache_size) / (main_model_cache_size + draft_model_cache_size);
 
         size_t main_cache_size = main_scheduler_config.cache_size * (1 - k),
-               draft_cache_size = main_scheduler_config.cache_size * k;
+               draft_cache_size = main_scheduler_config.cache_size - main_cache_size;
         if (draft_cache_size == 0) {
             main_cache_size -= main_cache_size > 1 ? 1 : 0;
             draft_cache_size = 1;
@@ -158,6 +158,10 @@ void ContinuousBatchingPipeline::SpeculativeDecodingImpl::step() {
             m_draft_generations.erase(request_id);
         }
         auto updated_seq_info = update_sequence_info[request_id];
+        // several prompt phase
+        if (updated_seq_info.inserted_tokens_cnt == 0) {
+            continue;
+        }
         float acceptance_rate = 1 - static_cast<float>(updated_seq_info.removed_tokens_cnt) / updated_seq_info.inserted_tokens_cnt;
         m_sd_metrics.update_acceptance_rate(request_id, acceptance_rate * 100);
         m_sd_metrics.update_draft_accepted_tokens(request_id, (updated_seq_info.inserted_tokens_cnt - updated_seq_info.removed_tokens_cnt));
@@ -203,6 +207,10 @@ ContinuousBatchingPipeline::SpeculativeDecodingImpl::generate(const std::vector<
     while (has_non_finished_requests() && continue_generation) {
         step();
         if (streamer_ptr) {
+            // not generated tokens like several prompt phase
+            if (!main_generations.at(0).get()->can_read()) {
+                continue;
+            }
             std::unordered_map<uint64_t, GenerationOutput> token = main_generations.at(0).get()->back();
             OPENVINO_ASSERT(1 <= token.size());
             OPENVINO_ASSERT(1 <= token.begin()->second.generated_ids.size());