Stream-Ordered Memory Allocation

cudaMalloc/cudaFree are synchronous and slow.

cudaMalloc/cudaFree are synchronous and slow. CUDA 11.2+ provides stream-ordered allocation from memory pools:

Libraries like RMM (RAPIDS Memory Manager) provide pool/arena allocators; PyTorch uses its own caching allocator.

float* tmp;
cudaMallocAsync(&tmp, bytes, stream);           // allocation ordered in the stream
kernel<<<grid, block, 0, stream>>>(tmp);
cudaFreeAsync(tmp, stream);                      // memory returned to the pool, reusable immediately

cudaMemPool_t pool;
cudaDeviceGetDefaultMemPool(&pool, 0);
uint64_t threshold = UINT64_MAX;
cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold);   // keep memory cached