Lab 01 — Caching GPU Allocator (Rust)

Phase: 05 — Runtime Systems | Difficulty: ⭐⭐⭐⭐☆ | Time: 5–8 hours Language: Rust (safe, stdlib only) | Hardware: none — a CPU-side model of the GPU original

Concept primer: ../WARMUP.md Ch. 2–5.

Run

cargo test                              # 12 correctness tests
cargo run --release --bin fragdemo      # the two teaching demos

0. The mission

Model the two layers every framework stacks on the driver:

  • Device = the "driver": a fixed address space, first-fit allocation, holes coalesce on free, and it counts free calls (each is a device sync in real CUDA). This is cudaMalloc/cudaFree.
  • CachingAllocator = the framework layer (the PyTorch design): size rounding into classes, binned free lists, block splitting, stream tagging, and flush-and-retry on OOM. This is what actually serves your tensors.

The tests prove the semantics; fragdemo proves why the design exists.

1. What fragdemo shows

== Demo 1: external fragmentation is a shape problem ==
  interleaved (mixed)    free=128 MB  largest hole=  8 MB  16MB alloc: FAILS (fragmented!)
  grouped (disciplined)  free=128 MB  largest hole=128 MB  16MB alloc: OK

== Demo 2: the caching allocator's reason to exist ==
  raw driver :  5000 mallocs,  5000 frees (each free = device sync)
  caching    :     1 device mallocs,  4999 bin hits, 0 device frees

Demo 1 is the day-3 OOM in 30 deterministic lines: the same 128 MB is free in both rows; only the layout differs. Mixed shapes strand free memory in 8 MB holes between permanent blocks, so a 16 MB request fails at 50% utilization — and since HBM can't be compacted (live pointers), the only fix in production is a restart. Shape discipline (grouping lifetimes) keeps one big contiguous hole. This is exactly why Phase 07 pages the KV cache.

Demo 2 is why the caching layer exists at all: 5000 alloc/free cycles become 5000 synchronizing driver frees raw, versus 1 device malloc + 4999 ~100 ns bin hits cached. The driver never sees a steady-state free.

2. Reading order (src/lib.rs)

  1. round_size + the Device model — the constraints (first-fit, coalescing, counted frees).
  2. CachingAllocator::alloc — the four-step path: bin hit → split larger → driver malloc → flush-and-retry. Each step maps to a WARMUP Ch. 4 bullet.
  3. free (goes to a bin, never the driver) and flush_cache (the empty_cache() semantics, including why a split block isn't returnable).
  4. The stream-tag logic in alloc — same-stream reuse is free; cross-stream reuse counts an event edge (cross_stream_syncs), the Ch. 7 safety rule.

Then trace a 1.2 MB allocation by hand and check yourself against test_split_accounting.

3. The tests, and what each proves

TestThe lesson it locks in
test_roundingsize classes (Ch. 4)
test_device_first_fit_and_coalescethe driver model's behavior
test_bin_reuse_fast_pathreuse never calls the driver (free_calls == 0)
test_split_accountingsplit math: 4MB → 1MB out + 3MB binned
test_cross_stream_sync_countedcross-stream reuse needs an event edge (Ch. 7)
test_same_stream_preferredsame-stream blocks chosen first (no edge)
test_reserved_vs_allocatedthe nvidia-smi-vs-framework gap (Ch. 4, Q2)
test_flush_and_retry_on_oomthe pre-OOM stall (Ch. 4)
test_split_block_not_returnablecan't return half a device segment

4. Extensions

  1. Coalescing (hard): when two adjacent free blocks exist, merge them back into one larger class — true defragmentation within the cache. Add a test that an alloc/free/alloc pattern which currently fails now succeeds.
  2. empty_cache policy: return only segments idle for N cycles; measure the reserved-memory sawtooth.
  3. Trace replayer: parse a CSV of (op, size, stream) lines and replay it, emitting reserved/allocated/largest-block per step — turn the allocator into a measurable system you can feed real PyTorch traces.
  4. Best-fit vs first-fit in the bins: measure fragmentation on the demo-1 workload; quantify the classic tradeoff.
  5. Prometheus metrics: export reserved/allocated/cached/largest-block — the Phase 10 fragmentation dashboard, fed from here.

5. Common pitfalls

  1. Returning a split remainder to the driver — you can't free half a segment (test_split_block_not_returnable enforces it; PyTorch has the same limit).
  2. Forgetting the stream tag on reuse — silent cross-stream UAF in the real thing; here it's a counted edge so you can see it.
  3. Treating reserved as a leak — it's cache; the leak signal is allocated trending up, not reserved.

6. What this lab proves about you

You can implement, test, and reason about the allocator at the heart of every ML framework — and translate its behavior into operational language (reserved vs allocated, fragmentation early-warning, the restart confession). Reading PyTorch's real CUDACachingAllocator.cpp is now a code-review, not an expedition.