refactor: refactor the priority and on/offline code.

weizhehuang0827 · weizhehuang0827 · commit 4aa32aa22a62 · 2025-08-30T00:59:01.000+08:00
diff --git a/xllm/core/common/global_flags.cpp b/xllm/core/common/global_flags.cpp
@@ -202,6 +202,6 @@ DEFINE_int32(heart_beat_interval, 3, "heart beat interval");
 
 DEFINE_string(priority_strategy, "FCFS", "priority strategy for requests");
 
-DEFINE_bool(enable_on_preempt_off,
+DEFINE_bool(enable_online_preempt_offline,
             true,
             "whether enable online preempt offline");
diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h
@@ -129,4 +129,4 @@ DECLARE_bool(use_zero_evict);
 
 DECLARE_string(priority_strategy);
 
-DECLARE_bool(enable_on_preempt_off);
+DECLARE_bool(enable_online_preempt_offline);
diff --git a/xllm/core/common/metrics.cpp b/xllm/core/common/metrics.cpp
@@ -88,13 +88,13 @@ DEFINE_GAUGE(num_running_requests, "Number of running requests in scheduler");
 DEFINE_GAUGE(num_waiting_requests, "Number of waiting requests in scheduler");
 DEFINE_GAUGE(num_preempted_requests,
              "Number of preempted requests in scheduler");
-DEFINE_GAUGE(num_offd_preempt_off_requests,
+DEFINE_GAUGE(num_offline_decode_preempt_offline_requests,
              "Number of offline decode preempt offline requests in scheduler");
-DEFINE_GAUGE(num_ond_preempt_on_requests,
+DEFINE_GAUGE(num_online_decode_preempt_online_requests,
              "Number of online decode preempt online requests in scheduler");
-DEFINE_GAUGE(num_onp_preempt_off_requests,
+DEFINE_GAUGE(num_online_prefill_preempt_offline_requests,
              "Number of online prefill preempt offline requests in scheduler");
-DEFINE_GAUGE(num_ond_preempt_off_requests,
+DEFINE_GAUGE(num_online_decode_preempt_offline_requests,
              "Number of online decode preempt offline requests in scheduler");
 
 DEFINE_GAUGE(num_running_sequences, "Number of running sequences");
diff --git a/xllm/core/common/metrics.h b/xllm/core/common/metrics.h
@@ -149,10 +149,10 @@ DECLARE_GAUGE(num_pending_requests);
 DECLARE_GAUGE(num_running_requests);
 DECLARE_GAUGE(num_waiting_requests);
 DECLARE_GAUGE(num_preempted_requests);
-DECLARE_GAUGE(num_offd_preempt_off_requests);
-DECLARE_GAUGE(num_ond_preempt_on_requests);
-DECLARE_GAUGE(num_onp_preempt_off_requests);
-DECLARE_GAUGE(num_ond_preempt_off_requests);
+DECLARE_GAUGE(num_offline_decode_preempt_offline_requests);
+DECLARE_GAUGE(num_online_decode_preempt_online_requests);
+DECLARE_GAUGE(num_online_prefill_preempt_offline_requests);
+DECLARE_GAUGE(num_online_decode_preempt_offline_requests);
 DECLARE_GAUGE(num_running_sequences);
 DECLARE_GAUGE(kv_cache_utilization_perc);
 DECLARE_GAUGE(num_blocks_in_prefix_cache);
diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h
@@ -116,7 +116,7 @@ class Options {
 
   PROPERTY(std::string, priority_strategy) = "FCFS";
 
-  PROPERTY(bool, enable_on_preempt_off) = true;
+  PROPERTY(bool, enable_online_preempt_offline) = true;
 };
 
 }  // namespace xllm
diff --git a/xllm/core/distributed_runtime/disagg_pd_service_impl.cpp b/xllm/core/distributed_runtime/disagg_pd_service_impl.cpp
@@ -90,14 +90,15 @@ std::shared_ptr<Request> DisaggPDServiceImpl::generate_request(
                          output_callback,
                          batch_output_callback);
 
-  auto new_request = std::make_shared<Request>(req.req_id(),
-                                               req.x_request_id(),
-                                               req.x_request_time(),
-                                               std::move(req_state),
-                                               req.service_req_id(),
-                                               req.offline(),
-                                               req.slo_ms(),
-                                               req.priority());
+  auto new_request = std::make_shared<Request>(
+      req.req_id(),
+      req.x_request_id(),
+      req.x_request_time(),
+      std::move(req_state),
+      req.service_req_id(),
+      req.offline(),
+      req.slo_ms(),
+      static_cast<xllm::RequestPriority>(req.priority()));
 
   // add one sequence, rest will be added by scheduler
   return new_request;
diff --git a/xllm/core/framework/request/request.cpp b/xllm/core/framework/request/request.cpp
@@ -36,7 +36,7 @@ Request::Request(const std::string& request_id,
                  const std::string& service_request_id,
                  bool offline,
                  int32_t slo_ms,
-                 xllm::proto::Priority priority)
+                 RequestPriority priority)
     : request_id_(request_id),
       service_request_id_(service_request_id),
       x_request_id_(x_request_id),
diff --git a/xllm/core/framework/request/request.h b/xllm/core/framework/request/request.h
@@ -31,6 +31,8 @@ limitations under the License.
 
 namespace xllm {
 
+enum class RequestPriority { DEFAULT = 0, HIGH = 1, NORMAL = 2, LOW = 3 };
+
 class Request {
  public:
   Request(const std::string& request_id,
@@ -40,7 +42,7 @@ class Request {
           const std::string& service_request_id = "",
           bool offline = false,
           int32_t slo_ms = 0,
-          xllm::proto::Priority priority = xllm::proto::Priority::NORMAL);
+          RequestPriority priority = RequestPriority::NORMAL);
 
   bool finished() const;
 
@@ -86,7 +88,7 @@ class Request {
 
   const bool offline() const { return offline_; }
   const int32_t slo_ms() const { return slo_ms_; }
-  const xllm::proto::Priority priority() const { return priority_; }
+  const RequestPriority priority() const { return priority_; }
 
   RequestState& state() { return state_; }
 
@@ -119,7 +121,7 @@ class Request {
 
   int32_t slo_ms_;
 
-  xllm::proto::Priority priority_;
+  RequestPriority priority_;
 
  private:
   void create_sequences_group();
diff --git a/xllm/core/framework/request/request_params.cpp b/xllm/core/framework/request/request_params.cpp
@@ -54,7 +54,7 @@ RequestParams::RequestParams(const proto::CompletionRequest& request,
     slo_ms = request.slo_ms();
   }
   if (request.has_priority()) {
-    priority = request.priority();
+    priority = static_cast<xllm::RequestPriority>(request.priority());
   }
 
   if (request.has_service_request_id()) {
@@ -203,7 +203,7 @@ void InitFromChatRequest(RequestParams& params, const ChatRequest& request) {
     params.slo_ms = request.slo_ms();
   }
   if (request.has_priority()) {
-    params.priority = request.priority();
+    params.priority = static_cast<xllm::RequestPriority>(request.priority());
   }
 
   if (request.has_service_request_id()) {
diff --git a/xllm/core/framework/request/request_params.h b/xllm/core/framework/request/request_params.h
@@ -29,6 +29,7 @@ limitations under the License.
 #include "core/common/types.h"
 #include "embedding.pb.h"
 #include "multimodal.pb.h"
+#include "request.h"
 #include "request_output.h"
 
 namespace xllm {
@@ -130,7 +131,7 @@ struct RequestParams {
 
   int32_t slo_ms = 0;
 
-  xllm::proto::Priority priority = xllm::proto::Priority::NORMAL;
+  RequestPriority priority = RequestPriority::NORMAL;
 };
 
 }  // namespace xllm
diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h
@@ -124,7 +124,7 @@ struct Options {
 
   PROPERTY(std::string, priority_strategy) = "FCFS";
 
-  PROPERTY(bool, enable_on_preempt_off) = true;
+  PROPERTY(bool, enable_online_preempt_offline) = true;
 };
 
 }  // namespace runtime
diff --git a/xllm/core/scheduler/chunked_prefill_scheduler.cpp b/xllm/core/scheduler/chunked_prefill_scheduler.cpp
@@ -222,7 +222,7 @@ void ChunkedPrefillScheduler::handle_running_queue_requests(
       if (request_to_preempt.get() != request.get()) {
         ++num_preempted_requests;
         block_manager_pool_->deallocate(request_to_preempt.get());
-        running_queue_.->pop_back();
+        running_queue_->pop_back();
         // add preemptable request to waiting priority queue
         request_to_preempt->set_preempted();
         waiting_priority_queue_.push(request_to_preempt);
diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp
@@ -96,11 +96,37 @@ void ContinuousScheduler::create_running_queue(const Options& options) {
   }
 }
 
+bool ContinuousScheduler::check_if_enough_to_evict(
+    DecodePriorityQueue* running_queue_to_evict,
+    Sequence* prefill_sequence,
+    size_t& num_request_to_evict) {
+  // check if it's enough when we evict this requests queue
+  auto block_size = block_manager_pool_->options().block_size();
+  const size_t num_blocks_needed =
+      (prefill_sequence->num_tokens() + block_size - 1) / block_size;
+  size_t num_blocks_can_evict = 0;
+  // count the number of blocks can be preempted
+  for (auto it = running_queue_to_evict->rbegin();
+       it != running_queue_to_evict->rend();
+       ++it) {
+    std::shared_ptr<Request> request_to_preempt = *it;
+    num_request_to_evict++;
+    // count the number of blocks belong to the request
+    for (const auto& seq : request_to_preempt->sequences()) {
+      num_blocks_can_evict += seq->kv_state().num_kv_blocks();
+    }
+    if (num_blocks_needed <= num_blocks_can_evict) {
+      return true;
+    }
+  }
+  return false;
+}
+
 void ContinuousScheduler::handle_prefill_requests(
     size_t& remaining_token_budget,
     size_t& remaining_seq_budget,
     RequestPriorityQueue& waiting_priority_queue,
-    size_t& num_onp_preempt_off_requests,
+    size_t& num_online_prefill_preempt_offline_requests,
     std::vector<std::shared_ptr<Request>>& finished_requests) {
   // Handle new request prompt first.
   // Include those requests that are preempted by others.
@@ -158,7 +184,7 @@ void ContinuousScheduler::handle_prefill_requests(
       // preempt offline decode
       if (!block_manager_pool_->allocate(prefill_sequence.get())) {
         can_schedule = false;
-        if (options_.enable_on_preempt_off() && !request->offline() &&
+        if (options_.enable_online_preempt_offline() && !request->offline() &&
             !running_queue_offline_->empty()) {
           size_t num_request_to_evict = 0;
           // according to the prefill_sequence num tokens to check if can
@@ -171,7 +197,7 @@ void ContinuousScheduler::handle_prefill_requests(
             for (size_t i = 0; i < num_request_to_evict; ++i) {
               std::shared_ptr<Request> request_to_preempt =
                   running_queue_offline_->back();
-              ++num_onp_preempt_off_requests;
+              ++num_online_prefill_preempt_offline_requests;
               block_manager_pool_->deallocate(request_to_preempt.get());
               running_queue_offline_->pop_back();
               // add preemptable request to waiting priority queue
@@ -245,9 +271,9 @@ void ContinuousScheduler::handle_prefill_requests(
 void ContinuousScheduler::handle_decode_requests(
     size_t& remaining_token_budget,
     size_t& remaining_seq_budget,
-    size_t& num_offd_preempt_off_requests,
-    size_t& num_ond_preempt_on_requests,
-    size_t& num_ond_preempt_off_requests,
+    size_t& num_offline_decode_preempt_offline_requests,
+    size_t& num_online_decode_preempt_online_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() &&
@@ -335,11 +361,11 @@ void ContinuousScheduler::handle_decode_requests(
     // continue to evict blocks until enough or no other requests that can be
     // preempted TO IMPROVE: preplan if is enough to evict, if not, then not
     // evict the offline request or online request with lowest priority
-    if (options_.enable_on_preempt_off() && !request->offline() &&
+    if (options_.enable_online_preempt_offline() && !request->offline() &&
         !running_queue_offline_->empty()) {
       std::shared_ptr<Request> request_to_preempt =
           running_queue_offline_->back();
-      ++num_ond_preempt_off_requests;
+      ++num_online_decode_preempt_offline_requests;
       block_manager_pool_->deallocate(request_to_preempt.get());
       running_queue_offline_->pop_back();
       // add preemptable request to waiting priority queue
@@ -350,9 +376,9 @@ void ContinuousScheduler::handle_decode_requests(
       std::shared_ptr<Request> request_to_preempt = running_queue->back();
       if (request_to_preempt.get() != request.get()) {
         if (request->offline()) {
-          ++num_offd_preempt_off_requests;
+          ++num_offline_decode_preempt_offline_requests;
         } else {
-          ++num_ond_preempt_on_requests;
+          ++num_online_decode_preempt_online_requests;
         }
         // TO IMPROVE: kv cache offload to cpu
         block_manager_pool_->deallocate(request_to_preempt.get());
@@ -520,23 +546,7 @@ std::vector<Batch> ContinuousScheduler::prepare_batch() {
       *it = nullptr;
     }
   }
-  // process previous batch
-  // insert running requests back to the running queue, iterating from
-  // the highest priority to the lowest
-  // insert running requests back to the running queue, iterating from
-  // the highest priority to the lowest
-  // 1. last step is prefill step:
-  // new prefill has high priority, but these requests has lower priority
-  // then existed requests in running_queue_ in decoding stage.
-  // so we need to push them to the back of running_queue_->
-  // 2. last step is decode step:
-  // We need to traverse running_requests_ array in reverse order.
-  // Because there may be some unexecuted requests with
-  // lower priorities remaining in the running_queue_->
-  // For the requests in running_requests_,
-  // their priorities are all higher than those of the
-  // remaining requests. Therefore, insert all requests to the front of
-  // running_queue_
+
   if (options_.priority_strategy() == "FCFS") {
     if (last_step_prefill_) {
       // insert all requests to the back of running_queue_
@@ -583,6 +593,7 @@ std::vector<Batch> ContinuousScheduler::prepare_batch() {
       }
     }
   } else {
+    // directly push running requests to the priority queue
     for (auto it = running_requests_.begin(); it != running_requests_.end();
          ++it) {
       if (*it == nullptr) {
@@ -606,20 +617,20 @@ std::vector<Batch> ContinuousScheduler::prepare_batch() {
   size_t remaining_token_budget = options_.max_tokens_per_batch();
   size_t remaining_seq_budget = std::max(options_.max_seqs_per_batch(), 1);
   size_t num_preempted_requests = 0;
-  size_t num_offd_preempt_off_requests = 0;
-  size_t num_ond_preempt_on_requests = 0;
-  size_t num_onp_preempt_off_requests = 0;
-  size_t num_ond_preempt_off_requests = 0;
+  size_t num_offline_decode_preempt_offline_requests = 0;
+  size_t num_online_decode_preempt_online_requests = 0;
+  size_t num_online_prefill_preempt_offline_requests = 0;
+  size_t num_online_decode_preempt_offline_requests = 0;
   // TO IMPROVE?: handle online decode request before prefill offline request
   handle_prefill_requests(remaining_token_budget,
                           remaining_seq_budget,
                           waiting_priority_queue_,
-                          num_onp_preempt_off_requests,
+                          num_online_prefill_preempt_offline_requests,
                           finished_requests);
   handle_prefill_requests(remaining_token_budget,
                           remaining_seq_budget,
                           waiting_priority_queue_offline_,
-                          num_onp_preempt_off_requests,
+                          num_online_prefill_preempt_offline_requests,
                           finished_requests);
 
   if (running_sequences_.empty()) {
@@ -628,21 +639,22 @@ std::vector<Batch> ContinuousScheduler::prepare_batch() {
     // queue
     handle_decode_requests(remaining_token_budget,
                            remaining_seq_budget,
-                           num_offd_preempt_off_requests,
-                           num_ond_preempt_on_requests,
-                           num_ond_preempt_off_requests,
+                           num_offline_decode_preempt_offline_requests,
+                           num_online_decode_preempt_online_requests,
+                           num_online_decode_preempt_offline_requests,
                            running_queue_);
     handle_decode_requests(remaining_token_budget,
                            remaining_seq_budget,
-                           num_offd_preempt_off_requests,
-                           num_ond_preempt_on_requests,
-                           num_ond_preempt_off_requests,
+                           num_offline_decode_preempt_offline_requests,
+                           num_online_decode_preempt_online_requests,
+                           num_online_decode_preempt_offline_requests,
                            running_queue_offline_);
   }
 
-  num_preempted_requests =
-      num_offd_preempt_off_requests + num_ond_preempt_on_requests +
-      num_ond_preempt_off_requests + num_onp_preempt_off_requests;
+  num_preempted_requests = num_offline_decode_preempt_offline_requests +
+                           num_online_decode_preempt_online_requests +
+                           num_online_decode_preempt_offline_requests +
+                           num_online_prefill_preempt_offline_requests;
   if (!finished_requests.empty()) {
     response_processor_->process_completed_requests(finished_requests);
   }
@@ -663,10 +675,14 @@ std::vector<Batch> ContinuousScheduler::prepare_batch() {
             waiting_priority_queue_.size() + running_queue_->size());
 
   GAUGE_ADD(num_preempted_requests, num_preempted_requests);
-  GAUGE_ADD(num_offd_preempt_off_requests, num_offd_preempt_off_requests);
-  GAUGE_ADD(num_ond_preempt_on_requests, num_ond_preempt_on_requests);
-  GAUGE_ADD(num_onp_preempt_off_requests, num_onp_preempt_off_requests);
-  GAUGE_ADD(num_ond_preempt_off_requests, num_ond_preempt_off_requests);
+  GAUGE_ADD(num_offline_decode_preempt_offline_requests,
+            num_offline_decode_preempt_offline_requests);
+  GAUGE_ADD(num_online_decode_preempt_online_requests,
+            num_online_decode_preempt_online_requests);
+  GAUGE_ADD(num_online_prefill_preempt_offline_requests,
+            num_online_prefill_preempt_offline_requests);
+  GAUGE_ADD(num_online_decode_preempt_offline_requests,
+            num_online_decode_preempt_offline_requests);
 
   GAUGE_SET(num_running_sequences, running_sequences_.size());
 
@@ -694,7 +710,8 @@ std::vector<Batch> ContinuousScheduler::schedule_request(
     }
 
     if (!waiting_priority_queue_.empty() || !running_queue_->empty() ||
-        !waiting_priority_queue_offline_.empty()) {
+        !waiting_priority_queue_offline_.empty() ||
+        !running_queue_offline_->empty()) {
       continue;
     }
 
diff --git a/xllm/core/scheduler/continuous_scheduler.h b/xllm/core/scheduler/continuous_scheduler.h
diff --git a/xllm/core/scheduler/continuous_scheduler_test.cpp b/xllm/core/scheduler/continuous_scheduler_test.cpp
diff --git a/xllm/core/scheduler/disagg_pd_scheduler.cpp b/xllm/core/scheduler/disagg_pd_scheduler.cpp
diff --git a/xllm/core/util/concurrent_queue.h b/xllm/core/util/concurrent_queue.h
diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp

Original file line number	Diff line number	Diff line change
`@@ -129,4 +129,4 @@ DECLARE_bool(use_zero_evict);`
`129`	`129`
`130`	`130`	`DECLARE_string(priority_strategy);`
`131`	`131`
`132`		`-DECLARE_bool(enable_on_preempt_off);`
	`132`	`+DECLARE_bool(enable_online_preempt_offline);`
Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ RequestParams::RequestParams(const proto::CompletionRequest& request,`
`54`	`54`	`slo_ms = request.slo_ms();`
`55`	`55`	`}`
`56`	`56`	`if (request.has_priority()) {`
`57`		`- priority = request.priority();`
	`57`	`+ priority = static_cast<xllm::RequestPriority>(request.priority());`
`58`	`58`	`}`
`59`	`59`
`60`	`60`	`if (request.has_service_request_id()) {`
`@@ -203,7 +203,7 @@ void InitFromChatRequest(RequestParams& params, const ChatRequest& request) {`
`203`	`203`	`params.slo_ms = request.slo_ms();`
`204`	`204`	`}`
`205`	`205`	`if (request.has_priority()) {`
`206`		`- params.priority = request.priority();`
	`206`	`+ params.priority = static_cast<xllm::RequestPriority>(request.priority());`
`207`	`207`	`}`
`208`	`208`
`209`	`209`	`if (request.has_service_request_id()) {`