diff --git a/python/sglang/srt/mem_cache/storage/nixl/README.md b/python/sglang/srt/mem_cache/storage/nixl/README.md index 277f7ec2d..8d1c542d9 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/README.md +++ b/python/sglang/srt/mem_cache/storage/nixl/README.md @@ -4,11 +4,12 @@ This directory contains the **NIXL (NVIDIA Inference Xfer Library)** integration NIXL provides a unified API for accessing various storage plugins, including but not limited to: +- POSIX for file based operations, including AIO / io_uring / POSIX AIO. - **Deepseek's 3FS APIs** for high-throughput file operations - **GPU Direct Storage (GDS)** for direct data movement between storage and GPU memory, bypassing CPU memory copies - **Amazon S3-compatible object storage** for key-value access patterns -Additional backend integrations are planned for future releases. +NIXL also supports additional backends such as **AZURE_BLOB**, **GUSLI**, and **UCX**. Additional backend integrations are planned for future releases. ## NIXL Resources @@ -22,6 +23,13 @@ The NIXL integration consists of two main files: - **`hicache_nixl.py`** - Main HiCache storage connector using NIXL - **`nixl_utils.py`** - Utility classes for backend selection, registration, and file management +At runtime, HiCache uses NIXL as a transfer layer between host memory and either: + +- **FILE-backed storage plugins** such as 3FS / POSIX / GDS / GDS_MT +- **OBJ-backed storage plugins** such as S3-compatible object stores + +The connector supports both the legacy tensor-oriented API (`get` / `set`) and the newer page-oriented API (`batch_get_v1` / `batch_set_v1`) used by modern HiCache backends. + ## Components ### HiCacheNixl @@ -29,6 +37,9 @@ The main storage connector that provides: - Single and batch tensor set/get operations - Automatic backend selection (3FS > POSIX > GDS_MT > GDS > OBJ) - High-performance file-based (or) object based storage access using NIXL +- Automatic zero-copy enablement when HiCache host memory layout is `page_first` or `page_first_direct` +- MLA-aware storage naming and backend-local MLA backup skipping on non-zero TP ranks +- Runtime diagnostics for mem-pool type, MLA mode, TP rank, and backup-skip state ### NixlUtils Consolidated utility classes: @@ -37,6 +48,8 @@ Consolidated utility classes: - **NixlRegistration** - Manages memory registration for tensors, files and objects - **NixlFileManager** - Handles file system operations and NIXL tuple creation +The current implementation performs per-transfer registration for file / object targets and explicitly closes FILE descriptors after registration / transfer setup to avoid descriptor leaks. + ## Using NIXL as the HiCache Storage Backend ### 1. How Backend Plugin Selection Works @@ -44,13 +57,19 @@ Consolidated utility classes: The NIXL backend can support **multiple storage plugins** (e.g., POSIX, GDS, GDS_MT, 3FS, object store, etc). * Each plugin has its own configuration section in the TOML file. +* The connector accepts configuration in two forms: + + * a **fully qualified** form such as `{"plugin": {"posix": {...}, "gds": {...}}}` + * a **flat** form such as `{"use_uring": "true"}`, which applies to the selected plugin * A plugin is considered **usable** if: * Its required library is available on the system (POSIX, GDS, GDS_MT are natively supported by NIXL). * Its configuration is valid. * It is marked as `active = true` in the configuration file (if applicable). * Some plugins (e.g., 3FS, GDS) require additional system libraries or hardware support. -* NIXL selects the backend based on **internal priority and availability**, if neither a config file nor an in-command-line config string is provided. +* If the config explicitly enables multiple plugins, the connector chooses the **first active plugin** in the config. +* If no plugin is explicitly selected in config, the connector falls back to the environment variable `SGLANG_HICACHE_NIXL_BACKEND_PLUGIN`, and finally to `auto`. +* In `auto` mode, NIXL selects the backend based on **internal priority and availability**. If a plugin is configured but its dependencies are missing, it will be skipped. @@ -63,6 +82,8 @@ For POSIX / GDS / GDS_MT file-based backends, the default storage location is `/ export SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR=/path/to/storage/dir ``` +This directory is used only for **FILE-backed** plugins. **OBJ-backed** plugins use object keys instead of local files. + ### 3. How to Provide Configuration for Backends There are three ways to specify configurations for the backends: default config, file based config, and command-line (JSON string based) config. @@ -86,6 +107,12 @@ python3 -m sglang.launch_server \ By default, NIXL will use its internal backend selection logic to choose an available storage plugin (and use default configs for the selected storage plugin). +For object storage backends, make sure the bucket is configured either in `--hicache-storage-backend-extra-config` or via: + +```bash +export AWS_DEFAULT_BUCKET= +``` + #### 2. Using a Configuration File (Recommended) @@ -125,7 +152,7 @@ This requires explicitly specifying the plugin type via an environment variable, The below example shows how to use command-line string to use the POSIX plugin where URING is enabled for async POSIX storage. ```bash -export SGLANG_HICACHE_NIXL_BACKEND_PLUGIN_TYPE=POSIX +export SGLANG_HICACHE_NIXL_BACKEND_PLUGIN=POSIX python3 -m sglang.launch_server \ --model-path \ @@ -143,6 +170,8 @@ python3 -m sglang.launch_server \ ⚠️ **Note**: This method is convenient for testing / experimenting. For production or multi-plugin setups, it is always recommended to use the config file based approach. +Also note that the flat inline config form is interpreted as plugin-specific parameters for the selected plugin. + ## Running Unit Tests @@ -177,20 +206,22 @@ Note: The `-o asyncio_mode=strict` flag is added to suppress warnings about asyn Tests for this integration, a test suite can be found at `test_hicache_nixl_storage.py` which covers: -### HiCache Integration Tests (4 tests) +### HiCache Integration Tests - Single tensor set/get operations - Batch tensor set/get operations - Mixed single and batch operations - Data integrity for various tensor types -### File Management Tests (5 tests) +### File Management Tests - Basic file operations - NIXL tuple creation - Error handling in file operations -### Registration Tests (2 tests) +### Registration and MLA / Query Tests - Tensor registration with memory type detection -- File registration using NIXL tuples +- File registration using file paths +- MLA backup-skip behavior for `batch_set_v1` +- Zero-copy `batch_exists()` accounting for MLA and MHA ## Expected Output @@ -213,16 +244,28 @@ If NIXL operations fail: - Check that NIXL is properly installed - Verify that required plugins are available - Ensure file permissions are correct for test directories +- For OBJ plugins, verify `bucket` or `AWS_DEFAULT_BUCKET` is set +- Check the NIXL diagnostic log emitted when the mem pool is registered; it includes: + - `mem_pool_device_type` + - `is_mla_model` + - `tp_rank` + - `backup_skip` + +### MLA Write Behavior +For MLA models, the NIXL backend now mirrors HF3FS's backend-local protection: +- TP rank 0 performs the actual storage write +- non-zero TP ranks skip backup writes locally in `batch_set` / `batch_set_v1` +- MLA storage names omit TP rank so all ranks refer to the same logical storage object or file ## File Structure -``` -python/sglang/srt/mem_cache/nixl/ -├── hicache_nixl.py # Main HiCache storage connector -├── nixl_utils.py # All NIXL utility classes -├── README.md # This file -└── tests/ - └── test_nixl_unified.py # All tests in one file +```text +python/sglang/srt/mem_cache/storage/nixl/ +├── hicache_nixl.py # Main HiCache storage connector +├── nixl_utils.py # NIXL utility classes +├── test_hicache_nixl_storage.py # Unit tests +├── nixl.config.toml.sample # Example configuration +└── README.md # This file ``` ## Dependencies @@ -239,9 +282,27 @@ python/sglang/srt/mem_cache/nixl/ - **Tensor side**: multi-dimensional tensors of all numeric types (int32, int64, float32, float64) are supported. - Tensors can be on CPU or GPU (as long as a GPU capable backend such as GDS_MT is available). - Currently each tensor is mapped to a file or key, but it can be extended to support multiple keys per file or key. + - The page-oriented `*_v1` path also supports zero-copy transfers using `(address, size)` metadata from the host memory pool. - **Storage side**: file and object are supported through their relevant backends (e.g., 3FS or OBJ). +### HiCache / NIXL Data Model + +- **FILE backends** use local file paths under `SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR` +- **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 +- In zero-copy mode: + - **MHA** expands each logical page into `_k` and `_v` entries + - **MLA** expands each logical page into a single `_k` entry because MLA stores one interleaved KV representation + +### Zero-Copy Behavior + +- Zero-copy is enabled automatically when the HiCache host layout is `page_first` or `page_first_direct` +- The connector uses `mem_pool_host.get_page_buffer_meta(...)` to obtain `(address, size)` metadata +- `batch_exists()` uses the same logical key expansion rules as `batch_get_v1()` / `batch_set_v1()` +- Non-zero MLA TP ranks skip `batch_set` / `batch_set_v1()` locally as a backend-side fallback guard + ### Backend Priority The NIXL backend selection follows this priority order: @@ -286,6 +347,8 @@ An example of the configuration is provided in [`nixl.config.toml.sample`](./nix * Internal backend priority rules * Unless otherwise stated, all configuration keys are **optional** and have sensible defaults. +For object storage, `bucket` may also be omitted from the config if `AWS_DEFAULT_BUCKET` is already defined in the environment. + ### 2. POSIX File System Backend (`plugin.posix`) @@ -467,4 +530,11 @@ Configures an object storage backend compatible with S3 APIs (e.g., AWS S3, MinI ## Note -This is v0 of the NIXL connector. Future versions will focus on further performance optimizations such as memory pre-registration (pre-allocating and registering memory buffers to reduce registration overhead during transfers) and block merging (combining related blocks as offsets within the same file to reduce file operations and improve throughput). These optimizations require changes at a higher layer, as the current HiCache API doesn't expose information like block relationships or hash patterns that would enable these optimizations. +This is v0 of the NIXL connector. The current implementation favors correctness and compatibility with the existing HiCache API: + +- file / object targets are registered per transfer +- FILE descriptors are explicitly cleaned up after registration / transfer setup +- MLA uses shared storage naming and backend-local write skipping on non-zero TP ranks +- zero-copy is driven by HiCache host-memory layout rather than a separate NIXL flag + +Future versions will focus on further performance optimizations such as memory pre-registration (pre-allocating and registering memory buffers to reduce registration overhead during transfers) and block merging (combining related blocks as offsets within the same file to reduce file operations and improve throughput). These optimizations require changes at a higher layer, as the current HiCache API doesn't expose information like block relationships or hash patterns that would enable these optimizations. diff --git a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py index c6fde8adf..5b961076a 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py +++ b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py @@ -64,6 +64,9 @@ class HiCacheNixl(HiCacheStorage): ) self.is_mla_model = storage_config.is_mla_model + self.is_zero_copy = False + self.storage_config = storage_config + self.backup_skip = self.is_mla_model and storage_config.tp_rank != 0 model_name = "-".join(model_name.split("/")) if model_name else "" @@ -284,6 +287,11 @@ class HiCacheNixl(HiCacheStorage): target_locations: Optional[List[int]] = None, target_sizes: Optional[List[int]] = 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 @@ -342,8 +350,8 @@ class HiCacheNixl(HiCacheStorage): if self.is_zero_copy: key_list = self._get_key_list_from_meta(keys) key_denominator = ( - 1 if not self.is_mla_model else 2 - ) # MLA model only has k buffer, no separate v buffer + 1 if self.is_mla_model else 2 + ) # MLA: 1 key per page (_k only), non-MLA: 2 NIXL keys per page (_k + _v) else: key_list = [self._get_suffixed_key(key) for key in keys] key_denominator = 1 @@ -603,6 +611,10 @@ class HiCacheNixl(HiCacheStorage): extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: + # skip on MLA backup rank + if self.backup_skip: + return [True] * len(keys) + if len(keys) == 0: return [] diff --git a/python/sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py b/python/sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py index cd558a759..9ddbbb24b 100755 --- a/python/sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py +++ b/python/sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py @@ -23,6 +23,7 @@ class TestNixlUnified(unittest.TestCase): # 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() @@ -42,9 +43,9 @@ class TestNixlUnified(unittest.TestCase): 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", - enable_storage_metrics=False, extra_config={"plugin": {"posix": {"active": True}}}, ) @@ -53,6 +54,7 @@ class TestNixlUnified(unittest.TestCase): 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") @@ -61,7 +63,7 @@ class TestNixlUnified(unittest.TestCase): if os.path.exists(self.test_dir): import shutil - shutil.rmtree(self.test_dir) + shutil.rmtree(self.test_dir, ignore_errors=True) @staticmethod def _open_fds() -> int: @@ -272,6 +274,42 @@ class TestNixlUnified(unittest.TestCase): 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()