[BugFix][ROCm] Prefer upstream PyTorch DLPack API in torch extension loader

[zihaomu]

Motivation:

• If PyTorch already provides __dlpack_c_exchange_api__, tvm-ffi uses it directly on all backends.
• The optional torch-c-dlpack-ext path is only used as a fallback for older PyTorch builds that do not provide the API.
• The fallback extension library name is selected with backend-aware detection for CPU, CUDA, and ROCm.

Tests are added for backend detection, ROCm short-circuit behavior when PyTorch already provides the API, and GPU tensor metadata through DLPack.

Related PR: tile-ai/tilelang#2179, I have finished the A/B test locally.

[gemini-code-assist]

Code Review

This pull request refactors the PyTorch backend detection logic into helper functions within the torch_c_dlpack_ext and tvm_ffi modules to better support ROCm alongside CUDA and CPU. It also adds an exception handler for OSError during extension loading to gracefully handle incompatible prebuilt libraries and introduces new tests for GPU tensor metadata. The review feedback highlights potential AttributeError risks when accessing version attributes directly on PyTorch builds missing specific backend support, recommending the use of getattr for safer detection.

[tqchen]

use the same code as

def _torch_extension_device(torch_module: Any) -> str:
    """Return the torch backend name used in the optional extension library name."""
    if torch_module.cuda.is_available():
        if torch_module.version.cuda is not None:
            return "cuda"
        if torch_module.version.hip is not None:
            return "rocm"
        raise ValueError("Cannot determine whether to build with CUDA or ROCm.")
    return "cpu" 

[tqchen]

assuming we detect correctly, is this still necessary?

[tqchen]

Mainly want to know if it is something as a temp measure that we can delete after a few version cycles, if so, please add ad comment

[tqchen]

Mainly want to know if it is something as a temp measure that we can delete after a few version cycles, if so, please add ad comment

[zihaomu]

Hi @tqchen. The scope of this PR has shifted from fixing ROCm-specific Torch C DLPack extension loading to making tvm-ffi avoid loading the optional out-of-tree torch-c-dlpack-ext package when upstream PyTorch already provides the C DLPack exchange API.

If newer PyTorch ROCm builds already expose a working __dlpack_c_exchange_api__, in that case, tvm-ffi should trust PyTorch and return early instead of forcing a local extension override. The local extension remains only as a compatibility fallback for older PyTorch versions

[tqchen]

thanks, looking good after fixing some minor nits

[tqchen]

would be good to use the following pattern, mainly also detect cuda version to be robust.

 if torch_module.cuda.is_available():
    if getattr(torch_module.version, "cuda", None) is not None:
       return "cuda"
    if getattr(torch_module.version, "hip", None) is not None:
       return "rocm"
  return "cuda"
return "cpu"

[zihaomu]

Done

[tqchen]

 if torch_module.cuda.is_available():
    if getattr(torch_module.version, "cuda", None) is not None:
       return "cuda"
    if getattr(torch_module.version, "hip", None) is not None:
       return "rocm"
  return "cuda"
return "cpu"

[zihaomu]

Done