feat: reduce KV cache allocation when enable_schedule_overlap is disabled

Author: DongheJin

xllm/core/scheduler/chunked_prefill_scheduler.cpp

@@ -53,7 +53,7 @@ void ChunkedPrefillScheduler::handle_running_queue_requests(
     bool& budget_exhausted,
     bool& blocks_exhausted) {
   while (!running_queue->empty() &&
-         remaining_token_budget > options_.num_speculative_tokens() * 2 &&
+         remaining_token_budget > min_speculative_tokens_required_ &&
          latency_budget > estimate_latency && remaining_seq_budget > 0) {
     std::shared_ptr<Request> request(running_queue->top());
     // TODO: check if request is timeout
@@ -96,7 +96,7 @@ void ChunkedPrefillScheduler::handle_running_queue_requests(
       size_t num_tokens_to_handle =
           sequence->is_prefill_stage()
               ? std::min(assume_max_tokens, num_tokens - kv_cache_tokens_num)
-              : 1 + options_.num_speculative_tokens() * 2;
+              : 1 + min_speculative_tokens_required_;
 
       if (allocated_seqs + 1 > remaining_seq_budget ||
           allocated_tokens + num_tokens_to_handle > remaining_token_budget) {
@@ -723,7 +723,7 @@ bool ChunkedPrefillScheduler::allocate_blocks_for(
     size_t token_budget,
     size_t* current_step_handle_tokens) {
   // token budget should be large enough for one speculative decoding step
-  CHECK_GT(token_budget, options_.num_speculative_tokens() * 2);
+  CHECK_GT(token_budget, min_speculative_tokens_required_);
 
   allocate_shared_blocks_for(sequence);
 
@@ -740,7 +740,7 @@ bool ChunkedPrefillScheduler::allocate_blocks_for(
   // if in decoding stage
   if (options_.num_speculative_tokens() > 0 && !sequence->is_prefill_stage() &&
       kv_cache_tokens_num > 0) {
-    max_handle_num_tokens += options_.num_speculative_tokens() * 2;
+    max_handle_num_tokens += min_speculative_tokens_required_;
   }
 
   // make sure the sequence proceeds forward

xllm/core/scheduler/continuous_scheduler.cpp

@@ -87,6 +87,13 @@ ContinuousScheduler::ContinuousScheduler(Engine* engine, const Options& options)
   instance_info_.name = options_.instance_name().value_or("");
   instance_info_.type = options_.instance_role().value().to_string();
   instance_info_.dp_size = options.dp_size();
+
+  if (options_.enable_schedule_overlap()) {
+    min_speculative_tokens_required_ = options_.num_speculative_tokens() * 2;
+  } else {
+    min_speculative_tokens_required_ = options_.num_speculative_tokens();
+  }
+
 }
 
 ContinuousScheduler::~ContinuousScheduler() { running_requests_.clear(); }
@@ -366,7 +373,7 @@ void ContinuousScheduler::handle_decode_requests(
     size_t& num_online_decode_preempt_offline_requests,
     std::unique_ptr<DecodePriorityQueue>& running_queue) {
   while (!running_queue->empty() &&
-         remaining_token_budget > options_.num_speculative_tokens() * 2 &&
+         remaining_token_budget > min_speculative_tokens_required_ &&
          latency_budget > estimate_latency && remaining_seq_budget > 0) {
     std::shared_ptr<Request> request = running_queue->top();
     // TODO: check if request is timeout
@@ -402,15 +409,15 @@ void ContinuousScheduler::handle_decode_requests(
           break;
         }
       }
-      if (allocated_tokens + options_.num_speculative_tokens() * 2 >=
+      if (allocated_tokens + min_speculative_tokens_required_ >=
               remaining_token_budget ||
           allocated_seqs >= remaining_seq_budget) {
         has_enough_budget = false;
         break;
       }
       // sequence token already appended
       size_t updated_num_tokens =
-          sequence->num_tokens() + options_.num_speculative_tokens() * 2;
+          sequence->num_tokens() + min_speculative_tokens_required_;
       // no blocks left
       if (!kv_cache_manager_->allocate(sequence.get(), updated_num_tokens)) {
         has_enough_blocks = false;
@@ -422,12 +429,12 @@ void ContinuousScheduler::handle_decode_requests(
       }
 
       // update the allocated tokens for the sequence
-      allocated_tokens += options_.num_speculative_tokens() * 2 + 1;
+      allocated_tokens += min_speculative_tokens_required_ + 1;
       allocated_seqs += 1;
       allocated_estimate_latency += seq_estimate_latency;
       candidate_sequences.emplace_back(sequence.get());
-      candidate_token_budgets.emplace_back(
-          options_.num_speculative_tokens() * 2 + 1);
+      candidate_token_budgets.emplace_back(min_speculative_tokens_required_ +
+                                           1);
     }
     CHECK(allocated_tokens <= remaining_token_budget);
     CHECK(allocated_seqs <= remaining_seq_budget);

xllm/core/scheduler/continuous_scheduler.h

@@ -246,6 +246,8 @@ class ContinuousScheduler : public Scheduler {
 
   InstanceInfo instance_info_;
 
+  int32_t min_speculative_tokens_required_ = 0;
+
   virtual void handle_prefill_requests(
       double& latency_budget,
       double& estimate_latency,