fix(engram): mlock in-place UVA page runs while registered (on by default, best effort)

atom/model_ops/engram/host_memory.py

@@ -21,7 +21,8 @@
 Ownership: each returned tensor's storage holds the lease on its memory, so
 the memory is released when the last tensor viewing it is collected -- never
 while one is alive, and not leaked when its owner is dropped without an
-explicit release.
+explicit release. In-place runs are also `mlock`ed while registered and
+unlocked with their last reference (`_PageRuns`).
 
 Release and the device: `hipHostFree` synchronizes the device before freeing,
 and an in-place registration synchronizes before unregistering, so kernels
@@ -33,9 +34,11 @@
 """
 
 import ctypes
+import errno
 import functools
 import logging
 import os
+import resource
 import sys
 import threading
 from concurrent.futures import ThreadPoolExecutor, wait
@@ -50,6 +53,7 @@
 _PAGE = 4096
 _MADV_DONTNEED = 4
 _MADV_POPULATE_READ = 22  # Linux 5.14+; elsewhere EINVAL, and nothing is lost
+_MLOCK_ONFAULT = 1
 # `hipHostRegisterReadOnly`: the GPU only reads the range. Needed for a range
 # the process cannot write (the shared read-only checkpoint mapping), which
 # the default flags would ask to map writable.
@@ -75,6 +79,25 @@
 _MAX_COPY_THREADS = 16
 
 
+@functools.cache
+def _libc():
+    """This process's libc, with the calls used here declared once."""
+    libc = ctypes.CDLL(None, use_errno=True)
+    libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+    libc.mlock2.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+    libc.munlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+    libc.prctl.argtypes = [ctypes.c_int] + [ctypes.c_ulong] * 4
+    return libc
+
+
+def _page_span(t: torch.Tensor) -> tuple[int, int]:
+    """`[lo, hi)`: the whole pages `t`'s bytes lie on."""
+    base = t.data_ptr()
+    lo = base - base % _PAGE
+    hi = -(-(base + t.numel() * t.element_size()) // _PAGE) * _PAGE
+    return lo, hi
+
+
 @functools.cache
 def _hip_runtime():
     """The HIP runtime this process already runs -- the `libamdhip64` torch
@@ -162,7 +185,6 @@ def __init__(self, hip, libc):
         ]
         self._free.argtypes = [ctypes.c_void_p]
         self._prctl = libc.prctl
-        self._prctl.argtypes = [ctypes.c_int] + [ctypes.c_ulong] * 4
 
     def alloc(self, nbytes: int) -> int | None:
         was = self._prctl(_PR_GET_THP_DISABLE, 0, 0, 0, 0)
@@ -193,7 +215,7 @@ def _host_allocator():
     if hip is None:
         return None
     try:
-        return _HipHostAllocator(hip, ctypes.CDLL(None, use_errno=True))
+        return _HipHostAllocator(hip, _libc())
     except (OSError, AttributeError):
         logger.warning("engram: hipHostMalloc unavailable", exc_info=True)
         return None
@@ -386,8 +408,7 @@ def _source_dropper(src: torch.Tensor):
         pagemap = os.open("/proc/self/pagemap", os.O_RDONLY)
     except OSError:
         return None
-    libc = ctypes.CDLL(None, use_errno=True)
-    libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+    libc = _libc()
 
     def unmodified(lo: int, hi: int) -> bool:
         first, count = lo // _PAGE, (hi - lo) // _PAGE
@@ -546,7 +567,8 @@ def _read_only(lo: int, hi: int) -> bool:
 
 
 class _PageRuns:
-    """This process's in-place page-locked runs, reference-counted.
+    """This process's in-place page-locked runs, reference-counted, each
+    `mlock`ed for as long as it is registered.
 
     `cudaHostRegister` needs page boundaries and refuses a range overlapping
     one already registered. Ranges stored back to back share their boundary
@@ -555,6 +577,17 @@ class _PageRuns:
     registers both whole -- so each range takes a reference on the runs that
     cover it and registers only the gaps between them. A run is unregistered
     with its last reference.
+
+    Why the `mlock`: a registration pins page-cache pages through a KFD
+    userptr BO, which the kernel still treats as reclaimable. Under host
+    memory pressure kswapd reclaim touches them, and every such MMU notifier
+    evicts all of the process's GPU queues until the BO is revalidated --
+    decode stalls of seconds, or a hang. `mlock` keeps reclaim off the pages;
+    kcompactd still migrates locked pages, firing the same notifier, unless
+    `vm.compact_unevictable_allowed=0`. Runs never overlap, so a run is
+    locked when it is created (`_fault_and_lock`, which also faults it in)
+    and unlocked when its last reference goes, before it is unregistered. A
+    refused lock warns; the registration stands.
     """
 
     def __init__(self):
@@ -563,8 +596,9 @@ def __init__(self):
 
     def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None:
         """References on the runs covering `[lo, hi)` (page-aligned), the
-        uncovered gaps registered first; None, holding nothing, if HIP
-        refuses one."""
+        uncovered gaps locked and registered first; None, holding nothing, if
+        HIP refuses one. A read-only range (`flags`) is faulted in even where
+        its lock is refused."""
         with self._lock:
             held: list[int] = []
             cursor = lo
@@ -574,8 +608,11 @@ def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None:
             for base in [*overlapping, None]:
                 gap_end = hi if base is None else base
                 if cursor < gap_end:
+                    read_only = bool(flags & _HOST_REGISTER_READ_ONLY)
+                    _fault_and_lock(cursor, gap_end, read_only)
                     if int(rt.cudaHostRegister(cursor, gap_end - cursor, flags)):
                         _clear_hip_error()
+                        _libc().munlock(cursor, gap_end - cursor)
                         self.release(rt, held)
                         return None
                     self._runs[cursor] = [gap_end, 1]
@@ -594,6 +631,7 @@ def release(self, rt, bases) -> None:
                 run[1] -= 1
                 if run[1] == 0:
                     del self._runs[base]
+                    _libc().munlock(base, run[0] - base)
                     if int(rt.cudaHostUnregister(base) or 0):
                         _clear_hip_error()
                         logger.warning("engram: cudaHostUnregister failed")
@@ -605,16 +643,29 @@ def __len__(self) -> int:
 _PAGE_RUNS = _PageRuns()
 
 
-def _populate(lo: int, hi: int) -> None:
-    """Fault `[lo, hi)` in ahead of registering it, on several threads.
-
-    Registration would otherwise fault every page itself, on one thread and
-    under the HIP runtime lock (which stalls GPU allocations elsewhere in the
-    process). Best effort: a kernel without `MADV_POPULATE_READ` declines and
-    registration faults the pages as before.
+def _fault_and_lock(lo: int, hi: int, populate: bool) -> None:
+    """Lock `[lo, hi)` and fault it in on several threads, ahead of
+    registering it.
+
+    `mlock2(MLOCK_ONFAULT)` marks the whole range locked in one call without
+    faulting it; `MADV_POPULATE_READ` per chunk then faults the pages in, in
+    parallel, and each is locked as it arrives. (A plain `mlock` faults on
+    the calling thread, and per-chunk `mlock`s scale poorly: on a warm 8 GiB
+    shared read-only mapping, 16 threads, this takes
+    0.27 s against 0.40 s for per-chunk `mlock` and 0.49 s for populate then
+    one `mlock`; populate alone 0.15 s.) Registration would otherwise fault
+    every page itself, on one thread and under the HIP runtime lock (which
+    stalls GPU allocations elsewhere in the process). Only a read-only range
+    (`populate`) is faulted in here, as before; the rest is locked as
+    registration faults it. A kernel without `MADV_POPULATE_READ` declines
+    and registration faults the pages as before. A refused lock warns once
+    per process and the range is still faulted in.
     """
-    libc = ctypes.CDLL(None, use_errno=True)
-    libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+    libc = _libc()
+    if libc.mlock2(lo, hi - lo, _MLOCK_ONFAULT) != 0:
+        _warn_unlocked(ctypes.get_errno(), hi - lo)
+    if not populate:
+        return
     pieces = _chunks(lo, hi - lo, _CHUNK_BYTES)
 
     def fault(piece) -> bool:
@@ -631,6 +682,27 @@ def fault(piece) -> bool:
         list(pool.map(fault, pieces))
 
 
+def _warn_unlocked(err: int, nbytes: int) -> None:
+    if _UNLOCKED_WARNED:
+        return
+    _UNLOCKED_WARNED.append(True)
+    hint = ""
+    soft, _ = resource.getrlimit(resource.RLIMIT_MEMLOCK)
+    if err in (errno.EPERM, errno.ENOMEM) and soft != resource.RLIM_INFINITY:
+        hint = f"; RLIMIT_MEMLOCK is {soft / 1024**3:.2f} GiB, raise it"
+    logger.warning(
+        "engram: mlock of %.2f GiB refused (%s)%s; reclaim can then stall "
+        "the GPU queues under host memory pressure",
+        nbytes / 1024**3,
+        os.strerror(err),
+        hint,
+    )
+
+
+# Set once a refused `mlock` has been logged (`_fault_and_lock`).
+_UNLOCKED_WARNED: list = []
+
+
 def _register(sources, leases):
     """Register the mapping in place; each range is widened to page
     boundaries, which registration requires (the extra bytes are neighbouring
@@ -644,11 +716,8 @@ def _register(sources, leases):
     for src in sources:
         base = src.data_ptr()
         nbytes = src.numel() * src.element_size()
-        lo = base - (base % _PAGE)
-        hi = -(-(base + nbytes) // _PAGE) * _PAGE
+        lo, hi = _page_span(src)
         flags = _HOST_REGISTER_READ_ONLY if _read_only(lo, hi) else 0
-        if flags:
-            _populate(lo, hi)
         runs = _PAGE_RUNS.acquire(rt, lo, hi, flags)
         if runs is None:
             logger.info("engram: page-lock registration of %d bytes refused", hi - lo)

docs/environment_variables.md

@@ -318,7 +318,7 @@ some layers and the device path for others, which is the confusing state.
 | Variable | Type | Default | Description |
 |----------|------|---------|-------------|
 | **ATOM_ENGRAM_UVA** | bool | 1 (true) | Page-lock this rank's shard of the hash tables (copied into page-locked memory by default, see `ATOM_ENGRAM_UVA_COPY`) and let a device kernel read the rows it needs across the bus, dequantizing there. No HBM for the table. `0` falls back to gathering the rows on the host, which returns the same rows but costs ~50 ms of CPU per decode step with the GPU idle behind it. Anything that would make the device path unsafe — no CUDA, more TP ranks than hash heads, a registration that will not fit — falls back on its own, so the switch is for taking the host path deliberately. The fallback is the whole TP group's: the lookup ends in an all-gather, so one rank that cannot register turns every rank around rather than leaving the others in a collective it never enters. |
-| **ATOM_ENGRAM_UVA_COPY** | bool | unset: 1, except under DP attention | How the UVA path page-locks this rank's shard. Unset, a TP or single-GPU rank copies (`1`), and a DP-attention rank -- whose TP group is one GPU, so it holds the whole table, the same one as every other DP rank of the node -- registers in place (`0`), copying only if that is refused: a copy per rank is ~183 GiB each for V4.1-Flash (~730 GiB on DEP4), the in-place registration pins the shared checkpoint page cache once per node. Set explicitly, it decides for every rank. `1` copies the rows into page-locked memory of exactly their size (`hipHostMalloc`, on base pages: with transparent huge pages it compacts memory once the page cache has fragmented it, ~197 s per TP2 rank instead of ~12 s): one `copy_` below 512 MiB, otherwise one thread per 128 MiB but at least 8, at most 16 and the CPUs this process may run on (two or three threads lose to one `copy_`); each copied chunk of the checkpoint is dropped from the process (`madvise(DONTNEED)`; it stays in the page cache, reclaimable), so the shard is resident once and the peak stays near one shard. `0` registers the checkpoint mapping in place (`hipHostRegister`). The Engram tables are mapped shared and read-only, so this pins the page cache itself -- the same pages for every process that maps the file, no copy -- faulted in on several threads first and registered read-only. Ranges that share a page (a layer's scales right before its values) share one reference-counted registration; a refused registration (or any failed HIP call) clears HIP's last error so the host-path fallback still runs. The page-locked memory is owned by the tensors the kernel reads and released with the last of them. Needs the HIP runtime -- the one torch loaded, found in `/proc/self/maps` -- and without it the shard is registered in place. `tools/engram/bench_uva_shard.py` reproduces per-rank time, memory and gather latency for either mode. |
+| **ATOM_ENGRAM_UVA_COPY** | bool | unset: 1, except under DP attention | How the UVA path page-locks this rank's shard. Unset, a TP or single-GPU rank copies (`1`), and a DP-attention rank -- whose TP group is one GPU, so it holds the whole table, the same one as every other DP rank of the node -- registers in place (`0`), copying only if that is refused: a copy per rank is ~183 GiB each for V4.1-Flash (~730 GiB on DEP4), the in-place registration pins the shared checkpoint page cache once per node. Set explicitly, it decides for every rank. `1` copies the rows into page-locked memory of exactly their size (`hipHostMalloc`, on base pages: with transparent huge pages it compacts memory once the page cache has fragmented it, ~197 s per TP2 rank instead of ~12 s): one `copy_` below 512 MiB, otherwise one thread per 128 MiB but at least 8, at most 16 and the CPUs this process may run on (two or three threads lose to one `copy_`); each copied chunk of the checkpoint is dropped from the process (`madvise(DONTNEED)`; it stays in the page cache, reclaimable), so the shard is resident once and the peak stays near one shard. `0` registers the checkpoint mapping in place (`hipHostRegister`). The Engram tables are mapped shared and read-only, so this pins the page cache itself -- the same pages for every process that maps the file, no copy -- faulted in on several threads first and registered read-only, and `mlock`ed while registered (faulted in and locked in one parallel pass, unlocked when the last tensor over it goes). The lock covers the shared checkpoint page cache node-wide, so host memory pressure ends in the OOM killer rather than in reclaim. A refused `mlock` (`RLIMIT_MEMLOCK` without `CAP_IPC_LOCK`) warns once and the registration stands. **Required host settings for in-place registration:** `vm.compact_unevictable_allowed=0` and `vm.compaction_proactiveness=0`; with the defaults kcompactd migrates the registered pages, locked or not, and each migration evicts the process's GPU queues (decode stalls of seconds, or a hang; watch `/sys/class/kfd/kfd/proc/<pid>/stats_*/evicted_ms`). Ranges that share a page (a layer's scales right before its values) share one reference-counted registration; a refused registration (or any failed HIP call) clears HIP's last error so the host-path fallback still runs. The page-locked memory is owned by the tensors the kernel reads and released (and unlocked) with the last of them. Needs the HIP runtime -- the one torch loaded, found in `/proc/self/maps` -- and without it the shard is registered in place. `tools/engram/bench_uva_shard.py` reproduces per-rank time, memory and gather latency for either mode. |
 | **ATOM_ENGRAM_CACHE_DIR** | path | `~/.cache/atom/engram` | Where the compressed-vocab table is cached between runs. The table is reproducible from the tokenizer, so this only trades startup time for disk; point it at shared storage to let several servers build it once. A truncated or stale cache is rebuilt rather than raised. |
 
 ## Attention side streams (DeepSeek-V4.1)