feat(hicache): Use NIXL path-mode (#27060)

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

@@ -51,6 +51,16 @@ class NixlRegistry:
         # from a single monotonic counter.
         self._obj_devid_lock = threading.Lock()
         self._obj_devid_next = 1
+        self.path_mode = mem_type == "FILE" and self._probe_path_mode()
+        if mem_type == "FILE" and self.path_mode:
+            logger.info("HiCacheNixl: path-mode FILE registration active.")
+        elif mem_type == "FILE":
+            # TODO: NIXL 1.3.0 adds path-mode support; remove this fd fallback once 1.3.0 is widely installed.
+            logger.info(
+                "HiCacheNixl: the installed NIXL build does not "
+                "support path-mode FILE registration; using legacy "
+                "fd registration."
+            )
 
     @contextmanager
     def _open_files(self, paths: List[str], create: bool):
@@ -95,6 +105,31 @@ class NixlRegistry:
                 except Exception as e:
                     logger.debug("deregister_memory skipped: %s", e)
 
+    def _probe_path_mode(self) -> bool:
+        """Probe whether NIXL honours path-mode metaInfo.
+
+        Register a FILE_SEG with a valid path-mode string pointing at a
+        nonexistent path (no 'create' flag). A path-mode-capable NIXL tries
+        to open() the path, fails with NIXL_ERR_BACKEND, and raises. A
+        pre-path-mode NIXL ignores metaInfo and returns NIXL_SUCCESS.
+        Error from register_memory => path mode supported.
+        """
+        reg_descs = self.agent.get_reg_descs(
+            [(0, 4096, 1, "rw:/nonexistent-nixl-probe")], "FILE"
+        )
+        if reg_descs is None:
+            return False
+        try:
+            reg = self.agent.register_memory(reg_descs)
+            if reg is not None:
+                try:
+                    self.agent.deregister_memory(reg)
+                except Exception:
+                    pass
+            return False
+        except Exception:
+            return True
+
     @contextmanager
     def storage(self, buffers, keys, direction):
         """Open + register the storage side; deregister and close fds on exit.
@@ -108,18 +143,32 @@ class NixlRegistry:
             return
 
         if self.mem_type == "FILE":
-            with self._open_files(keys, create=(direction == "WRITE")) as fds:
+            if self.path_mode:
-                if fds is None:
+                parts = ["rw", "create"] if direction == "WRITE" else ["ro"]
-                    yield None
+                if self.file_manager.use_direct_io:
-                    return
+                    parts.append("direct")
-                tuples = [(0, sizes[i], fds[i], keys[i]) for i in range(len(keys))]
+                spec = ",".join(parts)
+                tuples = [
+                    (0, sizes[i], i + 1, f"{spec}:{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(
+                    yield reg.trim()
-                        [(0, sizes[i], fds[i]) for i in range(len(fds))], "FILE"
+            else:
-                    )
+                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

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

@@ -3,7 +3,11 @@ import os
 from typing import Optional
 
 from sglang.srt.environ import envs
-from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key
+from sglang.srt.mem_cache.storage.nixl.nixl_routing import (
+    _BUCKET_MASK,
+    BUCKET_HEX_CHARS,
+    route_key,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -224,6 +228,7 @@ class NixlFileManager:
         else:
             for base in self.base_dirs:
                 os.makedirs(base, exist_ok=True)
+            self.ensure_all_bucket_dirs()
             logger.debug(
                 f"Initialized file manager with base directories: {self.base_dirs}. Direct I/O: {use_direct_io}"
             )
@@ -247,6 +252,19 @@ class NixlFileManager:
                 logger.error(f"Failed to clear base directory {base}: {e}")
         logger.debug(f"Cleared all files in base directories: {self.base_dirs}")
 
+    def ensure_all_bucket_dirs(self) -> None:
+        """Pre-create every possible bucket directory under each base dir.
+
+        Called once when path mode is active so NIXL O_CREAT writes never
+        fail due to a missing parent directory.
+        """
+        for base in self.base_dirs:
+            for i in range(_BUCKET_MASK + 1):
+                os.makedirs(
+                    os.path.join(base, f"{i:0{BUCKET_HEX_CHARS}x}"),
+                    exist_ok=True,
+                )
+
     def iter_all_base_dirs(self) -> list[str]:
         """Return base directories that may contain NIXL FILE cache entries."""
         return list(self.base_dirs)

test/registered/unit/mem_cache/test_hicache_nixl_storage.py

@@ -382,17 +382,36 @@ class TestNixlUnified(CustomTestCase):
                 num_pages * mock_host.page_size, dtype=torch.int64
             )
 
+            # Distinct data per key so a round trip that mixes keys up (e.g. a
+            # FILE registration that collapses every desc onto one file) is
+            # caught, not just the pass/fail return codes. Non-zero-copy only:
+            # the layer_first layout makes per-page seeding straightforward.
+            if not is_zero_copy_mode:
+                ps = mock_host.page_size
+                for p in range(num_pages):
+                    mock_host.kv_buffer[:, :, p * ps : (p + 1) * ps] = float(p + 1)
+                expected = mock_host.kv_buffer.clone()
+
             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}",
             )
 
+            if not is_zero_copy_mode:
+                mock_host.kv_buffer.zero_()
+
             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}",
             )
+
+            if not is_zero_copy_mode:
+                self.assertTrue(
+                    torch.equal(mock_host.kv_buffer, expected),
+                    "round trip corrupted or mixed up per-key data",
+                )
         finally:
             agent.get_reg_descs = orig_get_reg
             agent.register_memory = orig_register