Allow privateuse1 devices in FlexAttention device validation
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🚀 The feature, motivation and pitch
FlexAttention currently restricts supported devices via a hard-coded device-type allowlist in `_validate_device`, which excludes `privateuse1` devices.
This prevents FlexAttention from being used on third-party or experimental backends that integrate with PyTorch via the `privateuse1` mechanism, even when the backend provides compatible kernels or relies on Inductor/Triton lowering.
This feature request proposes allowing `privateuse1` devices to pass the device validation step, deferring actual support checks to the backend or kernel implementation.
The intention is not to claim official support for any private backend, but to avoid prematurely blocking valid execution paths at the Python layer and to better align FlexAttention with PyTorch’s extensibility model.
### Alternatives
An alternative would be to keep the current device allowlist and require private backends to modify FlexAttention locally.
However, this approach does not scale well for third-party backends and requires downstream forks or patches, whereas allowing `privateuse1` keeps the extensibility consistent with other PyTorch operators.
### Additional context
Relevant code:
https://github.com/pytorch/pytorch/blob/main/torch/nn/attention/flex_attention.py#L1318
```python
supported_devices = {"cuda", "cpu", "xpu", "hpu"}
cc @chauhang @penguinwu @Chillee @drisspg @yanboliang @BoyuanFeng
Contributor guide
Assessment
This issue has not been assessed yet.