feat: upgrade Mooncake to v0.3.6.

Author: Kang-Meng

CMakeLists.txt

@@ -3,6 +3,9 @@ set_property(GLOBAL PROPERTY USE_FOLDERS ON)
 
 option(USE_NPU "Enable NPU support" OFF)
 option(USE_MLU "Enable MLU support" OFF)
+add_compile_definitions(YLT_ENABLE_IBV)
+add_definitions(-DYLT_ENABLE_IBV)
+set(YLT_ENABLE_IBV ON)
 
 if(DEVICE_ARCH STREQUAL "ARM")
     set(CMAKE_SYSTEM_PROCESSOR aarch64)

third_party/Mooncake

@@ -26,7 +26,7 @@ index 8483085..9d263dd 100644
  add_subdirectory(benchmarks)
 \ No newline at end of file
 diff --git a/mooncake-store/include/offset_allocator/offset_allocator.hpp b/mooncake-store/include/offset_allocator/offset_allocator.hpp
-index fde978b..ac54f8c 100644
+index b6d55c8..2d80158 100644
 --- a/mooncake-store/include/offset_allocator/offset_allocator.hpp
 +++ b/mooncake-store/include/offset_allocator/offset_allocator.hpp
 @@ -6,7 +6,7 @@
@@ -38,40 +38,24 @@ index fde978b..ac54f8c 100644
 
  namespace mooncake::offset_allocator {
  typedef unsigned char uint8;
-diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h
-index a30f0c6..3bfeffb 100644
---- a/mooncake-store/include/storage_backend.h
-+++ b/mooncake-store/include/storage_backend.h
-@@ -4,8 +4,8 @@
- #include <vector>
- #include <mutex>
- #include <fstream>
--#include <types.h>
--#include <file_interface.h>
-+#include "types.h"
-+#include "file_interface.h"
- #include <filesystem>
- #include <thread>
- #include <chrono>
 diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h
-index d2830a3..39c3b27 100644
+index 077926b..b36862b 100644
 --- a/mooncake-store/include/types.h
 +++ b/mooncake-store/include/types.h
-@@ -10,8 +10,11 @@
- #include <variant>
+@@ -8,7 +8,10 @@
+ #include <unordered_map>
  #include <vector>
 
 -#include "Slab.h"
 +#include "cachelib_memory_allocator/Slab.h"
- #include "allocator.h"
 +namespace iguana {
 +using std::contiguous_iterator;
 +}
  #include "ylt/struct_json/json_reader.h"
  #include "ylt/struct_json/json_writer.h"
 
 diff --git a/mooncake-store/src/ha_helper.cpp b/mooncake-store/src/ha_helper.cpp
-index c1f4ded..e4b982b 100644
+index 796838a..8e72b80 100644
 --- a/mooncake-store/src/ha_helper.cpp
 +++ b/mooncake-store/src/ha_helper.cpp
 @@ -1,3 +1,6 @@
@@ -81,7 +65,7 @@ index c1f4ded..e4b982b 100644
  #include "ha_helper.h"
  #include "etcd_helper.h"
  #include "rpc_service.h"
-@@ -169,4 +172,6 @@ MasterServiceSupervisor::~MasterServiceSupervisor() {
+@@ -174,4 +177,6 @@ MasterServiceSupervisor::~MasterServiceSupervisor() {
      }
  }
 
@@ -113,7 +97,7 @@ index f515671..7a019b6 100644
 +#pragma GCC pop_options
 \ No newline at end of file
 diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp
-index b775851..736d61d 100644
+index 9678f57..f41eb10 100644
 --- a/mooncake-store/src/utils.cpp
 +++ b/mooncake-store/src/utils.cpp
 @@ -1,6 +1,6 @@

third_party/custom_patch/Mooncake.patch

@@ -247,7 +247,7 @@ bool BlockManagerPool::allocate(Sequence* sequence, size_t num_tokens) {
     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].push_back(
+      load_block_transfer_infos_[dp_rank].emplace_back(
           BlockTransferInfo(host_blocks[i].id(),
                             hbm_blocks[i].id(),
                             host_blocks[i].get_immutable_hash_value(),

xllm/core/framework/block/block_manager_pool.cpp

@@ -16,30 +16,30 @@ bool KVCacheStore::init(const StoreConfig& config,
   CHECK(!is_initialized_) << "KVCacheStore is initialized.";
   config_ = config;
   host_kv_caches_ = host_kv_caches;
+  std::optional<std::string> device_names = std::nullopt;
   if (config_.protocol == "rdma") {
-    if (getenv("DEVICE_NAME")) {
-      auto name = getenv("DEVICE_NAME");
-      LOG(INFO) << "device name: " << name;
-      args_ = mooncake::rdma_args(name);
+    if (getenv("DEVICE_NAMES")) {
+      device_names = getenv("DEVICE_NAMES");
+      LOG(INFO) << "device_names: " << device_names.value();
     } else {
       LOG(WARNING) << "env DEVICE_NAME not exist, set protocol as tcp";
       config_.protocol = "tcp";
-      args_ = nullptr;
     }
   }
 
   auto client_opt = mooncake::Client::Create(config_.localhost_name,
                                              config_.metadata_connstring,
                                              config_.protocol,
-                                             args_,
+                                             device_names,
                                              config_.master_server_entry);
 
   rep_config_.replica_num = config_.replica_num;
   // rep_config_.preferred_segment = config_.localhost_name;
 
   if (!client_opt.has_value()) {
-    LOG(FATAL) << "mooncake::Client::Create fail!";
-    return false;
+    LOG(FATAL) << "mooncake::Client::Create fail! Failed to create client with "
+                  "host_name: "
+               << config_.localhost_name;
   }
   client_ptr_ = client_opt.value();
 

xllm/core/framework/kv_cache/kv_cache_store.cpp

@@ -67,8 +67,6 @@ class KVCacheStore {
   StoreConfig config_;
   mooncake::ReplicateConfig rep_config_;
 
-  void** args_ = nullptr;
-
   std::vector<xllm::KVCache>* host_kv_caches_;
 
   uint64_t key_cache_size_per_layer_;