Fix for CI

Author: iefode

src/cpp/src/logit_processor.hpp

@@ -312,7 +312,6 @@ class LogitProcessor {
 
     // speculative decoding parameters
     float m_assistant_confidence_threshold = 0.f;
-    bool m_is_dynamic_speculative_decoding = false;
 
 
 public:
@@ -360,16 +359,11 @@ class LogitProcessor {
                 }
             }
             if (sampling_params.assistant_confidence_threshold > 0) {
-                m_is_dynamic_speculative_decoding = true;
                 m_assistant_confidence_threshold = sampling_params.assistant_confidence_threshold;
             }
         }
     }
 
-    bool is_dynamic_speculative_decoding() {
-        return m_is_dynamic_speculative_decoding;
-    }
-
     float get_assistant_confidence_threshold() {
         return m_assistant_confidence_threshold;
     }

src/cpp/src/sampler.cpp

@@ -565,23 +565,12 @@ void register_new_token(const Token& sampled_token_id,
                         Sequence::Ptr running_sequence,
                         LogitProcessor& logit_processor,
                         bool is_extend_sequence,
-                        bool is_update_len_logit_processor,
                         bool is_validation_mode_enabled) {
     logit_processor.register_new_generated_token(sampled_token_id.m_index);
-    size_t generated_len = logit_processor.get_generated_len();
     if (is_extend_sequence) {
         running_sequence->append_token(sampled_token_id.m_index, sampled_token_id.m_log_prob);
-    } else {
-        // just update the token log prob in case of successfully validated token
-        OPENVINO_ASSERT(generated_len < running_sequence->get_generated_len());
-        running_sequence->update_generated_log_prob(generated_len, sampled_token_id.m_log_prob);
-    }
-    // increment seq len only for one sequence in sequence group to sync them
-    if (is_update_len_logit_processor) {
-        logit_processor.update_generated_len(++generated_len);
     }
     if (!is_validation_mode_enabled &&
-        logit_processor.is_dynamic_speculative_decoding() &&
         std::fabs(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);
@@ -604,7 +593,7 @@ create_n_forked_sequences(SequenceGroup::Ptr sequence_group,
         const auto forked_sequence = sequence_group->fork_sequence(sequence_to_fork);
         const auto forked_seq_id = forked_sequence->get_id();
         forked_seq_ids.push_back(forked_seq_id);
-        register_new_token(sampled_tokens[i], forked_sequence, logit_processor, true, false, false);
+        register_new_token(sampled_tokens[i], forked_sequence, logit_processor, true, false);
     }
     return forked_seq_ids;
 }
@@ -616,6 +605,8 @@ stop_sample_tokens(Sequence::Ptr running_sequence,
                    size_t& max_removed_tokens_per_request) {
     running_sequence->remove_last_tokens(token_idx);
     max_removed_tokens_per_request = std::max(max_removed_tokens_per_request, token_idx);
+    running_sequence->set_status(SequenceStatus::FINISHED);
+    running_sequence->set_finish_reason(GenerationFinishReason::STOP);
 }
 
 void
@@ -742,13 +733,16 @@ SamplerOutput Sampler::sample(std::vector<SequenceGroup::Ptr> & sequence_groups,
                         }
                         // flag to add sampled token to generated sequence or extend logit processors only
                         bool is_extend_sequence = token_offset == 0 || is_generate_n_tokens,
-                             // flag to update generated length of sequence group in logit processor
-                             is_update_len_logit_processor = running_sequence_id == num_running_sequences - 1,
                              is_validation_passed = true;
                         if (is_validation_mode_enabled && !is_generate_n_tokens) {
                             is_validation_passed = validate_candidate(running_sequences[running_sequence_id], token_offset, sampled_token_id, is_extend_sequence, max_removed_tokens_per_request);
+                            // update log prob just while validation process
+                            if (!is_extend_sequence) {
+                                OPENVINO_ASSERT(generated_and_verified_len < running_sequences[running_sequence_id]->get_generated_len());
+                                running_sequence->update_generated_log_prob(generated_and_verified_len, sampled_token_id.m_log_prob);
+                            }
                         }
-                        register_new_token(sampled_token_id, running_sequences[running_sequence_id], logit_processor, is_extend_sequence, is_update_len_logit_processor, is_validation_mode_enabled);
+                        register_new_token(sampled_token_id, running_sequences[running_sequence_id], logit_processor, is_extend_sequence, is_validation_mode_enabled);
                         // to exit from sampling in case of failed token validation
                         if (!is_validation_passed) {
                             break;
@@ -794,6 +788,7 @@ SamplerOutput Sampler::sample(std::vector<SequenceGroup::Ptr> & sequence_groups,
             align_all_sequence_len(sequence_group, min_generated_len, logit_processor);
             auto min_processed_tokens = sequence_group->get_prompt_len() + min_generated_len - 1;
             sequence_group->update_processed_tokens_num(min_processed_tokens);
+            logit_processor.update_generated_len(min_processed_tokens);
         }
 
         // accumulate a number of processed tokens

src/cpp/src/speculative_decoding/continuous_batching_for_speculative_decoding_impl.cpp

@@ -308,6 +308,8 @@ void ContinuousBatchingPipeline::ContinuousBatchingForSpeculativeDecodingImpl::m
                 request->pause_generation(true);
             } else if (sampling_params.num_assistant_tokens <= generated_tokens_cnt) {
                 request->pause_generation(true);
+            } else if (request->get_num_processed_tokens() - request->get_prompt_len() + 1 >= sampling_params.max_new_tokens - 1) {
+                request->pause_generation(true);
             }
             to_generate |= request->can_generate_tokens();
         }

src/cpp/src/speculative_decoding/speculative_decoding_impl.cpp

@@ -181,8 +181,8 @@ ContinuousBatchingPipeline::SpeculativeDecodingImpl::generate(const std::vector<
 
         auto draft_sampling_params = sampling_params[request_id];
         // set the parameters do not stop draft generation without stopping of the same request for main pipeline
-        draft_sampling_params.max_new_tokens = SIZE_MAX - 1;
-        draft_sampling_params.min_new_tokens = SIZE_MAX - 1;
+        draft_sampling_params.max_new_tokens = draft_sampling_params.max_new_tokens + 1;
+        draft_sampling_params.min_new_tokens = draft_sampling_params.min_new_tokens + 1;
         draft_sampling_params.ignore_eos = true;
         draft_generations.push_back(m_draft_pipeline->add_request(request_id, input_ids[request_id], draft_sampling_params));
         // decrease generation len to generate last token by main model