Refactor NIXL hicache. Add O_DIRECT support (#25173)

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

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):
-            batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
+        for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE):
+            batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE]
             batch_host_indices = operation.host_indices[
                 i * self.page_size : (i + len(batch_hashes)) * self.page_size
             ]
@@ -978,11 +977,9 @@ class HiCacheController:
         hash_value = []
 
         for start in range(
-            0, len(tokens_to_fetch), self.page_size * self.storage_batch_size
+            0, len(tokens_to_fetch), self.page_size * STORAGE_BATCH_SIZE
         ):
-            end = min(
-                start + self.page_size * self.storage_batch_size, len(tokens_to_fetch)
-            )
+            end = min(start + self.page_size * STORAGE_BATCH_SIZE, len(tokens_to_fetch))
             batch_tokens = tokens_to_fetch[start:end]
             batch_hashes = []
             for i in range(0, len(batch_tokens), self.page_size):
@@ -1120,8 +1117,8 @@ class HiCacheController:
     def _page_backup(self, operation):
         # Backup batch by batch
         prefix_keys = operation.prefix_keys
-        for i in range(0, len(operation.hash_value), self.storage_batch_size):
-            batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
+        for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE):
+            batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE]
             batch_host_indices = operation.host_indices[
                 i * self.page_size : (i + len(batch_hashes)) * self.page_size
             ]

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 tensor
+        return alloc_mmap(dims, dtype)
 
 
 class HiSparseHostPoolMixin:
@@ -187,11 +189,16 @@ def alloc_with_host_register(
     """
     buffer = allocator.allocate(dims, dtype=dtype, device=device)
     if pin_memory:
-        ret = torch.cuda.cudart().cudaHostRegister(
-            buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0
-        )
-        if ret != 0:
-            raise RuntimeError(f"cudaHostRegister failed with error code {ret}")
+        cudart = torch.cuda.cudart()
+        n_bytes = buffer.numel() * buffer.element_size()
+        rc = cudart.cudaHostRegister(buffer.data_ptr(), n_bytes, 0)
+        if int(rc) != 0:
+            raise RuntimeError(
+                f"cudaHostRegister failed (rc={int(rc)}, "
+                f"{cudart.cudaGetErrorString(rc)}) for ptr={buffer.data_ptr():#x} "
+                f"size={n_bytes}; host buffer is not pinned and device transfers "
+                f"may silently return stale data."
+            )
     return buffer
 
 
@@ -324,6 +331,19 @@ class HostKVCache(abc.ABC):
         """
         raise NotImplementedError()
 
+    def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
+        """Return True if per-page strides are multiples of *page_size_bytes*.
+
+        Subclasses should override this with a layout-specific stride formula.
+        This base implementation logs a warning and returns False (safe default).
+        """
+        logger.warning(
+            "%s does not implement is_stride_page_aligned(); assuming not aligned. "
+            "O_DIRECT with a file-based NIXL backend will fall back to copy mode for this pool.",
+            type(self).__name__,
+        )
+        return False
+
     @synchronized
     def clear(self):
         # Initialize memory states and tracking structures.
@@ -850,6 +870,30 @@ class MHATokenToKVPoolHost(HostKVCache):
             raise ValueError(f"Unsupported layout: {self.layout}")
         return ptr_list, element_size_list
 
+    def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
+        """Return True if per-page strides are multiples of *page_size_bytes*.
+
+        When O_DIRECT is used with any file-based NIXL backend, every data pointer
+        passed to the kernel must be page-aligned.  In zero-copy mode the
+        pointer for KV page ``p`` is:
+
+            base_ptr + p * page_size * layer_num * head_num * head_dim * itemsize
+
+        For this to be page-aligned (given a page-aligned ``base_ptr``) the per-page
+        stride must itself be a multiple of the OS page size.
+        """
+        if self.layout not in ("page_first", "page_first_direct", "page_head"):
+            return False
+        stride = (
+            self.page_size
+            * self.layer_num
+            * self.head_num
+            * self.head_dim
+            * self.dtype.itemsize
+        )
+        base_aligned = self.kv_buffer.data_ptr() % page_size_bytes == 0
+        return base_aligned and stride % page_size_bytes == 0
+
 
 class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
     device_pool: MLATokenToKVPool
@@ -1250,6 +1294,26 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
             raise ValueError(f"Unsupported layout: {self.layout}")
         return ptr_list, element_size_list
 
+    def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
+        """Return True if per-page strides are multiples of *page_size_bytes*.
+
+        When O_DIRECT is used with any file-based NIXL backend, every data pointer
+        passed to the kernel must be page-aligned.  In zero-copy mode the
+        pointer for KV page ``p`` is:
+
+            base_ptr + p * page_size * layer_num * kv_cache_dim * itemsize
+
+        For this to be page-aligned (given a page-aligned ``base_ptr``) the per-page
+        stride must itself be a multiple of the OS page size.
+        """
+        if self.layout not in ("page_first", "page_first_direct"):
+            return False
+        stride = (
+            self.page_size * self.layer_num * self.kv_cache_dim * self.dtype.itemsize
+        )
+        base_aligned = self.kv_buffer.data_ptr() % page_size_bytes == 0
+        return base_aligned and stride % page_size_bytes == 0
+
 
 class MambaPoolHost(HostKVCache):
 

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 and NIXL tuple creation
+- **NixlFileManager** - Handles file system operations
+
+### NixlRegistry (`nixl_registry.py`)
+Owns the `(agent, mem_type, file_manager)` triple and exposes `host(...)` and `storage(...)` context managers that register on entry, yield the NIXL `xfer_descs`, and deregister + close fds on exit. Internally composes two single-resource primitives (`_open_files` and `_registered`) so leak-freeness is verifiable per primitive.
 
 The current implementation performs per-transfer registration for file / object targets and explicitly closes FILE descriptors after registration / transfer setup to avoid descriptor leaks.
 
@@ -149,7 +151,7 @@ For debugging or quick testing, you may pass a **JSON-style string** directly vi
 
 This requires explicitly specifying the plugin type via an environment variable, and this method can be applicable to **only a few** plugins (e.g., POSIX, GDS, GDS_MT)
 
-The below example shows how to use command-line string to use the POSIX plugin where URING is enabled for async POSIX storage.
+The below example shows how to use command-line string to use the POSIX plugin where URING is enabled for async POSIX storage, with O_DIRECT enabled (the default).
 
 ```bash
 export SGLANG_HICACHE_NIXL_BACKEND_PLUGIN=POSIX
@@ -164,7 +166,17 @@ python3 -m sglang.launch_server \
   --hicache-size 64 \
   --hicache-write-policy write_through \
   --hicache-storage-backend nixl \
-  --hicache-storage-backend-extra-config "{'use_uring': 'true'}"
+  --hicache-storage-backend-extra-config '{"use_uring": "true"}'
+```
+
+To disable O_DIRECT (e.g. for debugging or unsupported filesystems), set the top-level `use_direct_io` key:
+
+```bash
+export SGLANG_HICACHE_NIXL_BACKEND_PLUGIN=POSIX
+
+python3 -m sglang.launch_server \
+  ... \
+  --hicache-storage-backend-extra-config '{"use_direct_io": false, "use_uring": "true"}'
 ```
 
 ⚠️ **Note**:
@@ -349,6 +361,33 @@ An example of the configuration is provided in [`nixl.config.toml.sample`](./nix
 
 For object storage, `bucket` may also be omitted from the config if `AWS_DEFAULT_BUCKET` is already defined in the environment.
 
+### 1a. Top-Level Configuration Keys
+
+The following keys are placed at the **top level** of the config file (not inside any `[plugin.*]` section) and apply globally to the NIXL backend:
+
+| Key              | Type    | Default  | Description |
+| ---------------- | ------- | -------- | ----------- |
+| `use_direct_io`  | boolean | `true`   | Open cache files with `O_DIRECT` to bypass the OS page cache. Reduces memory pressure and improves NVMe throughput. Falls back to buffered I/O with a warning if `O_DIRECT` is unavailable on the current OS. Can also be overridden via the `SGLANG_HICACHE_NIXL_USE_DIRECT_IO` environment variable. |
+
+**Page-alignment and `O_DIRECT`**
+
+When `use_direct_io = true` with any file-based backend (POSIX, GDS, GDS_MT, 3FS), the kernel requires every I/O buffer pointer to be OS-page-aligned (4 KiB). SGLang handles this automatically:
+
+* **Zero-copy mode** (`page_first` / `page_first_direct` layout): the host memory pool is always mmap-backed and therefore page-aligned. If the per-page stride is also a multiple of 4 KiB, zero-copy transfers are used as-is.
+* **Copy mode** (all other layouts, or if stride alignment cannot be satisfied): SGLang pre-allocates page-aligned bounce buffers via `mmap` and falls back to copy mode, logging a warning. No user action is required -- this is fully automatic.
+
+To disable `O_DIRECT` (e.g. for debugging or when the filesystem does not support it):
+
+```toml
+use_direct_io = false
+
+[plugin.posix]
+use_uring = "true"
+active = true
+```
+
+or via environment variable: `SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0`.
+
 
 ### 2. POSIX File System Backend (`plugin.posix`)
 

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

@@ -1,27 +1,25 @@
 import logging
 import time
 import uuid
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional
 
 import torch
 
 from sglang.srt.environ import envs
 from sglang.srt.mem_cache.hicache_storage import (
+    STORAGE_BATCH_SIZE,
     HiCacheStorage,
     HiCacheStorageConfig,
     HiCacheStorageExtraInfo,
 )
 from sglang.srt.mem_cache.memory_pool_host import HostKVCache
+from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
 
-from .nixl_utils import (
-    NixlBackendConfig,
-    NixlBackendSelection,
-    NixlFileManager,
-    NixlRegistration,
-)
+from .nixl_registry import NixlRegistry
+from .nixl_utils import NixlBackendConfig, NixlBackendSelection, NixlFileManager
 
 try:
-    from nixl._api import nixl_agent, nixl_agent_config
+    from nixl._api import nixl_agent, nixl_agent_config, nixlBind
 except ImportError as e:
     raise ImportError(
         "Please install NIXL by following the instructions at "
@@ -48,15 +46,16 @@ class HiCacheNixl(HiCacheStorage):
         # select the NIXL backend plugin from extra_config or environment variable
         plugin = nixlconfig.get_specified_plugin()
 
+        use_direct_io = nixlconfig.get_use_direct_io()
+
         # Might be better to be unified across HiCache backends and moved to HiCacheController
         file_path = envs.SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR.get() or file_path
         self.file_manager = (
-            NixlFileManager(file_path)
+            NixlFileManager(file_path, use_direct_io=use_direct_io)
             if plugin not in NixlBackendSelection.OBJ_PLUGINS
             else None
         )
 
-        # Initialize suffix based on storage config
         tp_rank, tp_size, model_name = (
             storage_config.tp_rank,
             storage_config.tp_size,
@@ -75,253 +74,156 @@ class HiCacheNixl(HiCacheStorage):
         else:
             self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}"
 
+        sync_mode = getattr(
+            nixlBind, "NIXL_THREAD_SYNC_RW", nixlBind.NIXL_THREAD_SYNC_STRICT
+        )
         agent_config = nixl_agent_config(backends=[])
         self.agent_name = f"hicache_nixl_{str(uuid.uuid4())}"
         self.agent = nixl_agent(self.agent_name, agent_config)
+        bind_cfg = nixlBind.nixlAgentConfig()
+        bind_cfg.useProgThread = agent_config.enable_pthread
+        bind_cfg.useListenThread = agent_config.enable_listen
+        bind_cfg.listenPort = agent_config.port
+        bind_cfg.syncMode = sync_mode
+        bind_cfg.pthrDelay = 0
+        bind_cfg.lthrDelay = 100000
+        bind_cfg.captureTelemetry = agent_config.capture_telemetry
+        self.agent.agent = nixlBind.nixlAgent(self.agent_name, bind_cfg)
+        self.agent.plugin_list = self.agent.agent.getAvailPlugins()
 
         self.backend_selector = NixlBackendSelection(plugin, nixlconfig)
         if not self.backend_selector.create_backend(self.agent):
             raise RuntimeError("Failed to create NIXL backend")
 
-        self.registration = NixlRegistration(self.agent)
-        self.is_zero_copy = False
+        self.registry = NixlRegistry(
+            self.agent,
+            self.backend_selector.mem_type,
+            self.file_manager,
+        )
+        # O_DIRECT requires OS-page-aligned I/O buffers on all file-based backends
+        # (POSIX, GDS, GDS_MT, 3FS). OBJ backends never open files so they are exempt
+        # (file_manager is None for OBJ).
+        self.needs_page_alignment = use_direct_io and self.file_manager is not None
+        if self.needs_page_alignment:
+            logger.info(
+                "HiCacheNixl: O_DIRECT is active with a file-based backend (%s). "
+                "Page-aligned host buffers are required (needs_page_alignment=True).",
+                self.backend_selector.backend_name,
+            )
+        # Pre-registered host regions (set by register_mem_pool_host):
+        # zero-copy: one registration covering mem_pool_host.kv_buffer
+        # non-zero-copy: two registrations, one bounce buffer per direction
+        # (set/get) so the two storage threads never share slots.
+        self._host_regs: List[Any] = []
+        self._bounce_set: Optional[torch.Tensor] = None
+        self._bounce_get: Optional[torch.Tensor] = None
+        self._bounce_page_bytes: Optional[int] = None
 
     def _get_suffixed_key(self, key: str) -> str:
         return key + self.config_suffix
 
-    def register_buffers(
-        self, buffers: Union[torch.Tensor, List[torch.Tensor], List[tuple]]
-    ) -> Optional[Any]:
-        """Register tensor(s) or target locations in host memory (list of addr,len tuples) with NIXL."""
-        if isinstance(buffers[0], tuple):
-            tuples = [(x[0], x[1], 0, "") for x in buffers]
-            return self.registration._register_memory(tuples, "DRAM")
-        else:
-            return self.registration._register_memory(buffers)
+    def _create_query_tuple(self, key: str) -> tuple:
+        """Build the NIXL query_memory tuple for a single key."""
+        if self.backend_selector.mem_type == "FILE":
+            return (0, 0, 0, self.file_manager.get_file_path(key))
+        return (0, 0, 0, key)
 
-    def register_files(
-        self, file_paths: List[str], open_file: Optional[bool] = True
-    ) -> Optional[Any]:
-        """Register files with NIXL."""
-        tuples = self.file_manager.files_to_nixl_tuples(file_paths)
-        return self.registration._register_memory(tuples, "FILE")
-
-    def register_objects(
-        self, keys: List[str], sizes: Optional[List[int]] = None
-    ) -> Optional[Any]:
-        """Register objects with NIXL."""
-        if not keys:
-            return None
-        tuples = [(0, 0, key, "") for key in keys]
-        return self.registration._register_memory(tuples, "OBJ")
-
-    def _execute_transfer(
+    def _xfer_and_wait(
         self,
-        buffers: Optional[List[torch.Tensor | tuple]],
+        host_descs: Any,
+        storage_descs: Any,
+        direction: str,
+    ) -> bool:
+        """Initialize and poll a NIXL transfer to completion."""
+        try:
+            xfer_req = self.agent.initialize_xfer(
+                direction, host_descs, storage_descs, self.agent_name
+            )
+        except Exception as e:
+            logger.error(f"Failed to create transfer request: {e}")
+            return False
+
+        try:
+            state = self.agent.transfer(xfer_req)
+            while state != "DONE":
+                state = self.agent.check_xfer_state(xfer_req)
+                if state == "ERR":
+                    logger.error("Transfer failed")
+                    return False
+                # Best would be to have a better notification mechanism from NIXL,
+                # but we only have polling for now.
+                time.sleep(0.0001)
+            return True
+        except Exception as e:
+            logger.error(f"Failed to execute transfer: {e}")
+            import traceback
+
+            logger.error(f"Traceback: {traceback.format_exc()}")
+            return False
+        finally:
+            self.agent.release_xfer_handle(xfer_req)
+
+    def _xfer_pre_registered(
+        self,
+        host_buffers: List[tuple],
         keys: List[str],
         direction: str,
     ) -> bool:
-        if len(buffers) != len(keys):
-            logger.error("Mismatch between number of tensors/buffers and files/objects")
+        """Run a transfer where the host side is already pre-registered.
+
+        ``host_buffers`` is a list of ``(addr, size)`` tuples within the
+        pre-registered host region (kv_buffer for zero-copy, bounce buffer
+        otherwise). Only the storage side is registered per transfer.
+        """
+        if len(host_buffers) != len(keys):
+            logger.error("Mismatch between number of host buffers and keys")
             return False
 
-        # Registering file and object keys per transfer, to be updated when
-        # pre-registration for file and object is added to HiCache.
-        file_fds = []
-        try:
-            if self.backend_selector.mem_type == "FILE":
-                tuples = self.file_manager.files_to_nixl_tuples(keys)
-                file_fds = [t[2] for t in tuples]
-                if not tuples or not self.registration._register_memory(tuples, "FILE"):
-                    logger.error("Failed to prepare files for transfer")
-                    return False
-            else:  # mem_type == "OBJ"
-                tuples = [(0, 0, key, "") for key in keys]
-                if not tuples or not self.registration._register_memory(tuples, "OBJ"):
-                    logger.error("Failed to register objects")
-                    return False
+        host_descs = self.agent.get_xfer_descs(
+            [(addr, size, 0) for (addr, size) in host_buffers], "DRAM"
+        )
+        if host_descs is None:
+            logger.error("Failed to build host xfer descs")
+            return False
 
-            # Prepare transfer descriptors
-            if isinstance(buffers[0], torch.Tensor):
-                tensor_sizes = [
-                    tensor.element_size() * tensor.numel() for tensor in buffers
-                ]
-                storage_tuples = [(x[0], s, x[2]) for x, s in zip(tuples, tensor_sizes)]
-                host_descs = self.agent.get_xfer_descs(buffers)
-
-                if direction in ("READ", "WRITE"):
-                    # register buffer to avoid calling initialize_xfer twice due to missing registration
-                    self.register_buffers(buffers)
-
-            elif isinstance(buffers[0], tuple):
-                storage_tuples = [(x[0], y[1], x[2]) for x, y in zip(tuples, buffers)]
-                host_descs = self.agent.get_xfer_descs(
-                    [(x[0], x[1], 0) for x in buffers], "DRAM"
-                )
-
-                if direction in ("READ", "WRITE"):
-                    # register buffer to avoid calling initialize_xfer twice due to missing registration
-                    self.register_buffers(buffers)
-
-            else:
+        with self.registry.storage(host_buffers, keys, direction) as storage_descs:
+            if storage_descs is None:
                 return False
-
-            storage_descs = self.agent.get_xfer_descs(
-                storage_tuples, self.backend_selector.mem_type
-            )
-
-            if (host_descs is None) or (storage_descs is None):
-                logger.error("Failed to get transfer descriptors")
-                return False
-
-            # Initialize transfer, default assumption that tensor was registered
-
-            try:
-                xfer_req = self.agent.initialize_xfer(
-                    direction, host_descs, storage_descs, self.agent_name
-                )
-            except Exception:
-                # Check if it was due to missing pre-registration
-                if not self.register_buffers(buffers):
-                    logger.error("Failed to register tensors/buffers")
-                    return False
-
-                try:
-                    xfer_req = self.agent.initialize_xfer(
-                        direction, host_descs, storage_descs, self.agent_name
-                    )
-                except Exception as e:
-                    logger.error(f"Failed to create transfer request: {e}")
-                    return False
-
-            # Execute transfer and wait for its completion
-            try:
-                state = self.agent.transfer(xfer_req)
-                while state != "DONE":
-                    state = self.agent.check_xfer_state(xfer_req)
-                    if state == "ERR":
-                        self.agent.release_xfer_handle(xfer_req)
-                        logger.error("Transfer failed")
-                        return False
-                    time.sleep(
-                        0.0001
-                    )  # Can be changed to os.sched_yield() or parametrized
-
-                self.agent.release_xfer_handle(xfer_req)
-                return True
-
-            except Exception as e:
-                logger.error(f"Failed to execute transfer: {e}")
-                import traceback
-
-                logger.error(f"Traceback: {traceback.format_exc()}")
-                return False
-
-        finally:
-            for fd in file_fds:
-                self.file_manager.close_file(fd)
+            return self._xfer_and_wait(host_descs, storage_descs, direction)
 
     def get(
         self,
         key: str,
-        target_location: Optional[torch.Tensor | int] = None,
-        target_sizes: Optional[int] = None,
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
     ) -> torch.Tensor | None:
-        # To be removed, being compatible with the current API
-        if target_location is None:
-            return None
-        if target_sizes:
-            result = self.batch_get([key], [target_location], [target_sizes])
-        else:
-            result = self.batch_get([key], [target_location])
-        return result[0] if result else None
+        raise NotImplementedError("deprecated; use batch_get_v1")
 
     def batch_get(
         self,
         keys: List[str],
-        target_locations: Optional[List[torch.Tensor | int]] = None,
-        target_sizes: Optional[List[int]] = None,
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
     ) -> List[torch.Tensor | None]:
-        if not keys:
-            return []
-
-        # To be removed, being compatible with the current API
-        if not target_locations:
-            return [None] * len(keys)
-
-        if target_sizes and (len(target_sizes) != len(target_locations)):
-            logger.error("Mismatch between number of target_locations and target_sizes")
-            return [None] * len(keys)
-
-        if target_sizes:
-            dest = list(zip(target_locations, target_sizes))
-        else:
-            dest = target_locations
-
-        # Add suffix to keys
-        suffixed_keys = [self._get_suffixed_key(key) for key in keys]
-
-        if self.backend_selector.mem_type == "FILE":
-            file_paths = [self.file_manager.get_file_path(key) for key in suffixed_keys]
-            success = self._execute_transfer(dest, file_paths, "READ")
-        else:
-            success = self._execute_transfer(dest, suffixed_keys, "READ")
-        return target_locations if success and not target_sizes else [None] * len(keys)
+        raise NotImplementedError("deprecated; use batch_get_v1")
 
     def set(
         self,
         key: str,
-        value: Optional[torch.Tensor] = None,
-        target_location: Optional[int] = None,
-        target_sizes: Optional[int] = None,
+        value: Optional[Any] = None,
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
     ) -> bool:
-        if target_location and target_sizes:
-            return self.batch_set([key], None, [target_location], [target_sizes])
-        else:
-            return self.batch_set([key], [value])
+        raise NotImplementedError("deprecated; use batch_set_v1")
 
     def batch_set(
         self,
         keys: List[str],
-        values: Optional[List[torch.Tensor]] = None,
-        target_locations: Optional[List[int]] = None,
-        target_sizes: Optional[List[int]] = None,
+        values: Optional[Any] = None,
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
     ) -> bool:
-
-        # skip on MLA backup rank
-        if self.backup_skip:
-            return True
-
-        if not keys or (not values and (not target_locations or not target_sizes)):
-            logger.error("Keys or values were not passed")
-            return False
-
-        if not values:
-            values = list(zip(target_locations, target_sizes))
-
-        # Add suffix to keys
-        suffixed_keys = [self._get_suffixed_key(key) for key in keys]
-
-        if self.backend_selector.mem_type == "FILE":
-            file_paths = []
-            for key in suffixed_keys:
-                file_path = self.file_manager.get_file_path(key)
-                # New file per set, to be updated when partial writes is added to HiCache
-                if not self.file_manager.create_file(file_path):
-                    logger.error(f"Failed to create file {file_path}")
-                    return False
-                file_paths.append(file_path)
-            return self._execute_transfer(values, file_paths, "WRITE")
-        else:  # mem_type == "OBJ"
-            return self._execute_transfer(values, suffixed_keys, "WRITE")
-
-    ############################################################################
-    # batch_*_v1 functions
-    # zero copy + non-zero-copy version for get, set, exists, batch_exists
-    ############################################################################
-
-    def clear(self) -> None:
-        self.file_manager.clear()
+        raise NotImplementedError("deprecated; use batch_set_v1")
 
     def register_mem_pool_host(self, mem_pool_host: HostKVCache):
         super().register_mem_pool_host(mem_pool_host)
@@ -332,10 +234,96 @@ class HiCacheNixl(HiCacheStorage):
             "page_first_direct",
         ]
 
+        if self.needs_page_alignment and self.is_zero_copy:
+            # Check that the kv_buffer base AND per-page strides are multiples of
+            # the OS page size so every pointer passed to NIXL (base + p * stride)
+            # is page-aligned. The base is whatever torch.empty() happened to give
+            # us -- it is not guaranteed to be page-aligned. Fall back to copy mode
+            # if either condition fails.
+            # 4096: O_DIRECT alignment is FS-dependent (some allow 512 B); 4 KiB
+            # is the safe lower bound all known FSes accept, and real page-sizes meet it.
+            if not self.mem_pool_host.is_stride_page_aligned(4096):
+                logger.warning(
+                    "HiCacheNixl: O_DIRECT is active but the host kv_buffer is "
+                    "not OS-page-aligned (base or per-page stride). Falling back "
+                    "to copy mode for this pool."
+                )
+                self.is_zero_copy = False
+
+        if self.is_zero_copy:
+            kv = mem_pool_host.kv_buffer
+            self._pre_register_host(
+                kv.data_ptr(), kv.numel() * kv.element_size(), "kv_buffer"
+            )
+        else:
+            # One bounce buffer per direction so set/get run lock-free across
+            # the prefetch and backup threads. Sized from get_dummy_flat_data_page()
+            # so each slot matches what the v1 path would otherwise allocate.
+            sample = mem_pool_host.get_dummy_flat_data_page()
+            page_numel = sample.numel()
+            self._bounce_page_bytes = page_numel * sample.element_size()
+            del sample
+            pin_memory = bool(getattr(mem_pool_host, "pin_memory", False))
+            self._bounce_set = self._alloc_registered(
+                page_numel, mem_pool_host.dtype, pin_memory, "bounce_set"
+            )
+            self._bounce_get = self._alloc_registered(
+                page_numel, mem_pool_host.dtype, pin_memory, "bounce_get"
+            )
+
         logger.info(
-            f"HiCacheNixl: Registered mem_pool_host with layout {self.mem_pool_host.layout}, zero_copy set to {self.is_zero_copy}"
+            f"HiCacheNixl: pre-registered host regions for "
+            f"layout={mem_pool_host.layout} zero_copy={self.is_zero_copy}"
         )
 
+    def _alloc_registered(
+        self,
+        page_numel: int,
+        dtype: torch.dtype,
+        pin_memory: bool,
+        kind: str,
+    ) -> torch.Tensor:
+        """Allocate a ``(STORAGE_BATCH_SIZE, page_numel)`` bounce buffer and
+        pre-register it as a DRAM region with NIXL. Uses alloc_mmap so the
+        buffer is page-aligned -- required when O_DIRECT is on for any
+        file-based backend (POSIX/GDS/GDS_MT/3FS). pin_memory is currently
+        unused (alloc_mmap does not support it)."""
+        buf = alloc_mmap((STORAGE_BATCH_SIZE, page_numel), dtype)
+        self._pre_register_host(buf.data_ptr(), buf.numel() * buf.element_size(), kind)
+        return buf
+
+    def _pre_register_host(self, base_addr: int, total_size: int, kind: str) -> None:
+        """Register a single DRAM region up-front and remember the handle."""
+        reg_descs = self.agent.get_reg_descs([(base_addr, total_size, 0, "")], "DRAM")
+        if reg_descs is None:
+            raise RuntimeError(f"Failed to build reg descs for host {kind}")
+        try:
+            self._host_regs.append(self.agent.register_memory(reg_descs))
+        except Exception as e:
+            raise RuntimeError(f"Failed to pre-register host {kind} with NIXL") from e
+
+    def clear(self) -> None:
+        if self.file_manager is None:
+            return
+        self.file_manager.clear()
+
+    def close(self):
+        while self._host_regs:
+            reg = self._host_regs.pop()
+            try:
+                self.agent.deregister_memory(reg)
+            except Exception as e:
+                logger.debug("deregister of pre-registered host region failed: %s", e)
+        self._bounce_set = None
+        self._bounce_get = None
+        self._bounce_page_bytes = None
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
     def exists(self, key: str) -> bool:
         results = self.batch_exists([key])
         return results > 0
@@ -345,8 +333,6 @@ class HiCacheNixl(HiCacheStorage):
         keys: List[str],
         extra_info: Optional[HiCacheStorageExtraInfo] = None,
     ) -> int:
-        # Add suffix to key
-
         if self.is_zero_copy:
             key_list = self._get_key_list_from_meta(keys)
             key_denominator = (
@@ -356,14 +342,7 @@ class HiCacheNixl(HiCacheStorage):
             key_list = [self._get_suffixed_key(key) for key in keys]
             key_denominator = 1
 
-        # obtain list of tuples by calling self.registration.create_query_tuples()
-        tuples = []
-        for key in key_list:
-            tuples += self.registration.create_query_tuples(
-                key,
-                self.backend_selector.mem_type,
-                self.file_manager if self.backend_selector.mem_type == "FILE" else None,
-            )
+        tuples = [self._create_query_tuple(key) for key in key_list]
 
         query_res = self.agent.query_memory(
             tuples,
@@ -377,17 +356,14 @@ class HiCacheNixl(HiCacheStorage):
         return len(query_res) // key_denominator
 
     def _get_key_list_from_meta(self, keys: List[str]) -> List[str]:
-        # construct the key list for NIXL transfer based on the keys and the suffix, for each key, we will have one suffixed key for k buffer and one suffixed key for v buffer if it's not an MLA model, and only one suffixed key for k buffer if it's an MLA model, since MLA model only has k/v interleaved buffer
+        # Each key maps to a `_k` entry, plus a `_v` entry on non-MLA models
+        # (MLA stores k/v interleaved in a single buffer).
         key_list = []
-
-        for key_ in keys:
-            suffixed_key = self._get_suffixed_key(key_)
-            if self.is_mla_model:
-                key_list.append(f"{suffixed_key}_k")
-            else:
-                key_list.append(f"{suffixed_key}_k")
+        for key in keys:
+            suffixed_key = self._get_suffixed_key(key)
+            key_list.append(f"{suffixed_key}_k")
+            if not self.is_mla_model:
                 key_list.append(f"{suffixed_key}_v")
-
         return key_list
 
     def _get_location_and_size_list_from_meta(
@@ -403,99 +379,133 @@ class HiCacheNixl(HiCacheStorage):
             logger.error(
                 f"HiCacheNixl: mismatch between number of keys and number of buffer meta entries, keys: {len(keys)}, key_list: {len(key_list)}, buffer meta entries: {len(ptr_list)}"
             )
-            return [], [], [], []
+            return [], [], []
 
-        return key_list, [], ptr_list, element_size_list
+        return key_list, ptr_list, element_size_list
 
-    def _batch_get_preprocess(self, keys: List[str], host_indices: torch.Tensor):
-        page_num = len(host_indices) // self.mem_pool_host.page_size
+    def _bounce_slot_buffers(self, buf: torch.Tensor, page_num: int) -> List[tuple]:
+        """Return ``page_num`` ``(addr, size)`` tuples pointing at the first
+        ``page_num`` slots of ``buf``.
+        """
+        base = buf.data_ptr()
+        return [
+            (base + i * self._bounce_page_bytes, self._bounce_page_bytes)
+            for i in range(page_num)
+        ]
+
+    def _batch_preprocess(self, keys: List[str], host_indices: torch.Tensor, op: str):
+        """Build (key_list, host_buffers) for the v1 path.
+
+        For zero-copy: ``host_buffers`` are ``(addr, size)`` tuples inside the
+        pre-registered ``kv_buffer``.
+        For non-zero-copy: ``host_buffers`` are slots of the direction-specific
+        pre-registered bounce buffer (``_bounce_set`` for set, ``_bounce_get``
+        for get); for ``op == "set"`` we copy the host pages into those slots
+        here so the subsequent transfer reads from the bounce buffer.
+        Returns ``([], [])`` on validation failure.
+        """
+        page_size = self.mem_pool_host.page_size
+        page_num = len(host_indices) // page_size
 
         if len(keys) == 0 or len(keys) != page_num:
             logger.warning(
-                f"HiCacheNixl: empty keys or mismatch in keys and host_indices lengths. keys: {len(keys)}, host_indices: {len(host_indices)}, page_size: {self.mem_pool_host.page_size}"
+                f"HiCacheNixl: empty keys or mismatch in keys and host_indices lengths. keys: {len(keys)}, host_indices: {len(host_indices)}, page_size: {page_size}"
             )
-            return [], [], [], []
+            return [], []
 
         if self.is_zero_copy:
-            key_list, _, ptr_list, element_size_list = (
-                self._get_location_and_size_list_from_meta(keys, host_indices)
+            key_list, ptr_list, size_list = self._get_location_and_size_list_from_meta(
+                keys, host_indices
             )
-            return key_list, [], ptr_list, element_size_list
-        else:
-            # non zero copy: create contiguous, temporary tensors
-            target_tensors = [
-                self.mem_pool_host.get_dummy_flat_data_page() for i in range(page_num)
-            ]
+            host_buffers = list(zip(ptr_list, size_list))
+            return key_list, host_buffers
 
-            key_list = [self._get_suffixed_key(key) for key in keys]
-            ptr_list = [tensor.data_ptr() for tensor in target_tensors]
-            element_size_list = [
-                tensor.numel() * tensor.element_size() for tensor in target_tensors
-            ]
+        if page_num > STORAGE_BATCH_SIZE:
+            logger.error(
+                f"HiCacheNixl: batch size {page_num} exceeds bounce buffer capacity {STORAGE_BATCH_SIZE}"
+            )
+            return [], []
 
-            return key_list, target_tensors, ptr_list, element_size_list
+        bounce = self._bounce_set if op == "set" else self._bounce_get
+        if op == "set":
+            for i in range(page_num):
+                src = self.mem_pool_host.get_data_page(
+                    host_indices[i * page_size], flat=True
+                )
+                bounce[i].copy_(src)
 
-    def _batch_get_zero_copy_impl(
+        host_buffers = self._bounce_slot_buffers(bounce, page_num)
+        key_list = [self._get_suffixed_key(key) for key in keys]
+        return key_list, host_buffers
+
+    def _batch_xfer(
         self,
         keys: List[str],
         key_strs: List[str],
-        target_tensors: List[torch.Tensor],
-        target_locations: List[int],
-        target_sizes: List[int],
-    ) -> List[int]:
-
-        if not key_strs or not target_locations or not target_sizes:
+        host_buffers: List[tuple],
+        direction: str,
+    ) -> List[bool]:
+        """Run a batch READ or WRITE for the v1 path against the pre-registered
+        host region (no per-transfer host registration).
+        """
+        if not key_strs or not host_buffers:
             return [False] * len(keys)
 
-        if (len(key_strs) != len(target_locations)) or (
-            len(target_sizes) != len(target_locations)
-        ):
-            logger.error(
-                "Mismatch between number of key_strs, target_locations and target_sizes"
-            )
+        if len(key_strs) != len(host_buffers):
+            logger.error("Mismatch between number of key_strs and host_buffers")
             return [False] * len(keys)
 
-        if self.is_zero_copy:
-            dest = list(zip(target_locations, target_sizes))
-        else:
-            dest = target_tensors
-
         if self.backend_selector.mem_type == "FILE":
             file_paths = [self.file_manager.get_file_path(key) for key in key_strs]
-            success = self._execute_transfer(dest, file_paths, "READ")
-        else:
-            success = self._execute_transfer(dest, key_strs, "READ")
+            success = self._xfer_pre_registered(host_buffers, file_paths, direction)
+        else:  # mem_type == "OBJ"
+            success = self._xfer_pre_registered(host_buffers, key_strs, direction)
 
-        return [True] * len(key_strs) if success else [False] * len(key_strs)
+        # READ results are consumed by _batch_get_postprocess, which pairs
+        # entries 2*i / 2*i+1 for non-MLA zero-copy: it needs one bool per
+        # key_str (i.e. per `_k`/`_v` buffer). WRITE results map 1:1 to
+        # pages, i.e. to `keys`.
+        result_len = len(key_strs) if direction == "READ" else len(keys)
+        return [success] * result_len
 
     def _batch_get_postprocess(
         self,
         host_indices: torch.Tensor,
-        target_tensors: List[torch.Tensor],
         results: List[bool],
     ) -> List[bool]:
-
-        page_num = len(host_indices) // self.mem_pool_host.page_size
+        page_size = self.mem_pool_host.page_size
+        page_num = len(host_indices) // page_size
 
         if self.is_zero_copy:
             # zero copy: update final results based on the boolean results from NIXL transfer
             if self.is_mla_model:
                 return results
-            else:
-                results = [
-                    (results[2 * i] and results[2 * i + 1]) for i in range(page_num)
-                ]
-                return results
-        else:
-            # non zero copy: copy data from temporary tensors to mem_pool_host page by page
-            for i in range(page_num):
-                if not results[i]:
-                    break
-                self.mem_pool_host.set_from_flat_data_page(
-                    host_indices[i * self.mem_pool_host.page_size], target_tensors[i]
-                )
+            return [(results[2 * i] and results[2 * i + 1]) for i in range(page_num)]
 
-            return results
+        # non zero copy: copy data from the get-side bounce buffer to mem_pool_host
+        for i in range(page_num):
+            if not results[i]:
+                break
+            self.mem_pool_host.set_from_flat_data_page(
+                host_indices[i * page_size], self._bounce_get[i]
+            )
+        return results
+
+    def _log_xfer_stats(
+        self,
+        op_name: str,
+        num_keys: int,
+        host_indices: torch.Tensor,
+        buffer_sizes: List[int],
+        elapsed_ms: float,
+    ) -> None:
+        total_bytes = sum(s for s in buffer_sizes if s is not None)
+        bw = total_bytes / (elapsed_ms / 1000) / (1024 * 1024) if elapsed_ms else 0.0
+        logger.debug(
+            f"HiCacheNixl {op_name} transferred: {num_keys} keys (pages), "
+            f"{host_indices.numel()} host_indices, {total_bytes} bytes, "
+            f"total time: {elapsed_ms:.3f} ms, effective bandwidth: {bw:.2f} MB/s"
+        )
 
     def batch_get_v1(
         self,
@@ -503,106 +513,28 @@ class HiCacheNixl(HiCacheStorage):
         host_indices: torch.Tensor,
         extra_info: Optional[HiCacheStorageExtraInfo] = None,
     ) -> List[bool]:
-
-        key_strs, target_tensors, buffer_ptrs, buffer_sizes = (
-            self._batch_get_preprocess(keys, host_indices)
-        )
-
-        if not key_strs or not buffer_ptrs or not buffer_sizes:
+        if not self._host_regs:
             logger.error(
-                "HiCacheNixl batch_get_v1: preprocessing failed, empty key_strs, buffer_ptrs or buffer_sizes"
+                "HiCacheNixl batch_get_v1: register_mem_pool_host must be called first"
             )
             return [False] * len(keys)
 
+        key_strs, host_buffers = self._batch_preprocess(keys, host_indices, "get")
+        if not key_strs or not host_buffers:
+            return [False] * len(keys)
+
         start_time = time.perf_counter()
-
-        results_get = self._batch_get_zero_copy_impl(
-            keys, key_strs, target_tensors, buffer_ptrs, buffer_sizes
+        results = self._batch_xfer(keys, key_strs, host_buffers, "READ")
+        elapsed_ms = (time.perf_counter() - start_time) * 1000
+        self._log_xfer_stats(
+            "batch_get_v1",
+            len(keys),
+            host_indices,
+            [s for _, s in host_buffers],
+            elapsed_ms,
         )
 
-        end_time = time.perf_counter()
-        elapsed_time_ms = (end_time - start_time) * 1000
-        total_bytes = sum(s for s in buffer_sizes if s is not None)
-
-        logger.debug(
-            f"HiCacheNixl batch_get_v1 transferred: {len(keys)} keys (pages), {host_indices.numel()} host_indices, {total_bytes} bytes, total time: {elapsed_time_ms:.3f} ms, effective bandwidth: {total_bytes / (elapsed_time_ms / 1000) / (1024 * 1024):.2f} MB/s"
-        )
-
-        return self._batch_get_postprocess(host_indices, target_tensors, results_get)
-
-    def _batch_set_preprocess(self, keys: List[str], host_indices: torch.Tensor):
-
-        page_num = len(host_indices) // self.mem_pool_host.page_size
-
-        if len(keys) == 0 or len(keys) != page_num:
-            logger.warning(
-                f"HiCacheNixl: empty keys or mismatch in keys and host_indices lengths. keys: {len(keys)}, host_indices: {len(host_indices)}, page_size: {self.mem_pool_host.page_size}"
-            )
-            return [], [], [], []
-
-        if self.is_zero_copy:
-            key_list, _, ptr_list, element_size_list = (
-                self._get_location_and_size_list_from_meta(keys, host_indices)
-            )
-            return key_list, [], ptr_list, element_size_list
-        else:
-            # non zero copy: NIXL still requires contiguous tensors for transfer
-            target_tensors = [
-                self.mem_pool_host.get_data_page(
-                    host_indices[i * self.mem_pool_host.page_size], flat=False
-                ).contiguous()
-                for i in range(page_num)
-            ]
-
-            key_list = [self._get_suffixed_key(key) for key in keys]
-            ptr_list = [tensor.data_ptr() for tensor in target_tensors]
-            element_size_list = [
-                tensor.numel() * tensor.element_size() for tensor in target_tensors
-            ]
-
-            return key_list, target_tensors, ptr_list, element_size_list
-
-    def _batch_set_zero_copy_impl(
-        self,
-        keys: List[str],
-        key_strs: List[str],
-        target_tensors: List[torch.Tensor],
-        target_locations: List[int],
-        target_sizes: List[int],
-    ) -> List[bool]:
-
-        if not key_strs or not target_locations or not target_sizes:
-            return [False] * len(keys)
-
-        if (len(key_strs) != len(target_locations)) or (
-            len(target_sizes) != len(target_locations)
-        ):
-            logger.error(
-                "Mismatch between number of key_strs, target_locations and target_sizes"
-            )
-            return [False] * len(keys)
-
-        if self.is_zero_copy:
-            src = list(zip(target_locations, target_sizes))
-        else:
-            src = target_tensors
-
-        if self.backend_selector.mem_type == "FILE":
-            file_paths = []
-            for key in key_strs:
-                file_path = self.file_manager.get_file_path(key)
-                # New file per set, to be updated when partial writes is added to HiCache
-                if not self.file_manager.create_file(file_path):
-                    logger.error(
-                        f"******** Failed to create file {file_path} *********"
-                    )
-                    return [False] * len(keys)
-                file_paths.append(file_path)
-            success = self._execute_transfer(src, file_paths, "WRITE")
-        else:  # mem_type == "OBJ"
-            success = self._execute_transfer(src, key_strs, "WRITE")
-
-        return [True] * len(keys) if success else [False] * len(keys)
+        return self._batch_get_postprocess(host_indices, results)
 
     def batch_set_v1(
         self,
@@ -610,7 +542,6 @@ class HiCacheNixl(HiCacheStorage):
         host_indices: torch.Tensor,
         extra_info: Optional[HiCacheStorageExtraInfo] = None,
     ) -> List[bool]:
-
         # skip on MLA backup rank
         if self.backup_skip:
             return [True] * len(keys)
@@ -618,27 +549,25 @@ class HiCacheNixl(HiCacheStorage):
         if len(keys) == 0:
             return []
 
-        key_strs, target_tensors, buffer_ptrs, buffer_sizes = (
-            self._batch_set_preprocess(keys, host_indices)
-        )
-
-        if not key_strs or not buffer_ptrs or not buffer_sizes:
+        if not self._host_regs:
             logger.error(
-                "HiCacheNixl batch_set_v1: preprocessing failed, empty key_strs, buffer_ptrs or buffer_sizes"
+                "HiCacheNixl batch_set_v1: register_mem_pool_host must be called first"
             )
             return [False] * len(keys)
 
+        key_strs, host_buffers = self._batch_preprocess(keys, host_indices, "set")
+        if not key_strs or not host_buffers:
+            return [False] * len(keys)
+
         start_time = time.perf_counter()
-
-        results_set = self._batch_set_zero_copy_impl(
-            keys, key_strs, target_tensors, buffer_ptrs, buffer_sizes
+        results = self._batch_xfer(keys, key_strs, host_buffers, "WRITE")
+        elapsed_ms = (time.perf_counter() - start_time) * 1000
+        self._log_xfer_stats(
+            "batch_set_v1",
+            len(keys),
+            host_indices,
+            [s for _, s in host_buffers],
+            elapsed_ms,
         )
 
-        end_time = time.perf_counter()
-        elapsed_time_ms = (end_time - start_time) * 1000
-        total_bytes = sum(s for s in buffer_sizes if s is not None)
-        logger.debug(
-            f"HiCacheNixl batch_set_v1 transferred: {len(keys)} keys (pages), {host_indices.numel()} host_indices, {total_bytes} bytes, total time: {elapsed_time_ms:.3f} ms, effective bandwidth: {total_bytes / (elapsed_time_ms / 1000) / (1024 * 1024):.2f} MB/s"
-        )
-
-        return results_set
+        return results

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:
-        """Create a file if it doesn't exist."""
-        try:
-            os.makedirs(os.path.dirname(file_path), exist_ok=True)
-            if not os.path.exists(file_path):
-                with open(file_path, "wb") as f:
-                    pass  # Create empty file
-            return True
-        except Exception as e:
-            logger.error(f"Failed to create file {file_path}: {e}")
-            return False
+    def open_file(self, file_path: str, create: bool = False) -> Optional[int]:
+        """Open a file and return its file descriptor.
 
-    def open_file(self, file_path: str) -> Optional[int]:
-        """Open a file and return its file descriptor."""
+        If ``create`` is True, the file is created if it does not exist
+        (mode 0o644, no truncation). When ``self.use_direct_io`` is True,
+        the file is opened with ``O_DIRECT`` (bypasses the OS page cache);
+        falls back to buffered I/O with a warning if ``O_DIRECT`` is
+        unavailable on this platform.
+        """
+        flags = os.O_RDWR | os.O_CREAT if create else os.O_RDWR
+        if self.use_direct_io:
+            if hasattr(os, "O_DIRECT"):
+                flags |= os.O_DIRECT
+            else:
+                logger.warning(
+                    "use_direct_io is True, but O_DIRECT is not available on "
+                    "this system. Falling back to buffered I/O."
+                )
         try:
-            fd = os.open(file_path, os.O_RDWR)
-            return fd
+            return os.open(file_path, flags, 0o644)
         except Exception as e:
             logger.error(f"Failed to open file {file_path}: {e}")
             return None
@@ -283,17 +243,3 @@ class NixlFileManager:
         except Exception as e:
             logger.error(f"Failed to close file descriptor {fd}: {e}")
             return False
-
-    def files_to_nixl_tuples(
-        self, file_paths: List[str]
-    ) -> List[Tuple[int, int, int, str]]:
-        """Create NIXL tuples (offset, length, fd, file_path) for given files."""
-        tuples = []
-        for path in file_paths:
-            if (fd := self.open_file(path)) is None:
-                # Clean up on failure
-                for t in tuples:
-                    self.close_file(t[2])
-                return []
-            tuples.append((0, 0, fd, path))
-        return tuples

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