Add batched zero copy to NIXL backend (#18850)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

python/sglang/srt/managers/cache_controller.py

@@ -472,7 +472,7 @@ class HiCacheController:
             self.page_get_func = self._generic_page_get
             self.page_set_func = self._generic_page_set
 
-            if (self.storage_backend_type in ["hf3fs", "mooncake", "eic"]) or (
+            if (self.storage_backend_type in ["hf3fs", "mooncake", "eic", "nixl"]) or (
                 self.storage_backend_type == "dynamic"
                 and bool(self.storage_config.extra_config.get("interface_v1", 0))
             ):

python/sglang/srt/mem_cache/storage/backend_factory.py

@@ -161,7 +161,7 @@ class StorageBackendFactory:
         if backend_name == "file":
             return backend_class(storage_config)
         elif backend_name == "nixl":
-            return backend_class(storage_config)
+            return backend_class(storage_config, mem_pool_host)
         elif backend_name == "mooncake":
             backend = backend_class(storage_config, mem_pool_host)
             return backend

python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py

@@ -6,7 +6,12 @@ from typing import Any, List, Optional, Union
 import torch
 
 from sglang.srt.environ import envs
-from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig
+from sglang.srt.mem_cache.hicache_storage import (
+    HiCacheStorage,
+    HiCacheStorageConfig,
+    HiCacheStorageExtraInfo,
+)
+from sglang.srt.mem_cache.memory_pool_host import HostKVCache
 
 from .nixl_utils import (
     NixlBackendConfig,
@@ -52,14 +57,17 @@ class HiCacheNixl(HiCacheStorage):
         )
 
         # Initialize suffix based on storage config
-        tp_rank, tp_size, model_name, is_mla_model = (
+        tp_rank, tp_size, model_name = (
             storage_config.tp_rank,
             storage_config.tp_size,
             storage_config.model_name,
-            storage_config.is_mla_model,
         )
+
+        self.is_mla_model = storage_config.is_mla_model
+
         model_name = "-".join(model_name.split("/")) if model_name else ""
-        if is_mla_model:
+
+        if self.is_mla_model:
             self.config_suffix = f"_{model_name}"
         else:
             self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}"
@@ -133,11 +141,21 @@ class HiCacheNixl(HiCacheStorage):
             ]
             storage_tuples = [(x[0], s, x[2]) for x, s in zip(tuples, tensor_sizes)]
             host_descs = self.agent.get_xfer_descs(buffers)
+
+            if direction in ("READ", "WRITE"):
+                # register buffer to avoid calling initialize_xfer twice due to missing registration
+                self.register_buffers(buffers)
+
         elif isinstance(buffers[0], tuple):
             storage_tuples = [(x[0], y[1], x[2]) for x, y in zip(tuples, buffers)]
             host_descs = self.agent.get_xfer_descs(
                 [(x[0], x[1], 0) for x in buffers], "DRAM"
             )
+
+            if direction in ("READ", "WRITE"):
+                # register buffer to avoid calling initialize_xfer twice due to missing registration
+                self.register_buffers(buffers)
+
         else:
             return False
 
@@ -150,6 +168,7 @@ class HiCacheNixl(HiCacheStorage):
             return False
 
         # Initialize transfer, default assumption that tensor was registered
+
         try:
             xfer_req = self.agent.initialize_xfer(
                 direction, host_descs, storage_descs, self.agent_name
@@ -220,6 +239,7 @@ class HiCacheNixl(HiCacheStorage):
         if target_sizes and (len(target_sizes) != len(target_locations)):
             logger.error("Mismatch between number of target_locations and target_sizes")
             return [None] * len(keys)
+
         if target_sizes:
             dest = list(zip(target_locations, target_sizes))
         else:
@@ -277,21 +297,326 @@ class HiCacheNixl(HiCacheStorage):
         else:  # mem_type == "OBJ"
             return self._execute_transfer(values, suffixed_keys, "WRITE")
 
-    def exists(self, key: str) -> bool:
-        # Add suffix to key
-        suffixed_key = self._get_suffixed_key(key)
+    ############################################################################
+    # batch_*_v1 functions
+    # zero copy + non-zero-copy version for get, set, exists, batch_exists
+    ############################################################################
 
-        tuples = self.registration.create_query_tuples(
-            suffixed_key,
-            self.backend_selector.mem_type,
-            self.file_manager if self.backend_selector.mem_type == "FILE" else None,
+    def clear(self) -> None:
+        self.file_manager.clear()
+
+    def register_mem_pool_host(self, mem_pool_host: HostKVCache):
+        super().register_mem_pool_host(mem_pool_host)
+
+        # enable zero-copy automatically if mem layout is page_first or page_first_direct
+        self.is_zero_copy = self.mem_pool_host.layout in [
+            "page_first",
+            "page_first_direct",
+        ]
+
+        logger.info(
+            f"HiCacheNixl: Registered mem_pool_host with layout {self.mem_pool_host.layout}, zero_copy set to {self.is_zero_copy}"
         )
-        if not tuples:
-            return False
+
+    def exists(self, key: str) -> bool:
+        results = self.batch_exists([key])
+        return results > 0
+
+    def batch_exists(
+        self,
+        keys: List[str],
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> int:
+        # Add suffix to key
+
+        if self.is_zero_copy:
+            key_list = self._get_key_list_from_meta(keys)
+            key_denominator = (
+                1 if not self.is_mla_model else 2
+            )  # MLA model only has k buffer, no separate v buffer
+        else:
+            key_list = [self._get_suffixed_key(key) for key in keys]
+            key_denominator = 1
+
+        # obtain list of tuples by calling self.registration.create_query_tuples()
+        tuples = []
+        for key in key_list:
+            tuples += self.registration.create_query_tuples(
+                key,
+                self.backend_selector.mem_type,
+                self.file_manager if self.backend_selector.mem_type == "FILE" else None,
+            )
 
         query_res = self.agent.query_memory(
             tuples,
             self.backend_selector.backend_name,
             mem_type=self.backend_selector.mem_type,
         )
-        return query_res[0] is not None  # can be expanded to multiple keys
+
+        for i in range(len(query_res)):
+            if query_res[i] is None:
+                return i // key_denominator
+        return len(query_res) // key_denominator
+
+    def _get_key_list_from_meta(self, keys: List[str]) -> List[str]:
+        # construct the key list for NIXL transfer based on the keys and the suffix, for each key, we will have one suffixed key for k buffer and one suffixed key for v buffer if it's not an MLA model, and only one suffixed key for k buffer if it's an MLA model, since MLA model only has k/v interleaved buffer
+        key_list = []
+
+        for key_ in keys:
+            suffixed_key = self._get_suffixed_key(key_)
+            if self.is_mla_model:
+                key_list.append(f"{suffixed_key}_k")
+            else:
+                key_list.append(f"{suffixed_key}_k")
+                key_list.append(f"{suffixed_key}_v")
+
+        return key_list
+
+    def _get_location_and_size_list_from_meta(
+        self, keys: List[str], host_indices: torch.Tensor
+    ):
+        # zero copy: mem_pool_host.get_data_page() does not work due to non-contiguous tensors, causing issues for NIXL transfer
+        ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(
+            host_indices
+        )
+        key_list = self._get_key_list_from_meta(keys)
+
+        if len(key_list) != len(ptr_list):
+            logger.error(
+                f"HiCacheNixl: mismatch between number of keys and number of buffer meta entries, keys: {len(keys)}, key_list: {len(key_list)}, buffer meta entries: {len(ptr_list)}"
+            )
+            return [], [], [], []
+
+        return key_list, [], ptr_list, element_size_list
+
+    def _batch_get_preprocess(self, keys: List[str], host_indices: torch.Tensor):
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+
+        if len(keys) == 0 or len(keys) != page_num:
+            logger.warning(
+                f"HiCacheNixl: empty keys or mismatch in keys and host_indices lengths. keys: {len(keys)}, host_indices: {len(host_indices)}, page_size: {self.mem_pool_host.page_size}"
+            )
+            return [], [], [], []
+
+        if self.is_zero_copy:
+            key_list, _, ptr_list, element_size_list = (
+                self._get_location_and_size_list_from_meta(keys, host_indices)
+            )
+            return key_list, [], ptr_list, element_size_list
+        else:
+            # non zero copy: create contiguous, temporary tensors
+            target_tensors = [
+                self.mem_pool_host.get_dummy_flat_data_page() for i in range(page_num)
+            ]
+
+            key_list = [self._get_suffixed_key(key) for key in keys]
+            ptr_list = [tensor.data_ptr() for tensor in target_tensors]
+            element_size_list = [
+                tensor.numel() * tensor.element_size() for tensor in target_tensors
+            ]
+
+            return key_list, target_tensors, ptr_list, element_size_list
+
+    def _batch_get_zero_copy_impl(
+        self,
+        keys: List[str],
+        key_strs: List[str],
+        target_tensors: List[torch.Tensor],
+        target_locations: List[int],
+        target_sizes: List[int],
+    ) -> List[int]:
+
+        if not key_strs or not target_locations or not target_sizes:
+            return [False] * len(keys)
+
+        if (len(key_strs) != len(target_locations)) or (
+            len(target_sizes) != len(target_locations)
+        ):
+            logger.error(
+                "Mismatch between number of key_strs, target_locations and target_sizes"
+            )
+            return [False] * len(keys)
+
+        if self.is_zero_copy:
+            dest = list(zip(target_locations, target_sizes))
+        else:
+            dest = target_tensors
+
+        if self.backend_selector.mem_type == "FILE":
+            file_paths = [self.file_manager.get_file_path(key) for key in key_strs]
+            success = self._execute_transfer(dest, file_paths, "READ")
+        else:
+            success = self._execute_transfer(dest, key_strs, "READ")
+
+        return [True] * len(key_strs) if success else [False] * len(key_strs)
+
+    def _batch_get_postprocess(
+        self,
+        host_indices: torch.Tensor,
+        target_tensors: List[torch.Tensor],
+        results: List[bool],
+    ) -> List[bool]:
+
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+
+        if self.is_zero_copy:
+            # zero copy: update final results based on the boolean results from NIXL transfer
+            if self.is_mla_model:
+                return results
+            else:
+                results = [
+                    (results[2 * i] and results[2 * i + 1]) for i in range(page_num)
+                ]
+                return results
+        else:
+            # non zero copy: copy data from temporary tensors to mem_pool_host page by page
+            for i in range(page_num):
+                if not results[i]:
+                    break
+                self.mem_pool_host.set_from_flat_data_page(
+                    host_indices[i * self.mem_pool_host.page_size], target_tensors[i]
+                )
+
+            return results
+
+    def batch_get_v1(
+        self,
+        keys: List[str],
+        host_indices: torch.Tensor,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> List[bool]:
+
+        key_strs, target_tensors, buffer_ptrs, buffer_sizes = (
+            self._batch_get_preprocess(keys, host_indices)
+        )
+
+        if not key_strs or not buffer_ptrs or not buffer_sizes:
+            logger.error(
+                "HiCacheNixl batch_get_v1: preprocessing failed, empty key_strs, buffer_ptrs or buffer_sizes"
+            )
+            return [False] * len(keys)
+
+        start_time = time.perf_counter()
+
+        results_get = self._batch_get_zero_copy_impl(
+            keys, key_strs, target_tensors, buffer_ptrs, buffer_sizes
+        )
+
+        end_time = time.perf_counter()
+        elapsed_time_ms = (end_time - start_time) * 1000
+        total_bytes = sum(s for s in buffer_sizes if s is not None)
+
+        logger.debug(
+            f"HiCacheNixl batch_get_v1 transferred: {len(keys)} keys (pages), {host_indices.numel()} host_indices, {total_bytes} bytes, total time: {elapsed_time_ms:.3f} ms, effective bandwidth: {total_bytes / (elapsed_time_ms / 1000) / (1024 * 1024):.2f} MB/s"
+        )
+
+        return self._batch_get_postprocess(host_indices, target_tensors, results_get)
+
+    def _batch_set_preprocess(self, keys: List[str], host_indices: torch.Tensor):
+
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+
+        if len(keys) == 0 or len(keys) != page_num:
+            logger.warning(
+                f"HiCacheNixl: empty keys or mismatch in keys and host_indices lengths. keys: {len(keys)}, host_indices: {len(host_indices)}, page_size: {self.mem_pool_host.page_size}"
+            )
+            return [], [], [], []
+
+        if self.is_zero_copy:
+            key_list, _, ptr_list, element_size_list = (
+                self._get_location_and_size_list_from_meta(keys, host_indices)
+            )
+            return key_list, [], ptr_list, element_size_list
+        else:
+            # non zero copy: NIXL still requires contiguous tensors for transfer
+            target_tensors = [
+                self.mem_pool_host.get_data_page(
+                    host_indices[i * self.mem_pool_host.page_size], flat=False
+                ).contiguous()
+                for i in range(page_num)
+            ]
+
+            key_list = [self._get_suffixed_key(key) for key in keys]
+            ptr_list = [tensor.data_ptr() for tensor in target_tensors]
+            element_size_list = [
+                tensor.numel() * tensor.element_size() for tensor in target_tensors
+            ]
+
+            return key_list, target_tensors, ptr_list, element_size_list
+
+    def _batch_set_zero_copy_impl(
+        self,
+        keys: List[str],
+        key_strs: List[str],
+        target_tensors: List[torch.Tensor],
+        target_locations: List[int],
+        target_sizes: List[int],
+    ) -> List[bool]:
+
+        if not key_strs or not target_locations or not target_sizes:
+            return [False] * len(keys)
+
+        if (len(key_strs) != len(target_locations)) or (
+            len(target_sizes) != len(target_locations)
+        ):
+            logger.error(
+                "Mismatch between number of key_strs, target_locations and target_sizes"
+            )
+            return [False] * len(keys)
+
+        if self.is_zero_copy:
+            src = list(zip(target_locations, target_sizes))
+        else:
+            src = target_tensors
+
+        if self.backend_selector.mem_type == "FILE":
+            file_paths = []
+            for key in key_strs:
+                file_path = self.file_manager.get_file_path(key)
+                # New file per set, to be updated when partial writes is added to HiCache
+                if not self.file_manager.create_file(file_path):
+                    logger.error(
+                        f"******** Failed to create file {file_path} *********"
+                    )
+                    return [False] * len(keys)
+                file_paths.append(file_path)
+            success = self._execute_transfer(src, file_paths, "WRITE")
+        else:  # mem_type == "OBJ"
+            success = self._execute_transfer(src, key_strs, "WRITE")
+
+        return [True] * len(keys) if success else [False] * len(keys)
+
+    def batch_set_v1(
+        self,
+        keys: List[str],
+        host_indices: torch.Tensor,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> List[bool]:
+
+        if len(keys) == 0:
+            return []
+
+        key_strs, target_tensors, buffer_ptrs, buffer_sizes = (
+            self._batch_set_preprocess(keys, host_indices)
+        )
+
+        if not key_strs or not buffer_ptrs or not buffer_sizes:
+            logger.error(
+                "HiCacheNixl batch_set_v1: preprocessing failed, empty key_strs, buffer_ptrs or buffer_sizes"
+            )
+            return [False] * len(keys)
+
+        start_time = time.perf_counter()
+
+        results_set = self._batch_set_zero_copy_impl(
+            keys, key_strs, target_tensors, buffer_ptrs, buffer_sizes
+        )
+
+        end_time = time.perf_counter()
+        elapsed_time_ms = (end_time - start_time) * 1000
+        total_bytes = sum(s for s in buffer_sizes if s is not None)
+        logger.debug(
+            f"HiCacheNixl batch_set_v1 transferred: {len(keys)} keys (pages), {host_indices.numel()} host_indices, {total_bytes} bytes, total time: {elapsed_time_ms:.3f} ms, effective bandwidth: {total_bytes / (elapsed_time_ms / 1000) / (1024 * 1024):.2f} MB/s"
+        )
+
+        return results_set

python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py

@@ -63,7 +63,7 @@ class NixlBackendConfig:
             config_data = self.config
 
         for key, value in config_data.items():
-            initparams[key] = value
+            initparams[key] = str(value)
 
         return initparams
 
@@ -234,6 +234,22 @@ class NixlFileManager:
             os.makedirs(base_dir, exist_ok=True)
             logger.debug(f"Initialized file manager with base directory: {base_dir}")
 
+    def clear(self) -> None:
+        """Clear all files in the base directory."""
+        if self.base_dir == "":
+            logger.warning("Base directory is empty, skipping clear operation")
+            return
+
+        try:
+            for root, dirs, files in os.walk(self.base_dir):
+                for file in files:
+                    os.remove(os.path.join(root, file))
+            logger.debug(f"Cleared all files in base directory: {self.base_dir}")
+        except Exception as e:
+            logger.error(
+                f"Failed to clear files in base directory {self.base_dir}: {e}"
+            )
+
     def get_file_path(self, key: str) -> str:
         """Get full file path for a given key."""
         return os.path.join(self.base_dir, key)