Refactor NIXL hicache. Add O_DIRECT support (#25173)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

docs_new/docs/references/environment_variables.mdx

@@ -837,6 +837,16 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Decode-side incremental KV cache offload stride. Rounded down to a multiple of <code>--page-size</code> (min is <code>--page-size</code>). If unset/invalid/&lt;=0, it falls back to <code>--page-size</code>.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Not set (uses <code>--page-size</code>)</td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_HICACHE_NIXL_USE_DIRECT_IO</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable <code>O_DIRECT</code> 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 <code>{'{"use_direct_io": false}'}</code> in <code>--hicache-storage-backend-extra-config</code>. Falls back to buffered I/O with a warning when <code>O_DIRECT</code> is unavailable on the current OS.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_HUGEPAGE_SIZE</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use huge pages for host KV cache allocations (HiCache / disaggregation offload). Valid values: <code>2MB</code> (2 MiB pages via <code>MAP_HUGE_2MB</code>) or <code>1GB</code> (1 GiB pages via <code>MAP_HUGE_1GB</code>). Requires huge pages to be pre-allocated on the host OS (<code>/proc/sys/vm/nr_hugepages</code> or <code>/sys/kernel/mm/hugepages</code>). If the allocation fails, the allocator logs a warning and falls back to regular page-size mmap automatically.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Not set (uses OS default page size)</td>
+    </tr>
   </tbody>
 </table>
 

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)

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):
+        for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE):
-            batch_hashes = operation.hash_value[i : i + self.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(
+            end = min(start + self.page_size * STORAGE_BATCH_SIZE, len(tokens_to_fetch))
-                start + self.page_size * self.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):
+        for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE):
-            batch_hashes = operation.hash_value[i : i + self.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
             ]

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:

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 alloc_mmap(dims, dtype)
-        return tensor
 
 
 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(
+        cudart = torch.cuda.cudart()
-            buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0
+        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."
             )
-        if ret != 0:
-            raise RuntimeError(f"cudaHostRegister failed with error code {ret}")
     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):
 

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)

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
-- **NixlFileManager** - Handles file system operations and NIXL tuple creation
+
+### 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`)
 

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 (
+from .nixl_registry import NixlRegistry
-    NixlBackendConfig,
+from .nixl_utils import NixlBackendConfig, NixlBackendSelection, NixlFileManager
-    NixlBackendSelection,
-    NixlFileManager,
-    NixlRegistration,
-)
 
 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,117 +74,67 @@ 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.registry = NixlRegistry(
-        self.is_zero_copy = False
+            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(
+    def _create_query_tuple(self, key: str) -> tuple:
-        self, buffers: Union[torch.Tensor, List[torch.Tensor], List[tuple]]
+        """Build the NIXL query_memory tuple for a single key."""
-    ) -> Optional[Any]:
+        if self.backend_selector.mem_type == "FILE":
-        """Register tensor(s) or target locations in host memory (list of addr,len tuples) with NIXL."""
+            return (0, 0, 0, self.file_manager.get_file_path(key))
-        if isinstance(buffers[0], tuple):
+        return (0, 0, 0, key)
-            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 register_files(
+    def _xfer_and_wait(
-        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(
         self,
-        buffers: Optional[List[torch.Tensor | tuple]],
+        host_descs: Any,
-        keys: List[str],
+        storage_descs: Any,
         direction: str,
     ) -> bool:
-        if len(buffers) != len(keys):
+        """Initialize and poll a NIXL transfer to completion."""
-            logger.error("Mismatch between number of tensors/buffers and files/objects")
-            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
-
-            # 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:
-                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
@@ -194,134 +143,87 @@ class HiCacheNixl(HiCacheStorage):
             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(
+                # Best would be to have a better notification mechanism from NIXL,
-                        0.0001
+                # but we only have polling for now.
-                    )  # Can be changed to os.sched_yield() or parametrized
+                time.sleep(0.0001)
-
-                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.agent.release_xfer_handle(xfer_req)
-                self.file_manager.close_file(fd)
+
+    def _xfer_pre_registered(
+        self,
+        host_buffers: List[tuple],
+        keys: List[str],
+        direction: str,
+    ) -> bool:
+        """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
+
+        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
+
+        with self.registry.storage(host_buffers, keys, direction) as storage_descs:
+            if storage_descs is None:
+                return False
+            return self._xfer_and_wait(host_descs, storage_descs, direction)
 
     def get(
         self,
         key: str,
-        target_location: Optional[torch.Tensor | int] = None,
+        target_location: Optional[Any] = None,
-        target_sizes: Optional[int] = None,
+        target_sizes: Optional[Any] = None,
     ) -> torch.Tensor | None:
-        # To be removed, being compatible with the current API
+        raise NotImplementedError("deprecated; use batch_get_v1")
-        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
 
     def batch_get(
         self,
         keys: List[str],
-        target_locations: Optional[List[torch.Tensor | int]] = None,
+        target_locations: Optional[Any] = None,
-        target_sizes: Optional[List[int]] = None,
+        target_sizes: Optional[Any] = None,
     ) -> List[torch.Tensor | None]:
-        if not keys:
+        raise NotImplementedError("deprecated; use batch_get_v1")
-            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)
 
     def set(
         self,
         key: str,
-        value: Optional[torch.Tensor] = None,
+        value: Optional[Any] = None,
-        target_location: Optional[int] = None,
+        target_location: Optional[Any] = None,
-        target_sizes: Optional[int] = None,
+        target_sizes: Optional[Any] = None,
     ) -> bool:
-        if target_location and target_sizes:
+        raise NotImplementedError("deprecated; use batch_set_v1")
-            return self.batch_set([key], None, [target_location], [target_sizes])
-        else:
-            return self.batch_set([key], [value])
 
     def batch_set(
         self,
         keys: List[str],
-        values: Optional[List[torch.Tensor]] = None,
+        values: Optional[Any] = None,
-        target_locations: Optional[List[int]] = None,
+        target_locations: Optional[Any] = None,
-        target_sizes: Optional[List[int]] = None,
+        target_sizes: Optional[Any] = None,
     ) -> bool:
-
+        raise NotImplementedError("deprecated; use batch_set_v1")
-        # 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()
 
     def register_mem_pool_host(self, mem_pool_host: HostKVCache):
         super().register_mem_pool_host(mem_pool_host)
@@ -332,9 +234,95 @@ class HiCacheNixl(HiCacheStorage):
             "page_first_direct",
         ]
 
-        logger.info(
+        if self.needs_page_alignment and self.is_zero_copy:
-            f"HiCacheNixl: Registered mem_pool_host with layout {self.mem_pool_host.layout}, zero_copy set to {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: 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])
@@ -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 = [self._create_query_tuple(key) for key in key_list]
-        tuples = []
-        for key in key_list:
-            tuples += self.registration.create_query_tuples(
-                key,
-                self.backend_selector.mem_type,
-                self.file_manager if self.backend_selector.mem_type == "FILE" else None,
-            )
 
         query_res = self.agent.query_memory(
             tuples,
@@ -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:
-        for key_ in keys:
+            suffixed_key = self._get_suffixed_key(key)
-            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")
+            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,206 +379,162 @@ 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):
+    def _bounce_slot_buffers(self, buf: torch.Tensor, page_num: int) -> List[tuple]:
-        page_num = len(host_indices) // self.mem_pool_host.page_size
+        """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 = (
+            key_list, ptr_list, size_list = self._get_location_and_size_list_from_meta(
-                self._get_location_and_size_list_from_meta(keys, host_indices)
+                keys, host_indices
             )
-            return key_list, [], ptr_list, element_size_list
+            host_buffers = list(zip(ptr_list, size_list))
-        else:
+            return key_list, host_buffers
-            # non zero copy: create contiguous, temporary tensors
-            target_tensors = [
-                self.mem_pool_host.get_dummy_flat_data_page() for i in range(page_num)
-            ]
 
+        if page_num > STORAGE_BATCH_SIZE:
+            logger.error(
+                f"HiCacheNixl: batch size {page_num} exceeds bounce buffer capacity {STORAGE_BATCH_SIZE}"
+            )
+            return [], []
+
+        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)
+
+        host_buffers = self._bounce_slot_buffers(bounce, page_num)
         key_list = [self._get_suffixed_key(key) for key in keys]
-            ptr_list = [tensor.data_ptr() for tensor in target_tensors]
+        return key_list, host_buffers
-            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_xfer(
-
-    def _batch_get_zero_copy_impl(
         self,
         keys: List[str],
         key_strs: List[str],
-        target_tensors: List[torch.Tensor],
+        host_buffers: List[tuple],
-        target_locations: List[int],
+        direction: str,
-        target_sizes: List[int],
+    ) -> List[bool]:
-    ) -> List[int]:
+        """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 target_locations or not target_sizes:
+        """
+        if not key_strs or not host_buffers:
             return [False] * len(keys)
 
-        if (len(key_strs) != len(target_locations)) or (
+        if len(key_strs) != len(host_buffers):
-            len(target_sizes) != len(target_locations)
+            logger.error("Mismatch between number of key_strs and host_buffers")
-        ):
-            logger.error(
-                "Mismatch between number of key_strs, target_locations and target_sizes"
-            )
             return [False] * len(keys)
 
-        if self.is_zero_copy:
-            dest = list(zip(target_locations, target_sizes))
-        else:
-            dest = target_tensors
-
         if self.backend_selector.mem_type == "FILE":
             file_paths = [self.file_manager.get_file_path(key) for key in key_strs]
-            success = self._execute_transfer(dest, file_paths, "READ")
+            success = self._xfer_pre_registered(host_buffers, file_paths, direction)
-        else:
+        else:  # mem_type == "OBJ"
-            success = self._execute_transfer(dest, key_strs, "READ")
+            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_size = self.mem_pool_host.page_size
-        page_num = len(host_indices) // 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:
+            return [(results[2 * i] and results[2 * i + 1]) for i in range(page_num)]
-                results = [
+
-                    (results[2 * i] and results[2 * i + 1]) for i in range(page_num)
+        # non zero copy: copy data from the get-side bounce buffer to mem_pool_host
-                ]
-                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]
+                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,
         keys: List[str],
         host_indices: torch.Tensor,
         extra_info: Optional[HiCacheStorageExtraInfo] = None,
     ) -> List[bool]:
-
+        if not self._host_regs:
-        key_strs, target_tensors, buffer_ptrs, buffer_sizes = (
-            self._batch_get_preprocess(keys, host_indices)
-        )
-
-        if not key_strs or not buffer_ptrs or not buffer_sizes:
             logger.error(
-                "HiCacheNixl batch_get_v1: preprocessing failed, empty key_strs, buffer_ptrs or buffer_sizes"
+                "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 = self._batch_xfer(keys, key_strs, host_buffers, "READ")
-        results_get = self._batch_get_zero_copy_impl(
+        elapsed_ms = (time.perf_counter() - start_time) * 1000
-            keys, key_strs, target_tensors, buffer_ptrs, buffer_sizes
+        self._log_xfer_stats(
+            "batch_get_v1",
+            len(keys),
+            host_indices,
+            [s for _, s in host_buffers],
+            elapsed_ms,
         )
 
-        end_time = time.perf_counter()
+        return self._batch_get_postprocess(host_indices, results)
-        elapsed_time_ms = (end_time - start_time) * 1000
-        total_bytes = sum(s for s in buffer_sizes if s is not None)
-
-        logger.debug(
-            f"HiCacheNixl batch_get_v1 transferred: {len(keys)} keys (pages), {host_indices.numel()} host_indices, {total_bytes} bytes, total time: {elapsed_time_ms:.3f} ms, effective bandwidth: {total_bytes / (elapsed_time_ms / 1000) / (1024 * 1024):.2f} MB/s"
-        )
-
-        return self._batch_get_postprocess(host_indices, target_tensors, results_get)
-
-    def _batch_set_preprocess(self, keys: List[str], host_indices: torch.Tensor):
-
-        page_num = len(host_indices) // self.mem_pool_host.page_size
-
-        if len(keys) == 0 or len(keys) != page_num:
-            logger.warning(
-                f"HiCacheNixl: empty keys or mismatch in keys and host_indices lengths. keys: {len(keys)}, host_indices: {len(host_indices)}, page_size: {self.mem_pool_host.page_size}"
-            )
-            return [], [], [], []
-
-        if self.is_zero_copy:
-            key_list, _, ptr_list, element_size_list = (
-                self._get_location_and_size_list_from_meta(keys, host_indices)
-            )
-            return key_list, [], ptr_list, element_size_list
-        else:
-            # non zero copy: NIXL still requires contiguous tensors for transfer
-            target_tensors = [
-                self.mem_pool_host.get_data_page(
-                    host_indices[i * self.mem_pool_host.page_size], flat=False
-                ).contiguous()
-                for i in range(page_num)
-            ]
-
-            key_list = [self._get_suffixed_key(key) for key in keys]
-            ptr_list = [tensor.data_ptr() for tensor in target_tensors]
-            element_size_list = [
-                tensor.numel() * tensor.element_size() for tensor in target_tensors
-            ]
-
-            return key_list, target_tensors, ptr_list, element_size_list
-
-    def _batch_set_zero_copy_impl(
-        self,
-        keys: List[str],
-        key_strs: List[str],
-        target_tensors: List[torch.Tensor],
-        target_locations: List[int],
-        target_sizes: List[int],
-    ) -> List[bool]:
-
-        if not key_strs or not target_locations or not target_sizes:
-            return [False] * len(keys)
-
-        if (len(key_strs) != len(target_locations)) or (
-            len(target_sizes) != len(target_locations)
-        ):
-            logger.error(
-                "Mismatch between number of key_strs, target_locations and target_sizes"
-            )
-            return [False] * len(keys)
-
-        if self.is_zero_copy:
-            src = list(zip(target_locations, target_sizes))
-        else:
-            src = target_tensors
-
-        if self.backend_selector.mem_type == "FILE":
-            file_paths = []
-            for key in key_strs:
-                file_path = self.file_manager.get_file_path(key)
-                # New file per set, to be updated when partial writes is added to HiCache
-                if not self.file_manager.create_file(file_path):
-                    logger.error(
-                        f"******** Failed to create file {file_path} *********"
-                    )
-                    return [False] * len(keys)
-                file_paths.append(file_path)
-            success = self._execute_transfer(src, file_paths, "WRITE")
-        else:  # mem_type == "OBJ"
-            success = self._execute_transfer(src, key_strs, "WRITE")
-
-        return [True] * len(keys) if success else [False] * len(keys)
 
     def batch_set_v1(
         self,
@@ -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 = (
+        if not self._host_regs:
-            self._batch_set_preprocess(keys, host_indices)
-        )
-
-        if not key_strs or not buffer_ptrs or not buffer_sizes:
             logger.error(
-                "HiCacheNixl batch_set_v1: preprocessing failed, empty key_strs, buffer_ptrs or buffer_sizes"
+                "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 = self._batch_xfer(keys, key_strs, host_buffers, "WRITE")
-        results_set = self._batch_set_zero_copy_impl(
+        elapsed_ms = (time.perf_counter() - start_time) * 1000
-            keys, key_strs, target_tensors, buffer_ptrs, buffer_sizes
+        self._log_xfer_stats(
+            "batch_set_v1",
+            len(keys),
+            host_indices,
+            [s for _, s in host_buffers],
+            elapsed_ms,
         )
 
-        end_time = time.perf_counter()
+        return results
-        elapsed_time_ms = (end_time - start_time) * 1000
-        total_bytes = sum(s for s in buffer_sizes if s is not None)
-        logger.debug(
-            f"HiCacheNixl batch_set_v1 transferred: {len(keys)} keys (pages), {host_indices.numel()} host_indices, {total_bytes} bytes, total time: {elapsed_time_ms:.3f} ms, effective bandwidth: {total_bytes / (elapsed_time_ms / 1000) / (1024 * 1024):.2f} MB/s"
-        )
-
-        return results_set

python/sglang/srt/mem_cache/storage/nixl/nixl_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,
+                )

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:
+    def open_file(self, file_path: str, create: bool = False) -> Optional[int]:
-        """Create a file if it doesn't exist."""
+        """Open a file and return its file descriptor.
-        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) -> Optional[int]:
+        If ``create`` is True, the file is created if it does not exist
-        """Open a file and return its file descriptor."""
+        (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 os.open(file_path, flags, 0o644)
-            return fd
         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

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

@@ -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()

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()