[DeepSeek R1] Upgrade Mooncake to latest (#1829)

## Purpose
Changes to adoption the latest Mooncake Store API

Author: jerrychenhf

pd_xpyd/setup_env.sh

@@ -5,8 +5,8 @@ pip install colorlog
 
 echo "2. setting up mooncake mooncake-transfer-engine private build............."
 #Mooncake
-wget https://github.com/hlin99/Mooncake/releases/download/private_buildv2/mooncake_transfer_engine-0.1.0-cp310-cp310-manylinux2014_x86_64.whl
-pip install mooncake_transfer_engine-0.1.0-cp310-cp310-manylinux2014_x86_64.whl --force-reinstall
+wget https://github.com/hlin99/Mooncake/releases/download/private_buildv3/mooncake_transfer_engine-0.3.5-cp310-cp310-manylinux_2_17_x86_64.whl
+pip install mooncake_transfer_engine-0.3.5-cp310-cp310-manylinux_2_17_x86_64.whl --force-reinstall
 
 echo "3. setting up RDMA for mooncake ..................."
 #RDMA

pd_xpyd/start_etcd_mooncake_master.sh

@@ -29,9 +29,4 @@ etcd --listen-client-urls http://0.0.0.0:2379 \
      --advertise-client-urls http://localhost:2379 \
      >etcd.log 2>&1 &
 
-if [ "$BENCHMARK_MODE" == "1" ]; then
-  mooncake_master -max_threads 64 -port 50001 --v=1 >mooncake_master.log 2>&1 &
-else
-  mooncake_master -enable_gc true -max_threads 64 -port 50001 --v=1 >mooncake_master.log 2>&1 &
-fi
-
+mooncake_master -max_threads 64 -port 50001 -eviction_high_watermark_ratio 0.8 -eviction_ratio 0.2 >mooncake_master.log 2>&1 &

vllm/distributed/kv_transfer/kv_connector/mooncake_store_connector.py

@@ -251,9 +251,9 @@ def send_kv_caches_and_hidden_states_cpu(
             store_kvcache_key = f"{store_key_prefix}_{self.rank}"
             store_hidden_key = f"{store_key_prefix}_hidden_{self.rank}"
 
-            self.kv_store.put_unsafe(store_kvcache_key,
+            self.kv_store.put_tensor(store_kvcache_key,
                                      kv_caches_send_list[idx])
-            self.kv_store.put_unsafe(store_hidden_key, hidden_states_list[idx])
+            self.kv_store.put_tensor(store_hidden_key, hidden_states_list[idx])
         logger.info("[rank %d]: KV send DONE. send %d, takes %f s", self.rank,
                     len(input_tokens_list),
                     time.time() - start_time)
@@ -313,12 +313,12 @@ def send_kv_caches_and_hidden_states_hpu(
             kvcache_to_sent = keys.cpu()
             logger.debug("kv cache reshape time: %s", time.time() - start_time)
             store_kvcache_key = f"{store_key_prefix}_{self.rank}"
-            self.kv_store.put_unsafe(store_kvcache_key, kvcache_to_sent)
+            self.kv_store.put_tensor(store_kvcache_key, kvcache_to_sent)
 
             logger.debug("put kv cache key: %s", store_kvcache_key)
 
             hidden_key = f"{store_key_prefix}_hidden_{self.rank}"
-            self.kv_store.put(
+            self.kv_store.put_tensor(
                 hidden_key,
                 hidden_or_intermediate_states[idx].unsqueeze(0).cpu())
             # ==== graph should end here ======
@@ -396,15 +396,14 @@ def recv_kv_caches_and_hidden_states_hpu(
             load_key_prefix = self.tensor_hash(current_tokens)
             # For deepseek, we only need recv first rank
             load_kvcache_key = f"{load_key_prefix}_0"
-            shape = (61, num_blocks * 128, self.k_v_head_size)
             remote_kv = None
             if self._wait_for_key(load_kvcache_key):
-                remote_kv = self.kv_store.get_unsafe(load_kvcache_key, shape,
+                remote_kv = self.kv_store.get_tensor(load_kvcache_key,
                                                      self.dtype)
             hidden_key = f"{load_key_prefix}_hidden_0"
             hidden = None
             if self._wait_for_key(hidden_key):
-                hidden = self.kv_store.get(hidden_key)
+                hidden = self.kv_store.get_tensor(hidden_key)
 
             if remote_kv is None or hidden is None:
                 # didn't find any match.
@@ -477,13 +476,12 @@ def recv_kv_caches_and_hidden_states_cpu(
         load_hidden_key = f"{prefix}_hidden_0"
         remote_kv = None
         if self._wait_for_key(load_kvcache_key):
-            remote_kv = self.kv_store.get_unsafe(load_kvcache_key,
-                                                 shape=None,
+            remote_kv = self.kv_store.get_tensor(load_kvcache_key,
                                                  dtype=self.dtype)
         # hidden_states always use bf16.
         hidden = None
         if self._wait_for_key(load_hidden_key):
-            hidden = self.kv_store.get_unsafe(load_hidden_key, shape=(1, 7168))
+            hidden = self.kv_store.get_tensor(load_hidden_key)
 
         if remote_kv is None or hidden is None:
             # didn't find any match.

vllm/distributed/kv_transfer/kv_lookup_buffer/mooncake_store.py

@@ -84,7 +84,7 @@ def __init__(
     ):
 
         try:
-            from mooncake import MooncakeDistributedStore
+            from mooncake.store import MooncakeDistributedStore
         except ImportError as e:
             raise ImportError(
                 "Please install mooncake by following the instructions at "
@@ -199,47 +199,42 @@ def _get_impl(
 
         return None
 
-    def put_unsafe(
+    def put_tensor(
         self,
         key: str,
         value: Optional[torch.Tensor],
     ) -> None:
-        """Put KVCache to Mooncake Store"""
+        """Put tensor to Mooncake Store"""
         value = value.cpu()
-        start_serde = time.time()
-        data_ptr = value.data_ptr()
-        element_size = value.element_size()
-        numel = value.numel()
-        total_size = element_size * numel
-        end_serde = time.time()
+        start_put = time.time()
         try:
-            self.store.put_unsafe(key, data_ptr, total_size)
+            self.store.put_tensor(key, value)
         except TypeError as err:
-            logger.error("Failed to put value into Mooncake Store: %s", err)
+            logger.error("Failed to put tensor into Mooncake Store: %s", err)
             raise TypeError("Mooncake Store Put Type Error.") from err
         end_put = time.time()
-        logger.debug("contiguous time: %f, put time: %f",
-                     end_serde - start_serde, end_put - end_serde)
+        logger.debug("Put tensor to store. Time: %f", end_put - start_put)
 
-    def get_unsafe(self,
+    def get_tensor(self,
                    key: str,
-                   shape,
                    dtype=torch.bfloat16) -> Optional[torch.Tensor]:
-        """Get KVCache from Mooncake Store without type checking"""
+        """Get tensor from Mooncake Store"""
         start_get = time.time()
-        data = self.store.get(key)
+        try:
+            value = self.store.get_tensor(key)
+        except TypeError as err:
+            logger.error("Failed to get tensor from Mooncake Store: %s", err)
+            raise TypeError("Mooncake Store Get Type Error.") from err
         end_get = time.time()
-        if data:
-            tensor = torch.frombuffer(data, dtype=dtype)
-            shape = (61, -1, 1, 576) if shape is None else shape
-            tensor = tensor.view(shape)
-            end_from_buffer = time.time()
-            logger.debug("from buffer time: %f , get time: %f",
-                         end_from_buffer - end_get, end_get - start_get)
-
-            return tensor
-        return None
+        if value is None:
+            logger.error("Failed to get tensor from Mooncake Store: %s", key)
+            return None
+        # This is a workaround for get_tensor which returns wrong tensor type
+        # Remove when Mooncake get_tensor fixed the issue
+        value = value.view(dtype)
+        logger.debug("Get tensor from store. Time: %f", end_get - start_get)
+        return value
 
     def is_exist(self, key: str) -> bool:
         """Check if the key exists in the Mooncake Store"""
-        return self.store.isExist(key) == 1
+        return self.store.is_exist(key) == 1