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

atom/model_ops/engram/host.py

@@ -20,6 +20,7 @@
 import torch
 import torch.distributed as dist
 
+from atom.model_ops.engram import host_memory
 from atom.model_ops.engram.tables import HostEmbeddingTable
 from atom.utils import CpuGpuBuffer, envs
 
@@ -438,6 +439,11 @@ def _enable_uva(self, prefetcher, num_hash_heads):
             logger.info("engram: using the host path")
             self._release_shard()
             return False
+        # Only now that the group keeps the shard: a fallback above must not
+        # leave the layers registered so far locked.
+        host_memory.mlock_in_place(
+            t for table in self._pinned_tables for t in table.shard_tensors()
+        )
         logger.info(
             "engram: UVA device lookup active -- heads [%d, %d) of %d, "
             "%.1f GiB page-locked on this rank%s",

atom/model_ops/engram/host_memory.py

@@ -598,6 +598,18 @@ def release(self, rt, bases) -> None:
                         _clear_hip_error()
                         logger.warning("engram: cudaHostUnregister failed")
 
+    def covers(self, lo: int, hi: int) -> bool:
+        """Whether registered runs cover all of `[lo, hi)`."""
+        with self._lock:
+            cursor = lo
+            for base in sorted(
+                b for b, (e, _) in self._runs.items() if b < hi and e > lo
+            ):
+                if base > cursor:
+                    return False
+                cursor = max(cursor, self._runs[base][0])
+            return cursor >= hi
+
     def __len__(self) -> int:
         return len(self._runs)
 
@@ -631,6 +643,53 @@ def fault(piece) -> bool:
         list(pool.map(fault, pieces))
 
 
+def mlock_in_place(tensors) -> None:
+    """Lock in RAM the tensors `pin_rows` registered in place; copies (their
+    `hipHostMalloc` memory is locked already) are skipped.
+
+    An in-place registration pins page-cache pages through a KFD userptr BO,
+    which the kernel does not treat as locked: under memory pressure kswapd
+    reclaim or kcompactd migration still touches them, and every such MMU
+    notifier evicts all of the process's GPU queues until the BO is
+    revalidated (stalls of seconds, or a hang). `mlock` keeps reclaim off
+    these pages; set `vm.compact_unevictable_allowed=0` to keep compaction
+    off them as well.
+
+    Called once the shard is final (the TP group agreed to keep it), so a
+    fallback never leaves pages locked. Best effort: a refusal
+    (RLIMIT_MEMLOCK without CAP_IPC_LOCK) warns once and the registration
+    stands. Not undone on release -- ranges that share a page would unlock
+    each other; the lock goes with the mapping.
+    """
+    ranges = []
+    for t in tensors:
+        base = t.data_ptr()
+        lo = base - (base % _PAGE)
+        hi = -(-(base + t.numel() * t.element_size()) // _PAGE) * _PAGE
+        if hi > lo and _PAGE_RUNS.covers(lo, hi):
+            ranges.append((lo, hi))
+    if ranges:
+        _mlock(ranges)
+
+
+def _mlock(ranges) -> None:
+    """`mlock` each `[lo, hi)`; one warning for all that are refused."""
+    libc = ctypes.CDLL(None, use_errno=True)
+    libc.mlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+    refused, errno = 0, 0
+    for lo, hi in ranges:
+        if libc.mlock(lo, hi - lo) != 0:
+            refused += hi - lo
+            errno = ctypes.get_errno()
+    if refused:
+        logger.warning(
+            "engram: mlock of %d bytes failed (errno %d); reclaim may stall the "
+            "GPU queues under host memory pressure, raise RLIMIT_MEMLOCK",
+            refused,
+            errno,
+        )
+
+
 def _register(sources, leases):
     """Register the mapping in place; each range is widened to page
     boundaries, which registration requires (the extra bytes are neighbouring

atom/model_ops/engram/tables.py

@@ -159,6 +159,10 @@ def disable_uva(self) -> None:
         self._shard = ()
         self._uva = None
 
+    def shard_tensors(self) -> tuple[torch.Tensor, ...]:
+        """The page-locked rows (values, then scales); empty without a range."""
+        return self._shard
+
     def registered_shard(self):
         """`(weight, scales, row_start, row_end)` as the UVA kernel addresses them.
 

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, then `mlock`ed: the userptr registration alone leaves those pages reclaimable, and under host memory pressure each kswapd / kcompactd touch evicts all of the process's GPU queues until the registration is revalidated -- decode stalls of seconds, or a hang (watch `/sys/class/kfd/kfd/proc/<pid>/stats_*/evicted_ms`; set `vm.compact_unevictable_allowed=0` to keep compaction off them too). The lock covers the shared checkpoint page cache node-wide, so reclaim can no longer take it back: host memory pressure ends in the OOM killer rather than in a stall. A refused `mlock` (`RLIMIT_MEMLOCK` without `CAP_IPC_LOCK`) warns once and the registration stands. 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_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)

tests/model_ops/engram/test_host_path.py

@@ -1763,6 +1763,39 @@ def refused(tensor, op=None, group=None):
     rt.shutdown()
 
 
+@pytest.mark.parametrize("agree, copy", [(True, "0"), (False, "0"), (True, "1")])
+def test_in_place_shard_is_mlocked_only_once_the_group_keeps_it(
+    monkeypatch, agree, copy
+):
+    """The lock comes after the vote: an outvoted rank (or any fallback) must
+    not leave the layers it registered locked; copies are never mlocked."""
+    from atom.model_ops.engram import host as engram_host
+    from atom.model_ops.engram import host_memory
+
+    monkeypatch.setenv("ATOM_ENGRAM_UVA_COPY", copy)
+    allocator = FakeHostAllocator()
+    cudart = FakeCudart()
+    monkeypatch.setattr(host_memory, "_host_allocator", lambda: allocator)
+    monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart)
+    locked = []
+    monkeypatch.setattr(host_memory, "_mlock", lambda ranges: locked.extend(ranges))
+    patch_tp_group(monkeypatch, FakeGroup(world_size=2))
+
+    def vote(tensor, op=None, group=None):
+        if not agree:
+            tensor.fill_(0)
+
+    monkeypatch.setattr(engram_host.dist, "all_reduce", vote)
+    rt = make_runtime()
+    heads = rt.prefetcher._hash_mapping.config.num_hash_heads
+    assert rt._enable_uva(rt.prefetcher, heads) is agree
+    in_place = agree and copy == "0"
+    assert bool(locked) == in_place
+    if in_place:
+        assert len(locked) == sum(len(t.shard_tensors()) for t in rt._pinned_tables)
+    rt.shutdown()
+
+
 def test_a_table_that_raises_votes_no_instead_of_leaving_the_vote(monkeypatch):
     """An exception from one table's `enable_uva` (a table already holding
     another range, a bad range) must reach the group vote as a refusal: a

tests/model_ops/engram/test_mlock.py

@@ -0,0 +1,93 @@
+"""In-place UVA registrations are `mlock`ed: `_mlock` locks a shared read-only
+file mapping (what an in-place registration pins) and only warns when the
+kernel refuses; `mlock_in_place` locks only what is registered in place."""
+
+import ctypes
+import logging
+import mmap
+import resource
+import sys
+
+import numpy as np
+import pytest
+import torch
+
+from atom.model_ops.engram import host_memory
+
+pytestmark = pytest.mark.skipif(
+    not sys.platform.startswith("linux"), reason="mlock / VmLck are Linux-only"
+)
+
+
+def _vm_lck_kb() -> int:
+    with open("/proc/self/status") as f:
+        for line in f:
+            if line.startswith("VmLck:"):
+                return int(line.split()[1])
+    raise AssertionError("no VmLck in /proc/self/status")
+
+
+def _table(tmp_path, pages: int) -> np.memmap:
+    path = tmp_path / "table.bin"
+    path.write_bytes(b"\x01" * (pages * mmap.PAGESIZE))
+    return np.memmap(path, dtype=np.uint8, mode="r")
+
+
+def _mlock_allowed(lo: int, nbytes: int) -> bool:
+    """Whether this process may lock `nbytes` at `lo` (probe, then unlock)."""
+    libc = ctypes.CDLL(None, use_errno=True)
+    libc.mlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+    libc.munlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+    if libc.mlock(lo, nbytes) != 0:
+        return False
+    libc.munlock(lo, nbytes)
+    return True
+
+
+def test_mlock_locks_the_mapping(tmp_path):
+    table = _table(tmp_path, 4)
+    lo = table.ctypes.data
+    if not _mlock_allowed(lo, table.nbytes):
+        pytest.skip("RLIMIT_MEMLOCK too small for 4 pages")
+    before = _vm_lck_kb()
+    host_memory._mlock([(lo, lo + table.nbytes)])
+    assert _vm_lck_kb() - before >= table.nbytes // 1024
+
+
+def test_mlock_refusal_only_warns(tmp_path, caplog):
+    table = _table(tmp_path, 1)
+    lo = table.ctypes.data
+    soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK)
+    try:
+        resource.setrlimit(resource.RLIMIT_MEMLOCK, (0, hard))
+    except (ValueError, OSError):
+        pytest.skip("cannot lower RLIMIT_MEMLOCK")
+    try:
+        if _mlock_allowed(lo, table.nbytes):
+            pytest.skip("mlock allowed despite RLIMIT_MEMLOCK=0 (CAP_IPC_LOCK)")
+        before = _vm_lck_kb()
+        with caplog.at_level(logging.WARNING):
+            host_memory._mlock([(lo, lo + table.nbytes)] * 3)
+    finally:
+        resource.setrlimit(resource.RLIMIT_MEMLOCK, (soft, hard))
+    assert _vm_lck_kb() == before
+    warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
+    assert len(warnings) == 1 and "mlock" in warnings[0].getMessage()
+
+
+def test_mlock_in_place_skips_what_is_not_registered(monkeypatch):
+    """Only ranges `_PAGE_RUNS` holds (in place) are locked; a copy's memory
+    is locked by `hipHostMalloc` already."""
+    page = host_memory._PAGE
+    runs = host_memory._PageRuns()
+    monkeypatch.setattr(host_memory, "_PAGE_RUNS", runs)
+    locked = []
+    monkeypatch.setattr(host_memory, "_mlock", lambda ranges: locked.extend(ranges))
+    registered = torch.ones(3 * page, dtype=torch.uint8)
+    copied = torch.ones(3 * page, dtype=torch.uint8)
+    base = registered.data_ptr()
+    lo = base - base % page
+    hi = -(-(base + registered.numel()) // page) * page
+    runs._runs[lo] = [hi, 1]
+    host_memory.mlock_in_place([registered, copied])
+    assert locked == [(lo, hi)]