Remove paddings and add can_use_partial_preemption option for Scheduler

Author: sshlyapn

src/cpp/src/continuous_batching_pipeline.cpp

@@ -112,7 +112,14 @@ class ContinuousBatchingPipeline::Impl {
             updated_config.num_kv_blocks = device_config.get_num_kv_blocks();
         }
 
-        m_scheduler = std::make_shared<Scheduler>(updated_config);
+        bool can_use_partial_preemption = true;
+        if (device_config.get_device().find("GPU") != std::string::npos && !updated_config.dynamic_split_fuse) {
+            // in case of executing a `vLLM-like` pipeline, it's better not to use partial eviction on the GPU,
+            // as it may lead to performance slowdown
+            can_use_partial_preemption = false;
+        }
+
+        m_scheduler = std::make_shared<Scheduler>(updated_config, can_use_partial_preemption);
         // and finally create model runner
         m_model_runner = std::make_shared<ModelRunner>(infer_request, updated_config);
         m_sampler = std::make_shared<Sampler>(m_tokenizer);

src/cpp/src/scheduler.hpp

@@ -8,11 +8,14 @@
 #include <vector>
 
 #include "openvino/genai/scheduler_config.hpp"
+#include "device_config.hpp"
 #include "block_manager.hpp"
 #include "sequence_group.hpp"
 
 namespace ov::genai {
 class Scheduler {
+    bool m_can_use_partial_preemption;
+
     SchedulerConfig m_config;
     BlockManager m_block_manager;
 
@@ -32,8 +35,11 @@ class Scheduler {
         float m_cache_usage = 0.0;
     };
 
-    explicit Scheduler(const SchedulerConfig & config = {}) :
-        m_config(config), m_block_manager(m_config.num_kv_blocks, m_config.enable_prefix_caching, m_config.block_size) { }
+    explicit Scheduler(const SchedulerConfig & config = {}, bool can_use_partial_preemption = true) :
+        m_can_use_partial_preemption(can_use_partial_preemption),
+        m_config(config),
+        m_block_manager(m_config.num_kv_blocks, m_config.enable_prefix_caching, m_config.block_size) {
+    }
 
     Output schedule(std::vector<SequenceGroup::Ptr>& sequence_groups) {
         Output scheduler_output;
@@ -47,7 +53,6 @@ class Scheduler {
         } else {
             // vLLM case
             // schedule prompt phase using whole prompt's input_ids
-            // note, that we also apply padding, while need to be considered by model runner
 
             _schedule_prompt_phase_vllm(sequence_groups, scheduler_output);
 
@@ -105,7 +110,7 @@ class Scheduler {
         size_t preempted_tokens = 0;
         size_t num_blocks_occupied_by_sequence = m_block_manager.get_number_of_blocks_occupied_by_sequence(sequence_group);
 
-        if (num_blocks_occupied_by_sequence <= blocks_needed) {
+        if (num_blocks_occupied_by_sequence <= blocks_needed || !m_can_use_partial_preemption) {
             auto sequences = sequence_group->get_not_finished_sequences();
             for (size_t s = 0; s < sequences.size(); ++s) {
                 auto seq_id = sequences[s]->get_id();
@@ -115,7 +120,7 @@ class Scheduler {
             sequence_group->set_waiting();
             return m_block_manager.num_free_blocks() > prev_blocks_count;
         }
-        
+
         size_t logical_blocks_released;
         if (sequence_group->get_sampling_parameters().is_beam_search()) {
             logical_blocks_released = m_block_manager.free_partially_beam_search_group(sequence_group, blocks_needed);
@@ -126,7 +131,7 @@ class Scheduler {
 
         // calculate the number of preempted tokens
         auto tokens_in_last_block = processed_tokens % block_size;
-        if (tokens_in_last_block == 0) {    
+        if (tokens_in_last_block == 0) {
             tokens_in_last_block = block_size;
         }
         preempted_tokens = tokens_in_last_block + std::max<size_t>((int)logical_blocks_released - 1, 0) * block_size;
@@ -166,7 +171,7 @@ class Scheduler {
         while (!m_block_manager.can_append_slots(sequence_group)) {
             // let's run a sequence for eviction
             size_t evicted_sequence_group_id = _get_low_priority_sequence_group_id(sequence_groups);
-        
+
             if (evicted_sequence_group_id <= sequence_group_id) {
                 // we have a cycle when current group need to evict itself to be in a running state
                 break;
@@ -265,7 +270,7 @@ class Scheduler {
                     sequence_group->clear_scheduled_tokens();
                     continue;
                 }
-                
+
                 // allocate new slots
                 std::map<size_t, std::list<size_t>> copy_blocks_map = m_block_manager.append_slots(sequence_group);
 
@@ -311,19 +316,19 @@ class Scheduler {
         // TODO: it currently does not handle beam search, where beam width should contribute to total number of "num running sequences"
         size_t num_running_sequence_groups = _num_running_sequence_groups(sequence_groups);
 
-        for (size_t sequence_group_id = 0, num_scheduled_tokens = 0, max_sequence_len = 0; sequence_group_id < sequence_groups.size(); ++sequence_group_id) {
+        for (size_t sequence_group_id = 0; sequence_group_id < sequence_groups.size(); ++sequence_group_id) {
             SequenceGroup::Ptr sequence_group = sequence_groups[sequence_group_id];
-            if (!sequence_group->can_generate_tokens() && !sequence_group->is_waiting()) {
+            const bool recompute_evicted_sequences = sequence_group->get_num_processed_tokens() == 0 && !m_can_use_partial_preemption;
+            if ((!sequence_group->can_generate_tokens() || recompute_evicted_sequences) && !sequence_group->is_waiting()) {
                 size_t num_running_seqs = sequence_group->num_running_seqs();
                 // prompt phases can have a single running sequence
                 OPENVINO_ASSERT(num_running_seqs == 1);
                 // here we also assume that sequence must be scheduler in a single shot and has no already generated context
                 if (!m_config.enable_prefix_caching)
                     OPENVINO_ASSERT(sequence_group->get_context_len() == 0);
 
-                int64_t num_available_tokens_in_megabatch = m_config.max_num_batched_tokens - scheduler_output.m_total_num_scheduled_tokens;
+                size_t num_available_tokens_in_megabatch = m_config.max_num_batched_tokens - scheduler_output.m_total_num_scheduled_tokens;
                 size_t sequence_len = sequence_group->get_num_available_tokens_for_batching();
-                max_sequence_len = std::max(max_sequence_len, sequence_len);
 
                 // TODO: better handling
                 // e.g. return status that sequence is ignored and cannot be processed by current scheduling algorigthm
@@ -334,7 +339,7 @@ class Scheduler {
                     break;
 
                 // apply max num batched tokens limitation
-                if (num_available_tokens_in_megabatch < static_cast<int64_t>(max_sequence_len))
+                if (num_available_tokens_in_megabatch < sequence_len)
                     break;
 
                 // apply KV cache limitations
@@ -357,21 +362,20 @@ class Scheduler {
                     {
                         scheduler_output.m_scheduled_sequence_groups_ids.push_back(sequence_group_id);
                         scheduler_output.m_block_tables[seq_id] = m_block_manager.get_block_table(seq_id);
-                        scheduler_output.m_total_num_scheduled_tokens = max_sequence_len * scheduler_output.m_scheduled_sequence_groups_ids.size();
+                        scheduler_output.m_total_num_scheduled_tokens += sequence_len;
                     }
 
                     // update "is_prompt" flag
                     scheduler_output.is_prompt = true;
                 }
 
-                num_scheduled_tokens += sequence_len;
                 num_running_sequence_groups += 1;
             }
         }
     }
 
     void _clear_waiting_sequences(const std::vector<SequenceGroup::Ptr>& sequence_groups) {
-        for (size_t sequence_group_id = 0; sequence_group_id < sequence_groups.size(); ++sequence_group_id) { 
+        for (size_t sequence_group_id = 0; sequence_group_id < sequence_groups.size(); ++sequence_group_id) {
             sequence_groups[sequence_group_id]->clear_waiting_sequences();
         }
     }

tests/cpp/scheduler.cpp

@@ -675,4 +675,166 @@ TEST(TestScheduler, prefix_caching_with_max_new_tokens_equal_1) {
         }
     }
 
-}
+}
+
+TEST(TestScheduler, test_partially_preempted_prompt_not_allowed) {
+    SchedulerConfig scheduler_config;
+    scheduler_config.max_num_batched_tokens = 32;
+    scheduler_config.num_kv_blocks = 6;
+    scheduler_config.block_size = 4;
+    scheduler_config.dynamic_split_fuse = false;
+    scheduler_config.max_num_seqs = 5;
+
+    std::vector<uint64_t> tokens = {0,1,2,3,4,5,6,7,8,9,10,11};
+    SequenceGroup::Ptr sequence_group1 = std::make_shared<SequenceGroup>(0, ov::Tensor(ov::element::i64, {tokens.size()}, tokens.data()),
+                                                                            ov::genai::greedy(), scheduler_config.block_size, scheduler_config.enable_prefix_caching);
+    auto idx0 = (*sequence_group1)[0]->get_id();
+    SequenceGroup::Ptr sequence_group2 = std::make_shared<SequenceGroup>(1, ov::Tensor(ov::element::i64, {tokens.size()}, tokens.data()),
+                                                                            ov::genai::greedy(), scheduler_config.block_size, scheduler_config.enable_prefix_caching);
+    auto idx1 = (*sequence_group2)[0]->get_id();
+    std::vector<SequenceGroup::Ptr> requests = {sequence_group1, sequence_group2};
+
+
+    // schedule 2 sequence groups that use all available 2*3 kv blocks, we used all available kv-blocks.
+    const bool can_use_partial_preemption = false;
+    Scheduler scheduler = Scheduler(scheduler_config, can_use_partial_preemption);
+    auto out1 = scheduler.schedule(requests);
+
+    for (auto req : requests)
+        req->finish_iteration();
+
+    // sequence_group2 should be fully preempted
+    auto out2 = scheduler.schedule(requests);
+
+    // check that sequence_group1 has one more allocated block
+    auto block_table1 = scheduler.get_block_table(*(*sequence_group1)[0]);
+    ASSERT_EQ(block_table1.size(), 4);
+    ASSERT_EQ(block_table1[0]->get_index(), 0);
+    ASSERT_EQ(block_table1[1]->get_index(), 1);
+    ASSERT_EQ(block_table1[2]->get_index(), 2);
+    ASSERT_EQ(block_table1[3]->get_index(), 3);
+    ASSERT_EQ(out2.m_block_tables[idx0].size(), 4);
+    ASSERT_EQ(out2.m_block_tables[idx0][0]->get_index(), 0);
+    ASSERT_EQ(out2.m_block_tables[idx0][1]->get_index(), 1);
+    ASSERT_EQ(out2.m_block_tables[idx0][2]->get_index(), 2);
+    ASSERT_EQ(out2.m_block_tables[idx0][3]->get_index(), 3);
+
+    std::vector<uint64_t> ref_ids = {0};
+    ASSERT_EQ(out2.m_scheduled_sequence_groups_ids, ref_ids);
+    ASSERT_EQ(out2.m_total_num_scheduled_tokens, 1);
+
+    // for vllm case sequence_group2 is fully preempted
+    EXPECT_FALSE(scheduler.has_block_table(idx1));
+
+    for (auto req : requests)
+        req->finish_iteration();
+
+    // finish first sequence
+    requests[0]->get_running_sequences()[0]->set_status(SequenceStatus::FINISHED);
+    scheduler.free_sequence(idx0);
+    clear_finished_sequences(requests);
+
+    // sequence_group2 should be scheduled
+    auto out3 = scheduler.schedule(requests);
+
+    // prompt should be fully scheduled
+    ASSERT_EQ(out3.m_total_num_scheduled_tokens, 12);
+
+    ASSERT_EQ(out3.m_block_tables[idx1][0]->get_index(), 4);
+    ASSERT_EQ(out3.m_block_tables[idx1][1]->get_index(), 5);
+    ASSERT_EQ(out3.m_block_tables[idx1][2]->get_index(), 0);
+
+    auto block_table2 = scheduler.get_block_table(*(*sequence_group2)[0]);
+    ASSERT_EQ(block_table2.size(), 3);
+    ASSERT_EQ(block_table2[0]->get_index(), 4);
+    ASSERT_EQ(block_table2[1]->get_index(), 5);
+    ASSERT_EQ(block_table2[2]->get_index(), 0);
+
+    EXPECT_FALSE(scheduler.has_block_table(idx0));
+}
+
+TEST(TestScheduler, test_partially_preempted_prompt_not_allowed2) {
+    SchedulerConfig scheduler_config;
+    scheduler_config.max_num_batched_tokens = 32;
+    scheduler_config.num_kv_blocks = 6;
+    scheduler_config.block_size = 4;
+    scheduler_config.dynamic_split_fuse = false;
+    scheduler_config.max_num_seqs = 5;
+
+    std::vector<uint64_t> tokens = {0,1,2,3,4,5,6,7,8,9};
+    SequenceGroup::Ptr sequence_group1 = std::make_shared<SequenceGroup>(0, ov::Tensor(ov::element::i64, {tokens.size()}, tokens.data()),
+                                                                            ov::genai::greedy(), scheduler_config.block_size, scheduler_config.enable_prefix_caching);
+    auto idx0 = (*sequence_group1)[0]->get_id();
+    SequenceGroup::Ptr sequence_group2 = std::make_shared<SequenceGroup>(1, ov::Tensor(ov::element::i64, {tokens.size()}, tokens.data()),
+                                                                            ov::genai::greedy(), scheduler_config.block_size, scheduler_config.enable_prefix_caching);
+    auto idx1 = (*sequence_group2)[0]->get_id();
+    std::vector<SequenceGroup::Ptr> requests = {sequence_group1, sequence_group2};
+
+    // schedule 2 sequence groups that use all available 2*3 kv blocks, we used all available kv-blocks.
+    const bool can_use_partial_preemption = false;
+    Scheduler scheduler = Scheduler(scheduler_config, can_use_partial_preemption);
+    scheduler.schedule(requests);
+    for (auto req: requests)
+        req->finish_iteration();
+
+    scheduler.schedule(requests);
+    for (auto req: requests)
+        req->finish_iteration();
+
+    scheduler.schedule(requests);
+    for (auto req: requests)
+        req->finish_iteration();
+
+    // sequence_group2 should be fully preempted
+    scheduler.schedule(requests);
+    for (auto req: requests)
+        req->finish_iteration();
+
+    auto out2 = scheduler.schedule(requests);
+
+    // check that sequence_group1 has one more allocated block
+    auto block_table1 = scheduler.get_block_table(*(*sequence_group1)[0]);
+    ASSERT_EQ(block_table1.size(), 4);
+    ASSERT_EQ(block_table1[0]->get_index(), 0);
+    ASSERT_EQ(block_table1[1]->get_index(), 1);
+    ASSERT_EQ(block_table1[2]->get_index(), 2);
+    ASSERT_EQ(block_table1[3]->get_index(), 3);
+    ASSERT_EQ(out2.m_block_tables[idx0].size(), 4);
+    ASSERT_EQ(out2.m_block_tables[idx0][0]->get_index(), 0);
+    ASSERT_EQ(out2.m_block_tables[idx0][1]->get_index(), 1);
+    ASSERT_EQ(out2.m_block_tables[idx0][2]->get_index(), 2);
+    ASSERT_EQ(out2.m_block_tables[idx0][3]->get_index(), 3);
+
+    std::vector<uint64_t> ref_ids = {0};
+    ASSERT_EQ(out2.m_scheduled_sequence_groups_ids, ref_ids);
+    ASSERT_EQ(out2.m_total_num_scheduled_tokens, 1);
+
+    // for vllm case sequence_group2 is fully preempted
+    EXPECT_FALSE(scheduler.has_block_table(idx1));
+
+    for (auto req: requests)
+        req->finish_iteration();
+
+    // finish first sequence
+    requests[0]->get_running_sequences()[0]->set_status(SequenceStatus::FINISHED);
+    scheduler.free_sequence(idx0);
+    clear_finished_sequences(requests);
+
+    // sequence_group2 should be scheduled
+    auto out3 = scheduler.schedule(requests);
+
+    // prompt should be fully scheduled + generated tokens concatenated to prompt (10 + 2)
+    ASSERT_EQ(out3.m_total_num_scheduled_tokens, 12);
+
+    ASSERT_EQ(out3.m_block_tables[idx1][0]->get_index(), 4);
+    ASSERT_EQ(out3.m_block_tables[idx1][1]->get_index(), 5);
+    ASSERT_EQ(out3.m_block_tables[idx1][2]->get_index(), 0);
+
+    auto block_table2 = scheduler.get_block_table(*(*sequence_group2)[0]);
+    ASSERT_EQ(block_table2.size(), 3);
+    ASSERT_EQ(block_table2[0]->get_index(), 4);
+    ASSERT_EQ(block_table2[1]->get_index(), 5);
+    ASSERT_EQ(block_table2[2]->get_index(), 0);
+
+    EXPECT_FALSE(scheduler.has_block_table(idx0));
+}