feat: support SLO aware scheduling policy. (#20)

Author: JimHsiung

xllm_service/common/CMakeLists.txt

@@ -12,7 +12,7 @@ cc_library(
     macros.h
     slice.h
     threadpool.h
-    ttft_predictor.h
+    time_predictor.h
     types.h
     utils.h
     hash_util.h
@@ -23,7 +23,7 @@ cc_library(
     global_gflags.cpp
     json_reader.cpp
     threadpool.cpp
-    ttft_predictor.cpp
+    time_predictor.cpp
     utils.cpp
     hash_util.cpp
     xllm/uuid.cpp

xllm_service/common/global_gflags.cpp

@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "common/global_gflags.h"
 
+#include "brpc/reloadable_flags.h"
+
 DEFINE_string(server_host,
               "",
               "Server listen address, may be IPV4/IPV6/UDS."
@@ -83,9 +85,21 @@ DEFINE_int32(detect_disconnected_instance_interval,
              "The interval that server detect the disconnected instance.");
 
 DEFINE_int32(block_size,
-             16,
-             "Number of slots per kv cache block. Default is 16.");
+             128,
+             "Number of slots per kv cache block. Default is 128.");
 
 DEFINE_string(tokenizer_path, "", "tokenizer config path.");
 
 DEFINE_bool(enable_request_trace, false, "Whether to enable request trace");
+
+DEFINE_int32(target_ttft,
+             1000,
+             "Target Time to First Token (TTFT), in milliseconds.");
+
+BRPC_VALIDATE_GFLAG(target_ttft, brpc::NonNegativeInteger);
+
+DEFINE_int32(target_tpot,
+             50,
+             "Target Time Per Output Token (TPOT), in milliseconds.");
+
+BRPC_VALIDATE_GFLAG(target_tpot, brpc::NonNegativeInteger);

xllm_service/common/global_gflags.h

@@ -60,3 +60,7 @@ DECLARE_int32(block_size);
 DECLARE_string(tokenizer_path);
 
 DECLARE_bool(enable_request_trace);
+
+DECLARE_int32(target_ttft);
+
+DECLARE_int32(target_tpot);

xllm_service/common/time_predictor.cpp

@@ -0,0 +1,93 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm-service/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "time_predictor.h"
+
+static constexpr int32_t kDegree = 2;
+
+namespace xllm_service {
+
+TimePredictor::TimePredictor(
+    const std::vector<std::pair<int32_t, double>>& ttft_profiling_data,
+    const std::vector<std::tuple<int32_t, int32_t, double>>&
+        tpot_profiling_data) {
+  if (!ttft_profiling_data.empty()) {
+    // construct Vandermonde matrix
+    int32_t m = ttft_profiling_data.size();
+    int32_t n = kDegree + 1;
+    Eigen::MatrixXd matrix(m, n);
+    for (int32_t i = 0; i < m; ++i) {
+      for (int32_t j = 0; j < n; ++j) {
+        matrix(i, j) = std::pow(ttft_profiling_data[i].first, j);
+      }
+    }
+
+    // construct target vector
+    Eigen::VectorXd target(m);
+    for (int32_t i = 0; i < m; ++i) {
+      target(i) = ttft_profiling_data[i].second;
+    }
+
+    // get coefficients
+    ttft_coefficients_ = matrix.colPivHouseholderQr().solve(target);
+  } else {
+    ttft_coefficients_ = Eigen::VectorXd::Zero(1);
+  }
+
+  if (!tpot_profiling_data.empty()) {
+    int32_t m = tpot_profiling_data.size();
+    int32_t n = kDegree + 1;
+    Eigen::MatrixXd matrix(m, n);
+    for (int32_t i = 0; i < m; ++i) {
+      int32_t avg_length = std::get<0>(tpot_profiling_data[i]);
+      int32_t batch_size = std::get<1>(tpot_profiling_data[i]);
+
+      matrix(i, 0) = 1.0;  // the index 0 is always for constant
+      matrix(i, 1) = batch_size;
+      matrix(i, 2) = batch_size * (avg_length - 1);
+    }
+
+    // construct target vector
+    Eigen::VectorXd target(m);
+    for (int32_t i = 0; i < m; ++i) {
+      target(i) = std::get<2>(tpot_profiling_data[i]);
+    }
+
+    // get coefficients
+    tpot_coefficients_ = matrix.colPivHouseholderQr().solve(target);
+  } else {
+    ttft_coefficients_ = Eigen::VectorXd::Zero(3);
+  }
+}
+
+double TimePredictor::predict_ttft(int32_t length) {
+  double result = 0.0;
+  double power = 1.0;
+  for (int32_t i = 0; i < ttft_coefficients_.size(); ++i) {
+    result += ttft_coefficients_(i) * power;
+    power *= length;
+  }
+
+  return result;
+}
+
+double TimePredictor::predict_tpot(int32_t total_length, int32_t batch_size) {
+  double result = 0.0;
+  result = tpot_coefficients_(0) + tpot_coefficients_(1) * batch_size +
+           tpot_coefficients_(2) * total_length;
+  return result;
+}
+
+}  // namespace xllm_service

xllm_service/common/time_predictor.h

@@ -0,0 +1,40 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm-service/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <Eigen/Dense>
+
+namespace xllm_service {
+
+// Predictor for predicting TTFT and TPOT
+class TimePredictor final {
+ public:
+  TimePredictor(
+      const std::vector<std::pair<int32_t, double>>& ttft_profiling_data,
+      const std::vector<std::tuple<int32_t, int32_t, double>>&
+          tpot_profiling_data);
+  ~TimePredictor() = default;
+
+  double predict_ttft(int32_t length);
+
+  double predict_tpot(int32_t total_length, int32_t batch_size);
+
+ private:
+  Eigen::VectorXd ttft_coefficients_;
+  Eigen::VectorXd tpot_coefficients_;
+};
+
+}  // namespace xllm_service

xllm_service/common/ttft_predictor.cpp

@@ -74,6 +74,8 @@ enum class InstanceType : int8_t {
   PREFILL = 1,
   // decode instance
   DECODE = 2,
+  // mix instance
+  MIX = 3,
 };
 
 struct LoadMetrics {
@@ -114,6 +116,8 @@ struct LoadMetrics {
 
 // Record the latency monitoring metrics of the instance over the recent period
 struct LatencyMetrics {
+  LatencyMetrics() : recent_max_ttft(0), recent_max_tbt(0) {}
+
   LatencyMetrics(const int64_t& recent_max_ttft, const int64_t& recent_max_tbt)
       : recent_max_ttft(recent_max_ttft), recent_max_tbt(recent_max_tbt) {}
 
@@ -125,8 +129,9 @@ struct LatencyMetrics {
 enum class RequestAction : int32_t {
   SCHEDULE = 0,
   FINISH_PREFILL = 1,
-  FINISH_DECODE = 2,
-  CANCEL = 3,
+  GENERATE = 2,
+  FINISH_DECODE = 3,
+  CANCEL = 4,
 };
 
 // Record the request metrics of the instance
@@ -172,13 +177,19 @@ struct InstanceMetaInfo {
   std::vector<uint64_t> v_cache_ids;
   int32_t dp_size;
   // ttft profiling data
-  std::vector<std::pair<int32_t, int64_t>> ttft_profiling_data;
+  std::vector<std::pair<int32_t, double>> ttft_profiling_data;
+  // tpot profiling data
+  std::vector<std::tuple<int32_t, int32_t, double>> tpot_profiling_data;
 
   // latest heatbeat timestamp
   uint64_t latest_timestamp = 0;
 
   uint64_t instance_index = -1;
 
+  // Used to indicate the exact instance type of a MIX type instance currently,
+  // only used when the SLO Aware scheduling policy is enabled.
+  InstanceType current_type = InstanceType::PREFILL;
+
   nlohmann::json serialize_to_json() const {
     nlohmann::json json_val;
     json_val["name"] = name;
@@ -190,6 +201,7 @@ struct InstanceMetaInfo {
     json_val["v_cache_ids"] = v_cache_ids;
     json_val["dp_size"] = dp_size;
     json_val["ttft_profiling_data"] = ttft_profiling_data;
+    json_val["tpot_profiling_data"] = tpot_profiling_data;
     return json_val;
   }
 
@@ -230,6 +242,12 @@ struct InstanceMetaInfo {
         }
       }
 
+      for (const auto& item : json_value.at("tpot_profiling_data")) {
+        if (item.is_array() && item.size() == 3) {
+          tpot_profiling_data.emplace_back(item[0], item[1], item[2]);
+        }
+      }
+
       set_init_timestamp();
     } catch (const std::exception& e) {
       LOG(ERROR) << "json str:" << json_str

xllm_service/common/types.h

@@ -21,6 +21,7 @@ enum InstanceType {
   DEFAULT = 0;
   PREFILL = 1;
   DECODE = 2;
+  MIX = 3;
 }
 
 message WorkerKVAddr {

xllm_service/proto/xllm_rpc_service.proto

@@ -49,6 +49,9 @@ struct Request {
   // instance routing
   Routing routing;
 
+  // the number of generated tokens
+  int64_t num_generated_tokens = 0;
+
   // the estimated TTFT obtained from the TTFT predictor
   int64_t estimated_ttft = 0;
 

xllm_service/request/request.h

@@ -88,6 +88,8 @@ ErrorCode XllmRpcClient::register_instance(const InstanceMetaInfo& metainfo) {
     req.set_type(proto::InstanceType::PREFILL);
   } else if (metainfo.type == InstanceType::DECODE) {
     req.set_type(proto::InstanceType::DECODE);
+  } else if (metainfo.type == InstanceType::MIX) {
+    req.set_type(proto::InstanceType::MIX);
   } else {
     req.set_type(proto::InstanceType::DEFAULT);
   }