Warmup Guide — Runtime Systems: Allocators, Streams, Async Execution

Zero-to-expert primer for Phase 05. Assumes Phases 02 and 04 (the CUDA memory API and the driver beneath it). No allocator or concurrency background assumed. By the end you can design the runtime layer of an inference platform — and debug the ones you didn't design.

Table of Contents

Chapter 1: The Runtime Layer — What Sits Between Driver and Engine

The stack so far: silicon (Phase 01) ← kernels (02) ← compiler (03) ← driver (04). This phase is the layer that every framework and serving engine builds on top of the driver and underneath the model code:

model / engine code (vLLM, PyTorch, TensorRT)        Phase 07
─────────────────────────────────────────────
RUNTIME: memory pools, stream scheduling,            ← THIS PHASE
         graphs, staging pipelines
─────────────────────────────────────────────
CUDA runtime/driver API                              Phase 02/04

Why it exists at all: the driver's primitives (cudaMalloc, raw streams) have costs and semantics that are unacceptable to call in a hot loop. The runtime layer turns them into amortized, policy-bearing services. The two services that matter most — and that you'll build — are the allocator (Lab 01) and the execution scheduler (Lab 02).

Production significance: when this layer is wrong you get the worst class of incidents — not crashes, but slow decay: fragmentation OOM on day 3, throughput sag from false serialization, once-a-week corruption from a cross-stream lifetime bug. Nothing in a stack trace points here; you find these only if you know the layer exists.

Chapter 2: Why cudaMalloc Is Slow and Fragmenting

Three facts about the raw allocation API, each with a design consequence:

  1. cudaFree synchronizes the device (historically; and cudaMalloc can too). The driver must be sure no in-flight kernel still uses the region — lacking fine-grained tracking, it waits for everything. One stray cudaFree per request = your entire pipeline gains a global barrier. → Consequence: never call it in steady state; allocate from a pool.
  2. It's a driver round-trip (ioctl, page-table updates — Phase 04 Ch. 4): tens of µs vs ~100 ns for a pool hit. At thousands of tensor allocations per step, that's the difference between negligible and dominant. → Consequence: caching — keep freed blocks and reuse them.
  3. There is no paging and no compaction. Host malloc lives on virtual memory: fragmented physical pages don't matter because the VM layer remaps. GPU allocations are (effectively) physically contiguous reservations in a fixed-size HBM; nothing remaps behind your back, and you cannot move a live block (kernels hold raw pointers). → Consequence: fragmentation is permanent damage within a process lifetime; the allocator's only weapons are placement policy and shape discipline (Ch. 5).

(cudaMallocAsync, CUDA 11.2+, moved a pool into the driver — same design, different owner; the VMM API (cuMemCreate + address reservation) added remapping tricks that "expandable segments" exploit — Ch. 4.)

Chapter 3: Allocator Zoology — Bump, Buddy, Slab, Bins

The four designs every systems engineer should hold (Lab 01 implements the fourth):

  • Bump allocator: a pointer that only moves forward; free is a no-op (or whole-arena reset). O(1), zero fragmentation within an arena, but only fits phase-structured lifetimes. GPU relevance: activation memory in inference is exactly phase-structured — per-step arenas are a legitimate design (and what "static planning" compilers like TensorRT effectively do).
  • Buddy allocator: powers-of-two blocks, split on demand, coalesce with your "buddy" on free. O(log n), bounded external fragmentation, up to ~50% internal waste. The classic kernel-page-allocator design.
  • Slab allocator: per-size-class object caches. Brilliant for fixed-size objects — GPU relevance: KV-cache blocks (Phase 07) are fixed-size by design partly so a slab-style pool can manage them trivially.
  • Size-binned free lists with splitting (the caching allocator): free blocks indexed by size class; allocation finds the smallest fitting block, splits the remainder back into the lists. This is malloc's general shape and PyTorch's choice — flexible sizes, good reuse, but external fragmentation is back (Ch. 5).

The meta-lesson: allocator design = exploiting what you know about lifetimes and sizes. The more structure you can impose on the workload (fixed block sizes, phase lifetimes, stream affinity), the better the allocator can be. This is why Phase 07's PagedAttention is best understood as an allocator redesign, not an attention algorithm.

Chapter 4: The Caching Allocator — PyTorch's Design, Dissected

The reference design (Lab 01 is a faithful Rust model):

  • Two regimes by size: small (<1 MB) and large allocations get separate pools/bins (mixing tiny and huge requests is a fragmentation accelerant).
  • Rounding: sizes round up to multiples (512 B; larger granularity for big blocks) — bounded internal waste purchases massive reuse improvement, because "1000 different sizes" collapses to "30 size classes."
  • Allocation path: exact-ish bin hit → reuse (fast path, ~100 ns). Miss → take a larger cached block and split it (remainder returns to bins). Still nothing → cudaMalloc. OOM → free all cached-unsplit blocks, retry (the "sync-and-retry" you see as a mysterious stall right before a PyTorch OOM error).
  • Free path: blocks go back to the bins — the driver never sees the free (this is why nvidia-smi shows "used" memory that the framework considers free: reserved vs allocated, the #1 memory-confusion in every ML org — Q2 below).
  • Stream tagging: every block remembers the stream it was used on; reuse on a different stream requires an event dependency (or is simply not done) — preventing the cross-stream use-after-free (Ch. 7).
  • expandable_segments (modern): uses the CUDA VMM API to grow a virtual reservation with physical pages on demand — restoring a little of the "remapping" power CPUs get for free, specifically to fight fragmentation.

Misconception to kill: "the allocator caches because cudaMalloc is slow." Half right — the deeper reason is cudaFree's synchronization and the fragmentation ratchet. Even with a free cudaMalloc you'd still want the pool.

Chapter 5: Fragmentation — The Silent Fleet Killer

Internal fragmentation: waste inside blocks (rounding). Bounded, boring, fine.

External fragmentation: free memory exists but in pieces too small for the request. The serving-fleet scenario (Lab 01's fragdemo reproduces it):

  1. Day 0: requests of mixed sizes allocate/free happily; memory is a clean checkerboard.
  2. Day 2: long-lived allocations (a resident model, a stuck request's KV) sit between short-lived ones. Frees leave holes of assorted sizes.
  3. Day 3: a 2 GB request arrives. Free memory totals 9 GB. No hole exceeds 1.1 GB. OOM at 60% utilization. Restart fixes it (fresh address space) — which is why "we restart the pods nightly" is a real, embarrassing, industry-wide mitigation.

Because compaction is impossible (live raw pointers), the defenses are all preventive:

  • Shape discipline: bucket request sizes (pad to standard lengths) — fewer distinct sizes = checkerboard stays clean. (Same trick as batching same-length sequences — Phase 01 Ch. 4's divergence logic, reborn.)
  • Pool segregation: small/large pools; per-purpose pools (weights vs activations vs KV); per-stream pools.
  • Fixed-size paging for the dominant consumer: if one consumer (KV cache) dominates and grows unpredictably, give it fixed-size blocks + an indirection table = external fragmentation eliminated for that consumer by design. That sentence is PagedAttention (Phase 07), derived from first principles.
  • Metrics that see it coming (Phase 10): track reserved, allocated, and largest-free-block — the third is the early-warning signal nobody exports by default.

Chapter 6: Streams and Events — The Ordering Algebra

The execution model that Lab 02 implements, reduced to algebra:

  • A stream is a FIFO: operations issued to the same stream execute in issue order. Within a stream you never need explicit sync.
  • Different streams have no order — unless you create one.
  • An event is a marker recorded into a stream; streamWaitEvent(s2, e) makes s2 pause until e fires. Events are the edges of the DAG; streams are convenient paths through it.
  • cudaStreamSynchronize/cudaDeviceSynchronize join the CPU to the DAG — blunt instruments; in engine code, almost every device-wide sync is a bug or a missed event edge.
  • The legacy default stream synchronizes with all other streams (a global serializer hiding in every naive codebase — Lab 02 scenario 2 measures the damage; the fix is per-thread default streams or explicit streams everywhere, Phase 02 Ch. 9).

Mental model for review: draw the intended DAG, then check the code creates exactly those edges — extra edges = lost overlap (performance bug); missing edges = race (correctness bug). Most stream code review reduces to this.

Chapter 7: Stream-Ordered Lifetime — Why "Free" Is an Event

The bug class that motivates stream-aware allocators (Lab 02 scenario 3):

stream A: kernel K reads buffer B          (still running...)
host:     free(B)                          (CPU is ahead of the GPU!)
allocator: hands B's memory to stream C    -> C writes over K's input. Corruption.

The CPU runs ahead of the GPU (async launches, Phase 02 Ch. 1) — so a host-side free() happens-before the GPU work that uses the buffer. The fix is to make deallocation an operation in the stream: the block returns to the pool only when the stream reaches the free point (an event). Reuse on another stream inserts a wait-event edge first. cudaMallocAsync/ cudaFreeAsync are exactly this, productized; PyTorch's stream tagging (Ch. 4) is the same idea retrofitted.

The leadership phrasing: lifetime on a GPU is defined by stream position, not by wall-clock or host program order. Any API in your platform that forgets this will eventually corrupt a customer's tokens — and it will be unreproducible, because it needs a scheduling coincidence. (This is the bug class behind a good fraction of "model produced garbage once" tickets.)

Chapter 8: CUDA Graphs — Amortizing Launch Overhead

The numbers (Phase 02): each kernel launch costs ~5–10 µs of CPU-side work. An LLM decode step is hundreds of short kernels; at high batch, GPU work per kernel can be tens of µs — the CPU becomes the bottleneck issuing work, and the GPU starves between kernels.

CUDA Graphs: record the whole sequence of launches (+ their dependency edges) once — capture — then replay the graph with one API call. Launch cost amortizes from n × 8 µs to ~1 × 10 µs + tiny per-node cost. Serving engines (vLLM et al.) capture the decode step per batch-size bucket; it's worth 10–30% at small batch and is the standard answer to "why is my decode CPU-bound?"

Restrictions (the part people learn the hard way): the graph is static — shapes, pointers, and control flow are baked at capture; capture forbids synchronous APIs (cudaMalloc!) inside the captured region (stream-ordered allocation exists partly for graphs); dynamic batch = capture per bucket or use graph update APIs. Lab 02's capture/replay extension makes all of this tactile in 100 lines.

Chapter 9: Pinned Pools and Staging Pipelines

The host↔device data path, productized (composing Phase 02 Ch. 5/9 with this phase's pooling):

  • Pinned memory is mandatory for true async copies but expensive to allocate and steals unpageable RAM → pool it like device memory (a pinned staging-buffer pool is ~50 lines on top of Lab 01's design).
  • The standard pipeline: chunk → H2D copy on copy-stream → kernel on compute stream (event edge) → D2H on copy-stream (event edge) — double-buffered so chunk N+1's copy overlaps chunk N's compute. GPUs have separate copy engines precisely so this overlap is real.
  • Serving-specific instances you'll meet in Phase 07: weight loading at startup (pinned + multi-stream = minutes to seconds), KV-cache swap to host (the 50× bandwidth cliff makes the policy — what to swap when — the hard part, not the mechanism).

Lab Walkthrough Guidance

Order: Lab 01 → Lab 02 (the scheduler's stream-ordered-free scenario reuses allocator concepts).

Lab 01 (allocator, Rust):

  1. cargo test green first; then read src/lib.rs with Ch. 3–4 open. The types are small: Block, BinKey, CachingAllocator.
  2. Trace one allocation by hand: 1.2 MB request → which bin, which split, what's left where. Write the trace in comments; the test test_split_accounting checks your understanding.
  3. Run cargo run --release --bin fragdemo: the naive-vs-caching utilization table is the artifact. Explain the gap using Ch. 5's checkerboard story.
  4. Extensions in order: coalescing (hard correctness) → empty_cache → trace replayer (turns the allocator into a measurable system — feed it your own traces).

Lab 02 (scheduler, C++):

  1. make && ./scheduler_test; then read scheduler.hpp with Ch. 6–7 open — Stream is a worker+FIFO, Event is a counter+condvar, that's it.
  2. Scenario 1 (overlap): predict the speedup from the DAG before reading the assert. Scenario 2 (legacy default stream): observe the global serializer. Scenario 3 (lifetime): the unsafe path corrupts; the stream-ordered path can't — connect each line to Ch. 7's diagram.
  3. Extension: capture/replay. Measure issue-cost per op before/after — your own numbers for the Ch. 8 argument.

Success Criteria

  • You can give the three reasons raw cudaMalloc/Free are unusable in steady state, with the design consequence of each (Ch. 2)
  • You can sketch the caching allocator (bins, split, stream tags, sync-and-retry) and explain reserved vs allocated (Ch. 4)
  • You can tell the day-3 fragmentation story with numbers and name the four preventive defenses (Ch. 5)
  • You can review stream code by drawing its DAG and checking edges (Ch. 6)
  • You can explain why free must be stream-ordered and what bug class ignores it (Ch. 7)
  • You know what CUDA Graphs buy, their restrictions, and when serving engines use them (Ch. 8)
  • Both labs pass; one extension each implemented

Interview Q&A

Q1: Why does PyTorch implement its own allocator instead of calling cudaMalloc? A: Three driver-level facts force it: cudaFree device-synchronizes (a global barrier per free), cudaMalloc is a driver round-trip (tens of µs vs ~100 ns pool hit), and HBM has no paging/compaction so fragmentation is permanent. The caching allocator answers all three: it never returns memory to the driver in steady state (frees go to size-binned lists), reuses blocks in ~O(1), rounds sizes to bound fragmentation, splits larger cached blocks for fit, tags blocks by stream to keep reuse safe under async execution, and on apparent-OOM frees its caches and retries. cudaMallocAsync later moved the same design into the driver as stream-ordered pools.

Q2: nvidia-smi shows 78 GB used; the framework says 51 GB. Who's lying? A: Nobody — they measure different lines of the same ledger. nvidia-smi sees what the driver handed the process: reserved memory, which includes everything the caching allocator is holding in its free lists. The framework's 51 GB is allocated — blocks currently handed to live tensors. The 27 GB gap is cache (plus context/fragmentation overhead). It's reclaimable via empty_cache() (at a performance cost) and is not a leak. The follow-up metric a platform should actually watch is largest free block — the fragmentation early-warning — and allocated-vs-reserved divergence trend, which distinguishes cache growth from a real leak.

Q3: A serving process OOMs after 3 days at steady load. Hypotheses, in order? A: (1) External fragmentation — mixed request sizes + long-lived allocations checkerboarding HBM until no hole fits a big request; check largest-free-block trend and whether OOM happens at <70% allocated; mitigate with size bucketing, pool segregation, paged KV (or, honestly, scheduled restarts while you fix it). (2) Genuine slow leak — allocated (not just reserved) trending up: a cache without eviction, per-request tensors retained by a logger, growing CUDA Graph pools per shape bucket. (3) Pinned host memory exhaustion masquerading as device OOM in mixed stacks. (4) KV-cache admission policy allowing pathological long-context requests to accrete. The diagnostic discipline: reserved vs allocated vs largest-block, per pool, on a dashboard before day 3 (Phase 10).

Q4: Two streams, one buffer: kernel on stream A reads it, host frees it, stream C gets it and writes. How do you make this class of bug impossible? A: Make deallocation stream-ordered: free() doesn't return the block to the pool — it records an event at the free point in A; the block becomes reusable only when that event completes. Reuse on a different stream additionally inserts streamWaitEvent so C's first use happens-after A's last. This is exactly cudaFreeAsync semantics / PyTorch's stream tagging. The general principle: on an async device, lifetime is a position in the dependency DAG, not a moment in host time — any API that takes raw "free now" must be wrapped before it touches engine code. Detection for the residue: compute-sanitizer in CI and a debug allocator mode that poisons freed blocks.

Q5: When do CUDA Graphs help, and what are the sharp edges? A: They help when launch overhead is material: many short kernels per iteration with a stable structure — the LLM decode step is the canonical case (hundreds of kernels, µs-scale each); capture per batch-bucket and replay, typically 10–30% at small batch and a fix for CPU-bound issuance. Sharp edges: the graph freezes shapes, addresses, and topology (dynamic control flow can't be captured); no synchronous calls (cudaMalloc) inside capture — you need stream-ordered allocation; per-shape graph pools consume memory (a real OOM contributor at high bucket counts); and debugging a captured region is harder (prints/syncs change behavior). So: graphs for the steady-state inner loop, eager mode for everything else.

Q6 (leadership): What runtime-layer guarantees would you write into your platform's internal API contract? A: (1) No driver alloc/free in steady state — all hot-path memory from pools; CI fails on cudaMalloc in request paths. (2) Stream-ordered lifetime — the only free is an event-ordered free; raw frees don't exist in the API. (3) Pool segregation + metrics — weights/activations/KV in separate pools, each exporting reserved/allocated/largest-block (the fragmentation dashboard). (4) No device-wide syncs in engine code — event edges only; a lint catches deviceSynchronize. (5) Graph-compatible inner loops — steady-state paths must be capturable (no syncs, no allocs), which doubles as a cleanliness forcing-function. (6) Determinism escape hatch — a debug mode with serialized streams and poisoned frees for reproducing the "garbage tokens once" class. Each guarantee exists because its absence is a specific incident I can name — that's how runtime contracts should be justified.

References

  • PyTorch CUDA Caching Allocator notes (the design doc) — https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
  • PyTorch CUDACachingAllocator.cpp (read the real one after the lab) — https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDACachingAllocator.cpp
  • NVIDIA, "Using the NVIDIA CUDA Stream-Ordered Memory Allocator" (cudaMallocAsync) — https://developer.nvidia.com/blog/using-cuda-stream-ordered-memory-allocator-part-1/
  • NVIDIA, CUDA Graphs section of the Programming Guide + "Getting Started with CUDA Graphs" — https://developer.nvidia.com/blog/cuda-graphs/
  • NVIDIA VMM API ("Introducing Low-Level GPU Virtual Memory Management") — https://developer.nvidia.com/blog/introducing-low-level-gpu-virtual-memory-management/
  • Wilson et al., "Dynamic Storage Allocation: A Survey and Critical Review" (1995) — the allocator-design bible
  • Bonwick, "The Slab Allocator" (USENIX 1994)
  • vLLM source: cuda_graph usage and block manager — https://github.com/vllm-project/vllm
  • Cross-track: Phase 07 WARMUP (PagedAttention as allocator design), LLM Inference track Phase 09