I can’t figure out why this limitation exists. As far as I know, CUDA doesn’t forbid different cuda streams to access the memory allocated by cudaMalloc. And this limitation will disallow some free blocks allocated by one cuda stream to be reused by other cuda streams.
Example:
import torch
a = torch.rand((12, 12)).cuda() # allocated by default cuda stream
print(torch.cuda.memory_cached())
torch.cuda.synchronize()
s = torch.cuda.Stream()
with torch.cuda.stream(s):
b = torch.rand((12, 12)).cuda() # allocated by other cuda stream
torch.cuda.synchronize()
print(torch.cuda.memory_cached())
output is: 2097152 and 4194304.
import torch
a = torch.rand((12, 12)).cuda() # allocated by default cuda stream
print(torch.cuda.memory_cached())
b = torch.rand((12, 12)).cuda() # allocated by default cuda stream
print(torch.cuda.memory_cached())
As @albanD wrote, limiting CUDA allocations to a single stream reduces the number of CPU-GPU synchronizations necessary. CUDA kernels are asynchronous, so when an allocation is “freed” the kernel may not be finished (or may not have even started). Reusing the same allocation in a different stream could cause memory corruption because work in that stream may start before previously launched work in the original stream finishes.
It’s safe to immediately reuse the allocation in the same stream because operations within a stream are ordered sequentially. This is the strategy the caching allocator uses.
The CUDA memory API handles this differently: The cudaFree call synchronize all streams – the CPU waits until all streams finish all outstanding work before the cudaFree call completes. This ensures that subsequent uses of the memory are guaranteed to happen after the previous uses finish. However, this makes cudaFree a relatively expensive call. The primary goal of the caching allocator is to avoid this type of synchronization.