feat: optimize kv cache load/offload.

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)

setup.py

@@ -445,7 +445,7 @@ def check_and_install_pre_commit():
             print("Run 'pre-commit install' failed. Please install pre-commit: pip install pre-commit")
             exit(0)
 
-def run_git_command(command, cwd=None, check=True):
+def run_shell_command(command, cwd=None, check=True):
     try:
         subprocess.run(command, cwd=cwd, check=check, shell=True, capture_output=True, text=True)
         return True
@@ -492,15 +492,15 @@ def apply_patch_safely(patch_file_path, repo_path):
 
     if has_uncommitted_changes(repo_path):
         print(f"⚠️ Uncommitted changes detected. Running `git reset --hard` for {repo_path}")
-        if not run_git_command("git reset --hard", cwd=repo_path):
+        if not run_shell_command("git reset --hard", cwd=repo_path):
             print("❌ Failed to reset changes!")
             return False
 
     print(f"🛠️ Apply patch: {patch_file_path}")
-    apply_success = run_git_command(f"git apply --check {patch_file_path}", cwd=repo_path, check=False)
+    apply_success = run_shell_command(f"git apply --check {patch_file_path}", cwd=repo_path, check=False)
 
     if apply_success:
-        if not run_git_command(f"git apply {patch_file_path}", cwd=repo_path):
+        if not run_shell_command(f"git apply {patch_file_path}", cwd=repo_path):
             print("❌ apply patch fail!")
             apply_success = False
 
@@ -512,7 +512,7 @@ def apply_patch_safely(patch_file_path, repo_path):
         print(f"  cd {repo_path} && git apply {patch_file_path}")
         return False
 
-def apply_patch():
+def pre_build():
     if os.path.exists("third_party/custom_patch"):
         script_path = os.path.dirname(os.path.abspath(__file__))
         mooncake_repo_path = os.path.join(script_path, "third_party/Mooncake")
@@ -521,6 +521,9 @@ def apply_patch():
         cpprestsdk_repo_path = os.path.join(script_path, "third_party/cpprestsdk")
         if not apply_patch_safely("../custom_patch/cpprestsdk.patch", cpprestsdk_repo_path):
             exit(0)
+        if not run_shell_command("sh third_party/dependencies.sh", cwd=script_path):
+            print("❌ Failed to reset changes!")
+            exit(0)
 
 if __name__ == "__main__":
     device = 'a2'  # default
@@ -537,9 +540,10 @@ def apply_patch():
             del sys.argv[idx]
             del sys.argv[idx]
     if '--dry_run' not in sys.argv:
-        apply_patch()
+        pre_build()
     else:
         sys.argv.remove("--dry_run") 
+
     if '--install-xllm-kernels' in sys.argv:
         idx = sys.argv.index('--install-xllm-kernels')
         if idx + 1 < len(sys.argv):

third_party/custom_patch/Mooncake.patch

@@ -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/dependencies.sh

@@ -0,0 +1,96 @@
+#!/bin/bash
+
+# Color definitions
+GREEN="\033[0;32m"
+BLUE="\033[0;34m"
+YELLOW="\033[0;33m"
+RED="\033[0;31m"
+NC="\033[0m" # No Color
+
+# Configuration
+REPO_ROOT=`pwd`
+
+# Function to print section headers
+print_section() {
+    echo -e "\n${BLUE}=== $1 ===${NC}"
+}
+
+# Function to print success messages
+print_success() {
+    echo -e "${GREEN}✓ $1${NC}"
+}
+
+# Function to print error messages and exit
+print_error() {
+    echo -e "${RED}✗ ERROR: $1${NC}"
+    exit 1
+}
+
+# Function to check command success
+check_success() {
+    if [ $? -ne 0 ]; then
+        print_error "$1"
+    fi
+}
+
+if [ $(id -u) -ne 0 ]; then
+	print_error "Require root permission, try sudo ./dependencies.sh"
+fi
+
+
+# Install yalantinglibs
+print_section "Installing yalantinglibs"
+
+# Check if thirdparties directory exists
+if [ ! -d "${REPO_ROOT}/third_party/Mooncake/thirdparties" ]; then
+    mkdir -p "${REPO_ROOT}/third_party/Mooncake/thirdparties"
+    check_success "Failed to create Mooncake/thirdparties directory"
+fi
+
+# Change to thirdparties directory
+cd "${REPO_ROOT}/third_party/Mooncake/thirdparties"
+check_success "Failed to change to Mooncake/thirdparties directory"
+
+# Check if yalantinglibs is already installed
+if [ -d "yalantinglibs" ]; then
+    echo -e "${YELLOW}yalantinglibs directory already exists. Removing for fresh install...${NC}"
+    rm -rf yalantinglibs
+    check_success "Failed to remove existing yalantinglibs directory"
+fi
+
+# Clone yalantinglibs
+echo "Cloning yalantinglibs from https://github.com/alibaba/yalantinglibs.git"
+git clone https://github.com/alibaba/yalantinglibs.git
+check_success "Failed to clone yalantinglibs"
+
+# Build and install yalantinglibs
+cd yalantinglibs
+check_success "Failed to change to yalantinglibs directory"
+
+# Checkout version 0.5.5
+echo "Checking out yalantinglibs version 0.5.5..."
+git checkout 0.5.5
+check_success "Failed to checkout yalantinglibs version 0.5.5"
+
+mkdir -p build
+check_success "Failed to create build directory"
+
+cd build
+check_success "Failed to change to build directory"
+
+echo "Configuring yalantinglibs..."
+cmake .. -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DBUILD_UNIT_TESTS=OFF -DYLT_ENABLE_IBV=ON 
+check_success "Failed to configure yalantinglibs"
+
+echo "Building yalantinglibs (using $(nproc) cores)..."
+cmake --build . -j$(nproc)
+check_success "Failed to build yalantinglibs"
+
+echo "Installing yalantinglibs..."
+cmake --install .
+check_success "Failed to install yalantinglibs"
+
+sed -i '54s/target_link_libraries(${ylt_target_name} -libverbs)/target_link_libraries(${ylt_target_name} INTERFACE -libverbs)/' /usr/local/lib/cmake/yalantinglibs/config.cmake
+
+print_success "yalantinglibs installed successfully"
+

xllm/core/common/global_flags.cpp

@@ -332,14 +332,22 @@ DEFINE_string(store_protocol,
               "tcp",
               "KV cache store protocol(e.g. tcp, rdma).");
 
-DEFINE_string(store_master_server_entry,
+DEFINE_string(store_master_server_address,
               "",
               "The address information of the store master service.");
 
-DEFINE_string(store_metadata_connstring,
+DEFINE_string(store_metadata_server,
               "",
               "The address of the kv cache store metadata service.");
 
+DEFINE_string(store_local_hostname,
+              "",
+              "The local host name of the kv cache store client.");
+
+DEFINE_uint64(store_segment_size,
+              0,
+              "The contributed memory to kv cache store.");
+
 // --- computation communication parallel config ---
 
 DEFINE_bool(

xllm/core/common/global_flags.h

@@ -159,9 +159,13 @@ DECLARE_bool(enable_kvcache_store);
 
 DECLARE_string(store_protocol);
 
-DECLARE_string(store_master_server_entry);
+DECLARE_string(store_master_server_address);
 
-DECLARE_string(store_metadata_connstring);
+DECLARE_string(store_metadata_server);
+
+DECLARE_string(store_local_hostname);
+
+DECLARE_uint64(store_segment_size);
 
 DECLARE_bool(enable_multi_stream_parallel);
 

xllm/core/common/options.cpp

@@ -53,8 +53,10 @@ std::string Options::to_string() const {
      << ", enable_cache_upload: " << enable_cache_upload()
      << ", enable_kvcache_store: " << enable_kvcache_store()
      << ", store_protocol: " << store_protocol()
-     << ", store_master_server_entry: " << store_master_server_entry()
-     << ", store_metadata_connstring: " << store_metadata_connstring()
+     << ", store_master_server_address: " << store_master_server_address()
+     << ", store_metadata_server: " << store_metadata_server()
+     << ", store_local_hostname: " << store_local_hostname()
+     << ", store_segment_size: " << store_segment_size()
      << ", enable_multi_stream_parallel: " << enable_multi_stream_parallel()
      << ", enable_continuous_kvcache: " << enable_continuous_kvcache()
      << ", disable_ttft_profiling: " << disable_ttft_profiling()

xllm/core/common/options.h

@@ -143,9 +143,13 @@ class Options {
 
   PROPERTY(std::string, store_protocol) = "tcp";
 
-  PROPERTY(std::string, store_master_server_entry) = "";
+  PROPERTY(std::string, store_master_server_address) = "";
 
-  PROPERTY(std::string, store_metadata_connstring) = "";
+  PROPERTY(std::string, store_metadata_server) = "";
+
+  PROPERTY(std::string, store_local_hostname) = "";
+
+  PROPERTY(uint64_t, store_segment_size) = 0;
 
   PROPERTY(bool, enable_multi_stream_parallel) = false;