feat: remove redundant thread pool in WorkerImpl.

Kang-Meng · Kang-Meng · commit 88b12ae71d26 · 2025-11-14T12:22:28.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) {
@@ -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,8 @@ 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;
+  ThreadPool copy_threadpool_{5};
 };
 
 }  // 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
@@ -21,6 +21,7 @@ limitations under the License.
 #include <glog/logging.h>
 #include <torch/torch.h>
 #if defined(USE_NPU)
+#include "acl/acl.h"
 #include "kernels/npu/xllm_ops/replace_token.h"
 #elif defined(USE_MLU)
 #include <torch_mlu/csrc/framework/core/caching_allocator.h>
@@ -32,7 +33,6 @@ limitations under the License.
 #include <optional>
 #include <utility>
 
-#include "acl/acl.h"
 #include "common/device_monitor.h"
 #include "common/global_flags.h"
 #include "common/metrics.h"
@@ -148,6 +148,9 @@ bool WorkerImpl::allocate_host_kv_cache(
   int64_t host_bolck_size =
       device_kv_cache_shape[0][0] * options_.host_blocks_factor();
   host_kv_cache_shape[0][0] = num_layers;
+  CHECK(!host_kv_cache_shape[1].empty())
+      << "v cache shape should not be empty!";
+  // TODO(kangmeng3): support mlu kvcache
   host_kv_cache_shape[1][0] = num_layers;
 
   // create a KVCache shape: block_size * [layers, token, head, dim]
@@ -743,6 +746,7 @@ int64_t WorkerImpl::get_active_activation_memory() {
       .active_activation_memory;
 }
 
+// TODO(kangmeng): abstract this code(and the code below) into a new class here.
 uint32_t WorkerImpl::offload_kv_blocks(
     const std::vector<BlockTransferInfo>& block_transfer_info) {
   if (block_transfer_info.empty()) {
@@ -972,56 +976,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_;
 
@@ -262,6 +260,10 @@ class WorkerImpl {
   std::shared_ptr<KVCacheTransfer> kv_cache_transfer_;
   aclrtMemcpyBatchAttr h2d_attrs_;
   aclrtMemcpyBatchAttr d2h_attrs_;
+
+  mutable std::mutex mutex_;
+  std::unordered_map<uint64_t, std::shared_ptr<NPULayerSynchronizerImpl>>
+      layer_wise_load_synchronizer_;
 #endif
 
   uint64_t key_cache_size_per_layer_;
@@ -272,10 +274,6 @@ class WorkerImpl {
   Status status_ = Status::UNINITIALIZED;
 
   torch::Tensor expert_load_data_;
-
-  mutable std::mutex mutex_;
-  std::unordered_map<uint64_t, std::shared_ptr<NPULayerSynchronizerImpl>>
-      layer_wise_load_synchronizer_;
 };
 
 class AlignedTensorCreater {
diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp
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`