diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index e5e6d7189..e222daa0a 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -837,6 +837,16 @@ SGLang supports various environment variables that can be used to configure its Decode-side incremental KV cache offload stride. Rounded down to a multiple of --page-size (min is --page-size). If unset/invalid/<=0, it falls back to --page-size. Not set (uses --page-size) + + SGLANG_HICACHE_NIXL_USE_DIRECT_IO + Enable O_DIRECT for any file-based NIXL backend (POSIX, GDS, GDS_MT, 3FS) when opening cache files (bypasses the OS page cache, reducing memory pressure and improving throughput on NVMe). Can also be disabled via {'{"use_direct_io": false}'} in --hicache-storage-backend-extra-config. Falls back to buffered I/O with a warning when O_DIRECT is unavailable on the current OS. + true + + + SGLANG_HUGEPAGE_SIZE + Use huge pages for host KV cache allocations (HiCache / disaggregation offload). Valid values: 2MB (2 MiB pages via MAP_HUGE_2MB) or 1GB (1 GiB pages via MAP_HUGE_1GB). Requires huge pages to be pre-allocated on the host OS (/proc/sys/vm/nr_hugepages or /sys/kernel/mm/hugepages). If the allocation fails, the allocator logs a warning and falls back to regular page-size mmap automatically. + Not set (uses OS default page size) + diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9742bf441..86eb37e8f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -342,6 +342,11 @@ class Envs: SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None) SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None) SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None) + # Enable O_DIRECT when opening NIXL POSIX backend files (bypasses OS page cache). + # Disable with SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0 or via the + # "use_direct_io": false key in --hicache-storage-backend-extra-config. + SGLANG_HICACHE_NIXL_USE_DIRECT_IO = EnvBool(True) + SGLANG_HUGEPAGE_SIZE = EnvStr("") # Staging buffer for heterogeneous TP KV transfer SGLANG_DISAGG_STAGING_BUFFER = EnvBool(False) SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB = EnvInt(64) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index cbee102fe..65ef06ab3 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, List, NamedTuple, Optional import torch from sglang.srt.mem_cache.hicache_storage import ( + STORAGE_BATCH_SIZE, HiCacheStorageConfig, HiCacheStorageExtraInfo, ) @@ -507,8 +508,6 @@ class HiCacheController: self.prefetch_capacity_limit = max( 0, int(0.8 * (self.mem_pool_host.size - self.mem_pool_device.size)) ) - # granularity of batch storage IO operations, in number of pages - self.storage_batch_size = 128 # tracking the number of tokens locked in prefetching, updated by the main scheduler thread self.prefetch_tokens_occupied = 0 @@ -915,8 +914,8 @@ class HiCacheController: def _page_transfer(self, operation): # Transfer batch by batch prefix_keys = operation.prefix_keys - for i in range(0, len(operation.hash_value), self.storage_batch_size): - batch_hashes = operation.hash_value[i : i + self.storage_batch_size] + for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE): + batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE] batch_host_indices = operation.host_indices[ i * self.page_size : (i + len(batch_hashes)) * self.page_size ] @@ -978,11 +977,9 @@ class HiCacheController: hash_value = [] for start in range( - 0, len(tokens_to_fetch), self.page_size * self.storage_batch_size + 0, len(tokens_to_fetch), self.page_size * STORAGE_BATCH_SIZE ): - end = min( - start + self.page_size * self.storage_batch_size, len(tokens_to_fetch) - ) + end = min(start + self.page_size * STORAGE_BATCH_SIZE, len(tokens_to_fetch)) batch_tokens = tokens_to_fetch[start:end] batch_hashes = [] for i in range(0, len(batch_tokens), self.page_size): @@ -1120,8 +1117,8 @@ class HiCacheController: def _page_backup(self, operation): # Backup batch by batch prefix_keys = operation.prefix_keys - for i in range(0, len(operation.hash_value), self.storage_batch_size): - batch_hashes = operation.hash_value[i : i + self.storage_batch_size] + for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE): + batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE] batch_host_indices = operation.host_indices[ i * self.page_size : (i + len(batch_hashes)) * self.page_size ] diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index b384d4cdc..effbdde8e 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -16,6 +16,9 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +# Max pages per batched storage IO call. +STORAGE_BATCH_SIZE = 128 + @dataclass class HiCacheStorageConfig: diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index d5d9d5655..913918e16 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -37,6 +37,7 @@ from sglang.srt.mem_cache.memory_pool import ( MHATokenToKVPool, MLATokenToKVPool, ) +from sglang.srt.mem_cache.mmap_allocator import alloc_mmap from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu _is_cuda = is_cuda() @@ -78,18 +79,19 @@ def synchronized(func): return wrapper -class HostTensorAllocator(abc.ABC): +class HostTensorAllocator: def __init__(self): """Initialize the HostTensorAllocator.""" self.dtype = None self.dims = None def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor: - """Allocate a tensor of given dims and dtype on the memory.""" + assert ( + device == "cpu" + ), f"HostTensorAllocator only supports CPU allocations; got device={device!r}" self.dtype = dtype self.dims = dims - tensor = torch.empty(dims, dtype=dtype, device=device) - return tensor + return alloc_mmap(dims, dtype) class HiSparseHostPoolMixin: @@ -187,11 +189,16 @@ def alloc_with_host_register( """ buffer = allocator.allocate(dims, dtype=dtype, device=device) if pin_memory: - ret = torch.cuda.cudart().cudaHostRegister( - buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0 - ) - if ret != 0: - raise RuntimeError(f"cudaHostRegister failed with error code {ret}") + cudart = torch.cuda.cudart() + n_bytes = buffer.numel() * buffer.element_size() + rc = cudart.cudaHostRegister(buffer.data_ptr(), n_bytes, 0) + if int(rc) != 0: + raise RuntimeError( + f"cudaHostRegister failed (rc={int(rc)}, " + f"{cudart.cudaGetErrorString(rc)}) for ptr={buffer.data_ptr():#x} " + f"size={n_bytes}; host buffer is not pinned and device transfers " + f"may silently return stale data." + ) return buffer @@ -324,6 +331,19 @@ class HostKVCache(abc.ABC): """ raise NotImplementedError() + def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool: + """Return True if per-page strides are multiples of *page_size_bytes*. + + Subclasses should override this with a layout-specific stride formula. + This base implementation logs a warning and returns False (safe default). + """ + logger.warning( + "%s does not implement is_stride_page_aligned(); assuming not aligned. " + "O_DIRECT with a file-based NIXL backend will fall back to copy mode for this pool.", + type(self).__name__, + ) + return False + @synchronized def clear(self): # Initialize memory states and tracking structures. @@ -850,6 +870,30 @@ class MHATokenToKVPoolHost(HostKVCache): raise ValueError(f"Unsupported layout: {self.layout}") return ptr_list, element_size_list + def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool: + """Return True if per-page strides are multiples of *page_size_bytes*. + + When O_DIRECT is used with any file-based NIXL backend, every data pointer + passed to the kernel must be page-aligned. In zero-copy mode the + pointer for KV page ``p`` is: + + base_ptr + p * page_size * layer_num * head_num * head_dim * itemsize + + For this to be page-aligned (given a page-aligned ``base_ptr``) the per-page + stride must itself be a multiple of the OS page size. + """ + if self.layout not in ("page_first", "page_first_direct", "page_head"): + return False + stride = ( + self.page_size + * self.layer_num + * self.head_num + * self.head_dim + * self.dtype.itemsize + ) + base_aligned = self.kv_buffer.data_ptr() % page_size_bytes == 0 + return base_aligned and stride % page_size_bytes == 0 + class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): device_pool: MLATokenToKVPool @@ -1250,6 +1294,26 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): raise ValueError(f"Unsupported layout: {self.layout}") return ptr_list, element_size_list + def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool: + """Return True if per-page strides are multiples of *page_size_bytes*. + + When O_DIRECT is used with any file-based NIXL backend, every data pointer + passed to the kernel must be page-aligned. In zero-copy mode the + pointer for KV page ``p`` is: + + base_ptr + p * page_size * layer_num * kv_cache_dim * itemsize + + For this to be page-aligned (given a page-aligned ``base_ptr``) the per-page + stride must itself be a multiple of the OS page size. + """ + if self.layout not in ("page_first", "page_first_direct"): + return False + stride = ( + self.page_size * self.layer_num * self.kv_cache_dim * self.dtype.itemsize + ) + base_aligned = self.kv_buffer.data_ptr() % page_size_bytes == 0 + return base_aligned and stride % page_size_bytes == 0 + class MambaPoolHost(HostKVCache): diff --git a/python/sglang/srt/mem_cache/mmap_allocator.py b/python/sglang/srt/mem_cache/mmap_allocator.py new file mode 100644 index 000000000..f2aeeabfa --- /dev/null +++ b/python/sglang/srt/mem_cache/mmap_allocator.py @@ -0,0 +1,127 @@ +import ctypes +import ctypes.util +import logging +import math +import mmap +import os +import weakref + +import torch + +from sglang.srt.environ import envs + +logger = logging.getLogger(__name__) + +# Load libc once at module level so munmap is callable safely at GC/shutdown time. +# Resolve the SONAME via find_library so the allocator also works on systems +# whose libc is not named "libc.so.6" (e.g. musl / Alpine). +try: + _libc_name = ctypes.util.find_library("c") or "libc.so.6" + _libc = ctypes.CDLL(_libc_name, use_errno=True) + _libc.mmap.restype = ctypes.c_void_p + _libc.mmap.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_long, + ] + _libc.munmap.restype = ctypes.c_int + _libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t] +except OSError: + _libc = None + +# MAP_POPULATE is in Python's mmap module only since 3.11. +_MAP_POPULATE = getattr(mmap, "MAP_POPULATE", 0x08000) +# MAP_HUGETLB and MAP_HUGE_* are Linux-specific and not in Python's mmap module. +_MAP_HUGETLB = 0x40000 +_MAP_HUGE_2MB = 21 << 26 # 0x1400000 +_MAP_HUGE_1GB = 30 << 26 # 0x78000000 +_MAP_FAILED = ctypes.c_void_p(-1).value + + +def _alloc_hugepage(n_bytes: int, alloc_bytes: int, extra_flags: int) -> ctypes.Array: + """Call mmap via libc with hugepage flags and return an owning ctypes array. + + munmap fires automatically via weakref.finalize when the array is + garbage-collected (i.e. when the tensor that wraps it is freed). + """ + ptr = _libc.mmap( + None, + alloc_bytes, + mmap.PROT_READ | mmap.PROT_WRITE, + mmap.MAP_SHARED | mmap.MAP_ANONYMOUS | _MAP_POPULATE | extra_flags, + -1, + 0, + ) + if ptr is None or ptr == _MAP_FAILED: + errno = ctypes.get_errno() + raise OSError(errno, os.strerror(errno)) + array = (ctypes.c_uint8 * n_bytes).from_address(ptr) + weakref.finalize(array, _libc.munmap, ctypes.c_void_p(ptr), alloc_bytes) + return array + + +def alloc_mmap(dims: tuple, dtype: torch.dtype) -> torch.Tensor: + """Allocate a host tensor via anonymous mmap. Set SGLANG_HUGEPAGE_SIZE=2MB or 1GB for hugepages. + + MAP_SHARED + MAP_POPULATE are both required so cudaHostRegister pins real, + pre-faulted physical pages (otherwise pinning can race with COW or page + faults and the device ends up reading stale data). + + The tensor owns the mapping; munmap fires when the tensor is freed. + """ + # Re-read per call (not cached) so that envs.SGLANG_HUGEPAGE_SIZE.override() + # works correctly in tests. + hugepage_size = (envs.SGLANG_HUGEPAGE_SIZE.get() or "").strip().upper() + n_bytes = math.prod(dims) * torch.empty([], dtype=dtype).element_size() + + if hugepage_size == "": + page_size, extra_flags = mmap.PAGESIZE, 0 + elif hugepage_size == "2MB": + page_size, extra_flags = 2 * 1024 * 1024, _MAP_HUGETLB | _MAP_HUGE_2MB + elif hugepage_size == "1GB": + page_size, extra_flags = 1024 * 1024 * 1024, _MAP_HUGETLB | _MAP_HUGE_1GB + else: + logger.warning( + "Unrecognized SGLANG_HUGEPAGE_SIZE=%r; expected '2MB' or '1GB'. " + "Falling back to plain page-size mmap.", + envs.SGLANG_HUGEPAGE_SIZE.get(), + ) + page_size, extra_flags = mmap.PAGESIZE, 0 + + alloc_bytes = math.ceil(n_bytes / page_size) * page_size + + if extra_flags: + if _libc is None: + logger.error( + "Hugepage mmap requested but libc.so.6 could not be loaded; " + "falling back to plain mmap. SGLANG_HUGEPAGE_SIZE=%s will be ignored.", + hugepage_size, + ) + else: + try: + array = _alloc_hugepage(n_bytes, alloc_bytes, extra_flags) + return torch.frombuffer( + array, dtype=dtype, count=math.prod(dims) + ).reshape(dims) + except OSError as e: + logger.error( + "Hugepage mmap via libc failed (%s); falling back to plain mmap. " + "SGLANG_HUGEPAGE_SIZE=%s will be ignored.", + e, + hugepage_size, + ) + alloc_bytes = math.ceil(n_bytes / mmap.PAGESIZE) * mmap.PAGESIZE + + # Plain mmap path -- used directly when no hugepages requested, or as fallback. + # torch.frombuffer keeps a reference to mm inside the tensor storage, so mm + # stays alive until the tensor is freed and mmap.mmap.__del__ calls munmap. + mm = mmap.mmap( + -1, + alloc_bytes, + flags=mmap.MAP_SHARED | mmap.MAP_ANONYMOUS | _MAP_POPULATE, + prot=mmap.PROT_READ | mmap.PROT_WRITE, + ) + return torch.frombuffer(mm, dtype=dtype, count=math.prod(dims)).reshape(dims) diff --git a/python/sglang/srt/mem_cache/storage/nixl/README.md b/python/sglang/srt/mem_cache/storage/nixl/README.md index 8d1c542d9..7885f979f 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/README.md +++ b/python/sglang/srt/mem_cache/storage/nixl/README.md @@ -45,8 +45,10 @@ The main storage connector that provides: Consolidated utility classes: - **NixlBackendSelection** - Handles backend selection and creation - **NixlBackendConfig** - Handles backend configuration -- **NixlRegistration** - Manages memory registration for tensors, files and objects -- **NixlFileManager** - Handles file system operations and NIXL tuple creation +- **NixlFileManager** - Handles file system operations + +### NixlRegistry (`nixl_registry.py`) +Owns the `(agent, mem_type, file_manager)` triple and exposes `host(...)` and `storage(...)` context managers that register on entry, yield the NIXL `xfer_descs`, and deregister + close fds on exit. Internally composes two single-resource primitives (`_open_files` and `_registered`) so leak-freeness is verifiable per primitive. The current implementation performs per-transfer registration for file / object targets and explicitly closes FILE descriptors after registration / transfer setup to avoid descriptor leaks. @@ -149,7 +151,7 @@ For debugging or quick testing, you may pass a **JSON-style string** directly vi This requires explicitly specifying the plugin type via an environment variable, and this method can be applicable to **only a few** plugins (e.g., POSIX, GDS, GDS_MT) -The below example shows how to use command-line string to use the POSIX plugin where URING is enabled for async POSIX storage. +The below example shows how to use command-line string to use the POSIX plugin where URING is enabled for async POSIX storage, with O_DIRECT enabled (the default). ```bash export SGLANG_HICACHE_NIXL_BACKEND_PLUGIN=POSIX @@ -164,7 +166,17 @@ python3 -m sglang.launch_server \ --hicache-size 64 \ --hicache-write-policy write_through \ --hicache-storage-backend nixl \ - --hicache-storage-backend-extra-config "{'use_uring': 'true'}" + --hicache-storage-backend-extra-config '{"use_uring": "true"}' +``` + +To disable O_DIRECT (e.g. for debugging or unsupported filesystems), set the top-level `use_direct_io` key: + +```bash +export SGLANG_HICACHE_NIXL_BACKEND_PLUGIN=POSIX + +python3 -m sglang.launch_server \ + ... \ + --hicache-storage-backend-extra-config '{"use_direct_io": false, "use_uring": "true"}' ``` ⚠️ **Note**: @@ -349,6 +361,33 @@ An example of the configuration is provided in [`nixl.config.toml.sample`](./nix For object storage, `bucket` may also be omitted from the config if `AWS_DEFAULT_BUCKET` is already defined in the environment. +### 1a. Top-Level Configuration Keys + +The following keys are placed at the **top level** of the config file (not inside any `[plugin.*]` section) and apply globally to the NIXL backend: + +| Key | Type | Default | Description | +| ---------------- | ------- | -------- | ----------- | +| `use_direct_io` | boolean | `true` | Open cache files with `O_DIRECT` to bypass the OS page cache. Reduces memory pressure and improves NVMe throughput. Falls back to buffered I/O with a warning if `O_DIRECT` is unavailable on the current OS. Can also be overridden via the `SGLANG_HICACHE_NIXL_USE_DIRECT_IO` environment variable. | + +**Page-alignment and `O_DIRECT`** + +When `use_direct_io = true` with any file-based backend (POSIX, GDS, GDS_MT, 3FS), the kernel requires every I/O buffer pointer to be OS-page-aligned (4 KiB). SGLang handles this automatically: + +* **Zero-copy mode** (`page_first` / `page_first_direct` layout): the host memory pool is always mmap-backed and therefore page-aligned. If the per-page stride is also a multiple of 4 KiB, zero-copy transfers are used as-is. +* **Copy mode** (all other layouts, or if stride alignment cannot be satisfied): SGLang pre-allocates page-aligned bounce buffers via `mmap` and falls back to copy mode, logging a warning. No user action is required -- this is fully automatic. + +To disable `O_DIRECT` (e.g. for debugging or when the filesystem does not support it): + +```toml +use_direct_io = false + +[plugin.posix] +use_uring = "true" +active = true +``` + +or via environment variable: `SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0`. + ### 2. POSIX File System Backend (`plugin.posix`) 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 5b961076a..45088e0fc 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py +++ b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py @@ -1,27 +1,25 @@ import logging import time import uuid -from typing import Any, List, Optional, Union +from typing import Any, List, Optional import torch from sglang.srt.environ import envs from sglang.srt.mem_cache.hicache_storage import ( + STORAGE_BATCH_SIZE, HiCacheStorage, HiCacheStorageConfig, HiCacheStorageExtraInfo, ) from sglang.srt.mem_cache.memory_pool_host import HostKVCache +from sglang.srt.mem_cache.mmap_allocator import alloc_mmap -from .nixl_utils import ( - NixlBackendConfig, - NixlBackendSelection, - NixlFileManager, - NixlRegistration, -) +from .nixl_registry import NixlRegistry +from .nixl_utils import NixlBackendConfig, NixlBackendSelection, NixlFileManager try: - from nixl._api import nixl_agent, nixl_agent_config + from nixl._api import nixl_agent, nixl_agent_config, nixlBind except ImportError as e: raise ImportError( "Please install NIXL by following the instructions at " @@ -48,15 +46,16 @@ class HiCacheNixl(HiCacheStorage): # select the NIXL backend plugin from extra_config or environment variable plugin = nixlconfig.get_specified_plugin() + use_direct_io = nixlconfig.get_use_direct_io() + # Might be better to be unified across HiCache backends and moved to HiCacheController file_path = envs.SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR.get() or file_path self.file_manager = ( - NixlFileManager(file_path) + NixlFileManager(file_path, use_direct_io=use_direct_io) if plugin not in NixlBackendSelection.OBJ_PLUGINS else None ) - # Initialize suffix based on storage config tp_rank, tp_size, model_name = ( storage_config.tp_rank, storage_config.tp_size, @@ -75,253 +74,156 @@ class HiCacheNixl(HiCacheStorage): else: self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}" + sync_mode = getattr( + nixlBind, "NIXL_THREAD_SYNC_RW", nixlBind.NIXL_THREAD_SYNC_STRICT + ) agent_config = nixl_agent_config(backends=[]) self.agent_name = f"hicache_nixl_{str(uuid.uuid4())}" self.agent = nixl_agent(self.agent_name, agent_config) + bind_cfg = nixlBind.nixlAgentConfig() + bind_cfg.useProgThread = agent_config.enable_pthread + bind_cfg.useListenThread = agent_config.enable_listen + bind_cfg.listenPort = agent_config.port + bind_cfg.syncMode = sync_mode + bind_cfg.pthrDelay = 0 + bind_cfg.lthrDelay = 100000 + bind_cfg.captureTelemetry = agent_config.capture_telemetry + self.agent.agent = nixlBind.nixlAgent(self.agent_name, bind_cfg) + self.agent.plugin_list = self.agent.agent.getAvailPlugins() self.backend_selector = NixlBackendSelection(plugin, nixlconfig) if not self.backend_selector.create_backend(self.agent): raise RuntimeError("Failed to create NIXL backend") - self.registration = NixlRegistration(self.agent) - self.is_zero_copy = False + self.registry = NixlRegistry( + self.agent, + self.backend_selector.mem_type, + self.file_manager, + ) + # O_DIRECT requires OS-page-aligned I/O buffers on all file-based backends + # (POSIX, GDS, GDS_MT, 3FS). OBJ backends never open files so they are exempt + # (file_manager is None for OBJ). + self.needs_page_alignment = use_direct_io and self.file_manager is not None + if self.needs_page_alignment: + logger.info( + "HiCacheNixl: O_DIRECT is active with a file-based backend (%s). " + "Page-aligned host buffers are required (needs_page_alignment=True).", + self.backend_selector.backend_name, + ) + # Pre-registered host regions (set by register_mem_pool_host): + # zero-copy: one registration covering mem_pool_host.kv_buffer + # non-zero-copy: two registrations, one bounce buffer per direction + # (set/get) so the two storage threads never share slots. + self._host_regs: List[Any] = [] + self._bounce_set: Optional[torch.Tensor] = None + self._bounce_get: Optional[torch.Tensor] = None + self._bounce_page_bytes: Optional[int] = None def _get_suffixed_key(self, key: str) -> str: return key + self.config_suffix - def register_buffers( - self, buffers: Union[torch.Tensor, List[torch.Tensor], List[tuple]] - ) -> Optional[Any]: - """Register tensor(s) or target locations in host memory (list of addr,len tuples) with NIXL.""" - if isinstance(buffers[0], tuple): - tuples = [(x[0], x[1], 0, "") for x in buffers] - return self.registration._register_memory(tuples, "DRAM") - else: - return self.registration._register_memory(buffers) + def _create_query_tuple(self, key: str) -> tuple: + """Build the NIXL query_memory tuple for a single key.""" + if self.backend_selector.mem_type == "FILE": + return (0, 0, 0, self.file_manager.get_file_path(key)) + return (0, 0, 0, key) - def register_files( - self, file_paths: List[str], open_file: Optional[bool] = True - ) -> Optional[Any]: - """Register files with NIXL.""" - tuples = self.file_manager.files_to_nixl_tuples(file_paths) - return self.registration._register_memory(tuples, "FILE") - - def register_objects( - self, keys: List[str], sizes: Optional[List[int]] = None - ) -> Optional[Any]: - """Register objects with NIXL.""" - if not keys: - return None - tuples = [(0, 0, key, "") for key in keys] - return self.registration._register_memory(tuples, "OBJ") - - def _execute_transfer( + def _xfer_and_wait( self, - buffers: Optional[List[torch.Tensor | tuple]], + host_descs: Any, + storage_descs: Any, + direction: str, + ) -> bool: + """Initialize and poll a NIXL transfer to completion.""" + try: + xfer_req = self.agent.initialize_xfer( + direction, host_descs, storage_descs, self.agent_name + ) + except Exception as e: + logger.error(f"Failed to create transfer request: {e}") + return False + + try: + state = self.agent.transfer(xfer_req) + while state != "DONE": + state = self.agent.check_xfer_state(xfer_req) + if state == "ERR": + logger.error("Transfer failed") + return False + # Best would be to have a better notification mechanism from NIXL, + # but we only have polling for now. + time.sleep(0.0001) + return True + except Exception as e: + logger.error(f"Failed to execute transfer: {e}") + import traceback + + logger.error(f"Traceback: {traceback.format_exc()}") + return False + finally: + self.agent.release_xfer_handle(xfer_req) + + def _xfer_pre_registered( + self, + host_buffers: List[tuple], keys: List[str], direction: str, ) -> bool: - if len(buffers) != len(keys): - logger.error("Mismatch between number of tensors/buffers and files/objects") + """Run a transfer where the host side is already pre-registered. + + ``host_buffers`` is a list of ``(addr, size)`` tuples within the + pre-registered host region (kv_buffer for zero-copy, bounce buffer + otherwise). Only the storage side is registered per transfer. + """ + if len(host_buffers) != len(keys): + logger.error("Mismatch between number of host buffers and keys") return False - # Registering file and object keys per transfer, to be updated when - # pre-registration for file and object is added to HiCache. - file_fds = [] - try: - if self.backend_selector.mem_type == "FILE": - tuples = self.file_manager.files_to_nixl_tuples(keys) - file_fds = [t[2] for t in tuples] - if not tuples or not self.registration._register_memory(tuples, "FILE"): - logger.error("Failed to prepare files for transfer") - return False - else: # mem_type == "OBJ" - tuples = [(0, 0, key, "") for key in keys] - if not tuples or not self.registration._register_memory(tuples, "OBJ"): - logger.error("Failed to register objects") - return False + host_descs = self.agent.get_xfer_descs( + [(addr, size, 0) for (addr, size) in host_buffers], "DRAM" + ) + if host_descs is None: + logger.error("Failed to build host xfer descs") + return False - # Prepare transfer descriptors - if isinstance(buffers[0], torch.Tensor): - tensor_sizes = [ - tensor.element_size() * tensor.numel() for tensor in buffers - ] - 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: + with self.registry.storage(host_buffers, keys, direction) as storage_descs: + if storage_descs is None: return False - - storage_descs = self.agent.get_xfer_descs( - storage_tuples, self.backend_selector.mem_type - ) - - if (host_descs is None) or (storage_descs is None): - logger.error("Failed to get transfer descriptors") - 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 - ) - except Exception: - # Check if it was due to missing pre-registration - if not self.register_buffers(buffers): - logger.error("Failed to register tensors/buffers") - return False - - try: - xfer_req = self.agent.initialize_xfer( - direction, host_descs, storage_descs, self.agent_name - ) - except Exception as e: - logger.error(f"Failed to create transfer request: {e}") - return False - - # Execute transfer and wait for its completion - try: - state = self.agent.transfer(xfer_req) - while state != "DONE": - state = self.agent.check_xfer_state(xfer_req) - if state == "ERR": - self.agent.release_xfer_handle(xfer_req) - logger.error("Transfer failed") - return False - time.sleep( - 0.0001 - ) # Can be changed to os.sched_yield() or parametrized - - self.agent.release_xfer_handle(xfer_req) - return True - - except Exception as e: - logger.error(f"Failed to execute transfer: {e}") - import traceback - - logger.error(f"Traceback: {traceback.format_exc()}") - return False - - finally: - for fd in file_fds: - self.file_manager.close_file(fd) + return self._xfer_and_wait(host_descs, storage_descs, direction) def get( self, key: str, - target_location: Optional[torch.Tensor | int] = None, - target_sizes: Optional[int] = None, + target_location: Optional[Any] = None, + target_sizes: Optional[Any] = None, ) -> torch.Tensor | None: - # To be removed, being compatible with the current API - if target_location is None: - return None - if target_sizes: - result = self.batch_get([key], [target_location], [target_sizes]) - else: - result = self.batch_get([key], [target_location]) - return result[0] if result else None + raise NotImplementedError("deprecated; use batch_get_v1") def batch_get( self, keys: List[str], - target_locations: Optional[List[torch.Tensor | int]] = None, - target_sizes: Optional[List[int]] = None, + target_locations: Optional[Any] = None, + target_sizes: Optional[Any] = None, ) -> List[torch.Tensor | None]: - if not keys: - return [] - - # To be removed, being compatible with the current API - if not target_locations: - return [None] * len(keys) - - 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: - dest = target_locations - - # Add suffix to keys - suffixed_keys = [self._get_suffixed_key(key) for key in keys] - - if self.backend_selector.mem_type == "FILE": - file_paths = [self.file_manager.get_file_path(key) for key in suffixed_keys] - success = self._execute_transfer(dest, file_paths, "READ") - else: - success = self._execute_transfer(dest, suffixed_keys, "READ") - return target_locations if success and not target_sizes else [None] * len(keys) + raise NotImplementedError("deprecated; use batch_get_v1") def set( self, key: str, - value: Optional[torch.Tensor] = None, - target_location: Optional[int] = None, - target_sizes: Optional[int] = None, + value: Optional[Any] = None, + target_location: Optional[Any] = None, + target_sizes: Optional[Any] = None, ) -> bool: - if target_location and target_sizes: - return self.batch_set([key], None, [target_location], [target_sizes]) - else: - return self.batch_set([key], [value]) + raise NotImplementedError("deprecated; use batch_set_v1") def batch_set( self, keys: List[str], - values: Optional[List[torch.Tensor]] = None, - target_locations: Optional[List[int]] = None, - target_sizes: Optional[List[int]] = None, + values: Optional[Any] = None, + target_locations: Optional[Any] = None, + target_sizes: Optional[Any] = None, ) -> bool: - - # skip on MLA backup rank - if self.backup_skip: - return True - - if not keys or (not values and (not target_locations or not target_sizes)): - logger.error("Keys or values were not passed") - return False - - if not values: - values = list(zip(target_locations, target_sizes)) - - # Add suffix to keys - suffixed_keys = [self._get_suffixed_key(key) for key in keys] - - if self.backend_selector.mem_type == "FILE": - file_paths = [] - for key in suffixed_keys: - 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 - file_paths.append(file_path) - return self._execute_transfer(values, file_paths, "WRITE") - else: # mem_type == "OBJ" - return self._execute_transfer(values, suffixed_keys, "WRITE") - - ############################################################################ - # batch_*_v1 functions - # zero copy + non-zero-copy version for get, set, exists, batch_exists - ############################################################################ - - def clear(self) -> None: - self.file_manager.clear() + raise NotImplementedError("deprecated; use batch_set_v1") def register_mem_pool_host(self, mem_pool_host: HostKVCache): super().register_mem_pool_host(mem_pool_host) @@ -332,10 +234,96 @@ class HiCacheNixl(HiCacheStorage): "page_first_direct", ] + if self.needs_page_alignment and self.is_zero_copy: + # Check that the kv_buffer base AND per-page strides are multiples of + # the OS page size so every pointer passed to NIXL (base + p * stride) + # is page-aligned. The base is whatever torch.empty() happened to give + # us -- it is not guaranteed to be page-aligned. Fall back to copy mode + # if either condition fails. + # 4096: O_DIRECT alignment is FS-dependent (some allow 512 B); 4 KiB + # is the safe lower bound all known FSes accept, and real page-sizes meet it. + if not self.mem_pool_host.is_stride_page_aligned(4096): + logger.warning( + "HiCacheNixl: O_DIRECT is active but the host kv_buffer is " + "not OS-page-aligned (base or per-page stride). Falling back " + "to copy mode for this pool." + ) + self.is_zero_copy = False + + if self.is_zero_copy: + kv = mem_pool_host.kv_buffer + self._pre_register_host( + kv.data_ptr(), kv.numel() * kv.element_size(), "kv_buffer" + ) + else: + # One bounce buffer per direction so set/get run lock-free across + # the prefetch and backup threads. Sized from get_dummy_flat_data_page() + # so each slot matches what the v1 path would otherwise allocate. + sample = mem_pool_host.get_dummy_flat_data_page() + page_numel = sample.numel() + self._bounce_page_bytes = page_numel * sample.element_size() + del sample + pin_memory = bool(getattr(mem_pool_host, "pin_memory", False)) + self._bounce_set = self._alloc_registered( + page_numel, mem_pool_host.dtype, pin_memory, "bounce_set" + ) + self._bounce_get = self._alloc_registered( + page_numel, mem_pool_host.dtype, pin_memory, "bounce_get" + ) + logger.info( - f"HiCacheNixl: Registered mem_pool_host with layout {self.mem_pool_host.layout}, zero_copy set to {self.is_zero_copy}" + f"HiCacheNixl: pre-registered host regions for " + f"layout={mem_pool_host.layout} zero_copy={self.is_zero_copy}" ) + def _alloc_registered( + self, + page_numel: int, + dtype: torch.dtype, + pin_memory: bool, + kind: str, + ) -> torch.Tensor: + """Allocate a ``(STORAGE_BATCH_SIZE, page_numel)`` bounce buffer and + pre-register it as a DRAM region with NIXL. Uses alloc_mmap so the + buffer is page-aligned -- required when O_DIRECT is on for any + file-based backend (POSIX/GDS/GDS_MT/3FS). pin_memory is currently + unused (alloc_mmap does not support it).""" + buf = alloc_mmap((STORAGE_BATCH_SIZE, page_numel), dtype) + self._pre_register_host(buf.data_ptr(), buf.numel() * buf.element_size(), kind) + return buf + + def _pre_register_host(self, base_addr: int, total_size: int, kind: str) -> None: + """Register a single DRAM region up-front and remember the handle.""" + reg_descs = self.agent.get_reg_descs([(base_addr, total_size, 0, "")], "DRAM") + if reg_descs is None: + raise RuntimeError(f"Failed to build reg descs for host {kind}") + try: + self._host_regs.append(self.agent.register_memory(reg_descs)) + except Exception as e: + raise RuntimeError(f"Failed to pre-register host {kind} with NIXL") from e + + def clear(self) -> None: + if self.file_manager is None: + return + self.file_manager.clear() + + def close(self): + while self._host_regs: + reg = self._host_regs.pop() + try: + self.agent.deregister_memory(reg) + except Exception as e: + logger.debug("deregister of pre-registered host region failed: %s", e) + self._bounce_set = None + self._bounce_get = None + self._bounce_page_bytes = None + + def __del__(self): + try: + self.close() + except Exception: + pass + def exists(self, key: str) -> bool: results = self.batch_exists([key]) return results > 0 @@ -345,8 +333,6 @@ class HiCacheNixl(HiCacheStorage): 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 = ( @@ -356,14 +342,7 @@ class HiCacheNixl(HiCacheStorage): 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, - ) + tuples = [self._create_query_tuple(key) for key in key_list] query_res = self.agent.query_memory( tuples, @@ -377,17 +356,14 @@ class HiCacheNixl(HiCacheStorage): 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 + # Each key maps to a `_k` entry, plus a `_v` entry on non-MLA models + # (MLA stores k/v interleaved in a single 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") + for key in keys: + suffixed_key = self._get_suffixed_key(key) + key_list.append(f"{suffixed_key}_k") + if not self.is_mla_model: key_list.append(f"{suffixed_key}_v") - return key_list def _get_location_and_size_list_from_meta( @@ -403,99 +379,133 @@ class HiCacheNixl(HiCacheStorage): 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 [], [], [] - return key_list, [], ptr_list, element_size_list + 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 + def _bounce_slot_buffers(self, buf: torch.Tensor, page_num: int) -> List[tuple]: + """Return ``page_num`` ``(addr, size)`` tuples pointing at the first + ``page_num`` slots of ``buf``. + """ + base = buf.data_ptr() + return [ + (base + i * self._bounce_page_bytes, self._bounce_page_bytes) + for i in range(page_num) + ] + + def _batch_preprocess(self, keys: List[str], host_indices: torch.Tensor, op: str): + """Build (key_list, host_buffers) for the v1 path. + + For zero-copy: ``host_buffers`` are ``(addr, size)`` tuples inside the + pre-registered ``kv_buffer``. + For non-zero-copy: ``host_buffers`` are slots of the direction-specific + pre-registered bounce buffer (``_bounce_set`` for set, ``_bounce_get`` + for get); for ``op == "set"`` we copy the host pages into those slots + here so the subsequent transfer reads from the bounce buffer. + Returns ``([], [])`` on validation failure. + """ + page_size = self.mem_pool_host.page_size + page_num = len(host_indices) // 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}" + f"HiCacheNixl: empty keys or mismatch in keys and host_indices lengths. keys: {len(keys)}, host_indices: {len(host_indices)}, page_size: {page_size}" ) - return [], [], [], [] + return [], [] if self.is_zero_copy: - key_list, _, ptr_list, element_size_list = ( - self._get_location_and_size_list_from_meta(keys, host_indices) + key_list, ptr_list, 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) - ] + host_buffers = list(zip(ptr_list, size_list)) + return key_list, host_buffers - 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 - ] + if page_num > STORAGE_BATCH_SIZE: + logger.error( + f"HiCacheNixl: batch size {page_num} exceeds bounce buffer capacity {STORAGE_BATCH_SIZE}" + ) + return [], [] - return key_list, target_tensors, ptr_list, element_size_list + bounce = self._bounce_set if op == "set" else self._bounce_get + if op == "set": + for i in range(page_num): + src = self.mem_pool_host.get_data_page( + host_indices[i * page_size], flat=True + ) + bounce[i].copy_(src) - def _batch_get_zero_copy_impl( + host_buffers = self._bounce_slot_buffers(bounce, page_num) + key_list = [self._get_suffixed_key(key) for key in keys] + return key_list, host_buffers + + def _batch_xfer( 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: + host_buffers: List[tuple], + direction: str, + ) -> List[bool]: + """Run a batch READ or WRITE for the v1 path against the pre-registered + host region (no per-transfer host registration). + """ + if not key_strs or not host_buffers: 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" - ) + if len(key_strs) != len(host_buffers): + logger.error("Mismatch between number of key_strs and host_buffers") 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") + success = self._xfer_pre_registered(host_buffers, file_paths, direction) + else: # mem_type == "OBJ" + success = self._xfer_pre_registered(host_buffers, key_strs, direction) - return [True] * len(key_strs) if success else [False] * len(key_strs) + # READ results are consumed by _batch_get_postprocess, which pairs + # entries 2*i / 2*i+1 for non-MLA zero-copy: it needs one bool per + # key_str (i.e. per `_k`/`_v` buffer). WRITE results map 1:1 to + # pages, i.e. to `keys`. + result_len = len(key_strs) if direction == "READ" else len(keys) + return [success] * result_len 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 + page_size = self.mem_pool_host.page_size + page_num = len(host_indices) // 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[2 * i] and results[2 * i + 1]) for i in range(page_num)] - return results + # non zero copy: copy data from the get-side bounce buffer to mem_pool_host + for i in range(page_num): + if not results[i]: + break + self.mem_pool_host.set_from_flat_data_page( + host_indices[i * page_size], self._bounce_get[i] + ) + return results + + def _log_xfer_stats( + self, + op_name: str, + num_keys: int, + host_indices: torch.Tensor, + buffer_sizes: List[int], + elapsed_ms: float, + ) -> None: + total_bytes = sum(s for s in buffer_sizes if s is not None) + bw = total_bytes / (elapsed_ms / 1000) / (1024 * 1024) if elapsed_ms else 0.0 + logger.debug( + f"HiCacheNixl {op_name} transferred: {num_keys} keys (pages), " + f"{host_indices.numel()} host_indices, {total_bytes} bytes, " + f"total time: {elapsed_ms:.3f} ms, effective bandwidth: {bw:.2f} MB/s" + ) def batch_get_v1( self, @@ -503,106 +513,28 @@ class HiCacheNixl(HiCacheStorage): 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: + if not self._host_regs: logger.error( - "HiCacheNixl batch_get_v1: preprocessing failed, empty key_strs, buffer_ptrs or buffer_sizes" + "HiCacheNixl batch_get_v1: register_mem_pool_host must be called first" ) return [False] * len(keys) + key_strs, host_buffers = self._batch_preprocess(keys, host_indices, "get") + if not key_strs or not host_buffers: + 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 + results = self._batch_xfer(keys, key_strs, host_buffers, "READ") + elapsed_ms = (time.perf_counter() - start_time) * 1000 + self._log_xfer_stats( + "batch_get_v1", + len(keys), + host_indices, + [s for _, s in host_buffers], + elapsed_ms, ) - 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) + return self._batch_get_postprocess(host_indices, results) def batch_set_v1( self, @@ -610,7 +542,6 @@ class HiCacheNixl(HiCacheStorage): host_indices: torch.Tensor, extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: - # skip on MLA backup rank if self.backup_skip: return [True] * len(keys) @@ -618,27 +549,25 @@ class HiCacheNixl(HiCacheStorage): 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: + if not self._host_regs: logger.error( - "HiCacheNixl batch_set_v1: preprocessing failed, empty key_strs, buffer_ptrs or buffer_sizes" + "HiCacheNixl batch_set_v1: register_mem_pool_host must be called first" ) return [False] * len(keys) + key_strs, host_buffers = self._batch_preprocess(keys, host_indices, "set") + if not key_strs or not host_buffers: + 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 + results = self._batch_xfer(keys, key_strs, host_buffers, "WRITE") + elapsed_ms = (time.perf_counter() - start_time) * 1000 + self._log_xfer_stats( + "batch_set_v1", + len(keys), + host_indices, + [s for _, s in host_buffers], + elapsed_ms, ) - 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 + return results diff --git a/python/sglang/srt/mem_cache/storage/nixl/nixl_registry.py b/python/sglang/srt/mem_cache/storage/nixl/nixl_registry.py new file mode 100644 index 000000000..f9f341a05 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/nixl/nixl_registry.py @@ -0,0 +1,144 @@ +"""NIXL memory-registration helpers, exposed as context managers. + +A ``NixlRegistry`` instance bundles the agent, the memory type, and +(optionally) the file manager. Its ``storage(...)`` method is a context +manager that performs the entire register-and-build-descs sequence for +the storage side of a transfer on entry, yields the ``xfer_descs`` (or +None on failure), and unwinds ``agent.deregister_memory`` plus any +``os.close(fd)`` on exit. + +The host side is pre-registered up front by ``HiCacheNixl`` and is not +touched per transfer. +""" + +import logging +import threading +from contextlib import contextmanager +from typing import List, Optional + +from .nixl_utils import NixlFileManager + +logger = logging.getLogger(__name__) + + +def _buffer_sizes(buffers) -> Optional[List[int]]: + """Per-buffer byte sizes for ``(addr, len)`` tuple inputs.""" + if not buffers or not isinstance(buffers[0], tuple): + return None + return [b[1] for b in buffers] + + +class NixlRegistry: + """Owns the (agent, mem_type, file_manager) triple and provides a + context manager for the storage side of a transfer. + + A single instance is created once per HiCacheNixl in __init__ and + reused for every transfer. + """ + + def __init__( + self, + agent, + mem_type: str, + file_manager: Optional[NixlFileManager] = None, + ): + self.agent = agent + self.mem_type = mem_type + self.file_manager = file_manager + # OBJ devIds key a process-wide map in the NIXL OBJ plugin + # (devIdToObjKey_) that is not protected by a lock, so concurrent + # OBJ registrations must use disjoint devId ranges. Allocate them + # from a single monotonic counter. + self._obj_devid_lock = threading.Lock() + self._obj_devid_next = 1 + + @contextmanager + def _open_files(self, paths: List[str], create: bool): + """Open fds for ``paths``; close all of them on exit. + + Yields the list of fds, or None if any open fails (already-opened + fds are closed before returning by the same ``finally``). + """ + fds: List[int] = [] + try: + for path in paths: + fd = self.file_manager.open_file(path, create=create) + if fd is None: + yield None + return + fds.append(fd) + yield fds + finally: + for fd in fds: + self.file_manager.close_file(fd) + + @contextmanager + def _registered(self, items: List[tuple], mem_type: str): + """Register ``items`` with NIXL; deregister on exit. + + Yields the registration handle, or None if registration fails. + """ + reg = None + if items: + reg_descs = self.agent.get_reg_descs(items, mem_type) + if reg_descs is not None: + try: + reg = self.agent.register_memory(reg_descs) + except Exception as e: + logger.error(f"Failed to register memory of type {mem_type}: {e}") + try: + yield reg + finally: + if reg is not None: + try: + self.agent.deregister_memory(reg) + except Exception as e: + logger.debug("deregister_memory skipped: %s", e) + + @contextmanager + def storage(self, buffers, keys, direction): + """Open + register the storage side; deregister and close fds on exit. + + Yields the storage xfer_descs, or None on failure. For the FILE + backend, files are created (O_CREAT) when ``direction == "WRITE"``. + """ + sizes = _buffer_sizes(buffers) + if sizes is None: + yield None + return + + if self.mem_type == "FILE": + with self._open_files(keys, create=(direction == "WRITE")) as fds: + if fds is None: + yield None + return + tuples = [(0, sizes[i], fds[i], keys[i]) for i in range(len(keys))] + with self._registered(tuples, "FILE") as reg: + if reg is None: + yield None + return + yield self.agent.get_xfer_descs( + [(0, sizes[i], fds[i]) for i in range(len(fds))], "FILE" + ) + else: # OBJ + # Reg tuple: (addr=0, size, devId, metaInfo=key). + # Xfer tuple: (addr=0, size, devId). devId links each xfer desc + # back to its registered object's metaInfo, so devIds must be + # unique within the list AND globally unique across concurrent + # storage() calls (the OBJ plugin's devIdToObjKey_ map is shared + # and unlocked). NIXL's pybind layer requires position 3 to be + # int, hence the key goes in metaInfo (position 4). + n = len(keys) + with self._obj_devid_lock: + base = self._obj_devid_next + self._obj_devid_next += n + dev_ids = list(range(base, base + n)) + tuples = [(0, sizes[i], dev_ids[i], keys[i]) for i in range(n)] + with self._registered(tuples, "OBJ") as reg: + if reg is None: + yield None + return + yield self.agent.get_xfer_descs( + [(0, sizes[i], dev_ids[i]) for i in range(n)], + self.mem_type, + ) 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 9742cf3f3..6e874b63c 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py +++ b/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py @@ -1,8 +1,8 @@ import logging import os -from typing import Any, List, Optional, Tuple, Union +from typing import Optional -import torch +from sglang.srt.environ import envs logger = logging.getLogger(__name__) @@ -23,6 +23,17 @@ class NixlBackendConfig: """ self.config = config or {} + def get_use_direct_io(self) -> bool: + """Return True if O_DIRECT should be requested when opening files. + + Checks the top-level ``use_direct_io`` key in the long-form JSON config first, + then falls back to the ``SGLANG_HICACHE_NIXL_USE_DIRECT_IO`` environment variable + (default: enabled). + """ + if "use_direct_io" in self.config: + return bool(self.config["use_direct_io"]) + return envs.SGLANG_HICACHE_NIXL_USE_DIRECT_IO.get() + def get_specified_plugin(self) -> str: """decide which plugin to use: either config or SGLANG_HICACHE_NIXL_BACKEND_PLUGIN specifies the plugin, if not, use "auto" """ @@ -90,11 +101,6 @@ class NixlBackendSelection: self.mem_type = None self.nixlconfig = nixlconfig - def set_bucket(self, bucket_name: str) -> None: - """Set AWS bucket name in environment variable.""" - os.environ["AWS_DEFAULT_BUCKET"] = bucket_name - logger.debug(f"Set AWS bucket name to: {bucket_name}") - def create_backend(self, agent) -> bool: """Create the appropriate NIXL backend based on configuration.""" try: @@ -162,77 +168,28 @@ class NixlBackendSelection: return False -class NixlRegistration: - """Handles NIXL memory registration.""" - - def __init__(self, agent): - self.agent = agent - - def create_query_tuples( - self, key: str, mem_type: str, file_manager=None - ) -> List[Tuple]: - """Create NIXL tuples for querying memory. - Args: - key: Key to query (file path for FILE or object key for OBJ) - mem_type: Memory type ("FILE" or "OBJ") - file_manager: Optional NixlFileManager for FILE memory type - Returns: - List of NIXL tuples for querying - """ - if mem_type == "FILE": - if file_manager is None: - logger.error("file_manager required for FILE memory type") - return [] - return [(0, 0, 0, file_manager.get_file_path(key))] - else: # OBJ - return [(0, 0, 0, key)] - - def _register_memory( - self, - items: Union[List[tuple], torch.Tensor, List[torch.Tensor]], - mem_type: Optional[str] = None, - ) -> Optional[Any]: - """Common registration logic for files, objects, and buffers. - Args: - items: List of tuples or tensors to register - mem_type: Memory type ("FILE", "OBJ") or None for tensor or list of tensors - """ - if isinstance(items, list) and not items: - return None - - reg_descs = self.agent.get_reg_descs(items, mem_type) - if reg_descs is None: - logger.error("Failed to create registration descriptors") - return None - - try: - registered_memory = self.agent.register_memory(reg_descs) - return registered_memory # Could be None in case of error - except Exception as e: - if not mem_type: - logger.error(f"Failed to register Tensors with NIXL: {e}") - else: - logger.error( - f"Failed to register memory of type {mem_type} with NIXL: {e}" - ) - return None - - class NixlFileManager: """Handles file system operations for NIXL.""" - def __init__(self, base_dir: str): + def __init__(self, base_dir: str, use_direct_io: bool = True): """ Initialize file manager. Args: base_dir: Base directory for storing tensor files + use_direct_io: If True, open files with O_DIRECT (bypasses OS page cache). + Falls back to buffered I/O with a warning when O_DIRECT is unavailable. """ self.base_dir = base_dir + self.use_direct_io = use_direct_io if base_dir == "": - logger.debug(f"Initialized file manager without a base directory") + logger.debug( + f"Initialized file manager without a base directory. Direct I/O: {use_direct_io}" + ) else: os.makedirs(base_dir, exist_ok=True) - logger.debug(f"Initialized file manager with base directory: {base_dir}") + logger.debug( + f"Initialized file manager with base directory: {base_dir}. Direct I/O: {use_direct_io}" + ) def clear(self) -> None: """Clear all files in the base directory.""" @@ -254,23 +211,26 @@ class NixlFileManager: """Get full file path for a given key.""" return os.path.join(self.base_dir, key) - def create_file(self, file_path: str) -> bool: - """Create a file if it doesn't exist.""" - try: - os.makedirs(os.path.dirname(file_path), exist_ok=True) - if not os.path.exists(file_path): - with open(file_path, "wb") as f: - pass # Create empty file - return True - except Exception as e: - logger.error(f"Failed to create file {file_path}: {e}") - return False + def open_file(self, file_path: str, create: bool = False) -> Optional[int]: + """Open a file and return its file descriptor. - def open_file(self, file_path: str) -> Optional[int]: - """Open a file and return its file descriptor.""" + If ``create`` is True, the file is created if it does not exist + (mode 0o644, no truncation). When ``self.use_direct_io`` is True, + the file is opened with ``O_DIRECT`` (bypasses the OS page cache); + falls back to buffered I/O with a warning if ``O_DIRECT`` is + unavailable on this platform. + """ + flags = os.O_RDWR | os.O_CREAT if create else os.O_RDWR + if self.use_direct_io: + if hasattr(os, "O_DIRECT"): + flags |= os.O_DIRECT + else: + logger.warning( + "use_direct_io is True, but O_DIRECT is not available on " + "this system. Falling back to buffered I/O." + ) try: - fd = os.open(file_path, os.O_RDWR) - return fd + return os.open(file_path, flags, 0o644) except Exception as e: logger.error(f"Failed to open file {file_path}: {e}") return None @@ -283,17 +243,3 @@ class NixlFileManager: except Exception as e: logger.error(f"Failed to close file descriptor {fd}: {e}") return False - - def files_to_nixl_tuples( - self, file_paths: List[str] - ) -> List[Tuple[int, int, int, str]]: - """Create NIXL tuples (offset, length, fd, file_path) for given files.""" - tuples = [] - for path in file_paths: - if (fd := self.open_file(path)) is None: - # Clean up on failure - for t in tuples: - self.close_file(t[2]) - return [] - tuples.append((0, 0, fd, path)) - return tuples diff --git a/python/sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py b/python/sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py deleted file mode 100755 index 9ddbbb24b..000000000 --- a/python/sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py +++ /dev/null @@ -1,315 +0,0 @@ -#!/usr/bin/env python3 - -import os -import unittest -from typing import List -from unittest.mock import MagicMock - -import torch - -from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig -from sglang.srt.mem_cache.storage.nixl.hicache_nixl import HiCacheNixl -from sglang.srt.mem_cache.storage.nixl.nixl_utils import ( - NixlFileManager, - NixlRegistration, -) - - -class TestNixlUnified(unittest.TestCase): - """Unified test suite for all NIXL components.""" - - def setUp(self): - """Set up test environment.""" - # Create test directories - self.test_dir = "/tmp/test_nixl_unified" - os.makedirs(self.test_dir, exist_ok=True) - os.environ["SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR"] = self.test_dir - - # Mock NIXL agent for registration tests - self.mock_agent = MagicMock() - self.mock_agent.get_reg_descs.return_value = "mock_reg_descs" - self.mock_agent.register_memory.return_value = "mock_registered_memory" - - # Create instances - self.file_manager = NixlFileManager(self.test_dir) - self.registration = NixlRegistration(self.mock_agent) - - # Create storage config for testing - self.storage_config = HiCacheStorageConfig( - tp_rank=0, - tp_size=2, - pp_rank=0, - pp_size=1, - attn_cp_rank=0, - attn_cp_size=1, - is_mla_model=False, - enable_storage_metrics=False, - is_page_first_layout=False, - model_name="test_model", - extra_config={"plugin": {"posix": {"active": True}}}, - ) - - try: - self.hicache = HiCacheNixl( - storage_config=self.storage_config, - file_path=self.test_dir, - ) - self.hicache = HiCacheNixl(storage_config=self.storage_config) - except ImportError: - self.skipTest("NIXL not available, skipping NIXL storage tests") - - def tearDown(self): - """Clean up test directories.""" - if os.path.exists(self.test_dir): - import shutil - - shutil.rmtree(self.test_dir, ignore_errors=True) - - @staticmethod - def _open_fds() -> int: - return len(os.listdir("/proc/self/fd")) - - def delete_test_file(self, file_path: str) -> bool: - """Helper method to delete a test file. - - Args: - file_path: Path to the file to delete - - Returns: - bool: True if file was deleted or didn't exist, False on error - """ - try: - if os.path.exists(file_path): - os.remove(file_path) - return True - except Exception as e: - return False - - def verify_tensors_equal(self, expected: torch.Tensor, actual: torch.Tensor): - """Helper to verify tensor equality.""" - self.assertIsNotNone(actual, "Retrieved tensor is None") - self.assertTrue( - torch.allclose(expected, actual, atol=1e-6), - f"Tensors not equal:\nExpected: {expected}\nActual: {actual}", - ) - - def verify_tensor_lists_equal( - self, expected: List[torch.Tensor], actual: List[torch.Tensor] - ): - """Helper to verify lists of tensors are equal.""" - self.assertEqual(len(expected), len(actual), "Lists have different lengths") - for exp, act in zip(expected, actual): - self.verify_tensors_equal(exp, act) - - # ============================================================================ - # HiCache Integration Tests - # ============================================================================ - - def test_single_set_get(self): - """Test single tensor set/get operations.""" - key = "test_key" - value = torch.randn(10, 10, device="cpu") - dst_tensor = torch.zeros_like(value, device="cpu") - - # Test set - self.assertTrue(self.hicache.set(key, value)) - self.assertTrue(self.hicache.exists(key)) - - # Test get - retrieved = self.hicache.get(key, dst_tensor) - self.verify_tensors_equal(value, dst_tensor) - self.verify_tensors_equal(value, retrieved) - - # Same test in addr,len mode with another key and dst_tensor - key2 = "test_key2" - dst_tensor2 = torch.zeros_like(value, device="cpu") - src_addr, src_len = value.data_ptr(), value.numel() * value.element_size() - dst_addr, dst_len = ( - dst_tensor2.data_ptr(), - dst_tensor2.numel() * dst_tensor2.element_size(), - ) - - # Test set - self.assertTrue(self.hicache.set(key, None, src_addr, src_len)) - self.assertTrue(self.hicache.exists(key)) - - # Test get - retrieved2 = self.hicache.get(key, dst_addr, dst_len) - self.assertTrue(retrieved2 is None) - self.verify_tensors_equal(value, dst_tensor2) - - def test_batch_set_get(self): - """Test batch tensor set/get operations.""" - keys = ["key1", "key2", "key3"] - values = [ - torch.randn(5, 5, device="cpu"), - torch.randn(3, 3, device="cpu"), - torch.randn(7, 7, device="cpu"), - ] - dst_tensors = [torch.zeros_like(v, device="cpu") for v in values] - - # Test batch set - self.assertTrue(self.hicache.batch_set(keys, values)) - self.assertTrue(all(self.hicache.exists(key) for key in keys)) - - # Test batch get - retrieved = self.hicache.batch_get(keys, dst_tensors) - self.verify_tensor_lists_equal(values, retrieved) - - # Same test in addr,len mode with another key and dst_tensor - keys2 = ["key4", "key5", "key6"] - dst_tensors2 = [torch.zeros_like(v, device="cpu") for v in values] - src_addrs = [v.data_ptr() for v in values] - src_lens = [v.numel() * v.element_size() for v in values] - dst_addrs = [dt.data_ptr() for dt in dst_tensors2] - dst_lens = [dt.numel() * dt.element_size() for dt in dst_tensors2] - - # Test batch set - self.assertTrue(self.hicache.batch_set(keys2, None, src_addrs, src_lens)) - self.assertTrue(all(self.hicache.exists(key) for key in keys2)) - - # Test batch get - retrieved2 = self.hicache.batch_get(keys, dst_addrs, dst_lens) - self.assertTrue(all(ret is None for ret in retrieved2)) - self.verify_tensor_lists_equal(values, dst_tensors2) - - def test_mixed_operations(self): - """Test mixing single and batch operations.""" - # Test interleaved set/get operations - key1, key2 = "key1", "key2" - value1 = torch.randn(4, 4, device="cpu") - value2 = torch.randn(6, 6, device="cpu") - dst1 = torch.zeros_like(value1) - dst2 = torch.zeros_like(value2) - - # Single set/get; baseline after first set absorbs any one-time NIXL internals - self.assertTrue(self.hicache.set(key1, value1)) - fds = self._open_fds() - retrieved1 = self.hicache.get(key1, dst1) - self.verify_tensors_equal(value1, retrieved1) - self.assertEqual(self._open_fds(), fds, "fd leak after get") - - # Batch set/get - self.assertTrue(self.hicache.batch_set([key2], [value2])) - self.assertEqual(self._open_fds(), fds, "fd leak after batch_set") - retrieved2 = self.hicache.batch_get([key2], [dst2]) - self.verify_tensors_equal(value2, retrieved2[0]) - self.assertEqual(self._open_fds(), fds, "fd leak after batch_get") - - def test_data_integrity(self): - """Test data integrity across operations.""" - # Test with various tensor types and sizes - test_cases = [ - ("float32", torch.randn(10, 10, dtype=torch.float32)), - ("float64", torch.randn(5, 5, dtype=torch.float64)), - ("int32", torch.randint(-100, 100, (8, 8), dtype=torch.int32)), - ("int64", torch.randint(-100, 100, (6, 6), dtype=torch.int64)), - ("bool", torch.randint(0, 2, (4, 4)).bool()), - ] - - for name, tensor in test_cases: - with self.subTest(tensor_type=name): - key = f"test_{name}" - dst_tensor = torch.zeros_like(tensor) - - # Set and immediately get - self.assertTrue(self.hicache.set(key, tensor)) - retrieved1 = self.hicache.get(key, dst_tensor) - self.verify_tensors_equal(tensor, retrieved1) - - # Get again to verify persistence - dst_tensor.zero_() - retrieved2 = self.hicache.get(key, dst_tensor) - self.verify_tensors_equal(tensor, retrieved2) - - def test_basic_file_operations(self): - """Test basic file operations.""" - test_file = os.path.join(self.test_dir, "test_file.bin") - self.file_manager.create_file(test_file) - self.assertTrue(os.path.exists(test_file)) - self.assertEqual(os.path.getsize(test_file), 0) # Empty file - - # Test file deletion - self.assertTrue(self.delete_test_file(test_file)) - self.assertFalse(os.path.exists(test_file)) - - def test_create_nixl_tuples(self): - """Test creation of NIXL tuples.""" - test_file = os.path.join(self.test_dir, "test_file.bin") - self.file_manager.create_file(test_file) - - # Test tuple creation - tuples = self.file_manager.files_to_nixl_tuples([test_file]) - self.assertIsNotNone(tuples) - self.assertTrue(len(tuples) > 0) - - def test_error_handling(self): - """Test error handling in file operations.""" - # Test non-existent file - self.assertTrue( - self.delete_test_file("nonexistent_file.bin") - ) # Returns True if file doesn't exist - - # Test invalid file path - self.assertFalse(self.file_manager.create_file("")) # Empty path should fail - - def test_register_buffers(self): - """Test registration of memory buffers.""" - # Create test tensor - tensor = torch.randn(10, 10) - - # Test buffer registration - self.assertIsNotNone(self.hicache.register_buffers(tensor)) - - # Test batch registration - tensors = [torch.randn(5, 5) for _ in range(3)] - self.assertIsNotNone(self.hicache.register_buffers(tensors)) - - def test_register_files(self): - """Test registration of files with NIXL.""" - files = [os.path.join(self.test_dir, f"test_file_{i}.bin") for i in range(3)] - for file in files: - self.file_manager.create_file(file) - - result = self.hicache.register_files(files) - self.assertIsNotNone(result) - - def test_batch_set_v1_skips_on_nonzero_mla_rank(self): - """Test batch_set_v1 is a no-op on nonzero MLA backup ranks.""" - self.hicache.storage_config.is_mla_model = True - self.hicache.storage_config.tp_rank = 1 - self.hicache.backup_skip = True - self.hicache._batch_set_preprocess = MagicMock( - side_effect=AssertionError("batch_set_v1 should have been skipped") - ) - - results = self.hicache.batch_set_v1(["key1", "key2"], torch.tensor([0, 1])) - - self.assertEqual(results, [True, True]) - self.hicache._batch_set_preprocess.assert_not_called() - - def test_batch_exists_zero_copy_mla_uses_single_key_denominator(self): - """Test zero-copy MLA batch_exists counts one storage key per logical key.""" - self.hicache.is_zero_copy = True - self.hicache.is_mla_model = True - self.hicache.agent.query_memory = MagicMock(return_value=[object(), None]) - - result = self.hicache.batch_exists(["key1", "key2"]) - - self.assertEqual(result, 1) - - def test_batch_exists_zero_copy_mha_uses_two_key_denominator(self): - """Test zero-copy MHA batch_exists counts k/v pairs per logical key.""" - self.hicache.is_zero_copy = True - self.hicache.is_mla_model = False - self.hicache.agent.query_memory = MagicMock( - return_value=[object(), object(), None, None] - ) - - result = self.hicache.batch_exists(["key1", "key2"]) - - self.assertEqual(result, 1) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py new file mode 100644 index 000000000..a6ea710a8 --- /dev/null +++ b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py @@ -0,0 +1,777 @@ +"""Unit tests for the NIXL HiCache storage backend -- no server, no model loading.""" + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-a", runner_config="1-gpu-small") + +import os +import shutil +import socket +import subprocess +import tempfile +import threading +import time +import unittest +from typing import List + +import torch + +from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig +from sglang.srt.mem_cache.storage.nixl.hicache_nixl import HiCacheNixl +from sglang.test.test_utils import CustomTestCase + +# Stress tests are opt-in: CI never sets this; set locally to exercise them. +STRESS_ENABLED = bool(os.environ.get("SGLANG_RUN_NIXL_STRESS")) + + +class MockMemPoolHost: + """Minimal MHA-style HostKVCache stand-in supporting the v1 paths. + + zero_copy mode uses ``page_first`` so ``get_page_buffer_meta`` returns + valid (k, v) pointers into ``kv_buffer``. Non-zero-copy uses + ``layer_first`` so the slow path uses ``get_data_page`` / + ``set_from_flat_data_page`` against the same buffer. + """ + + def __init__( + self, + is_zero_copy_mode: bool, + page_size: int = 2, + layer_num: int = 2, + head_num: int = 2, + head_dim: int = 4, + num_pages: int = 4, + dtype: torch.dtype = torch.float32, + ): + self.layout = "page_first" if is_zero_copy_mode else "layer_first" + self.page_size = page_size + self.layer_num = layer_num + self.head_num = head_num + self.head_dim = head_dim + self.dtype = dtype + self.num_pages = num_pages + self.size = page_size * num_pages + self.pin_memory = False + if is_zero_copy_mode: + # page_first: (2, size, layer, head, head_dim) + self.kv_buffer = torch.zeros( + (2, self.size, layer_num, head_num, head_dim), dtype=dtype + ) + else: + # layer_first: (2, layer, size, head, head_dim) + self.kv_buffer = torch.zeros( + (2, layer_num, self.size, head_num, head_dim), dtype=dtype + ) + + def get_page_buffer_meta(self, indices): + ptr_list = [] + base = self.kv_buffer.data_ptr() + v_offset = ( + self.layer_num + * self.size + * self.head_num + * self.head_dim + * self.dtype.itemsize + ) + idx_list = indices.tolist() + for i in range(0, len(idx_list), self.page_size): + k_ptr = base + idx_list[i] * ( + self.layer_num * self.head_num * self.head_dim * self.dtype.itemsize + ) + ptr_list.append(k_ptr) + ptr_list.append(k_ptr + v_offset) + element_size = ( + self.layer_num + * self.dtype.itemsize + * self.page_size + * self.head_num + * self.head_dim + ) + return ptr_list, [element_size] * len(ptr_list) + + def get_dummy_flat_data_page(self): + return torch.zeros( + (2, self.layer_num, self.page_size, self.head_num, self.head_dim), + dtype=self.dtype, + ).flatten() + + def get_data_page(self, index, flat=True): + if hasattr(index, "item"): + index = int(index.item()) + page = self.kv_buffer[:, :, index : index + self.page_size, :, :] + return page.flatten() if flat else page + + def set_from_flat_data_page(self, index, data_page): + if hasattr(index, "item"): + index = int(index.item()) + self.kv_buffer[:, :, index : index + self.page_size, :, :] = data_page.reshape( + 2, self.layer_num, self.page_size, self.head_num, self.head_dim + ) + + def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool: + # Test tensors are too small to satisfy 4 KiB stride alignment; the + # O_DIRECT path correctly falls back to copy mode in this case. + return False + + +class MinioFixture: + """Spin up a single-node MinIO server on localhost and create a bucket. + + Relies on MinIO's default ``minioadmin``/``minioadmin`` root credentials + so no env vars need to be plumbed through. + """ + + user = "minioadmin" + password = "minioadmin" + + def __init__(self, bucket: str = "hicache-test"): + self.bucket = bucket + self.api_port = self._find_free_port() + self.data_dir = tempfile.mkdtemp(prefix="nixl_minio_") + self.proc: subprocess.Popen | None = None + + @property + def endpoint(self) -> str: + return f"127.0.0.1:{self.api_port}" + + @staticmethod + def _find_free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + @staticmethod + def _minio_bin() -> str | None: + path = shutil.which("minio") or "/usr/local/bin/minio" + if os.path.isfile(path) and os.access(path, os.X_OK): + return path + return None + + @classmethod + def is_available(cls) -> bool: + """True iff a minio binary and boto3 are both importable.""" + if cls._minio_bin() is None: + return False + try: + import boto3 # noqa: F401 + except ImportError: + return False + return True + + def start(self) -> None: + minio_bin = self._minio_bin() + if minio_bin is None: + raise FileNotFoundError("minio binary not available") + + self.proc = subprocess.Popen( + [minio_bin, "server", "--address", self.endpoint, self.data_dir], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + deadline = time.time() + 15.0 + while time.time() < deadline: + if self.proc.poll() is not None: + raise RuntimeError(f"minio exited early with rc={self.proc.returncode}") + try: + with socket.create_connection( + ("127.0.0.1", self.api_port), timeout=0.5 + ): + break + except OSError: + time.sleep(0.1) + else: + self.stop() + raise RuntimeError("minio did not become ready within 15s") + + import boto3 + from botocore.config import Config + + s3 = boto3.client( + "s3", + endpoint_url=f"http://{self.endpoint}", + aws_access_key_id=self.user, + aws_secret_access_key=self.password, + config=Config(s3={"addressing_style": "path"}, signature_version="s3v4"), + ) + s3.create_bucket(Bucket=self.bucket) + + def stop(self) -> None: + if self.proc and self.proc.poll() is None: + self.proc.terminate() + try: + self.proc.wait(timeout=5) + except subprocess.TimeoutExpired: + self.proc.kill() + self.proc.wait(timeout=5) + shutil.rmtree(self.data_dir, ignore_errors=True) + + +class TestNixlUnified(CustomTestCase): + """Unified test suite for all NIXL components.""" + + def setUp(self): + """Set up test environment.""" + self.test_dir = "/tmp/test_nixl_unified" + os.makedirs(self.test_dir, exist_ok=True) + + # Disable O_DIRECT here: these tests use small, arbitrarily-aligned + # tensors that do not satisfy the sector-alignment constraints required + # by O_DIRECT. O_DIRECT-specific behaviour is exercised in + # TestNixlDirectIO below. + self.storage_config = HiCacheStorageConfig( + tp_rank=0, + tp_size=2, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=False, + is_page_first_layout=False, + model_name="test_model", + enable_storage_metrics=False, + extra_config={ + "plugin": {"posix": {"active": True}}, + "use_direct_io": False, + }, + ) + + try: + self.hicache = HiCacheNixl( + storage_config=self.storage_config, + file_path=self.test_dir, + ) + except ImportError: + self.skipTest("NIXL not available, skipping NIXL storage tests") + + def tearDown(self): + """Clean up test directories.""" + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir, ignore_errors=True) + + @staticmethod + def _open_fds() -> int: + return len(os.listdir("/proc/self/fd")) + + def test_storage_register_failure_closes_fds(self): + """If NIXL register_memory raises after fds are opened, all fds are still closed.""" + files = [os.path.join(self.test_dir, f"fail_{i}.bin") for i in range(3)] + buffers = [(0, 64) for _ in range(3)] + + fds_before = self._open_fds() + + orig = self.hicache.agent.register_memory + + def boom(*args, **kwargs): + raise RuntimeError("simulated register_memory failure") + + self.hicache.agent.register_memory = boom + try: + with self.hicache.registry.storage(buffers, files, "WRITE") as descs: + self.assertIsNone( + descs, "storage CM should yield None on register failure" + ) + finally: + self.hicache.agent.register_memory = orig + + self.assertEqual( + self._open_fds(), + fds_before, + "fd leak after register_memory failure mid-storage", + ) + + def _assert_host_addrs_pre_registered( + self, is_zero_copy_mode: bool, hicache: HiCacheNixl = None + ): + """Exercise the v1 path and assert every host xfer addr lies within a + currently-registered host (DRAM/tensor) region. + + Spies are installed BEFORE ``register_mem_pool_host`` so the up-front + pre-registration is captured too. + """ + if hicache is None: + hicache = self.hicache + agent = hicache.agent + + # Map registration-handle id -> [(addr, size, mem_type), ...] + active_regs: dict = {} + # Capture items list per get_reg_descs call so we can attribute them + # to the registration handle returned by the next register_memory call. + pending: list = [] + + orig_get_reg = agent.get_reg_descs + + def spy_get_reg(items, mem_type=None): + # NIXL's register_memory calls get_reg_descs internally with an + # already-built nixlRegDList; iterating that pybind11 type is + # unsafe, so only record entries when the input is a plain list. + if isinstance(items, list) and items: + entries = [] + for it in items: + if isinstance(it, torch.Tensor): + entries.append( + (it.data_ptr(), it.numel() * it.element_size(), None) + ) + elif isinstance(it, tuple): + entries.append((it[0], it[1], mem_type)) + pending.append(entries) + return orig_get_reg(items, mem_type) + + orig_register = agent.register_memory + + def spy_register(reg_descs): + reg = orig_register(reg_descs) + entries = pending.pop(0) if pending else [] + active_regs[id(reg)] = entries + return reg + + orig_dereg = agent.deregister_memory + + def spy_dereg(reg): + active_regs.pop(id(reg), None) + return orig_dereg(reg) + + last_host_xfer: list = [] + + orig_get_xfer = agent.get_xfer_descs + + def spy_get_xfer(items, mem_type=None): + if mem_type in (None, "DRAM"): + ranges = [] + for it in items: + if isinstance(it, torch.Tensor): + ranges.append((it.data_ptr(), it.numel() * it.element_size())) + elif isinstance(it, tuple): + ranges.append((it[0], it[1])) + last_host_xfer.clear() + last_host_xfer.extend(ranges) + return orig_get_xfer(items, mem_type) + + violations: list = [] + orig_init = agent.initialize_xfer + + def spy_init(direction, local, remote, agent_name): + host_regs = [ + (a, s) + for entries in active_regs.values() + for (a, s, mt) in entries + if mt in (None, "DRAM") + ] + for a, s in last_host_xfer: + if not any(ra <= a and a + s <= ra + rs for (ra, rs) in host_regs): + violations.append((a, s, dict(host_regs=host_regs))) + last_host_xfer.clear() + return orig_init(direction, local, remote, agent_name) + + agent.get_reg_descs = spy_get_reg + agent.register_memory = spy_register + agent.deregister_memory = spy_dereg + agent.get_xfer_descs = spy_get_xfer + agent.initialize_xfer = spy_init + try: + mock_host = MockMemPoolHost(is_zero_copy_mode) + hicache.register_mem_pool_host(mock_host) + # Force the requested mode regardless of how register_mem_pool_host derives it. + hicache.is_zero_copy = is_zero_copy_mode + + num_pages = 3 + keys = [ + f"compliance_{int(is_zero_copy_mode)}_{i}" for i in range(num_pages) + ] + host_indices = torch.arange( + num_pages * mock_host.page_size, dtype=torch.int64 + ) + + set_results = hicache.batch_set_v1(keys, host_indices) + self.assertTrue( + all(set_results), + f"batch_set_v1 failed (zero_copy={is_zero_copy_mode}): {set_results}", + ) + + get_results = hicache.batch_get_v1(keys, host_indices) + self.assertTrue( + all(get_results), + f"batch_get_v1 failed (zero_copy={is_zero_copy_mode}): {get_results}", + ) + finally: + agent.get_reg_descs = orig_get_reg + agent.register_memory = orig_register + agent.deregister_memory = orig_dereg + agent.get_xfer_descs = orig_get_xfer + agent.initialize_xfer = orig_init + + self.assertEqual( + violations, + [], + f"Host xfer addrs not covered by registration (zero_copy={is_zero_copy_mode}): {violations}", + ) + + def test_nixl_api_contract_host_addrs_within_registered_region_zero_copy(self): + """All host xfer addrs must lie within a registered region -- zero-copy.""" + self._assert_host_addrs_pre_registered(is_zero_copy_mode=True) + + def test_nixl_api_contract_host_addrs_within_registered_region_non_zero_copy(self): + """All host xfer addrs must lie within a registered region -- non-zero-copy.""" + self._assert_host_addrs_pre_registered(is_zero_copy_mode=False) + + def _make_obj_hicache(self) -> HiCacheNixl: + """Start a MinIO server (cleaned up via addCleanup) and return a + HiCacheNixl wired to its OBJ backend. Skips the test if the backend + cannot be constructed.""" + minio = MinioFixture() + minio.start() + self.addCleanup(minio.stop) + + obj_config = HiCacheStorageConfig( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=False, + is_page_first_layout=False, + model_name="test_model", + enable_storage_metrics=False, + extra_config={ + "plugin": { + "obj": { + "active": True, + "endpoint_override": f"http://{minio.endpoint}", + "use_virtual_addressing": "false", + "access_key": minio.user, + "secret_key": minio.password, + "bucket": minio.bucket, + } + } + }, + ) + try: + return HiCacheNixl(storage_config=obj_config, file_path="") + except Exception as e: + self.skipTest(f"NIXL OBJ backend unavailable: {e}") + + @unittest.skipUnless( + MinioFixture.is_available(), "minio binary or boto3 not available" + ) + def test_nixl_api_contract_host_addrs_within_registered_region_obj(self): + """Same property over the OBJ backend (MinIO fixture).""" + self._assert_host_addrs_pre_registered( + is_zero_copy_mode=False, hicache=self._make_obj_hicache() + ) + + def test_batch_set_v1_skips_on_nonzero_mla_rank(self): + """batch_set_v1 is a no-op on nonzero MLA backup ranks. + + With backup_skip=True the early-return must fire before the host-regs + check, so calling without register_mem_pool_host still returns all-True + (the host-regs check would otherwise return all-False). + """ + self.hicache.backup_skip = True + results = self.hicache.batch_set_v1( + ["key1", "key2"], torch.tensor([0, 1], dtype=torch.int64) + ) + self.assertEqual(results, [True, True]) + + def test_batch_exists_zero_copy_mla_uses_single_key_denominator(self): + """Zero-copy MLA batch_exists counts one storage key per logical key.""" + self.hicache.is_zero_copy = True + self.hicache.is_mla_model = True + self.hicache.agent.query_memory = lambda *a, **kw: [object(), None] + + self.assertEqual(self.hicache.batch_exists(["key1", "key2"]), 1) + + def test_batch_exists_zero_copy_mha_uses_two_key_denominator(self): + """Zero-copy non-MLA batch_exists counts k/v pairs per logical key.""" + self.hicache.is_zero_copy = True + self.hicache.is_mla_model = False + self.hicache.agent.query_memory = lambda *a, **kw: [ + object(), + object(), + None, + None, + ] + + self.assertEqual(self.hicache.batch_exists(["key1", "key2"]), 1) + + def _run_concurrent_stress( + self, is_zero_copy_mode: bool, hicache: HiCacheNixl = None + ): + """One getter thread + one setter thread share the same HiCacheNixl + for ``is_zero_copy_mode``. Defaults to ``self.hicache`` (FILE backend); + pass ``hicache`` to exercise a different backend (e.g. OBJ). + + Phase 1 pre-seeds N preset pages and stores them under fixed keys. + Phase 2 runs the getter (reads the presets back and verifies content) + concurrently with the setter (writes a stream of fresh distinct keys + from a disjoint source region). The kv_buffer regions touched by the + two threads are disjoint so any data corruption observed is from the + backend's shared state (bounce buffers, devId maps, fd pool). + """ + if hicache is None: + hicache = self.hicache + + # 8 preset pages, 8 getter dst pages, 8 setter src pages -> 24 in use. + mock_host = MockMemPoolHost(is_zero_copy_mode=is_zero_copy_mode, num_pages=32) + hicache.register_mem_pool_host(mock_host) + hicache.is_zero_copy = is_zero_copy_mode + + page_size = mock_host.page_size + dtype = mock_host.dtype + num_pages = 8 + + # Disjoint per-thread regions in kv_buffer (indexed by token index). + preset_src = (0, num_pages) + getter_dst = (num_pages, 2 * num_pages) + setter_src = (2 * num_pages, 3 * num_pages) + + # zero_copy=page_first uses dim 1 for the token axis; non-zero-copy= + # layer_first uses dim 2. All buffer accesses below go through this so + # the rest of the harness stays layout-agnostic. + def token_index(start_token: int, n_tokens: int): + s = slice(start_token, start_token + n_tokens) + if is_zero_copy_mode: + return (slice(None), s, slice(None), slice(None), slice(None)) + return (slice(None), slice(None), s, slice(None), slice(None)) + + def page_index(start_page: int, n_pages: int): + return token_index(start_page * page_size, n_pages * page_size) + + def fill_pages(start_page: int, n_pages: int, value_fn): + """value_fn(i) -> scalar value for page i.""" + for i in range(n_pages): + idx = page_index(start_page + i, 1) + shape = mock_host.kv_buffer[idx].shape + mock_host.kv_buffer[idx] = torch.full( + shape, float(value_fn(i)), dtype=dtype + ) + + # Phase 1: distinct value per preset page so a wrong-page result is + # detectable; setter source is constant (value irrelevant to the + # test, just needs to be valid). + fill_pages(preset_src[0], num_pages, lambda i: i + 1) + fill_pages(setter_src[0], num_pages, lambda i: -1.0) + + preset_keys = [f"preset_{int(is_zero_copy_mode)}_{i}" for i in range(num_pages)] + preset_indices = torch.arange( + preset_src[0] * page_size, + preset_src[1] * page_size, + dtype=torch.int64, + ) + self.assertTrue( + all(hicache.batch_set_v1(preset_keys, preset_indices)), + "phase 1: presetting keys failed", + ) + + # Expected per-page-i payload after a successful get into getter_dst. + expected_pages = [ + mock_host.kv_buffer[page_index(preset_src[0] + i, 1)].clone() + for i in range(num_pages) + ] + + # Phase 2. + stop = threading.Event() + errors: List[str] = [] + errors_lock = threading.Lock() + + def record_error(msg: str): + with errors_lock: + errors.append(msg) + + def getter_loop(): + dst_indices = torch.arange( + getter_dst[0] * page_size, + getter_dst[1] * page_size, + dtype=torch.int64, + ) + loops = 0 + while not stop.is_set(): + # Zero the dst pages so a no-op get is observable. + mock_host.kv_buffer[page_index(getter_dst[0], num_pages)] = 0.0 + ok = hicache.batch_get_v1(preset_keys, dst_indices) + if not all(ok): + record_error(f"getter loop {loops}: batch_get_v1 returned {ok}") + return + for i in range(num_pages): + got = mock_host.kv_buffer[page_index(getter_dst[0] + i, 1)] + if not torch.equal(got, expected_pages[i]): + record_error(f"getter loop {loops}: preset page {i} corrupted") + return + loops += 1 + + def setter_loop(): + src_indices = torch.arange( + setter_src[0] * page_size, + setter_src[1] * page_size, + dtype=torch.int64, + ) + loops = 0 + while not stop.is_set(): + keys = [ + f"setter_{int(is_zero_copy_mode)}_{loops}_{i}" + for i in range(num_pages) + ] + ok = hicache.batch_set_v1(keys, src_indices) + if not all(ok): + record_error(f"setter loop {loops}: batch_set_v1 returned {ok}") + return + loops += 1 + + t_get = threading.Thread(target=getter_loop, daemon=True) + t_set = threading.Thread(target=setter_loop, daemon=True) + t_get.start() + t_set.start() + + # Bounded run: long enough to interleave many ops under NIXL I/O + # GIL release, short enough for a unit test. + time.sleep(3.0) + stop.set() + t_get.join(timeout=10) + t_set.join(timeout=10) + + self.assertFalse( + t_get.is_alive() or t_set.is_alive(), + "stress threads failed to stop", + ) + self.assertEqual(errors, [], f"concurrency errors: {errors}") + + @unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run") + def test_concurrent_getter_setter_file_zero_copy(self): + """Stress: concurrent getter+setter, FILE backend, zero-copy.""" + self._run_concurrent_stress(is_zero_copy_mode=True) + + @unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run") + def test_concurrent_getter_setter_file_non_zero_copy(self): + """Stress: concurrent getter+setter, FILE backend, non-zero-copy.""" + self._run_concurrent_stress(is_zero_copy_mode=False) + + @unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run") + @unittest.skipUnless( + MinioFixture.is_available(), "minio binary or boto3 not available" + ) + def test_concurrent_getter_setter_obj_zero_copy(self): + """Stress: concurrent getter+setter, OBJ backend (MinIO), zero-copy.""" + self._run_concurrent_stress( + is_zero_copy_mode=True, hicache=self._make_obj_hicache() + ) + + @unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run") + @unittest.skipUnless( + MinioFixture.is_available(), "minio binary or boto3 not available" + ) + def test_concurrent_getter_setter_obj_non_zero_copy(self): + """Stress: concurrent getter+setter, OBJ backend (MinIO), non-zero-copy.""" + self._run_concurrent_stress( + is_zero_copy_mode=False, hicache=self._make_obj_hicache() + ) + + +@unittest.skipUnless(hasattr(os, "O_DIRECT"), "O_DIRECT not available on this platform") +class TestNixlDirectIO(CustomTestCase): + """Tests for the O_DIRECT file I/O path in NixlFileManager and HiCacheNixl.""" + + def setUp(self): + self.test_dir = "/tmp/test_nixl_direct_io" + os.makedirs(self.test_dir, exist_ok=True) + + def tearDown(self): + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir, ignore_errors=True) + + def test_open_file_sets_o_direct(self): + """open_file sets O_DIRECT on the file descriptor when use_direct_io=True.""" + import fcntl + + from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager + + fm = NixlFileManager(self.test_dir, use_direct_io=True) + test_file = os.path.join(self.test_dir, "test_odirect.bin") + fd = fm.open_file(test_file, create=True) + try: + self.assertTrue(fcntl.fcntl(fd, fcntl.F_GETFL) & os.O_DIRECT) + finally: + os.close(fd) + + def test_open_file_no_o_direct(self): + """open_file does not set O_DIRECT when use_direct_io=False.""" + import fcntl + + from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager + + fm = NixlFileManager(self.test_dir, use_direct_io=False) + test_file = os.path.join(self.test_dir, "test_buffered.bin") + fd = fm.open_file(test_file, create=True) + try: + self.assertFalse(fcntl.fcntl(fd, fcntl.F_GETFL) & os.O_DIRECT) + finally: + os.close(fd) + + def _make_direct_io_hicache(self) -> HiCacheNixl: + """Return a HiCacheNixl configured for O_DIRECT (default) with the POSIX backend.""" + storage_config = HiCacheStorageConfig( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=False, + is_page_first_layout=False, + model_name="test_model", + enable_storage_metrics=False, + extra_config={"plugin": {"posix": {"active": True}}}, + # use_direct_io defaults to True (env var) + ) + try: + return HiCacheNixl(storage_config=storage_config, file_path=self.test_dir) + except ImportError: + self.skipTest("NIXL not available") + + def test_needs_page_alignment_true_for_file_backend(self): + """File-based backend + use_direct_io=True must set needs_page_alignment.""" + hicache = self._make_direct_io_hicache() + self.assertTrue(hicache.needs_page_alignment) + + def test_odirect_unaligned_pool_falls_back_to_copy(self): + """O_DIRECT with non-aligned pool strides falls back to copy mode.""" + hicache = self._make_direct_io_hicache() + + mock_host = MockMemPoolHost(is_zero_copy_mode=True) + hicache.register_mem_pool_host(mock_host) + + # MockMemPoolHost.is_stride_page_aligned() returns False, so even though + # the layout would otherwise enable zero-copy, the backend must fall back. + self.assertFalse(hicache.is_zero_copy) + self.assertIsNotNone(hicache._bounce_set) + self.assertIsNotNone(hicache._bounce_get) + + def test_odirect_disabled_via_config(self): + """Top-level use_direct_io=false in extra_config disables O_DIRECT.""" + storage_config = HiCacheStorageConfig( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=False, + is_page_first_layout=False, + model_name="test_model", + enable_storage_metrics=False, + extra_config={ + "plugin": {"posix": {"active": True}}, + "use_direct_io": False, + }, + ) + try: + hicache = HiCacheNixl( + storage_config=storage_config, file_path=self.test_dir + ) + except ImportError: + self.skipTest("NIXL not available") + self.assertFalse(hicache.needs_page_alignment) + self.assertFalse(hicache.file_manager.use_direct_io) + + +if __name__ == "__main__": + unittest.main()