flagos-ai / flagos-ai/FlagTree

Ascend backend: decouple torch_npu so PrivateUse1 plugins can use it

Open
#1,046 0 comments 0 reactions 0 assignees View on GitHub
ascend
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.