(tvm-ffi) branch main updated: [FIX] Use DLPack v1 dtype mapping for Torch allocation (#914)

[tqchen.apache.org via commits]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git


The following commit(s) were added to refs/heads/main by this push:
     new 2cb94ecf [FIX] Use DLPack v1 dtype mapping for Torch allocation (#914)
2cb94ecf is described below

commit 2cb94ecf730d8e404d4845047328d6dcef80e4fa
Author: Lei Wang <34334180+leiwang1...@users.noreply.github.com>
AuthorDate: Wed Nov 4 20:10:38 2026 +0800

    [FIX] Use DLPack v1 dtype mapping for Torch allocation (#914)
    
    The Torch addon's `ManagedTensorAllocator` still calls PyTorch's
    `at::toScalarType`, while tensor import already uses the addon's
    `toScalarTypeForDLPackv1`. In PyTorch 2.8, the former does not recognize
    the FP8 DLPack codes, so allocating an FP8 output fails with
    `Unsupported code` even though the addon can import tensors of that
    dtype.
    
    Use the existing DLPack v1 mapper for output allocation as well.
    Ordinary scalar types retain their mappings, and the allocator can use
    the FP8 and packed FP4 mappings already supported by the addon.
    
    Also keep the ordinary `lanes == 1` check active on older PyTorch
    versions. It is currently enclosed in a PyTorch 2.8+ guard, so switching
    the allocator to this helper alone would allow inputs such as `float32`
    with `lanes=2` to be silently treated as scalar `float32` on PyTorch
    2.4–2.7. Remove that outer guard; the FP4 enum comparison uses the
    bundled DLPack header, while the branch referencing PyTorch's packed FP4
    type remains version-gated and checks `bits=4, lanes=2`.
    
    Existing cached addon binaries need to be rebuilt to pick up this
    change.
---
 python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py 
b/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py
index 38d3ed0b..70381159 100644
--- a/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py
+++ b/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py
@@ -255,12 +255,10 @@ static Device getATenDeviceForDLPackv1(DLDeviceType type, 
c10::DeviceIndex index
 
 ScalarType toScalarTypeForDLPackv1(const DLDataType& dtype) {
   ScalarType stype = ScalarType::Undefined;
-#if TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 8
   if (dtype.code != DLDataTypeCode::kDLFloat4_e2m1fn) {
     TORCH_CHECK(dtype.lanes == 1, "ATen does not support lanes != 1 for dtype 
code",
                 std::to_string(dtype.code));
   }
-#endif
   switch (dtype.code) {
     case DLDataTypeCode::kDLUInt:
       switch (dtype.bits) {
@@ -547,7 +545,7 @@ struct TorchDLPackExchangeAPI : public DLPackExchangeAPI {
       at::IntArrayRef shape(prototype->shape, prototype->shape + 
prototype->ndim);
       at::TensorOptions options =
           at::TensorOptions()
-              .dtype(at::toScalarType(prototype->dtype))
+              .dtype(at::toScalarTypeForDLPackv1(prototype->dtype))
               
.device(at::getATenDeviceForDLPackv1(prototype->device.device_type,
                                                    
prototype->device.device_id));
       at::Tensor tensor = at::empty(shape, options);