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¶
使用 DLPack 从 CUDA 内存指针创建 PyTorch 张量。
- 参数:
ptr – CUDA 内存指针地址，整数类型
shape – 期望的张量形状
dtype – PyTorch 数据类型
device_id – CUDA 设备 ID
- 返回值:
封装 CUDA 内存的 PyTorch 张量