refactor: change host KV cache memory layout from layer-wise to block-wise.

Author: Kang-Meng

xllm/core/framework/kv_cache/kv_cache_store.cpp

@@ -43,29 +43,24 @@ bool KVCacheStore::init(const StoreConfig& config,
   }
   client_ptr_ = client_opt.value();
 
-  auto key_tensor_one_layer = host_kv_caches_->at(0).get_k_cache();
-  auto value_tensor_one_layer = host_kv_caches_->at(0).get_v_cache();
+  auto k_tensor_one_block = host_kv_caches_->at(0).get_k_cache();
+  auto v_tensor_one_block = host_kv_caches_->at(0).get_v_cache();
 
-  key_cache_size_per_layer_ =
-      key_tensor_one_layer[0].numel() * key_tensor_one_layer[0].element_size();
-  value_cache_size_per_layer_ = value_tensor_one_layer[0].numel() *
-                                value_tensor_one_layer[0].element_size();
+  k_cache_size_per_block_ =
+      k_tensor_one_block.numel() * k_tensor_one_block.element_size();
+  v_cache_size_per_block_ =
+      v_tensor_one_block.numel() * v_tensor_one_block.element_size();
 
-  auto key_cache_host_size =
-      key_tensor_one_layer.numel() * key_tensor_one_layer.element_size();
-  auto value_cache_host_size =
-      value_tensor_one_layer.numel() * value_tensor_one_layer.element_size();
-
-  LOG(INFO) << "key_cache_size_per_layer: " << key_cache_size_per_layer_;
-  LOG(INFO) << "value_cache_size_per_layer: " << value_cache_size_per_layer_;
+  LOG(INFO) << "k_cache_size_per_block: " << k_cache_size_per_block_;
+  LOG(INFO) << "v_cache_size_per_block: " << v_cache_size_per_block_;
 
   if (config_.protocol == "rdma") {
-    for (int layer = 0; layer < host_kv_caches_->size(); layer++) {
+    for (int block = 0; block < host_kv_caches_->size(); block++) {
       void* key_cache = static_cast<char*>(
-          host_kv_caches_->at(layer).get_k_cache().data_ptr());
+          host_kv_caches_->at(block).get_k_cache().data_ptr());
 
       auto register_k_result = client_ptr_->RegisterLocalMemory(
-          key_cache, key_cache_host_size, "cpu:0", false, false);
+          key_cache, k_cache_size_per_block_, "cpu:0", false, false);
 
       if (!register_k_result.has_value()) {
         LOG(ERROR) << "Failed to register local memory for key cache: "
@@ -74,10 +69,10 @@ bool KVCacheStore::init(const StoreConfig& config,
       }
 
       void* value_cache = static_cast<char*>(
-          host_kv_caches_->at(layer).get_v_cache().data_ptr());
+          host_kv_caches_->at(block).get_v_cache().data_ptr());
 
       auto register_v_result = client_ptr_->RegisterLocalMemory(
-          value_cache, value_cache_host_size, "cpu:0", false, false);
+          value_cache, v_cache_size_per_block_, "cpu:0", false, false);
 
       if (!register_v_result.has_value()) {
         LOG(ERROR) << "Failed to register local memory for value cache: "
@@ -119,23 +114,14 @@ uint32_t KVCacheStore::batch_put(
 
     str_keys.emplace_back(str_key);
 
-    std::vector<mooncake::Slice> slice;
-    slice.reserve(host_kv_caches_->size() * 2);
-    for (int layer = 0; layer < host_kv_caches_->size(); layer++) {
-      void* key_cache =
-          static_cast<char*>(
-              host_kv_caches_->at(layer).get_k_cache().data_ptr()) +
-          block_info.dst_block_id * key_cache_size_per_layer_;
-      slice.emplace_back(mooncake::Slice{key_cache, key_cache_size_per_layer_});
-
-      void* value_cache =
-          static_cast<char*>(
-              host_kv_caches_->at(layer).get_v_cache().data_ptr()) +
-          block_info.dst_block_id * value_cache_size_per_layer_;
-      slice.emplace_back(
-          mooncake::Slice{value_cache, value_cache_size_per_layer_});
-    }
-    slices.emplace_back(std::move(slice));
+    void* k_cache =
+        host_kv_caches_->at(block_info.dst_block_id).get_k_cache().data_ptr();
+    void* v_cache =
+        host_kv_caches_->at(block_info.dst_block_id).get_k_cache().data_ptr();
+
+    slices.emplace_back(std::vector<mooncake::Slice>{
+        mooncake::Slice{k_cache, k_cache_size_per_block_},
+        mooncake::Slice{v_cache, v_cache_size_per_block_}});
   }
 
   if (str_keys.size() == 0) {
@@ -177,24 +163,16 @@ uint32_t KVCacheStore::batch_get(
 
     str_keys.emplace_back(str_key);
 
-    slices.insert(std::make_pair(str_key, std::vector<mooncake::Slice>()));
-
-    slices[str_key].reserve(host_kv_caches_->size() * 2);
-    for (int layer = 0; layer < host_kv_caches_->size(); layer++) {
-      void* key_cache =
-          static_cast<char*>(
-              host_kv_caches_->at(layer).get_k_cache().data_ptr()) +
-          block_info.dst_block_id * key_cache_size_per_layer_;
-      slices[str_key].emplace_back(
-          mooncake::Slice{key_cache, key_cache_size_per_layer_});
-
-      void* value_cache =
-          static_cast<char*>(
-              host_kv_caches_->at(layer).get_v_cache().data_ptr()) +
-          block_info.dst_block_id * value_cache_size_per_layer_;
-      slices[str_key].emplace_back(
-          mooncake::Slice{value_cache, value_cache_size_per_layer_});
-    }
+    void* k_cache =
+        host_kv_caches_->at(block_info.dst_block_id).get_k_cache().data_ptr();
+    void* v_cache =
+        host_kv_caches_->at(block_info.dst_block_id).get_k_cache().data_ptr();
+
+    slices.insert(
+        std::make_pair(str_key,
+                       std::vector<mooncake::Slice>{
+                           mooncake::Slice{k_cache, k_cache_size_per_block_},
+                           mooncake::Slice{v_cache, v_cache_size_per_block_}}));
   }
 
   if (str_keys.size() == 0) {

xllm/core/framework/kv_cache/kv_cache_store.h

@@ -69,8 +69,8 @@ class KVCacheStore {
 
   std::vector<xllm::KVCache>* host_kv_caches_;
 
-  uint64_t key_cache_size_per_layer_;
-  uint64_t value_cache_size_per_layer_;
+  uint64_t k_cache_size_per_block_;
+  uint64_t v_cache_size_per_block_;
 
   std::shared_ptr<mooncake::Client> client_ptr_;
 };

xllm/core/framework/request/sequence.h

@@ -239,10 +239,10 @@ class Sequence final {
 
   void sync_result() {
     if (futures_.has_value()) {
-      auto success_cnt = host_kv_state_.num_kv_blocks();
+      uint32_t success_cnt = host_kv_state_.num_kv_blocks();
       for (auto& future : futures_.value()) {
         if (future.isReady()) {
-          success_cnt = std::min(success_cnt, size_t(future.value()));
+          success_cnt = std::min(success_cnt, future.value());
         } else {
           return;
         }

xllm/core/runtime/worker_impl.cpp

@@ -132,17 +132,19 @@ bool WorkerImpl::allocate_host_kv_cache(
 
   CHECK(model_ != nullptr) << "Model is not initialized.";
   CHECK(host_kv_caches_.empty()) << "KV caches are already initialized.";
+  CHECK(device_kv_cache_shape[0][0] == device_kv_cache_shape[1][0]);
 
   std::vector<std::vector<int64_t>> host_kv_cache_shape = device_kv_cache_shape;
-  host_kv_cache_shape[0][0] =
+  const int64_t num_layers = context_.get_model_args().n_layers();
+  int64_t host_bolck_size =
       device_kv_cache_shape[0][0] * options_.host_blocks_factor();
-  host_kv_cache_shape[1][0] =
-      device_kv_cache_shape[1][0] * options_.host_blocks_factor();
+  host_kv_cache_shape[0][0] = num_layers;
+  host_kv_cache_shape[1][0] = num_layers;
 
-  // create a KVCache for each layer
-  const int64_t num_layers = context_.get_model_args().n_layers();
-  host_kv_caches_.reserve(num_layers);
-  for (int64_t i = 0; i < num_layers; ++i) {
+  // create a KVCache shape: block_size * [layers, token, head, dim]
+  host_kv_caches_.reserve(host_bolck_size);
+
+  for (int64_t i = 0; i < host_bolck_size; ++i) {
     torch::Tensor key_cache, value_cache;
     key_cache = torch::empty(host_kv_cache_shape[0],
                              torch::dtype(dtype_).device(torch::kCPU))
@@ -152,8 +154,7 @@ bool WorkerImpl::allocate_host_kv_cache(
                       .pin_memory();
     host_kv_caches_.emplace_back(key_cache, value_cache);
   }
-  LOG(INFO) << "Initializing host k cache size: " << host_kv_cache_shape[0][0];
-  LOG(INFO) << "Initializing host v cache size: " << host_kv_cache_shape[1][0];
+  LOG(INFO) << "Initializing host kv block size: " << host_bolck_size;
 
   int32_t device_id = device_.index();
   h2d_attrs_.dstLoc.id = device_id;
@@ -688,22 +689,8 @@ uint32_t WorkerImpl::transfer_kv_blocks(
 
   switch (block_transfer_info[0].transfer_type) {
     case TransferType::G2H: {
-      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));
-          });
-
-      try {
-        auto timeout = std::chrono::seconds(KVSTORE_TIMEOUT);
-        return std::move(future).wait(timeout);
-      } catch (const folly::FutureTimeout& e) {
-        LOG(WARNING) << "BatchGet operation timed out";
-        return 0;
-      }
+      Slice<BlockTransferInfo> info_slice{block_transfer_info};
+      return load_from_store(info_slice);
     }
     case TransferType::D2G:
       return offload_kv_blocks(block_transfer_info);
@@ -793,23 +780,7 @@ uint32_t WorkerImpl::offload_kv_blocks(
                                promise = std::move(promise),
                                slice = std::move(slice)]() mutable {
       bool ret = d2h_batch_copy(slice);
-      uint32_t success_cnt = 0;
-
-      folly::Promise<uint32_t> store_promise;
-      auto future = store_promise.getSemiFuture();
-
-      batchput_threadpool_.schedule(
-          [this, &slice, promise = std::move(store_promise)]() mutable {
-            promise.setValue(KVCacheStore::get_instance().batch_put(slice));
-          });
-
-      try {
-        auto timeout = std::chrono::seconds(KVSTORE_TIMEOUT);
-        success_cnt = std::move(future).wait(timeout);
-      } catch (const folly::FutureTimeout& e) {
-        LOG(WARNING) << "BatchPut operation timed out";
-      }
-
+      auto success_cnt = offload_to_store(slice);
       if (success_cnt != slice.size()) {
         LOG(WARNING) << "KVCacheStore not all put success: " << success_cnt
                      << "/" << slice.size();
@@ -895,6 +866,7 @@ bool WorkerImpl::d2h_batch_copy(Slice<BlockTransferInfo>& block_transfer_info) {
 #if defined(USE_NPU)
   CHECK(copy_stream_.count(std::this_thread::get_id()) != 0)
       << "WorkerImpl::d2h_batch_copy can only be called in copy_threadpool_.";
+
   const int64_t num_layers = context_.get_model_args().n_layers();
   uint32_t num_batches = block_transfer_info.size() * num_layers * 2;
   void** srcs = new void*[num_batches];
@@ -904,26 +876,25 @@ bool WorkerImpl::d2h_batch_copy(Slice<BlockTransferInfo>& block_transfer_info) {
   size_t attrs_indexes[1] = {0};
   size_t fail_index;
   uint32_t curr_index = 0;
-  for (int layer_id = 0; layer_id < num_layers; layer_id++) {
-    auto src_k_cache = kv_caches_.at(layer_id).get_k_cache();
-    auto dst_k_cache = host_kv_caches_.at(layer_id).get_k_cache();
-    auto src_v_cache = kv_caches_.at(layer_id).get_v_cache();
-    auto dst_v_cache = host_kv_caches_.at(layer_id).get_v_cache();
-
-    for (int idx = 0; idx < block_transfer_info.size(); idx++) {
-      srcs[curr_index] =
-          src_k_cache[block_transfer_info[idx].src_block_id].data_ptr();
-      dsts[curr_index] =
-          dst_k_cache[block_transfer_info[idx].dst_block_id].data_ptr();
 
+  for (const auto& info : block_transfer_info) {
+    auto dst_k_cache = host_kv_caches_.at(info.dst_block_id).get_k_cache();
+    auto dst_v_cache = host_kv_caches_.at(info.dst_block_id).get_v_cache();
+
+    for (int layer_id = 0; layer_id < num_layers; layer_id++) {
+      auto src_k_cache = kv_caches_.at(layer_id).get_k_cache();
+      auto src_v_cache = kv_caches_.at(layer_id).get_v_cache();
+
+      srcs[curr_index] = src_k_cache[info.src_block_id].data_ptr();
+      dsts[curr_index] = dst_k_cache[layer_id].data_ptr();
       copy_size[curr_index] = key_cache_size_per_layer_;
+
       curr_index++;
 
-      srcs[curr_index] =
-          src_v_cache[block_transfer_info[idx].src_block_id].data_ptr();
-      dsts[curr_index] =
-          dst_v_cache[block_transfer_info[idx].dst_block_id].data_ptr();
+      srcs[curr_index] = src_v_cache[info.src_block_id].data_ptr();
+      dsts[curr_index] = dst_v_cache[layer_id].data_ptr();
       copy_size[curr_index] = value_cache_size_per_layer_;
+
       curr_index++;
     }
   }
@@ -961,6 +932,7 @@ bool WorkerImpl::h2d_batch_copy(Slice<BlockTransferInfo>& block_transfer_info) {
 #if defined(USE_NPU)
   CHECK(copy_stream_.count(std::this_thread::get_id()) != 0)
       << "WorkerImpl::h2d_batch_copy can only be called in copy_threadpool_.";
+
   const int64_t num_layers = context_.get_model_args().n_layers();
   uint32_t num_batches = block_transfer_info.size() * num_layers * 2;
   void** srcs = new void*[num_batches];
@@ -971,24 +943,21 @@ bool WorkerImpl::h2d_batch_copy(Slice<BlockTransferInfo>& block_transfer_info) {
   size_t fail_index;
   uint32_t curr_index = 0;
 
-  for (int layer_id = 0; layer_id < num_layers; layer_id++) {
-    auto src_k_cache = host_kv_caches_.at(layer_id).get_k_cache();
-    auto dst_k_cache = kv_caches_.at(layer_id).get_k_cache();
-    auto src_v_cache = host_kv_caches_.at(layer_id).get_v_cache();
-    auto dst_v_cache = kv_caches_.at(layer_id).get_v_cache();
-
-    for (int idx = 0; idx < block_transfer_info.size(); idx++) {
-      srcs[curr_index] =
-          src_k_cache[block_transfer_info[idx].src_block_id].data_ptr();
-      dsts[curr_index] =
-          dst_k_cache[block_transfer_info[idx].dst_block_id].data_ptr();
+  for (const auto& info : block_transfer_info) {
+    auto src_k_cache = host_kv_caches_.at(info.src_block_id).get_k_cache();
+    auto src_v_cache = host_kv_caches_.at(info.src_block_id).get_v_cache();
+
+    for (int layer_id = 0; layer_id < num_layers; layer_id++) {
+      auto dst_k_cache = kv_caches_.at(layer_id).get_k_cache();
+      auto dst_v_cache = kv_caches_.at(layer_id).get_v_cache();
+
+      srcs[curr_index] = src_k_cache[layer_id].data_ptr();
+      dsts[curr_index] = dst_k_cache[info.dst_block_id].data_ptr();
       copy_size[curr_index] = key_cache_size_per_layer_;
       curr_index++;
 
-      srcs[curr_index] =
-          src_v_cache[block_transfer_info[idx].src_block_id].data_ptr();
-      dsts[curr_index] =
-          dst_v_cache[block_transfer_info[idx].dst_block_id].data_ptr();
+      srcs[curr_index] = src_v_cache[layer_id].data_ptr();
+      dsts[curr_index] = dst_v_cache[info.dst_block_id].data_ptr();
       copy_size[curr_index] = value_cache_size_per_layer_;
       curr_index++;
     }
@@ -1022,4 +991,64 @@ bool WorkerImpl::h2d_batch_copy(Slice<BlockTransferInfo>& block_transfer_info) {
   return false;
 }
 
+uint32_t WorkerImpl::offload_to_store(
+    Slice<BlockTransferInfo>& block_transfer_info) {
+  if (!options_.enable_kvcache_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));
+      });
+
+  auto timeout = std::chrono::seconds(KVSTORE_TIMEOUT);
+  return std::move(future)
+      .via(folly::getGlobalCPUExecutor())
+      .within(timeout)
+      .thenTry([](folly::Try<uint32_t>&& t) -> uint32_t {
+        if (t.hasValue()) {
+          return t.value();
+        } else {
+          LOG(WARNING) << "BatchPut operation timed out";
+          return 0u;
+        }
+      })
+      .get();
+}
+
+uint32_t WorkerImpl::load_from_store(
+    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));
+      });
+
+  auto timeout = std::chrono::seconds(KVSTORE_TIMEOUT);
+  return std::move(future)
+      .via(folly::getGlobalCPUExecutor())
+      .within(timeout)
+      .thenTry([](folly::Try<uint32_t>&& t) -> uint32_t {
+        if (t.hasValue()) {
+          return t.value();
+        } else {
+          LOG(WARNING) << "BatchGet operation timed out";
+          return 0u;
+        }
+      })
+      .get();
+}
+
 }  // namespace xllm