pytorch / pytorch/pytorch

Allow privateuse1 devices in FlexAttention device validation

Open
#173,071 4 comments 0 reactions 0 assignees View on GitHub
module: flex attention oncall: pt2 triaged
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.