feat: support contributing memory to distributed KV cache storage.

Author: Kang-Meng

xllm/core/common/global_flags.cpp

@@ -344,6 +344,10 @@ DEFINE_string(store_local_hostname,
               "",
               "The local host name of the kv cache store client.");
 
+DEFINE_uint64(store_segment_size,
+              0,
+              "The contributed memory to kv cache store.");
+
 // --- computation communication parallel config ---
 
 DEFINE_bool(

xllm/core/common/global_flags.h

@@ -165,6 +165,8 @@ DECLARE_string(store_metadata_server);
 
 DECLARE_string(store_local_hostname);
 
+DECLARE_uint64(store_segment_size);
+
 DECLARE_bool(enable_multi_stream_parallel);
 
 DECLARE_int32(micro_batch_num);

xllm/core/common/options.cpp

@@ -56,6 +56,7 @@ std::string Options::to_string() const {
      << ", store_master_server_address: " << store_master_server_address()
      << ", store_metadata_server: " << store_metadata_server()
      << ", store_local_hostname: " << store_local_hostname()
+     << ", store_segment_size: " << store_segment_size()
      << ", enable_multi_stream_parallel: " << enable_multi_stream_parallel()
      << ", enable_continuous_kvcache: " << enable_continuous_kvcache()
      << ", disable_ttft_profiling: " << disable_ttft_profiling()

xllm/core/common/options.h

@@ -149,6 +149,8 @@ class Options {
 
   PROPERTY(std::string, store_local_hostname) = "";
 
+  PROPERTY(uint64_t, store_segment_size) = 0;
+
   PROPERTY(bool, enable_multi_stream_parallel) = false;
 
   PROPERTY(bool, enable_profile_step_time) = false;

xllm/core/framework/kv_cache/kv_cache_store.cpp

@@ -69,14 +69,42 @@ bool KVCacheStore::init(const StoreConfig& config,
                  << ", and data ptr: " << uint64_t(config_.tensor_data);
     }
   }
+
+  if (config_.segment_size != 0) {
+    global_segment_ptr_ =
+        mooncake::allocate_buffer_allocator_memory(config_.segment_size);
+    if (global_segment_ptr_ == nullptr) {
+      LOG(FATAL) << "allocate_buffer_allocator_memory " << config_.segment_size
+                 << " bytes failed!";
+    }
+    auto mount_result =
+        client_ptr_->MountSegment(global_segment_ptr_, config_.segment_size);
+    if (!mount_result.has_value()) {
+      LOG(ERROR) << "Failed to mount segment: "
+                 << toString(mount_result.error());
+    }
+    LOG(INFO) << "Segment mounted successfully";
+  }
+
   is_initialized_ = true;
   return true;
 }
 
 KVCacheStore::~KVCacheStore() {
+  if (client_ptr_ && global_segment_ptr_) {
+    if (!client_ptr_->UnmountSegment(global_segment_ptr_, config_.segment_size)
+             .has_value()) {
+      LOG(ERROR) << "Failed to unmount segment";
+    }
+  }
+
   if (client_ptr_) {
     client_ptr_.reset();
   }
+
+  if (global_segment_ptr_) {
+    free(global_segment_ptr_);
+  }
 }
 
 uint32_t KVCacheStore::batch_put(

xllm/core/framework/kv_cache/kv_cache_store.h

@@ -19,6 +19,7 @@ struct StoreConfig {
   std::string master_server_address = "";
   int replica_num = 1;
   uint32_t tp_rank = 0;
+  size_t segment_size = 0;
   size_t total_size = 0;
   void* tensor_data = nullptr;
 };
@@ -75,6 +76,8 @@ class KVCacheStore {
   uint64_t v_cache_size_per_block_;
 
   std::shared_ptr<mooncake::Client> client_ptr_;
+
+  void* global_segment_ptr_ = nullptr;
 };
 
 }  // namespace xllm

xllm/core/runtime/master.cpp

@@ -201,6 +201,7 @@ Master::Master(const Options& options, EngineType type) : options_(options) {
         .store_master_server_address(options_.store_master_server_address())
         .store_metadata_server(options_.store_metadata_server())
         .store_local_hostname(options_.store_local_hostname())
+        .store_segment_size(options_.store_segment_size())
         .enable_continuous_kvcache(options_.enable_continuous_kvcache())
         .enable_offline_inference(options_.enable_offline_inference())
         .spawn_worker_path(options_.spawn_worker_path())

xllm/core/runtime/options.h

@@ -155,6 +155,9 @@ struct Options {
   //  value used if port is not included)
   PROPERTY(std::string, store_local_hostname) = "";
 
+  //  the contributed memory to kv cache store
+  PROPERTY(uint64_t, store_segment_size) = 0;
+
   // dit
   // max requests per batch
   PROPERTY(int, max_requests_per_batch) = 0;

xllm/core/runtime/worker_impl.cpp

@@ -178,6 +178,7 @@ bool WorkerImpl::allocate_host_kv_cache(
     config.tp_rank = options_.dp_size() > 1
                          ? options_.node_rank() % options_.dp_size()
                          : options_.node_rank();
+    config.segment_size = options_.store_segment_size();
     config.total_size = aligned_tensor_creater_->get_total_size();
     config.tensor_data = aligned_tensor_creater_->get_base_ptr();
 

xllm/xllm.cpp

@@ -171,6 +171,7 @@ int run() {
       .store_master_server_address(FLAGS_store_master_server_address)
       .store_metadata_server(FLAGS_store_metadata_server)
       .store_local_hostname(FLAGS_store_local_hostname)
+      .store_segment_size(FLAGS_store_segment_size)
       .enable_multi_stream_parallel(FLAGS_enable_multi_stream_parallel)
       .enable_profile_step_time(FLAGS_enable_profile_step_time)
       .enable_profile_token_budget(FLAGS_enable_profile_token_budget)