bugfix: resolve prefetch early termination issue in multi-tprank scenarios.

Kang-Meng · Kang-Meng · commit 97d2df6eca83 · 2025-12-16T18:16:59.000+08:00
diff --git a/xllm/core/distributed_runtime/comm_channel.cpp b/xllm/core/distributed_runtime/comm_channel.cpp
@@ -372,13 +372,13 @@ void CommChannel::transfer_kv_blocks(
 
 class ClientStreamReceiver : public brpc::StreamInputHandler {
  private:
-  std::shared_ptr<std::atomic<bool>> termination_flag_;
+  std::shared_ptr<std::atomic<int32_t>> termination_flag_;
   std::shared_ptr<std::atomic<uint32_t>> success_cnt_;
   std::promise<void> close_promise_;
   std::atomic<bool> promise_set_{false};
 
  public:
-  ClientStreamReceiver(std::shared_ptr<std::atomic<bool>> termination_flag,
+  ClientStreamReceiver(std::shared_ptr<std::atomic<int32_t>> termination_flag,
                        std::shared_ptr<std::atomic<uint32_t>> success_cnt)
       : termination_flag_(termination_flag), success_cnt_(success_cnt) {}
 
@@ -398,10 +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) > 0) {
         success_cnt_->fetch_add(success_cnt, std::memory_order_relaxed);
       } else {
-        termination_flag_->store(true, std::memory_order_release);
+        termination_flag_->fetch_sub(1, std::memory_order_release);
         brpc::StreamClose(id);
         if (!promise_set_.exchange(true)) {
           close_promise_.set_value();
@@ -427,7 +427,7 @@ class ClientStreamReceiver : public brpc::StreamInputHandler {
 
 void CommChannel::prefetch_from_storage(
     const std::vector<BlockTransferInfo>& block_transfer_info,
-    std::shared_ptr<std::atomic<bool>> flag,
+    std::shared_ptr<std::atomic<int32_t>> 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,
diff --git a/xllm/core/distributed_runtime/comm_channel.h b/xllm/core/distributed_runtime/comm_channel.h
@@ -99,7 +99,7 @@ class CommChannel {
 
   virtual void prefetch_from_storage(
       const std::vector<BlockTransferInfo>& block_transfer_info,
-      std::shared_ptr<std::atomic<bool>> flag,
+      std::shared_ptr<std::atomic<int32_t>> flag,
       std::shared_ptr<std::atomic<uint32_t>> success_cnt);
 
   virtual bool get_last_step_result_async(
diff --git a/xllm/core/distributed_runtime/engine.h b/xllm/core/distributed_runtime/engine.h
@@ -97,7 +97,7 @@ class Engine {
   virtual void prefetch_from_storage(
       const uint32_t dp_rank,
       const std::vector<BlockTransferInfo>& block_transfer_info,
-      std::shared_ptr<std::atomic<bool>> flag,
+      std::shared_ptr<std::atomic<int32_t>> flag,
       std::vector<std::shared_ptr<std::atomic<uint32_t>>>* prefetch_results) {
     LOG(FATAL) << " prefetch_from_storage is not implemented!";
   };
diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp
@@ -518,9 +518,10 @@ void LLMEngine::transfer_kv_blocks(
 void LLMEngine::prefetch_from_storage(
     const uint32_t dp_rank,
     const std::vector<BlockTransferInfo>& block_transfer_info,
-    std::shared_ptr<std::atomic<bool>> flag,
+    std::shared_ptr<std::atomic<int32_t>> flag,
     std::vector<std::shared_ptr<std::atomic<uint32_t>>>* prefetch_results) {
   prefetch_results->reserve(dp_local_tp_size_);
+  flag->store(dp_local_tp_size_, std::memory_order_acquire);
   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]
diff --git a/xllm/core/distributed_runtime/llm_engine.h b/xllm/core/distributed_runtime/llm_engine.h
@@ -83,7 +83,7 @@ class LLMEngine : public Engine {
   void prefetch_from_storage(
       const uint32_t dp_rank,
       const std::vector<BlockTransferInfo>& block_transfer_info,
-      std::shared_ptr<std::atomic<bool>> flag,
+      std::shared_ptr<std::atomic<int32_t>> flag,
       std::vector<std::shared_ptr<std::atomic<uint32_t>>>* prefetch_results)
       override;
 
diff --git a/xllm/core/distributed_runtime/remote_worker.cpp b/xllm/core/distributed_runtime/remote_worker.cpp
@@ -314,7 +314,7 @@ void RemoteWorker::transfer_kv_blocks(
 
 void RemoteWorker::prefetch_from_storage(
     const std::vector<BlockTransferInfo>& block_transfer_info,
-    std::shared_ptr<std::atomic<bool>> flag,
+    std::shared_ptr<std::atomic<int32_t>> flag,
     std::shared_ptr<std::atomic<uint32_t>> success_cnt) {
   copy_threadpool_.schedule(
       [this,
diff --git a/xllm/core/distributed_runtime/remote_worker.h b/xllm/core/distributed_runtime/remote_worker.h
@@ -121,7 +121,7 @@ class RemoteWorker : public WorkerClient {
 
   virtual void prefetch_from_storage(
       const std::vector<BlockTransferInfo>& block_transfer_info,
-      std::shared_ptr<std::atomic<bool>> flag,
+      std::shared_ptr<std::atomic<int32_t>> flag,
       std::shared_ptr<std::atomic<uint32_t>> success_cnt) override;
 
   // Run the model and return the output.
diff --git a/xllm/core/framework/block/hierarchy_block_manager_pool.cpp b/xllm/core/framework/block/hierarchy_block_manager_pool.cpp
@@ -59,7 +59,7 @@ void HierarchyBlockManagerPool::deallocate(Sequence* sequence) {
   auto* blocks = sequence->kv_state().mutable_kv_blocks();
   auto* host_blocks = sequence->host_kv_state().mutable_kv_blocks();
 
-  if (blocks->size() == 0 || host_blocks->size() >= blocks->size()) {
+  if (blocks->size() == 0 || host_blocks->size() > blocks->size()) {
     return;
   }
 
@@ -148,12 +148,14 @@ void HierarchyBlockManagerPool::prefetch_from_storage(
             prefill_sequence->tokens());
     prefill_sequence->add_shared_host_kv_blocks(std::move(shared_blocks));
 
-    const size_t num_blocks = prefill_sequence->host_kv_state().num_kv_blocks();
     // round down to the nearest block number
-    const size_t block_size = options_.block_size();
+    size_t shared_blocks_num =
+        prefill_sequence->host_kv_state().shared_kv_blocks_num();
     const size_t num_additional_blocks =
-        prefill_sequence->num_tokens() / block_size - num_blocks;
-    if (num_additional_blocks <= 0) {
+        (prefill_sequence->num_tokens() + options_.block_size() - 1) /
+            options_.block_size() -
+        shared_blocks_num;
+    if (num_additional_blocks <= 1) {
       return;
     }
 
@@ -165,20 +167,19 @@ void HierarchyBlockManagerPool::prefetch_from_storage(
     prefill_sequence->host_kv_state().add_kv_blocks(host_blocks);
     PrefixCache::compute_hash_keys(
         prefill_sequence->tokens(),
-        *prefill_sequence->host_kv_state().mutable_kv_blocks());
+        *prefill_sequence->host_kv_state().mutable_kv_blocks(),
+        shared_blocks_num);
 
-    if (num_additional_blocks > 0) {
+    if (num_additional_blocks > 1) {
       const auto host_blocks = prefill_sequence->host_kv_state().kv_blocks();
       std::vector<BlockTransferInfo> block_transfer_infos;
       block_transfer_infos.reserve(num_additional_blocks);
-      for (int i = host_blocks.size() - num_additional_blocks;
-           i < host_blocks.size();
-           i++) {
-        block_transfer_infos.emplace_back(
-            BlockTransferInfo(-1,
-                              host_blocks[i].id(),
-                              host_blocks[i].get_immutable_hash_value(),
-                              TransferType::G2H));
+      for (int i = 0; i < num_additional_blocks - 1; i++) {
+        block_transfer_infos.emplace_back(BlockTransferInfo(
+            -1,
+            host_blocks[shared_blocks_num + i].id(),
+            host_blocks[shared_blocks_num + i].get_immutable_hash_value(),
+            TransferType::G2H));
       }
 
       engine_->prefetch_from_storage(prefill_sequence->dp_rank(),
@@ -198,8 +199,21 @@ bool HierarchyBlockManagerPool::update_prefetch_result(
 
   bool prefetch_result = true;
   for (auto& prefill_sequence : request->sequences()) {
-    prefetch_result &= prefill_sequence->update_prefetch_result(timeout);
+    uint32_t success_cnt = 0;
+    prefetch_result &=
+        prefill_sequence->update_prefetch_result(timeout, success_cnt);
+
+    if (success_cnt > 0) {
+      int32_t dp_rank = BlockManagerPool::get_dp_rank(prefill_sequence.get());
+      auto host_blocks = prefill_sequence->host_kv_state().kv_blocks();
+      auto cached_blocks =
+          prefill_sequence->host_kv_state().shared_kv_blocks_num();
+
+      host_block_managers_[dp_rank]->cache(
+          host_blocks.slice(cached_blocks - success_cnt, cached_blocks));
+    }
   }
+
   return prefetch_result;
 }
 
diff --git a/xllm/core/framework/prefix_cache/prefix_cache.cpp b/xllm/core/framework/prefix_cache/prefix_cache.cpp
@@ -125,6 +125,11 @@ size_t PrefixCache::insert(const Slice<int32_t>& token_ids,
 }
 
 size_t PrefixCache::insert(const std::vector<Block>& blocks) {
+  Slice<Block> slice(blocks);
+  return insert(slice);
+}
+
+size_t PrefixCache::insert(Slice<Block>& blocks) {
   std::vector<Murmur3Key> insert_keys;
   return insert(blocks, &insert_keys);
 }
@@ -197,7 +202,7 @@ size_t PrefixCache::insert(const Slice<int32_t>& token_ids,
   return n_tokens;
 }
 
-size_t PrefixCache::insert(const std::vector<Block>& blocks,
+size_t PrefixCache::insert(Slice<Block>& blocks,
                            std::vector<Murmur3Key>* insert_keys) {
   const int64_t now = absl::ToUnixMicros(absl::Now());
   DNodeList node_list;
@@ -279,7 +284,8 @@ size_t PrefixCache::evict(size_t n_blocks,
 }
 
 uint32_t PrefixCache::compute_hash_keys(const Slice<int32_t>& token_ids,
-                                        std::vector<Block>& blocks) {
+                                        std::vector<Block>& blocks,
+                                        const size_t cached_blocks) {
   if (blocks.size() == 0) {
     return 0;
   }
@@ -289,8 +295,10 @@ uint32_t PrefixCache::compute_hash_keys(const Slice<int32_t>& token_ids,
     LOG(ERROR) << "token ids do not cover the allocate block.";
     return 0;
   }
+  size_t full_block_size =
+      std::min(token_ids.size() / block_size, blocks.size());
 
-  for (size_t i = 0; i < token_ids.size() / block_size; i++) {
+  for (size_t i = cached_blocks; i < full_block_size; i++) {
     if (i == 0) {
       murmur_hash3(nullptr,
                    token_ids.slice(i * block_size, (i + 1) * block_size),
@@ -302,7 +310,7 @@ uint32_t PrefixCache::compute_hash_keys(const Slice<int32_t>& token_ids,
     }
   }
 
-  return token_ids.size() / block_size;
+  return full_block_size;
 }
 
 }  // namespace xllm
diff --git a/xllm/core/framework/prefix_cache/prefix_cache.h b/xllm/core/framework/prefix_cache/prefix_cache.h
@@ -73,6 +73,7 @@ class PrefixCache {
                         std::vector<Block>& blocks);
 
   // insert the blocks with hash key into the prefix tree
+  virtual size_t insert(Slice<Block>& blocks);
   virtual size_t insert(const std::vector<Block>& blocks);
 
   // evict blocks hold by the prefix cache
@@ -97,15 +98,15 @@ class PrefixCache {
   virtual KvCacheEvent* get_upload_kvcache_events() { return nullptr; }
 
   static uint32_t compute_hash_keys(const Slice<int32_t>& token_ids,
-                                    std::vector<Block>& blocks);
+                                    std::vector<Block>& blocks,
+                                    const size_t cached_blocks = 0);
 
  protected:
   size_t insert(const Slice<int32_t>& token_ids,
                 std::vector<Block>& blocks,
                 std::vector<Murmur3Key>* insert_keys);
 
-  size_t insert(const std::vector<Block>& blocks,
-                std::vector<Murmur3Key>* insert_keys);
+  size_t insert(Slice<Block>& blocks, std::vector<Murmur3Key>* insert_keys);
 
   size_t evict(size_t n_blocks, std::vector<Murmur3Key>* evict_keys);
 
diff --git a/xllm/core/framework/prefix_cache/prefix_cache_with_upload.cpp b/xllm/core/framework/prefix_cache/prefix_cache_with_upload.cpp
@@ -40,6 +40,11 @@ size_t PrefixCacheWithUpload::insert(const Slice<int32_t>& token_ids,
 }
 
 size_t PrefixCacheWithUpload::insert(const std::vector<Block>& blocks) {
+  Slice<Block> slice(blocks);
+  return insert(slice);
+}
+
+size_t PrefixCacheWithUpload::insert(Slice<Block>& blocks) {
   std::vector<Murmur3Key> insert_keys;
   auto n_tokens = PrefixCache::insert(blocks, &insert_keys);
   save_event_async(true, insert_keys);
diff --git a/xllm/core/framework/prefix_cache/prefix_cache_with_upload.h b/xllm/core/framework/prefix_cache/prefix_cache_with_upload.h
@@ -20,6 +20,7 @@ class PrefixCacheWithUpload final : public PrefixCache {
 
   // insert the blocks with hash key into the prefix tree
   size_t insert(const std::vector<Block>& blocks) override;
+  size_t insert(Slice<Block>& blocks) override;
 
   // evict blocks hold by the prefix cache
   // return the actual number of evicted blocks
diff --git a/xllm/core/framework/request/sequence.cpp b/xllm/core/framework/request/sequence.cpp
@@ -45,7 +45,7 @@ Sequence::Sequence(size_t index,
       latest_generate_time_(absl::Now()),
       sequence_params_(seq_params),
       decoder_(std::move(decoder)),
-      termination_flag_(std::make_shared<std::atomic<bool>>(false)) {
+      termination_flag_(std::make_shared<std::atomic<int32_t>>(INT32_MAX)) {
   CHECK(!prompt_token_ids.empty()) << "empty prompt token ids";
   auto capacity = sequence_params_.seq_capacity;
   CHECK_GT(capacity, prompt_token_ids.size()) << "capacity too small";
@@ -95,7 +95,7 @@ Sequence::Sequence(const Sequence& other)
       cur_generated_token_idx_(other.cur_generated_token_idx_),
       first_token_(other.first_token_),
       is_pre_scheduled_step_prefill_(other.is_pre_scheduled_step_prefill_),
-      termination_flag_(std::make_shared<std::atomic<bool>>(false)) {
+      termination_flag_(std::make_shared<std::atomic<int32_t>>(INT32_MAX)) {
   logprob_state_ = std::make_unique<LogprobState>(*other.logprob_state_);
 }
 
@@ -459,12 +459,12 @@ Slice<int32_t> Sequence::get_generated_tokens() const {
   return {tokens_.data(), 0};
 }
 
-bool Sequence::update_prefetch_result(uint32_t timeout) {
+bool Sequence::update_prefetch_result(uint32_t timeout, uint32_t& success_cnt) {
   if (prefetch_results_.empty()) {
     return true;
   }
 
-  if (timeout != 0 && !termination_flag_->load(std::memory_order_acquire)) {
+  if (timeout != 0 && termination_flag_->load(std::memory_order_acquire) > 0) {
     if (!is_timeout_set_) {
       timer_.reset();
       is_timeout_set_ = true;
@@ -476,14 +476,15 @@ bool Sequence::update_prefetch_result(uint32_t timeout) {
     }
   }
 
-  termination_flag_->store(true, std::memory_order_release);
-  uint32_t success_cnt = host_kv_state_.kv_blocks().size();
+  termination_flag_->store(0, std::memory_order_release);
+  success_cnt = host_kv_state_.kv_blocks().size();
   for (auto& cnt : prefetch_results_) {
     success_cnt = std::min(success_cnt, cnt->load());
   }
   if (success_cnt > 0) {
     host_kv_state_.incr_kv_cache_tokens_num(
         success_cnt * host_kv_state_.kv_blocks()[0].size());
+    host_kv_state_.incr_shared_kv_blocks_num(success_cnt);
   }
   prefetch_results_.clear();
   return true;
diff --git a/xllm/core/framework/request/sequence.h b/xllm/core/framework/request/sequence.h
@@ -243,14 +243,14 @@ class Sequence final {
       const Tokenizer& tokenizer,
       std::optional<std::vector<LogProb>>& out_logprobs);
 
-  std::shared_ptr<std::atomic<bool>> get_termination_flag() {
+  std::shared_ptr<std::atomic<int32_t>> get_termination_flag() {
     return termination_flag_;
   }
   std::vector<std::shared_ptr<std::atomic<uint32_t>>>* get_prefetch_results() {
     return &prefetch_results_;
   }
 
-  bool update_prefetch_result(uint32_t timeout);
+  bool update_prefetch_result(uint32_t timeout, uint32_t& success_cnt);
 
   void reset();
 
@@ -362,7 +362,7 @@ class Sequence final {
   std::atomic<bool> cancelled_{false};
 
   // kvcache store copy async result
-  std::shared_ptr<std::atomic<bool>> termination_flag_;
+  std::shared_ptr<std::atomic<int32_t>> termination_flag_;
   std::vector<std::shared_ptr<std::atomic<uint32_t>>> prefetch_results_;
 
   Timer timer_;
diff --git a/xllm/core/framework/request/sequence_kv_state.cpp b/xllm/core/framework/request/sequence_kv_state.cpp
@@ -53,6 +53,11 @@ void KVCacheState::add_kv_blocks(const std::vector<Block>& new_blocks) {
   blocks_.insert(blocks_.end(), new_blocks.begin(), new_blocks.end());
 }
 
+void KVCacheState::incr_shared_kv_blocks_num(size_t num) {
+  CHECK(num_owned_shared_blocks_ + num <= num_kv_blocks());
+  num_owned_shared_blocks_ += num;
+}
+
 void KVCacheState::add_shared_kv_blocks(std::vector<Block>&& blocks,
                                         size_t current_total_num_tokens) {
   if (blocks.empty()) {
diff --git a/xllm/core/framework/request/sequence_kv_state.h b/xllm/core/framework/request/sequence_kv_state.h
@@ -36,6 +36,7 @@ class KVCacheState {
   void add_kv_blocks(const std::vector<Block>& new_blocks);
   void add_shared_kv_blocks(std::vector<Block>&& blocks,
                             size_t current_total_num_tokens);
+  void incr_shared_kv_blocks_num(size_t num);
 
   size_t current_max_tokens_capacity() const;
 
diff --git a/xllm/core/runtime/worker_client.cpp b/xllm/core/runtime/worker_client.cpp
@@ -167,7 +167,7 @@ folly::SemiFuture<uint32_t> WorkerClient::transfer_kv_blocks(
 
 void WorkerClient::prefetch_from_storage(
     const std::vector<BlockTransferInfo>& block_transfer_info,
-    std::shared_ptr<std::atomic<bool>> flag,
+    std::shared_ptr<std::atomic<int32_t>> flag,
     std::shared_ptr<std::atomic<uint32_t>> success_cnt) {
   LOG(FATAL) << "WorkerClient Method prefetch_from_storage is UnImplemented.";
 }
diff --git a/xllm/core/runtime/worker_client.h b/xllm/core/runtime/worker_client.h
@@ -118,7 +118,7 @@ class WorkerClient {
 
   virtual void prefetch_from_storage(
       const std::vector<BlockTransferInfo>& block_transfer_info,
-      std::shared_ptr<std::atomic<bool>> flag,
+      std::shared_ptr<std::atomic<int32_t>> flag,
       std::shared_ptr<std::atomic<uint32_t>> success_cnt);
 
   // Run the model on the given input. async call
diff --git a/xllm/core/scheduler/disagg_pd_scheduler.cpp b/xllm/core/scheduler/disagg_pd_scheduler.cpp
@@ -267,6 +267,8 @@ bool DisaggPDScheduler::add_request(std::shared_ptr<Request>& request) {
   CHECK(request != nullptr);
   CHECK(!request->sequences().empty());
 
+  kv_cache_manager_->prefetch_from_storage(request);
+
   if (request->offline()) {
     // offline request, push to offline queue
     prefill_request_queue_offline_.enqueue(request);

Original file line number	Diff line number	Diff line change
`@@ -125,6 +125,11 @@ size_t PrefixCache::insert(const Slice<int32_t>& token_ids,`
`125`	`125`	`}`
`126`	`126`
`127`	`127`	`size_t PrefixCache::insert(const std::vector<Block>& blocks) {`
	`128`	`+ Slice<Block> slice(blocks);`
	`129`	`+ return insert(slice);`
	`130`	`+}`
	`131`	`+`
	`132`	`+size_t PrefixCache::insert(Slice<Block>& blocks) {`
`128`	`133`	`std::vector<Murmur3Key> insert_keys;`
`129`	`134`	`return insert(blocks, &insert_keys);`
`130`	`135`	`}`
`@@ -197,7 +202,7 @@ size_t PrefixCache::insert(const Slice<int32_t>& token_ids,`
`197`	`202`	`return n_tokens;`
`198`	`203`	`}`
`199`	`204`
`200`		`-size_t PrefixCache::insert(const std::vector<Block>& blocks,`
	`205`	`+size_t PrefixCache::insert(Slice<Block>& blocks,`
`201`	`206`	`std::vector<Murmur3Key>* insert_keys) {`
`202`	`207`	`const int64_t now = absl::ToUnixMicros(absl::Now());`
`203`	`208`	`DNodeList node_list;`
`@@ -279,7 +284,8 @@ size_t PrefixCache::evict(size_t n_blocks,`
`279`	`284`	`}`
`280`	`285`
`281`	`286`	`uint32_t PrefixCache::compute_hash_keys(const Slice<int32_t>& token_ids,`
`282`		`- std::vector<Block>& blocks) {`
	`287`	`+ std::vector<Block>& blocks,`
	`288`	`+ const size_t cached_blocks) {`
`283`	`289`	`if (blocks.size() == 0) {`
`284`	`290`	`return 0;`
`285`	`291`	`}`
`@@ -289,8 +295,10 @@ uint32_t PrefixCache::compute_hash_keys(const Slice<int32_t>& token_ids,`
`289`	`295`	`LOG(ERROR) << "token ids do not cover the allocate block.";`
`290`	`296`	`return 0;`
`291`	`297`	`}`
	`298`	`+ size_t full_block_size =`
	`299`	`+ std::min(token_ids.size() / block_size, blocks.size());`
`292`	`300`
`293`		`- for (size_t i = 0; i < token_ids.size() / block_size; i++) {`
	`301`	`+ for (size_t i = cached_blocks; i < full_block_size; i++) {`
`294`	`302`	`if (i == 0) {`
`295`	`303`	`murmur_hash3(nullptr,`
`296`	`304`	`token_ids.slice(i * block_size, (i + 1) * block_size),`
`@@ -302,7 +310,7 @@ uint32_t PrefixCache::compute_hash_keys(const Slice<int32_t>& token_ids,`
`302`	`310`	`}`
`303`	`311`	`}`
`304`	`312`
`305`		`- return token_ids.size() / block_size;`
	`313`	`+ return full_block_size;`
`306`	`314`	`}`
`307`	`315`
`308`	`316`	`} // namespace xllm`