From 443b1a88d105d6d3a322cc9436eb9651089d00f5 Mon Sep 17 00:00:00 2001 From: hxie Date: Wed, 18 Feb 2026 16:31:02 -0800 Subject: [PATCH] Add batched zero copy to NIXL backend (#18850) Co-authored-by: Zhiqiang Xie --- .../sglang/srt/managers/cache_controller.py | 2 +- .../srt/mem_cache/storage/backend_factory.py | 2 +- .../mem_cache/storage/nixl/hicache_nixl.py | 353 +++++++++++++++++- .../srt/mem_cache/storage/nixl/nixl_utils.py | 18 +- 4 files changed, 358 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 845dffe96..34f78c889 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/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)) ): diff --git a/python/sglang/srt/mem_cache/storage/backend_factory.py b/python/sglang/srt/mem_cache/storage/backend_factory.py index eaa5a3e18..2eb6afe24 100644 --- a/python/sglang/srt/mem_cache/storage/backend_factory.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py index abb26c4ed..78addf514 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py b/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py index 960ff5cf1..9742cf3f3 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py +++ b/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)