← Back to Table of Contents

Chapter 10 — Memory Management

Chapter 9 catalogued the hardware: SRAM, DRAM, HBM, the hierarchy of speed and capacity. But having fast memory is only half the problem. The other half is managing it — making sure the right data is in the right level at the right time, at the right address, visible to the right processor. If you’ve ever wondered why tensor.contiguous() matters, why pin_memory=True in a DataLoader speeds things up, why NCHW and NHWC have different performance, or why NUMA topology affects multi-GPU training even though you haven’t written a single line of C — this chapter has the answers.

Memory management is the software/hardware interface that makes the memory hierarchy invisible to programs while ensuring it performs well. It operates at several layers simultaneously: virtual memory (translating addresses, protecting processes from each other), caching (transparently migrating hot data toward the CPU), and coherence (ensuring multiple cores and GPUs see a consistent view of shared state). Each layer has failure modes that appear as mysterious slowdowns in ML workloads.

The one-sentence version: Memory management is the full hardware+OS stack that translates program addresses to physical locations, caches hot data automatically, and maintains consistency across cores — and its failure modes (TLB misses, cache thrash, coherence traffic, NUMA penalties) are real performance hazards in large-model training and inference.

Why We Need Memory Management

Consider a server with 4 GPUs, 96 CPU cores, 1.5 TB of DDR5, and a 70B-parameter model. Several problems arise immediately without memory management:

1. Addresses would conflict. With multiple processes and the OS running simultaneously, two programs might try to use the same physical address. A bug in one program would corrupt another.
2. Programs cannot always fit in physical memory. A model larger than VRAM must be partitioned or swapped.
3. Without caching, every memory access would pay full DRAM latency (~70 ns) — L1 cache at 1–4 ns is a 17–70× speedup.
4. Multiple cores write the same data. Without coherence, Core 0 might read a stale value after Core 1 updated it.
5. Memory on different NUMA nodes has different latency. Ignoring this causes 10–40% throughput losses in multi-socket servers.

Each of these problems has a hardware/OS solution. We cover them in order.

Virtual Memory

The Problem: Physical Address Chaos

Early computers ran one program at a time; the program controlled all physical memory. Modern systems run hundreds of processes simultaneously. Each process believes it owns the entire address space — in 64-bit systems, a 48-bit (256 TB) virtual address space.

Virtual memory is the abstraction that makes this possible. Each process has its own virtual address space: addresses 0 through $2^{48}-1$. The hardware + OS transparently maps these virtual addresses to physical addresses (the actual DRAM locations) at runtime.

Paging and Page Tables

The virtual address space is divided into fixed-size pages (typically 4 KB, but also 2 MB or 1 GB — “huge pages”). Physical memory is divided into page frames of the same size. A page table (a data structure in DRAM, maintained by the OS) maps each virtual page number to a physical frame number.

For a 48-bit virtual address with 4 KB pages, the layout is:

1
2
3
4
5
6
7
Virtual Address (48-bit):
┌─────────────────────────────────────────────────────┐
│  [47:39]  │  [38:30]  │  [29:21]  │  [20:12]  │  [11:0]  │
│  PGD idx  │  PUD idx  │  PMD idx  │  PT idx   │  offset  │
│  (9 bits) │  (9 bits) │  (9 bits) │  (9 bits) │ (12 bits)│
└─────────────────────────────────────────────────────┘
     L4 table    L3 table    L2 table    L1 table    offset in page

This is a 4-level page table (x86-64 uses exactly this). The page walk traverses 4 levels of tables stored in physical memory, each at a different address — potentially 4 DRAM accesses, each ~70 ns, totaling ~280 ns just to translate one address. This would be catastrophic without caching.

The MMU and TLB

The Memory Management Unit (MMU) is the hardware component inside the CPU/GPU that performs address translation. It contains:

• TLB (Translation Lookaside Buffer): A small, fully-associative cache of recent virtual→physical translations. On a TLB hit (the translation is cached), translation costs 1–3 cycles. On a TLB miss, the MMU performs the full page-table walk — 4 memory accesses, ~100–300 cycles.

A modern CPU’s L2 STLB with 4096 entries covers only 16 MB of 4 KB pages. A 70B model in BF16 is 140 GB — accessing it with 4 KB pages would cause millions of TLB misses per second if not addressed (pun intended).

Huge Pages to the Rescue

Huge pages (2 MB on x86; 1 GB with 1GB pages) map much larger regions with a single TLB entry. Instead of 140 GB / 4 KB = 36 million TLB entries, 140 GB / 2 MB = 70,000 entries — fitting comfortably in the L2 STLB.

PyTorch uses mmap with MAP_HUGETLB for large weight tensors when huge pages are available. Linux Transparent Huge Pages (THP) attempts to do this automatically. For large model inference servers, explicitly enabling 1GB huge pages for VRAM-sized allocations can reduce TLB pressure and improve throughput by 5–15%.

1
2
3
4
5
6
# Enable 1GB huge pages at boot (Linux)
# /etc/default/grub: add "hugepagesz=1G hugepages=128" to GRUB_CMDLINE_LINUX
# After reboot:
grep HugePages /proc/meminfo
# HugePages_Total:     128
# HugePages_Free:       64

Caches in Depth

Cache Lines

The CPU does not load a single byte from DRAM on a cache miss — it loads a cache line: a contiguous block of 64 bytes (8 × float64, 16 × float32, 32 × fp16/bf16, 64 × int8). This exploits spatial locality: if you access address $A$, you’ll likely soon access $A+4$, $A+8$, etc. Loading a full line amortizes the DRAM access cost over 64 bytes.

Consequence for ML: accessing a tensor row-by-row (column-major for a row-major layout) loads 64 bytes, uses 4–8 of them, then discards the rest — catastrophically wasteful. Accessing in the tensor’s storage order (row-major for the default C-contiguous PyTorch layout) uses all 64 bytes before eviction.

Associativity

A cache must decide where to store each cache line. Three main policies:

A modern L3 cache might be 16-way set associative with 512 MB: 512 MB / 64 B line / 16 ways = 512K sets, each holding 16 cache lines. The index field of an address selects one of the 512K sets; the tag field identifies which specific line is stored in each way.

Replacement Policies

When all ways in a set are full and a new line must be loaded, the cache evicts one. Common policies:

• LRU (Least Recently Used): Evict the line accessed longest ago. Near-optimal for temporal locality. Expensive to implement exactly for large caches — approximated with pseudo-LRU trees.
• PLRU (Pseudo-LRU): Uses a binary tree to approximate LRU. Hardware-efficient; used in most L2/L3 caches.
• RRIP (Re-Reference Interval Prediction): Smarter; predicts how soon each line will be reused. Used in Intel’s L3 caches.

For ML: if your model weights are larger than the L3 cache (almost always true), the cache behaves as a streaming buffer — each line is used once and evicted. LRU/PLRU makes no difference. What matters is spatial locality (cache-line fill efficiency) and prefetching.

Write Policies

• Write-through: Every store writes to both the cache and DRAM immediately. Cache always matches DRAM; consistent, but every write incurs a DRAM write.
• Write-back: Stores update only the cache; the line is marked dirty. DRAM is updated only when the dirty line is evicted. Lower DRAM write traffic, but the OS/other devices may see stale data.
• Write-allocate: On a write miss, load the line into cache first, then write. Pairs with write-back; standard for CPUs.

L1/L2/L3 caches on modern CPUs are write-back with write-allocate. GPU L1/L2 are typically write-back as well.

The Three C’s of Cache Misses

Every cache miss falls into one of three categories:

Spatial and Temporal Locality

• Temporal locality: A recently accessed datum will likely be accessed again soon. Caches exploit this by keeping recently-used lines.
• Spatial locality: If address $A$ is accessed, $A+4$, $A+8$, … are likely to be accessed soon. Cache lines (64 bytes) exploit this by loading a block.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import numpy as np
import time

N = 4096  # 4096×4096 float32 matrix = 64 MB (>> typical L3 cache)
A = np.random.rand(N, N).astype(np.float32)  # row-major (C order)
B = np.random.rand(N, N).astype(np.float32)

# --- Cache-FRIENDLY: row-major sequential scan (spatial locality) ---
t0 = time.perf_counter()
total = 0.0
for i in range(N):
    total += A[i, :].sum()   # access full row: sequential in memory
t1 = time.perf_counter()
print(f"Row scan (friendly): {(t1-t0)*1e3:.1f} ms")

# --- Cache-UNFRIENDLY: column-major scan (stride = 4096 × 4 bytes = 16 KB) ---
t0 = time.perf_counter()
total = 0.0
for j in range(N):
    total += A[:, j].sum()   # access full column: stride 16 KB — one element per cache line
t1 = time.perf_counter()
print(f"Col scan (unfriendly): {(t1-t0)*1e3:.1f} ms")
# Expected: column scan is 3–10× slower depending on CPU cache size

Cache-Friendly Matmul: Tiling / Blocking

Naïve matmul $C = AB$ with $A \in \mathbb{R}^{M \times K}$, $B \in \mathbb{R}^{K \times N}$:

1
2
3
4
5
6
7
8
9
10
# SLOW: naive triple loop — accesses B column-by-column (poor locality for row-major B)
def matmul_naive(A, B):
    M, K = A.shape
    K2, N = B.shape
    C = np.zeros((M, N), dtype=np.float32)
    for i in range(M):
        for j in range(N):
            for k in range(K):
                C[i, j] += A[i, k] * B[k, j]   # B[k,j]: stride-N access
    return C

Blocked (tiled) matmul partitions the matrices into tiles that fit in L2/L3 cache, dramatically improving reuse:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def matmul_blocked(A, B, block=64):
    """
    Block size 64: tile = 64×64 float32 = 16 KB — fits in L2 cache (256 KB).
    Each tile of A is reused N/block times; each tile of B is reused M/block times.
    """
    M, K = A.shape
    K2, N = B.shape
    C = np.zeros((M, N), dtype=A.dtype)
    for i0 in range(0, M, block):
        for j0 in range(0, N, block):
            for k0 in range(0, K, block):
                # These three slices are all in cache simultaneously
                Atile = A[i0:i0+block, k0:k0+block]   # block × block = 16 KB
                Btile = B[k0:k0+block, j0:j0+block]   # block × block = 16 KB
                C[i0:i0+block, j0:j0+block] += Atile @ Btile
    return C

This is conceptually what cuBLAS, OneDNN, and Flash Attention do — tile computations to maximize cache reuse and minimize DRAM traffic. On GPU, the “cache” is shared memory (32–64 KB per SM), and tiling is done explicitly in the CUDA kernel.

Cache Coherence

The Problem

A modern server might have 96 CPU cores, each with its own L1 and L2 cache, all sharing the same physical DRAM. Suppose:

1. Core 0 loads address $X$ into its L1 cache (value = 5).
2. Core 1 loads address $X$ into its L1 cache (value = 5).
3. Core 0 writes $X = 7$ — updates its L1 (dirty line).
4. Core 1 reads $X$ — gets 5 from its own L1 (stale!).

Without coherence, parallel programs are undefined. Hardware must enforce: if Core 1 reads $X$ after Core 0 wrote $X$, it gets the new value.

The MESI Protocol

MESI (Modified–Exclusive–Shared–Invalid) is the most widely used cache coherence protocol. Each cache line has a state:

Snooping vs Directory Coherence

Snooping: Every cache monitors (snoops) a shared bus or interconnect. When a core writes, it broadcasts an invalidate message; every cache holding that line transitions to Invalid. Simple; used in small (<16 core) systems.

Directory coherence: A centralized (or distributed) directory tracks which caches hold each physical line. On a write, the directory sends invalidates only to the caches that have the line — much more scalable. Used in large multi-socket servers (EPYC, Xeon) with hundreds of cores.

False Sharing

A pernicious coherence problem: two threads on different cores each modify different fields that happen to occupy the same 64-byte cache line. Every write from one core triggers an invalidate to the other, forcing a cache miss on the next access — despite no logical data sharing.

1
2
3
4
5
6
7
8
import numpy as np

# Array where adjacent elements are in the same cache line (64 bytes = 16 float32s)
# Thread 0 writes index 0; Thread 1 writes index 1 → same cache line → false sharing
data = np.zeros(1024, dtype=np.float32)

# Mitigation: pad each "thread-private" element to a full cache line (64 bytes = 16 floats)
# data_padded[thread_id * 16] avoids false sharing between adjacent threads

In PyTorch: gradient accumulation buffers, counters, and locks can suffer from false sharing. Libraries like torch.distributed and deepspeed are carefully designed to avoid this.

Memory Consistency Models

Coherence says “every core eventually sees a write to address $X$.” Consistency says “in what order do writes to different addresses become visible?”

• Sequential Consistency (SC): The ideal model. All operations appear to happen in some global total order. Easy to reason about; expensive to implement.
• Total Store Order (TSO): x86’s model. Stores are buffered (store buffer) before becoming globally visible; loads from the same core bypass the store buffer. Relaxes SC while remaining easy to reason about.
• Weak/Release Consistency: ARM, POWER, RISC-V. Stores and loads can be reordered aggressively by hardware for performance; memory barriers (DMB/DSB on ARM, FENCE on RISC-V) enforce ordering at specific points.

For ML: PyTorch’s torch.distributed operations, CUDA’s __threadfence() and __threadfence_system(), and Python’s threading.Lock() all insert the appropriate barriers. You generally don’t need to manage this manually — but understanding it explains why GPU kernels need __syncthreads() after shared memory writes, and why naïve atomicAdd accesses are slower on relaxed-consistency architectures.

NUMA — Non-Uniform Memory Access

The Problem with Multi-Socket Systems

A server with two CPU sockets (e.g., dual-socket AMD EPYC with 96 cores per socket) has two NUMA nodes. Each NUMA node has its own local DRAM controller and memory banks, connected by an inter-socket interconnect (AMD’s Infinity Fabric, Intel’s UPI).

Remote NUMA access is 1.5–3× slower than local. If CPU processes that handle a specific GPU are scheduled on the non-local NUMA node, every data transfer from host to GPU pays the inter-socket overhead.

NUMA in Multi-GPU Systems

In a DGX-H100 with 8 GPUs and 2 CPU sockets:

1
2
3
4
5
6
7
NUMA Node 0 (Socket 0, 48 cores, 768 GB DDR5):
  ├── GPU 0, 1, 2, 3 (via PCIe 5.0 / NVLink)
  └── NVLink switch 0

NUMA Node 1 (Socket 1, 48 cores, 768 GB DDR5):
  ├── GPU 4, 5, 6, 7 (via PCIe 5.0 / NVLink)
  └── NVLink switch 1

For data-parallel training: the DataLoader workers pinning memory and copying to GPU should be NUMA-local to their GPU. Using os.sched_setaffinity() or numactl --cpunodebind=0 --membind=0 pins processes to NUMA node 0 for GPU 0–3, avoiding remote DRAM.

1
2
3
4
5
6
7
8
9
10
11
# PyTorch DataLoader: pin_memory=True allocates host tensors in pinned (page-locked)
# NUMA-local memory for fast DMA to the GPU.
from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,
    pin_memory=True,          # allocates in pinned memory → DMA-friendly
    pin_memory_device="cuda:0",  # pins to NUMA node local to GPU 0 (PyTorch ≥2.1)
)

Pinned memory (page-locked, non-pageable) lets the GPU DMA engine copy directly from host RAM to GPU VRAM without a CPU intermediate copy — critical for avoiding PCIe transfer overhead when streaming training data.

Tensor Layout, Contiguity, and Memory Management

Why tensor.contiguous() Exists

PyTorch tensors have a storage (a flat 1D array in physical memory) and metadata (shape, strides, offset) that describe how the logical N-D tensor maps onto the 1D storage.

1
2
3
4
5
6
7
8
import torch

x = torch.randn(3, 4, 5)  # shape [3,4,5], strides (20, 5, 1) — contiguous, row-major
y = x.transpose(0, 2)      # shape [5,4,3], strides (1, 5, 20) — same storage, different view
print(y.is_contiguous())   # False — stride order doesn't match dim order

z = y.contiguous()         # Allocates new storage, copies data in new stride order
print(z.is_contiguous())   # True

A non-contiguous tensor accesses memory with non-unit strides — skipping over memory locations. For the transposed y with stride (1, 5, 20), iterating over the last dimension accesses elements 20 apart in memory — one element per cache line (at float32: 20 × 4 = 80 bytes > 64 byte line). Calling .contiguous() creates a new allocation with sequential access, making it cache-friendly.

NCHW vs NHWC: Cache Implications

Convolutional layers operate on 4D tensors with shape [batch, channels, height, width]. Two common layouts:

1
2
3
4
5
6
7
8
9
10
11
import torch
import time

# Create NCHW and NHWC tensors
N, C, H, W = 32, 256, 64, 64
x_nchw = torch.randn(N, C, H, W, device='cpu')
x_nhwc = x_nchw.to(memory_format=torch.channels_last)  # PyTorch channels_last

# On GPU, cuDNN will auto-convert if it detects channels_last input
# Pass memory_format=torch.channels_last to model layers for best GPU conv perf:
conv = torch.nn.Conv2d(256, 256, 3, padding=1).to(memory_format=torch.channels_last)

Why This Matters for Model Optimization

The Full Picture: One Forward Pass

A single transformer forward pass on a multi-GPU node involves all the memory management concepts in this chapter:

1
2
3
4
5
6
7
1. DataLoader workers (pinned memory, NUMA-local) → H2D DMA (virtual→physical, TLB walk)
2. Embedding lookup → L2/L3 cache misses (compulsory), DRAM/HBM access
3. Attention matmul → tiled kernel, tiles in GPU shared memory (SRAM), HBM for K/V/Q
4. FFN matmul → cache-friendly NHWC layout; weight tensor too large for L3 (capacity misses)
5. Gradient AllReduce → inter-GPU via NVLink; CPU gradient buffers under MESI coherence
6. Optimizer step → reads/writes parameter + optimizer state tensors; NUMA-aware pinning
7. Checkpoint write → DMA to pinned host memory → async copy to NVMe SSD

Each step is a negotiation with the memory hierarchy. Understanding these layers is why ML systems engineers spend so much time on memory layout, pinning, affinity, and tiling.

Key Takeaways

• Virtual memory gives each process its own address space; the MMU translates virtual → physical addresses using page tables. The TLB caches recent translations (~1 cycle hit vs ~100–300 cycle miss on a page-table walk).
• Huge pages (2 MB, 1 GB) dramatically reduce TLB pressure for large allocations — important for model weights and KV caches exceeding tens of gigabytes.
• Cache lines are 64 bytes; accessing tensors in non-contiguous or column-major order wastes most of each cache line and slows performance 3–10× — the hardware basis for tensor.contiguous() and NHWC layout.
• The three C’s — compulsory, capacity, and conflict misses — classify why a cache miss occurs; capacity misses dominate for model weights larger than L3.
• Tiling/blocking is the universal technique for fitting working sets into fast SRAM (CPU L2 or GPU shared memory) to maximize reuse and minimize DRAM traffic — the core idea behind cuBLAS, Triton, and Flash Attention.
• MESI cache coherence ensures multiple cores see a consistent memory view via Modified/Exclusive/Shared/Invalid states; false sharing (two threads modifying different fields in the same cache line) can silently destroy parallel performance.
• NUMA topology means that remote-socket DRAM accesses are 1.5–3× slower; GPU DataLoaders and DMA transfers should be NUMA-local to their target GPU. pin_memory=True enables zero-copy DMA and is almost always a free win.

Next: Chapter 11 — Memory Wall & Bandwidth, where we quantify exactly how badly memory limits compute using the roofline model and arithmetic intensity — the tools that explain why almost every large-model inference workload is memory-bound.

← Back to Table of Contents