From 7d9de81cda08d7fb63cf5f56266cd0ed82a96541 Mon Sep 17 00:00:00 2001 From: Lukas Humbel Date: Wed, 1 Jul 2026 11:59:45 +0200 Subject: [PATCH] feat(hicache): Use NIXL path-mode (#27060) --- .../mem_cache/storage/nixl/nixl_registry.py | 65 ++++++++++++++++--- .../srt/mem_cache/storage/nixl/nixl_utils.py | 20 +++++- .../mem_cache/test_hicache_nixl_storage.py | 19 ++++++ 3 files changed, 95 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/mem_cache/storage/nixl/nixl_registry.py b/python/sglang/srt/mem_cache/storage/nixl/nixl_registry.py index f9f341a05..ecbe2e50a 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/nixl_registry.py +++ b/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 fds is None: - yield None - return - tuples = [(0, sizes[i], fds[i], keys[i]) for i in range(len(keys))] + if self.path_mode: + parts = ["rw", "create"] if direction == "WRITE" else ["ro"] + if self.file_manager.use_direct_io: + parts.append("direct") + 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( - [(0, sizes[i], fds[i]) for i in range(len(fds))], "FILE" - ) + yield reg.trim() + 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 diff --git a/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py b/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py index 05a31b57d..ed826eb84 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py +++ b/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) diff --git a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py index 564fc3449..841133420 100644 --- a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py +++ b/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