We have 3 design choices here

1. Inheritance from Allocator: Currently, both StreamSafeCUDAAllocator and CUDAMallocAsyncAllocator inherit directly from the Allocator class. However, only StreamSafeCUDAAllocator and CUDAMallocAsyncAllocator possess stream-related methods, while other allocator types do not.
   • Pros: This approach aligns well with the conceptual model where CUDAMallocAsyncAllocator is seen as a alternative to StreamSafeCUDAAllocator.
   • Cons: We must upcast to Allocator and then downcast back to CUDAMallocAsyncAllocator or StreamSafeCUDAAllocator.
2. Centralizing Stream-Related Methods: Another approach is to move stream-related methods into the base Allocator class. For allocators that do not support these methods, they would trigger a runtime error.
   • Cons: It complicates the base class with methods that are irrelevant for some of its subclasses, violating the principle of interface segregation.
3. Inheriting CUDAMallocAsyncAllocator from StreamSafeCUDAAllocator: The third option considers making CUDAMallocAsyncAllocator inherit from StreamSafeCUDAAllocator. This approach implies a direct relationship between the two, with one being a more specific version of the other.
   • Cons: This design is conceptually awkward as it positions CUDAMallocAsyncAllocator as a subtype of StreamSafeCUDAAllocator, despite it being intended as a replacement. It might imply a false relationship or hierarchy, leading to confusion and potentially misused inheritance.

This API has been renamed for clarity.

If input_tensor.grad is None, this may cause a hanging issue.

Add some comments？

Fixed

Where is memory_stream_ used?

Currently, the memory_stream_ serves no immediate function and is being retained for potential future applications (see code here). The original intention behind its design was to simplify the memory management process by utilizing a single stream specialize for both malloc and free operations. This approach aims to eliminate the need for complex host-side blocking mechanisms related to CUDA events that used in StreamSafeCUDAAllocator.

When introducing CUDAMallocAsyncAllocator with stream-ordered semantics, why do we still need a complex CUDA event releated mechanism similar to StreamSafeCUDAAllocator involving EventRecord, EventQuery and unfreed Allocation management? Since in Paddle the Allocator malloc and free are dispatched to the same stream as their relevant OP kernel, and the cross-stream synchronization for kernel is guaranteed by upstream code, can we completely transfer all management responsibilities to CUDA, and just do some simple CUDAMalloc/FreeAsync calls in Allocator?

The development of a mechanism designed to fully offload memory management responsibilities to CUDA is currently in progress, utilizing memory streams for this specific aim. This initiative remains under construction and could see significant advancements in forthcoming pull requests (Note here). At present, the necessity for employing APIs akin to EventRecord, EventQuery, and the management of unfreed allocations persists, primarily because the stream passed into the Allocator is the default stream, and the RecordStream is employed to annotate the specific stream that the memory block operates on, this semantic is not fully compatible with stream-order allocator. Our ongoing efforts are focused on addressing these issues, with a commitment to refining and enhancing this aspect in future updates.