pinned-allocator

A PyTorch CUDA allocator backed by CPU pinned (page-locked) memory.

Tensors allocated through this allocator report device='cuda' but their data physically lives in CPU RAM, consuming zero GPU HBM. The GPU accesses them via CUDA mapped memory (cudaHostRegisterMapped / UVA) over PCIe.

Use case

When GPU HBM is scarce, offload large weight or buffer tensors to CPU pinned memory while keeping them accessible as normal cuda tensors. Useful for inference with oversized models, multi-tenant GPU sharing, or any workload where a subset of tensors is accessed infrequently.

Installation

pip install -e /path/to/pinned-allocator

The C++ extension (pinned_alloc_ext.so) is compiled automatically on first import using g++ and the system CUDA toolkit. CUDA_HOME defaults to /usr/local/cuda.

Critical requirement

import pinned_allocator must happen before any CUDA context is initialized — before any torch.*(..., device='cuda') call, torch.cuda.current_device(), etc. PyTorch raises an error if you call change_current_allocator after the default allocator is already initialized.

# CORRECT
import pinned_allocator          # first
import torch                     # fine to import torch before, but no cuda ops yet
x = torch.randn(1024, device="cuda")

# WRONG — will raise at import
x = torch.randn(1024, device="cuda")
import pinned_allocator          # too late

API

After import, the allocator is installed but defaults to normal cudaMalloc mode. Switch to pinned mode explicitly:

import pinned_allocator
import torch

# Enable pinned memory (data lives in CPU RAM, zero HBM)
pinned_allocator.use_pinned(True)
w = torch.randn(4096, 4096, device="cuda")   # zero HBM consumed

# Switch back to GPU HBM (cudaMalloc)
pinned_allocator.use_pinned(False)
x = torch.randn(1024, 1024, device="cuda")   # uses HBM

# Query current mode
print(pinned_allocator.is_pinned_mode())     # False

# Allocate a single tensor in pinned mode regardless of current mode
buf = pinned_allocator.pin_empty(4096, 4096, dtype=torch.float16)

# Inspect the physical CPU address of a pinned tensor (returns 0 for HBM tensors)
host_ptr = pinned_allocator.host_ptr_of(w)
print(hex(host_ptr))

Mode changes affect only tensors allocated after the call. Existing tensors retain their original backing and are freed correctly regardless of the current mode.

Functions

Function	Description
use_pinned(enabled: bool)	Switch allocation backend. True = CPU pinned, False = GPU HBM.
is_pinned_mode() -> bool	Return current allocation mode.
pin_empty(*size, dtype, device) -> Tensor	Allocate a tensor always backed by pinned memory, temporarily overriding the current mode.
host_ptr_of(tensor) -> int	Return the CPU physical address of a pinned tensor, or 0 for HBM tensors.

Verified HBM behavior

Measured on a single GPU with 84 GB of logical cuda tensors:

Step	Action	HBM used
Baseline	—	0 GB
Step 1	Allocate 40 GB pinned tensor	~0.59 GB (CUDA context overhead only)
Step 2	Allocate 40 GB HBM tensor	~40.59 GB
Step 3	Allocate 4 GB pinned tensor	~40.60 GB (unchanged)
Cleanup	del + empty_cache	~0.51 GB

Pin memory tensors consume CPU RAM, not GPU HBM.

Caveats

PCIe bandwidth. GPU access to pinned memory goes over PCIe (~32 GB/s) rather than HBM (~3 TB/s). Pinned tensors are suitable for weights loaded infrequently or streamed once, but not for activations or compute-intensive operations that need high bandwidth.

Memory stats APIs. torch.cuda.memory_allocated(), memory_reserved(), and related functions do not work with CUDAPluggableAllocator. This library patches them to return safe defaults (0 / empty dict). Use nvidia-smi to query real HBM usage:

nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits

CUDA_VISIBLE_DEVICES remapping. Setting CUDA_VISIBLE_DEVICES=5 remaps physical GPU 5 to cuda:0 inside the process. Pass --id 5 (physical index) to nvidia-smi.

Internals

The C++ extension (pinned_alloc_ext.cpp) implements two allocation paths behind a single pinned_malloc / pinned_free interface required by CUDAPluggableAllocator:

• Pinned path (g_use_pinned == true): posix_memalign → cudaHostRegister(cudaHostRegisterMapped) → cudaHostGetDevicePointer. Returns the device pointer; the host pointer is stored in g_dev_to_host.
• GPU path (g_use_pinned == false): plain cudaMalloc. A nullptr sentinel is stored in g_dev_to_host to distinguish this allocation at free time.
• Free: looks up g_dev_to_host to determine which path was used, then calls cudaHostUnregister + free (pinned) or cudaFree (GPU). Thread-safe via std::mutex.
• Mode switch: std::atomic<bool> — lock-free, safe to call from any thread.