From de198b24197841ca7e52f7f114db94f66ff08382 Mon Sep 17 00:00:00 2001 From: yuanxiaolan Date: Thu, 22 Jan 2026 13:47:45 +0800 Subject: [PATCH 1/2] fix tvm-ffi path in from_dlpack --- python/CuTeDSL/cutlass/cute/runtime.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/CuTeDSL/cutlass/cute/runtime.py b/python/CuTeDSL/cutlass/cute/runtime.py index 87082656..d4f8c554 100644 --- a/python/CuTeDSL/cutlass/cute/runtime.py +++ b/python/CuTeDSL/cutlass/cute/runtime.py @@ -141,6 +141,11 @@ class _Tensor(Tensor): # If tensor is already a DLPack object, use it directly if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"): self._dlpack_data = tensor.__dlpack_device__() + elif enable_tvm_ffi: + import tvm_ffi + + self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor) + self._dlpack_data = self._tvm_ffi_tensor.__dlpack__() else: try: # we expect no stream sync. Because torch has different default behavior @@ -149,11 +154,6 @@ class _Tensor(Tensor): self._dlpack_data = tensor.__dlpack__(stream=-1) except Exception: self._dlpack_data = tensor.__dlpack__() - if enable_tvm_ffi: - import tvm_ffi - - self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor) - self._dlpack_data = self._tvm_ffi_tensor.__dlpack__() self._dltensor_wrapper = None self._assumed_align = assumed_align From de161925a50e7d96bb36ed549a09bf2769e64cc5 Mon Sep 17 00:00:00 2001 From: yuanxiaolan Date: Tue, 3 Feb 2026 11:55:54 +0800 Subject: [PATCH 2/2] pass in stream=-1 --- python/CuTeDSL/cutlass/cute/runtime.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/CuTeDSL/cutlass/cute/runtime.py b/python/CuTeDSL/cutlass/cute/runtime.py index d4f8c554..4019f1b6 100644 --- a/python/CuTeDSL/cutlass/cute/runtime.py +++ b/python/CuTeDSL/cutlass/cute/runtime.py @@ -144,7 +144,7 @@ class _Tensor(Tensor): elif enable_tvm_ffi: import tvm_ffi - self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor) + self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor, stream=-1) self._dlpack_data = self._tvm_ffi_tensor.__dlpack__() else: try: