dlpack issue with PyTorch

[wxianghao]

Description

Hello, the newest release (25.10.0) says dlpack is supported. I was attempting to importing torch tensor on GPU to cupynumeric, but got an error. Also, the document misses the usage of dlpack support.

Environment

torch: 2.8.0
legate: 25.10.0
nvidia-cupynumeric: 25.10.0

Code and running command

'''Within one GPU'''

import cupylma
import cupynumeric as np
import torch
import argparse

from cupynumeric import from_dlpack
from cupylma import get_available_gpus

device = get_available_gpus()[0]

# Create tensor on torch
num_bytes = 5_000_000_000
bytes_per_element = 4
num_elements = (num_bytes + bytes_per_element - 1) // bytes_per_element
tensor_shape = (num_elements,)
tensor = torch.empty(tensor_shape, dtype=torch.float32, device=device)

# Send tensor to cupynumeric
cp_array = from_dlpack(tensor, copy=True)

legate --gpus 1 test.py

Error

Traceback (most recent call last):
  File "/home/659/xw6261/test-dlpack/send.py", line 21, in <module>
    cp_array = from_dlpack(tensor, copy=True)
TypeError: Tensor.__dlpack__() got an unexpected keyword argument 'dl_device'

[manopapad]

I believe this should be fixed on the upcoming 25.11 (and the fix should be available on the latest nightlies). @bryevdv could you please verify? I think you also added dlpack (and other missing methods) to the docs?

[bryevdv]

I am unable to run the code provided, more information about "cupylma" is needed

dev ❯ legate --gpus 1 ~/tmp/foo.py
Traceback (most recent call last):
  File "/home/bryan/tmp/foo.py", line 7, in <module>
    from cupylma import get_available_gpus
ImportError: cannot import name 'get_available_gpus' from 'cupylma' (/home/bryan/anaconda3/envs/dev/lib/python3.13/site-packages/cupylma/__init__.py)

This is the version I obtained from pip install cupylma

cupylma                    0.1.2            pypi_0           pypi

[bryevdv]

I guess it needs a pre-release version. @manopapad I get the following different error with the latest dlpack work merged

dev ❯ legate --gpus 1 ~/tmp/foo.py
Traceback (most recent call last):
  File "/home/bryan/tmp/foo.py", line 19, in <module>
    cp_array = from_dlpack(tensor, copy=True)
  File "/home/bryan/anaconda3/envs/dev/lib/python3.13/site-packages/torch/_tensor.py", line 1726, in __dlpack__
    raise BufferError(
    ...<3 lines>...
    )
BufferError: Can't export tensors on a different CUDA device index. Expected: 1. Current device: 0.

[bryevdv]

If I manually override the device from cuda:1 to cuda:0 then see another different error:

dev ❯ legate --gpus 1 ~/tmp/foo.py
Traceback (most recent call last):
  File "/home/bryan/tmp/foo.py", line 21, in <module>
    cp_array = from_dlpack(tensor, copy=True)
RuntimeError: Unsupported device in DLTensor.

Which seems like torch passing a device arg that legate does not expect.

[manopapad]

@bryevdv can you also try setting CUDA_VISBILE_DEVICES=0, to (hopefully" ensure that pytorch only uses the first device? --gpus 1 constrains Legate but not PyTorch.

[bryevdv]

With that env set pytorch still picks up cuda:1 instead of cuda:0

torch ❯ CUDA_VISBILE_DEVICES=0 legate --gpus 1 ~/tmp/foo.py
cuda:1
Traceback (most recent call last):
  File "/home/bryan/tmp/foo.py", line 20, in <module>
    cp_array = from_dlpack(tensor, copy=True)
  File "<cupynumeric internals>", line 200, in from_dlpack_implemented_wrapper
  File "/home/bryan/work/cupynumeric.internal/cupynumeric/_dlpack/__init__.py", line 55, in from_dlpack
    store = from_dlpack(x, device=device, copy=copy)
  File "legate/core/_lib/utilities/detail/dlpack/from_dlpack.pyx", line 93, in legate.core._lib.utilities.detail.dlpack.from_dlpack.from_dlpack
  File "/home/bryan/anaconda3/envs/torch/lib/python3.13/site-packages/torch/_tensor.py", line 1726, in __dlpack__
    raise BufferError(
    ...<3 lines>...
    )
BufferError: Can't export tensors on a different CUDA device index. Expected: 1. Current device: 0.

[manopapad]

OK, then I'd ask @wang-xianghao to try with latest nightly packages and tell us if it solves the original problem. We seem to be stumped by pytorch's (or possibly cupylma's) behavior.

[wxianghao]

Hello @manopapad. Sorry for the previous confusion brought by cupylma and here is the code depending only on torch and cupynumeric. I tested on the newest version (25.10.00, cuda12_py313_gpu_g66d872d2_0) and got an different error by running legate --gpus 1 send.py.

Traceback (most recent call last):
  File "/home/659/xw6261/test-dlpack/send.py", line 19, in <module>
    cp_array = from_dlpack(tensor, copy=True)
RuntimeError: Unsupported device in DLTensor.

'''Within one GPU'''

import cupynumeric as np
import torch
import argparse

from cupynumeric import from_dlpack

device = torch.device('cuda:0')

# Create tensor on torch
num_bytes = 5_000_000
bytes_per_element = 4
num_elements = (num_bytes + bytes_per_element - 1) // bytes_per_element
tensor_shape = (num_elements,)
tensor = torch.empty(tensor_shape, dtype=torch.float32, device=device)

# Send tensor to cupynumeric
cp_array = from_dlpack(tensor, copy=True)

[wxianghao]

Also, cupynumeric and torch are both hosted on the same GPU device. If I change torch.device('cuda:0') to torch.device('coda:1') which puts the torch tensor and the cupynumeric tensor on different devices, I got an error raised by torch.

From this behavior It seems that even if dlpack is supported, it only makes sense when both torch and cupynumeric target at the same device.

Since cupynumeric highlight multi-GPU acceleration, I was wondering if in the future cupynumeric supports to smoothly transfer and distribute data from other framework to multiple GPUs.

Traceback (most recent call last):
  File "/home/659/xw6261/test-dlpack/send.py", line 19, in <module>
    cp_array = from_dlpack(tensor, copy=True)
  File "/scratch/um09/xw6261/miniconda3/envs/cupynumeric-latest/lib/python3.13/site-packages/torch/_tensor.py", line 1726, in __dlpack__
    raise BufferError(
    ...<3 lines>...
    )
BufferError: Can't export tensors on a different CUDA device index. Expected: 1. Current device: 0.

[manopapad]

@wang-xianghao could you please try with latest nightly build? Your example seems to work for me with:

conda create -n 2025-11-26 -c conda-forge -c legate-nightly cupynumeric

From this behavior It seems that even if dlpack is supported, it only makes sense when both torch and cupynumeric target at the same device.

Some explanation on what I believe is happening here. By asking PyTorch to use torch.device('coda:1'), it instantiates the data on GPU 1. Then later cuPyNumeric is asking for the data to be exported on GPU 0. PyTorch has not currently implemented cross-device export (which requires PyTorch to manage the copying between GPUs). Therefore, the only thing we can potentially do is query on which device PyTorch has placed the data (I assume PyTorch exposes __dlpack_device__), and make sure we ask it to be exported on that device only. Wdyt @Jacobfaib ?

Since cupynumeric highlight multi-GPU acceleration, I was wondering if in the future cupynumeric supports to smoothly transfer and distribute data from other framework to multiple GPUs.

Please note, once the data is imported to cuPyNumeric we can re-shard as necessary. The issue here is the initial handshake between PyTorch and cuPyNumeric through dlpack, which is currently not being coordinated properly in all cases.

[magnatelee]

@wang-xianghao looking at the issue closely, I believe this is a bug in PyTorch's dlpack implementation not abiding by the passed dl_device; this line checks the need for cross-device export not based on the dl_device but the current device, which I believe is incorrect. Legate's from_dlpack() implementation simply passes what's returned from __dlpack_device__() to __dlpack__(), which shouldn't incur any cross-device exports. A workaround is wrapping the from_dlpack() call in your original two-GPU code with a context manager like this:

device = torch.device('cuda:1')

...

with torch.cuda.device('cuda:1'):
  cp_array = from_dlpack(tensor, copy=True)

@Jacobfaib @manopapad though this isn't an issue in Legate, we could potentially put in this PyTorch-specific workaround to our code. what do you think?