Memory Allocation and Virtual Memory

This file is the conceptual ancestor of PagedAttention (Section V.10). Read it carefully; when you meet paged KV cache you will recognize every idea here.

1. What is it?#

Virtual memory gives every process the illusion of a large, contiguous, private address space. The hardware (MMU) and the OS translate virtual addresses to physical ones, page by page (usually 4 KB).

Allocation is the software layer on top: malloc, jemalloc, PyTorch’s caching allocator — carving up big regions from the OS into the sizes programs actually request.

Process view (virtual)         Reality (physical)
┌──────────────┐               ┌──────────────┐
│ 0x0000...    │──────┐        │ frame 42     │
│ contiguous!  │      ├───────►│ frame 7      │  scattered
│ 0xFFFF...    │──────┘        │ frame 1893   │
└──────────────┘               └──────────────┘
                 page table

2. Why does it exist?#

Three problems, one mechanism:

1. Fragmentation. Without indirection, allocating and freeing variable-size blocks leaves unusable gaps. With paging, physical memory need not be contiguous, so external fragmentation disappears.
2. Isolation. Each process’s page table maps only its own frames.
3. Overcommit. You can promise more memory than exists and back it lazily.

Every one of those three motivations reappears verbatim in KV cache management. The KV cache suffered exactly problem 1 — variable-length sequences in a contiguous allocator wasted 60-80% of memory — and PagedAttention fixed it with exactly this mechanism.

3. Simple analogy#

A library’s card catalogue. Book “Chapter 7” is logically after “Chapter 6”, but physically they can be on different floors. The catalogue (page table) maps logical to physical. You can shelve a new book in any free slot, anywhere — no need for a contiguous run of empty shelf.

Without the catalogue, you would need every multi-volume set stored contiguously, and after a year of additions and removals you’d have plenty of free shelf space but nowhere to put a 12-volume encyclopedia. That is external fragmentation, and it is precisely what killed pre-PagedAttention KV allocators.

4. Tiny example#

Fragmentation, made concrete:

Memory: 16 units.  Allocate A(4), B(4), C(4), D(4)  → full
        [AAAA][BBBB][CCCC][DDDD]

Free B and D:
        [AAAA][....][CCCC][....]        8 units free

Now allocate E(6):
        FAILS — 8 units free, but the largest contiguous run is 4.

This is external fragmentation. Now the paged version:

Page size 2. Memory: 8 pages.
A→pages{0,1}  B→{2,3}  C→{4,5}  D→{6,7}

Free B and D:  free pages = {2,3,6,7}

Allocate E(6) → needs 3 pages → gets {2,3,6}
        SUCCEEDS. Physical discontiguity is invisible to E.

That is the entire idea of PagedAttention, four sections early. A sequence’s KV cache is stored in fixed-size blocks (typically 16 tokens each) that need not be contiguous, with a block table playing the role of the page table.

5. Technical explanation#

Address translation and the TLB#

Translating every access through a multi-level page table would be ruinous (4-5 memory accesses per access). The TLB (Translation Lookaside Buffer) caches translations:

L1 dTLB:  ~64 entries    →  64 × 4 KB = 256 KB of coverage
L2 TLB:   ~1500-2000     →  ~8 MB of coverage

Notice the problem: a 14 GB model’s weights need 3.5 million 4 KB pages. Your TLB covers 0.05% of that. Every weight access risks a TLB miss (a page walk: ~20-100 cycles).

Huge pages fix it:

4 KB pages:   14 GB needs 3,500,000 entries
2 MB pages:   14 GB needs     7,000 entries   ← now L2 TLB covers 4 GB
1 GB pages:   14 GB needs        14 entries

Enable transparent huge pages, or allocate explicitly:

cat /sys/kernel/mm/transparent_hugepage/enabled   # [always] madvise never
# For CPU inference of large models, 'always' or explicit madvise can give 5-20%

This matters for CPU inference and for large pinned host buffers. GPUs have their own MMU and NVIDIA’s driver already uses large pages for device memory.

Overcommit and the OOM killer#

Linux by default lets you malloc more than exists, because most programs don’t touch it all. When they do:

cat /proc/sys/vm/overcommit_memory   # 0=heuristic 1=always 2=strict

If physical memory runs out, the kernel’s OOM killer picks a victim by oom_score. In a container, exceeding the memory cgroup limit triggers a cgroup OOM kill — your process dies with no Python traceback, just exit code 137. If your inference server “randomly disappears,” check dmesg | grep -i oom before anything else.

Allocators#

malloc (glibc)      general purpose; arenas per thread; can fragment badly
jemalloc            better fragmentation behavior; used by many servers
tcmalloc            fast thread-caching; good for many small allocations
PyTorch CUDA caching allocator   never returns memory to the driver; sub-allocates from
                                 large cudaMalloc'd segments

PyTorch’s caching allocator is the one you’ll fight with. Key facts:

• cudaMalloc is slow (~100 µs) and synchronizing, so PyTorch grabs big segments and reuses.
• Freed tensors return to PyTorch’s pool, not to the GPU. nvidia-smi still shows the memory used.
• Fragmentation within the pool causes “CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 79.15 GiB total capacity; 60.10 GiB already allocated; 1.20 GiB free …)” — note the gap between “free” and “total minus allocated”: that gap is fragmentation.
• Tunable: PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True reduces fragmentation for varying shapes — genuinely useful for inference with variable sequence lengths.

torch.cuda.memory_allocated()      # bytes in live tensors
torch.cuda.memory_reserved()       # bytes held by the allocator from the driver
torch.cuda.empty_cache()           # return unused segments to the driver
print(torch.cuda.memory_summary()) # detailed breakdown — read this when debugging OOM

Pinned (page-locked) memory#

Normal host memory can be swapped out or moved, so the DMA engine cannot safely read it. CUDA therefore copies through a staging buffer, halving effective bandwidth. Pinned memory is locked in place and can be DMA’d directly:

buf = torch.empty(size, pin_memory=True)     # ~2x faster H2D
tensor.to('cuda', non_blocking=True)         # only truly async from pinned memory

Costs: pinning is slow to allocate, and pinning too much starves the OS. Use a pool of reusable pinned buffers, not per-request allocation.

6. Under the hood#

# Process memory map
cat /proc/<pid>/status | grep -E 'VmRSS|VmSize|VmSwap|HugetlbPages'
pmap -x <pid> | tail -3

# Page faults
perf stat -e page-faults,minor-faults,major-faults ./prog
# major faults = went to disk. Any major faults in steady state = trouble.

# TLB misses
perf stat -e dTLB-load-misses,dtlb_load_misses.walk_completed ./prog

# Huge page usage
grep -i huge /proc/meminfo

VmRSS (resident) is what you actually occupy; VmSize (virtual) includes reservations and is usually alarming and meaningless — a process mapping a 140 GB model file has a huge VmSize and a modest RSS.

7. Performance implications#

8. Production implications#

• Disable swap on inference nodes. A swapped-out page in the hot path is a multi-millisecond stall. Kubernetes disables it by default; verify.
• Preallocate the KV cache pool at startup. vLLM does this (gpu_memory_utilization), claiming a fixed fraction of the GPU up front so that steady-state allocation never fails. This is the production answer to fragmentation: allocate once, manage yourself.
• Set PYTORCH_CUDA_ALLOC_CONF deliberately for variable-shape workloads.
• Use a pinned-buffer pool for any host↔device streaming.
• Monitor memory_reserved vs memory_allocated. A growing gap is fragmentation.
• Exit code 137 means OOM-killed. Put that in your runbook.

9. Common mistakes#

Interpreting nvidia-smi memory as “tensors in use.” It shows the allocator’s reservation.

Calling empty_cache() in the hot path. It synchronizes and forces future cudaMallocs. Use it once after model load or between phases, never per request.

Allocating pinned memory per request. cudaHostAlloc is slow and serializing. Pool it.

Assuming OOM means “not enough memory.” Often it means fragmentation. Read memory_summary().

Letting variable batch shapes fragment the pool. Bucket your shapes, or use expandable_segments.

Ignoring the container memory limit. The GPU has 80 GB but your pod may be limited to 32 GB of host RAM, and weight loading needs host memory too.

10. Hands-on exercise#

A. Reproduce fragmentation. In PyTorch on GPU, allocate many tensors of random sizes, free every other one, then try to allocate a large one. Trigger an OOM with plenty of free memory. Print memory_summary(). Then re-run with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True and compare.

B. Measure pinned vs pageable. Time 1 GB H2D transfers from pageable and pinned memory. Report both bandwidths. Record in numbers.md.

C. Huge pages. Run a CPU-side memory-intensive benchmark with THP always vs never. Measure dTLB-load-misses and wall time for each.

D. Build the paged allocator. Implement, in ~80 lines of Go, a block allocator: fixed 2-unit blocks, a free list, allocate/free by sequence id, and a block table per sequence. Show it succeeding on the E(6) case from section 4 where a contiguous allocator fails. Keep this code — Project 09 extends it into a real KV cache manager.

11. Interview questions#

1. What problem does virtual memory solve, and how does PagedAttention borrow it?
2. What is a TLB miss and why do huge pages help large models?
3. Why does nvidia-smi show more memory than torch.cuda.memory_allocated()?
4. You get a CUDA OOM but the error says 5 GB free. Explain and give two fixes.
5. What is pinned memory, when do you need it, and what does it cost?
6. Your container dies with exit code 137 and no traceback. What happened?
7. Explain internal vs external fragmentation with a KV cache example of each.

12. Further reading#

• [FUNDAMENTAL] Operating Systems: Three Easy Pieces, virtualization section — free online
• [REFERENCE] PyTorch CUDA memory management docs; PYTORCH_CUDA_ALLOC_CONF
• [ESTABLISHED] Kwon et al., PagedAttention (SOSP 2023) §4 — read after this file
• Next: 07 — Storage and model loading