Refactor NIXL hicache. Add O_DIRECT support (#25173)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Lukas Humbel
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
89feb18eb9
commit
d8a5a25c36
@@ -837,6 +837,16 @@ SGLang supports various environment variables that can be used to configure its
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Decode-side incremental KV cache offload stride. Rounded down to a multiple of <code>--page-size</code> (min is <code>--page-size</code>). If unset/invalid/<=0, it falls back to <code>--page-size</code>.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Not set (uses <code>--page-size</code>)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_HICACHE_NIXL_USE_DIRECT_IO</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable <code>O_DIRECT</code> for any file-based NIXL backend (POSIX, GDS, GDS_MT, 3FS) when opening cache files (bypasses the OS page cache, reducing memory pressure and improving throughput on NVMe). Can also be disabled via <code>{'{"use_direct_io": false}'}</code> in <code>--hicache-storage-backend-extra-config</code>. Falls back to buffered I/O with a warning when <code>O_DIRECT</code> is unavailable on the current OS.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_HUGEPAGE_SIZE</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use huge pages for host KV cache allocations (HiCache / disaggregation offload). Valid values: <code>2MB</code> (2 MiB pages via <code>MAP_HUGE_2MB</code>) or <code>1GB</code> (1 GiB pages via <code>MAP_HUGE_1GB</code>). Requires huge pages to be pre-allocated on the host OS (<code>/proc/sys/vm/nr_hugepages</code> or <code>/sys/kernel/mm/hugepages</code>). If the allocation fails, the allocator logs a warning and falls back to regular page-size mmap automatically.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Not set (uses OS default page size)</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
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."
|
||||
)
|
||||
if ret != 0:
|
||||
raise RuntimeError(f"cudaHostRegister failed with error code {ret}")
|
||||
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):
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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`)
|
||||
|
||||
|
||||
@@ -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,117 +74,67 @@ 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]],
|
||||
keys: List[str],
|
||||
host_descs: Any,
|
||||
storage_descs: Any,
|
||||
direction: str,
|
||||
) -> bool:
|
||||
if len(buffers) != len(keys):
|
||||
logger.error("Mismatch between number of tensors/buffers and files/objects")
|
||||
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
|
||||
|
||||
# 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:
|
||||
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
|
||||
|
||||
"""Initialize and poll a NIXL transfer to completion."""
|
||||
try:
|
||||
xfer_req = self.agent.initialize_xfer(
|
||||
direction, host_descs, storage_descs, self.agent_name
|
||||
@@ -194,134 +143,87 @@ class HiCacheNixl(HiCacheStorage):
|
||||
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)
|
||||
# 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:
|
||||
for fd in file_fds:
|
||||
self.file_manager.close_file(fd)
|
||||
self.agent.release_xfer_handle(xfer_req)
|
||||
|
||||
def _xfer_pre_registered(
|
||||
self,
|
||||
host_buffers: List[tuple],
|
||||
keys: List[str],
|
||||
direction: str,
|
||||
) -> bool:
|
||||
"""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
|
||||
|
||||
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
|
||||
|
||||
with self.registry.storage(host_buffers, keys, direction) as storage_descs:
|
||||
if storage_descs is None:
|
||||
return False
|
||||
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,9 +234,95 @@ class HiCacheNixl(HiCacheStorage):
|
||||
"page_first_direct",
|
||||
]
|
||||
|
||||
logger.info(
|
||||
f"HiCacheNixl: Registered mem_pool_host with layout {self.mem_pool_host.layout}, zero_copy set to {self.is_zero_copy}"
|
||||
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: 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])
|
||||
@@ -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:
|
||||
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,206 +379,162 @@ 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
|
||||
|
||||
if page_num > STORAGE_BATCH_SIZE:
|
||||
logger.error(
|
||||
f"HiCacheNixl: batch size {page_num} exceeds bounce buffer capacity {STORAGE_BATCH_SIZE}"
|
||||
)
|
||||
return [], []
|
||||
|
||||
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)
|
||||
|
||||
host_buffers = self._bounce_slot_buffers(bounce, 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, host_buffers
|
||||
|
||||
return key_list, target_tensors, ptr_list, element_size_list
|
||||
|
||||
def _batch_get_zero_copy_impl(
|
||||
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
|
||||
return [(results[2 * i] and results[2 * i + 1]) for i in range(page_num)]
|
||||
|
||||
# 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 * self.mem_pool_host.page_size], target_tensors[i]
|
||||
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,
|
||||
keys: List[str],
|
||||
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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -0,0 +1,777 @@
|
||||
"""Unit tests for the NIXL HiCache storage backend -- no server, no model loading."""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, stage="base-a", runner_config="1-gpu-small")
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
|
||||
from sglang.srt.mem_cache.storage.nixl.hicache_nixl import HiCacheNixl
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# Stress tests are opt-in: CI never sets this; set locally to exercise them.
|
||||
STRESS_ENABLED = bool(os.environ.get("SGLANG_RUN_NIXL_STRESS"))
|
||||
|
||||
|
||||
class MockMemPoolHost:
|
||||
"""Minimal MHA-style HostKVCache stand-in supporting the v1 paths.
|
||||
|
||||
zero_copy mode uses ``page_first`` so ``get_page_buffer_meta`` returns
|
||||
valid (k, v) pointers into ``kv_buffer``. Non-zero-copy uses
|
||||
``layer_first`` so the slow path uses ``get_data_page`` /
|
||||
``set_from_flat_data_page`` against the same buffer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
is_zero_copy_mode: bool,
|
||||
page_size: int = 2,
|
||||
layer_num: int = 2,
|
||||
head_num: int = 2,
|
||||
head_dim: int = 4,
|
||||
num_pages: int = 4,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
):
|
||||
self.layout = "page_first" if is_zero_copy_mode else "layer_first"
|
||||
self.page_size = page_size
|
||||
self.layer_num = layer_num
|
||||
self.head_num = head_num
|
||||
self.head_dim = head_dim
|
||||
self.dtype = dtype
|
||||
self.num_pages = num_pages
|
||||
self.size = page_size * num_pages
|
||||
self.pin_memory = False
|
||||
if is_zero_copy_mode:
|
||||
# page_first: (2, size, layer, head, head_dim)
|
||||
self.kv_buffer = torch.zeros(
|
||||
(2, self.size, layer_num, head_num, head_dim), dtype=dtype
|
||||
)
|
||||
else:
|
||||
# layer_first: (2, layer, size, head, head_dim)
|
||||
self.kv_buffer = torch.zeros(
|
||||
(2, layer_num, self.size, head_num, head_dim), dtype=dtype
|
||||
)
|
||||
|
||||
def get_page_buffer_meta(self, indices):
|
||||
ptr_list = []
|
||||
base = self.kv_buffer.data_ptr()
|
||||
v_offset = (
|
||||
self.layer_num
|
||||
* self.size
|
||||
* self.head_num
|
||||
* self.head_dim
|
||||
* self.dtype.itemsize
|
||||
)
|
||||
idx_list = indices.tolist()
|
||||
for i in range(0, len(idx_list), self.page_size):
|
||||
k_ptr = base + idx_list[i] * (
|
||||
self.layer_num * self.head_num * self.head_dim * self.dtype.itemsize
|
||||
)
|
||||
ptr_list.append(k_ptr)
|
||||
ptr_list.append(k_ptr + v_offset)
|
||||
element_size = (
|
||||
self.layer_num
|
||||
* self.dtype.itemsize
|
||||
* self.page_size
|
||||
* self.head_num
|
||||
* self.head_dim
|
||||
)
|
||||
return ptr_list, [element_size] * len(ptr_list)
|
||||
|
||||
def get_dummy_flat_data_page(self):
|
||||
return torch.zeros(
|
||||
(2, self.layer_num, self.page_size, self.head_num, self.head_dim),
|
||||
dtype=self.dtype,
|
||||
).flatten()
|
||||
|
||||
def get_data_page(self, index, flat=True):
|
||||
if hasattr(index, "item"):
|
||||
index = int(index.item())
|
||||
page = self.kv_buffer[:, :, index : index + self.page_size, :, :]
|
||||
return page.flatten() if flat else page
|
||||
|
||||
def set_from_flat_data_page(self, index, data_page):
|
||||
if hasattr(index, "item"):
|
||||
index = int(index.item())
|
||||
self.kv_buffer[:, :, index : index + self.page_size, :, :] = data_page.reshape(
|
||||
2, self.layer_num, self.page_size, self.head_num, self.head_dim
|
||||
)
|
||||
|
||||
def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
|
||||
# Test tensors are too small to satisfy 4 KiB stride alignment; the
|
||||
# O_DIRECT path correctly falls back to copy mode in this case.
|
||||
return False
|
||||
|
||||
|
||||
class MinioFixture:
|
||||
"""Spin up a single-node MinIO server on localhost and create a bucket.
|
||||
|
||||
Relies on MinIO's default ``minioadmin``/``minioadmin`` root credentials
|
||||
so no env vars need to be plumbed through.
|
||||
"""
|
||||
|
||||
user = "minioadmin"
|
||||
password = "minioadmin"
|
||||
|
||||
def __init__(self, bucket: str = "hicache-test"):
|
||||
self.bucket = bucket
|
||||
self.api_port = self._find_free_port()
|
||||
self.data_dir = tempfile.mkdtemp(prefix="nixl_minio_")
|
||||
self.proc: subprocess.Popen | None = None
|
||||
|
||||
@property
|
||||
def endpoint(self) -> str:
|
||||
return f"127.0.0.1:{self.api_port}"
|
||||
|
||||
@staticmethod
|
||||
def _find_free_port() -> int:
|
||||
with socket.socket() as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
@staticmethod
|
||||
def _minio_bin() -> str | None:
|
||||
path = shutil.which("minio") or "/usr/local/bin/minio"
|
||||
if os.path.isfile(path) and os.access(path, os.X_OK):
|
||||
return path
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def is_available(cls) -> bool:
|
||||
"""True iff a minio binary and boto3 are both importable."""
|
||||
if cls._minio_bin() is None:
|
||||
return False
|
||||
try:
|
||||
import boto3 # noqa: F401
|
||||
except ImportError:
|
||||
return False
|
||||
return True
|
||||
|
||||
def start(self) -> None:
|
||||
minio_bin = self._minio_bin()
|
||||
if minio_bin is None:
|
||||
raise FileNotFoundError("minio binary not available")
|
||||
|
||||
self.proc = subprocess.Popen(
|
||||
[minio_bin, "server", "--address", self.endpoint, self.data_dir],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
|
||||
deadline = time.time() + 15.0
|
||||
while time.time() < deadline:
|
||||
if self.proc.poll() is not None:
|
||||
raise RuntimeError(f"minio exited early with rc={self.proc.returncode}")
|
||||
try:
|
||||
with socket.create_connection(
|
||||
("127.0.0.1", self.api_port), timeout=0.5
|
||||
):
|
||||
break
|
||||
except OSError:
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
self.stop()
|
||||
raise RuntimeError("minio did not become ready within 15s")
|
||||
|
||||
import boto3
|
||||
from botocore.config import Config
|
||||
|
||||
s3 = boto3.client(
|
||||
"s3",
|
||||
endpoint_url=f"http://{self.endpoint}",
|
||||
aws_access_key_id=self.user,
|
||||
aws_secret_access_key=self.password,
|
||||
config=Config(s3={"addressing_style": "path"}, signature_version="s3v4"),
|
||||
)
|
||||
s3.create_bucket(Bucket=self.bucket)
|
||||
|
||||
def stop(self) -> None:
|
||||
if self.proc and self.proc.poll() is None:
|
||||
self.proc.terminate()
|
||||
try:
|
||||
self.proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
self.proc.kill()
|
||||
self.proc.wait(timeout=5)
|
||||
shutil.rmtree(self.data_dir, ignore_errors=True)
|
||||
|
||||
|
||||
class TestNixlUnified(CustomTestCase):
|
||||
"""Unified test suite for all NIXL components."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test environment."""
|
||||
self.test_dir = "/tmp/test_nixl_unified"
|
||||
os.makedirs(self.test_dir, exist_ok=True)
|
||||
|
||||
# Disable O_DIRECT here: these tests use small, arbitrarily-aligned
|
||||
# tensors that do not satisfy the sector-alignment constraints required
|
||||
# by O_DIRECT. O_DIRECT-specific behaviour is exercised in
|
||||
# TestNixlDirectIO below.
|
||||
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,
|
||||
is_page_first_layout=False,
|
||||
model_name="test_model",
|
||||
enable_storage_metrics=False,
|
||||
extra_config={
|
||||
"plugin": {"posix": {"active": True}},
|
||||
"use_direct_io": False,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
self.hicache = HiCacheNixl(
|
||||
storage_config=self.storage_config,
|
||||
file_path=self.test_dir,
|
||||
)
|
||||
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):
|
||||
shutil.rmtree(self.test_dir, ignore_errors=True)
|
||||
|
||||
@staticmethod
|
||||
def _open_fds() -> int:
|
||||
return len(os.listdir("/proc/self/fd"))
|
||||
|
||||
def test_storage_register_failure_closes_fds(self):
|
||||
"""If NIXL register_memory raises after fds are opened, all fds are still closed."""
|
||||
files = [os.path.join(self.test_dir, f"fail_{i}.bin") for i in range(3)]
|
||||
buffers = [(0, 64) for _ in range(3)]
|
||||
|
||||
fds_before = self._open_fds()
|
||||
|
||||
orig = self.hicache.agent.register_memory
|
||||
|
||||
def boom(*args, **kwargs):
|
||||
raise RuntimeError("simulated register_memory failure")
|
||||
|
||||
self.hicache.agent.register_memory = boom
|
||||
try:
|
||||
with self.hicache.registry.storage(buffers, files, "WRITE") as descs:
|
||||
self.assertIsNone(
|
||||
descs, "storage CM should yield None on register failure"
|
||||
)
|
||||
finally:
|
||||
self.hicache.agent.register_memory = orig
|
||||
|
||||
self.assertEqual(
|
||||
self._open_fds(),
|
||||
fds_before,
|
||||
"fd leak after register_memory failure mid-storage",
|
||||
)
|
||||
|
||||
def _assert_host_addrs_pre_registered(
|
||||
self, is_zero_copy_mode: bool, hicache: HiCacheNixl = None
|
||||
):
|
||||
"""Exercise the v1 path and assert every host xfer addr lies within a
|
||||
currently-registered host (DRAM/tensor) region.
|
||||
|
||||
Spies are installed BEFORE ``register_mem_pool_host`` so the up-front
|
||||
pre-registration is captured too.
|
||||
"""
|
||||
if hicache is None:
|
||||
hicache = self.hicache
|
||||
agent = hicache.agent
|
||||
|
||||
# Map registration-handle id -> [(addr, size, mem_type), ...]
|
||||
active_regs: dict = {}
|
||||
# Capture items list per get_reg_descs call so we can attribute them
|
||||
# to the registration handle returned by the next register_memory call.
|
||||
pending: list = []
|
||||
|
||||
orig_get_reg = agent.get_reg_descs
|
||||
|
||||
def spy_get_reg(items, mem_type=None):
|
||||
# NIXL's register_memory calls get_reg_descs internally with an
|
||||
# already-built nixlRegDList; iterating that pybind11 type is
|
||||
# unsafe, so only record entries when the input is a plain list.
|
||||
if isinstance(items, list) and items:
|
||||
entries = []
|
||||
for it in items:
|
||||
if isinstance(it, torch.Tensor):
|
||||
entries.append(
|
||||
(it.data_ptr(), it.numel() * it.element_size(), None)
|
||||
)
|
||||
elif isinstance(it, tuple):
|
||||
entries.append((it[0], it[1], mem_type))
|
||||
pending.append(entries)
|
||||
return orig_get_reg(items, mem_type)
|
||||
|
||||
orig_register = agent.register_memory
|
||||
|
||||
def spy_register(reg_descs):
|
||||
reg = orig_register(reg_descs)
|
||||
entries = pending.pop(0) if pending else []
|
||||
active_regs[id(reg)] = entries
|
||||
return reg
|
||||
|
||||
orig_dereg = agent.deregister_memory
|
||||
|
||||
def spy_dereg(reg):
|
||||
active_regs.pop(id(reg), None)
|
||||
return orig_dereg(reg)
|
||||
|
||||
last_host_xfer: list = []
|
||||
|
||||
orig_get_xfer = agent.get_xfer_descs
|
||||
|
||||
def spy_get_xfer(items, mem_type=None):
|
||||
if mem_type in (None, "DRAM"):
|
||||
ranges = []
|
||||
for it in items:
|
||||
if isinstance(it, torch.Tensor):
|
||||
ranges.append((it.data_ptr(), it.numel() * it.element_size()))
|
||||
elif isinstance(it, tuple):
|
||||
ranges.append((it[0], it[1]))
|
||||
last_host_xfer.clear()
|
||||
last_host_xfer.extend(ranges)
|
||||
return orig_get_xfer(items, mem_type)
|
||||
|
||||
violations: list = []
|
||||
orig_init = agent.initialize_xfer
|
||||
|
||||
def spy_init(direction, local, remote, agent_name):
|
||||
host_regs = [
|
||||
(a, s)
|
||||
for entries in active_regs.values()
|
||||
for (a, s, mt) in entries
|
||||
if mt in (None, "DRAM")
|
||||
]
|
||||
for a, s in last_host_xfer:
|
||||
if not any(ra <= a and a + s <= ra + rs for (ra, rs) in host_regs):
|
||||
violations.append((a, s, dict(host_regs=host_regs)))
|
||||
last_host_xfer.clear()
|
||||
return orig_init(direction, local, remote, agent_name)
|
||||
|
||||
agent.get_reg_descs = spy_get_reg
|
||||
agent.register_memory = spy_register
|
||||
agent.deregister_memory = spy_dereg
|
||||
agent.get_xfer_descs = spy_get_xfer
|
||||
agent.initialize_xfer = spy_init
|
||||
try:
|
||||
mock_host = MockMemPoolHost(is_zero_copy_mode)
|
||||
hicache.register_mem_pool_host(mock_host)
|
||||
# Force the requested mode regardless of how register_mem_pool_host derives it.
|
||||
hicache.is_zero_copy = is_zero_copy_mode
|
||||
|
||||
num_pages = 3
|
||||
keys = [
|
||||
f"compliance_{int(is_zero_copy_mode)}_{i}" for i in range(num_pages)
|
||||
]
|
||||
host_indices = torch.arange(
|
||||
num_pages * mock_host.page_size, dtype=torch.int64
|
||||
)
|
||||
|
||||
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}",
|
||||
)
|
||||
|
||||
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}",
|
||||
)
|
||||
finally:
|
||||
agent.get_reg_descs = orig_get_reg
|
||||
agent.register_memory = orig_register
|
||||
agent.deregister_memory = orig_dereg
|
||||
agent.get_xfer_descs = orig_get_xfer
|
||||
agent.initialize_xfer = orig_init
|
||||
|
||||
self.assertEqual(
|
||||
violations,
|
||||
[],
|
||||
f"Host xfer addrs not covered by registration (zero_copy={is_zero_copy_mode}): {violations}",
|
||||
)
|
||||
|
||||
def test_nixl_api_contract_host_addrs_within_registered_region_zero_copy(self):
|
||||
"""All host xfer addrs must lie within a registered region -- zero-copy."""
|
||||
self._assert_host_addrs_pre_registered(is_zero_copy_mode=True)
|
||||
|
||||
def test_nixl_api_contract_host_addrs_within_registered_region_non_zero_copy(self):
|
||||
"""All host xfer addrs must lie within a registered region -- non-zero-copy."""
|
||||
self._assert_host_addrs_pre_registered(is_zero_copy_mode=False)
|
||||
|
||||
def _make_obj_hicache(self) -> HiCacheNixl:
|
||||
"""Start a MinIO server (cleaned up via addCleanup) and return a
|
||||
HiCacheNixl wired to its OBJ backend. Skips the test if the backend
|
||||
cannot be constructed."""
|
||||
minio = MinioFixture()
|
||||
minio.start()
|
||||
self.addCleanup(minio.stop)
|
||||
|
||||
obj_config = HiCacheStorageConfig(
|
||||
tp_rank=0,
|
||||
tp_size=1,
|
||||
pp_rank=0,
|
||||
pp_size=1,
|
||||
attn_cp_rank=0,
|
||||
attn_cp_size=1,
|
||||
is_mla_model=False,
|
||||
is_page_first_layout=False,
|
||||
model_name="test_model",
|
||||
enable_storage_metrics=False,
|
||||
extra_config={
|
||||
"plugin": {
|
||||
"obj": {
|
||||
"active": True,
|
||||
"endpoint_override": f"http://{minio.endpoint}",
|
||||
"use_virtual_addressing": "false",
|
||||
"access_key": minio.user,
|
||||
"secret_key": minio.password,
|
||||
"bucket": minio.bucket,
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
try:
|
||||
return HiCacheNixl(storage_config=obj_config, file_path="")
|
||||
except Exception as e:
|
||||
self.skipTest(f"NIXL OBJ backend unavailable: {e}")
|
||||
|
||||
@unittest.skipUnless(
|
||||
MinioFixture.is_available(), "minio binary or boto3 not available"
|
||||
)
|
||||
def test_nixl_api_contract_host_addrs_within_registered_region_obj(self):
|
||||
"""Same property over the OBJ backend (MinIO fixture)."""
|
||||
self._assert_host_addrs_pre_registered(
|
||||
is_zero_copy_mode=False, hicache=self._make_obj_hicache()
|
||||
)
|
||||
|
||||
def test_batch_set_v1_skips_on_nonzero_mla_rank(self):
|
||||
"""batch_set_v1 is a no-op on nonzero MLA backup ranks.
|
||||
|
||||
With backup_skip=True the early-return must fire before the host-regs
|
||||
check, so calling without register_mem_pool_host still returns all-True
|
||||
(the host-regs check would otherwise return all-False).
|
||||
"""
|
||||
self.hicache.backup_skip = True
|
||||
results = self.hicache.batch_set_v1(
|
||||
["key1", "key2"], torch.tensor([0, 1], dtype=torch.int64)
|
||||
)
|
||||
self.assertEqual(results, [True, True])
|
||||
|
||||
def test_batch_exists_zero_copy_mla_uses_single_key_denominator(self):
|
||||
"""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 = lambda *a, **kw: [object(), None]
|
||||
|
||||
self.assertEqual(self.hicache.batch_exists(["key1", "key2"]), 1)
|
||||
|
||||
def test_batch_exists_zero_copy_mha_uses_two_key_denominator(self):
|
||||
"""Zero-copy non-MLA 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 = lambda *a, **kw: [
|
||||
object(),
|
||||
object(),
|
||||
None,
|
||||
None,
|
||||
]
|
||||
|
||||
self.assertEqual(self.hicache.batch_exists(["key1", "key2"]), 1)
|
||||
|
||||
def _run_concurrent_stress(
|
||||
self, is_zero_copy_mode: bool, hicache: HiCacheNixl = None
|
||||
):
|
||||
"""One getter thread + one setter thread share the same HiCacheNixl
|
||||
for ``is_zero_copy_mode``. Defaults to ``self.hicache`` (FILE backend);
|
||||
pass ``hicache`` to exercise a different backend (e.g. OBJ).
|
||||
|
||||
Phase 1 pre-seeds N preset pages and stores them under fixed keys.
|
||||
Phase 2 runs the getter (reads the presets back and verifies content)
|
||||
concurrently with the setter (writes a stream of fresh distinct keys
|
||||
from a disjoint source region). The kv_buffer regions touched by the
|
||||
two threads are disjoint so any data corruption observed is from the
|
||||
backend's shared state (bounce buffers, devId maps, fd pool).
|
||||
"""
|
||||
if hicache is None:
|
||||
hicache = self.hicache
|
||||
|
||||
# 8 preset pages, 8 getter dst pages, 8 setter src pages -> 24 in use.
|
||||
mock_host = MockMemPoolHost(is_zero_copy_mode=is_zero_copy_mode, num_pages=32)
|
||||
hicache.register_mem_pool_host(mock_host)
|
||||
hicache.is_zero_copy = is_zero_copy_mode
|
||||
|
||||
page_size = mock_host.page_size
|
||||
dtype = mock_host.dtype
|
||||
num_pages = 8
|
||||
|
||||
# Disjoint per-thread regions in kv_buffer (indexed by token index).
|
||||
preset_src = (0, num_pages)
|
||||
getter_dst = (num_pages, 2 * num_pages)
|
||||
setter_src = (2 * num_pages, 3 * num_pages)
|
||||
|
||||
# zero_copy=page_first uses dim 1 for the token axis; non-zero-copy=
|
||||
# layer_first uses dim 2. All buffer accesses below go through this so
|
||||
# the rest of the harness stays layout-agnostic.
|
||||
def token_index(start_token: int, n_tokens: int):
|
||||
s = slice(start_token, start_token + n_tokens)
|
||||
if is_zero_copy_mode:
|
||||
return (slice(None), s, slice(None), slice(None), slice(None))
|
||||
return (slice(None), slice(None), s, slice(None), slice(None))
|
||||
|
||||
def page_index(start_page: int, n_pages: int):
|
||||
return token_index(start_page * page_size, n_pages * page_size)
|
||||
|
||||
def fill_pages(start_page: int, n_pages: int, value_fn):
|
||||
"""value_fn(i) -> scalar value for page i."""
|
||||
for i in range(n_pages):
|
||||
idx = page_index(start_page + i, 1)
|
||||
shape = mock_host.kv_buffer[idx].shape
|
||||
mock_host.kv_buffer[idx] = torch.full(
|
||||
shape, float(value_fn(i)), dtype=dtype
|
||||
)
|
||||
|
||||
# Phase 1: distinct value per preset page so a wrong-page result is
|
||||
# detectable; setter source is constant (value irrelevant to the
|
||||
# test, just needs to be valid).
|
||||
fill_pages(preset_src[0], num_pages, lambda i: i + 1)
|
||||
fill_pages(setter_src[0], num_pages, lambda i: -1.0)
|
||||
|
||||
preset_keys = [f"preset_{int(is_zero_copy_mode)}_{i}" for i in range(num_pages)]
|
||||
preset_indices = torch.arange(
|
||||
preset_src[0] * page_size,
|
||||
preset_src[1] * page_size,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
self.assertTrue(
|
||||
all(hicache.batch_set_v1(preset_keys, preset_indices)),
|
||||
"phase 1: presetting keys failed",
|
||||
)
|
||||
|
||||
# Expected per-page-i payload after a successful get into getter_dst.
|
||||
expected_pages = [
|
||||
mock_host.kv_buffer[page_index(preset_src[0] + i, 1)].clone()
|
||||
for i in range(num_pages)
|
||||
]
|
||||
|
||||
# Phase 2.
|
||||
stop = threading.Event()
|
||||
errors: List[str] = []
|
||||
errors_lock = threading.Lock()
|
||||
|
||||
def record_error(msg: str):
|
||||
with errors_lock:
|
||||
errors.append(msg)
|
||||
|
||||
def getter_loop():
|
||||
dst_indices = torch.arange(
|
||||
getter_dst[0] * page_size,
|
||||
getter_dst[1] * page_size,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
loops = 0
|
||||
while not stop.is_set():
|
||||
# Zero the dst pages so a no-op get is observable.
|
||||
mock_host.kv_buffer[page_index(getter_dst[0], num_pages)] = 0.0
|
||||
ok = hicache.batch_get_v1(preset_keys, dst_indices)
|
||||
if not all(ok):
|
||||
record_error(f"getter loop {loops}: batch_get_v1 returned {ok}")
|
||||
return
|
||||
for i in range(num_pages):
|
||||
got = mock_host.kv_buffer[page_index(getter_dst[0] + i, 1)]
|
||||
if not torch.equal(got, expected_pages[i]):
|
||||
record_error(f"getter loop {loops}: preset page {i} corrupted")
|
||||
return
|
||||
loops += 1
|
||||
|
||||
def setter_loop():
|
||||
src_indices = torch.arange(
|
||||
setter_src[0] * page_size,
|
||||
setter_src[1] * page_size,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
loops = 0
|
||||
while not stop.is_set():
|
||||
keys = [
|
||||
f"setter_{int(is_zero_copy_mode)}_{loops}_{i}"
|
||||
for i in range(num_pages)
|
||||
]
|
||||
ok = hicache.batch_set_v1(keys, src_indices)
|
||||
if not all(ok):
|
||||
record_error(f"setter loop {loops}: batch_set_v1 returned {ok}")
|
||||
return
|
||||
loops += 1
|
||||
|
||||
t_get = threading.Thread(target=getter_loop, daemon=True)
|
||||
t_set = threading.Thread(target=setter_loop, daemon=True)
|
||||
t_get.start()
|
||||
t_set.start()
|
||||
|
||||
# Bounded run: long enough to interleave many ops under NIXL I/O
|
||||
# GIL release, short enough for a unit test.
|
||||
time.sleep(3.0)
|
||||
stop.set()
|
||||
t_get.join(timeout=10)
|
||||
t_set.join(timeout=10)
|
||||
|
||||
self.assertFalse(
|
||||
t_get.is_alive() or t_set.is_alive(),
|
||||
"stress threads failed to stop",
|
||||
)
|
||||
self.assertEqual(errors, [], f"concurrency errors: {errors}")
|
||||
|
||||
@unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run")
|
||||
def test_concurrent_getter_setter_file_zero_copy(self):
|
||||
"""Stress: concurrent getter+setter, FILE backend, zero-copy."""
|
||||
self._run_concurrent_stress(is_zero_copy_mode=True)
|
||||
|
||||
@unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run")
|
||||
def test_concurrent_getter_setter_file_non_zero_copy(self):
|
||||
"""Stress: concurrent getter+setter, FILE backend, non-zero-copy."""
|
||||
self._run_concurrent_stress(is_zero_copy_mode=False)
|
||||
|
||||
@unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run")
|
||||
@unittest.skipUnless(
|
||||
MinioFixture.is_available(), "minio binary or boto3 not available"
|
||||
)
|
||||
def test_concurrent_getter_setter_obj_zero_copy(self):
|
||||
"""Stress: concurrent getter+setter, OBJ backend (MinIO), zero-copy."""
|
||||
self._run_concurrent_stress(
|
||||
is_zero_copy_mode=True, hicache=self._make_obj_hicache()
|
||||
)
|
||||
|
||||
@unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run")
|
||||
@unittest.skipUnless(
|
||||
MinioFixture.is_available(), "minio binary or boto3 not available"
|
||||
)
|
||||
def test_concurrent_getter_setter_obj_non_zero_copy(self):
|
||||
"""Stress: concurrent getter+setter, OBJ backend (MinIO), non-zero-copy."""
|
||||
self._run_concurrent_stress(
|
||||
is_zero_copy_mode=False, hicache=self._make_obj_hicache()
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(hasattr(os, "O_DIRECT"), "O_DIRECT not available on this platform")
|
||||
class TestNixlDirectIO(CustomTestCase):
|
||||
"""Tests for the O_DIRECT file I/O path in NixlFileManager and HiCacheNixl."""
|
||||
|
||||
def setUp(self):
|
||||
self.test_dir = "/tmp/test_nixl_direct_io"
|
||||
os.makedirs(self.test_dir, exist_ok=True)
|
||||
|
||||
def tearDown(self):
|
||||
if os.path.exists(self.test_dir):
|
||||
shutil.rmtree(self.test_dir, ignore_errors=True)
|
||||
|
||||
def test_open_file_sets_o_direct(self):
|
||||
"""open_file sets O_DIRECT on the file descriptor when use_direct_io=True."""
|
||||
import fcntl
|
||||
|
||||
from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager
|
||||
|
||||
fm = NixlFileManager(self.test_dir, use_direct_io=True)
|
||||
test_file = os.path.join(self.test_dir, "test_odirect.bin")
|
||||
fd = fm.open_file(test_file, create=True)
|
||||
try:
|
||||
self.assertTrue(fcntl.fcntl(fd, fcntl.F_GETFL) & os.O_DIRECT)
|
||||
finally:
|
||||
os.close(fd)
|
||||
|
||||
def test_open_file_no_o_direct(self):
|
||||
"""open_file does not set O_DIRECT when use_direct_io=False."""
|
||||
import fcntl
|
||||
|
||||
from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager
|
||||
|
||||
fm = NixlFileManager(self.test_dir, use_direct_io=False)
|
||||
test_file = os.path.join(self.test_dir, "test_buffered.bin")
|
||||
fd = fm.open_file(test_file, create=True)
|
||||
try:
|
||||
self.assertFalse(fcntl.fcntl(fd, fcntl.F_GETFL) & os.O_DIRECT)
|
||||
finally:
|
||||
os.close(fd)
|
||||
|
||||
def _make_direct_io_hicache(self) -> HiCacheNixl:
|
||||
"""Return a HiCacheNixl configured for O_DIRECT (default) with the POSIX backend."""
|
||||
storage_config = HiCacheStorageConfig(
|
||||
tp_rank=0,
|
||||
tp_size=1,
|
||||
pp_rank=0,
|
||||
pp_size=1,
|
||||
attn_cp_rank=0,
|
||||
attn_cp_size=1,
|
||||
is_mla_model=False,
|
||||
is_page_first_layout=False,
|
||||
model_name="test_model",
|
||||
enable_storage_metrics=False,
|
||||
extra_config={"plugin": {"posix": {"active": True}}},
|
||||
# use_direct_io defaults to True (env var)
|
||||
)
|
||||
try:
|
||||
return HiCacheNixl(storage_config=storage_config, file_path=self.test_dir)
|
||||
except ImportError:
|
||||
self.skipTest("NIXL not available")
|
||||
|
||||
def test_needs_page_alignment_true_for_file_backend(self):
|
||||
"""File-based backend + use_direct_io=True must set needs_page_alignment."""
|
||||
hicache = self._make_direct_io_hicache()
|
||||
self.assertTrue(hicache.needs_page_alignment)
|
||||
|
||||
def test_odirect_unaligned_pool_falls_back_to_copy(self):
|
||||
"""O_DIRECT with non-aligned pool strides falls back to copy mode."""
|
||||
hicache = self._make_direct_io_hicache()
|
||||
|
||||
mock_host = MockMemPoolHost(is_zero_copy_mode=True)
|
||||
hicache.register_mem_pool_host(mock_host)
|
||||
|
||||
# MockMemPoolHost.is_stride_page_aligned() returns False, so even though
|
||||
# the layout would otherwise enable zero-copy, the backend must fall back.
|
||||
self.assertFalse(hicache.is_zero_copy)
|
||||
self.assertIsNotNone(hicache._bounce_set)
|
||||
self.assertIsNotNone(hicache._bounce_get)
|
||||
|
||||
def test_odirect_disabled_via_config(self):
|
||||
"""Top-level use_direct_io=false in extra_config disables O_DIRECT."""
|
||||
storage_config = HiCacheStorageConfig(
|
||||
tp_rank=0,
|
||||
tp_size=1,
|
||||
pp_rank=0,
|
||||
pp_size=1,
|
||||
attn_cp_rank=0,
|
||||
attn_cp_size=1,
|
||||
is_mla_model=False,
|
||||
is_page_first_layout=False,
|
||||
model_name="test_model",
|
||||
enable_storage_metrics=False,
|
||||
extra_config={
|
||||
"plugin": {"posix": {"active": True}},
|
||||
"use_direct_io": False,
|
||||
},
|
||||
)
|
||||
try:
|
||||
hicache = HiCacheNixl(
|
||||
storage_config=storage_config, file_path=self.test_dir
|
||||
)
|
||||
except ImportError:
|
||||
self.skipTest("NIXL not available")
|
||||
self.assertFalse(hicache.needs_page_alignment)
|
||||
self.assertFalse(hicache.file_manager.use_direct_io)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user