[FEAT] Add Ascend stream support to the Torch C-DLPack addon

[voidLitchi]

Motivation

TVM-FFI uses the DLPack C exchange API’s current_work_stream callback to propagate the framework’s current execution stream into an FFI call. This allows a kernel launched through TVM-FFI to follow the caller’s stream ordering without requiring an explicit synchronization.

On Ascend, torch_npu registers the NPU backend through PyTorch’s PrivateUse1 mechanism. These tensors are currently represented as kDLExtDev when exchanged through DLPack. The optional Torch C-DLPack addon can already convert such tensors, but its current_work_stream implementation previously handled only CUDA and ROCm. An Ascend request therefore returned a null stream and caused the caller to fall back to the default NPU stream instead of using the active torch.npu stream.

This PR extends the optional addon so that Ascend kernels launched through TVM-FFI receive the correct current NPU stream.

Changes

Backend selection and addon loading

• Extend the Torch addon backend selector with an ascend result.
• Keep the existing selection order: CUDA and ROCm are checked first, followed by Ascend, then CPU.
• Detect Ascend only on Linux. Windows and macOS fall back to the CPU addon because neither torch_npu nor the Ascend runtime supports those platforms.
• Add the backend name to the existing ABI-aware addon cache key, preventing an Ascend addon from colliding with a CPU addon built against the same PyTorch package.
• Force the TVM-FFI addon to be used for Ascend even when torch.Tensor already provides __dlpack_c_exchange_api__. The upstream PyTorch API can convert the tensor, but it does not provide the required Ascend current_work_stream behavior.

Ascend addon build support

• Add a mutually exclusive --build-with-ascend build option alongside the existing CUDA and ROCm options.
• Reject the Ascend build option on Windows and macOS with an explicit argument error.
• Discover the installed torch_npu package and add its header and library directories to the build.
• Define BUILD_WITH_ASCEND only for the Ascend addon.
• Link against libtorch_npu, which exports the c10_npu stream symbols used by the addon.

Current NPU stream propagation

When built with Ascend support, the addon now handles kDLExtDev in its current_work_stream callback and returns:

c10_npu::getCurrentNPUStream(device_id).stream()

This keeps TVM-FFI kernel launches on the active torch.npu stream, including non-default streams selected through torch.npu.stream(...).

Why Ascend detection has a separate torch.npu branch

ROCm follows PyTorch’s AMD compatibility model and intentionally reuses the torch.cuda Python API. CUDA and ROCm can therefore share the same initial availability check and be distinguished using torch.version.cuda and torch.version.hip, as documented in the [PyTorch ROCm semantics](https://docs.pytorch.org/docs/stable/notes/hip.html).

torch_npu follows a different integration model. It is an out-of-tree extension that registers Ascend through PrivateUse1 and exposes its runtime API as torch.npu, rather than reusing torch.cuda. This model is also described in PyTorch’s [device-extension autoload RFC](pytorch/pytorch#122468).

As a result, Ascend cannot be folded into the CUDA/ROCm branch without changing the supported PyTorch/extension compatibility model. A generic torch.accelerator-based implementation was considered, but adopting it here would broaden the change across existing backends and supported PyTorch versions. The guarded:

hasattr(torch, "npu") and torch.npu.is_available()

check is therefore kept as an Ascend-specific compatibility path. CUDA and ROCm retain priority when their shared torch.cuda interface is active.

DLPack device type follow-up

This PR intentionally continues to use kDLExtDev for Ascend.

A dedicated kDLAscend device type will be proposed to the DLPack project so that Ascend memory is no longer represented by the generic extension-device value. That migration is not included here because it requires coordinated changes across DLPack, TVM-FFI, PyTorch/torch_npu, and downstream consumers.

At the time of writing:

• TVM-FFI pins the released [DLPack v1.3](https://github.com/dmlc/dlpack/releases/tag/v1.3) commit 84d107b.
• DLPack main is three commits ahead at [77aafa4](https://github.com/dmlc/dlpack/commit/77aafa4d3b0f80feffce9ad4c718dd26751ee0e4) and has already allocated values 19 and 20 to kDLTPU and kDLTPUHost.
• Consequently, the numeric value and semantics of kDLAscend should be standardized in DLPack first instead of being provisionally assigned by TVM-FFI.

After such an enum is accepted, TVM-FFI would need to update more than the vendored header. The required work includes:

• Updating the DLPack submodule.
• Extending C++ device parsing and formatting.
• Extending the Python DLDeviceType enum, name mappings, type stubs, and device tests.
• Updating the Torch addon’s PrivateUse1 ↔ DLPack conversion logic.
• Updating stream lookup to recognize the new type.
• Providing a transition period in which both kDLExtDev and kDLAscend are accepted, because existing torch_npu versions will continue to produce kDLExtDev.

Switching this PR directly to a new enum would therefore break compatibility with existing producers without completing the corresponding ecosystem migration. Keeping kDLExtDev here makes the stream fix usable immediately while leaving the enum standardization to a separate follow-up.

Testing

The PR adds coverage for:

• CPU, CUDA, ROCm, and Ascend backend selection.
• CUDA/ROCm precedence when an Ascend namespace is also present.
• Windows and macOS fallback behavior.
• Reuse of PyTorch’s existing exchange API on non-Ascend backends.
• Forced addon selection on Ascend.
• Mutually exclusive CUDA, ROCm, and Ascend build options.
• Rejection of Ascend builds on unsupported platforms.
• Building and loading the addon from a fresh per-test output directory.
• An Ascend runtime test that extracts the installed DLPack exchange API and compares current_work_stream(kDLExtDev, device_id) against a non-default torch.npu.Stream handle.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[tqchen]

Thanks @voidLitchi . We are increasingly moving towards upstream to bring up the convention which is more robust in terms of ABI version and use the optional plugin mainly for backward compact cases.

Ideally the best places the DLPack exchange API location definition should be contributed to the upstream package(in this case torch_npu package as part of its dlpack module).

the overall mechanism is torch_npu get autoloaded as a torch backend. It can bundle this variant of the C DLPack Exchange API in then override torch.Tensor.__dlpack_c_exchange_api__ during its autoloading process

[tqchen]

BUILD_WITH_TORCH_NPU

[voidLitchi]

Thank you for reviewing! I have already renamed all ascend to torch_npu as you suggested in 1403e7f.

[tqchen]

build_with_torch_npu

[voidLitchi]

Updated to build_with_torch_npu in the same commit.

[tqchen]

thanks @voidLitchi just want to follow up on my comment in implementation goes to torch_npu, if that is possible it would be more ideal and likely we won't need these extra changes here

[voidLitchi]

Thanks @voidLitchi . We are increasingly moving towards upstream to bring up the convention which is more robust in terms of ABI version and use the optional plugin mainly for backward compact cases.

Ideally the best places the DLPack exchange API location definition should be contributed to the upstream package(in this case torch_npu package as part of its dlpack module).

the overall mechanism is torch_npu get autoloaded as a torch backend. It can bundle this variant of the C DLPack Exchange API in then override torch.Tensor.__dlpack_c_exchange_api__ during its autoloading process

Thanks for the suggestion. We discussed this with the Ascend community, and we agree that torch_npu is the more appropriate home for this capability.
We will investigate the integration on the torch_npu side and work toward implementing it there. In the meantime, we would prefer to keep this PR open as a reference while we validate that approach and confirm that it fully covers the relevant downstream use cases. If that path proves sufficient, we can close this PR in favor of the torch_npu implementation.