feat: adapt host block manager for PD.

Author: Kang-Meng

xllm/core/framework/block/block.h

@@ -67,7 +67,7 @@ class Block final {
     memcpy(hash_value_, hash_value, MURMUR_HASH3_VALUE_LEN);
   }
 
-  uint32_t get_hash_value_len() { return MURMUR_HASH3_VALUE_LEN; }
+  uint32_t get_hash_value_len() const { return MURMUR_HASH3_VALUE_LEN; }
 
  private:
   // increase reference count

xllm/core/framework/block/block_manager.h

@@ -54,8 +54,6 @@ class BlockManager {
 
   virtual void deallocate(const Slice<Block>& blocks) = 0;
 
-  virtual void deallocate(std::vector<Block>& blocks) = 0;
-
   virtual std::vector<Block> allocate(size_t num_blocks) = 0;
 
   virtual std::vector<Block> allocate_shared(

xllm/core/framework/block/block_manager_impl.cpp

@@ -93,12 +93,6 @@ void BlockManagerImpl::deallocate(const Slice<Block>& blocks) {
   }
 }
 
-void BlockManagerImpl::deallocate(std::vector<Block>& blocks) {
-  Slice<Block> slice(blocks);
-  deallocate(slice);
-  blocks.clear();
-}
-
 bool BlockManagerImpl::has_enough_blocks(uint32_t num_blocks) {
   if (num_blocks <= num_free_blocks_) {
     return true;

xllm/core/framework/block/block_manager_impl.h

@@ -35,8 +35,6 @@ class BlockManagerImpl : public BlockManager {
 
   void deallocate(const Slice<Block>& blocks) override;
 
-  void deallocate(std::vector<Block>& blocks) override;
-
   // allocate shared blocks when enable prefix cache
   std::vector<Block> allocate_shared(
       const Slice<int32_t>& tokens_ids,

xllm/core/framework/block/concurrent_block_manager_impl.cpp

@@ -30,11 +30,6 @@ void ConcurrentBlockManagerImpl::deallocate(const Slice<Block>& blocks) {
   BlockManagerImpl::deallocate(blocks);
 }
 
-void ConcurrentBlockManagerImpl::deallocate(std::vector<Block>& blocks) {
-  std::lock_guard<std::mutex> lock(mutex_);
-  BlockManagerImpl::deallocate(blocks);
-}
-
 std::vector<Block> ConcurrentBlockManagerImpl::allocate_shared(
     const Slice<int32_t>& tokens_ids,
     const Slice<Block>& existed_shared_blocks) {

xllm/core/framework/block/concurrent_block_manager_impl.h

@@ -30,8 +30,6 @@ class ConcurrentBlockManagerImpl : public BlockManagerImpl {
 
   void deallocate(const Slice<Block>& blocks) override;
 
-  void deallocate(std::vector<Block>& blocks) override;
-
   // try to share blocks among sequences with the same prefix
   std::vector<Block> allocate_shared(
       const Slice<int32_t>& tokens_ids,

xllm/core/framework/block/hierarchy_block_manager_pool.cpp

@@ -47,9 +47,7 @@ HierarchyBlockManagerPool::HierarchyBlockManagerPool(
   }
 
   load_block_transfer_infos_.resize(host_block_managers_.size());
-  offload_block_transfer_infos_.resize(host_block_managers_.size());
-  saved_host_blocks_.resize(host_block_managers_.size());
-  saved_device_blocks_.resize(host_block_managers_.size());
+  offload_block_pair_queues_.resize(host_block_managers_.size());
 }
 
 void HierarchyBlockManagerPool::deallocate(Sequence* sequence) {
@@ -68,10 +66,6 @@ void HierarchyBlockManagerPool::deallocate(Sequence* sequence) {
   size_t cached_block_num =
       sequence->host_kv_state().kv_cache_tokens_num() / options_.block_size();
 
-  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();
 
@@ -88,14 +82,9 @@ void HierarchyBlockManagerPool::deallocate(Sequence* sequence) {
     }
 
     host_blocks->at(i).set_hash_value(blocks->at(i).get_immutable_hash_value());
-    saved_host_blocks_[dp_rank].emplace_back(std::move(host_blocks->at(i)));
-    saved_device_blocks_[dp_rank].emplace_back(std::move(blocks->at(i)));
-    offload_block_transfer_infos_[dp_rank].emplace_back(BlockTransferInfo(
-        saved_device_blocks_[dp_rank].back().id(),
-        saved_host_blocks_[dp_rank].back().id(),
-        saved_host_blocks_[dp_rank].back().get_immutable_hash_value(),
-        saved_host_blocks_[dp_rank].back().get_hash_value_len(),
-        TransferType::D2G));
+    auto block_pair = std::make_shared<OffloadBlockPair>(
+        std::move(blocks->at(i)), std::move(host_blocks->at(i)));
+    offload_block_pair_queues_[dp_rank].enqueue(std::move(block_pair));
   }
   host_block_managers_[dp_rank]->cache(
       *sequence->host_kv_state().mutable_kv_blocks());
@@ -235,36 +224,50 @@ void HierarchyBlockManagerPool::transfer_blocks(
   }
 
   // offload blocks from device to host and kvcache store
-  for (int i = 0; i < offload_block_transfer_infos_.size(); i++) {
-    if (!offload_block_transfer_infos_[i].empty()) {
-      folly::collectAll(std::move(engine_->transfer_kv_blocks(
-                            i, std::move(offload_block_transfer_infos_[i]))))
+  for (int i = 0; i < offload_block_pair_queues_.size(); i++) {
+    std::vector<BlockTransferInfo> transfer_infos;
+    std::vector<Block> src_blocks;
+    std::vector<Block> dst_blocks;
+
+    std::shared_ptr<OffloadBlockPair> block_pair;
+    while (offload_block_pair_queues_[i].try_dequeue(block_pair)) {
+      src_blocks.emplace_back(std::move(block_pair->src));
+      dst_blocks.emplace_back(std::move(block_pair->dst));
+      transfer_infos.emplace_back(
+          BlockTransferInfo(src_blocks.back().id(),
+                            dst_blocks.back().id(),
+                            dst_blocks.back().get_immutable_hash_value(),
+                            TransferType::D2G));
+      block_pair.reset();
+    }
+
+    if (!transfer_infos.empty()) {
+      folly::collectAll(
+          std::move(engine_->transfer_kv_blocks(i, std::move(transfer_infos))))
           .via(folly::getGlobalCPUExecutor())
-          .thenValue([host_blocks = std::move(saved_host_blocks_[i]),
-                      device_blocks = std::move(saved_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) {
+          .thenValue([device_blocks = std::move(src_blocks),
+                      host_blocks = std::move(dst_blocks),
+                      device_block_mgr_ptr = block_managers_[i].get(),
+                      host_block_mgr_ptr = host_block_managers_[i].get()](
+                         std::vector<folly::Try<uint32_t>>&& results) mutable {
             for (auto&& result : results) {
               if (result.value() != host_blocks.size()) {
                 LOG(FATAL) << "Offload copy fail, expected "
                            << host_blocks.size() << ", got " << result.value();
               }
             }
+
+            device_block_mgr_ptr->deallocate({device_blocks});
+            device_blocks.clear();
+
             host_block_mgr_ptr->cache(host_blocks);
             host_block_mgr_ptr->deallocate({host_blocks});
-            device_block_mgr_ptr->deallocate({device_blocks});
+            host_blocks.clear();
+
             return 0;
           });
     }
   }
-
-  offload_block_transfer_infos_.clear();
-  saved_host_blocks_.clear();
-  saved_device_blocks_.clear();
-  offload_block_transfer_infos_.resize(host_block_managers_.size());
-  saved_host_blocks_.resize(host_block_managers_.size());
-  saved_device_blocks_.resize(host_block_managers_.size());
 }
 
 void HierarchyBlockManagerPool::get_merged_kvcache_event(

xllm/core/framework/block/hierarchy_block_manager_pool.h

@@ -17,13 +17,31 @@ limitations under the License.
 
 #include "block_manager_pool.h"
 #include "runtime/engine.h"
+#include "util/blockingconcurrentqueue.h"
 
 namespace xllm {
 
 class Engine;
 
+struct OffloadBlockPair {
+  OffloadBlockPair(Block& s, Block& d) : src(s), dst(d) {}
+
+  OffloadBlockPair(Block&& s, Block&& d)
+      : src(std::move(s)), dst(std::move(d)) {}
+
+  OffloadBlockPair(Block& s) : src(s) {}
+
+  OffloadBlockPair(Block&& s) : src(std::move(s)) {}
+
+  Block src;
+  Block dst;
+};
+
 class HierarchyBlockManagerPool : public BlockManagerPool {
  public:
+  using OffloadBlockPairQueue =
+      moodycamel::BlockingConcurrentQueue<std::shared_ptr<OffloadBlockPair>>;
+
   explicit HierarchyBlockManagerPool(const BlockManagerPool::Options& options,
                                      Engine* engine,
                                      int32_t dp_size = 1);
@@ -51,9 +69,7 @@ class HierarchyBlockManagerPool : public BlockManagerPool {
 
   // BlockTransferInfo per step
   std::vector<std::vector<BlockTransferInfo>> load_block_transfer_infos_;
-  std::vector<std::vector<BlockTransferInfo>> offload_block_transfer_infos_;
-  std::vector<std::vector<Block>> saved_host_blocks_;
-  std::vector<std::vector<Block>> saved_device_blocks_;
+  std::vector<OffloadBlockPairQueue> offload_block_pair_queues_;
 };
 
 }  // namespace xllm

xllm/core/framework/model/model_input_params.h

@@ -24,6 +24,7 @@ limitations under the License.
 #include "framework/batch/batch_forward_type.h"
 #include "framework/request/mm_data.h"
 #include "npu_dp_ep_padding.h"
+#include "util/hash_util.h"
 #include "util/tensor_helper.h"
 
 namespace xllm {
@@ -38,47 +39,60 @@ enum class TransferType : uint8_t {
 struct BlockTransferInfo {
   int32_t src_block_id = -1;
   int32_t dst_block_id = -1;
-  uint8_t* hash_key = nullptr;
+  uint8_t hash_key[MURMUR_HASH3_VALUE_LEN];
   TransferType transfer_type;
-  uint32_t hash_key_len = -1;
 
   BlockTransferInfo(int32_t src_block_id, int32_t dst_block_id) {
     this->src_block_id = src_block_id;
     this->dst_block_id = dst_block_id;
   }
 
-  BlockTransferInfo(int32_t src_block_id,
-                    int32_t dst_block_id,
-                    const uint8_t* hash_key,
-                    TransferType transfer_type) {
-    this->src_block_id = src_block_id;
-    this->dst_block_id = dst_block_id;
-    this->hash_key = const_cast<uint8_t*>(hash_key);
-    this->transfer_type = transfer_type;
+  BlockTransferInfo(int32_t src_id,
+                    int32_t dst_id,
+                    const uint8_t* key,
+                    TransferType type)
+      : src_block_id(src_id), dst_block_id(dst_id), transfer_type(type) {
+    memcpy(hash_key, key, MURMUR_HASH3_VALUE_LEN);
   }
 
-  BlockTransferInfo(int32_t src_block_id,
-                    int32_t dst_block_id,
-                    const uint8_t* hash_key,
-                    uint32_t hash_key_len,
-                    TransferType transfer_type) {
-    this->src_block_id = src_block_id;
-    this->dst_block_id = dst_block_id;
-    this->hash_key = new uint8_t[hash_key_len];
-    memcpy(this->hash_key, hash_key, hash_key_len);
-    this->transfer_type = transfer_type;
+  BlockTransferInfo(const BlockTransferInfo& other)
+      : src_block_id(other.src_block_id),
+        dst_block_id(other.dst_block_id),
+        transfer_type(other.transfer_type) {
+    memcpy(hash_key, other.hash_key, MURMUR_HASH3_VALUE_LEN);
   }
 
-  ~BlockTransferInfo() {
-    if (hash_key_len != -1 && hash_key != nullptr) {
-      delete[] hash_key;
-    }
+  BlockTransferInfo(BlockTransferInfo&& other)
+      : src_block_id(other.src_block_id),
+        dst_block_id(other.dst_block_id),
+        transfer_type(other.transfer_type) {
+    memcpy(hash_key, other.hash_key, MURMUR_HASH3_VALUE_LEN);
+
+    other.src_block_id = -1;
+    other.dst_block_id = -1;
+  }
+
+  BlockTransferInfo& operator=(const BlockTransferInfo& other) {
+    src_block_id = other.src_block_id;
+    dst_block_id = other.dst_block_id;
+    transfer_type = other.transfer_type;
+    memcpy(hash_key, other.hash_key, MURMUR_HASH3_VALUE_LEN);
+  }
+
+  BlockTransferInfo& operator=(BlockTransferInfo&& other) {
+    src_block_id = other.src_block_id;
+    dst_block_id = other.dst_block_id;
+    transfer_type = other.transfer_type;
+    memcpy(hash_key, other.hash_key, MURMUR_HASH3_VALUE_LEN);
+
+    other.src_block_id = -1;
+    other.dst_block_id = -1;
   }
 
   std::string to_string() const {
     std::string rt = ", has_key:";
     for (int i = 0; i < 16; i++) {
-      rt += std::to_string(int64_t(*(hash_key + i))) + " ";
+      rt += std::to_string(int64_t(hash_key[i])) + " ";
     }
     return std::to_string(src_block_id) + "->" + std::to_string(dst_block_id) +
            ", " + std::to_string(uint32_t(transfer_type)) + rt;

xllm/core/runtime/params_utils.cpp

@@ -723,7 +723,6 @@ uint64_t proto_to_block_transfer_info(
         pb_block_transfer_info.transfer_infos(i).dst_block_id(),
         reinterpret_cast<const uint8_t*>(
             pb_block_transfer_info.transfer_infos(i).hash_key().data()),
-        pb_block_transfer_info.transfer_infos(i).hash_key().size(),
         TransferType(pb_block_transfer_info.transfer_type()));
   }
 
@@ -737,11 +736,6 @@ bool block_transfer_info_to_proto(
       block_transfer_info.size());
   auto transfer_type = block_transfer_info[0].transfer_type;
   for (const BlockTransferInfo info : block_transfer_info) {
-    if (info.hash_key == nullptr) {
-      LOG(ERROR) << "Convert to BlockTransferInfos fail, hash key is nullptr!";
-      return false;
-    }
-
     if (transfer_type != info.transfer_type) {
       LOG(ERROR) << "Convert to BlockTransferInfos fail, TransferType must be "
                     "same, but got "