Heterogeneous Memory Management (HMM): The Data Structures Behind CPU–GPU Page Migration

1. Introduction

Every time a CUDA, ROCm, or oneAPI program calls malloc() and hands the pointer straight to the GPU without ever calling cudaMallocManaged() or pinning a buffer, something has to reconcile two completely different memory systems: a CPU page table walked by the MMU, and a GPU page table walked by hardware that has never heard of struct page. That reconciliation is the job of HMM — Heterogeneous Memory Management — a kernel subsystem that lives almost entirely in mm/hmm.c and mm/migrate_device.c.

HMM is not a scheduler policy and not a NUMA balancing heuristic. It is a set of data structures and synchronization primitives that let a device driver ask the mm subsystem, “what physical memory currently backs this range of a process’s address space, and can I safely migrate it to my device?” The interesting engineering is entirely in how that question is answered correctly while the answer can change out from under you at any instant — because the CPU can unmap, fork, swap, or collapse a THP in the middle of a GPU page fault.

This lesson builds a faithful userspace model of that race, breaks it deliberately, catches the breakage with ThreadSanitizer and AddressSanitizer, and then fixes it using the exact synchronization pattern HMM uses in the kernel: sequence-counter retry plus page pinning.

2. Historical Background

Before HMM, heterogeneous memory on Linux meant one of two unattractive options. Either the application used a special allocator (cudaMallocManaged, or vendor-specific unified memory APIs) that carved out a dedicated address range the driver fully owned and mirrored, or it pinned ordinary pages with get_user_pages() and DMA’d them in place, permanently, for the lifetime of the mapping. Pinning works but doesn’t scale: pinned memory can’t be swapped, can’t be moved by compaction, and can’t participate in normal reclaim. A GPU workload that touches gigabytes of working set but only needs megabytes resident at once has no good answer under that model.

HMM, merged incrementally between Linux 4.14 and 4.20, generalized the page fault itself. Jérôme Glisse’s original design let a device register a “mirror” of a process’s address space and receive faults through the same path the CPU MMU uses, backed by mmu_notifier (later mmu_interval_notifier) infrastructure that had already existed for KVM and RDMA. The payoff: ordinary malloc()‘d memory becomes GPU-accessible on demand, migrates to device memory when the GPU touches it, and migrates back when the CPU touches it — all without the application doing anything special. ZONE_DEVICE, added around the same era, gave device memory its own struct page representation so the rest of the mm subsystem (reference counting, migrate_pages(), even /proc/pid/smaps in some configurations) could reason about it using existing machinery instead of a parallel one.

The API itself went through a notable simplification. The earliest HMM mirroring interface exposed a much wider surface — drivers registered fault and invalidate callbacks directly and managed a fair amount of bookkeeping themselves. By the time HMM stabilized, the API had collapsed down to essentially two entry points a driver author needs to understand deeply: hmm_range_fault() for turning an address range into a PFN list, and the migrate_vma_*() family for actually moving pages between host and device memory. The mmu_interval_notifier conversion (replacing the older, coarser mmu_notifier in this role) was itself a correctness-driven rewrite — the older interface made it easy for drivers to accidentally miss an invalidation window, exactly the class of bug this lesson reproduces.

It’s worth being precise about what HMM is not. It is not a NUMA migration policy — that’s the job of subsystems like DAMON-based migration, numa_balancing, or explicit move_pages() calls, which decide when and whether migration should happen based on access patterns. HMM sits one layer below: it is the mechanism that makes any migration safe to perform at all against a live, mutable address space. A NUMA balancer and an HMM-based GPU driver can both be moving pages around the same process concurrently, and the correctness burden HMM exists to satisfy is exactly what stops those two movers from stepping on each other.

3. The Systems-Level Problem

The problem HMM solves has a name in the concurrent-programming literature: TOCTOU (time-of-check to time-of-use), applied to page tables. A GPU driver’s fault handler does roughly:

1. Look up which host pages back a faulting address range.

2. Do something slow with that information — set up a DMA transfer, migrate the pages to device memory, install device page-table entries.

3. Let the GPU proceed, now trusting that the device PTEs point at valid memory.

Between steps 1 and 2, the kernel can invalidate that mapping for reasons that have nothing to do with the GPU: the process calls munmap(), a fork() triggers copy-on-write, the reclaim path swaps a page out, or a transparent huge page gets split or collapsed. If the driver’s snapshot from step 1 is stale by the time it commits in step 3, it can install a device-visible mapping to memory that has been freed and reallocated for something else entirely — a heap-use-after-free with a hardware DMA engine as the writer, which is about as unpleasant as a race condition gets.

HMM’s entire data-structure design is aimed at making that race unrepresentable if a driver author follows the API contract, and it is very representable if they don’t — which is a real bug class in tree, not a hypothetical.

4. Linux Kernel Architecture

The relevant pieces span three layers, shown in the diagram above:

• Userspace: the application never opts in explicitly; a CUDA/ROCm runtime calling regular malloc() is enough. The GPU user-mode driver (UMD) submits work and receives device page faults from the kernel-mode driver.

• Syscall interface: mmap() establishes the VMA the driver will mirror; ioctl() calls into the DRM subsystem register the mirror and handle fault notifications; madvise(MADV_DONTNEED) and friends can trigger the invalidation paths HMM has to survive.

• Kernel: mm/hmm.c owns the range-fault API and works directly against mm_struct, vm_area_struct, and the page tables. mm/migrate_device.c owns the actual page migration once HMM has identified which pages need to move. The mmu_interval_notifier (generalized from the older mmu_notifier) is the synchronization primitive gluing them together — it is the object that tells a driver “your snapshot might be stale, recheck before you trust it.”

A driver’s involvement starts well before any fault happens. At probe/init time it calls mmu_interval_notifier_insert() to register interest in a [start, end) range against a target mm_struct, supplying an ops table whose .invalidate() callback the mm subsystem will call synchronously whenever something in that range changes. That callback’s job is deliberately narrow: bump the sequence counter and, if the invalidation type demands it (e.g. an actual unmap rather than a permissions change), block until any pins the driver is holding on the affected pages are released. Everything downstream of that — deciding whether to migrate, which pages, in which direction — is driver policy layered on top of a kernel mechanism that only guarantees one thing: you will find out, cheaply and reliably, if your snapshot went stale.

This layering is why the same mmu_interval_notifier machinery serves GPU drivers, RDMA on-demand paging, and even KVM’s original mmu_notifier use case for shadow page tables — the correctness problem (”a hardware page table mirrors a software one that can change”) is identical regardless of what’s on the other end of the mirror.

5. Internal Working

The two structures that matter most:

/* include/linux/hmm.h (simplified) */
struct hmm_range {
    struct mmu_interval_notifier *notifier;
    unsigned long        notifier_seq;   /* snapshot at read_begin() */
    unsigned long        start;
    unsigned long        end;
    unsigned long        *hmm_pfns;      /* per-page PFN + flags array */
    unsigned long        default_flags;
    unsigned long        pfn_flags_mask;
    void                 *dev_private_owner;
};

notifier_seq is the field this entire lesson is about. It is captured with mmu_interval_read_begin(notifier) before the driver walks the page tables, and it must be rechecked with mmu_interval_read_retry(notifier, seq) after the driver has done its slow work but before it commits anything device-visible. If the sequence changed, the correct action is not “proceed carefully” — it is “throw the result away and refault.” There is no safe way to patch up a stale snapshot; the only correct move is to redo it.

The second half of the mechanism is page pinning. A stale sequence number alone only tells you that something changed; it doesn’t prevent the underlying page from being freed while you’re still looking at it. HMM-based drivers take a reference (conceptually get_page()/folio_get()) on the pages they’re inspecting, and the invalidation path is required to wait for outstanding references to drop before it actually reclaims the page — it can announce the invalidation (bump the sequence, so readers know to discard their snapshot) without having to block every reader synchronously, because the pin keeps the memory itself alive during the (short) drop-and-refault window.

6. Step-by-Step Execution Flow

Github link:

https://github.com/sysdr/howtech-p/tree/main/Data_Structures_Heterogeneous/hmm-demo

1. GPU hardware raises a page fault on an address the device page table doesn’t have mapped.

2. The kernel-mode driver’s fault handler calls into hmm_range_fault(), which walks the CPU page tables for the faulting range.

3. Before the walk, mmu_interval_read_begin() captures notifier_seq.

4. The driver pins the discovered pages and begins migration via migrate_vma_setup() / migrate_vma_pages() — a DMA copy from host DRAM to device HBM, non-trivial in duration.

5. Concurrently, anything in the kernel that changes the mapping (munmap, COW fork, THP split, swap-out) fires the registered mmu_interval_notifier invalidate callback, which bumps the sequence counter.

6. After the migration work completes, the driver calls mmu_interval_read_retry(notifier, seq).

7. Decision node: if the sequence changed, drop the pin and refault — go back to step 3 with a fresh snapshot. If unchanged, proceed.

8. The driver calls migrate_vma_finalize(), installing the device page-table entry.

9. The pin is dropped.

10. The GPU resumes execution against the now-valid device mapping; control returns to userspace only once the fault is resolved from the GPU’s perspective.

7. Kernel Data Structures

Beyond struct hmm_range, the structures worth knowing by name:

• struct mmu_interval_notifier — registers a [start, end) range of interest against an mm_struct and carries the ops table (.invalidate()) the mm subsystem calls into.

• struct migrate_vma — driven by migrate_vma_setup(); carries src and dst PFN arrays that describe the in-flight migration on a per-page basis, plus the same VMA range HMM already walked.

• struct page / struct folio with ZONE_DEVICE — device-resident memory gets real struct page entries via devm_memremap_pages(), so put_page(), reference counting, and even some reclaim paths work unmodified against device memory instead of needing a parallel type system.

• hmm_pfns[] — the per-page flags array in hmm_range, encoding whether a page is valid, needs a fault, is write-protected, or is already device-resident (HMM_PFN_VALID, HMM_PFN_WRITE, HMM_PFN_ERROR, and friends).

The per-page flags matter because a single hmm_range_fault() call can return a mix of outcomes across the range — some pages already resident and mappable immediately, some requiring the caller to fault them in (e.g. they’re currently swapped out), and some simply erroring out (e.g. the address isn’t backed by anything, or belongs to a VMA type HMM doesn’t support like a device-special mapping). A driver has to walk this array and handle each case, not assume uniform success across the range:

/* Simplified caller-side pattern after hmm_range_fault() returns */
for (i = 0; i < npages; i++) {
    unsigned long pfn_flags = range->hmm_pfns[i];

    if (pfn_flags & HMM_PFN_ERROR) {
        /* unrecoverable for this page: SIGBUS-equivalent to the device */
        continue;
    }
    if (!(pfn_flags & HMM_PFN_VALID)) {
        /* page not resident; caller must retry with fault-in requested */
        need_refault = true;
        continue;
    }
    /* pfn_flags encodes the PFN itself plus HMM_PFN_WRITE etc. */
    device_pfn = hmm_pfn_to_pfn(pfn_flags);
}

This is a direct structural analog of what the demo’s hmm_pfns-equivalent (g_host_pages[] plus the pin count) is doing, simplified down to the single property this lesson is about: can you trust the pointer you’re holding, and is it still backed by live memory.

8. CPU-Level Behaviour

The synchronization pattern is deliberately lock-light on the fast path. mmu_interval_read_begin()/read_retry() is a sequence-counter (seqcount-style) pattern, not a mutex: readers never block writers, and the common case (no concurrent invalidation) costs one atomic load at the start and one at the end, with no cache-line contention against other faulting threads. This matters because GPU fault handling is latency-sensitive — thousands of faults per second under real workloads — and a coarse-grained lock across the whole migration would serialize unrelated ranges against each other for no reason.

The cost is pushed onto the rare path: a genuine invalidation means discarded work and a refault, which is strictly correct but can show up as measurable overhead under pathological access patterns (e.g., a CPU thread and GPU kernel ping-ponging writes to the same page, which is a real anti-pattern user code can hit and one HMM cannot fix for you).

9. Performance Analysis

The demo below models exactly this pattern: one thread continuously performs range-fault-and-migrate, another continuously invalidates at the same granularity, deliberately adversarial to maximize the refault rate for demonstration. In the corrected implementation, run against 20,000 fault rounds across 64 page slots:

hmm_sim_after: completed 20000 rounds on 64 slots, 17665 refaults

An 88% refault rate under this artificial, maximally-contended workload is expected — it’s a stress test, not a representative access pattern. Real GPU workloads see refault rates close to zero because invalidation of a range a GPU kernel is actively resident in is rare; the mechanism is sized for correctness under worst-case concurrency, not tuned assuming it will fire often. Running the same demo three times shows the refault count moving with scheduling noise rather than converging to a fixed value, which is itself the expected signature of a genuine race window rather than a deterministic algorithm:

Build variant Refaults (of 20,000) Sanitizer result -O2 strict release build 17,665 n/a (no sanitizer) ThreadSanitizer build 6,638 clean ASan + UBSan build 17,392 clean Valgrind (-O0) 19,624 0 errors, 0 leaks

The spread — from roughly a third to nearly all rounds refaulting — comes entirely from each instrumentation’s effect on relative thread scheduling and timing, not from any change in the underlying algorithm; the usleep(50) window in the fault path is fixed, but how much CPU time the invalidate thread gets to race into that window varies with how much overhead the sanitizer adds around each memory access. This is a useful intuition to carry into kernel debugging generally: a race’s reproduction rate under a given tool is not a measure of the race’s real-world frequency, only of how that tool happens to perturb scheduling.

The actionable performance lesson for driver code: keep the window between read_begin() and read_retry() as short as possible, since every microsecond in that window is exposure to a refault, and refaults on real hardware mean a repeated DMA setup, not just a cheap retry loop. Production GPU drivers batch this — a single hmm_range_fault() call typically covers many pages at once specifically to amortize the fixed cost of a snapshot-and-retry cycle across as much migrated data as possible, rather than doing it one page at a time the way this simplified demo does for clarity.

10. Debugging Techniques

Three tools catch three different facets of this bug class:

ThreadSanitizer catches the race itself — two threads touching the same memory without a happens-before edge — even in the (surprisingly common) case where the racing writes don’t cause immediately visible corruption:

WARNING: ThreadSanitizer: data race (pid=956)
  Write of size 2 at 0x720400000000 by thread T1:
    #0 hmm_range_fault_and_migrate hmm_sim_before.c:103
    #1 fault_worker hmm_sim_before.c:112
  Previous write of size 8 at 0x720400000000 by thread T2 (mutexes: write M0):
    #0 malloc <libtsan interceptor>
    #1 alloc_host_page hmm_sim_before.c:67
    #2 invalidate_worker hmm_sim_before.c:129

AddressSanitizer catches the consequence — a genuine heap-use-after-free, because the “invalidate” side in this model actually frees and reallocates the backing memory, mirroring real page reclaim:

==967==ERROR: AddressSanitizer: heap-use-after-free on address 0x502000000010
WRITE of size 2 at 0x502000000010 thread T1
    #0 hmm_range_fault_and_migrate hmm_sim_before.c:103
freed by thread T2 here:
    #0 free
    #1 invalidate_worker hmm_sim_before.c:127
previously allocated by thread T0 here:
    #0 malloc
    #1 alloc_host_page hmm_sim_before.c:67

Valgrind’s memcheck is the right tool for the fixed version specifically because it validates the absence of leaks introduced by the pin/unpin bookkeeping — every pin needs a matching unpin on both the commit path and the refault path, and a leaked pin count would silently disable the fix without crashing anything.

On real kernel code, the equivalent debugging surface is CONFIG_DEBUG_ATOMIC_SLEEP (catches sleeping inside the notifier invalidate callback, which must not block), lockdep annotations on the mmu_interval_notifier machinery, and ftrace events around mmu_notifier_invalidate_range_start/end for tracing actual invalidation timing against driver fault handling in production.

Two things are worth calling out about applying userspace sanitizers to this style of kernel-adjacent logic. First, the race only reproduces reliably because the demo’s usleep(50) widens the window enough for the scheduler to interleave the two threads inside it — remove that call and the same bug exists but may take millions of iterations to surface, which is the userspace analog of a kernel race that only shows up on specific hardware timing. Don’t mistake “sanitizer didn’t fire in N runs” for “no race exists”; absence of a report under a fixed number of iterations is evidence, not proof. Second, TSan and ASan are catching genuinely different failure modes here even though they’re triggered by the same root cause: TSan flags the unsynchronized access regardless of whether it’s harmful in a given run, while ASan only fires once the access actually lands on freed memory. Running both is not redundant — a version of this bug that raced on non-freed, merely stale data would still be caught by TSan and missed entirely by ASan.

11. Production Failure Scenarios

Three real-world triggers for the exact race modeled here:

• fork() under active GPU use. A COW fork of a process with GPU-resident mappings triggers exactly the kind of invalidation a naive driver can race against — this is precisely why HMM’s documentation calls out fork handling explicitly as a case drivers must test.

• THP collapse/split racing a fault. Transparent huge pages can be split or collapsed by khugepaged concurrently with a device fault walking the same range; a driver that snapshots PFNs without the retry check can migrate a page whose backing has just changed shape underneath it.

• Swap-out under memory pressure. If system memory pressure triggers reclaim on a page mid-fault, the naive driver is racing the exact free-then-reuse pattern this lesson’s buggy version reproduces deterministically under TSAN.

• Multi-GPU contention on shared address ranges. In multi-accelerator systems, two devices can both hold mmu_interval_notifier registrations against overlapping ranges of the same process. An invalidation triggered by device A’s migration can race device B’s in-flight fault handler exactly as the CPU-vs-GPU case does here — the notifier mechanism doesn’t care which side of the race is “the CPU”; it only cares that a range changed while someone else was mid-snapshot. Driver authors who tested only single-GPU configurations have historically missed this until multi-GPU topologies exposed it.

• Userspace calling madvise(MADV_DONTNEED) concurrently with GPU compute. This is a directly reachable, non-exotic trigger: any application that frees or resets a buffer while a kernel is still running on the GPU is exercising this exact invalidation path, which is precisely why kernel selftests for HMM (tools/testing/selftests/mm/hmm-tests.c) explicitly include concurrent-invalidation test cases rather than relying only on single-threaded correctness checks.

What makes this bug class particularly dangerous in production rather than merely academic is that all three triggers above are routine operations from the CPU side — nothing about fork(), THP management, memory pressure, or madvise() is unusual or attacker-controlled. A driver with this bug doesn’t need a hostile workload to fail; it needs an ordinary Linux system doing ordinary Linux things at the wrong moment relative to GPU activity, which is exactly why the failure often surfaces only after a driver has shipped and accumulated enough real-world usage hours to hit the timing window.

Working demo link:

12. Real-World Production Use Cases

Nouveau’s SVM (Shared Virtual Memory) support and AMDGPU’s KFD SVM path are the two upstream, shipping consumers of this exact API surface, both built directly on hmm_range_fault() and mmu_interval_notifier. Outside GPUs, RDMA on-demand-paging (ODP) uses the same underlying mmu_interval_notifier infrastructure to let InfiniBand hardware fault in memory registrations lazily instead of requiring the whole region pinned up front — a different device, the identical TOCTOU problem, the identical fix.

13. Hands-on Lab

setup.sh builds both the intentionally-buggy and the corrected version of the simulator, runs the full validation gauntlet, and shows you the failure and the fix side by side. Run it with:

chmod +x startup.sh --docker
./startup.sh --docker

Expect it to: detect your distro and kernel version, install build-essential and valgrind if missing, compile the buggy version under TSan and ASan (both should report the race/UAF), compile the fixed version under -Wall -Wextra -Werror -O2, TSan, ASan+UBSan, and Valgrind (all should pass clean), and print a summary.

The exercise worth doing by hand afterward: open hmm_sim_before.c, find the missing notifier_retry() check documented in the comment above hmm_range_fault_and_migrate(), and add it yourself — following the same pattern already implemented in hmm_sim_after.c — then rerun the TSan and ASan builds against your patched version. Watching the exact same reports from Section 10 disappear once you’ve added the three lines that constitute the fix is a more durable way to internalize the API contract than reading about it, and it mirrors the actual code-review question a kernel maintainer would ask of a first HMM driver patch: “where’s your retry check, and what happens if it fires?”

14. Best Practices

• Never trust a page-table snapshot across any operation that can sleep or take meaningful time; always pair mmu_interval_read_begin() with a read_retry() immediately before committing device-visible state, with as little work as possible in between.

• Pin what you’re actively migrating, and make sure every code path — including error and refault paths — drops the pin. A pin leak is a correctness bug that manifests as a hang or a stuck reclaim, not a crash, and is easy to miss in testing.

• Treat a positive read_retry() as the normal case, not an exceptional one, in your capacity planning — write the refault loop assuming it will fire under contention, because it will.

• Keep invalidate callbacks non-blocking. The whole design assumes the invalidate side is cheap (bump a counter); if you make it wait synchronously on unrelated work, you reintroduce the kind of stall the seqcount pattern was chosen specifically to avoid.

• Test with an adversarial invalidator, not just a quiet one. This lesson’s demo deliberately runs the fault path and the invalidate path at matched, maximal frequency because a driver that only gets exercised against occasional, well-spaced invalidations in CI will never hit the retry path enough times to prove it works. If your test suite’s refault rate is near zero, that’s a sign your test isn’t contending hard enough to be meaningful, not a sign your driver is fast.

• Don’t confuse “the sanitizer didn’t fire” with “the code is correct.” As this lesson’s own before/after comparison shows, the same race reproduces at wildly different rates depending on instrumentation and scheduling noise — a clean run proves nothing about a race window that a different kernel, a different core count, or a different sanitizer might expose.

15. Summary

HMM’s contribution isn’t a clever allocator or a scheduling trick — it’s a correct answer to a hard concurrency problem: how does a driver safely act on a page-table snapshot that the CPU is free to invalidate at any moment, without either serializing every access behind a lock or accepting the possibility of a device DMA engine writing into freed memory. The mechanism is a sequence-counter retry check plus page pinning, and this lesson’s demo makes the failure mode and the fix both fully reproducible: omit the retry check and TSan/ASan will show you a real race and a real use-after-free within seconds; add it back correctly, and the identical workload passes a full sanitizer and Valgrind gauntlet clean. That gap — between “compiles and usually works” and “provably correct under adversarial scheduling” — is exactly the gap HMM’s API contract exists to close for every driver author who has to implement it.