flagos-ai / flagos-ai/FlagTree
fix(nvidia): nvidia driver is_active() fires on Ascend NPU under torch_npu TORCH_TRANSFER_TO_NPU=1
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
### Bug description
`third_party/nvidia/backend/driver.py` `CUDABackend.is_active()` reports the nvidia backend active on machines that have **no NVIDIA GPU** when `torch_npu` is loaded with the transfer-to-NPU shim (`TORCH_TRANSFER_TO_NPU=1`):
```python
return torch.cuda.is_available() and (torch.version.hip is None)
```
On Ascend NPU hosts the shim makes `torch.cuda.is_available()` return `True` (it redirects CUDA API calls to the NPU), and the torch build is not a HIP build, so `torch.version.hip is None` also holds. The nvidia driver therefore reports active on a pure NPU host. The Ascend `NPUDriver.is_active()` (bisheng `npucompiler` probe) is *also* active on that host, so driver selection sees two active backends and `triton.runtime.driver._create_driver()` fails.
### Environment
- Ascend 910B, CANN 9.0.0, `torch_npu` with `TORCH_TRANSFER_TO_NPU=1`
- FlagTree runtime wheel (Triton-based), nvidia driver at `third_party/nvidia/backend/driver.py:876` (main @ 52678a8b7)
### Steps to reproduce
Run any Triton workload on an Ascend host with `TORCH_TRANSFER_TO_NPU=1` set (the standard torch_npu shim mode). Driver probing enumerates backends via `is_active()`; nvidia returns `True` (CUDA-available faked by the shim, not a HIP torch) while the ascend driver also returns `True`, so driver selection errors instead of resolving to `NPUDriver`.
### Expected behavior
`is_active()` should test for a *real* CUDA torch build, not just CUDA-API availability. A genuine CUDA build always carries a non-`None` `torch.version.cuda`; the shim does not fake it. On the Ascend host the check should return `False` (observed after the fix below: `amd=False, ascend=True, nvidia=False` → `NPUDriver` selected).
### Proposed fix
```python
return (
torch.cuda.is_available()
and torch.version.cuda is not None
and (torch.version.hip is None)
)
```
This matches the guard upstream triton uses in its nvidia driver. No behavior change on real CUDA/HIP hosts (a CUDA torch always has `version.cuda` set).
Contributor guide
Research direction
Start in third_party/nvidia/backend/driver.py at CUDABackend.is_active(), then compare its guard with the upstream Triton NVIDIA driver guard mentioned in the issue. Reproduce with torch_npu and TORCH_TRANSFER_TO_NPU=1 on an Ascend host; done means the NVIDIA backend reports inactive and driver selection resolves to NPUDriver.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100