From f108e69152181e357d3d5bb32f42f2d02387b19c Mon Sep 17 00:00:00 2001 From: "Verma, Karan" Date: Tue, 3 Nov 2026 00:25:56 +0000 Subject: [PATCH 1/7] fix(engram): serialize UVA registration across DP ranks and register in 1 GiB chunks Four DP-attention ranks registering the shared read-only Engram mapping at once hang in the KFD ioctl, and a single ~92 GiB cudaHostRegister sits uninterruptible. Hold a cross-process flock (ATOM_ENGRAM_REG_LOCK) around registration and register each gap in ATOM_ENGRAM_REG_CHUNK_MB pieces (default 1024; 0 keeps one registration per gap). ATOM_ENGRAM_MLOCK=1 pins the faulted-in table first so a cold reclaim cannot push registration back to disk. Co-authored-by: Cursor --- atom/model_ops/engram/host_memory.py | 93 +++++++++++++++++++++++++--- 1 file changed, 86 insertions(+), 7 deletions(-) diff --git a/atom/model_ops/engram/host_memory.py b/atom/model_ops/engram/host_memory.py index c729772e68..6fd24a09d0 100644 --- a/atom/model_ops/engram/host_memory.py +++ b/atom/model_ops/engram/host_memory.py @@ -32,7 +32,9 @@ during a capture is skipped, with a warning, rather than breaking it. """ +import contextlib import ctypes +import fcntl import functools import logging import os @@ -545,6 +547,36 @@ def _read_only(lo: int, hi: int) -> bool: return False +# One cudaHostRegister of a whole Engram table is 90 GiB or more. Four DP +# ranks doing that together never return (the thread sits in D state inside +# the ioctl, and the "registered N GiB" log never prints). Default to 1 GiB +# pieces. ATOM_ENGRAM_REG_CHUNK_MB=0 keeps the old one-BO-per-gap behavior. +_REG_CHUNK = ( + (int(os.environ.get("ATOM_ENGRAM_REG_CHUNK_MB", "1024") or "0") << 20) + // _PAGE + * _PAGE +) +_REG_LOCK_PATH = os.environ.get( + "ATOM_ENGRAM_REG_LOCK", "/dev/shm/engram-uva-register.lock" +) + + +@contextlib.contextmanager +def _engram_register_lock(): + """Hold an exclusive lock so one process at a time page-locks Engram tables. + + The tables are one shared read-only checkpoint mapping. Concurrent + cudaHostRegister calls on it deadlock in the HIP/KFD ioctl. + """ + os.makedirs(os.path.dirname(_REG_LOCK_PATH) or ".", exist_ok=True) + with open(_REG_LOCK_PATH, "a+") as handle: + fcntl.flock(handle, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(handle, fcntl.LOCK_UN) + + class _PageRuns: """This process's in-place page-locked runs, reference-counted. @@ -574,12 +606,30 @@ 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: - if int(rt.cudaHostRegister(cursor, gap_end - cursor, flags)): - _clear_hip_error() - self.release(rt, held) - return None - self._runs[cursor] = [gap_end, 1] - held.append(cursor) + # Register in _REG_CHUNK pieces (ATOM_ENGRAM_REG_CHUNK_MB). + step = _REG_CHUNK if _REG_CHUNK > 0 else gap_end - cursor + c = cursor + while c < gap_end: + e = min(gap_end, c + step) + nbytes = e - c + if nbytes >= (1 << 30): + logger.info( + "engram: cudaHostRegister %.2f GiB", + nbytes / 2**30, + ) + if int(rt.cudaHostRegister(c, nbytes, flags)): + _clear_hip_error() + self.release(rt, held) + return None + self._runs[c] = [e, 1] + held.append(c) + c = e + if gap_end - cursor >= (1 << 30): + logger.info( + "engram: registered %.2f GiB in-place as %d userptr BO(s)", + (gap_end - cursor) / 2**30, + -(-(gap_end - cursor) // step), + ) if base is None: break self._runs[base][1] += 1 @@ -631,15 +681,42 @@ def fault(piece) -> bool: list(pool.map(fault, pieces)) +# ATOM_ENGRAM_MLOCK=1 pins the faulted-in table before registering it, so +# the kernel cannot reclaim its pages between _populate and cudaHostRegister +# (a cold reclaim makes registration fault the table back in from disk). The +# process needs an unlimited (or table-sized) RLIMIT_MEMLOCK. +_MLOCK = os.environ.get("ATOM_ENGRAM_MLOCK", "0") == "1" + + +def _mlock(lo: int, hi: int) -> None: + libc = ctypes.CDLL(None, use_errno=True) + libc.mlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + if libc.mlock(lo, hi - lo) != 0: + logger.warning( + "engram: mlock of %d bytes failed: errno %d", hi - lo, ctypes.get_errno() + ) + elif hi - lo >= (1 << 30): + logger.info("engram: mlocked %.2f GiB", (hi - lo) / 2**30) + + def _register(sources, leases): """Register the mapping in place; each range is widened to page boundaries, which registration requires (the extra bytes are neighbouring rows the kernel never addresses), and shares any page another range of this process already holds (`_PageRuns`). A read-only mapping (the shared checkpoint mapping) is faulted in on several threads first and registered - read-only.""" + read-only. + + The whole function holds a cross-process lock. DP attention registers the + same checkpoint pages on every rank, and overlapping cudaHostRegister + calls hang. + """ rt = torch.cuda.cudart() + with _engram_register_lock(): + return _register_locked(rt, sources, leases) + +def _register_locked(rt, sources, leases): out = [] for src in sources: base = src.data_ptr() @@ -649,6 +726,8 @@ def _register(sources, leases): flags = _HOST_REGISTER_READ_ONLY if _read_only(lo, hi) else 0 if flags: _populate(lo, hi) + if _MLOCK: + _mlock(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) From a9c69ae9d524446580804b72b80cff3f93cbfb40 Mon Sep 17 00:00:00 2001 From: "Verma, Karan" Date: Tue, 3 Nov 2026 01:06:29 +0000 Subject: [PATCH 2/7] fix(engram): read registration knobs from envs, log lock waits, lock unregistration - ATOM_ENGRAM_REG_CHUNK_MB, ATOM_ENGRAM_REG_LOCK and ATOM_ENGRAM_MLOCK move to envs.py with docs rows. The chunk size uses _nonnegative_int_env, so an empty value reads as the default instead of turning chunking off. - A lock file that cannot be opened or locked is logged and skipped instead of raising, which read as a refused registration and sent the rank to the copy path. - A rank that has to wait for the lock logs the wait and its length. - Unregistration takes the same lock, before _PageRuns._lock as registration does. The lock is reentrant within the process for the rollback path. - Tests cover the chunk split, a refusal part way through a gap, chunk size 0 and empty, an unopenable lock file, and the lock around unregistration. Co-authored-by: Cursor --- atom/model_ops/engram/host_memory.py | 122 +++++++++++++++++------ atom/utils/envs.py | 15 +++ docs/environment_variables.md | 3 + tests/model_ops/engram/test_host_path.py | 108 +++++++++++++++++++- 4 files changed, 214 insertions(+), 34 deletions(-) diff --git a/atom/model_ops/engram/host_memory.py b/atom/model_ops/engram/host_memory.py index 6fd24a09d0..ee4166ace8 100644 --- a/atom/model_ops/engram/host_memory.py +++ b/atom/model_ops/engram/host_memory.py @@ -40,6 +40,7 @@ import os import sys import threading +import time from concurrent.futures import ThreadPoolExecutor, wait import numpy as np @@ -547,34 +548,92 @@ def _read_only(lo: int, hi: int) -> bool: return False -# One cudaHostRegister of a whole Engram table is 90 GiB or more. Four DP -# ranks doing that together never return (the thread sits in D state inside -# the ioctl, and the "registered N GiB" log never prints). Default to 1 GiB -# pieces. ATOM_ENGRAM_REG_CHUNK_MB=0 keeps the old one-BO-per-gap behavior. -_REG_CHUNK = ( - (int(os.environ.get("ATOM_ENGRAM_REG_CHUNK_MB", "1024") or "0") << 20) - // _PAGE - * _PAGE -) -_REG_LOCK_PATH = os.environ.get( - "ATOM_ENGRAM_REG_LOCK", "/dev/shm/engram-uva-register.lock" -) +def _reg_chunk_bytes() -> int: + """Largest single `cudaHostRegister` (`ATOM_ENGRAM_REG_CHUNK_MB`), in whole + pages; 0 registers each gap in one call. + + One registration of a whole Engram table is 90 GiB or more, and four DP + ranks doing that together never return (the thread sits in D state inside + the KFD ioctl). + """ + return (envs.ATOM_ENGRAM_REG_CHUNK_MB << 20) // _PAGE * _PAGE + + +# This process's hold on the registration lock. Reentrant: a refused +# registration unregisters what it already took while still holding it. +_REG_LOCK_THREADS = threading.RLock() +_reg_lock_depth = 0 + + +def _flock_logging_wait(fd: int, path: str) -> None: + """Exclusive `flock` on `fd`; a wait is logged, so a stalled holder shows + in every waiting rank's log.""" + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + return + except BlockingIOError: + pass + logger.info("engram: waiting for %s; another process is page-locking", path) + start = time.monotonic() + fcntl.flock(fd, fcntl.LOCK_EX) + logger.info("engram: took %s after %.1f s", path, time.monotonic() - start) + + +@contextlib.contextmanager +def _register_lock_file(): + """Hold the node-wide lock file `ATOM_ENGRAM_REG_LOCK`. + + A lock file that cannot be opened or locked is logged and skipped: that is + registration as it was before the lock, while raising here would read as a + refused registration and send every rank to the ~183 GiB-per-rank copy. + """ + path = envs.ATOM_ENGRAM_REG_LOCK + try: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o666) + except OSError as exc: + logger.warning( + "engram: cannot open registration lock %s (%s); page-locking without it", + path, + exc, + ) + yield + return + try: + try: + _flock_logging_wait(fd, path) + except OSError as exc: + logger.warning( + "engram: cannot lock %s (%s); page-locking without it", path, exc + ) + yield + return + try: + yield + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) @contextlib.contextmanager def _engram_register_lock(): - """Hold an exclusive lock so one process at a time page-locks Engram tables. + """Page-lock and unlock Engram ranges one process at a time, node-wide. - The tables are one shared read-only checkpoint mapping. Concurrent - cudaHostRegister calls on it deadlock in the HIP/KFD ioctl. + Every DP-attention rank page-locks the same shared read-only checkpoint + mapping, and concurrent `cudaHostRegister` calls on it hang in the KFD + ioctl. Unregistration takes the same lock, so ranks also release the + shared pages one at a time. Reentrant within the process. """ - os.makedirs(os.path.dirname(_REG_LOCK_PATH) or ".", exist_ok=True) - with open(_REG_LOCK_PATH, "a+") as handle: - fcntl.flock(handle, fcntl.LOCK_EX) + global _reg_lock_depth + with _REG_LOCK_THREADS, contextlib.ExitStack() as stack: + if _reg_lock_depth == 0: + stack.enter_context(_register_lock_file()) + _reg_lock_depth += 1 try: yield finally: - fcntl.flock(handle, fcntl.LOCK_UN) + _reg_lock_depth -= 1 class _PageRuns: @@ -606,8 +665,7 @@ 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: - # Register in _REG_CHUNK pieces (ATOM_ENGRAM_REG_CHUNK_MB). - step = _REG_CHUNK if _REG_CHUNK > 0 else gap_end - cursor + step = _reg_chunk_bytes() or gap_end - cursor c = cursor while c < gap_end: e = min(gap_end, c + step) @@ -638,7 +696,9 @@ def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None: return held def release(self, rt, bases) -> None: - with self._lock: + # The node-wide lock first, as in `_register`, so the two locks are + # always taken in the same order. + with _engram_register_lock(), self._lock: for base in bases: run = self._runs[base] run[1] -= 1 @@ -681,14 +741,11 @@ def fault(piece) -> bool: list(pool.map(fault, pieces)) -# ATOM_ENGRAM_MLOCK=1 pins the faulted-in table before registering it, so -# the kernel cannot reclaim its pages between _populate and cudaHostRegister -# (a cold reclaim makes registration fault the table back in from disk). The -# process needs an unlimited (or table-sized) RLIMIT_MEMLOCK. -_MLOCK = os.environ.get("ATOM_ENGRAM_MLOCK", "0") == "1" - - def _mlock(lo: int, hi: int) -> None: + """Pin `[lo, hi)` (`ATOM_ENGRAM_MLOCK`) so the pages `_populate` faulted in + cannot be reclaimed before `cudaHostRegister`, which would otherwise fault + them back in from disk. Needs an unlimited or table-sized RLIMIT_MEMLOCK; + a failure only warns.""" libc = ctypes.CDLL(None, use_errno=True) libc.mlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] if libc.mlock(lo, hi - lo) != 0: @@ -707,9 +764,8 @@ def _register(sources, leases): checkpoint mapping) is faulted in on several threads first and registered read-only. - The whole function holds a cross-process lock. DP attention registers the - same checkpoint pages on every rank, and overlapping cudaHostRegister - calls hang. + Runs under `_engram_register_lock`: DP attention registers the same + checkpoint pages on every rank, and overlapping registrations hang. """ rt = torch.cuda.cudart() with _engram_register_lock(): @@ -726,7 +782,7 @@ def _register_locked(rt, sources, leases): flags = _HOST_REGISTER_READ_ONLY if _read_only(lo, hi) else 0 if flags: _populate(lo, hi) - if _MLOCK: + if envs.ATOM_ENGRAM_MLOCK: _mlock(lo, hi) runs = _PAGE_RUNS.acquire(rt, lo, hi, flags) if runs is None: diff --git a/atom/utils/envs.py b/atom/utils/envs.py index b6897485d1..70be76461a 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1258,6 +1258,21 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # rather than copied once per rank -- and copies only if that is refused. # `tools/engram/bench_uva_shard.py` measures both. "ATOM_ENGRAM_UVA_COPY": lambda: os.getenv("ATOM_ENGRAM_UVA_COPY", "1") == "1", + # In-place registration (ATOM_ENGRAM_UVA_COPY=0, the DP-attention default): + # the largest single cudaHostRegister in MiB; 0 registers each gap whole. + # One registration of a whole V4.1-Flash table is ~92 GiB, and four DP + # ranks doing that at once hang in the KFD ioctl. + "ATOM_ENGRAM_REG_CHUNK_MB": lambda: _nonnegative_int_env( + "ATOM_ENGRAM_REG_CHUNK_MB", 1024 + ), + # Lock file that serializes in-place registration and unregistration + # across the processes of a node. Unopenable means "register without it". + "ATOM_ENGRAM_REG_LOCK": lambda: os.getenv("ATOM_ENGRAM_REG_LOCK", "") + or "/dev/shm/engram-uva-register.lock", + # mlock each faulted-in range before registering it, so a cold reclaim + # cannot push registration back to disk. Needs RLIMIT_MEMLOCK at least the + # table size (e.g. docker --ulimit memlock=-1:-1); a failure only warns. + "ATOM_ENGRAM_MLOCK": lambda: _flag_env("ATOM_ENGRAM_MLOCK", "0"), # 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. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index d719c58b77..3962d966de 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -319,6 +319,9 @@ some layers and the device path for others, which is the confusing state. |----------|------|---------|-------------| | **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_REG_CHUNK_MB** | int | 1024 | In-place registration (`ATOM_ENGRAM_UVA_COPY=0`): the largest single `cudaHostRegister`, in MiB. Each uncovered gap is registered in pieces of this size, each its own reference-counted run, and every piece of 1 GiB or more is logged first (`engram: cudaHostRegister 1.00 GiB`). One registration of a whole V4.1-Flash table is ~92 GiB, and four DP-attention ranks doing that at once hang in the KFD ioctl. `0` registers each gap in one call. Unset, empty or invalid reads as the default. | +| **ATOM_ENGRAM_REG_LOCK** | path | `/dev/shm/engram-uva-register.lock` | Lock file that makes in-place registration and unregistration of Engram ranges run one process at a time on the node, since every DP-attention rank page-locks the same checkpoint pages. A rank that has to wait logs it and how long it waited. A lock file that cannot be opened or locked is logged and skipped, and the rank registers without it. | +| **ATOM_ENGRAM_MLOCK** | bool | 0 (false) | `mlock` each range after it is faulted in and before it is registered in place, so the kernel cannot reclaim the pages in between and send registration back to disk (cold cache: ~8 min per V4.1-Flash rank instead of ~30 s). Needs `RLIMIT_MEMLOCK` of at least the table size, e.g. `docker run --ulimit memlock=-1:-1`; a failure is logged and registration proceeds. | | **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) diff --git a/tests/model_ops/engram/test_host_path.py b/tests/model_ops/engram/test_host_path.py index b90e8698b4..3edd1c9701 100644 --- a/tests/model_ops/engram/test_host_path.py +++ b/tests/model_ops/engram/test_host_path.py @@ -604,7 +604,7 @@ def test_staging_rejects_padding_that_truncates_or_exceeds_capacity(): @pytest.fixture(autouse=True) -def _no_real_host_allocator(monkeypatch): +def _no_real_host_allocator(monkeypatch, tmp_path): """Keep every test off the real `hipHostMalloc`: the registration tests script `cudaHostRegister`, and the copy tests install `FakeHostAllocator` themselves.""" @@ -615,6 +615,8 @@ def _no_real_host_allocator(monkeypatch): # copy policy as a single-GPU or TP rank sees it by default. monkeypatch.setattr(host_memory, "_PAGE_RUNS", host_memory._PageRuns()) monkeypatch.delenv("ATOM_ENGRAM_UVA_COPY", raising=False) + # A per-test registration lock file, so tests never share /dev/shm. + monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", str(tmp_path / "register.lock")) class FakeCudart: @@ -691,6 +693,110 @@ def test_enabling_the_same_range_twice_registers_once(monkeypatch): assert cudart.calls == 2 +WIDE_VALUES = torch.ones(4, 4096) + + +def wide_table() -> HostEmbeddingTable: + """One values range of 64 KiB, so a chunk of a page or two splits it.""" + return HostEmbeddingTable(WIDE_VALUES, num_rows=4, head_dim=4096) + + +def registered_pages(values: torch.Tensor) -> int: + from atom.model_ops.engram import host_memory + + base = values.data_ptr() + nbytes = values.numel() * values.element_size() + lo = base - base % host_memory._PAGE + hi = -(-(base + nbytes) // host_memory._PAGE) * host_memory._PAGE + return (hi - lo) // host_memory._PAGE + + +def test_registration_is_split_into_chunks(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: 2 * host_memory._PAGE) + table = wide_table() + assert table.enable_uva(0, 4) is True + assert cudart.calls == -(-registered_pages(WIDE_VALUES) // 2) + table.disable_uva() + assert cudart.registered == [] + + +def test_refused_chunk_gives_back_the_chunks_before_it(monkeypatch): + """A refusal part way through a gap releases the pieces already taken. + + The release runs while the registration lock is held, so this also checks + that the lock is reentrant within the process. + """ + from atom.model_ops.engram import host_memory + + cudart = FakeCudart(refuse_call=3) + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: host_memory._PAGE) + table = wide_table() + assert table.enable_uva(0, 4) is False + assert cudart.registered == [] + assert table._uva is None + + +def test_chunk_size_zero_registers_each_gap_whole(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setenv("ATOM_ENGRAM_REG_CHUNK_MB", "0") + assert host_memory._reg_chunk_bytes() == 0 + table = wide_table() + assert table.enable_uva(0, 4) is True + assert cudart.calls == 1 + + +def test_empty_chunk_size_reads_as_the_default(monkeypatch): + from atom.model_ops.engram import host_memory + + monkeypatch.setenv("ATOM_ENGRAM_REG_CHUNK_MB", "") + assert host_memory._reg_chunk_bytes() == 1024 << 20 + + +def test_unopenable_lock_file_still_registers_in_place(monkeypatch, tmp_path): + """A lock file that cannot be opened is skipped, not treated as a refused + registration (which would send the rank to the copy path).""" + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + blocker = tmp_path / "not-a-dir" + blocker.write_text("") + monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", str(blocker / "register.lock")) + table = wide_table() + assert table.enable_uva(0, 4) is True + assert len(cudart.registered) >= 1 + + +def test_registration_and_unregistration_take_the_lock(monkeypatch): + import contextlib + + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + taken = [] + + @contextlib.contextmanager + def counting_lock_file(): + taken.append(len(cudart.registered)) + yield + + monkeypatch.setattr(host_memory, "_register_lock_file", counting_lock_file) + table = wide_table() + assert table.enable_uva(0, 4) is True + table.disable_uva() + # Once around registering (nothing held yet), once around releasing. + assert taken[0] == 0 + assert len(taken) >= 2 + assert cudart.registered == [] + + class FakeHostAllocator: """Page-locked allocations backed by ctypes buffers kept alive here.""" From 529fc7da3872c7d6642cb423bd33fe26ee8305e5 Mon Sep 17 00:00:00 2001 From: "Verma, Karan" Date: Tue, 3 Nov 2026 02:07:40 +0000 Subject: [PATCH 3/7] fix(engram): munlock with the last reference, state the lock order, log the lock path - ATOM_ENGRAM_MLOCK pins move into _PageRuns.acquire. Runs record whether they are mlocked once the range holds them, and are munlocked with their last reference; a refused registration unpins the pages no live run covers. - The release comment states the actual invariant: every path takes the node-wide lock before _PageRuns._lock, including acquire's rollback. - The resolved lock path is logged once per process, and the docs note that /dev/shm is per container unless shared. - Tests: munlock on release, and on a refusal part way through a range. Co-authored-by: Cursor --- atom/model_ops/engram/host_memory.py | 67 +++++++++++++++++++----- atom/utils/envs.py | 2 + docs/environment_variables.md | 4 +- tests/model_ops/engram/test_host_path.py | 48 +++++++++++++++++ 4 files changed, 105 insertions(+), 16 deletions(-) diff --git a/atom/model_ops/engram/host_memory.py b/atom/model_ops/engram/host_memory.py index ee4166ace8..8ad3e8b083 100644 --- a/atom/model_ops/engram/host_memory.py +++ b/atom/model_ops/engram/host_memory.py @@ -563,6 +563,7 @@ def _reg_chunk_bytes() -> int: # registration unregisters what it already took while still holding it. _REG_LOCK_THREADS = threading.RLock() _reg_lock_depth = 0 +_reg_lock_path_logged = False def _flock_logging_wait(fd: int, path: str) -> None: @@ -587,6 +588,7 @@ def _register_lock_file(): registration as it was before the lock, while raising here would read as a refused registration and send every rank to the ~183 GiB-per-rank copy. """ + global _reg_lock_path_logged path = envs.ATOM_ENGRAM_REG_LOCK try: os.makedirs(os.path.dirname(path) or ".", exist_ok=True) @@ -599,6 +601,9 @@ def _register_lock_file(): ) yield return + if not _reg_lock_path_logged: + _reg_lock_path_logged = True + logger.info("engram: page-locking serialized on %s", path) try: try: _flock_logging_wait(fd, path) @@ -645,22 +650,28 @@ class _PageRuns: shard, and a rank holding every hash head (one DP-attention rank per GPU) 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. + with its last reference, and munlocked then if it was mlocked: the runs + partition the registered pages, so no live range still needs them. """ def __init__(self): - self._runs: dict[int, list[int]] = {} # base -> [end, references] + # base -> [end, references, mlocked] + self._runs: dict[int, list] = {} self._lock = threading.RLock() - def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None: + def acquire( + self, rt, lo: int, hi: int, flags: int = 0, mlock: bool = False + ) -> list[int] | None: """References on the runs covering `[lo, hi)` (page-aligned), the uncovered gaps registered first; None, holding nothing, if HIP - refuses one.""" + refuses one. `mlock` pins the whole range first (`_mlock`); on a + refusal the pages no live run covers are unpinned again.""" with self._lock: + locked = mlock and _mlock(lo, hi) held: list[int] = [] cursor = lo overlapping = sorted( - b for b, (e, _) in self._runs.items() if b < hi and e > lo + b for b, (e, *_) in self._runs.items() if b < hi and e > lo ) for base in [*overlapping, None]: gap_end = hi if base is None else base @@ -678,8 +689,10 @@ def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None: if int(rt.cudaHostRegister(c, nbytes, flags)): _clear_hip_error() self.release(rt, held) + if locked: + self._munlock_uncovered(lo, hi) return None - self._runs[c] = [e, 1] + self._runs[c] = [e, 1, False] held.append(c) c = e if gap_end - cursor >= (1 << 30): @@ -693,11 +706,16 @@ def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None: self._runs[base][1] += 1 held.append(base) cursor = max(cursor, self._runs[base][0]) + # Marked only once the whole range holds its runs, so a refusal + # leaves the unpinning to `_munlock_uncovered` alone. + for base in held: + self._runs[base][2] = self._runs[base][2] or locked return held def release(self, rt, bases) -> None: - # The node-wide lock first, as in `_register`, so the two locks are - # always taken in the same order. + # Every path takes the node-wide lock before `self._lock`: `_register` + # holds it around `acquire`, so acquire's rollback arrives here with + # both already held (both are reentrant), and teardown takes it here. with _engram_register_lock(), self._lock: for base in bases: run = self._runs[base] @@ -707,6 +725,18 @@ def release(self, rt, bases) -> None: if int(rt.cudaHostUnregister(base) or 0): _clear_hip_error() logger.warning("engram: cudaHostUnregister failed") + if run[2]: + _munlock(base, run[0]) + + def _munlock_uncovered(self, lo: int, hi: int) -> None: + """Unpin the pages of `[lo, hi)` that no live run covers.""" + cursor = lo + for base in sorted(b for b, (e, *_) in self._runs.items() if b < hi and e > lo): + if cursor < base: + _munlock(cursor, base) + cursor = max(cursor, self._runs[base][0]) + if cursor < hi: + _munlock(cursor, hi) def __len__(self) -> int: return len(self._runs) @@ -741,19 +771,30 @@ def fault(piece) -> bool: list(pool.map(fault, pieces)) -def _mlock(lo: int, hi: int) -> None: +def _mlock(lo: int, hi: int) -> bool: """Pin `[lo, hi)` (`ATOM_ENGRAM_MLOCK`) so the pages `_populate` faulted in cannot be reclaimed before `cudaHostRegister`, which would otherwise fault them back in from disk. Needs an unlimited or table-sized RLIMIT_MEMLOCK; - a failure only warns.""" + a failure only warns. True if the range is pinned.""" libc = ctypes.CDLL(None, use_errno=True) libc.mlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] if libc.mlock(lo, hi - lo) != 0: logger.warning( "engram: mlock of %d bytes failed: errno %d", hi - lo, ctypes.get_errno() ) - elif hi - lo >= (1 << 30): + return False + if hi - lo >= (1 << 30): logger.info("engram: mlocked %.2f GiB", (hi - lo) / 2**30) + return True + + +def _munlock(lo: int, hi: int) -> None: + libc = ctypes.CDLL(None, use_errno=True) + libc.munlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + if libc.munlock(lo, hi - lo) != 0: + logger.warning( + "engram: munlock of %d bytes failed: errno %d", hi - lo, ctypes.get_errno() + ) def _register(sources, leases): @@ -782,9 +823,7 @@ def _register_locked(rt, sources, leases): flags = _HOST_REGISTER_READ_ONLY if _read_only(lo, hi) else 0 if flags: _populate(lo, hi) - if envs.ATOM_ENGRAM_MLOCK: - _mlock(lo, hi) - runs = _PAGE_RUNS.acquire(rt, lo, hi, flags) + runs = _PAGE_RUNS.acquire(rt, lo, hi, flags, mlock=envs.ATOM_ENGRAM_MLOCK) if runs is None: logger.info("engram: page-lock registration of %d bytes refused", hi - lo) return None diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 70be76461a..c0c805cd10 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1267,6 +1267,8 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: ), # Lock file that serializes in-place registration and unregistration # across the processes of a node. Unopenable means "register without it". + # /dev/shm is per container unless it is shared (docker --ipc=host uses + # the host's), so ranks in separate containers need a path they share. "ATOM_ENGRAM_REG_LOCK": lambda: os.getenv("ATOM_ENGRAM_REG_LOCK", "") or "/dev/shm/engram-uva-register.lock", # mlock each faulted-in range before registering it, so a cold reclaim diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 3962d966de..6a74730d99 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -320,8 +320,8 @@ some layers and the device path for others, which is the confusing state. | **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_REG_CHUNK_MB** | int | 1024 | In-place registration (`ATOM_ENGRAM_UVA_COPY=0`): the largest single `cudaHostRegister`, in MiB. Each uncovered gap is registered in pieces of this size, each its own reference-counted run, and every piece of 1 GiB or more is logged first (`engram: cudaHostRegister 1.00 GiB`). One registration of a whole V4.1-Flash table is ~92 GiB, and four DP-attention ranks doing that at once hang in the KFD ioctl. `0` registers each gap in one call. Unset, empty or invalid reads as the default. | -| **ATOM_ENGRAM_REG_LOCK** | path | `/dev/shm/engram-uva-register.lock` | Lock file that makes in-place registration and unregistration of Engram ranges run one process at a time on the node, since every DP-attention rank page-locks the same checkpoint pages. A rank that has to wait logs it and how long it waited. A lock file that cannot be opened or locked is logged and skipped, and the rank registers without it. | -| **ATOM_ENGRAM_MLOCK** | bool | 0 (false) | `mlock` each range after it is faulted in and before it is registered in place, so the kernel cannot reclaim the pages in between and send registration back to disk (cold cache: ~8 min per V4.1-Flash rank instead of ~30 s). Needs `RLIMIT_MEMLOCK` of at least the table size, e.g. `docker run --ulimit memlock=-1:-1`; a failure is logged and registration proceeds. | +| **ATOM_ENGRAM_REG_LOCK** | path | `/dev/shm/engram-uva-register.lock` | Lock file that makes in-place registration and unregistration of Engram ranges run one process at a time on the node, since every DP-attention rank page-locks the same checkpoint pages. A rank that has to wait logs it and how long it waited. A lock file that cannot be opened or locked is logged and skipped, and the rank registers without it. The resolved path is logged once per process. The default lives in `/dev/shm`, which is private to a container unless shared (`docker run --ipc=host` uses the host's), so ranks split across containers that do not share it must point this at a common path, e.g. a bind-mounted directory. | +| **ATOM_ENGRAM_MLOCK** | bool | 0 (false) | `mlock` each range after it is faulted in and before it is registered in place, so the kernel cannot reclaim the pages in between and send registration back to disk (cold cache: ~8 min per V4.1-Flash rank instead of ~30 s). Pages are unpinned with the last reference to their registration, or right away if the registration is refused. Needs `RLIMIT_MEMLOCK` of at least the table size, e.g. `docker run --ulimit memlock=-1:-1`; a failure is logged and registration proceeds. | | **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) diff --git a/tests/model_ops/engram/test_host_path.py b/tests/model_ops/engram/test_host_path.py index 3edd1c9701..b7647d10a9 100644 --- a/tests/model_ops/engram/test_host_path.py +++ b/tests/model_ops/engram/test_host_path.py @@ -797,6 +797,54 @@ def counting_lock_file(): assert cudart.registered == [] +def pin_recorder(monkeypatch): + """Record `_mlock` / `_munlock` ranges instead of calling libc.""" + from atom.model_ops.engram import host_memory + + locked, unlocked = [], [] + monkeypatch.setenv("ATOM_ENGRAM_MLOCK", "1") + monkeypatch.setattr( + host_memory, "_mlock", lambda lo, hi: locked.append((lo, hi)) or True + ) + monkeypatch.setattr( + host_memory, "_munlock", lambda lo, hi: unlocked.append((lo, hi)) + ) + return locked, unlocked + + +def covered_bytes(ranges) -> int: + return sum(hi - lo for lo, hi in ranges) + + +def test_mlocked_runs_are_munlocked_with_their_last_reference(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: 2 * host_memory._PAGE) + locked, unlocked = pin_recorder(monkeypatch) + table = wide_table() + assert table.enable_uva(0, 4) is True + assert covered_bytes(locked) == registered_pages(WIDE_VALUES) * host_memory._PAGE + assert unlocked == [] + table.disable_uva() + assert cudart.registered == [] + assert covered_bytes(unlocked) == covered_bytes(locked) + + +def test_refused_mlocked_registration_unpins_its_pages(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart(refuse_call=3) + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: host_memory._PAGE) + locked, unlocked = pin_recorder(monkeypatch) + table = wide_table() + assert table.enable_uva(0, 4) is False + assert cudart.registered == [] + assert covered_bytes(unlocked) == covered_bytes(locked) + + class FakeHostAllocator: """Page-locked allocations backed by ctypes buffers kept alive here.""" From c7feddded3e02f8bd68f530ddb9e5f2c45dc2f57 Mon Sep 17 00:00:00 2001 From: "Verma, Karan" Date: Tue, 3 Nov 2026 16:12:29 +0000 Subject: [PATCH 4/7] fix(engram): narrow the registration lock and drop the extra knobs Take the lock only around shared cudaHostRegister and unregister, pin when RLIMIT_MEMLOCK allows, and keep the 1 GiB chunk size fixed so a refused overlap cannot leak pins. Co-authored-by: Cursor --- atom/model_ops/engram/host_memory.py | 402 +++++++++++++++-------- atom/utils/envs.py | 20 +- docs/environment_variables.md | 6 +- tests/model_ops/engram/test_host_path.py | 258 +++++++++++++-- 4 files changed, 496 insertions(+), 190 deletions(-) diff --git a/atom/model_ops/engram/host_memory.py b/atom/model_ops/engram/host_memory.py index 8ad3e8b083..d631fce807 100644 --- a/atom/model_ops/engram/host_memory.py +++ b/atom/model_ops/engram/host_memory.py @@ -34,10 +34,13 @@ import contextlib import ctypes +import dataclasses import fcntl import functools +import hashlib import logging import os +import resource import sys import threading import time @@ -363,6 +366,23 @@ def _file_mapped(addr: int, nbytes: int) -> bool: # swapped, 61 file page (or shared anonymous). _PM_PRESENT, _PM_SWAPPED, _PM_FILE = 1 << 63, 1 << 62, 1 << 61 +# Largest single cudaHostRegister. One ioctl covering a whole Engram table +# does not return when several ranks issue it together. +_REG_CHUNK_BYTES = 1 << 30 +# How long a rank waits for the registration lock before giving up. +_REG_LOCK_TIMEOUT_S = 600 +_logged_lock_paths: set[str] = set() + + +@functools.cache +def _libc(): + """This process's libc, with the calls made here typed once.""" + libc = ctypes.CDLL(None, use_errno=True) + libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] + libc.mlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + libc.munlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + return libc + def _source_dropper(src: torch.Tensor): """`drop(offset, nbytes)` unmapping that byte range of `src`'s checkpoint @@ -389,8 +409,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 @@ -495,7 +514,7 @@ def pin_rows( if how == "copy": out = _copy(allocator, sources, leases) else: - out = _register(sources, leases) + out = _register(sources, leases, shared=shared) except Exception: # any failure is a fallback logger.warning("engram: page-locking the UVA shard failed", exc_info=True) out = None @@ -549,26 +568,49 @@ def _read_only(lo: int, hi: int) -> bool: def _reg_chunk_bytes() -> int: - """Largest single `cudaHostRegister` (`ATOM_ENGRAM_REG_CHUNK_MB`), in whole - pages; 0 registers each gap in one call. + """Largest single `cudaHostRegister`, in bytes.""" + return _REG_CHUNK_BYTES - One registration of a whole Engram table is 90 GiB or more, and four DP - ranks doing that together never return (the thread sits in D state inside - the KFD ioctl). - """ - return (envs.ATOM_ENGRAM_REG_CHUNK_MB << 20) // _PAGE * _PAGE +def _mapping_path(addr: int) -> str | None: + """Regular file mapped at `addr` (`/proc/self/maps`), or None.""" + try: + with open("/proc/self/maps") as maps: + for line in maps: + fields = line.split(maxsplit=5) + lo, hi = (int(x, 16) for x in fields[0].split("-")) + if not lo <= addr < hi: + continue + name = fields[5].strip() if len(fields) > 5 else "" + if name.startswith("/") and not name.startswith("/dev/"): + return name + return None + except (OSError, ValueError, IndexError): + return None + return None -# This process's hold on the registration lock. Reentrant: a refused -# registration unregisters what it already took while still holding it. -_REG_LOCK_THREADS = threading.RLock() -_reg_lock_depth = 0 -_reg_lock_path_logged = False + +def _registration_lock_path(sources) -> str: + """Lock file for `sources`. + + `ATOM_ENGRAM_REG_LOCK`, when set. Otherwise one file per checkpoint under + `/dev/shm`, so two servers loading different checkpoints do not share it. + """ + explicit = envs.ATOM_ENGRAM_REG_LOCK + if explicit: + return explicit + mapped = _mapping_path(sources[0].data_ptr()) if sources else None + key = os.path.realpath(mapped) if mapped else "anonymous" + digest = hashlib.sha256(key.encode()).hexdigest()[:16] + return f"/dev/shm/engram-uva-{digest}.lock" def _flock_logging_wait(fd: int, path: str) -> None: - """Exclusive `flock` on `fd`; a wait is logged, so a stalled holder shows - in every waiting rank's log.""" + """Exclusive `flock` on `fd`. + + A wait is logged. Past `_REG_LOCK_TIMEOUT_S` this raises `TimeoutError` + (an `OSError` subclass; callers must not swallow it with `except OSError`). + """ try: fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) return @@ -576,69 +618,83 @@ def _flock_logging_wait(fd: int, path: str) -> None: pass logger.info("engram: waiting for %s; another process is page-locking", path) start = time.monotonic() - fcntl.flock(fd, fcntl.LOCK_EX) - logger.info("engram: took %s after %.1f s", path, time.monotonic() - start) + deadline = start + _REG_LOCK_TIMEOUT_S + while True: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + if time.monotonic() >= deadline: + logger.error( + "engram: gave up on %s after %d s; the holder is still " + "page-locking", + path, + _REG_LOCK_TIMEOUT_S, + ) + raise TimeoutError( + f"engram registration lock {path} still held after " + f"{_REG_LOCK_TIMEOUT_S}s" + ) + time.sleep(0.05) + else: + logger.info("engram: took %s after %.1f s", path, time.monotonic() - start) + return @contextlib.contextmanager -def _register_lock_file(): - """Hold the node-wide lock file `ATOM_ENGRAM_REG_LOCK`. +def _register_lock_file(path: str): + """Hold the registration lock file `path`. - A lock file that cannot be opened or locked is logged and skipped: that is - registration as it was before the lock, while raising here would read as a - refused registration and send every rank to the ~183 GiB-per-rank copy. + A file that cannot be opened or locked is logged and skipped. The `yield` + is outside the `except`, so an error during registration is not chained + onto that failure. A wait past `_REG_LOCK_TIMEOUT_S` raises `TimeoutError`. """ - global _reg_lock_path_logged - path = envs.ATOM_ENGRAM_REG_LOCK + fd = None + held = False + open_error = None try: os.makedirs(os.path.dirname(path) or ".", exist_ok=True) fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o666) except OSError as exc: + open_error = exc + if open_error is not None: logger.warning( "engram: cannot open registration lock %s (%s); page-locking without it", path, - exc, + open_error, ) - yield - return - if not _reg_lock_path_logged: - _reg_lock_path_logged = True - logger.info("engram: page-locking serialized on %s", path) - try: + else: + if path not in _logged_lock_paths: + _logged_lock_paths.add(path) + logger.info("engram: page-locking serialized on %s", path) try: _flock_logging_wait(fd, path) + except TimeoutError: + os.close(fd) + fd = None + raise except OSError as exc: logger.warning( "engram: cannot lock %s (%s); page-locking without it", path, exc ) - yield - return - try: - yield - finally: - fcntl.flock(fd, fcntl.LOCK_UN) + else: + held = True + try: + yield finally: - os.close(fd) + if held: + fcntl.flock(fd, fcntl.LOCK_UN) + if fd is not None: + os.close(fd) -@contextlib.contextmanager -def _engram_register_lock(): - """Page-lock and unlock Engram ranges one process at a time, node-wide. +@dataclasses.dataclass +class _Run: + """One in-place registration: `[base, end)`, its references, and whether + this process pinned those pages with `mlock`.""" - Every DP-attention rank page-locks the same shared read-only checkpoint - mapping, and concurrent `cudaHostRegister` calls on it hang in the KFD - ioctl. Unregistration takes the same lock, so ranks also release the - shared pages one at a time. Reentrant within the process. - """ - global _reg_lock_depth - with _REG_LOCK_THREADS, contextlib.ExitStack() as stack: - if _reg_lock_depth == 0: - stack.enter_context(_register_lock_file()) - _reg_lock_depth += 1 - try: - yield - finally: - _reg_lock_depth -= 1 + end: int + refs: int + mlocked: bool class _PageRuns: @@ -646,98 +702,129 @@ class _PageRuns: `cudaHostRegister` needs page boundaries and refuses a range overlapping one already registered. Ranges stored back to back share their boundary - page -- a layer's scales sit right before its values in the checkpoint - shard, and a rank holding every hash head (one DP-attention rank per GPU) - 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, and munlocked then if it was mlocked: the runs - partition the registered pages, so no live range still needs them. + page, 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, and munlocked then if `mlocked`: the runs partition the + registered pages. """ def __init__(self): - # base -> [end, references, mlocked] - self._runs: dict[int, list] = {} + self._runs: dict[int, _Run] = {} self._lock = threading.RLock() def acquire( - self, rt, lo: int, hi: int, flags: int = 0, mlock: bool = False + self, rt, lo: int, hi: int, flags: int = 0, mlocked: bool = False ) -> list[int] | None: - """References on the runs covering `[lo, hi)` (page-aligned), the - uncovered gaps registered first; None, holding nothing, if HIP - refuses one. `mlock` pins the whole range first (`_mlock`); on a - refusal the pages no live run covers are unpinned again.""" + """References on the runs covering `[lo, hi)` (page-aligned). + + Uncovered gaps are registered in `_REG_CHUNK_BYTES` pieces. None, + holding nothing, if HIP refuses one. `mlocked` means the caller + already pinned `[lo, hi)`; on a refusal the pages no already-pinned + run covers are unpinned. Rollback calls `_release_locked` and does + not take the registration lock; a shared registration already holds it. + """ with self._lock: - locked = mlock and _mlock(lo, hi) held: list[int] = [] + created = 0 + created_bytes = 0 cursor = lo overlapping = sorted( - b for b, (e, *_) in self._runs.items() if b < hi and e > lo + b for b, run in self._runs.items() if b < hi and run.end > lo ) for base in [*overlapping, None]: gap_end = hi if base is None else base if cursor < gap_end: - step = _reg_chunk_bytes() or gap_end - cursor + step = _reg_chunk_bytes() c = cursor while c < gap_end: e = min(gap_end, c + step) nbytes = e - c - if nbytes >= (1 << 30): - logger.info( + if nbytes >= _REG_CHUNK_BYTES: + logger.debug( "engram: cudaHostRegister %.2f GiB", - nbytes / 2**30, + nbytes / _REG_CHUNK_BYTES, ) if int(rt.cudaHostRegister(c, nbytes, flags)): _clear_hip_error() - self.release(rt, held) - if locked: + self._release_locked(rt, held) + if mlocked: self._munlock_uncovered(lo, hi) return None - self._runs[c] = [e, 1, False] + self._runs[c] = _Run(e, 1, False) held.append(c) + created += 1 + created_bytes += nbytes c = e - if gap_end - cursor >= (1 << 30): - logger.info( - "engram: registered %.2f GiB in-place as %d userptr BO(s)", - (gap_end - cursor) / 2**30, - -(-(gap_end - cursor) // step), - ) if base is None: break - self._runs[base][1] += 1 + self._runs[base].refs += 1 held.append(base) - cursor = max(cursor, self._runs[base][0]) + cursor = max(cursor, self._runs[base].end) # Marked only once the whole range holds its runs, so a refusal # leaves the unpinning to `_munlock_uncovered` alone. for base in held: - self._runs[base][2] = self._runs[base][2] or locked + run = self._runs[base] + run.mlocked = run.mlocked or mlocked + if created and created_bytes >= _REG_CHUNK_BYTES: + logger.info( + "engram: registered %.2f GiB in-place as %d userptr BO(s)", + created_bytes / _REG_CHUNK_BYTES, + created, + ) return held - def release(self, rt, bases) -> None: - # Every path takes the node-wide lock before `self._lock`: `_register` - # holds it around `acquire`, so acquire's rollback arrives here with - # both already held (both are reentrant), and teardown takes it here. - with _engram_register_lock(), self._lock: - for base in bases: - run = self._runs[base] - run[1] -= 1 - if run[1] == 0: - del self._runs[base] - if int(rt.cudaHostUnregister(base) or 0): - _clear_hip_error() - logger.warning("engram: cudaHostUnregister failed") - if run[2]: - _munlock(base, run[0]) + def release(self, rt, bases, lock_path: str | None = None) -> None: + """Drop references taken by `acquire`. + + Lock order, when `lock_path` is set: that file, then `self._lock`. + """ + if lock_path: + with _register_lock_file(lock_path), self._lock: + self._release_locked(rt, bases) + else: + with self._lock: + self._release_locked(rt, bases) + + def _release_locked(self, rt, bases) -> None: + """`self._lock` is held by the caller. Does not take the registration lock.""" + for base in bases: + run = self._runs[base] + run.refs -= 1 + if run.refs == 0: + del self._runs[base] + if int(rt.cudaHostUnregister(base) or 0): + _clear_hip_error() + logger.warning("engram: cudaHostUnregister failed") + if run.mlocked: + _munlock(base, run.end) def _munlock_uncovered(self, lo: int, hi: int) -> None: - """Unpin the pages of `[lo, hi)` that no live run covers.""" + """Unpin pages of `[lo, hi)` that no already-pinned run covers. + + `mlock` pins the whole range, including pages of an older run whose + own pin failed (`mlocked` false). That run will not unpin them on + release. `munlock` is not reference-counted, so a run that is already + pinned is left alone: this call's `mlock` was a no-op on its pages. + """ cursor = lo - for base in sorted(b for b, (e, *_) in self._runs.items() if b < hi and e > lo): + overlapping = sorted( + b for b, run in self._runs.items() if b < hi and run.end > lo + ) + for base in overlapping: + run = self._runs[base] + if not run.mlocked: + continue if cursor < base: - _munlock(cursor, base) - cursor = max(cursor, self._runs[base][0]) + _munlock(cursor, min(base, hi)) + cursor = max(cursor, run.end) if cursor < hi: _munlock(cursor, hi) + def unpin(self, lo: int, hi: int) -> None: + """Undo an `mlock` of `[lo, hi)` on pages no live pinned run covers.""" + with self._lock: + self._munlock_uncovered(lo, hi) + def __len__(self) -> int: return len(self._runs) @@ -753,8 +840,7 @@ def _populate(lo: int, hi: int) -> None: process). Best effort: a kernel without `MADV_POPULATE_READ` declines and registration faults the pages as before. """ - libc = ctypes.CDLL(None, use_errno=True) - libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] + libc = _libc() pieces = _chunks(lo, hi - lo, _CHUNK_BYTES) def fault(piece) -> bool: @@ -771,50 +857,61 @@ def fault(piece) -> bool: list(pool.map(fault, pieces)) +def _memlock_fits(nbytes: int) -> bool: + """Whether `RLIMIT_MEMLOCK` can pin `nbytes`. + + Unlimited is the deployment that wants this (`--ulimit memlock=-1`). A + finite limit smaller than the range skips the call. + """ + try: + soft, _hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) + except (ValueError, OSError): + return False + if soft == resource.RLIM_INFINITY: + return True + return soft >= nbytes + + def _mlock(lo: int, hi: int) -> bool: - """Pin `[lo, hi)` (`ATOM_ENGRAM_MLOCK`) so the pages `_populate` faulted in - cannot be reclaimed before `cudaHostRegister`, which would otherwise fault - them back in from disk. Needs an unlimited or table-sized RLIMIT_MEMLOCK; - a failure only warns. True if the range is pinned.""" - libc = ctypes.CDLL(None, use_errno=True) - libc.mlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] - if libc.mlock(lo, hi - lo) != 0: + """Pin `[lo, hi)`. False, after a warning, if the call fails.""" + if _libc().mlock(lo, hi - lo) != 0: logger.warning( "engram: mlock of %d bytes failed: errno %d", hi - lo, ctypes.get_errno() ) return False - if hi - lo >= (1 << 30): - logger.info("engram: mlocked %.2f GiB", (hi - lo) / 2**30) + if hi - lo >= _REG_CHUNK_BYTES: + logger.info("engram: mlocked %.2f GiB", (hi - lo) / _REG_CHUNK_BYTES) return True def _munlock(lo: int, hi: int) -> None: - libc = ctypes.CDLL(None, use_errno=True) - libc.munlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] - if libc.munlock(lo, hi - lo) != 0: + if _libc().munlock(lo, hi - lo) != 0: logger.warning( "engram: munlock of %d bytes failed: errno %d", hi - lo, ctypes.get_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 - rows the kernel never addresses), and shares any page another range of - this process already holds (`_PageRuns`). A read-only mapping (the shared - checkpoint mapping) is faulted in on several threads first and registered - read-only. +def _pin_range(lo: int, hi: int) -> bool: + """`mlock` `[lo, hi)` when the memlock limit can hold it.""" + if hi <= lo or not _memlock_fits(hi - lo): + return False + return _mlock(lo, hi) - Runs under `_engram_register_lock`: DP attention registers the same - checkpoint pages on every rank, and overlapping registrations hang. - """ - rt = torch.cuda.cudart() - with _engram_register_lock(): - return _register_locked(rt, sources, leases) +def _register(sources, leases, *, shared: bool): + """Register each mapping in place. -def _register_locked(rt, sources, leases): - out = [] + Ranges are widened to page boundaries and share a page another range of + this process already holds (`_PageRuns`). A read-only mapping is faulted + in first. Fault-in and `mlock` stay outside the registration lock. That + lock is taken only for a `shared` table: around the register calls here, + and around the matching unregister. + """ + if not sources: + return () + rt = torch.cuda.cudart() + prepared = [] + pinned_ranges = [] for src in sources: base = src.data_ptr() nbytes = src.numel() * src.element_size() @@ -823,15 +920,42 @@ def _register_locked(rt, sources, leases): 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, mlock=envs.ATOM_ENGRAM_MLOCK) + pinned = _pin_range(lo, hi) + if pinned: + pinned_ranges.append((lo, hi)) + prepared.append((src, base, nbytes, lo, hi, flags, pinned)) + lock_path = _registration_lock_path(sources) if shared else None + # Ranges still pinned and not yet handed to `acquire`. A refusal unpins + # its own range; this covers a lock timeout and anything not attempted. + pending = pinned_ranges + result = None + try: + if lock_path: + with _register_lock_file(lock_path): + result = _register_prepared(rt, prepared, leases, lock_path, pending) + else: + result = _register_prepared(rt, prepared, leases, None, pending) + return result + finally: + if result is None: + for lo, hi in pending: + _PAGE_RUNS.unpin(lo, hi) + + +def _register_prepared(rt, prepared, leases, lock_path, pending): + out = [] + for src, base, nbytes, lo, hi, flags, pinned in prepared: + runs = _PAGE_RUNS.acquire(rt, lo, hi, flags, mlocked=pinned) + if pinned: + pending.remove((lo, hi)) if runs is None: logger.info("engram: page-lock registration of %d bytes refused", hi - lo) return None - def unregister(_ptr, runs=runs): + def unregister(_ptr, runs=runs, lock_path=lock_path): if torch.cuda.is_initialized(): torch.cuda.synchronize() - _PAGE_RUNS.release(rt, runs) + _PAGE_RUNS.release(rt, runs, lock_path=lock_path) lease = _Lease(lo, hi - lo, unregister, keep=src) leases.append(lease) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index c0c805cd10..d0829c8846 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1258,23 +1258,9 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # rather than copied once per rank -- and copies only if that is refused. # `tools/engram/bench_uva_shard.py` measures both. "ATOM_ENGRAM_UVA_COPY": lambda: os.getenv("ATOM_ENGRAM_UVA_COPY", "1") == "1", - # In-place registration (ATOM_ENGRAM_UVA_COPY=0, the DP-attention default): - # the largest single cudaHostRegister in MiB; 0 registers each gap whole. - # One registration of a whole V4.1-Flash table is ~92 GiB, and four DP - # ranks doing that at once hang in the KFD ioctl. - "ATOM_ENGRAM_REG_CHUNK_MB": lambda: _nonnegative_int_env( - "ATOM_ENGRAM_REG_CHUNK_MB", 1024 - ), - # Lock file that serializes in-place registration and unregistration - # across the processes of a node. Unopenable means "register without it". - # /dev/shm is per container unless it is shared (docker --ipc=host uses - # the host's), so ranks in separate containers need a path they share. - "ATOM_ENGRAM_REG_LOCK": lambda: os.getenv("ATOM_ENGRAM_REG_LOCK", "") - or "/dev/shm/engram-uva-register.lock", - # mlock each faulted-in range before registering it, so a cold reclaim - # cannot push registration back to disk. Needs RLIMIT_MEMLOCK at least the - # table size (e.g. docker --ulimit memlock=-1:-1); a failure only warns. - "ATOM_ENGRAM_MLOCK": lambda: _flag_env("ATOM_ENGRAM_MLOCK", "0"), + # Optional override of the file that serializes shared in-place + # registration. Empty: host_memory names one from the checkpoint path. + "ATOM_ENGRAM_REG_LOCK": lambda: os.getenv("ATOM_ENGRAM_REG_LOCK", ""), # 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. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 6a74730d99..8fbc7e7694 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -318,10 +318,8 @@ 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_REG_CHUNK_MB** | int | 1024 | In-place registration (`ATOM_ENGRAM_UVA_COPY=0`): the largest single `cudaHostRegister`, in MiB. Each uncovered gap is registered in pieces of this size, each its own reference-counted run, and every piece of 1 GiB or more is logged first (`engram: cudaHostRegister 1.00 GiB`). One registration of a whole V4.1-Flash table is ~92 GiB, and four DP-attention ranks doing that at once hang in the KFD ioctl. `0` registers each gap in one call. Unset, empty or invalid reads as the default. | -| **ATOM_ENGRAM_REG_LOCK** | path | `/dev/shm/engram-uva-register.lock` | Lock file that makes in-place registration and unregistration of Engram ranges run one process at a time on the node, since every DP-attention rank page-locks the same checkpoint pages. A rank that has to wait logs it and how long it waited. A lock file that cannot be opened or locked is logged and skipped, and the rank registers without it. The resolved path is logged once per process. The default lives in `/dev/shm`, which is private to a container unless shared (`docker run --ipc=host` uses the host's), so ranks split across containers that do not share it must point this at a common path, e.g. a bind-mounted directory. | -| **ATOM_ENGRAM_MLOCK** | bool | 0 (false) | `mlock` each range after it is faulted in and before it is registered in place, so the kernel cannot reclaim the pages in between and send registration back to disk (cold cache: ~8 min per V4.1-Flash rank instead of ~30 s). Pages are unpinned with the last reference to their registration, or right away if the registration is refused. Needs `RLIMIT_MEMLOCK` of at least the table size, e.g. `docker run --ulimit memlock=-1:-1`; a failure is logged and registration proceeds. | +| **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, outside the registration lock, and registered read-only in 1 GiB pieces: one `cudaHostRegister` covering a whole table (~92 GiB), issued by several DP ranks at once, does not return. The faulted-in range is `mlock`ed when `RLIMIT_MEMLOCK` is unlimited or at least the range (`docker run --ulimit memlock=-1:-1` — a smaller limit skips the call, and a failure only warns), so a cold page cache cannot be reclaimed before registration; pages are unpinned with the last reference. 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_REG_LOCK** | path | per checkpoint, under `/dev/shm` | Lock file for in-place registration of a shared Engram table (DP attention: every rank page-locks the same checkpoint). Unset, the path is `/dev/shm/engram-uva-.lock`, hashed from that checkpoint, so two servers loading different checkpoints do not queue on each other; set it to force one path. Taken only around `cudaHostRegister` / `cudaHostUnregister`, and only for that shared registration — a TP rank registering its own shard does not take it. A wait is logged; after 600 s the rank gives up and the registration falls back to copying. A lock file that cannot be opened or locked is logged and skipped, and the rank registers without it. `/dev/shm` is private to a container unless shared (`docker run --ipc=host`), so ranks that do not share it must set this to a path they all see. | | **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) diff --git a/tests/model_ops/engram/test_host_path.py b/tests/model_ops/engram/test_host_path.py index b7647d10a9..eddedd8b6f 100644 --- a/tests/model_ops/engram/test_host_path.py +++ b/tests/model_ops/engram/test_host_path.py @@ -5,7 +5,14 @@ which only happens if the prime search runs in exactly the reference order. """ +import contextlib +import os +import subprocess +import sys +import threading +import time from concurrent.futures import ThreadPoolExecutor +from itertools import pairwise import numpy as np import pytest @@ -727,39 +734,21 @@ def test_registration_is_split_into_chunks(monkeypatch): def test_refused_chunk_gives_back_the_chunks_before_it(monkeypatch): """A refusal part way through a gap releases the pieces already taken. - The release runs while the registration lock is held, so this also checks - that the lock is reentrant within the process. + The rollback runs while this call holds the registration lock, and must + not take it again (that flock is not reentrant). """ from atom.model_ops.engram import host_memory cudart = FakeCudart(refuse_call=3) monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: host_memory._PAGE) + monkeypatch.setattr(host_memory, "_REG_LOCK_TIMEOUT_S", 2) table = wide_table() - assert table.enable_uva(0, 4) is False + assert table.enable_uva(0, 4, shared=True) is False assert cudart.registered == [] assert table._uva is None -def test_chunk_size_zero_registers_each_gap_whole(monkeypatch): - from atom.model_ops.engram import host_memory - - cudart = FakeCudart() - monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) - monkeypatch.setenv("ATOM_ENGRAM_REG_CHUNK_MB", "0") - assert host_memory._reg_chunk_bytes() == 0 - table = wide_table() - assert table.enable_uva(0, 4) is True - assert cudart.calls == 1 - - -def test_empty_chunk_size_reads_as_the_default(monkeypatch): - from atom.model_ops.engram import host_memory - - monkeypatch.setenv("ATOM_ENGRAM_REG_CHUNK_MB", "") - assert host_memory._reg_chunk_bytes() == 1024 << 20 - - def test_unopenable_lock_file_still_registers_in_place(monkeypatch, tmp_path): """A lock file that cannot be opened is skipped, not treated as a refused registration (which would send the rank to the copy path).""" @@ -769,13 +758,26 @@ def test_unopenable_lock_file_still_registers_in_place(monkeypatch, tmp_path): blocker.write_text("") monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", str(blocker / "register.lock")) table = wide_table() - assert table.enable_uva(0, 4) is True + assert table.enable_uva(0, 4, shared=True) is True assert len(cudart.registered) >= 1 -def test_registration_and_unregistration_take_the_lock(monkeypatch): - import contextlib +def test_registration_exception_is_not_chained_to_a_skipped_lock(tmp_path): + from atom.model_ops.engram import host_memory + + blocker = tmp_path / "not-a-dir" + blocker.write_text("x") + path = str(blocker / "register.lock") + with ( + pytest.raises(RuntimeError) as caught, + host_memory._register_lock_file(path), + ): + raise RuntimeError("during registration") + assert caught.value.__cause__ is None + assert caught.value.__context__ is None + +def test_registration_and_unregistration_take_the_lock(monkeypatch): from atom.model_ops.engram import host_memory cudart = FakeCudart() @@ -783,26 +785,175 @@ def test_registration_and_unregistration_take_the_lock(monkeypatch): taken = [] @contextlib.contextmanager - def counting_lock_file(): - taken.append(len(cudart.registered)) + def counting_lock_file(path): + taken.append((path, len(cudart.registered))) yield monkeypatch.setattr(host_memory, "_register_lock_file", counting_lock_file) table = wide_table() - assert table.enable_uva(0, 4) is True + assert table.enable_uva(0, 4, shared=True) is True table.disable_uva() # Once around registering (nothing held yet), once around releasing. - assert taken[0] == 0 + assert taken[0][1] == 0 assert len(taken) >= 2 assert cudart.registered == [] +def test_unshared_registration_skips_the_lock(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + taken = [] + + @contextlib.contextmanager + def counting_lock_file(path): + taken.append(path) + yield + + monkeypatch.setattr(host_memory, "_register_lock_file", counting_lock_file) + table = wide_table() + assert table.enable_uva(0, 4) is True + table.disable_uva() + assert taken == [] + assert cudart.registered == [] + + +def test_lock_wait_is_bounded(tmp_path, monkeypatch): + from atom.model_ops.engram import host_memory + + monkeypatch.setattr(host_memory, "_REG_LOCK_TIMEOUT_S", 0.3) + path = str(tmp_path / "register.lock") + held, release = threading.Event(), threading.Event() + + def holder(): + with host_memory._register_lock_file(path): + held.set() + release.wait(10) + + thread = threading.Thread(target=holder) + thread.start() + assert held.wait(5) + try: + started = time.monotonic() + with pytest.raises(TimeoutError), host_memory._register_lock_file(path): + pass + assert time.monotonic() - started < 2 + finally: + release.set() + thread.join(5) + assert not thread.is_alive() + + +def test_lock_path_uses_the_override_or_the_mapped_file(monkeypatch, tmp_path): + import mmap + + from atom.model_ops.engram import host_memory + + monkeypatch.delenv("ATOM_ENGRAM_REG_LOCK", raising=False) + anon = torch.ones(4) + anonymous = host_memory._registration_lock_path([anon]) + assert anonymous.startswith("/dev/shm/engram-uva-") + assert host_memory._registration_lock_path([torch.ones(4)]) == anonymous + monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", str(tmp_path / "forced.lock")) + assert host_memory._registration_lock_path([anon]) == str(tmp_path / "forced.lock") + + monkeypatch.delenv("ATOM_ENGRAM_REG_LOCK", raising=False) + with contextlib.ExitStack() as stack: + paths = [] + for name in ("a.bin", "b.bin"): + raw = tmp_path / name + raw.write_bytes(b"\0" * 4096) + handle = stack.enter_context(open(raw, "rb")) + mapped = stack.enter_context( + mmap.mmap(handle.fileno(), 4096, prot=mmap.PROT_READ) + ) + paths.append( + host_memory._registration_lock_path( + [torch.frombuffer(mapped, dtype=torch.uint8)] + ) + ) + assert paths[0] != paths[1] + assert paths[0] != anonymous + + +_LOCK_WORKER = r""" +import os, sys, time, traceback +os.environ["ATOM_ENGRAM_REG_LOCK"] = os.environ["LOCK_PATH"] +os.environ["ATOM_ENGRAM_UVA_COPY"] = "0" +try: + import torch + from atom.model_ops.engram import host_memory + from atom.model_ops.engram.tables import HostEmbeddingTable + host_memory._memlock_fits = lambda _n: False + hold = float(os.environ.get("HOLD_S", "0.4")) + + class _Rt: + def cudaHostRegister(self, ptr, size, flags): + start = time.monotonic() + time.sleep(hold) + with open(os.environ["RECORD"], "a") as fh: + fh.write(f"{start:.6f} {time.monotonic():.6f}\n") + return 0 + + def cudaHostUnregister(self, ptr): + return 0 + + torch.cuda.cudart = lambda: _Rt() + values = torch.ones(4, 4096) + table = HostEmbeddingTable(values, num_rows=4, head_dim=4096) + if not table.enable_uva(0, 4, shared=True): + sys.exit("enable_uva failed") + table.disable_uva() +except Exception: + traceback.print_exc() + sys.exit(1) +""" + + +def test_two_processes_do_not_register_at_once(tmp_path): + env = os.environ.copy() + env["LOCK_PATH"] = str(tmp_path / "register.lock") + env["HOLD_S"] = "0.4" + env["HIP_VISIBLE_DEVICES"] = "" + procs = [] + records = [] + for i in range(2): + record = tmp_path / f"rec{i}" + records.append(record) + child = env.copy() + child["RECORD"] = str(record) + procs.append( + subprocess.Popen( + [sys.executable, "-c", _LOCK_WORKER], + env=child, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + ) + for proc in procs: + _out, err = proc.communicate(timeout=60) + assert proc.returncode == 0, err + spans = [] + for record in records: + text = record.read_text().strip().splitlines() + assert text, record + for line in text: + start, end = (float(part) for part in line.split()) + spans.append((start, end)) + spans.sort() + assert len(spans) >= 2 + for (_start, end), (nxt, _end) in pairwise(spans): + assert end <= nxt + + def pin_recorder(monkeypatch): """Record `_mlock` / `_munlock` ranges instead of calling libc.""" from atom.model_ops.engram import host_memory locked, unlocked = [], [] - monkeypatch.setenv("ATOM_ENGRAM_MLOCK", "1") + monkeypatch.setattr(host_memory, "_memlock_fits", lambda _n: True) monkeypatch.setattr( host_memory, "_mlock", lambda lo, hi: locked.append((lo, hi)) or True ) @@ -845,6 +996,53 @@ def test_refused_mlocked_registration_unpins_its_pages(monkeypatch): assert covered_bytes(unlocked) == covered_bytes(locked) +def _aligned_buffer(pages: int): + from atom.model_ops.engram import host_memory + + page = host_memory._PAGE + buf = torch.empty((pages + 1) * page, dtype=torch.uint8) + ptr = buf.data_ptr() + lo = ptr + (-ptr % page) + assert lo + pages * page <= ptr + buf.numel() + return buf, lo, page + + +def test_refused_overlap_unpins_pages_an_older_run_does_not_own(monkeypatch): + """A refused registration mlocked the whole range, including pages of an + older run whose own mlock failed. That run will not unpin them.""" + from atom.model_ops.engram import host_memory + + buf, lo, page = _aligned_buffer(4) + rt = FakeCudart() + runs = host_memory._PAGE_RUNS + assert runs.acquire(rt, lo, lo + page, mlocked=False) == [lo] + assert runs._runs[lo].mlocked is False + _, unlocked = pin_recorder(monkeypatch) + rt.refuse_call = rt.calls + 1 + assert runs.acquire(rt, lo, lo + 3 * page, mlocked=True) is None + assert unlocked == [(lo, lo + 3 * page)] + assert runs._runs[lo].mlocked is False + assert rt.registered == [lo] + del buf + + +def test_refused_overlap_keeps_pages_an_older_run_already_pinned(monkeypatch): + from atom.model_ops.engram import host_memory + + buf, lo, page = _aligned_buffer(4) + rt = FakeCudart() + runs = host_memory._PAGE_RUNS + assert runs.acquire(rt, lo, lo + page, mlocked=True) == [lo] + assert runs._runs[lo].mlocked is True + _, unlocked = pin_recorder(monkeypatch) + rt.refuse_call = rt.calls + 1 + assert runs.acquire(rt, lo, lo + 3 * page, mlocked=True) is None + assert unlocked == [(lo + page, lo + 3 * page)] + assert runs._runs[lo].mlocked is True + assert rt.registered == [lo] + del buf + + class FakeHostAllocator: """Page-locked allocations backed by ctypes buffers kept alive here.""" From a1c146d18af889220548e4b1873a3020f53b94e1 Mon Sep 17 00:00:00 2001 From: "Verma, Karan" Date: Tue, 3 Nov 2026 17:40:32 +0000 Subject: [PATCH 5/7] fix(engram): release without the lock after a timeout, latch timed-out lock paths A release whose lock file times out now unregisters anyway, and a path that timed out once fails at once for later tables instead of waiting again. Co-authored-by: Cursor --- atom/model_ops/engram/host_memory.py | 33 +++++++++++++++++++----- tests/model_ops/engram/test_host_path.py | 23 +++++++++++++++++ 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/atom/model_ops/engram/host_memory.py b/atom/model_ops/engram/host_memory.py index b8b2959bea..b1fd5a429c 100644 --- a/atom/model_ops/engram/host_memory.py +++ b/atom/model_ops/engram/host_memory.py @@ -12,6 +12,9 @@ page runs (`_PageRuns`): ranges are widened to page boundaries, so ranges stored back to back (a layer's scales right before its values) share a page, and HIP refuses to register a range overlapping one already registered. +Gaps are registered in 1 GiB pieces, and a table every DP-attention rank +shares registers and unregisters under a per-checkpoint file lock +(`_register_lock_file`), one process at a time. Every refused or failed HIP call clears HIP's last error (`_clear_hip_error`): HIP also leaves it as the runtime's sticky last error, which torch raises at @@ -85,9 +88,11 @@ # Largest single cudaHostRegister. One ioctl covering a whole Engram table # does not return when several ranks issue it together. _REG_CHUNK_BYTES = 1 << 30 -# How long a rank waits for the registration lock before giving up. +# How long a rank waits for the registration lock before giving up. Once a +# path has timed out it is not waited on again in this process. _REG_LOCK_TIMEOUT_S = 600 _logged_lock_paths: set[str] = set() +_timed_out_lock_paths: set[str] = set() @functools.cache @@ -656,8 +661,11 @@ def _register_lock_file(path: str): A file that cannot be opened or locked is logged and skipped. The `yield` is outside the `except`, so an error during registration is not chained - onto that failure. A wait past `_REG_LOCK_TIMEOUT_S` raises `TimeoutError`. + onto that failure. A wait past `_REG_LOCK_TIMEOUT_S` raises `TimeoutError`, + and so does every later call for that path, without waiting. """ + if path in _timed_out_lock_paths: + raise TimeoutError(f"engram registration lock {path} timed out earlier") fd = None held = False open_error = None @@ -679,6 +687,7 @@ def _register_lock_file(path: str): try: _flock_logging_wait(fd, path) except TimeoutError: + _timed_out_lock_paths.add(path) os.close(fd) fd = None raise @@ -782,11 +791,21 @@ def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None: def release(self, rt, bases, lock_path: str | None = None) -> None: """Drop references taken by `acquire`. Lock order: the registration - lock file `lock_path` (a shared registration), then `self._lock`.""" - with ( - _register_lock_file(lock_path) if lock_path else contextlib.nullcontext(), - self._lock, - ): + lock file `lock_path` (a shared registration), then `self._lock`. + + Best effort: if the lock file times out the runs are released without + it, so they never stay registered and locked for the process's life. + """ + if lock_path: + try: + with _register_lock_file(lock_path), self._lock: + self._release_locked(rt, bases) + return + except TimeoutError: + logger.warning( + "engram: unregistering without %s, which timed out", lock_path + ) + with self._lock: self._release_locked(rt, bases) def _release_locked(self, rt, bases) -> None: diff --git a/tests/model_ops/engram/test_host_path.py b/tests/model_ops/engram/test_host_path.py index 9b631e094a..b70ee9c0b6 100644 --- a/tests/model_ops/engram/test_host_path.py +++ b/tests/model_ops/engram/test_host_path.py @@ -624,6 +624,7 @@ def _no_real_host_allocator(monkeypatch, tmp_path): monkeypatch.delenv("ATOM_ENGRAM_UVA_COPY", raising=False) # A per-test registration lock file, so tests never share /dev/shm. monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", str(tmp_path / "register.lock")) + monkeypatch.setattr(host_memory, "_timed_out_lock_paths", set()) class FakeCudart: @@ -839,12 +840,34 @@ def holder(): with pytest.raises(TimeoutError), host_memory._register_lock_file(path): pass assert time.monotonic() - started < 2 + # Latched: the next table does not wait out the timeout again. + monkeypatch.setattr(host_memory, "_REG_LOCK_TIMEOUT_S", 30) + started = time.monotonic() + with pytest.raises(TimeoutError), host_memory._register_lock_file(path): + pass + assert time.monotonic() - started < 1 finally: release.set() thread.join(5) assert not thread.is_alive() +def test_release_after_a_lock_timeout_still_unregisters(monkeypatch, tmp_path): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + path = str(tmp_path / "register.lock") + monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", path) + table = wide_table() + assert table.enable_uva(0, 4, shared=True) is True + assert cudart.registered + host_memory._timed_out_lock_paths.add(path) + table.disable_uva() + assert cudart.registered == [] + assert len(host_memory._PAGE_RUNS) == 0 + + def test_lock_path_uses_the_override_or_the_mapped_file(monkeypatch, tmp_path): import mmap From 9ebdec47a7b329d0b33c7055ae2022bbf222dd1b Mon Sep 17 00:00:00 2001 From: "Verma, Karan" Date: Tue, 3 Nov 2026 19:45:38 +0000 Subject: [PATCH 6/7] fix(engram): time out a silent lock holder, not a slow one The holder touches the lock file after every chunk, and a waiter gives up only after 600 s without a touch, so cold-cache ranks behind a slow holder no longer fall back to copying. Document that a timed-out path is never waited on again. Co-authored-by: Cursor --- atom/model_ops/engram/host_memory.py | 43 +++++++++++++++++------- docs/environment_variables.md | 2 +- tests/model_ops/engram/test_host_path.py | 28 +++++++++++++++ 3 files changed, 60 insertions(+), 13 deletions(-) diff --git a/atom/model_ops/engram/host_memory.py b/atom/model_ops/engram/host_memory.py index b1fd5a429c..f2c8bd3313 100644 --- a/atom/model_ops/engram/host_memory.py +++ b/atom/model_ops/engram/host_memory.py @@ -88,11 +88,21 @@ # Largest single cudaHostRegister. One ioctl covering a whole Engram table # does not return when several ranks issue it together. _REG_CHUNK_BYTES = 1 << 30 -# How long a rank waits for the registration lock before giving up. Once a -# path has timed out it is not waited on again in this process. +# How long a waiter tolerates a registration lock whose holder shows no +# progress (the holder touches the file after every chunk) before giving up. +# Once a path has timed out it is never waited on again in this process. _REG_LOCK_TIMEOUT_S = 600 _logged_lock_paths: set[str] = set() _timed_out_lock_paths: set[str] = set() +# Descriptors of the registration lock files this process holds. +_held_lock_fds: list[int] = [] + + +def _lock_heartbeat() -> None: + """Touch the held registration lock file, so waiters see progress.""" + for fd in _held_lock_fds: + with contextlib.suppress(OSError): + os.utime(fd) @functools.cache @@ -623,8 +633,10 @@ def _registration_lock_path(sources) -> str: def _flock_logging_wait(fd: int, path: str) -> None: """Exclusive `flock` on `fd`. - A wait is logged. Past `_REG_LOCK_TIMEOUT_S` this raises `TimeoutError` - (an `OSError` subclass; callers must not swallow it with `except OSError`). + A wait is logged. It is bounded by the holder's progress, not by its + total hold time: the holder touches the file after every chunk, and only + `_REG_LOCK_TIMEOUT_S` without a touch raises `TimeoutError` (an `OSError` + subclass; callers must not swallow it with `except OSError`). """ try: fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) @@ -633,21 +645,22 @@ def _flock_logging_wait(fd: int, path: str) -> None: pass logger.info("engram: waiting for %s; another process is page-locking", path) start = time.monotonic() - deadline = start + _REG_LOCK_TIMEOUT_S + wait_from = time.time() while True: try: fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) except BlockingIOError: - if time.monotonic() >= deadline: + last = max(wait_from, os.fstat(fd).st_mtime) + if time.time() - last >= _REG_LOCK_TIMEOUT_S: logger.error( - "engram: gave up on %s after %d s; the holder is still " - "page-locking", + "engram: gave up on %s: its holder made no progress for " + "%d s; not waiting on it again in this process", path, _REG_LOCK_TIMEOUT_S, ) raise TimeoutError( - f"engram registration lock {path} still held after " - f"{_REG_LOCK_TIMEOUT_S}s" + f"engram registration lock {path}: holder made no progress " + f"for {_REG_LOCK_TIMEOUT_S}s" ) time.sleep(0.05) else: @@ -661,8 +674,9 @@ def _register_lock_file(path: str): A file that cannot be opened or locked is logged and skipped. The `yield` is outside the `except`, so an error during registration is not chained - onto that failure. A wait past `_REG_LOCK_TIMEOUT_S` raises `TimeoutError`, - and so does every later call for that path, without waiting. + onto that failure. A holder that makes no progress for + `_REG_LOCK_TIMEOUT_S` raises `TimeoutError`, and after that the path is + never waited on again in this process: every later call raises at once. """ if path in _timed_out_lock_paths: raise TimeoutError(f"engram registration lock {path} timed out earlier") @@ -697,10 +711,13 @@ def _register_lock_file(path: str): ) else: held = True + _held_lock_fds.append(fd) + _lock_heartbeat() try: yield finally: if held: + _held_lock_fds.remove(fd) fcntl.flock(fd, fcntl.LOCK_UN) if fd is not None: os.close(fd) @@ -775,6 +792,7 @@ def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None: held.append(c) created += 1 created_bytes += e - c + _lock_heartbeat() c = e if base is None: break @@ -819,6 +837,7 @@ def _release_locked(self, rt, bases) -> None: if int(rt.cudaHostUnregister(base) or 0): _clear_hip_error() logger.warning("engram: cudaHostUnregister failed") + _lock_heartbeat() def __len__(self) -> int: return len(self._runs) diff --git a/docs/environment_variables.md b/docs/environment_variables.md index c08e9d0520..60fdb824f3 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -319,7 +319,7 @@ some layers and the device path for others, which is the confusing state. |----------|------|---------|-------------| | **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 in 1 GiB pieces (one `cudaHostRegister` covering a whole ~92 GiB table, issued by several DP ranks at once, does not return), 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//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_REG_LOCK** | path | per checkpoint, under `/dev/shm` | Lock file for in-place registration of a shared Engram table (DP attention: every rank page-locks the same checkpoint). Unset, the path is `/dev/shm/engram-uva-.lock`, hashed from that checkpoint, so two servers loading different checkpoints do not queue on each other; set it to force one path. Taken only around `cudaHostRegister` / `cudaHostUnregister` (fault-in runs before it), and only for that shared registration — a TP rank registering its own shard does not take it. A wait is logged; after 600 s the rank gives up and the registration falls back to copying. A lock file that cannot be opened or locked is logged and skipped, and the rank registers without it. `/dev/shm` is private to a container unless shared (`docker run --ipc=host`), so ranks that do not share it must set this to a path they all see. | +| **ATOM_ENGRAM_REG_LOCK** | path | per checkpoint, under `/dev/shm` | Lock file for in-place registration of a shared Engram table (DP attention: every rank page-locks the same checkpoint). Unset, the path is `/dev/shm/engram-uva-.lock`, hashed from that checkpoint, so two servers loading different checkpoints do not queue on each other; set it to force one path. Taken only around `cudaHostRegister` / `cudaHostUnregister` (fault-in runs before it), and only for that shared registration — a TP rank registering its own shard does not take it. A wait is logged. The holder touches the file after every 1 GiB chunk, so a slow but progressing holder (a cold page cache, minutes per rank) is waited on for as long as it takes; only a holder with no progress for 600 s makes the waiter give up and fall back to copying. After that the path is never waited on again in that process: its later tables copy at once, and its releases unregister without the lock. A lock file that cannot be opened or locked is logged and skipped, and the rank registers without it. `/dev/shm` is private to a container unless shared (`docker run --ipc=host`), so ranks that do not share it must set this to a path they all see. | | **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) diff --git a/tests/model_ops/engram/test_host_path.py b/tests/model_ops/engram/test_host_path.py index b70ee9c0b6..522c120618 100644 --- a/tests/model_ops/engram/test_host_path.py +++ b/tests/model_ops/engram/test_host_path.py @@ -852,6 +852,34 @@ def holder(): assert not thread.is_alive() +def test_a_progressing_holder_is_waited_on_past_the_timeout(tmp_path, monkeypatch): + """The timeout bounds a holder's silence, not its hold time: a cold-cache + rank holding the lock for minutes, chunk by chunk, must not send the + ranks behind it to the copy path.""" + from atom.model_ops.engram import host_memory + + monkeypatch.setattr(host_memory, "_REG_LOCK_TIMEOUT_S", 0.3) + path = str(tmp_path / "register.lock") + held = threading.Event() + + def holder(): + with host_memory._register_lock_file(path): + held.set() + for _ in range(10): # ~1 s of chunks, three timeouts long + time.sleep(0.1) + host_memory._lock_heartbeat() + + thread = threading.Thread(target=holder) + thread.start() + assert held.wait(5) + started = time.monotonic() + with host_memory._register_lock_file(path): + waited = time.monotonic() - started + thread.join(5) + assert waited > 0.6 + assert path not in host_memory._timed_out_lock_paths + + def test_release_after_a_lock_timeout_still_unregisters(monkeypatch, tmp_path): from atom.model_ops.engram import host_memory From cc331a462ac32071fa7b249ea43908ad0ae96380 Mon Sep 17 00:00:00 2001 From: "Verma, Karan" Date: Fri, 6 Nov 2026 02:23:02 +0000 Subject: [PATCH 7/7] fix(engram): heartbeat during fault-in; guard the held lock fds - _populate touches any held registration lock after each faulted-in piece. acquire faults in under the lock what reclaim took back, and on a cold cache that could outlast _REG_LOCK_TIMEOUT_S with no chunk registered, so waiters saw no progress. - _held_lock_fds and every touch are under _held_lock_fds_lock: fault-in workers now touch the files while the main thread adds and removes them, and a descriptor leaves the list before it is closed, so a closed or reused descriptor is never touched. - Tests: a held lock's mtime moves during _populate, once per piece; a heartbeat waits for the list lock. Co-authored-by: Cursor --- atom/model_ops/engram/host_memory.py | 43 +++++++++++++------- tests/model_ops/engram/test_host_path.py | 51 ++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 14 deletions(-) diff --git a/atom/model_ops/engram/host_memory.py b/atom/model_ops/engram/host_memory.py index f2c8bd3313..a49a87206c 100644 --- a/atom/model_ops/engram/host_memory.py +++ b/atom/model_ops/engram/host_memory.py @@ -89,20 +89,25 @@ # does not return when several ranks issue it together. _REG_CHUNK_BYTES = 1 << 30 # How long a waiter tolerates a registration lock whose holder shows no -# progress (the holder touches the file after every chunk) before giving up. -# Once a path has timed out it is never waited on again in this process. +# progress (the holder touches the file after every faulted-in piece and every +# registered chunk) before giving up. Once a path has timed out it is never +# waited on again in this process. _REG_LOCK_TIMEOUT_S = 600 _logged_lock_paths: set[str] = set() _timed_out_lock_paths: set[str] = set() -# Descriptors of the registration lock files this process holds. +# Descriptors of the registration lock files this process holds. Fault-in +# worker threads touch them, so the list and each touch are under the lock: a +# descriptor leaves the list before it is closed and is never touched after. _held_lock_fds: list[int] = [] +_held_lock_fds_lock = threading.Lock() def _lock_heartbeat() -> None: - """Touch the held registration lock file, so waiters see progress.""" - for fd in _held_lock_fds: - with contextlib.suppress(OSError): - os.utime(fd) + """Touch the held registration lock files, so waiters see progress.""" + with _held_lock_fds_lock: + for fd in _held_lock_fds: + with contextlib.suppress(OSError): + os.utime(fd) @functools.cache @@ -634,9 +639,10 @@ def _flock_logging_wait(fd: int, path: str) -> None: """Exclusive `flock` on `fd`. A wait is logged. It is bounded by the holder's progress, not by its - total hold time: the holder touches the file after every chunk, and only - `_REG_LOCK_TIMEOUT_S` without a touch raises `TimeoutError` (an `OSError` - subclass; callers must not swallow it with `except OSError`). + total hold time: the holder touches the file after every faulted-in piece + and every registered chunk, and only `_REG_LOCK_TIMEOUT_S` without a + touch raises `TimeoutError` (an `OSError` subclass; callers must not + swallow it with `except OSError`). """ try: fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) @@ -711,13 +717,15 @@ def _register_lock_file(path: str): ) else: held = True - _held_lock_fds.append(fd) + with _held_lock_fds_lock: + _held_lock_fds.append(fd) _lock_heartbeat() try: yield finally: if held: - _held_lock_fds.remove(fd) + with _held_lock_fds_lock: + _held_lock_fds.remove(fd) fcntl.flock(fd, fcntl.LOCK_UN) if fd is not None: os.close(fd) @@ -871,13 +879,20 @@ def _fault_and_lock(lo: int, hi: int, populate: bool) -> None: def _populate(lo: int, hi: int) -> None: - """`MADV_POPULATE_READ` over `[lo, hi)`, one chunk per thread.""" + """`MADV_POPULATE_READ` over `[lo, hi)`, one chunk per thread. + + Each finished chunk touches any registration lock this process holds: + `acquire` faults in under the lock what reclaim took back, and on a cold + cache that can outlast `_REG_LOCK_TIMEOUT_S` with no chunk registered. + """ libc = _libc() pieces = _chunks(lo, hi - lo, _CHUNK_BYTES) def fault(piece) -> bool: o, n = piece - return libc.madvise(lo + o, n, _MADV_POPULATE_READ) == 0 + ok = libc.madvise(lo + o, n, _MADV_POPULATE_READ) == 0 + _lock_heartbeat() + return ok threads = min(copy_threads(hi - lo), len(pieces)) if threads <= 1: diff --git a/tests/model_ops/engram/test_host_path.py b/tests/model_ops/engram/test_host_path.py index 522c120618..1075707f09 100644 --- a/tests/model_ops/engram/test_host_path.py +++ b/tests/model_ops/engram/test_host_path.py @@ -880,6 +880,57 @@ def holder(): assert path not in host_memory._timed_out_lock_paths +def test_fault_in_touches_a_held_lock(tmp_path, monkeypatch): + """`acquire` faults in under the lock what reclaim took back; on a cold + cache that can outlast the timeout before a chunk is registered, so every + faulted-in piece touches the held lock file.""" + from atom.model_ops.engram import host_memory + + advised = [] + + class Libc: + def madvise(self, addr, nbytes, advice): + advised.append(nbytes) + return 0 + + monkeypatch.setattr(host_memory, "_libc", lambda: Libc()) + beats = [] + heartbeat = host_memory._lock_heartbeat + monkeypatch.setattr( + host_memory, "_lock_heartbeat", lambda: (beats.append(1), heartbeat()) + ) + path = str(tmp_path / "register.lock") + pieces = 4 + lo = 1 << 30 + with host_memory._register_lock_file(path): + os.utime(path, (0, 0)) + beats.clear() + host_memory._populate(lo, lo + pieces * host_memory._CHUNK_BYTES) + assert os.stat(path).st_mtime > 0 + assert len(advised) == pieces + assert len(beats) == pieces + + +def test_heartbeat_touches_under_the_fd_list_lock(): + """Fault-in workers touch the held lock files while the main thread adds + and removes them: a touch waits for the list's lock, so a descriptor that + is being removed and closed is never touched.""" + from atom.model_ops.engram import host_memory + + done = threading.Event() + + def beat(): + host_memory._lock_heartbeat() + done.set() + + with host_memory._held_lock_fds_lock: + thread = threading.Thread(target=beat) + thread.start() + assert not done.wait(0.2) + assert done.wait(5) + thread.join(5) + + def test_release_after_a_lock_timeout_still_unregisters(monkeypatch, tmp_path): from atom.model_ops.engram import host_memory