feat: introduce BThreadPool and remove redundant thread pool in WorkerImpl.

Kang-Meng · Kang-Meng · commit e5ba51fb6177 · 2025-11-13T00:19:10.000+08:00
diff --git a/xllm/core/distributed_runtime/worker_service.cpp b/xllm/core/distributed_runtime/worker_service.cpp
@@ -34,6 +34,8 @@ limitations under the License.
 
 namespace xllm {
 
+constexpr uint32_t COPY_BATCH_SIZE = 1;
+
 WorkerService::WorkerService(runtime::Options options,
                              const torch::Device& device)
     : options_(options), device_(device), initialized_(false) {
@@ -490,7 +492,7 @@ void WorkerService::PrefetchFromStorage(
   std::vector<BlockTransferInfo> block_transfer_info;
   proto_to_block_transfer_info(*req, block_transfer_info);
 
-  copy_threadpool_.schedule(
+  BThreadPool::get_instance().schedule(
       [this,
        block_transfer_info = std::move(block_transfer_info),
        stream_id = std::move(stream_id),
@@ -499,15 +501,14 @@ void WorkerService::PrefetchFromStorage(
         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_) {
+        for (size_t i = 0; i < transfer_slice.size(); i += COPY_BATCH_SIZE) {
           auto current_slice = transfer_slice.slice(
-              i, std::min(i + stream_copy_batch_size_, transfer_slice.size()));
+              i, std::min(i + 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()) {
+              i + COPY_BATCH_SIZE >= transfer_slice.size()) {
             is_completed = true;
           }
 
diff --git a/xllm/core/distributed_runtime/worker_service.h b/xllm/core/distributed_runtime/worker_service.h
@@ -154,9 +154,6 @@ class WorkerService : public proto::DistributeWorker {
   std::unique_ptr<std::thread> polling_thread_;
 
   std::unique_ptr<ThreadPool> threadpool_;
-  ThreadPool copy_threadpool_{5};
-
-  uint32_t stream_copy_batch_size_ = 2;
 };
 
 }  // namespace xllm
diff --git a/xllm/core/framework/block/block_manager_pool.cpp b/xllm/core/framework/block/block_manager_pool.cpp
@@ -118,7 +118,7 @@ void BlockManagerPool::deallocate(Sequence* sequence) {
   int32_t dp_rank = get_dp_rank(sequence);
   cache(sequence);
   if (!host_block_managers_.empty()) {
-    record_offload_blocks(sequence);
+    save_offload_blocks(sequence);
   }
   block_managers_[dp_rank]->deallocate(sequence->kv_state().kv_blocks());
   // release the blocks after prefix cache insertion
@@ -140,7 +140,7 @@ BlockManagerPool::get_load_block_transfer_infos() {
   return &load_block_transfer_infos_;
 }
 
-void BlockManagerPool::set_offload_callback(
+void BlockManagerPool::postprocess_offload(
     std::vector<std::vector<folly::SemiFuture<uint32_t>>>& futures) {
   DCHECK(futures.size() == block_managers_.size());
   for (int i = 0; i < futures.size(); i++) {
@@ -356,7 +356,7 @@ void BlockManagerPool::allocate_host_shared(Sequence* sequence) {
   }
 }
 
-void BlockManagerPool::record_offload_blocks(Sequence* sequence) {
+void BlockManagerPool::save_offload_blocks(Sequence* sequence) {
   DCHECK(sequence != nullptr);
 
   auto* blocks = sequence->kv_state().mutable_kv_blocks();
diff --git a/xllm/core/framework/block/block_manager_pool.h b/xllm/core/framework/block/block_manager_pool.h
@@ -64,7 +64,7 @@ class BlockManagerPool final : public KVCacheManager {
   get_offload_block_transfer_infos() override;
   std::vector<std::vector<BlockTransferInfo>>* get_load_block_transfer_infos()
       override;
-  void set_offload_callback(
+  void postprocess_offload(
       std::vector<std::vector<folly::SemiFuture<uint32_t>>>& futures) override;
   void reset_transfer_infos() override;
 
@@ -77,6 +77,9 @@ class BlockManagerPool final : public KVCacheManager {
   std::vector<size_t> num_free_blocks() const override;
   std::vector<size_t> num_used_blocks() const override;
   double kv_cache_utilization() const override;
+  bool allow_host_block_extend() override {
+    return !host_block_managers_.empty();
+  };
 
   // get the options for the block manager
   const Options& options() const { return options_; }
@@ -86,7 +89,7 @@ class BlockManagerPool final : public KVCacheManager {
   int32_t get_dp_rank(Sequence* sequence) const;
 
   void allocate_host_shared(Sequence* sequence);
-  void record_offload_blocks(Sequence* sequence);
+  void save_offload_blocks(Sequence* sequence);
 
   void process_beam_search(Sequence* sequence, bool need_swap = false);
 
diff --git a/xllm/core/framework/block/kv_cache_manager.h b/xllm/core/framework/block/kv_cache_manager.h
@@ -51,7 +51,7 @@ class KVCacheManager {
   virtual std::vector<std::vector<BlockTransferInfo>>*
   get_load_block_transfer_infos() = 0;
 
-  virtual void set_offload_callback(
+  virtual void postprocess_offload(
       std::vector<std::vector<folly::SemiFuture<uint32_t>>>& futures) = 0;
 
   virtual void reset_transfer_infos() = 0;
@@ -62,6 +62,7 @@ class KVCacheManager {
   virtual std::vector<size_t> num_free_blocks() const = 0;
   virtual std::vector<size_t> num_used_blocks() const = 0;
   virtual double kv_cache_utilization() const = 0;
+  virtual bool allow_host_block_extend() { return false; };
 
  protected:
   KVCacheManager() = default;
diff --git a/xllm/core/framework/kv_cache/kv_cache_store.cpp b/xllm/core/framework/kv_cache/kv_cache_store.cpp
@@ -64,9 +64,10 @@ bool KVCacheStore::init(const StoreConfig& config,
         return false;
       }
     } else {
-      LOG(FATAL) << "rdma must RegisterLocalMemory, but got register size: "
+      LOG(ERROR) << "rdma must RegisterLocalMemory, but got register size: "
                  << config_.total_size
                  << ", and data ptr: " << uint64_t(config_.tensor_data);
+      return false;
     }
   }
   is_initialized_ = true;
diff --git a/xllm/core/framework/request/request.h b/xllm/core/framework/request/request.h
@@ -73,12 +73,7 @@ class Request : public RequestBase {
 
   size_t total_num_blocks();
 
-  void set_preempted() {
-    state_.preempted = true;
-    for (auto& seq : sequences_group_->sequences()) {
-      seq->preempted();
-    }
-  }
+  void set_preempted() { state_.preempted = true; }
 
   bool preempted() const { return state_.preempted; }
 
diff --git a/xllm/core/framework/request/sequence.cpp b/xllm/core/framework/request/sequence.cpp
@@ -383,7 +383,6 @@ void Sequence::reset() {
   kv_state_.reset();
   host_kv_state_.reset();
   volatile_num_prompt_tokens_ = num_tokens_;
-  preempted_ = false;
 }
 
 void Sequence::add_shared_kv_blocks(std::vector<Block>&& blocks) {
diff --git a/xllm/core/framework/request/sequence.h b/xllm/core/framework/request/sequence.h
@@ -192,9 +192,6 @@ class Sequence final {
   void close() { closed_ = true; }
   bool is_closed() const { return closed_; }
 
-  void preempted() { preempted_ = true; }
-  bool is_preempted() const { return preempted_; }
-
   // time between two tokens
   int64_t tbt(const absl::Time& now);
   // set sequence ttft
@@ -324,9 +321,6 @@ class Sequence final {
   // is the sequence closed.
   bool closed_ = false;
 
-  // is the sequence preempted.
-  bool preempted_ = false;
-
   // dp_rank
   int32_t dp_rank_ = -1;
 
diff --git a/xllm/core/framework/xtensor/xtensor_manager_pool.h b/xllm/core/framework/xtensor/xtensor_manager_pool.h
@@ -74,9 +74,9 @@ class XTensorManagerPool final : public KVCacheManager {
                   "manager pool";
   }
 
-  void set_offload_callback(
+  void postprocess_offload(
       std::vector<std::vector<folly::SemiFuture<uint32_t>>>& futures) override {
-    LOG(FATAL) << "set_offload_callback is not implemented for page "
+    LOG(FATAL) << "postprocess_offload is not implemented for page "
                   "manager pool";
   }
 
diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp
@@ -740,6 +740,8 @@ int64_t WorkerImpl::get_active_activation_memory() {
       .active_activation_memory;
 }
 
+// TODO(kangmeng): Abstract a new class to reduce coupling between the
+// WorkerImpl and cache offload/load operations.
 uint32_t WorkerImpl::offload_kv_blocks(
     const std::vector<BlockTransferInfo>& block_transfer_info) {
   if (block_transfer_info.empty()) {
@@ -969,56 +971,15 @@ uint32_t WorkerImpl::offload_to_store(
     return block_transfer_info.size();
   }
 
-  folly::Promise<uint32_t> promise;
-  auto future = promise.getSemiFuture();
-
-  batchput_threadpool_.schedule(
-      [this, &block_transfer_info, promise = std::move(promise)]() mutable {
-        promise.setValue(
-            KVCacheStore::get_instance().batch_put(block_transfer_info));
-      });
-
-  return std::move(future)
-      .via(folly::getGlobalCPUExecutor())
-      .within(std::chrono::seconds(TIMEOUT_S))
-      .thenTry([](folly::Try<uint32_t>&& t) -> uint32_t {
-        if (t.hasValue()) {
-          return t.value();
-        } else {
-          LOG(WARNING) << "BatchPut operation timed out";
-          return 0u;
-        }
-      })
-      .get();
+  return KVCacheStore::get_instance().batch_put(block_transfer_info);
 }
 
 uint32_t WorkerImpl::prefetch_from_storage(
     Slice<BlockTransferInfo>& block_transfer_info) {
   if (!options_.enable_kvcache_store()) {
     return 0;
   }
-
-  folly::Promise<uint32_t> promise;
-  auto future = promise.getSemiFuture();
-
-  batchget_threadpool_.schedule(
-      [this, &block_transfer_info, promise = std::move(promise)]() mutable {
-        promise.setValue(
-            KVCacheStore::get_instance().batch_get(block_transfer_info));
-      });
-
-  return std::move(future)
-      .via(folly::getGlobalCPUExecutor())
-      .within(std::chrono::seconds(TIMEOUT_S))
-      .thenTry([](folly::Try<uint32_t>&& t) -> uint32_t {
-        if (t.hasValue()) {
-          return t.value();
-        } else {
-          LOG(WARNING) << "BatchGet operation timed out";
-          return 0u;
-        }
-      })
-      .get();
+  return KVCacheStore::get_instance().batch_get(block_transfer_info);
 }
 
 AlignedTensorCreater::AlignedTensorCreater(
diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h
@@ -215,8 +215,6 @@ class WorkerImpl {
   // working thread for data copy
   std::unique_ptr<ThreadPool> h2d_threadpool_;
   std::unique_ptr<ThreadPool> d2h_threadpool_;
-  ThreadPool batchget_threadpool_{5};
-  ThreadPool batchput_threadpool_{2};
   // copy streams only can be used in h2d_threadpool_ and d2h_threadpool_
   moodycamel::BlockingConcurrentQueue<std::unique_ptr<Stream>> copy_stream_;
 
diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp
@@ -254,11 +254,11 @@ void ContinuousScheduler::handle_prefill_requests(
               std::shared_ptr<Request> request_to_preempt =
                   running_queue_offline_->back();
               ++num_online_prefill_preempt_offline_requests;
-              request_to_preempt->set_preempted();
               kv_cache_manager_->deallocate(request_to_preempt.get());
               running_queue_offline_->pop_back();
               // add preemptable request to waiting priority queue
               // TO IMPROVE?: not process this offline request in current batch
+              request_to_preempt->set_preempted();
               waiting_priority_queue_offline_.push(request_to_preempt);
             }
             if (!kv_cache_manager_->allocate(prefill_sequence.get())) {
@@ -475,20 +475,20 @@ void ContinuousScheduler::handle_decode_requests(
       std::shared_ptr<Request> request_to_preempt =
           running_queue_offline_->back();
       ++num_online_decode_preempt_offline_requests;
-      request_to_preempt->set_preempted();
       kv_cache_manager_->deallocate(request_to_preempt.get());
       running_queue_offline_->pop_back();
       // add preemptable request to waiting priority queue
+      request_to_preempt->set_preempted();
       waiting_priority_queue_offline_.push(request_to_preempt);
       continue;
     } else if (running_queue->size() > 1) {
       std::shared_ptr<Request> request_to_preempt = running_queue->back();
       if (request_to_preempt.get() != request.get()) {
         // TO IMPROVE: kv cache offload to cpu
-        request_to_preempt->set_preempted();
         kv_cache_manager_->deallocate(request_to_preempt.get());
         running_queue->pop_back();
         // add preemptable request to waiting priority queue
+        request_to_preempt->set_preempted();
         if (request_to_preempt->offline()) {
           ++num_offline_decode_preempt_offline_requests;
           waiting_priority_queue_offline_.push(request_to_preempt);
@@ -801,38 +801,44 @@ std::vector<Batch> ContinuousScheduler::prepare_batch() {
     // only update the scheduling latency when there are requests to process
     COUNTER_ADD(scheduling_latency_seconds, timer.elapsed_seconds());
 
-    auto* load_block_transfer_infos =
-        kv_cache_manager_->get_load_block_transfer_infos();
-
-    for (int i = 0; i < batches.size(); i++) {
-      if (!load_block_transfer_infos->at(i).empty()) {
-        batches[i].set_batch_id();
-        engine_->transfer_kv_blocks(
-            i,
-            batches[i].batch_id(),
-            std::move(load_block_transfer_infos->at(i)));
+    // TODO(kangmeng): : Abstract a new class to reduce coupling between the
+    // scheduler and block offload/load operations.
+    if (kv_cache_manager_->allow_host_block_extend()) {
+      auto* load_block_transfer_infos =
+          kv_cache_manager_->get_load_block_transfer_infos();
+
+      for (int i = 0; i < batches.size(); i++) {
+        if (!load_block_transfer_infos->at(i).empty()) {
+          batches[i].set_batch_id();
+          engine_->transfer_kv_blocks(
+              i,
+              batches[i].batch_id(),
+              std::move(load_block_transfer_infos->at(i)));
+        }
       }
     }
   }
 
-  auto* offload_block_transfer_infos =
-      kv_cache_manager_->get_offload_block_transfer_infos();
+  if (kv_cache_manager_->allow_host_block_extend()) {
+    auto* offload_block_transfer_infos =
+        kv_cache_manager_->get_offload_block_transfer_infos();
 
-  bool is_all_dp_copy_info_empty = true;
-  std::vector<std::vector<folly::SemiFuture<uint32_t>>> futures;
-  futures.resize(offload_block_transfer_infos->size());
+    bool is_all_dp_copy_info_empty = true;
+    std::vector<std::vector<folly::SemiFuture<uint32_t>>> futures;
+    futures.resize(offload_block_transfer_infos->size());
 
-  for (int i = 0; i < futures.size(); i++) {
-    if (!offload_block_transfer_infos->at(i).empty()) {
-      futures[i] = std::move(engine_->transfer_kv_blocks(
-          i, std::move(offload_block_transfer_infos->at(i))));
+    for (int i = 0; i < futures.size(); i++) {
+      if (!offload_block_transfer_infos->at(i).empty()) {
+        futures[i] = std::move(engine_->transfer_kv_blocks(
+            i, std::move(offload_block_transfer_infos->at(i))));
 
-      is_all_dp_copy_info_empty = false;
+        is_all_dp_copy_info_empty = false;
+      }
     }
-  }
 
-  if (!is_all_dp_copy_info_empty) {
-    kv_cache_manager_->set_offload_callback(futures);
+    if (!is_all_dp_copy_info_empty) {
+      kv_cache_manager_->postprocess_offload(futures);
+    }
   }
 
   GAUGE_SET(num_pending_requests,
diff --git a/xllm/core/util/threadpool.h b/xllm/core/util/threadpool.h

Original file line number	Diff line number	Diff line change
`@@ -118,7 +118,7 @@ void BlockManagerPool::deallocate(Sequence* sequence) {`
`118`	`118`	`int32_t dp_rank = get_dp_rank(sequence);`
`119`	`119`	`cache(sequence);`
`120`	`120`	`if (!host_block_managers_.empty()) {`
`121`		`- record_offload_blocks(sequence);`
	`121`	`+ save_offload_blocks(sequence);`
`122`	`122`	`}`
`123`	`123`	`block_managers_[dp_rank]->deallocate(sequence->kv_state().kv_blocks());`
`124`	`124`	`// release the blocks after prefix cache insertion`
`@@ -140,7 +140,7 @@ BlockManagerPool::get_load_block_transfer_infos() {`
`140`	`140`	`return &load_block_transfer_infos_;`
`141`	`141`	`}`
`142`	`142`
`143`		`-void BlockManagerPool::set_offload_callback(`
	`143`	`+void BlockManagerPool::postprocess_offload(`
`144`	`144`	`std::vector<std::vector<folly::SemiFuture<uint32_t>>>& futures) {`
`145`	`145`	`DCHECK(futures.size() == block_managers_.size());`
`146`	`146`	`for (int i = 0; i < futures.size(); i++) {`
`@@ -356,7 +356,7 @@ void BlockManagerPool::allocate_host_shared(Sequence* sequence) {`
`356`	`356`	`}`
`357`	`357`	`}`
`358`	`358`
`359`		`-void BlockManagerPool::record_offload_blocks(Sequence* sequence) {`
	`359`	`+void BlockManagerPool::save_offload_blocks(Sequence* sequence) {`
`360`	`360`	`DCHECK(sequence != nullptr);`
`361`	`361`
`362`	`362`	`auto* blocks = sequence->kv_state().mutable_kv_blocks();`
Original file line number	Diff line number	Diff line change
`@@ -64,9 +64,10 @@ bool KVCacheStore::init(const StoreConfig& config,`
`64`	`64`	`return false;`
`65`	`65`	`}`
`66`	`66`	`} else {`
`67`		`- LOG(FATAL) << "rdma must RegisterLocalMemory, but got register size: "`
	`67`	`+ LOG(ERROR) << "rdma must RegisterLocalMemory, but got register size: "`
`68`	`68`	`<< config_.total_size`
`69`	`69`	`<< ", and data ptr: " << uint64_t(config_.tensor_data);`
	`70`	`+ return false;`
`70`	`71`	`}`
`71`	`72`	`}`
`72`	`73`	`is_initialized_ = true;`
Original file line number	Diff line number	Diff line change
`@@ -383,7 +383,6 @@ void Sequence::reset() {`
`383`	`383`	`kv_state_.reset();`
`384`	`384`	`host_kv_state_.reset();`
`385`	`385`	`volatile_num_prompt_tokens_ = num_tokens_;`
`386`		`- preempted_ = false;`
`387`	`386`	`}`
`388`	`387`
`389`	`388`	`void Sequence::add_shared_kv_blocks(std::vector<Block>&& blocks) {`
Original file line number	Diff line number	Diff line change
`@@ -74,9 +74,9 @@ class XTensorManagerPool final : public KVCacheManager {`
`74`	`74`	`"manager pool";`
`75`	`75`	`}`
`76`	`76`
`77`		`- void set_offload_callback(`
	`77`	`+ void postprocess_offload(`
`78`	`78`	`std::vector<std::vector<folly::SemiFuture<uint32_t>>>& futures) override {`
`79`		`- LOG(FATAL) << "set_offload_callback is not implemented for page "`
	`79`	`+ LOG(FATAL) << "postprocess_offload is not implemented for page "`
`80`	`80`	`"manager pool";`
`81`	`81`	`}`
`82`	`82`