[yhl-amd]

Summary

mlock the page runs the Engram UVA path registers in place, for as long as they are registered. On by default, best effort: a refused mlock warns once and the registration stands. No new environment variable.

Problem

An in-place registration (ATOM_ENGRAM_UVA_COPY=0, or the DP-attention default) pins the checkpoint page cache through a KFD userptr BO. Reclaim still treats those pages as reclaimable. Once host memory gets tight, kswapd reclaim or kcompactd migration touches them. Each MMU notifier evicts all of the process's GPU queues until the BO is revalidated.

On DSV4.1-Flash DEP4 (183 GiB registered in place, MI455X), a 1 h MLPerf Endpoints c512 run:

• started stalling at 14.6 min;
• hung at 31.5 min, with all 4 ranks in stream.synchronize and the GPUs idle.

The baseline config without our other changes did the same. /sys/class/kfd/kfd/proc/<pid>/stats_*/evicted_ms grew with every stall. The trigger was host memory filling up: a benchmark client left ~160 GB in /dev/shm.

Fix

• The lock lives in _PageRuns: acquire locks each run it creates, release unlocks a run when its last reference goes, before cudaHostUnregister. A refused registration unlocks what that acquire took. A shard the TP group refuses, and shutdown(), unlock with the release.
• Locking and faulting are one parallel pass: mlock2(MLOCK_ONFAULT) over the run, then MADV_POPULATE_READ per chunk on the existing thread pool (pages arrive locked). If the lock is refused, the read-only range is still faulted in.
• The refusal warning names the errno, gives the RLIMIT_MEMLOCK hint only for EPERM/ENOMEM under a finite limit, and logs GiB.
• Side effect, in the ATOM_ENGRAM_UVA_COPY docs: the shared checkpoint page cache is locked node-wide, so host memory pressure ends in the OOM killer instead of a stall.
• Required host settings for in-place registration (documented): vm.compact_unevictable_allowed=0 and vm.compaction_proactiveness=0. With the defaults kcompactd migrates locked pages too, which fires the same notifier.
• The copy path is unaffected: its hipHostMalloc memory is already locked.

Validation

1 h c512 runs (same machine and config as the failure, mlock on plus both sysctls), two runs:

• 0 stalls and no hang in either run;
• evicted_ms grew only during startup, before the sysctls were set, and by a few ms over the whole hour of load;
• 0 failed requests.

Those runs had both sysctls set, so they do not separate mlock from the sysctls. Whether mlock alone (sysctls at their defaults) removes the stalls is not tested yet; until it is, the docs make both sysctls required.

Startup cost, host side only (no GPU, no HIP registration): an 8 GiB file, warm in the page cache, mapped shared read-only like the checkpoint and mapped afresh per trial. 16 threads (copy_threads), 64 CPUs, kernel 6.8. 6–10 trials each:

per 8 GiB	time
main: parallel populate, no lock	0.13–0.16 s
413eeb0: parallel populate, then one serial mlock	0.49–0.53 s
per-chunk mlock in the pool	0.38–0.43 s
this PR: mlock2(ONFAULT) + parallel populate	0.25–0.29 s

At 183 GiB per rank, that is ≈2.7 s added over main, down from ≈7.8 s. VmLck was the full 8 GiB in each variant.

Unit tests, tests/model_ops/engram/test_mlock.py (7 cases; pass as root, in a container without CAP_IPC_LOCK, and as an unprivileged user with ulimit -l 64; black and ruff clean):

• acquire locks exactly the runs it creates (only the uncovered tail of an overlap), faults read-only runs in chunks, and leaves writable ones to registration;
• release unlocks a run with its last reference only;
• a refused registration unlocks the runs that acquire created, and leaves runs held by other ranges locked;
• a refused lock still faults the range in and warns once, with strerror, GiB and the rlimit hint only for EPERM;
• real mlock of a file mapping: VmLck grows by exactly its size and returns to the baseline on release;
• real refusal under RLIMIT_MEMLOCK=0: VmLck unchanged, one warning; skipped only under CAP_IPC_LOCK.

🤖 Generated with Claude Code

[github-actions]

🏷️ 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 2898 --add-label <label>

[valar-cop]

some changes needed..

ci green.

[valar-cop]

mlock before acquire succeeds? if acquire returns None the range stays locked forever, move it after or munlock on the None path

[yhl-amd]

Moved: _register now mlocks the shard's ranges only after every range is registered, so a refused registration (host-path fallback) leaves nothing locked. Covered by test_register_locks_only_a_complete_registration. 0b6d34e

[valar-cop]

skip on empty caplog means deleting the warning makes this pass.. probe mlock permission instead

[yhl-amd]

Now probes: lower RLIMIT_MEMLOCK to 0, try a raw mlock; skip only if it still succeeds (CAP_IPC_LOCK). Otherwise the test asserts VmLck is unchanged and exactly one warning is logged, so deleting the warning fails it.

[valar-cop]

nothing tests the env gate itself, and the autouse fixture should delenv ATOM_ENGRAM_MLOCK so local shells don't change behavior

[yhl-amd]

The switch is gone, mlock is the default now, so there is no gate to test or env to clear. Added a test for the register-then-lock order instead.

[valar-cop]

worth one line: this pins the shared checkpoint page cache node-wide, reclaim pressure becomes oom-kill

[yhl-amd]

Added to the ATOM_ENGRAM_UVA_COPY row (the ATOM_ENGRAM_MLOCK row is gone with the switch): the lock covers the shared checkpoint page cache node-wide, so memory pressure ends in the OOM killer rather than a stall.

[yhl-amd]

Pushed 0b6d34e (merged current main first):

• ATOM_ENGRAM_MLOCK removed; in-place registrations are mlocked by default, best effort (one warning on refusal, startup unaffected).
• Lock only after the whole shard is registered, so a fallback leaves nothing locked.
• Tests probe the mlock permission instead of skipping on an empty log; new test for the lock order.
• Node-wide effect documented under ATOM_ENGRAM_UVA_COPY.

@valarLip @valar-cop please re-review, thanks.

[valar-cop]

some changes needed..

ci green.

[valar-cop]

lock is per table, not per rank shard.. earlier layers stay mlocked when a later one falls back. maybe thread the ranges out and mlock once in _register_shard?

[yhl-amd]

Right, moved in bc01784: _register no longer locks. _enable_uva calls host_memory.mlock_in_place once, after the TP vote, over every pinned table; it locks only what _PAGE_RUNS holds in place (copies skipped). New test_in_place_shard_is_mlocked_only_once_the_group_keeps_it covers agree / outvoted / copy.

[yhl-amd]

Pushed bc01784: the lock now happens once per rank, after the TP group keeps the shard (host_memory.mlock_in_place from _enable_uva), so no fallback leaves layers locked. @valar-cop please re-review.

[valar-cop]

a few small things..

ci green.

[valar-cop]

bench tool registers in place without mlock now, so its VmLck report no longer matches server. maybe call mlock_in_place there too, or note it only measures registration?

[yhl-amd]

Done in 413eeb0: the bench calls mlock_in_place after enabling the tables, as _enable_uva does, so its VmLck and timing match the server.

[yhl-amd]

413eeb0: the bench tool now mlocks like the server. All checks are green, including the ci:atom accuracy runs (11/11).

@valar-cop re-review please. @valarLip could you take a look when you have time? thanks.

[valarLip]

Follow-up to my approval of 413eeb0, after a closer read and a small host-side measurement; tests, ruff and black pass (the 6 mlock cases included). Locking the in-place pages is the right fix and I'm not withdrawing the approval, but the lock lives at the wrong layer, and that's why four call sites had to be added around it. Each item ends with the change I'm asking for.

Main

1. Tie the lock to _PageRuns, not to a post-vote pass over tensors (engram/host_memory.py:646, :548-612). _PageRuns already holds exactly the in-place ranges: page-aligned, disjoint, reference-counted. acquire registers only uncovered gaps as new runs, and release unregisters a run when its count reaches 0. So mlock when acquire creates a run and munlock when release drops it to 0 keep the lock's lifetime identical to the registration's. The docstring's reason for never unlocking ("ranges that share a page would unlock each other") doesn't hold at run granularity, since runs never overlap. It also answers valar-cop's fallback concern by construction: a shard the TP group refuses is released, and its runs unlock with it. Because the lock isn't on the runs, the PR needed covers(), TableShard.shard_tensors(), the post-vote call in host.py:442-446 and the same call in tools/engram/bench_uva_shard.py:117-120. And it still misses one case: after shutdown() / _release_shard() unregister, the locked pages (183 GiB per rank at DEP4) stay locked for as long as the checkpoint memmap lives (an engine rebuilt in the same process, tests).
→ Change: in _PageRuns.acquire, lock each new run ([cursor, gap_end)) after cudaHostRegister succeeds. In release, munlock the run before cudaHostUnregister when its count hits 0. Then delete mlock_in_place, covers(), shard_tensors() and the two call sites.

2. The post-vote pass is a serial startup cost with no number in the PR (engram/host.py:444). mlock faults pages in itself, which is what _populate's MADV_POPULATE_READ already did. On an 8 GiB read-only shared file mapping here, populate took 0.33 s and an mlock after it another 0.26–0.42 s. A direct mlock that does the faulting took 0.53 s, about the same as populate plus mlock. At 183 GiB per rank (DEP4), that's roughly 6–9.5 s of single-threaded work added after the vote.
→ Change: with item 1, do the lock inside _register / _populate's existing chunked thread pool (mlock per chunk, falling back to MADV_POPULATE_READ when refused), so the pages are faulted and locked in one parallel pass. Please post the startup time before and after.

3. The evidence doesn't separate mlock from the sysctls (docs/environment_variables.md:321). Both 1 h validation runs had vm.compact_unevictable_allowed=0 and vm.compaction_proactiveness=0 set. With the default compact_unevictable_allowed=1, kcompactd still migrates mlocked pages, and that fires the same MMU notifier and queue eviction this PR describes. So it's open whether mlock alone, which this PR makes the default, removes the stalls, or whether the sysctls did. The docs mention only the first sysctl.
→ Change: run the c512 1 h case once with mlock on and both sysctls at their defaults. If it still stalls, document both sysctls as required (including compaction_proactiveness) rather than as optional.

Smaller

4. ctypes.CDLL(None, use_errno=True) plus argtypes is now the fourth copy in this file (host_memory.py:196, :389, :628, :677), and test_mlock.py's _mlock_allowed adds a fifth. Each call re-dlopens libc and re-declares the signature.
→ Change: one @functools.cache _libc() that declares madvise, mlock, munlock and prctl once.

5. The page-alignment arithmetic is written three times (host_memory.py:663-664, _register at :706-707, tests/model_ops/engram/test_mlock.py:88-89).
→ Change: a _page_span(t) -> (lo, hi) helper; with item 1 the copy in mlock_in_place goes away.

6. Docs don't match the code, and the reason is written four times. The ATOM_ENGRAM_UVA_COPY row still says the page-locked memory is "released with the last of them", but the mlock is never undone. The why of mlock appears in that row, in mlock_in_place's docstring, in host.py and in the bench comment, and the module docstring's Ownership section doesn't mention mlock.
→ Change: keep the reason once, in the docstring of whatever owns the lock after item 1. Have the doc row state the behaviour and the required sysctls, add a line to the Ownership section, and drop the host.py and bench comments.

7. The tests fake registrations through private state, and one can't fail for the right reason (tests/model_ops/engram/test_mlock.py:87). runs._runs[lo] = [hi, 1] writes _PageRuns' internals and repeats the alignment arithmetic. test_host_path asserts only len(locked) equals the tensor count, so locking the wrong ranges (e.g. a copy's address) still passes. test_mlock_locks_the_mapping locks without unlocking.
→ Change: register through acquire and assert the exact (lo, hi) set locked. With item 1, add the release → munlock case, and unlock in every test that locks.

8. The warning reports one errno and always blames the rlimit (host_memory.py:685). It keeps only the last errno but always says "raise RLIMIT_MEMLOCK". EAGAIN (some pages couldn't be locked) and an ENOMEM that isn't the rlimit point elsewhere. It also prints raw bytes (196608000000) where the rest of the module logs GiB.
→ Change: name the errno (os.strerror), give the rlimit hint only for EPERM/ENOMEM under a finite RLIMIT_MEMLOCK, and log GiB.

Also, the history adds ATOM_ENGRAM_MLOCK and then removes it, so please squash on merge.

[yhl-amd]

@valarLip thanks for the follow-up. Pushed 38fd2bb, on top of a merge of current main. Item by item:

1. Lock on _PageRuns. Done. acquire locks each new run before cudaHostRegister. If the registration is refused, it unlocks that gap, then the runs this acquire took are released and unlocked. release munlocks a run when its count reaches 0, before cudaHostUnregister. mlock_in_place, covers(), shard_tensors() and the host.py / bench call sites are gone. host.py, tables.py and the bench tool now match main. So a refused shard, shutdown() and a rebuilt engine all unlock with the release.

2. Startup cost. The lock now runs inside the existing chunked pool. I tried the per-chunk mlock first, but it scaled poorly, so the code calls mlock2(MLOCK_ONFAULT) once per run (this marks the range locked without faulting it) and then runs MADV_POPULATE_READ per chunk on the pool. Pages arrive already locked, so it is still one parallel pass. If the lock is refused, a read-only range is still populated. I measured this host side only (no GPU, no HIP registration), on the same kind of mapping you used: an 8 GiB shared read-only file mapping, warm cache, mapped afresh per trial, 16 threads, 64 CPUs, kernel 6.8, 6–10 trials each:

• main (populate only): 0.13–0.16 s
• 413eeb0 (populate, then one serial mlock): 0.49–0.53 s
• per-chunk mlock in the pool: 0.38–0.43 s
• now (mlock2 ONFAULT + parallel populate): 0.25–0.29 s

At 183 GiB per rank, that's about +2.7 s over main instead of +7.8 s. VmLck was the full 8 GiB in every variant. I don't have a GPU startup number for this round.

3. mlock vs the sysctls. I haven't run the 1 h c512 case with the sysctls at their defaults. So, as in your second option, the ATOM_ENGRAM_UVA_COPY row now lists vm.compact_unevictable_allowed=0 and vm.compaction_proactiveness=0 as required for in-place registration. The PR description says the existing runs had both set, and that mlock's effect on its own is still to be checked. I'll follow up with that run.

4. A single @functools.cache _libc() now declares madvise, mlock2, munlock and prctl. The allocator, _source_dropper and the lock path all use it.

5. Added _page_span(t). _register uses it, and the mlock_in_place copy is gone.

6. The reason is now written once, in _PageRuns' docstring. The env row states only the behaviour and the two required sysctls, and "released (and unlocked) with the last of them" matches the code. The Ownership section has a new line. The host.py and bench comments are gone.

7. The tests register through acquire (with a fake cudart) and assert the exact spans locked and unlocked:

• the new runs only;
• release → munlock on the last reference;
• a refused gap unlocks what that acquire created and leaves other ranges' runs locked;
• the real-mapping test unlocks through release and checks VmLck goes back to the baseline.

None of them writes _runs directly.

8. The warning now uses os.strerror, gives the rlimit hint only for EPERM/ENOMEM under a finite RLIMIT_MEMLOCK, and logs GiB. It is logged once per process.

I'll squash on merge as you asked.

@valar-cop re-review please.

[valar-cop]

LGTM

ci failed: Pre Checkin, could not tell why yet..

[yhl-amd]

Note on the red Run non-GPU unit tests: collection fails on tests/model_ops/test_dp_gather_routed_rows.py (No module named 'triton.language.extra'). main's pre-checks hit the same error on every run since af4d581, so it isn't from this PR. I'll re-run once main is fixed.

[valar-cop]

ci triage for 38fd2bb:

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

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

Pre Checkin / Run non-GPU unit tests failed — main is already broken the same way (medium): The non-GPU unit test job is already failing repeatedly on main (4 consecutive failures starting ~2026-11-03T14:17, after successes earlier that day), so the PR inherits a pre-existing break rather than introducing it. No job log was provided to confirm the failure mode matches, but the main-side failure cluster is the dominant signal.

Same job on main (newest first): ['2026-11-03T14:52 failure', '2026-11-03T14:40 failure', '2026-11-03T14:39 failure', '2026-11-03T14:17 failure', '2026-11-03T13:46 success', '2026-11-03T08:20 success', '2026-11-03T08:13 success', '2026-11-03T05:49 success']
Same job on other recent runs: {'failure': 4, 'success': 1, 'cancelled': 2}
Infra signatures matched by regex: none
next: Do not chase this in the PR: open/inspect the failing 'Run non-GPU unit tests' job on main, bisect the commit merged between the 13:46 success and the 14:17 failure to find the breaking change, and rerun this PR once main is green. If the PR's new tests/model_ops/engram/test_mlock.py is itself the culprit on main, fix that file. https://github.com/ROCm/ATOM/actions/runs/41659665374/job/125999210900