add xpu support

[mgrabban]

Summary

Adds xpu support so all tests, benchmarks etc. run on XPUs or Intel GPUs.

Details

infer_device() function is moved to a separate file and in any file where previously "cuda" was needed, infer_device is imported and "cuda" is replaced with return value of a call to infer_device()

Testing Done

A100 80GB PCIe, RTX 3060, Intel Data Center GPU Max 1550

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[faaany]

I ran the UTs on XPU, but got "Segmentation fault (core dumped)" at one test, under investigation.

[mgrabban]

I ran the UTs on XPU, but got "Segmentation fault (core dumped)" at one test, under investigation.

Which specific Intel GPU did you test on?
Also does the test run if you just change "cuda" to "xpu" manually without using this PR?

Also: I just added xpu support to simpo_loss (which was added later on and still had "cuda" hard coded).

[faaany]

I ran the UTs on XPU, but got "Segmentation fault (core dumped)" at one test, under investigation.

Which specific Intel GPU did you test on? Also does the test run if you just change "cuda" to "xpu" manually without using this PR?

Also: I just added xpu support to simpo_loss (which was added later on and still had "cuda" hard coded).

I use "Intel Data Center GPU Max 1550".

And I tested your latest code. All tests pass except "pytest -rA test/transformers/test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-2-128-512]", but this one is a known issue and got fixed in the latest pytorch-triton-xpu. Don't you have this issue?

[mgrabban]

And I tested your latest code. All tests pass except "pytest -rA test/transformers/test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-2-128-512]", but this one is a known issue and got fixed in the latest pytorch-triton-xpu. Don't you have this issue?

I don't have this issue. It could be because I use nightly intel-xpu-backend-for-triton.

test_rms_norm.py::test_correctness[True-LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                           [  3%]
test_rms_norm.py::test_correctness[True-LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                           [  6%]
test_rms_norm.py::test_correctness[True-LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-2-128-512] PASSED                                                                               [  9%]
test_rms_norm.py::test_correctness[True-LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-5-123-123] PASSED                                                                               [ 12%]
test_rms_norm.py::test_correctness[True-GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                           [ 15%]
test_rms_norm.py::test_correctness[True-GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                           [ 18%]
test_rms_norm.py::test_correctness[True-GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-2-128-512] PASSED                                                                               [ 21%]
test_rms_norm.py::test_correctness[True-GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-5-123-123] PASSED                                                                               [ 25%]
test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                             [ 28%]
test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                             [ 31%]
test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-2-128-512] PASSED                                                                                 [ 34%]
test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-5-123-123] PASSED                                                                                 [ 37%]
test_rms_norm.py::test_correctness[False-LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                          [ 40%]
test_rms_norm.py::test_correctness[False-LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                          [ 43%]
test_rms_norm.py::test_correctness[False-LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-2-128-512] PASSED                                                                              [ 46%]
test_rms_norm.py::test_correctness[False-LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-5-123-123] PASSED                                                                              [ 50%]
test_rms_norm.py::test_correctness[False-GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                          [ 53%]
test_rms_norm.py::test_correctness[False-GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                          [ 56%]
test_rms_norm.py::test_correctness[False-GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-2-128-512] PASSED                                                                              [ 59%]
test_rms_norm.py::test_correctness[False-GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-5-123-123] PASSED                                                                              [ 62%]
test_rms_norm.py::test_correctness[False-BaseRMSNorm-0.0-none-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                            [ 65%]
test_rms_norm.py::test_correctness[False-BaseRMSNorm-0.0-none-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                            [ 68%]
test_rms_norm.py::test_correctness[False-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-2-128-512] PASSED                                                                                [ 71%]
test_rms_norm.py::test_correctness[False-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-5-123-123] PASSED                                                                                [ 75%]
test_rms_norm.py::test_correctness_functional[LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-2-2-8] PASSED                                                                         [ 78%]
test_rms_norm.py::test_correctness_functional[LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-9-7-41] PASSED                                                                        [ 81%]
test_rms_norm.py::test_correctness_functional[LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-2-2-8] PASSED                                                                             [ 84%]
test_rms_norm.py::test_correctness_functional[LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-9-7-41] PASSED                                                                            [ 87%]
test_rms_norm.py::test_correctness_functional[GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-2-2-8] PASSED                                                                         [ 90%]
test_rms_norm.py::test_correctness_functional[GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-9-7-41] PASSED                                                                        [ 93%]
test_rms_norm.py::test_correctness_functional[GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-2-2-8] PASSED                                                                             [ 96%]
test_rms_norm.py::test_correctness_functional[GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-9-7-41] PASSED                                                                            [100%]

[faaany]

And I tested your latest code. All tests pass except "pytest -rA test/transformers/test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-2-128-512]", but this one is a known issue and got fixed in the latest pytorch-triton-xpu. Don't you have this issue?

I don't have this issue. It could be because I use nightly intel-xpu-backend-for-triton.

test_rms_norm.py::test_correctness[True-LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                           [  3%]
test_rms_norm.py::test_correctness[True-LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                           [  6%]
test_rms_norm.py::test_correctness[True-LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-2-128-512] PASSED                                                                               [  9%]
test_rms_norm.py::test_correctness[True-LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-5-123-123] PASSED                                                                               [ 12%]
test_rms_norm.py::test_correctness[True-GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                           [ 15%]
test_rms_norm.py::test_correctness[True-GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                           [ 18%]
test_rms_norm.py::test_correctness[True-GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-2-128-512] PASSED                                                                               [ 21%]
test_rms_norm.py::test_correctness[True-GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-5-123-123] PASSED                                                                               [ 25%]
test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                             [ 28%]
test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                             [ 31%]
test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-2-128-512] PASSED                                                                                 [ 34%]
test_rms_norm.py::test_correctness[True-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-5-123-123] PASSED                                                                                 [ 37%]
test_rms_norm.py::test_correctness[False-LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                          [ 40%]
test_rms_norm.py::test_correctness[False-LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                          [ 43%]
test_rms_norm.py::test_correctness[False-LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-2-128-512] PASSED                                                                              [ 46%]
test_rms_norm.py::test_correctness[False-LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-5-123-123] PASSED                                                                              [ 50%]
test_rms_norm.py::test_correctness[False-GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                          [ 53%]
test_rms_norm.py::test_correctness[False-GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                          [ 56%]
test_rms_norm.py::test_correctness[False-GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-2-128-512] PASSED                                                                              [ 59%]
test_rms_norm.py::test_correctness[False-GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-5-123-123] PASSED                                                                              [ 62%]
test_rms_norm.py::test_correctness[False-BaseRMSNorm-0.0-none-dtype0-0.0001-1e-06-2-128-512] PASSED                                                                            [ 65%]
test_rms_norm.py::test_correctness[False-BaseRMSNorm-0.0-none-dtype0-0.0001-1e-06-5-123-123] PASSED                                                                            [ 68%]
test_rms_norm.py::test_correctness[False-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-2-128-512] PASSED                                                                                [ 71%]
test_rms_norm.py::test_correctness[False-BaseRMSNorm-0.0-none-dtype1-0.2-0.02-5-123-123] PASSED                                                                                [ 75%]
test_rms_norm.py::test_correctness_functional[LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-2-2-8] PASSED                                                                         [ 78%]
test_rms_norm.py::test_correctness_functional[LlamaRMSNorm-0.0-llama-dtype0-0.0001-1e-06-9-7-41] PASSED                                                                        [ 81%]
test_rms_norm.py::test_correctness_functional[LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-2-2-8] PASSED                                                                             [ 84%]
test_rms_norm.py::test_correctness_functional[LlamaRMSNorm-0.0-llama-dtype1-0.2-0.02-9-7-41] PASSED                                                                            [ 87%]
test_rms_norm.py::test_correctness_functional[GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-2-2-8] PASSED                                                                         [ 90%]
test_rms_norm.py::test_correctness_functional[GemmaRMSNorm-1.0-gemma-dtype0-0.0001-1e-06-9-7-41] PASSED                                                                        [ 93%]
test_rms_norm.py::test_correctness_functional[GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-2-2-8] PASSED                                                                             [ 96%]
test_rms_norm.py::test_correctness_functional[GemmaRMSNorm-1.0-gemma-dtype1-0.2-0.02-9-7-41] PASSED                                                                            [100%]

Thanks for the update. This PR looks good to me.

[ByronHsu]

@faaany @mgrabban can you fix the conflict and we can merge this ASAP?

[ByronHsu]

Looks good to me! @mgrabban i just invited you as the collab of this repo, can you check the email? After acceptance, can you create a new branch in the main repo, and create a new PR based on that branch? Our CI has issues currently, so any PR from external folks cannot run CI. Thanks in advance!!

[mgrabban]

Looks good to me! @mgrabban i just invited you as the collab of this repo, can you check the email? After acceptance, can you create a new branch in the main repo, and create a new PR based on that branch? Our CI has issues currently, so any PR from external folks cannot run CI. Thanks in advance!!

This is done now. See #407