[Misc] add store intf with tensor addr ptr (#288)

* add store intf with tensor addr ptr

* fix interface doxy

Author: mag1c-h

ucm/store/dramstore/dramstore_connector.py

@@ -165,6 +165,28 @@ def dump(
         logger.debug(f"dump block {block_ids} finished.")
         return task
 
+    def fetch_data(
+        self,
+        block_ids: List[str],
+        offset: List[int],
+        dst_addr: List[int],
+        size: List[int],
+    ) -> Task:
+        raise NotImplementedError(
+            "Method(fetch_data) not yet implemented in this version"
+        )
+
+    def dump_data(
+        self,
+        block_ids: List[str],
+        offset: List[int],
+        src_addr: List[int],
+        size: List[int],
+    ) -> Task:
+        raise NotImplementedError(
+            "Method(dump_data) not yet implemented in this version"
+        )
+
     def wait(self, task: DramTask) -> int:
         """
         wait kv cache kv transfer task finished.

ucm/store/mooncakestore/mooncake_connector.py

@@ -259,6 +259,28 @@ async def _dump_impl(
                 raise TypeError("Mooncake Store Put Type Error.") from err
         return 0
 
+    def fetch_data(
+        self,
+        block_ids: List[str],
+        offset: List[int],
+        dst_addr: List[int],
+        size: List[int],
+    ) -> Task:
+        raise NotImplementedError(
+            "Method(fetch_data) not yet implemented in this version"
+        )
+
+    def dump_data(
+        self,
+        block_ids: List[str],
+        offset: List[int],
+        src_addr: List[int],
+        size: List[int],
+    ) -> Task:
+        raise NotImplementedError(
+            "Method(dump_data) not yet implemented in this version"
+        )
+
     def wait(self, task: Task) -> int:
         """
         wait kv cache kv transfer task finished.

ucm/store/nfsstore/nfsstore_connector.py

@@ -88,6 +88,26 @@ def dump(
         )
         return NfsTask(task_id=task_id)
 
+    def fetch_data(
+        self,
+        block_ids: List[str],
+        offset: List[int],
+        dst_addr: List[int],
+        size: List[int],
+    ) -> Task:
+        task_id = self.store.LoadToDevice(block_ids, offset, dst_addr, size)
+        return NfsTask(task_id=task_id)
+
+    def dump_data(
+        self,
+        block_ids: List[str],
+        offset: List[int],
+        src_addr: List[int],
+        size: List[int],
+    ) -> Task:
+        task_id = self.store.DumpFromDevice(block_ids, offset, src_addr, size)
+        return NfsTask(task_id=task_id)
+
     def wait(self, task: Task) -> int:
         return self.store.Wait(task.task_id)
 

ucm/store/ucmstore.py

@@ -113,7 +113,7 @@ def dump(
         self, block_ids: List[str], offset: List[int], src_tensor: List[torch.Tensor]
     ) -> Task:
         """
-        dump kv cache to device.
+        dump kv cache from device.
 
         Args:
             block_ids (List[str]): vLLM block hash.
@@ -124,6 +124,48 @@ def dump(
         """
         pass
 
+    @abstractmethod
+    def fetch_data(
+        self,
+        block_ids: List[str],
+        offset: List[int],
+        dst_addr: List[int],
+        size: List[int],
+    ) -> Task:
+        """
+        load kv cache data to device.
+
+        Args:
+            block_ids (List[str]): vLLM block hash.
+            offset(List[int]): tp > 1 scene
+            dst_addr: List[int]: device tensor addr ptr.
+            size: List[int]: device tensor size.
+        Returns:
+            task(Task).
+        """
+        pass
+
+    @abstractmethod
+    def dump_data(
+        self,
+        block_ids: List[str],
+        offset: List[int],
+        src_addr: List[int],
+        size: List[int],
+    ) -> Task:
+        """
+        dump kv cache data from device.
+
+        Args:
+            block_ids (List[str]): vLLM block hash.
+            offset(List[int]): tp > 1 scene
+            src_addr: List[int]: device tensor addr ptr.
+            size: List[int]: device tensor size.
+        Returns:
+            task(Task).
+        """
+        pass
+
     @abstractmethod
     def wait(self, task: Task) -> int:
         """