feat: support prefetch with timeout.

Author: Kang-Meng

xllm/core/distributed_runtime/comm_channel.cpp

@@ -372,13 +372,13 @@ void CommChannel::transfer_kv_blocks(
 
 class ClientStreamReceiver : public brpc::StreamInputHandler {
  private:
-  const std::atomic<bool>& termination_flag_;
+  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,
+  ClientStreamReceiver(std::atomic<bool>* termination_flag,
                        std::shared_ptr<std::atomic<uint32_t>>& success_cnt)
       : termination_flag_(termination_flag), success_cnt_(success_cnt) {}
 
@@ -398,9 +398,10 @@ class ClientStreamReceiver : public brpc::StreamInputHandler {
       int32_t success_cnt = std::stoi(msg_str);
 
       if (success_cnt > 0 &&
-          !termination_flag_.load(std::memory_order_acquire)) {
+          !termination_flag_->load(std::memory_order_acquire)) {
         success_cnt_->fetch_add(success_cnt, std::memory_order_relaxed);
       } else {
+        termination_flag_->store(true, std::memory_order_release);
         brpc::StreamClose(id);
         if (!promise_set_.exchange(true)) {
           close_promise_.set_value();
@@ -425,8 +426,8 @@ class ClientStreamReceiver : public brpc::StreamInputHandler {
 };
 
 void CommChannel::prefetch_from_storage(
-    const std::atomic<bool>& flag,
     const std::vector<BlockTransferInfo>& block_transfer_info,
+    std::atomic<bool>* flag,
     std::shared_ptr<std::atomic<uint32_t>>& success_cnt) {
   proto::BlockTransferInfos pb_block_transfer_info;
   if (!block_transfer_info_to_proto(block_transfer_info,

xllm/core/distributed_runtime/comm_channel.h

@@ -98,8 +98,8 @@ class CommChannel {
       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::atomic<bool>* flag,
       std::shared_ptr<std::atomic<uint32_t>>& success_cnt);
 
   virtual bool get_last_step_result_async(

xllm/core/distributed_runtime/remote_worker.cpp

@@ -313,15 +313,15 @@ void RemoteWorker::transfer_kv_blocks(
 }
 
 void RemoteWorker::prefetch_from_storage(
-    const std::atomic<bool>& flag,
     const std::vector<BlockTransferInfo>& block_transfer_info,
+    std::atomic<bool>* flag,
     std::shared_ptr<std::atomic<uint32_t>>& success_cnt) {
   copy_threadpool_.schedule(
       [this,
-       flag = &flag,
+       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);
+        channel_->prefetch_from_storage(block_transfer_info, flag, success_cnt);
       });
 }
 

xllm/core/distributed_runtime/remote_worker.h

@@ -120,8 +120,8 @@ class RemoteWorker : public WorkerClient {
       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::atomic<bool>* flag,
       std::shared_ptr<std::atomic<uint32_t>>& success_cnt) override;
 
   // Run the model and return the output.

xllm/core/framework/request/sequence.cpp

@@ -381,6 +381,7 @@ void Sequence::add_host_kv_blocks(const std::vector<Block>& blocks) {
 void Sequence::reset() {
   kv_state_.reset();
   host_kv_state_.reset();
+  timeout_checker_.reset();
   volatile_num_prompt_tokens_ = num_tokens_;
 }
 
@@ -455,9 +456,14 @@ Slice<int32_t> Sequence::get_generated_tokens() const {
   return {tokens_.data(), 0};
 }
 
-void Sequence::update_prefetch_result() {
+bool Sequence::update_prefetch_result() {
   if (prefetch_results_.empty()) {
-    return;
+    return true;
+  }
+
+  if (!termination_flag_.load(std::memory_order_acquire) &&
+      timeout_checker_.check_timeout()) {
+    return false;
   }
 
   termination_flag_.store(true, std::memory_order_release);
@@ -470,6 +476,7 @@ void Sequence::update_prefetch_result() {
         success_cnt * host_kv_state_.kv_blocks()[0].size());
   }
   prefetch_results_.clear();
+  return true;
 }
 
 }  // namespace xllm

xllm/core/framework/request/sequence.h

@@ -19,6 +19,7 @@ limitations under the License.
 #include <absl/time/time.h>
 #include <folly/futures/Future.h>
 
+#include <chrono>
 #include <cstdint>
 #include <vector>
 
@@ -81,6 +82,44 @@ struct SequenceParams {
   StoppingChecker* stopping_checker;  // not owned
 };
 
+static uint32_t timeout_ms = 0;
+class TimeoutChecker {
+ private:
+  std::chrono::steady_clock::time_point timeout_start_;
+  bool is_timeout_set_ = false;
+
+ public:
+  TimeoutChecker() { init(); }
+
+  bool check_timeout() {
+    if (!is_timeout_set_) {
+      timeout_start_ = std::chrono::steady_clock::now();
+      is_timeout_set_ = true;
+
+      return false;
+    } else {
+      auto now = std::chrono::steady_clock::now();
+      auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
+          now - timeout_start_);
+
+      return elapsed.count() >= timeout_ms;
+    }
+  }
+
+  void reset() { is_timeout_set_ = false; }
+
+ private:
+  static void init_timeout() {
+    const char* env_str = std::getenv("PREFETCH_TIMEOUT_MS");
+    timeout_ms = env_str ? std::strtoul(env_str, nullptr, 10) : 0;
+    LOG(INFO) << "Prefetch timeout set as: " << timeout_ms;
+  }
+  static void init() {
+    static std::once_flag flag_;
+    std::call_once(flag_, init_timeout);
+  }
+};
+
 class Sequence final {
  public:
   Sequence(size_t index,
@@ -242,12 +281,12 @@ class Sequence final {
       const Tokenizer& tokenizer,
       std::optional<std::vector<LogProb>>& out_logprobs);
 
-  const std::atomic<bool>& get_termination_flag() { return termination_flag_; }
+  std::atomic<bool>* get_termination_flag() { return &termination_flag_; }
   std::vector<std::shared_ptr<std::atomic<uint32_t>>>* get_prefetch_results() {
     return &prefetch_results_;
   }
 
-  void update_prefetch_result();
+  bool update_prefetch_result();
 
   void reset();
 
@@ -361,6 +400,8 @@ class Sequence final {
   // kvcache store copy async result
   std::atomic<bool> termination_flag_{false};
   std::vector<std::shared_ptr<std::atomic<uint32_t>>> prefetch_results_;
+
+  TimeoutChecker timeout_checker_;
 };
 
 }  // namespace xllm

xllm/core/runtime/engine.h

@@ -96,8 +96,8 @@ class Engine {
 
   virtual void prefetch_from_storage(
       const uint32_t dp_rank,
-      const std::atomic<bool>& flag,
       const std::vector<BlockTransferInfo>& block_transfer_info,
+      std::atomic<bool>* flag,
       std::vector<std::shared_ptr<std::atomic<uint32_t>>>* prefetch_results) {
     LOG(FATAL) << " prefetch_from_storage is not implemented!";
   };

xllm/core/runtime/llm_engine.cpp

@@ -510,15 +510,15 @@ void LLMEngine::transfer_kv_blocks(
 
 void LLMEngine::prefetch_from_storage(
     const uint32_t dp_rank,
-    const std::atomic<bool>& flag,
     const std::vector<BlockTransferInfo>& block_transfer_info,
+    std::atomic<bool>* flag,
     std::vector<std::shared_ptr<std::atomic<uint32_t>>>* prefetch_results) {
-  prefetch_results->resize(dp_local_tp_size_,
-                           std::make_shared<std::atomic<uint32_t>>(0));
+  prefetch_results->reserve(dp_local_tp_size_);
   for (auto tp_rank = 0; tp_rank < dp_local_tp_size_; ++tp_rank) {
+    prefetch_results->emplace_back(std::make_shared<std::atomic<uint32_t>>(0));
     worker_clients_[tp_rank + dp_local_tp_size_ * dp_rank]
         ->prefetch_from_storage(
-            flag, block_transfer_info, prefetch_results->at(tp_rank));
+            block_transfer_info, flag, prefetch_results->at(tp_rank));
   }
 }
 

xllm/core/runtime/llm_engine.h

@@ -82,8 +82,8 @@ class LLMEngine : public Engine {
 
   void prefetch_from_storage(
       const uint32_t dp_rank,
-      const std::atomic<bool>& flag,
       const std::vector<BlockTransferInfo>& block_transfer_info,
+      std::atomic<bool>* flag,
       std::vector<std::shared_ptr<std::atomic<uint32_t>>>* prefetch_results)
       override;
 

xllm/core/runtime/worker_client.cpp

@@ -166,8 +166,8 @@ folly::SemiFuture<uint32_t> WorkerClient::transfer_kv_blocks(
 }
 
 void WorkerClient::prefetch_from_storage(
-    const std::atomic<bool>& flag,
     const std::vector<BlockTransferInfo>& block_transfer_info,
+    std::atomic<bool>* flag,
     std::shared_ptr<std::atomic<uint32_t>>& success_cnt) {
   LOG(FATAL) << "WorkerClient Method prefetch_from_storage is UnImplemented.";
 }