Add bucketed multi-dir layout for NIXL file storage (#27672)

Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>

python/sglang/srt/mem_cache/storage/nixl/README.md

@@ -81,10 +81,12 @@ If a plugin is configured but its dependencies are missing, it will be skipped.
 For POSIX / GDS / GDS_MT file-based backends, the default storage location is `/tmp/hicache_storage`. However, you can customize where cached data is stored:
 
 ```bash
-export SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR=/path/to/storage/dir
+# When specifying multiple storage directories. SGLang routes each cache object to one
+# directory with a stable hash.
+export SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR=/path/to/storage/dir1,/path/to/storage/dir2,/path/to/storage/dir3
 ```
 
-This directory is used only for **FILE-backed** plugins. **OBJ-backed** plugins use object keys instead of local files.
+These directories are used only for **FILE-backed** plugins. **OBJ-backed** plugins use object keys instead of local files.
 
 ### 3. How to Provide Configuration for Backends
 
@@ -300,7 +302,7 @@ python/sglang/srt/mem_cache/storage/nixl/
 
 ### HiCache / NIXL Data Model
 
-- **FILE backends** use local file paths under `SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR`
+- **FILE backends** use local file paths under `SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR`. When multiple comma-separated directories are configured, each logical cache key is routed to one base directory with a stable hash and stored as `base_dir/<bucket>/<key>`.
 - **OBJ backends** use object keys directly
 - **MHA naming** includes TP rank and TP size, so each rank stores its own KV data
 - **MLA naming** omits TP rank, so all ranks refer to one shared logical KV object / file

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

@@ -1,4 +1,5 @@
 import logging
+import os
 import time
 import uuid
 from typing import Any, List, Optional
@@ -30,6 +31,26 @@ except ImportError as e:
 logger = logging.getLogger(__name__)
 
 
+def _parse_storage_dirs(raw: Optional[str]) -> List[str]:
+    """Split NIXL FILE storage directory config into ordered unique paths."""
+    if not raw:
+        return []
+    candidates = [path.strip() for path in raw.split(",")]
+    candidates = [path for path in candidates if path]
+    seen: dict[str, str] = {}
+    ordered: List[str] = []
+    for path in candidates:
+        real_path = os.path.realpath(path)
+        if real_path in seen:
+            raise ValueError(
+                "SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR contains duplicate "
+                f"path {path!r} (same mount as {seen[real_path]!r})."
+            )
+        seen[real_path] = path
+        ordered.append(path)
+    return ordered
+
+
 class HiCacheNixl(HiCacheStorage):
     """HiCacheNixl provides high-performance storage using NIXL plugins."""
 
@@ -49,9 +70,11 @@ class HiCacheNixl(HiCacheStorage):
         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
+        storage_dirs = _parse_storage_dirs(
+            envs.SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR.get() or file_path
+        )
         self.file_manager = (
-            NixlFileManager(file_path, use_direct_io=use_direct_io)
+            NixlFileManager(storage_dirs, use_direct_io=use_direct_io)
             if plugin not in NixlBackendSelection.OBJ_PLUGINS
             else None
         )

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

@@ -0,0 +1,29 @@
+"""Deterministic path routing for NIXL FILE-backed HiCache storage."""
+
+import hashlib
+
+_BUCKET_HEX_CHARS = 2
+_BUCKET_MASK = (1 << (4 * _BUCKET_HEX_CHARS)) - 1
+
+
+def stable_key_hash(key: str) -> int:
+    """Return a process-stable 64-bit hash for a NIXL storage key."""
+    return int.from_bytes(
+        hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest(), "big"
+    )
+
+
+def route_key(key: str, num_disks: int) -> tuple[int, str]:
+    """Return the storage disk index and bucket directory for a storage key."""
+    if num_disks <= 0:
+        raise ValueError("num_disks must be positive")
+    key_hash = stable_key_hash(key)
+    return (
+        (key_hash >> 16) % num_disks,
+        f"{key_hash & _BUCKET_MASK:0{_BUCKET_HEX_CHARS}x}",
+    )
+
+
+def route_disk(key: str, num_disks: int) -> int:
+    """Return the storage disk index for a storage key."""
+    return route_key(key, num_disks)[0]

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

@@ -3,6 +3,7 @@ import os
 from typing import Optional
 
 from sglang.srt.environ import envs
+from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key
 
 logger = logging.getLogger(__name__)
 
@@ -171,45 +172,60 @@ class NixlBackendSelection:
 class NixlFileManager:
     """Handles file system operations for NIXL."""
 
-    def __init__(self, base_dir: str, use_direct_io: bool = True):
+    def __init__(self, base_dir: "list[str] | str", use_direct_io: bool = True):
         """
         Initialize file manager.
         Args:
-            base_dir: Base directory for storing tensor files
+            base_dir: Base directory or ordered base directories for 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
+        if isinstance(base_dir, str):
+            self.base_dirs = [base_dir] if base_dir else []
+        else:
+            self.base_dirs = [d for d in base_dir if d]
         self.use_direct_io = use_direct_io
-        if base_dir == "":
+        self._created_bucket_dirs: set[str] = set()
+        if not self.base_dirs:
             logger.debug(
                 f"Initialized file manager without a base directory. Direct I/O: {use_direct_io}"
             )
         else:
-            os.makedirs(base_dir, exist_ok=True)
+            for base in self.base_dirs:
+                os.makedirs(base, exist_ok=True)
             logger.debug(
-                f"Initialized file manager with base directory: {base_dir}. Direct I/O: {use_direct_io}"
+                f"Initialized file manager with base directories: {self.base_dirs}. Direct I/O: {use_direct_io}"
             )
 
     def clear(self) -> None:
-        """Clear all files in the base directory."""
-        if self.base_dir == "":
-            logger.warning("Base directory is empty, skipping clear operation")
+        """Clear all files below every configured base directory."""
+        if not self.base_dirs:
+            logger.warning("Base directories are empty, skipping clear operation")
             return
 
-        try:
-            for root, dirs, files in os.walk(self.base_dir):
-                for file in files:
-                    os.remove(os.path.join(root, file))
-            logger.debug(f"Cleared all files in base directory: {self.base_dir}")
-        except Exception as e:
-            logger.error(
-                f"Failed to clear files in base directory {self.base_dir}: {e}"
-            )
+        for base in self.base_dirs:
+            try:
+                for root, _dirs, files in os.walk(base):
+                    for file in files:
+                        file_path = os.path.join(root, file)
+                        try:
+                            os.remove(file_path)
+                        except OSError as e:
+                            logger.warning(f"Failed to remove file {file_path}: {e}")
+            except Exception as e:
+                logger.error(f"Failed to clear base directory {base}: {e}")
+        logger.debug(f"Cleared all files in base directories: {self.base_dirs}")
+
+    def iter_all_base_dirs(self) -> list[str]:
+        """Return base directories that may contain NIXL FILE cache entries."""
+        return list(self.base_dirs)
 
     def get_file_path(self, key: str) -> str:
         """Get full file path for a given key."""
-        return os.path.join(self.base_dir, key)
+        if not self.base_dirs:
+            return key
+        disk_idx, bucket = route_key(key, len(self.base_dirs))
+        return os.path.join(self.base_dirs[disk_idx], bucket, key)
 
     def open_file(self, file_path: str, create: bool = False) -> Optional[int]:
         """Open a file and return its file descriptor.
@@ -230,6 +246,11 @@ class NixlFileManager:
                     "this system. Falling back to buffered I/O."
                 )
         try:
+            if create:
+                parent = os.path.dirname(file_path)
+                if parent and parent not in self._created_bucket_dirs:
+                    os.makedirs(parent, exist_ok=True)
+                    self._created_bucket_dirs.add(parent)
             return os.open(file_path, flags, 0o644)
         except Exception as e:
             logger.error(f"Failed to open file {file_path}: {e}")

test/registered/unit/mem_cache/test_hicache_nixl_storage.py

@@ -773,5 +773,77 @@ class TestNixlDirectIO(CustomTestCase):
         self.assertFalse(hicache.file_manager.use_direct_io)
 
 
+class TestNixlFileLayout(CustomTestCase):
+    """Tests for deterministic NIXL FILE storage path layout."""
+
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp(prefix="test_nixl_layout_")
+
+    def tearDown(self):
+        shutil.rmtree(self.test_dir, ignore_errors=True)
+
+    def test_route_key_is_stable_and_bucketed(self):
+        from sglang.srt.mem_cache.storage.nixl.nixl_routing import (
+            _BUCKET_HEX_CHARS,
+            route_key,
+        )
+
+        self.assertEqual(route_key("page-123", 4), route_key("page-123", 4))
+        disk_idx, bucket = route_key("page-123", 4)
+        self.assertGreaterEqual(disk_idx, 0)
+        self.assertLess(disk_idx, 4)
+        self.assertEqual(len(bucket), _BUCKET_HEX_CHARS)
+        self.assertRegex(bucket, rf"^[0-9a-f]{{{_BUCKET_HEX_CHARS}}}$")
+
+    def test_route_key_rejects_empty_disk_set(self):
+        from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key
+
+        with self.assertRaises(ValueError):
+            route_key("page-123", 0)
+
+    def test_file_manager_routes_to_bucketed_base_dir(self):
+        from sglang.srt.mem_cache.storage.nixl.nixl_routing import (
+            route_disk,
+            route_key,
+        )
+        from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager
+
+        base_dirs = [os.path.join(self.test_dir, f"disk{i}") for i in range(3)]
+        fm = NixlFileManager(base_dirs, use_direct_io=False)
+        key = "page-123"
+
+        disk_idx, bucket = route_key(key, len(base_dirs))
+        self.assertEqual(route_disk(key, len(base_dirs)), disk_idx)
+        self.assertEqual(
+            fm.get_file_path(key), os.path.join(base_dirs[disk_idx], bucket, key)
+        )
+        self.assertEqual(fm.iter_all_base_dirs(), base_dirs)
+
+    def test_open_file_creates_bucket_directory(self):
+        from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager
+
+        fm = NixlFileManager(self.test_dir, use_direct_io=False)
+        file_path = fm.get_file_path("page-123")
+        fd = fm.open_file(file_path, create=True)
+        try:
+            self.assertIsNotNone(fd)
+            self.assertTrue(os.path.exists(file_path))
+        finally:
+            if fd is not None:
+                os.close(fd)
+
+    def test_clear_removes_nested_bucket_files(self):
+        from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager
+
+        fm = NixlFileManager(self.test_dir, use_direct_io=False)
+        file_path = fm.get_file_path("page-123")
+        fd = fm.open_file(file_path, create=True)
+        os.close(fd)
+
+        fm.clear()
+
+        self.assertFalse(os.path.exists(file_path))
+
+
 if __name__ == "__main__":
     unittest.main()