CUDA Memory Pool, in depth: the stream-ordered allocator, release thresholds, cross-stream reuse, CUDA graphs and PyTorch's caching allocator

Why pools exist

Plain cudaMalloc and cudaFree are expensive for two reasons. First, they involve the driver changing the GPU's virtual memory mappings, which costs far more than handing out a pointer from a free list. Second, cudaFree synchronises the device: it waits for outstanding work so that memory still in use by a kernel is not released under it. A training step that allocates and frees hundreds of activation and gradient buffers would stall the GPU repeatedly if every buffer went through these calls.

A pool fixes both. Freed blocks go back to the pool instead of the driver, and new requests are served from the pool when a large-enough block is available. The driver is only called when the pool must grow or is asked to shrink. The cost is that memory held by the pool is not available to other processes, and that a pool can hold plenty of free memory in pieces too small or too scattered to satisfy a large request, which is fragmentation. Where the memory goes inside a step, weights, gradients, optimizer state and activations, is covered in Anatomy of One GPU Training Step.

The stream-ordered allocator

CUDA 11.2 introduced cudaMallocAsync and cudaFreeAsync. Both take a stream argument, and that is the key idea: allocation and free become operations in the stream, ordered with the kernels around them. A free is enqueued like a kernel. Memory freed in stream S is safe to reuse for any later allocation in S, without waiting, because stream order already guarantees that the earlier kernels using it run first. No device-wide synchronisation is needed.

Each device has a default pool, which cudaMallocAsync uses, and you can create explicit pools with their own properties and allocate from them with cudaMallocFromPoolAsync. The pool's reserved memory is the physical memory it has mapped; its used memory is the portion currently handed out. The difference is cached free memory waiting for reuse.

Creating and tuning a pool

The following program creates an explicit pool on device 0, keeps up to 4 GiB of freed memory across synchronisations, runs an allocate-compute-free loop on one stream, and reads the pool's statistics. Error checking is reduced to a macro for brevity.

#include <cuda_runtime.h>
#include <cstdio>
#include <cstdint>

#define CK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \
    printf("%s failed: %s\n", #x, cudaGetErrorString(e)); return 1; } } while (0)

__global__ void scale(float* p, size_t n) {
    size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x;
    if (i < n) p[i] = p[i] * 0.5f + 1.0f;
}

int main() {
    cudaMemPoolProps props = {};
    props.allocType = cudaMemAllocationTypePinned;
    props.location.type = cudaMemLocationTypeDevice;
    props.location.id = 0;

    cudaMemPool_t pool;
    CK(cudaMemPoolCreate(&pool, &props));

    uint64_t keep = 4ull << 30;                       // keep up to 4 GiB cached at sync
    CK(cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &keep));

    cudaStream_t s;
    CK(cudaStreamCreate(&s));
    for (int step = 0; step < 100; ++step) {
        size_t n = (size_t)(1 + step % 4) << 24;       // varying sizes, 64 to 256 MiB
        float* buf;
        CK(cudaMallocFromPoolAsync((void**)&buf, n * sizeof(float), pool, s));
        CK(cudaMemsetAsync(buf, 0, n * sizeof(float), s));
        scale<<<(unsigned)((n + 255) / 256), 256, 0, s>>>(buf, n);
        CK(cudaFreeAsync(buf, s));                     // ordered after the kernel
        CK(cudaStreamSynchronize(s));                  // a sync point: trimming may happen here
    }

    uint64_t reserved = 0, used = 0, high = 0;
    cudaMemPoolGetAttribute(pool, cudaMemPoolAttrReservedMemCurrent, &reserved);
    cudaMemPoolGetAttribute(pool, cudaMemPoolAttrUsedMemCurrent, &used);
    cudaMemPoolGetAttribute(pool, cudaMemPoolAttrReservedMemHigh, &high);
    printf("reserved %llu used %llu reserved_high %llu\n",
           (unsigned long long)reserved, (unsigned long long)used, (unsigned long long)high);

    CK(cudaMemPoolTrimTo(pool, 0));                    // hand unused memory back to the driver
    CK(cudaStreamDestroy(s));
    CK(cudaMemPoolDestroy(pool));
    return 0;
}

The release threshold is the most important knob and its default surprises people. By default it is zero, which means all unused memory in the pool is released back to the operating system at every synchronisation on a stream, event or device. A loop like the one above that synchronises every iteration would then remap memory every iteration and lose most of the pool's benefit. Setting the threshold to the working set, or to UINT64_MAX for a process that owns the GPU, keeps the memory cached. cudaMemPoolTrimTo(pool, bytesToKeep) releases unused memory above bytesToKeep explicitly, which is how a long-lived server gives memory back after a burst.

Reuse across streams and devices

Reuse within one stream is free of hazards. Reuse across streams is where the design gets careful. If stream A frees a block and stream B allocates immediately, B's kernels might run while A's last kernel is still reading the block. The pool has three reuse policies, all enabled by default and controllable per pool through attributes:

• cudaMemPoolReuseFollowEventDependencies: if B has waited on an event recorded in A after the free, the pool knows the free has happened from B's point of view and can hand the block to B.
• cudaMemPoolReuseAllowOpportunistic: the pool may reuse a block whose free has already completed on the GPU, even without an explicit dependency. Whether that happens depends on timing, so the memory footprint can vary from run to run.
• cudaMemPoolReuseAllowInternalDependencies: the pool may insert a dependency itself, making B's allocation wait for A's free, rather than grow the pool.

The practical rule: if you need reproducible peak memory, disable opportunistic reuse and express cross-stream ordering with events. If you need the smallest footprint and can tolerate variation, leave the defaults. Either way, a pointer allocated in one stream and used in another must be ordered with events before the free, exactly as for any other data hazard.

Pools are also the unit of sharing. cudaMemPoolSetAccess grants peer devices access to a pool's allocations. For inter-process sharing, a pool created with a shareable handle type in props.handleTypes can be exported with cudaMemPoolExportToShareableHandle and individual allocations with cudaMemPoolExportPointer; the default pool is not created for IPC, so this needs an explicit pool.

Pools and CUDA graphs

When a stream is captured into a CUDA graph, cudaMallocAsync and cudaFreeAsync calls are recorded as memory allocation and free nodes in the graph. That lets the graph own its temporary buffers and lets the runtime plan their placement for the whole graph, overlapping lifetimes that never coexist. If an allocation made inside a graph is not freed inside it, the graph cannot be relaunched until that memory is freed, unless it is instantiated with cudaGraphInstantiateFlagAutoFreeOnLaunch.

PyTorch takes a related approach: torch.cuda.graph captures into a private memory pool so the addresses baked into the graph stay valid on replay, and graphs that run strictly one after another can share a pool through torch.cuda.graph_pool_handle(). The consequence for capacity planning is that memory held by captured graphs is not available to the eager allocator, so an inference server with many captured batch sizes can hold a surprising amount of reserved memory.

PyTorch's caching allocator

PyTorch's default allocator, called native in its configuration, is a caching allocator built directly on cudaMalloc. It rounds requests up, caches freed blocks per stream, splits large cached blocks to serve smaller requests, and on an allocation failure frees its cached blocks and retries before raising an out-of-memory error. In its source at the time of writing, small requests (up to 1 MB) are carved from 2 MB segments and mid-size ones from 20 MB segments; the documentation exposes the large segment size as large_segment_size_mb with a default of 20.

Configuration goes through the PYTORCH_ALLOC_CONF environment variable, with PYTORCH_CUDA_ALLOC_CONF kept as a backward-compatible alias. The options that matter most:

Two other mechanisms explain common surprises. A tensor used on a stream other than the one that allocated it must be registered with tensor.record_stream(s), or the allocator may reuse its block while the other stream still reads it. And torch.cuda.empty_cache() returns unused cached memory to the driver so other processes can use it; it does not free tensors and does not make more memory available to the same process.

Diagnosing fragmentation: a worked example

Diagnose with two numbers before changing any setting. torch.cuda.memory_allocated() is memory occupied by live tensors; torch.cuda.memory_reserved() is everything the allocator holds, including its cache. The gap is cached but unused memory. If an out-of-memory error reports large reserved-but-unallocated memory, the problem is fragmentation, not capacity.

import torch

torch.cuda.memory._record_memory_history(max_entries=100_000)
for step, batch in enumerate(loader):
    loss = model(batch).loss
    loss.backward()
    opt.step(); opt.zero_grad(set_to_none=True)
    if step % 50 == 0:
        alloc = torch.cuda.memory_allocated() / 2**30
        resv = torch.cuda.memory_reserved() / 2**30
        peak = torch.cuda.max_memory_allocated() / 2**30
        frag = 1 - alloc / resv if resv else 0.0
        print(f"step {step}: allocated {alloc:.1f} GiB reserved {resv:.1f} GiB "
              f"peak {peak:.1f} GiB cached-unused {frag:.0%}")
torch.cuda.memory._dump_snapshot("mem_snapshot.pickle")   # inspect at pytorch.org/memory_viz

A worked example with illustrative numbers. A fine-tuning job on an 80 GB GPU uses variable-length batches. After a few hundred steps it fails with an out-of-memory error trying to allocate 2.1 GiB, while the log shows 61 GiB allocated and 76 GiB reserved: 15 GiB cached but unused, about 20 percent of the reserve. The snapshot shows the cache is dozens of blocks between 200 MB and 1.5 GB, each left behind by a different sequence length, none big enough for the 2.1 GiB request. The fix order is: pad or bucket sequence lengths so sizes repeat; try expandable_segments:True, which lets segments grow contiguously instead of leaving fixed-size holes; and only then lower the batch size. Measure the cached-unused share after each change rather than assuming it helped. For the activation side of the budget, Activation Memory Math gives the formulas, and nvidia-smi shows the process-level view, which counts reserved memory plus the CUDA context, not allocated memory.

Failure modes

• Default release threshold of zero. A loop that synchronises often remaps memory every iteration. Set the threshold to the working set.
• Missing cross-stream ordering. Using a buffer in stream B after freeing it in stream A without an event is a race. In PyTorch, the equivalent bug is a missing record_stream.
• Reading nvidia-smi as tensor usage. It reports reserved memory plus the CUDA context, so a cache looks like a leak. Compare allocated and reserved inside the process.
• Calling empty_cache in the training loop. It forces the next steps to go back to the driver and slows the job, without preventing fragmentation.
• Two allocators in one process. A library using its own pool and the framework using another each cache memory the other cannot use. Route libraries through the framework allocator where they support it, or cap one of them.
• Unbounded graph pools. Capturing many shapes each with private pools multiplies reserved memory. Share pools across graphs that never run concurrently.

Trade-offs

What to do next

1. Log allocated, reserved and peak memory every N steps in one real job and compute the cached-unused share.
2. Record one memory snapshot around a failing or near-failing step and identify which sizes fill the cache.
3. Bucket or pad variable shapes so allocation sizes repeat.
4. Try expandable_segments:True on that job and compare peak reserved memory and step time.
5. In CUDA code, replace hot cudaMalloc/cudaFree pairs with cudaMallocAsync/cudaFreeAsync and set a release threshold.
6. Audit every cross-stream buffer for an event or record_stream before its free.