Cythonize `Buffer`

[leofang]

From @shwina

In [15]: mr = DeviceMemoryResource(Device().device_id)

In [16]: %timeit mr.allocate(10)
3.35 μs ± 95.1 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [17]: %timeit cp.empty(10, dtype=cp.int8)
1.48 μs ± 11.4 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

[shwina]

Unfortunately, the problem goes way beyond just Python overheads due to Buffer construction. If we actually sync the device, we see a huge perf gap:

In [1]: from cuda.core.experimental import Device, DeviceMemoryResource

In [2]: import cupy as cp

In [3]: mr = DeviceMemoryResource(Device().device_id)

In [4]: Device().set_current()

In [5]: %timeit cp.empty(10); cp.cuda.runtime.deviceSynchronize()
1.93 μs ± 116 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [6]: %timeit mr.allocate(10); Device().sync()
646 μs ± 1.14 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Using cuda.bindings directly:

In [22]: %timeit (_, ptr) = cudaMallocAsync(10, None); cudaFreeAsync(ptr, None)
625 ns ± 3.99 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

[kkraus14]

Using cuda.bindings directly:

In [22]: %timeit (_, ptr) = cudaMallocAsync(10, None); cudaFreeAsync(ptr, None)
625 ns ± 3.99 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

I think this still runs asynchronously where if you synchronize the device with cudaDeviceSynchronize() in the timeit loop it will show similar performance to using mr.allocate?

I think the reason that cupy is so much faster here is because they don't use the driver managed memory pool by default, they have their own memory pool implementation that they default to. @leofang probably knows the implementation details here, but I'm guessing they allocate in a blocksize and suballocate from blocks, whereas the async pool doesn't.

We could try using cudaMemPoolSetAttribute to set cudaMemPoolAttrReleaseThreshold to a large number and then making a large allocation to prime the pool before running this timeit loop. I think this is what RMM effectively does: https://github.com/rapidsai/rmm/blob/branch-25.08/cpp/include/rmm/mr/device/cuda_async_memory_resource.hpp#L133-L147

[shwina]

I think this still runs asynchronously where if you synchronize the device with cudaDeviceSynchronize() in the timeit loop it will show similar performance to using mr.allocate?

Adding a sync brings the perf closer to CuPy's empty:

In [27]: %timeit (_, ptr) = cudaMallocAsync(10, None); cudaFreeAsync(ptr, None); Device().sync()
1.94 μs ± 128 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)