flagos-ai / flagos-ai/FlagTree
Ascend backend: decouple torch_npu so PrivateUse1 plugins can use it
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
## 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**:
| Location | Emits |
| --- | --- |
| `backend_register.py:227` | `-I{torch_npu}/include`, `-L{torch_npu}/lib`, `-ltorch_npu` |
| `backend_register.py:293` | `#include ` |
| `backend_register.py:316` | `at_npu::native::allocate_workspace(...)` |
| `backend_register.py:340` | `at_npu::native::OpCommand cmd;` |
| `backend_register.py:305` | allocates on `at::kPrivateUse1` |
| `backend_register.py:270-278` | `torch_npu._C._npu_getCurrentRawStream[NoWait]` (compiled extension) |
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("")` 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.
Contributor guide
Research direction
Start with backend_register.py, utils.py, driver.py, and compiler/hint_manager.py, then trace BackendStrategyRegistry and the listed torch_npu-generated headers, symbols, and link flags. Establish whether a registered non-torch_npu policy can cover the dispatched methods and be selected through TRITON_BACKEND; done means the accepted scope and policy boundary are clear without relying on an unresolved design choice.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python, pytorch
- Domain
- ai, backend, compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100