Address review comments

Author: michalkulakowski

src/cpp/src/continuous_batching_impl.cpp

@@ -522,7 +522,7 @@ void ContinuousBatchingPipeline::ContinuousBatchingImpl::_fill_prompt_log_probs(
         }
         currently_processed_tokens += output_seq_len * num_running_sequences;
         // For max_new_tokens == 0, we don't reach sampling so need to notify handle separately
-        if(sequence_group->get_sampling_parameters().max_new_tokens == 0) {
+        if(sequence_group->get_sampling_parameters().get_max_new_tokens(sequence_group->get_prompt_len()) == 0) {
             sequence_group->notify_handle_echo_only();
         }
     }

src/cpp/src/prompt_lookup/continuous_batching_for_prompt_lookup.cpp

@@ -67,7 +67,7 @@ void ContinuousBatchingPipeline::ContinuousBatchingForPromptLookupImpl::generate
             const auto sampling_params = request->get_sampling_parameters();
             {
                 const auto generated_len = running_sequence->get_generated_len();
-                const auto left_generated_len = std::min(sampling_params.max_new_tokens, sampling_params.max_length) - generated_len - 1;
+                const auto left_generated_len = std::min(sampling_params.get_max_new_tokens(request->get_prompt_len()), sampling_params.max_length) - generated_len - 1;
                 min_num_assistant_tokens = std::min(sampling_params.num_assistant_tokens, left_generated_len);
             }
             TokenIds candidates = generate_candidates(full_input_ids, min_num_assistant_tokens, sampling_params.max_ngram_size);

src/cpp/src/sampler.cpp

@@ -215,6 +215,7 @@ Sampler::GroupBeamSearcher::GroupBeamSearcher(SequenceGroup::Ptr sequence_group,
         // to avoid selecting the same tokens for beams within group, let's just initialize score
         // for the front one
         group.ongoing.front().m_score = 0.0f;
+        group.prompt_len = this->m_sequence_group->get_prompt_len();
     }
 }
 
@@ -408,7 +409,7 @@ void Sampler::GroupBeamSearcher::select_next_tokens(const ov::Tensor& logits,
         }
 
         // check whether group has finished
-        group.is_done(m_parameters, this->m_sequence_group->get_prompt_len());
+        group.is_done(m_parameters);
 
         // group cannot continue if there are no valid child beams
         if (child_beams_per_group[group_id].size() == 0) {
@@ -956,7 +957,7 @@ int64_t Sampler::GroupBeamSearcher::Group::finish(Beam beam, const ov::genai::Ge
     return preeempted_sequence_id;
 }
 
-void Sampler::GroupBeamSearcher::Group::is_done(const ov::genai::GenerationConfig& sampling_params, size_t prompt_len) {
+void Sampler::GroupBeamSearcher::Group::is_done(const ov::genai::GenerationConfig& sampling_params) {
     assert(sampling_params.num_beams % sampling_params.num_beam_groups == 0 &&
         "number of beams should be divisible by number of groups");
     size_t group_size = sampling_params.num_beams / sampling_params.num_beam_groups;

src/cpp/src/sampler.hpp

@@ -111,10 +111,11 @@ class Sampler::GroupBeamSearcher {
     struct Group {
         std::vector<Beam> ongoing;  // Best beams in front
         std::vector<Beam> min_heap;  // The worst of the best completed beams is the first
+        size_t prompt_len;
         bool done = false;
 
         int64_t finish(Beam beam, const ov::genai::GenerationConfig& sampling_params);
-        void is_done(const ov::genai::GenerationConfig& sampling_params, size_t prompt_len);
+        void is_done(const ov::genai::GenerationConfig& sampling_params);
     };
 
     SequenceGroup::Ptr m_sequence_group;

src/cpp/src/sequence_group.hpp

@@ -457,7 +457,7 @@ class SequenceGroup  : public std::enable_shared_from_this<SequenceGroup> {
     }
 
     bool requires_sampling() const {
-        return get_context_len() >= get_prompt_len() && get_context_len() > m_max_content_len && m_sampling_params.max_new_tokens > 0;
+        return get_context_len() >= get_prompt_len() && get_context_len() > m_max_content_len && m_sampling_params.get_max_new_tokens(this->get_prompt_len()) > 0;
     }
 
     void schedule_tokens(size_t num_tokens) {

src/cpp/src/speculative_decoding/continuous_batching_for_speculative_decoding_impl.cpp

@@ -261,7 +261,7 @@ ContinuousBatchingPipeline::ContinuousBatchingForSpeculativeDecodingImpl::update
         const size_t num_processed_tokens = request->get_num_processed_tokens(),
                      prompt_len = request->get_prompt_len(),
                      updated_context_len = min_candidate_len + prompt_len,
-                     max_new_tokens = request->get_sampling_parameters().max_new_tokens;
+                     max_new_tokens = request->get_sampling_parameters().get_max_new_tokens(request->get_prompt_len());
         size_t generated_len = request->get_context_len() >= request->get_prompt_len() ? request->get_context_len() - request->get_prompt_len() + 1 : 0;
         if (generated_len > 0 && result.removed_tokens_cnt > 0) {
             request->update_processed_tokens_num(num_processed_tokens - result.removed_tokens_cnt + 1);
@@ -324,13 +324,13 @@ void ContinuousBatchingPipeline::ContinuousBatchingForSpeculativeDecodingImpl::m
                 // generate only one token in case of non speculative decoding
                 request->pause_generation(true);
             } else if (request->get_num_processed_tokens() >= request->get_prompt_len() &&
-                (request->get_num_processed_tokens() - request->get_prompt_len() + 1) >= sampling_params.max_new_tokens - 1) {
+                (request->get_num_processed_tokens() - request->get_prompt_len() + 1) >= sampling_params.get_max_new_tokens(request->get_prompt_len()) - 1) {
                 request->pause_generation(true);
             } else if (request->get_num_processed_tokens() == 0 && sampling_params.num_return_sequences > 1) {
                 request->pause_generation(true);
             } else if (sampling_params.num_assistant_tokens <= generated_tokens_cnt && sampling_params.assistant_confidence_threshold == 0.f) {
                 request->pause_generation(true);
-            } else if (sampling_params.max_new_tokens == 0) {
+            } else if (sampling_params.get_max_new_tokens(request->get_prompt_len()) == 0) {
                 request->pause_generation(true);
             } else if (request->get_num_processed_tokens() == request->get_prompt_len()) {
                 request->pause_generation(true);