Can I construct a PyTorch device tensor with cudart.cudaMalloc pointer?

[MARD1NO]

If I malloc a buffer by using cudart.cudaMalloc, it will return a dev_ptr.

And I want to print it in my python program, so I want to construct a Torch tensor with this dev_ptr, is this achieveable?...

[rwgk]

Trying to be helpful, although I'm new here:

• https://chatgpt.com/share/67e8405e-4768-8008-a076-fc92ada875bb

That suggests:

import torch
import dlpack   # hypothetical—depends on your CUDA binding

dl_tensor = dlpack.wrap(dev_ptr, shape, dtype, device='cuda')
torch_tensor = torch.utils.dlpack.from_dlpack(dl_tensor)

[MARD1NO]

it seems dlpack do not have wrap function, I change to use cupy:

import cupy as cp 
cupy_memory = cp.cuda.Memory(your_buffer_size)
cupy_memory.ptr = your_buffer

memptr = cp.cuda.MemoryPointer(cupy_memory, 0)

cupy_array = cp.ndarray(
    your_shape,
    dtype=cp.float16,
    memptr=memptr,
)

[leofang]

Yes, using CuPy is one option, and then you can "launder" it to a torch.Tensor through either torch.as_tensor() or torch.from_dlpack().

However, if your code uses PyTorch heavily, I recommend not to do that because it could lead to unnecessarily out of memory. Each GPU project has its own memory allocator, and so if you really need to allocate raw GPU memory, you should try to use the low-level APIs provided by your framework of choice, ex: https://pytorch.org/docs/stable/generated/torch.cuda.caching_allocator_alloc.html.