Ascend backend: decouple torch_npu so PrivateUse1 plugins can use it #1046

Summary

The Ascend backend is coupled to torch_npu at the C++/link level, not just at the Python import level. This makes it unusable from PyTorch out-of-tree device plugins that claim the PrivateUse1 backend key themselves, because torch_npu claims PrivateUse1 on import and only one owner is possible per process.

I would like to use FlagTree's Ascend backend as the Triton layer under torch.compile for such a plugin, and currently cannot. triton-ascend works there precisely because it does not require torch_npu.

All references below are against triton_v3.5.x at d2063b0.

The coupling is not just the imports

The obvious reading is that this is three stray imports (driver.py:231, utils.py:48, backend_register.py:87) which a stub module could satisfy. That is not the case — the torch_npu backend policy also selects generated C++ and link flags:

So the launcher .so genuinely needs libtorch_npu.so and its headers. A Python-level shim cannot work, and backend_register.py:305 targeting at::kPrivateUse1 is itself the conflict rather than a workaround.

Why this blocks plugin-based integrations

PyTorch allows exactly one out-of-tree backend to occupy PrivateUse1. A plugin that calls torch.utils.rename_privateuse1_backend("<name>") and registers its own kernels, allocator and streams cannot coexist with torch_npu in one process — whichever imports first wins, and the other fails to register.

Such a plugin already has everything the Ascend backend needs from torch_npu: current device, current stream as an aclrtStream, and a workspace allocator. What it cannot do is be torch_npu.

What I'd like to ask

The registry already anticipates more than one host runtime — mindspore is a full second policy alongside torch_npu, and BackendStrategyRegistry.register(category, method) accepts any category string. That seems like the right seam. Concretely:

1. Would you accept a third, non-torch_npu policy covering the 11 methods driver.py dispatches (get_current_device, set_current_device, get_current_stream, get_device_interface, get_empty_tensor, get_tensor_params_shape, header_file, get_cc_cmd, allocate_memory, allocate_sync_block_lock, pre_launch, async_launch)? A generic torch_plugin policy that parameterizes the device name and drops the at_npu:: calls would serve any PrivateUse1 plugin, not just mine.

2. Could utils.get_backend_func stop hard-coding the accepted policy names? utils.py:42-44 whitelists only torch_npu and mindspore, so even a correctly registered third policy cannot be selected via TRITON_BACKEND. Accepting any registered category, or at least consulting the registry, would make the extension point usable without patching.

3. Could the two at_npu::native:: uses fall back to CANN APIs? allocate_sync_block_lock and async_launch are the only places that need more than ATen. If the workspace could come from a plain at::empty on the active device and the task-queue path were optional, the policy would need no vendor-specific C++ at all.

4. Smaller, independent: driver.py:229-232 imports torch_npu only so torch.device("npu", ...) is constructible, and its one caller (compiler/hint_manager.py:95) just reads device.type. If the device name were taken from the policy, this import could go away on its own.

I'm happy to prepare a PR for (1) and (2) if that direction is acceptable — please say which shape you'd prefer before I start. Thanks for the work on this backend.