feat: implement batch prefetch from store.

Author: Kang-Meng

xllm/core/common/global_flags.cpp

@@ -338,14 +338,18 @@ DEFINE_string(store_protocol,
               "tcp",
               "KV cache store protocol(e.g. tcp, rdma).");
 
-DEFINE_string(store_master_server_entry,
+DEFINE_string(store_master_server_address,
               "",
               "The address information of the store master service.");
 
-DEFINE_string(store_metadata_connstring,
+DEFINE_string(store_metadata_server,
               "",
               "The address of the kv cache store metadata service.");
 
+DEFINE_string(store_local_hostname,
+              "",
+              "The local host name of the kv cache store client.");
+
 // --- computation communication parallel config ---
 
 DEFINE_bool(

xllm/core/common/global_flags.h

@@ -163,9 +163,11 @@ DECLARE_bool(enable_kvcache_store);
 
 DECLARE_string(store_protocol);
 
-DECLARE_string(store_master_server_entry);
+DECLARE_string(store_master_server_address);
 
-DECLARE_string(store_metadata_connstring);
+DECLARE_string(store_metadata_server);
+
+DECLARE_string(store_local_hostname);
 
 DECLARE_bool(enable_multi_stream_parallel);
 

xllm/core/common/options.cpp

@@ -53,8 +53,9 @@ std::string Options::to_string() const {
      << ", enable_cache_upload: " << enable_cache_upload()
      << ", enable_kvcache_store: " << enable_kvcache_store()
      << ", store_protocol: " << store_protocol()
-     << ", store_master_server_entry: " << store_master_server_entry()
-     << ", store_metadata_connstring: " << store_metadata_connstring()
+     << ", store_master_server_address: " << store_master_server_address()
+     << ", store_metadata_server: " << store_metadata_server()
+     << ", store_local_hostname: " << store_local_hostname()
      << ", 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

@@ -141,9 +141,11 @@ class Options {
 
   PROPERTY(std::string, store_protocol) = "tcp";
 
-  PROPERTY(std::string, store_master_server_entry) = "";
+  PROPERTY(std::string, store_master_server_address) = "";
 
-  PROPERTY(std::string, store_metadata_connstring) = "";
+  PROPERTY(std::string, store_metadata_server) = "";
+
+  PROPERTY(std::string, store_local_hostname) = "";
 
   PROPERTY(bool, enable_multi_stream_parallel) = false;
 

xllm/core/distributed_runtime/comm_channel.cpp

@@ -18,6 +18,8 @@ limitations under the License.
 #include <brpc/controller.h>
 #include <glog/logging.h>
 
+#include <future>
+
 namespace xllm {
 
 bool CommChannel::init_brpc(const std::string& server_address) {
@@ -335,6 +337,94 @@ void CommChannel::transfer_kv_blocks(
   stub_->TransferBlocks(&cntl, &pb_block_transfer_info, &response, nullptr);
 }
 
+class ClientStreamReceiver : public brpc::StreamInputHandler {
+ private:
+  const std::atomic<bool>& termination_flag_;
+  std::shared_ptr<std::atomic<uint32_t>> success_cnt_;
+  std::promise<void> close_promise_;
+  std::atomic<bool> promise_set_{false};
+
+ public:
+  ClientStreamReceiver(const std::atomic<bool>& termination_flag,
+                       std::shared_ptr<std::atomic<uint32_t>>& success_cnt)
+      : termination_flag_(termination_flag), success_cnt_(success_cnt) {}
+
+  ~ClientStreamReceiver() {
+    if (!promise_set_.exchange(true)) {
+      try {
+        close_promise_.set_value();
+      } catch (const std::exception& e) {
+        LOG(WARNING) << "Exception in destructor: " << e.what();
+      }
+    }
+  }
+
+  std::future<void> get_close_future() { return close_promise_.get_future(); }
+
+  int on_received_messages(brpc::StreamId id,
+                           butil::IOBuf* const messages[],
+                           size_t size) override {
+    for (size_t i = 0; i < size; ++i) {
+      std::string msg_str = messages[i]->to_string();
+      int32_t success_cnt = std::stoi(msg_str);
+
+      if (success_cnt > 0 &&
+          !termination_flag_.load(std::memory_order_acquire)) {
+        success_cnt_->fetch_add(success_cnt, std::memory_order_relaxed);
+      } else {
+        brpc::StreamClose(id);
+        if (!promise_set_.exchange(true)) {
+          close_promise_.set_value();
+        }
+        break;
+      }
+    }
+    return 0;
+  }
+
+  virtual void on_idle_timeout(brpc::StreamId id) override {
+    if (!promise_set_.exchange(true)) {
+      close_promise_.set_value();
+    }
+  }
+
+  virtual void on_closed(brpc::StreamId id) override {
+    if (!promise_set_.exchange(true)) {
+      close_promise_.set_value();
+    }
+  }
+};
+
+void CommChannel::prefetch_from_storage(
+    const std::atomic<bool>& flag,
+    const std::vector<BlockTransferInfo>& block_transfer_info,
+    std::shared_ptr<std::atomic<uint32_t>>& success_cnt) {
+  proto::BlockTransferInfos pb_block_transfer_info;
+  if (!block_transfer_info_to_proto(
+          0x0, block_transfer_info, &pb_block_transfer_info)) {
+    return;
+  }
+  ClientStreamReceiver receiver(flag, success_cnt);
+  brpc::Controller cntl;
+  brpc::StreamOptions stream_options;
+  brpc::StreamId stream_id;
+  proto::Status response;
+  stream_options.handler = &receiver;
+  if (brpc::StreamCreate(&stream_id, cntl, &stream_options) != 0) {
+    LOG(ERROR) << "Failed to create stream";
+    return;
+  }
+
+  stub_->PrefetchFromStorage(
+      &cntl, &pb_block_transfer_info, &response, nullptr);
+
+  if (cntl.Failed()) {
+    LOG(ERROR) << "Fail to connect stream, " << cntl.ErrorText();
+  }
+
+  receiver.get_close_future().wait();
+}
+
 bool CommChannel::get_last_step_result_async(
     folly::Promise<std::optional<RawForwardOutput>>& promise) {
   proto::Empty req;

xllm/core/distributed_runtime/comm_channel.h

@@ -87,10 +87,6 @@ class CommChannel {
       const uint64_t kv_cache_size,
       const std::vector<std::vector<int64_t>>& kv_cache_shape);
 
-  virtual bool load_kv_blocks_from_store_async(
-      const std::vector<CacheBlockInfo>& cache_block_info,
-      folly::Promise<uint32_t>& promise);
-
   virtual void transfer_kv_blocks(
       const std::vector<BlockTransferInfo>& block_transfer_info,
       folly::Promise<uint32_t>& promise);
@@ -99,6 +95,11 @@ class CommChannel {
       const uint64_t batch_id,
       const std::vector<BlockTransferInfo>& block_transfer_info);
 
+  virtual void prefetch_from_storage(
+      const std::atomic<bool>& flag,
+      const std::vector<BlockTransferInfo>& block_transfer_info,
+      std::shared_ptr<std::atomic<uint32_t>>& success_cnt);
+
   virtual bool get_last_step_result_async(
       folly::Promise<std::optional<RawForwardOutput>>& promise);
 

xllm/core/distributed_runtime/remote_worker.cpp

@@ -35,6 +35,7 @@ limitations under the License.
 #include "util/hash_util.h"
 
 namespace xllm {
+
 RemoteWorker::RemoteWorker(int32_t global_rank,
                            const std::string& server_address,
                            const torch::Device& d,
@@ -286,7 +287,7 @@ folly::SemiFuture<uint32_t> RemoteWorker::transfer_kv_blocks(
     const std::vector<BlockTransferInfo>& block_transfer_info) {
   folly::Promise<uint32_t> promise;
   auto future = promise.getSemiFuture();
-  general_threadpool_.schedule(
+  copy_threadpool_.schedule(
       [this,
        block_transfer_info = std::move(block_transfer_info),
        promise = std::move(promise)]() mutable {
@@ -298,14 +299,27 @@ folly::SemiFuture<uint32_t> RemoteWorker::transfer_kv_blocks(
 void RemoteWorker::transfer_kv_blocks(
     const uint64_t batch_id,
     const std::vector<BlockTransferInfo>& block_transfer_info) {
-  general_threadpool_.schedule(
+  copy_threadpool_.schedule(
       [this,
        batch_id = batch_id,
        block_transfer_info = std::move(block_transfer_info)]() mutable {
         channel_->transfer_kv_blocks(batch_id, block_transfer_info);
       });
 }
 
+void RemoteWorker::prefetch_from_storage(
+    const std::atomic<bool>& flag,
+    const std::vector<BlockTransferInfo>& block_transfer_info,
+    std::shared_ptr<std::atomic<uint32_t>>& success_cnt) {
+  copy_threadpool_.schedule(
+      [this,
+       flag = &flag,
+       block_transfer_info = std::move(block_transfer_info),
+       success_cnt = success_cnt]() mutable {
+        channel_->prefetch_from_storage(flag, block_transfer_info, success_cnt);
+      });
+}
+
 const torch::Device& RemoteWorker::device() const {
   LOG(ERROR) << "RemoteWorker Method device is UnImplemented.";
 }

xllm/core/distributed_runtime/remote_worker.h

@@ -117,6 +117,11 @@ class RemoteWorker : public WorkerClient {
       const uint64_t batch_id,
       const std::vector<BlockTransferInfo>& block_transfer_info) override;
 
+  virtual void prefetch_from_storage(
+      const std::atomic<bool>& flag,
+      const std::vector<BlockTransferInfo>& block_transfer_info,
+      std::shared_ptr<std::atomic<uint32_t>>& success_cnt) override;
+
   // Run the model and return the output.
   virtual folly::SemiFuture<std::optional<ForwardOutput>> step_async(
       const ForwardInput& inputs) override;
@@ -144,9 +149,8 @@ class RemoteWorker : public WorkerClient {
   // connection resource
   std::unique_ptr<CommChannel> channel_;
   ThreadPool threadpool_;
-  // general working thread
-  // do some overlap work with model execute
-  ThreadPool general_threadpool_{4};
+  // copy working thread
+  ThreadPool copy_threadpool_{4};
   const torch::Device device_;
 };
 }  // namespace xllm

xllm/core/distributed_runtime/worker_service.cpp

@@ -416,13 +416,12 @@ void WorkerService::PullKVCache(::google::protobuf::RpcController* controller,
 
 void WorkerService::TransferBlocks(
     ::google::protobuf::RpcController* controller,
-    const ::xllm::proto::BlockTransferInfos* req,
-    ::xllm::proto::TransferStatus* resp,
+    const proto::BlockTransferInfos* req,
+    proto::TransferStatus* resp,
     ::google::protobuf::Closure* done) {
   brpc::ClosureGuard done_guard(done);
   std::vector<BlockTransferInfo> block_transfer_info;
-  uint64_t batch_id;
-  proto_to_block_transfer_info(*req, batch_id, block_transfer_info);
+  uint64_t batch_id = proto_to_block_transfer_info(*req, block_transfer_info);
 
   if (batch_id == 0x0) {
     resp->set_success_cnt(worker_->transfer_kv_blocks(block_transfer_info));
@@ -432,6 +431,114 @@ void WorkerService::TransferBlocks(
   return;
 }
 
+class ServerStreamHandler : public brpc::StreamInputHandler {
+ private:
+  std::promise<void> close_promise_;
+  std::atomic<bool> promise_set_{false};
+
+ public:
+  ~ServerStreamHandler() {
+    if (!promise_set_.exchange(true)) {
+      try {
+        close_promise_.set_value();
+      } catch (const std::exception& e) {
+        LOG(WARNING) << "Exception in destructor: " << e.what();
+      }
+    }
+  }
+
+  std::future<void> get_close_future() { return close_promise_.get_future(); }
+
+  int on_received_messages(brpc::StreamId id,
+                           butil::IOBuf* const messages[],
+                           size_t size) override {
+    LOG(WARNING) << "ServerStreamHandler::on_received_messages not implement.";
+    return 0;
+  }
+
+  void on_closed(brpc::StreamId id) override {
+    if (!promise_set_.exchange(true)) {
+      close_promise_.set_value();
+    }
+  }
+
+  void on_idle_timeout(brpc::StreamId id) override {
+    if (!promise_set_.exchange(true)) {
+      LOG(WARNING) << "Stream idle timeout: " << id;
+      close_promise_.set_value();
+    }
+  }
+};
+
+void WorkerService::PrefetchFromStorage(
+    google::protobuf::RpcController* controller,
+    const proto::BlockTransferInfos* req,
+    proto::Status* resp,
+    google::protobuf::Closure* done) {
+  brpc::ClosureGuard done_guard(done);
+  brpc::Controller* cntl = static_cast<brpc::Controller*>(controller);
+
+  auto stream_handler = std::make_unique<ServerStreamHandler>();
+  auto stream_id = std::make_unique<brpc::StreamId>();
+  brpc::StreamOptions stream_options;
+  stream_options.handler = stream_handler.get();
+  if (brpc::StreamAccept(stream_id.get(), *cntl, &stream_options) != 0) {
+    resp->set_ok(false);
+    LOG(ERROR) << "Failed to accept stream!";
+    return;
+  }
+
+  std::vector<BlockTransferInfo> block_transfer_info;
+  proto_to_block_transfer_info(*req, block_transfer_info);
+
+  copy_threadpool_.schedule(
+      [this,
+       block_transfer_info = std::move(block_transfer_info),
+       stream_id = std::move(stream_id),
+       stream_handler = std::move(stream_handler)]() mutable {
+        Slice<BlockTransferInfo> transfer_slice{block_transfer_info};
+        auto close_future = stream_handler->get_close_future();
+        bool is_completed = false;
+
+        for (size_t i = 0; i < transfer_slice.size();
+             i += stream_copy_batch_size_) {
+          auto current_slice = transfer_slice.slice(
+              i, std::min(i + stream_copy_batch_size_, transfer_slice.size()));
+
+          auto success_cnt = worker_->prefetch_from_storage(current_slice);
+
+          if (success_cnt != current_slice.size() ||
+              i + stream_copy_batch_size_ >= transfer_slice.size()) {
+            is_completed = true;
+          }
+
+          butil::IOBuf buf;
+          buf.append(std::to_string(success_cnt));
+          if (brpc::StreamWrite(*stream_id.get(), buf) != 0) {
+            brpc::StreamClose(*stream_id.get());
+            is_completed = false;
+            break;
+          }
+
+          if (is_completed) {
+            if (success_cnt != 0) {
+              butil::IOBuf buf_end;
+              buf_end.append("0");
+              brpc::StreamWrite(*stream_id.get(), buf_end);
+            }
+            break;
+          }
+        }
+        if (is_completed) {
+          close_future.wait();
+        }
+        brpc::StreamClose(*stream_id.get());
+      });
+
+  resp->set_ok(true);
+  return;
+}
+
 void WorkerService::GetDeviceInfo(::google::protobuf::RpcController* controller,
                                   const proto::Empty* req,
                                   proto::DeviceInfo* resp,