add support for Intel XPU platform

[kaixuanliu]

No description provided.

[kaixuanliu]

@alihassanijr @lartpang , Hi, we are from Intel and want to add Intel XPU platform support for this project. Can you help review this PR? Many thanks!!

[yao-matrix]

@alihassanijr @lartpang , for more context, Intel XPU(which is Intel's GPU) is a built-in device of PyTorch since PyTorch 2.5 and with Flex Attention also enabled and validated. So, we'd like to enable natten on XPU platform to make XPU users can run natten without friction on it. Thx very much.

[alihassanijr]

Thank you for your contribution.

Could you please confirm if you've also run the Flex unit tests on XPU?
I don't have an xpu device myself to verify correctness.

Could you also share the version of torch you're using? If xpu support in Flex was added in a specific version we'd want to guard against it.

[yao-matrix]

Thank you for your contribution.

Could you please confirm if you've also run the Flex unit tests on XPU? I don't have an xpu device myself to verify correctness.

Could you also share the version of torch you're using? If xpu support in Flex was added in a specific version we'd want to guard against it.

Sure. @kaixuanliu pls run flex-atten related UT on XPU and report results here, and pls add version guard for flex attention.

[kaixuanliu]

Hi @alihassanijr , I can run Flex unit tests successfully on XPU w/ adjustment in L125-L126, I need to change reference_backend to flex-fna and reference_fmha_backend to flex-fmha, and need some other minor device-specific adjustment. And the pytorch version I use is 2.9.0+xpu, it can be installed by pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/xpu. And Flex Attn should be supported after pytorch 2.9 for Intel XPU platform, I have added related version check in this PR, pls help review, thx!

[yao-matrix]

@alihassanijr , could you pls take a second look? Thx very much

[alihassanijr]

@yao-matrix @kaixuanliu apologies for the delay.

It looks like you're testing the flex backend against itself if you changed only those lines.

The only alternative I can think of, since our reference is also in CUDA C++, is to test flex XPU against flex CPU, and that might require some additional work in the test suite.

Could you take a shot at that and push the tests as well?

Also, feel free to make a separate test for just xpu like test_flex_xpu.py -- that might be easier to make and maintain in the long run.

[kaixuanliu]

@alihassanijr Thx! Compare with flex-fna CPU version is a good advice, I have updated the testing part following your suggestion, and tests passed in my environment. Pls have a review again.

[yao-matrix]

@alihassanijr , happy new year, could you help take a second look at this PR? Thx very much.

[yao-matrix]

@kaixuanliu , could you pls rebase the PR? Thx.

[kaixuanliu]

Sure, have done.

[yao-matrix]

@alihassanijr , could you help take a second look? Thx very much.

[yao-matrix]

@alihassanijr , could you help take a second look? Thx very much.