The caching allocator

The allocator is holding 8 GB of reserved memory, but no single contiguous run of 2 GB inside one segment is free. Device allocations must be contiguous, so the request fails despite the total.

How you get there: the allocator obtained segments over time and split them into blocks of the sizes you asked for. Free those blocks and they coalesce — but only with neighbours inside the same segment. Two 1 GB holes in different segments cannot serve a 2 GB request. A workload alternating between size classes accelerates this: request 1.5 GB, split a 2 GB block, leave a 0.5 GB remainder; free it; request 1.8 GB, which does not fit the 0.5 GB remainder, so take a new segment. Repeat, and reserved memory climbs while usable contiguity falls.

The diagnostic sequence:

1. Check reserved − allocated. A large gap is fragmentation; a small gap means you genuinely need more memory and no allocator setting will help.
2. Record a snapshot with _record_memory_history(), reproduce, _dump_snapshot(), and view it. The timeline shows segments, the blocks inside them, and the allocation call stacks — it makes the holes literally visible.
3. Look for size-class churn in the trace: many distinct sizes is the signature.

The fixes, in order:

• expandable_segments:True. This is the real answer and it is nearly free — see the next probe.
• Reduce shape diversity. Bucket sequence lengths, pad to multiples, fix the batch size. This helps compilation too (Module 10), which is not a coincidence.
• max_split_size_mb to stop large blocks being carved into unusable remainders.
• empty_cache() only as a last resort, and understand you are paying synchronisation to reset the pool.

They separate virtual address reservation from physical memory backing, using CUDA’s virtual memory management API rather than plain cudaMalloc.

A classic segment is a fixed slab of physical memory at a fixed size. To grow past it you allocate a new segment, and the two are not adjacent, so their free space cannot merge — that non-adjacency is fragmentation’s root cause.

An expandable segment reserves a large virtual address range up front and maps physical pages into it on demand. Because the virtual range is contiguous and reserved, the allocator can grow the segment by mapping more pages at the end, and shrink it by unmapping. There is effectively one growable region per pool per stream instead of many fixed islands, so a large request can almost always be satisfied by extending rather than by needing a pre-existing hole of the right size.

Why it is not simply the default everywhere: it relies on the driver’s VMM API, interacts with features that assume stable physical mappings (some IPC and peer-access paths), and adds cuMemMap/cuMemUnmap calls on growth. In practice it is a large win for varying-shape workloads and roughly neutral for fixed-shape ones, and it has been becoming the recommended default.

The conceptual point worth keeping: fragmentation is a contiguity problem, not a capacity problem, and virtual memory is the general solution to contiguity problems. CPU allocators have used this for decades; PyTorch’s allocator predates the CUDA API that makes it possible on device.

Because freeing a tensor is a CPU-side event and the GPU work using it may not have run yet.

del x drops the refcount and returns the block to the pool immediately, on the CPU. But the kernel reading x was enqueued asynchronously and may still be pending. If the allocator handed that block to another allocation whose kernel runs on a different stream, the two kernels could execute concurrently and one would overwrite memory the other is reading. A race, producing wrong numbers non-deterministically.

The allocator’s rule: a block freed on stream A is reusable on stream A without ceremony, because stream A is ordered with respect to itself. To reuse it on stream B, the allocator must record an event on A and make B wait on it. PyTorch does this only when you tell it to, via tensor.record_stream(s) — and forgetting that call in multi-stream code is a classic source of rare, load-dependent corruption.

The practical consequences:

• Multi-stream code needs record_stream anywhere a tensor allocated on one stream is consumed on another, which includes many overlapped-communication patterns.
• Per-stream pools cost memory. Blocks freed on a busy stream are unavailable to others, so effective capacity falls as you add streams.
• torch.cuda.Stream is not free. Overlapping compute and copy has a real memory cost that is invisible until you OOM at higher concurrency.

The general shape: PyTorch’s allocator is fast because it makes decisions synchronously on the CPU about memory used asynchronously on the GPU, and every hazard in this module comes from that gap.

memory_summary() gives aggregates — allocated, reserved, per-pool counts, peak. It tells you that you have fragmentation. It cannot tell you what caused it.

The snapshot records every allocation and free with its Python stack trace and timestamp, plus the segment and offset. Loaded into the viewer at pytorch.org/memory_viz you get two things that change the investigation:

• The segment view, showing each segment as a bar with its blocks laid out by address. Holes are visible directly, and you can see which allocations are pinning a segment open — very often a single long-lived small tensor sitting in the middle of an otherwise-free large segment, which aggregates can never reveal.
• The timeline, showing allocated memory over the step with each event attributed to a stack frame. Peaks become attributable: you can see the exact line whose activation caused the maximum, which is usually somewhere you did not expect.

The workflow that resolves most real cases:

torch.cuda.memory._record_memory_history(max_entries=100_000)
run_until_it_ooms()
torch.cuda.memory._dump_snapshot("snap.pickle")

Then look for: a long-lived allocation inside a large segment (move it, or allocate it first during warm-up so it lands at the bottom); many distinct size classes (bucket your shapes); and whether the peak comes from activations, from the optimiser state, or from a cached-but-unused gap.

The habit worth building: peak memory is a property of a schedule, not of a model. The snapshot is the only tool that shows you the schedule.

Because planning requires knowing the whole program in advance, and eager mode’s entire premise is that it does not.

XLA compiles a complete graph with static shapes, so it can compute exact buffer lifetimes, solve for a near-optimal assignment, and emit code that indexes into one preallocated arena with no allocator at run time at all. That is strictly better — when it applies. It requires static shapes, no data-dependent control flow, and a complete graph, which is the TF1 bargain from Module 1.

PyTorch’s allocator is what you build when the next allocation is genuinely unknown: fast enough to call thousands of times per step, adaptive to any sequence of requests, and correct under arbitrary Python. It gives up optimality for generality, and the fragmentation in this module is precisely the price.

The interesting part is that the two are converging rather than competing. Inductor plans within a compiled region (Module 6). CUDA Graphs use a private memory pool whose addresses are baked into the replay, so a captured graph has statically-assigned memory (Module 9). vLLM preallocates the KV cache as one arena and manages it itself with paging (Module 14) — reinventing planned allocation where it can prove the shapes.

So the answer to “why not plan” is: PyTorch plans exactly where it can prove it is safe, and the scope of that proof has been growing for three years. The allocator remains for everything outside it, which is still most of a real program.