flashinfer.comm.mnnvl.create_tensor_from_cuda_memory

flashinfer.comm.mnnvl.create_tensor_from_cuda_memory(ptr: int, shape: tuple, dtype: dtype, device_id: int) → Tensor

Create a PyTorch tensor from a CUDA memory pointer using DLPack.

Parameters:

  • ptr – CUDA memory pointer address as integer

  • shape – Desired tensor shape

  • dtype – PyTorch data type

  • device_id – CUDA device ID

Returns:

PyTorch tensor that wraps the CUDA memory