refactor: abstract multi-tier block management from BlockManagerPool.

Author: Kang-Meng

xllm/core/framework/block/CMakeLists.txt

@@ -11,11 +11,13 @@ cc_library(
     block_manager_pool.h
     block_manager_impl.h
     concurrent_block_manager_impl.h
+    multi_tier_block_manager_pool.h
   SRCS 
     block.cpp
     block_manager_pool.cpp
     concurrent_block_manager_impl.cpp
     block_manager_impl.cpp
+    multi_tier_block_manager_pool.cpp
   DEPS
     $<$<BOOL:${USE_NPU}>:torch_npu>
     $<$<BOOL:${USE_NPU}>:graph>

xllm/core/framework/block/block_manager_pool.cpp

@@ -24,7 +24,6 @@ BlockManagerPool::BlockManagerPool(const Options& options, int32_t dp_size)
     : options_(options) {
   CHECK(dp_size > 0) << "dp_size must be greater than 0";
   block_managers_.reserve(dp_size);
-  host_block_managers_.reserve(dp_size);
 
   BlockManager::Options npu_options;
   npu_options.num_blocks(options_.num_blocks())
@@ -33,31 +32,16 @@ BlockManagerPool::BlockManagerPool(const Options& options, int32_t dp_size)
       .enable_disagg_pd(options_.enable_disagg_pd())
       .enable_cache_upload(options_.enable_cache_upload());
 
-  BlockManager::Options host_options = npu_options;
-  host_options.num_blocks(options_.host_num_blocks())
-      .enable_cache_upload(false);
-
   for (int32_t i = 0; i < dp_size; ++i) {
     if (options.enable_disagg_pd() || options_.enable_kvcache_store()) {
       block_managers_.emplace_back(
           std::make_unique<ConcurrentBlockManagerImpl>(npu_options));
-      if (options_.host_num_blocks() > 0) {
-        host_block_managers_.emplace_back(
-            std::make_unique<ConcurrentBlockManagerImpl>(host_options));
-      }
     } else {
       block_managers_.emplace_back(
           std::make_unique<BlockManagerImpl>(npu_options));
-      if (options_.host_num_blocks() > 0) {
-        host_block_managers_.emplace_back(
-            std::make_unique<BlockManagerImpl>(host_options));
-      }
     }
   }
   reset_transfer_infos();
-  offload_block_transfer_infos_.resize(block_managers_.size());
-  released_host_blocks_.resize(block_managers_.size());
-  released_device_blocks_.resize(block_managers_.size());
 }
 
 int32_t BlockManagerPool::get_manager_with_max_free_blocks() const {
@@ -89,16 +73,6 @@ int32_t BlockManagerPool::get_dp_rank(Sequence* sequence) const {
   return dp_rank;
 }
 
-BlockManager* BlockManagerPool::get_block_manager(Sequence* sequence,
-                                                  bool is_host) {
-  int32_t dp_rank = get_dp_rank(sequence);
-  if (is_host) {
-    return host_block_managers_[dp_rank].get();
-  } else {
-    return block_managers_[dp_rank].get();
-  }
-}
-
 void BlockManagerPool::deallocate(Request* request) {
   DCHECK(request != nullptr);
   for (auto& sequence : request->sequences()) {
@@ -117,9 +91,6 @@ void BlockManagerPool::deallocate(Sequence* sequence) {
   // add blocks to the prefix cache
   int32_t dp_rank = get_dp_rank(sequence);
   cache(sequence);
-  if (!host_block_managers_.empty()) {
-    save_offload_blocks(sequence);
-  }
   block_managers_[dp_rank]->deallocate(sequence->kv_state().kv_blocks());
   // release the blocks after prefix cache insertion
   sequence->reset();
@@ -130,57 +101,9 @@ BlockManagerPool::get_swap_block_transfer_infos() {
   return &swap_block_transfer_infos_;
 }
 
-std::vector<std::vector<BlockTransferInfo>>*
-BlockManagerPool::get_offload_block_transfer_infos() {
-  return &offload_block_transfer_infos_;
-}
-
-std::vector<std::vector<BlockTransferInfo>>*
-BlockManagerPool::get_load_block_transfer_infos() {
-  return &load_block_transfer_infos_;
-}
-
-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++) {
-    if (futures[i].empty()) {
-      continue;
-    }
-    // TODO(kangmeng): add timeout
-    folly::collectAll(std::move(futures[i]))
-        .via(folly::getGlobalCPUExecutor())
-        .thenValue([host_blocks = std::move(released_host_blocks_[i]),
-                    device_blocks = std::move(released_device_blocks_[i]),
-                    host_block_mgr_ptr = host_block_managers_[i].get(),
-                    device_block_mgr_ptr = block_managers_[i].get()](
-                       std::vector<folly::Try<uint32_t>>&& results) {
-          for (auto&& result : results) {
-            if (result.value() != host_blocks.size()) {
-              LOG(FATAL) << "Offload copy fail, expected " << host_blocks.size()
-                         << ", got " << result.value();
-            }
-          }
-          host_block_mgr_ptr->cache(host_blocks);
-          host_block_mgr_ptr->deallocate({host_blocks});
-          device_block_mgr_ptr->deallocate({device_blocks});
-          return 0;
-        });
-  }
-
-  offload_block_transfer_infos_.clear();
-  released_host_blocks_.clear();
-  released_device_blocks_.clear();
-  offload_block_transfer_infos_.resize(block_managers_.size());
-  released_host_blocks_.resize(block_managers_.size());
-  released_device_blocks_.resize(block_managers_.size());
-}
-
 void BlockManagerPool::reset_transfer_infos() {
   swap_block_transfer_infos_.clear();
   swap_block_transfer_infos_.resize(block_managers_.size());
-  load_block_transfer_infos_.clear();
-  load_block_transfer_infos_.resize(block_managers_.size());
 }
 
 bool BlockManagerPool::allocate(Sequence* sequence) {
@@ -206,9 +129,6 @@ bool BlockManagerPool::allocate(Sequence* sequence, size_t num_tokens) {
   // first try to allocate shared blocks
   if (sequence->kv_state().num_kv_blocks() == 0) {
     allocate_shared(sequence);
-    if (sequence->host_kv_state().num_kv_blocks() == 0) {
-      allocate_host_shared(sequence);
-    }
   }
 
   const size_t num_blocks = sequence->kv_state().num_kv_blocks();
@@ -232,25 +152,6 @@ bool BlockManagerPool::allocate(Sequence* sequence, size_t num_tokens) {
 
   sequence->add_kv_blocks(blocks);
 
-  size_t hbm_cache_token_num = sequence->kv_state().kv_cache_tokens_num();
-  size_t host_cache_token_num = sequence->host_kv_state().kv_cache_tokens_num();
-  if (hbm_cache_token_num < host_cache_token_num) {
-    auto hbm_blocks = sequence->kv_state().kv_blocks();
-    auto host_blocks = sequence->host_kv_state().kv_blocks();
-
-    for (int i = hbm_cache_token_num / options_.block_size();
-         i < host_cache_token_num / options_.block_size();
-         i++) {
-      load_block_transfer_infos_[dp_rank].emplace_back(
-          BlockTransferInfo(host_blocks[i].id(),
-                            hbm_blocks[i].id(),
-                            host_blocks[i].get_immutable_hash_value(),
-                            TransferType::H2D));
-    }
-    sequence->kv_state().incr_kv_cache_tokens_num(host_cache_token_num -
-                                                  hbm_cache_token_num);
-  }
-
   return true;
 }
 
@@ -289,39 +190,6 @@ bool BlockManagerPool::process_beam_search(Sequence* sequence, bool need_swap) {
   return true;
 }
 
-uint32_t BlockManagerPool::pre_allocate(Sequence* sequence) {
-  DCHECK(sequence != nullptr);
-
-  if (!options_.enable_kvcache_store() ||
-      sequence->kv_state().num_kv_blocks() != 0 ||
-      sequence->host_kv_state().num_kv_blocks() != 0) {
-    return 0;
-  }
-
-  int32_t dp_rank = get_dp_rank(sequence);
-  allocate_host_shared(sequence);
-
-  const size_t num_blocks = sequence->host_kv_state().num_kv_blocks();
-  // round down to the nearest block number
-  const size_t block_size = options_.block_size();
-  const size_t num_additional_blocks =
-      sequence->num_tokens() / block_size - num_blocks;
-  if (num_additional_blocks <= 0) {
-    return 0;
-  }
-
-  auto host_blocks =
-      host_block_managers_[dp_rank]->allocate(num_additional_blocks);
-  if (host_blocks.size() != num_additional_blocks) {
-    return 0;
-  }
-  sequence->host_kv_state().add_kv_blocks(host_blocks);
-  PrefixCache::compute_hash_keys(
-      sequence->tokens(), *sequence->host_kv_state().mutable_kv_blocks());
-
-  return num_additional_blocks;
-}
-
 void BlockManagerPool::allocate_shared(Sequence* sequence) {
   // only allocate shared blocks for prefill sequences
   if (options_.enable_prefix_cache()) {
@@ -344,71 +212,6 @@ void BlockManagerPool::cache(Sequence* sequence) {
   block_managers_[dp_rank]->cache(token_ids, *blocks);
 }
 
-void BlockManagerPool::allocate_host_shared(Sequence* sequence) {
-  // only allocate shared blocks for prefill sequences
-  if (sequence->host_kv_state().num_kv_blocks() != 0 ||
-      host_block_managers_.size() != block_managers_.size()) {
-    return;
-  }
-
-  if (options_.enable_prefix_cache()) {
-    int32_t dp_rank = get_dp_rank(sequence);
-    std::vector<Block> shared_blocks =
-        host_block_managers_[dp_rank]->allocate_shared(sequence->tokens());
-    sequence->add_shared_host_kv_blocks(std::move(shared_blocks));
-  }
-}
-
-void BlockManagerPool::save_offload_blocks(Sequence* sequence) {
-  DCHECK(sequence != nullptr);
-
-  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()) {
-    return;
-  }
-
-  int cached_block_num =
-      sequence->host_kv_state().kv_cache_tokens_num() / options_.block_size();
-
-  int32_t dp_rank = get_dp_rank(sequence);
-
-  if (host_blocks->size() > 0) {
-    host_block_managers_[dp_rank]->cache(sequence->tokens(), *host_blocks);
-  }
-
-  size_t needed_block_num =
-      sequence->num_tokens() / options_.block_size() - host_blocks->size();
-
-  if (needed_block_num == 0) {
-    return;
-  }
-
-  sequence->host_kv_state().add_kv_blocks(
-      host_block_managers_[dp_rank]->allocate(needed_block_num));
-
-  for (int i = cached_block_num; i < host_blocks->size(); i++) {
-    if (blocks->at(i).ref_count() != 2) {
-      continue;
-    }
-
-    host_blocks->at(i).set_hash_value(blocks->at(i).get_immutable_hash_value());
-    released_host_blocks_[dp_rank].emplace_back(std::move(host_blocks->at(i)));
-    released_device_blocks_[dp_rank].emplace_back(std::move(blocks->at(i)));
-    offload_block_transfer_infos_[dp_rank].emplace_back(BlockTransferInfo(
-        released_device_blocks_[dp_rank].back().id(),
-        released_host_blocks_[dp_rank].back().id(),
-        released_host_blocks_[dp_rank].back().get_immutable_hash_value(),
-        released_host_blocks_[dp_rank].back().get_hash_value_len(),
-        TransferType::D2G));
-  }
-  host_block_managers_[dp_rank]->cache(
-      *sequence->host_kv_state().mutable_kv_blocks());
-  host_block_managers_[dp_rank]->deallocate(
-      sequence->host_kv_state().kv_blocks());
-}
-
 void BlockManagerPool::get_merged_kvcache_event(KvCacheEvent* event) const {
   for (int32_t i = 0; i < block_managers_.size(); ++i) {
     block_managers_[i]->get_merged_kvcache_event(event);