Align PromptLookupDecoding with greedy when dynamic_split_fuse works (#2360)

[CVS-169291](https://jira.devtools.intel.com/browse/CVS-169291)

when dynamic_split_fuse is true and max_num_batched_tokens less them
prompt len, then PromptLookupDecoding Pipeline starts to generate
candidates before full prompt will be put to cache, which causes to
difference with greedy pipeline, let's avoid that

Author: sbalandi

src/cpp/src/prompt_lookup/continuous_batching_for_prompt_lookup.cpp

@@ -60,6 +60,9 @@ void ContinuousBatchingPipeline::ContinuousBatchingForPromptLookupImpl::generate
         size_t max_validation_len = 0;
         for (auto& running_sequence : request->get_running_sequences()) {
             const auto generated_tokens = running_sequence->get_generated_ids();
+            if (generated_tokens.empty()) {
+                continue;
+            }
             TokenIds full_input_ids = prompt;
             full_input_ids.insert(full_input_ids.end(), generated_tokens.begin(), generated_tokens.end());
 

src/cpp/src/sampling/sampler.cpp

@@ -764,13 +764,14 @@ SequenceGroupSamplingInfo Sampler::sample_from_sequence_group(SequenceGroup::Ptr
     const size_t output_seq_len = sequence_group->get_output_seq_len();
     // get number of tokens to be validated
     size_t num_tokens_to_process = sequence_group->get_num_tokens_to_validate();
+    size_t num_generated_tokens_to_validate = num_tokens_to_process;
 
     if (num_tokens_to_process > output_seq_len - 1) {
         auto delta = num_tokens_to_process - (output_seq_len - 1);
         assisting_pipeline_info.updated_validation_len = std::max(assisting_pipeline_info.updated_validation_len, delta);
         num_tokens_to_process -= delta;
     }
-    
+
     if (sampling_params.is_greedy_decoding() || sampling_params.is_multinomial()) {
         std::vector<Sequence::Ptr> running_sequences = sequence_group->get_running_sequences();
         size_t num_running_sequences = sequence_group->num_running_seqs();
@@ -786,25 +787,25 @@ SequenceGroupSamplingInfo Sampler::sample_from_sequence_group(SequenceGroup::Ptr
                     break;
                 sg_sampling_info.sampler_output.num_generated_tokens++;
                 // calculate token offset from the end of logit
-                size_t token_offset = num_tokens_to_process - i;
+                size_t logit_token_offset = num_tokens_to_process - i;
+                size_t generated_seq_token_offset = num_generated_tokens_to_validate - i;
                 // max counter of needed to be sampled tokens
-                OPENVINO_ASSERT(running_sequence->get_generated_len() >= token_offset);
-                size_t generated_and_verified_len = running_sequence->get_generated_len() - token_offset;
+                OPENVINO_ASSERT(running_sequence->get_generated_len() >= generated_seq_token_offset);
+                size_t generated_and_verified_len = running_sequence->get_generated_len() - generated_seq_token_offset;
                 OPENVINO_ASSERT(sequence_group->get_max_new_tokens() >= generated_and_verified_len);
                 size_t max_num_sampled_token = sequence_group->get_max_new_tokens() - generated_and_verified_len;
                 if (max_num_sampled_token == 0) {
-                    stop_sample_tokens(running_sequence, token_offset, max_num_sampled_token, assisting_pipeline_info.max_removed_tokens_per_request);
+                    stop_sample_tokens(running_sequence, generated_seq_token_offset, max_num_sampled_token, assisting_pipeline_info.max_removed_tokens_per_request);
                     break;
                 }
-                
                 // do sampling only for token validation/generation.
                 // continue in case of extending draft model sequences by main model generated tokens which
                 // should be taken to KV cache without validation
-                if (!is_validation_mode_enabled && token_offset > 0) {
+                if (!is_validation_mode_enabled && generated_seq_token_offset > 0) {
                     continue;
                 }
 
-                auto logit_vector = _get_logit_vector(sequence_group_logits, running_sequence_id, token_offset);
+                auto logit_vector = _get_logit_vector(sequence_group_logits, running_sequence_id, logit_token_offset);
                 logit_processor.apply(logit_vector);
 
                 Token sampled_token;
@@ -826,8 +827,8 @@ SequenceGroupSamplingInfo Sampler::sample_from_sequence_group(SequenceGroup::Ptr
                     sampled_token = sampled_token_ids.front();
                     // make `_speculative_sampling` in case of previous token was not accepted in speculative decoding
                     if (!is_validation_passed) {
-                        float p_prime = get_p_prime(running_sequence, sampled_token, token_offset + 1);
-                        assisting_pipeline_info.max_removed_tokens_per_request = std::max(assisting_pipeline_info.max_removed_tokens_per_request, token_offset);
+                        float p_prime = get_p_prime(running_sequence, sampled_token, generated_seq_token_offset + 1);
+                        assisting_pipeline_info.max_removed_tokens_per_request = std::max(assisting_pipeline_info.max_removed_tokens_per_request, generated_seq_token_offset);
                         // update prob only in case candidate prob > sampled token prob
                         if (p_prime > 0.f) {
                             auto prob = std::exp(sampled_token.m_log_prob);
@@ -837,12 +838,13 @@ SequenceGroupSamplingInfo Sampler::sample_from_sequence_group(SequenceGroup::Ptr
                     }
                 }
                 // flag to add sampled token to generated sequence or extend logit processors only
-                bool is_extend_sequence = token_offset == 0 || is_generate_n_tokens || !is_validation_passed;
+                bool is_extend_sequence = logit_token_offset == 0 || is_generate_n_tokens || !is_validation_passed;
                 if (is_validation_mode_enabled && !is_extend_sequence) {
-                    is_validation_passed = validate_candidate(running_sequences[running_sequence_id], token_offset, sampled_token,
-                                                              is_extend_sequence, assisting_pipeline_info.max_removed_tokens_per_request,
+                    is_validation_passed = validate_candidate(running_sequences[running_sequence_id], generated_seq_token_offset,
+                                                              sampled_token, is_extend_sequence, assisting_pipeline_info.max_removed_tokens_per_request,
                                                               sampling_params.do_sample, !sampling_params.is_prompt_lookup());
-                    // doing resample in case of non accepted tokens in specualtive sampling, if candidates have real logits
+
+                    // doing resample in case of non accepted tokens in speculative sampling
                     if (!is_validation_passed && sampling_params.do_sample && !sampling_params.is_prompt_lookup()) {
                         continue;
                     }
@@ -897,7 +899,7 @@ SequenceGroupSamplingInfo Sampler::sample_from_sequence_group(SequenceGroup::Ptr
     }
     // Notify handle after sampling is done. 
     // For non-streaming this is effective only when the generation is finished.
-    OPENVINO_ASSERT(num_tokens_to_process >= assisting_pipeline_info.max_removed_tokens_per_request);
+    OPENVINO_ASSERT(num_generated_tokens_to_validate >= assisting_pipeline_info.max_removed_tokens_per_request);
     sequence_group->notify_handle();
     return sg_sampling_info;
 }

src/cpp/src/speculative_decoding/speculative_decoding_impl.cpp

@@ -74,6 +74,9 @@ ContinuousBatchingPipeline::SpeculativeDecodingImpl::SpeculativeDecodingImpl(con
 
         main_scheduler_config_updated.cache_size = main_cache_size;
         draft_scheduler_config.cache_size = draft_cache_size;
+    } else {
+        draft_scheduler_config.dynamic_split_fuse = main_scheduler_config_updated.dynamic_split_fuse;
+        draft_scheduler_config.max_num_batched_tokens = main_scheduler_config_updated.max_num_batched_tokens;
     }
 
     ov::AnyMap draft_properties = draft_model_desc.properties.empty() ? main_model_desc.properties : draft_model_desc.properties;

tests/python_tests/test_continuous_batching.py

@@ -4,6 +4,7 @@
 import os
 import pytest
 import math
+import sys
 
 from pathlib import Path
 from shutil import rmtree
@@ -414,6 +415,30 @@ def test_preemption_with_multinomial_n_seq(dynamic_split_fuse):
                          scheduler_config=scheduler_config)
 
 
+@pytest.mark.parametrize("pipeline_type", [PipelineType.PROMPT_LOOKUP_DECODING])
+@pytest.mark.precommit
+def test_dynamic_split_fuse_doesnt_affect_generated_text(pipeline_type):
+    model_id : str = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    _, _, models_path = download_and_convert_model(model_id)
+
+    scheduler_config_ref = dict_to_scheduler_config({"dynamic_split_fuse": False, "max_num_batched_tokens": sys.maxsize})
+    cb_pipe_ref = create_ov_pipeline(models_path, scheduler_config=scheduler_config_ref, pipeline_type=pipeline_type)
+
+    scheduler_config_target = dict_to_scheduler_config({"dynamic_split_fuse": True, "max_num_batched_tokens": 5})
+    cb_pipe_target = create_ov_pipeline(models_path, scheduler_config=scheduler_config_target, pipeline_type=pipeline_type)
+
+    generation_config = GenerationConfig(do_sample=False, max_new_tokens=20, eos_token_id=cb_pipe_ref.get_tokenizer().get_eos_token_id())
+
+    generation_config = prepare_generation_config_by_pipe_type(generation_config=generation_config, pipeline_type=pipeline_type)
+    cb_pipe_ref.set_generation_config(generation_config)
+    cb_pipe_target.set_generation_config(generation_config)
+
+    question = "Why is the Sun yellow?"
+    reference = cb_pipe_ref.generate(question, generation_config=generation_config)
+    generated = cb_pipe_target.generate(question, generation_config=generation_config)
+    assert generated == reference
+
+
 def get_data_by_pipeline_type(model_path: Path, pipeline_type: str, generation_config: GenerationConfig):
     device = "CPU"
     prompt = "Prompt example is"