fix(engram): serialize UVA registration across DP ranks and register in 1 GiB chunks by karverma-amd · Pull Request #2932 · ROCm/ATOM

karverma-amd · 2026-11-03T00:28:43Z

On the Engram UVA path every DP-attention rank page-locks the same read-only checkpoint mapping in place. With four ranks that registration hangs: each rank sits in cudaHostRegister (D state inside the KFD ioctl) and the server never finishes loading. Shared registrations now take a per-checkpoint file lock around that ioctl only, and each ioctl is at most 1 GiB.

atom/model_ops/engram/host_memory.py: the lock is taken only when the table is shared (DP attention). It covers cudaHostRegister and cudaHostUnregister, in that order relative to the in-process run lock. Fault-in and mlock run first, outside it, so ranks read the page cache concurrently. The lock file defaults to /dev/shm/engram-uva-<hash>.lock, hashed from the checkpoint path, so two servers loading different checkpoints do not queue on each other. ATOM_ENGRAM_REG_LOCK overrides it. A wait is logged. The holder touches the lock file after every 1 GiB chunk, and a waiter gives up only after 600 s with no touch, so a cold-cache holder that takes minutes is waited on; a holder stuck in D state is not. A path that timed out 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. A refused registration releases what it took without taking the lock again.
Each uncovered gap is registered in 1 GiB pieces (_REG_CHUNK_BYTES). One info line reports the bytes and the number of runs actually created; per-chunk lines are debug.
mlock is #2898's, merged in: each run is locked when it is created and unlocked with its last reference. A chunk is its own run, so it is unlocked on its own. A refused chunk unlocks the rest of its gap, and the chunks before it are released through _release_locked. For a shared table the multi-threaded fault-in runs once before the file lock, and the mlock + fault-in inside acquire only brings back what reclaim took meanwhile.
atom/utils/envs.py, docs/environment_variables.md: ATOM_ENGRAM_REG_LOCK is the only new knob. Chunk size is a constant.
tests/model_ops/engram/test_host_path.py: chunk split, refusal part way through a gap (including under the lock), unopenable lock, an error during registration not chained onto a skipped lock, lock around shared register and unregister, no lock for an unshared shard, bounded wait, lock path from the mapped file, two processes whose registration intervals do not overlap, and chunked runs unlocked on release and on refusal. fix(engram): mlock in-place UVA page runs while registered (on by default, best effort) #2898's test_mlock.py passes unchanged.

Motivation

DeepSeek-V4.1-Flash TP+EP4 with DP attention on MI455X (gfx1250), 20261102 image. Each rank registers ~91.6 GiB of Engram table. Before this change, concurrency 128 and 256 servers hung in Engram setup: the log stopped at engram: loaded compressed-vocab table, and py-spy showed all four ranks inside cudaHostRegister at once. docker rm -f could not stop the stuck processes for about a minute.

With a cold page cache, registration also faults the table in from disk at ~180–270 MB/s, about 8 minutes per rank, because pages faulted in ahead of cudaHostRegister can be reclaimed before that call.

These are different failures. The hang is concurrent cudaHostRegister on the shared mapping. The multi-minute startup is the cold cache, and mlock is what removes it. This PR does not have a run of each mechanism alone; see the test result for what was actually compared.

Technical Details

Lock order is the registration file, then _PageRuns._lock. Acquire's rollback calls _release_locked and does not lock again.
A run is a small dataclass (end, refs) rather than a list.
No change when Engram UVA is off or in copy mode. A TP rank registering its own shard does not take the lock.

Test Plan

Serve DeepSeek-V4.1-Flash TP+EP4 + DP attention on 4x MI455X, concurrency 128 / 256 / 1024.
Unlimited RLIMIT_MEMLOCK (--ulimit memlock=-1:-1), and once with the default limit to check that mlock is skipped or warns.
tests/model_ops/engram/test_host_path.py and tests/test_envs.py.

Test Result

Without the change: all four ranks hang in cudaHostRegister (concurrency 128 and 256). py-spy showed them in that ioctl together.
With the lock and 1 GiB chunks, before mlock: every rank logs engram: registered 91.55 GiB in-place as 92 userptr BO(s) and UVA device lookup active -- heads [0, 24) of 24, 183.1 GiB page-locked. Concurrency 128, 256 and 1024 ran to completion with 0 failed samples. Concurrency 1024: 47,758 tok/s (11,939 tok/s/GPU), inline accuracy 61.29%. On a cold cache this still took about 8 minutes per rank.
With mlock (unlimited memlock; then this branch's own pin, now fix(engram): mlock in-place UVA page runs while registered (on by default, best effort) #2898's): registration dropped to ~25–30 s per rank. mlock was not run without the lock, and the lock was not run without chunking.
test_host_path.py + test_mlock.py + tests/test_envs.py: 207 passed, 8 skipped, 8 failed; main fails the same 8 in this CPU-only environment (ModuleNotFoundError: No module named 'jax' while importing the device op). Black and ruff (0.17.3, what CI installs) clean.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

…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 <cursoragent@cursor.com>

github-actions · 2026-11-03T00:29:15Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2932 --add-label <label>

valar-cop

some changes needed..

ci green.

valar-cop · 2026-11-03T00:49:34Z

+# 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)


use envs.py for these three, and add docs rows.. this file already uses envs.ATOM_ENGRAM_UVA_COPY

Done in a9c69ae. ATOM_ENGRAM_REG_CHUNK_MB, ATOM_ENGRAM_REG_LOCK and ATOM_ENGRAM_MLOCK are now in envs.py next to ATOM_ENGRAM_UVA_COPY, read through envs.*, with rows in docs/environment_variables.md.

valar-cop · 2026-11-03T00:49:34Z

+# 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)


hand-rolled parsing here.. VAR= becomes 0 and turns chunking off, use the envs int helper

Fixed in a9c69ae: the chunk size uses _nonnegative_int_env, so unset or empty reads as 1024 and an invalid or negative value warns and falls back to it. Only an explicit 0 registers each gap whole. test_empty_chunk_size_reads_as_the_default and test_chunk_size_zero_registers_each_gap_whole cover both.

valar-cop · 2026-11-03T00:49:34Z

+    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)


lock open error falls back to 183GiB copy per rank? that's not a registration refusal..

Right, any exception in _register sends the rank to the copy plan. Fixed in a9c69ae: a lock file that cannot be opened or locked is logged as a warning and the rank registers without the lock (the behavior before this PR), instead of raising. test_unopenable_lock_file_still_registers_in_place points the lock under a regular file and checks the range still registers in place.

valar-cop · 2026-11-03T00:49:34Z

+    """
+    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)


flock blocks with no log.. a recurrence will be silent again, log the wait?

Done in a9c69ae: the lock is tried with LOCK_NB first; a rank that has to wait logs engram: waiting for <path>; another process is page-locking and then took <path> after N s.

valar-cop · 2026-11-03T00:49:34Z

-                    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


chunk loop never tested, all fake ranges are below 1 GiB.. set _REG_CHUNK to a page or two in test?

Added in a9c69ae: test_registration_is_split_into_chunks sets the chunk to two pages on a 64 KiB range and checks the call count and that disable_uva releases every piece; test_refused_chunk_gives_back_the_chunks_before_it refuses the third one-page chunk and checks nothing stays registered (it also exercises the lock being reentrant, since the rollback unregisters while holding it).

valar-cop · 2026-11-03T00:49:34Z

+    calls hang.
+    """
    rt = torch.cuda.cudart()
+    with _engram_register_lock():


unregister on shared pages from all ranks concurrently, same hazard? any evidence it's safe..

No evidence it is safe: in our runs the servers were stopped by killing the containers, so a graceful concurrent unregister was never exercised. a9c69ae puts cudaHostUnregister under the same lock, taken before _PageRuns._lock in both release and _register so the order never inverts, and reentrant within the process for the rollback path. test_registration_and_unregistration_take_the_lock checks both paths take it.

…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 <cursoragent@cursor.com>

valar-cop

a few small things..

cold-cache 4-rank startup total with the node-wide lock?

ci Offline inference smoke test failed — not sure yet (low): The job log under .valar/ci is empty (0 bytes), so nothing can be attributed. The same job is green on all 8 recent main runs and on 2 of 4 other recent PR runs, so flakiness cannot be excluded either.

next: Re-run the job and attach its log/traceback before attributing it to this PR. https://github.com/ROCm/ATOM/actions/runs/41557566952/job/125664860298

valar-cop · 2026-11-03T01:35:04Z

+    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:


nothing ever munlocks, pages stay pinned after release/disable_uva.. munlock in _PageRuns.release when refcount hits 0?

Fixed in 529fc7d. The mlock moved into _PageRuns.acquire; each run records whether it is mlocked (set only once the whole range holds its runs) and is munlocked right after its cudaHostUnregister when the refcount reaches 0. Runs partition the registered pages, so no live range still needs those pages. A refused registration unpins the pages of its range that no live run covers. Tests: test_mlocked_runs_are_munlocked_with_their_last_reference, test_refused_mlocked_registration_unpins_its_pages.

valar-cop · 2026-11-03T01:35:04Z


    def release(self, rt, bases) -> None:
-        with self._lock:
+        # The node-wide lock first, as in `_register`, so the two locks are


comment is wrong on rollback path, release is called under _lock there.. state the real invariant or return held from acquire?

Agreed, the comment was misleading. 529fc7d states the invariant: every path takes the node-wide lock before self._lock; _register holds it around acquire, so the rollback reaches release with both already held (both reentrant), and teardown takes it in release.

valar-cop · 2026-11-03T01:35:04Z

+    ),
+    # 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", "")


default /dev/shm lock is per-container, ranks in different containers lock different files.. say it in docs and log the resolved path?

Done in 529fc7d: the resolved path is logged once per process (engram: page-locking serialized on <path>), and the docs row and envs comment say /dev/shm is per container unless shared (docker run --ipc=host uses the host's), so ranks split across containers need a common path such as a bind-mounted directory. In our runs all four DP ranks were in one container.

…og 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 <cursoragent@cursor.com>

karverma-amd · 2026-11-03T02:07:59Z

Addressed the three inline points in 529fc7d (tests: test_host_path.py 91 passed with the same 8 environment failures as main; tests/test_envs.py 106 passed; black and ruff clean).

Cold-cache 4-rank startup with the node-wide lock: ranks register one after another, so the total is the sum of the per-rank times.

With ATOM_ENGRAM_MLOCK=1: first registration to the last rank's UVA device lookup active took 6 min 14 s and 8 min 9 s in two concurrency-1024 runs (all 4 ranks, ~91.6 GiB each). The page cache was not dropped before those runs, so they are not strictly cold.
Without mlock on a cold page cache: one run had registered only 1 of 4 ranks after 37 min before it was stopped. Registration faulted the table in from disk at ~180–270 MB/s because pages _populate faulted in were reclaimed before cudaHostRegister. That is why ATOM_ENGRAM_MLOCK exists; I would set it whenever the node has the RAM.

Smoke test: the failing Offline inference smoke test job stopped at "Download model" with a 401 (Illegal header value b'Bearer ', an empty HF token on that runner), before any ATOM code ran. The other smoke job on this PR (job 125659945204) completed "Run offline inference" successfully.

valarLip

Thanks for chasing this hang. I reviewed HEAD 529fc7da3. test_host_path.py passes locally except test_a_refused_registration_leaves_the_host_path_working, which fails identically at the merge base (an environment issue here, not from this PR); the 8 new tests pass. Findings were checked against the code; the reentrancy self-deadlock (:634) holds by mechanism but has a narrow window and isn't reproduced.

Main point: the PR adds three mechanisms at once (node flock, 1 GiB chunking, optional mlock) and compares only all-on vs all-off, so it's unclear which fixes the hang — and the most likely root-cause fix, mlock, is off by default. Please measure each alone and remove what doesn't help (host_memory.py:812).

Then: the lock is too broad (non-shared TP shards, plus populate/mlock inside it), the reentrant locking has a self-deadlock window, mlock bookkeeping can leak locked pages, the lock file is global with no timeout, and the exception chaining in _register_lock_file misleads. Please also remove the two unnecessary env vars (envs.py:1265) and the five-fold repeated rationale (:552). The rest are cleanups (run record type, libc handle, log spam), a missing cross-process test, and docs that don't mention the lock.

valarLip · 2026-11-03T03:56:59Z

+    checkpoint pages on every rank, and overlapping registrations hang.
+    """
    rt = torch.cuda.cudart()
+    with _engram_register_lock():


Three mechanisms at once, with no per-mechanism comparison. This PR adds a node-wide flock, 1 GiB chunked registration and optional mlock together, and the test result only compares "all on" vs "all off" — so it's unknown which one actually fixes the hang. My guess (unverified): the 92 GiB of file-backed pages get reclaimed between _populate and cudaHostRegister, so the userptr registration keeps invalidating and retrying. The PR itself notes "pages faulted by _populate can be reclaimed before cudaHostRegister", and turning on mlock took the cold-cache case from ~8 min to ~30 s — which both point at mlock, yet mlock is off by default. If chunking or mlock alone is enough, the lock just serializes startup for nothing. Could you measure lock-only, chunking-only and mlock-only separately, and remove whatever has no effect?

Separately, the lock is taken for every in-place registration, not just the DP-attention case where ranks share one table. With TP>1 and ATOM_ENGRAM_UVA_COPY=0 (or no HIP runtime), each TP rank registers a disjoint shard — the conflict the PR describes can't happen — but they're now serialized, and the vote all_reduce at host.py:432 waits for the slowest rank. pin_rows already has a shared parameter; please take the lock only when shared=True.

The hang and the cold-cache stall are different failures, and this revision keeps both fixes for that reason.

The hang, before any of this: py-spy showed all four DP ranks inside cudaHostRegister at once, D state in the KFD ioctl, on one registration of the whole ~92 GiB table. That is the shared mapping being registered concurrently. The lock is what keeps those ioctls from overlapping. I have not run lock-only, chunk-only, and mlock-only as a factorial, so I can't say the 1 GiB split is still required once the calls are serialized. The call that hung was that single 92 GiB ioctl, and I have not measured an ioctl of that size with the lock held, so the split stays.

The cold cache: with the lock and the 1 GiB split already in place, mlock took the same registration from about 8 min/rank down to about 25–30 s. That is the reclaim between fault-in and cudaHostRegister. mlock is no longer a knob. It runs when RLIMIT_MEMLOCK can hold the range, which is what --ulimit memlock=-1 selects.

The lock is now taken only for shared=True, the DP-attention case where every rank page-locks the same table. A TP rank registering its own shard does not take it, so it no longer lines up behind the others at the all_reduce in host.py.

c7feddd

valarLip · 2026-11-03T03:56:59Z

@@ -649,7 +823,7 @@ def _register(sources, leases):
        flags = _HOST_REGISTER_READ_ONLY if _read_only(lo, hi) else 0
        if flags:


_populate (multithreaded page-in) and _mlock also run inside the node lock, so far more is serialized than the single cudaHostRegister that needs it. On a cold cache one rank's populate takes ~8 min (per the PR), so under DEP4 the later ranks queue behind earlier ranks reading the disk — while the page cache is shared and concurrent populates of the same file would read it once. Please move populate/mlock outside the lock and wrap only register/unregister.

Done. _populate and mlock run before _register_lock_file. The lock is held for cudaHostRegister in _PageRuns.acquire and for cudaHostUnregister in release.

c7feddd

valarLip · 2026-11-03T03:56:59Z

+    shared pages one at a time. Reentrant within the process.
+    """
+    global _reg_lock_depth
+    with _REG_LOCK_THREADS, contextlib.ExitStack() as stack:


The reentrancy here takes three layers (RLock, a global depth counter, ExitStack), and depth is only incremented after the flock is acquired, leaving a self-deadlock window: with fd1 flocked by the outer call but depth still 0, if the same thread reaches release via _Lease.__del__, it sees depth == 0, opens fd2 and flocks it — and flock excludes across separate open file descriptions, so the thread waits on itself forever. (The mechanism holds; the window is narrow and I haven't reproduced it.) Simpler and safe: take the lock only in _register and the public release, and have acquire's rollback call an internal _release_locked that doesn't. Then reentrancy, depth, _REG_LOCK_THREADS, the ExitStack and the nested lock in release (~L719) can all go.

The reentrancy is gone: no RLock around the file lock, no depth counter, no ExitStack. The file lock is taken in _register (shared tables only) and in release. Acquire's rollback calls _release_locked, which requires self._lock and does not touch the file lock.

test_refused_chunk_gives_back_the_chunks_before_it runs that rollback while the file lock is held, with the wait capped at 2 s, so taking the lock again fails the test instead of deadlocking the thread.

c7feddd

valarLip · 2026-11-03T03:56:59Z

                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:


mlock bookkeeping is per run, but mlock (~L670) locks the whole [lo, hi), including pages already covered by existing runs. On a refused registration, _munlock_uncovered (~L731) only unlocks pages not covered by a live run; if that older run's own mlock had failed (RLIMIT, ENOMEM → run[2] = False), the overlapping pages locked this time are never munlocked, since release only munlocks when run[2] is true. This path (overlap + mlock + refusal) has no test.

_munlock_uncovered now unpins every page of the refused range that no already-pinned run covers. An older run whose own mlock failed stays mlocked=False and will not unpin on release, so this call undoes those pages too. munlock is not reference-counted, so a run that is already pinned is left alone: this attempt's mlock was a no-op there.

test_refused_overlap_unpins_pages_an_older_run_does_not_own and test_refused_overlap_keeps_pages_an_older_run_already_pinned cover the two cases.

c7feddd

valarLip · 2026-11-03T03:56:59Z

+    # 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(


Please drop two of the three new env vars. ATOM_ENGRAM_REG_CHUNK_MB should be a module constant (nothing motivates "0 means don't chunk"), and ATOM_ENGRAM_MLOCK can be decided from RLIMIT_MEMLOCK, or simply always attempted since failure only warns. Defaults should be the right behaviour, not something deployments have to tune (and document).

Dropped ATOM_ENGRAM_REG_CHUNK_MB and ATOM_ENGRAM_MLOCK. Chunk size is _REG_CHUNK_BYTES (1 GiB). mlock is attempted when RLIMIT_MEMLOCK is unlimited or at least the range; a smaller limit skips the call, and a failed mlock still only warns. ATOM_ENGRAM_REG_LOCK stays as the path override.

A deployment that set ATOM_ENGRAM_MLOCK=1 can drop it. --ulimit memlock=-1 is what turns pinning on. The variable is now unused and ignored.

c7feddd

valarLip · 2026-11-03T03:56:59Z

+    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)


_mlock / _munlock create a new ctypes.CDLL(None) and reset argtypes on every call — once per 1 GiB chunk released — and the same pattern already exists in _populate and _source_dropper (four copies now). A functools.cached _libc() with argtypes set once would replace them.

_libc() is functools.cached and sets madvise / mlock / munlock argtypes once. _populate, _source_dropper, _mlock, and _munlock all use it.

c7feddd

valarLip · 2026-11-03T03:56:59Z

+                    while c < gap_end:
+                        e = min(gap_end, c + step)
+                        nbytes = e - c
+                        if nbytes >= (1 << 30):


Every 1 GiB chunk logs an INFO line, so a 92 GiB table prints 92 "cudaHostRegister 1.00 GiB" lines per rank — hundreds under DEP8. 1 << 30 appears three times, and the summary's chunk count is recomputed with ceil rather than counted in the loop (and the summary is per gap, while the description says it reports the total). Please drop the per-chunk line to DEBUG (or one progress summary) and count chunks from the runs actually created.

The per-chunk cudaHostRegister line is debug. The info line is emitted once per range and reports the bytes and the count of runs this call created, rather than a ceil over the gap. The 1 GiB size is _REG_CHUNK_BYTES, written once.

c7feddd

valarLip · 2026-11-03T03:56:59Z



+def _reg_chunk_bytes() -> int:
+    """Largest single `cudaHostRegister` (`ATOM_ENGRAM_REG_CHUNK_MB`), in whole


Please trim the repeated rationale. The same motivation ("92 GiB in one registration, four DP ranks stuck in the KFD ioctl") is written five times: this docstring, the envs comment, the docs table, _engram_register_lock and _register. The comments are nearly as long as the code and purely explanatory; the motivation belongs in the docs table and PR description, and the code needs only the one constraint (lock ordering).

The hang write-up now lives in the ATOM_ENGRAM_UVA_COPY / ATOM_ENGRAM_REG_LOCK docs rows and in the PR description. The code comments state the lock order and what _REG_CHUNK_BYTES bounds.

c7feddd

valarLip · 2026-11-03T03:56:59Z

+    assert len(cudart.registered) >= 1
+
+
+def test_registration_and_unregistration_take_the_lock(monkeypatch):


The core change is cross-process serialization, but there's no cross-process test — the lock tests only stub _register_lock_file and count calls. If flock were replaced by a no-op, or the reentrancy broken (a second fd flocking deadlocks itself), only the refusal test would surface it, and by hanging rather than by an assertion failure. Please add a test that spawns two processes contending for the same lock file and asserts their registration intervals don't overlap. Also the local import contextlib here should move to the top of the file.

test_two_processes_do_not_register_at_once spawns two processes against one lock file. Each sleeps inside cudaHostRegister, and the test asserts those intervals do not overlap. import contextlib is at the top of the test module.

c7feddd

valarLip · 2026-11-03T03:56:59Z

@@ -32,12 +32,15 @@
 during a capture is skipped, with a warning, rather than breaking it.
 """


Docs vs code: the module docstring's in-place registration description (L11-14) and the ATOM_ENGRAM_UVA_COPY row in the docs don't mention that in-place registration now takes a node-wide lock and registers in chunks. And the description says "all of one rank's registrations complete while the other ranks wait", but the lock is taken per pin_rows call (per table), so ranks interleave. Readers miss the cross-process blocking point, and startup-time estimates based on the description will be off.

valar-cop · 2026-11-03T04:08:14Z

ci triage for 529fc7d:

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/41575208093

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/41575208122

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/41575208158

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 <cursoragent@cursor.com>

Resolve against ROCm#2898, which mlocks each in-place run while registered. Its per-run lock replaces this branch's pin bookkeeping; the chunked runs, the shared-only registration lock and the pre-lock fault-in sit on top of it. Co-authored-by: Cursor <cursoragent@cursor.com>

valar-cop

some changes needed..

per-mechanism isolation numbers (lock-only, chunk-only, mlock-only)? and any evidence concurrent unregister on the shared mapping is unsafe?

ci green.

valar-cop · 2026-11-03T17:31:32Z

+        """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(),


if release times out the run stays registered and mlocked forever. make release best effort, catch TimeoutError and call _release_locked without the lock?

release is best effort now. If the lock file times out it logs a warning and runs _release_locked under self._lock alone, so the runs are always unregistered and unlocked. Covered by test_release_after_a_lock_timeout_still_unregisters.

a1c146d

valar-cop · 2026-11-03T17:31:32Z

+# 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


timeout is per table, not startup. with two tables that's 20 min before fallback. latch lock-unavailable after first timeout?

Latched. A path that has timed out once goes into _timed_out_lock_paths, and every later _register_lock_file call for it raises TimeoutError at once. The second table falls back to copying immediately, and releases skip the lock. test_lock_wait_is_bounded now checks that the second attempt returns in under a second with the timeout raised to 30 s.

a1c146d

valar-cop · 2026-11-03T17:31:32Z

+                        )
+                        if int(rt.cudaHostRegister(c, e - c, flags)):
+                            _clear_hip_error()
+                            _libc().munlock(c, gap_end - c)


those two tests don't exist at head. re-point the thread at test_refused_registration_leaves_nothing_locked?

Right, those two went with the merge of #2898: its per-run mlock replaced this branch's pin bookkeeping, so the overlap case they tested no longer exists. The refusal paths are covered by test_refused_registration_leaves_nothing_locked (#2898's test_mlock.py, a refused gap with an older run in the middle), and here by test_refused_chunk_leaves_nothing_locked (refusal on the third 1-page chunk) and test_chunked_runs_are_unlocked_with_their_last_reference.

valar-cop · 2026-11-03T17:31:32Z

@@ -33,14 +33,19 @@
 during a capture is skipped, with a warning, rather than breaking it.
 """


docstring still says in-place registration, no lock no chunk. add one sentence..

Added: 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.

a1c146d

…t 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 <cursoragent@cursor.com>

karverma-amd · 2026-11-03T17:41:06Z

Replies to the two questions in the review:

Per-mechanism numbers. I don't have lock-only, chunk-only or mlock-only runs. The four MI455X GPUs on that box are running an MLPerf sweep right now. What was measured:

No lock and one ~92 GiB registration per rank: all four DP ranks hung inside cudaHostRegister together (py-spy), at concurrency 128 and 256.
Lock and 1 GiB chunks: every rank registered and the servers completed, but a cold cache took about 8 min per rank.
Lock, chunks and mlock: about 25–30 s per rank.

So mlock is what fixed the slow start. The hang went away when the lock and the chunks went in together, and I can't say which of the two did it. When the GPUs are free I can run chunk-only (no lock) and lock-only (one call per gap) on a cold cache under DEP4 and post the result. If chunking alone is enough, the lock can come out.

Concurrent unregister. I have no evidence that concurrent unregister on the shared mapping is unsafe. The lock is on that path because a TP group falling back at load unregisters while other ranks may still be registering the same pages, and concurrent register is the combination that hangs. As of a1c146d it is best effort: a timeout releases without the lock. If you'd rather not take it on unregister until there's evidence, I'll drop it there.

valar-cop

some changes needed..

cold-cache dep4 run showing each rank's lock hold time and whether ranks 2-4 hit the 600s timeout?

ci green.

valar-cop · 2026-11-03T18:04:56Z

+# 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()


latching the timeout means ranks 3-4 never wait again and copy ~183 GiB each, that cold-cache dep4 case was never run..

You're right: with an absolute 600 s, a cold-cache DEP4 start without mlock (about 8 min per rank) would have timed out ranks 3 and 4 and sent them to the ~183 GiB copy. That case was never run.

The timeout now bounds the holder's silence rather than its hold time. The holder touches the lock file when it takes the lock and after every 1 GiB chunk it registers or unregisters. A waiter gives up only after 600 s with no touch. A slow cold-cache holder is waited on for as long as it keeps registering chunks, so the ranks behind it don't time out. A holder stuck in D state stops touching the file, the waiter times out, and only then is the path latched. test_a_progressing_holder_is_waited_on_past_the_timeout holds the lock for over three timeouts while touching it, and the waiter gets the lock rather than a TimeoutError.

9ebdec4

valar-cop · 2026-11-03T18:04:56Z

 | **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, and `mlock`ed while registered (faulted in and locked in one parallel pass, unlocked when the last tensor over it goes). The lock covers the shared checkpoint page cache node-wide, so host memory pressure ends in the OOM killer rather than in reclaim. A refused `mlock` (`RLIMIT_MEMLOCK` without `CAP_IPC_LOCK`) warns once and the registration stands. **Required host settings for in-place registration:** `vm.compact_unevictable_allowed=0` and `vm.compaction_proactiveness=0`; with the defaults kcompactd migrates the registered pages, locked or not, and each migration evicts the process's GPU queues (decode stalls of seconds, or a hang; watch `/sys/class/kfd/kfd/proc/<pid>/stats_*/evicted_ms`). Ranges that share a page (a layer's scales right before its values) share one reference-counted registration; a refused registration (or any failed HIP call) clears HIP's last error so the host-path fallback still runs. The page-locked memory is owned by the tensors the kernel reads and released (and unlocked) with the last of them. Needs the HIP runtime -- the one torch loaded, found in `/proc/self/maps` -- and without it the shard is registered in place. `tools/engram/bench_uva_shard.py` reproduces per-rank time, memory and gather latency for either mode. |
+| **ATOM_ENGRAM_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/<pid>/stats_*/evicted_ms`). Ranges that share a page (a layer's scales right before its values) share one reference-counted registration; a refused registration (or any failed HIP call) clears HIP's last error so the host-path fallback still runs. The page-locked memory is owned by the tensors the kernel reads and released (and unlocked) with the last of them. Needs the HIP runtime -- the one torch loaded, found in `/proc/self/maps` -- and without it the shard is registered in place. `tools/engram/bench_uva_shard.py` reproduces per-rank time, memory and gather latency for either mode. |
+| **ATOM_ENGRAM_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-<hash>.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. |


after this a timed-out path is never waited on again, please say that here and in the docstring..

Done, in the docs row, the _register_lock_file docstring, and the comment on _REG_LOCK_TIMEOUT_S: after a timeout the path is never waited on again in that process. Its later tables copy at once, and its releases unregister without the lock.

9ebdec4

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 <cursoragent@cursor.com>

karverma-amd · 2026-11-03T19:46:01Z

On the review question, a cold-cache DEP4 run with each rank's lock hold time: I haven't run it. The box is needed for other work right now. The harness is ready. It drops the page cache, starts DEP4, and samples each rank's state, D-state threads, and bytes read every 30 s until all four log UVA device lookup active. I'll post per-rank hold times for all-on, lock-only, chunk-only, mlock-only, and none when it runs.

Until then, ranks 2-4 can no longer hit the timeout just by queueing behind a slow holder. As of 9ebdec4 it fires only after 600 s with no progress from the holder.

valar-cop

a few small things..

per-mechanism isolation on a cold-cache DEP4 start, and lock-hold time per rank with whether 2-4 latch? both still unrun.

ci Offline inference smoke test (ATOM Test) failed — not sure yet (low): The log under .valar/ci is 0 bytes, so neither the failing step nor an infra cause can be identified. The same job is green on the six most recent main runs. An earlier instance of this job on this PR failed at 'Download model' with a 401 from an empty HF token, which points at the runner, but that is a different run.

(log file is empty, 0 bytes)
next: Re-run the job and attach its log or traceback before attributing it to this PR. https://github.com/ROCm/ATOM/actions/runs/41706310578/job/126162643790

valar-cop · 2026-11-03T20:06:49Z

+                        held.append(c)
+                        created += 1
+                        created_bytes += e - c
+                        _lock_heartbeat()


heartbeat only fires per chunk, the fault-in under the lock emits none? maybe call _lock_heartbeat from _populate's fault worker..

Done in cc331a4. _populate's fault worker calls _lock_heartbeat() after each faulted-in piece (256 MiB), so the fault-in acquire runs under the lock (_fault_and_lock, for what reclaim took back) now touches the file throughout. New test_fault_in_touches_a_held_lock holds a lock, backdates its mtime, runs _populate over 4 pieces, and checks the mtime moved and the heartbeat ran once per piece. It fails on the previous head (mtime stays 0) and passes now. The timeout comment and _flock_logging_wait docstring now say the holder touches after every faulted-in piece and every registered chunk.

valar-cop · 2026-11-03T20:06:49Z

+            )
+        else:
+            held = True
+            _held_lock_fds.append(fd)


_held_lock_fds is global and unlocked, heartbeat iterates it while another thread appends/removes. guard it with its own lock?

Done in cc331a4. _held_lock_fds has its own _held_lock_fds_lock. _register_lock_file appends and removes under it, and _lock_heartbeat holds it for the whole touch loop, not just a snapshot. So a descriptor leaves the list before it is closed, and a fault worker can never utime one that is closed or reused. New test_heartbeat_touches_under_the_fd_list_lock checks a heartbeat waits while the lock is held. tests/model_ops/engram/test_host_path.py: 114 passed (112 before plus the 2 new). Black and ruff clean.

karverma-amd · 2026-11-06T00:57:24Z

On the review question: a cold-cache DEP4 start with each rank's lock time, and whether ranks 2-4 hit the 600 s timeout.

Setup: 4x gfx1250, DeepSeek-V4.1-Flash, the recipe's serve_dep4.sh (-tp 4 --enable-expert-parallel --enable-dp-attention, ATOM_ENGRAM_UVA=1 ATOM_ENGRAM_UVA_COPY=0, so tables are registered in place). The code is the rocm/fw-bringup-mlperf:gfx1250-atom-20261105-dsv41 image's ATOM with this PR's host_memory.py (head 9ebdec4) and the rest of its diff applied. The host page cache was dropped (echo 3 > /proc/sys/vm/drop_caches, 20.6 GiB cached after) right before the server started.

Result: ready 220 s after start. No timeout and no skipped lock: the log has no timed out, cannot lock, cannot open registration lock or unregistering without line, so nothing latched and every rank registered in place (each logs 183.1 GiB page-locked on this rank).

	rank (in lock order)	waited for the lock	held it (register 91.55 + 2.86 GiB)
table 1 lock	1st	0 s	~2.5 s
	2nd	2.4 s	~2.5 s
	3rd	4.9 s	~2.5 s
	4th	7.4 s	~2.5 s
table 2 lock	each rank	0.1-2.5 s	~2.5 s

All four ranks finished registering 13 s after the first took the lock (00:54:56 -> 00:55:09). The longest wait, 7.4 s, is 1.2% of the 600 s timeout. At shutdown each rank's unregister took the lock within 0.1-0.2 s.

Not run: the per-mechanism isolation (lock only, chunks only, mlock only). Chunking and mlock are no longer switches in this PR, and reproducing the original all-ranks-in-cudaHostRegister hang on a shared box risked a wedged node, so I left that out.

- _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 <cursoragent@cursor.com>

valar-cop reviewed Nov 3, 2026

View reviewed changes

valarLip reviewed Nov 3, 2026

View reviewed changes

karverma-amd requested a review from valarLip November 3, 2026 16:15

valar-cop reviewed Nov 3, 2026

View reviewed changes

fix(engram): release without the lock after a timeout, latch timed-ou…

a1c146d

…t 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 <cursoragent@cursor.com>

valar-cop reviewed Nov 3, 2026

View reviewed changes

		@@ -649,7 +823,7 @@ def _register(sources, leases):
		flags = _HOST_REGISTER_READ_ONLY if _read_only(lo, hi) else 0
		if flags:



		def _reg_chunk_bytes() -> int:
		"""Largest single `cudaHostRegister` (`ATOM_ENGRAM_REG_CHUNK_MB`), in whole

		assert len(cudart.registered) >= 1


		def test_registration_and_unregistration_take_the_lock(monkeypatch):

		@@ -32,12 +32,15 @@
		during a capture is skipped, with a warning, rather than breaking it.
		"""

		@@ -33,14 +33,19 @@
		during a capture is skipped, with a warning, rather than breaking it.
		"""

Conversation

karverma-amd commented Nov 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Nov 3, 2026

🏷️ CI Guide

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

karverma-amd commented Nov 3, 2026

Uh oh!

valarLip left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

karverma-amd commented Nov 3, 2026 •

edited

Loading