enable OffloadedCache on XPU from PyTorch 2.7

[yao-matrix]

XPU are aligning features in PyTorch w/ CUDA. Since PyTorch 2.6, an device agnostic torch.Stream is supported and XPU support this API. So, I enabled OffloadedCache on XPU.

Why start from 2.7? The reason is OffloadedCache needs StreamContext, but the PR to support __enter__ attribute of StreamContext is not merged in 2.6, but will be in 2.7.

Tested w/ PyTorch 2.7 dev package(pip install --pre torch==2.7.0.dev20250306 --index-url https://download.pytorch.org/whl/nightly/xpu).

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the Ready for review button (at the bottom of the PR page).

[yao-matrix]

the ci failed cases seems irrelevant to my changes.

[ydshieh]

Hi @yao-matrix Thank you for make this supported.

Hi @n17s, are you interested to take a first look? cc @gante

[SunMarc]

Looks fine to me overall !

[n17s]

Looks good to me

[gante]

LGTM, thank you for adding support! 🤗

Added a minor nit with a more recent import guard practice, happy to merge when it's sorted

[SunMarc]

Thanks ! LGTM !

[loadams]

Hi @yao-matrix and @SunMarc - it looks like running this PR with torch 2.5.0a0+b465a5843b.nv24.9 (from nvcr.io/nvidia/pytorch:24.09-py3) I see the following error:

ImportError: cannot import name 'Replicate' from 'torch.distributed.tensor' (/usr/local/lib/python3.10/dist-packages/torch/distributed/tensor/__init__.py)

Perhaps the guards are on the wrong version of pytorch?

[yao-matrix]

Hi @yao-matrix and @SunMarc - it looks like running this PR with torch 2.5.0a0+b465a5843b.nv24.9 (from nvcr.io/nvidia/pytorch:24.09-py3) I see the following error:

ImportError: cannot import name 'Replicate' from 'torch.distributed.tensor' (/usr/local/lib/python3.10/dist-packages/torch/distributed/tensor/__init__.py)

Perhaps the guards are on the wrong version of pytorch?

it's weird, my version checks on 2.7, which means if version >= 2.7, goes the new API; else the old. But I can see in your PR you changed the pytorch from 2.5 to 2.6, both versions go the old path.

[loadams]

Hi @yao-matrix and @SunMarc - it looks like running this PR with torch 2.5.0a0+b465a5843b.nv24.9 (from nvcr.io/nvidia/pytorch:24.09-py3) I see the following error:

ImportError: cannot import name 'Replicate' from 'torch.distributed.tensor' (/usr/local/lib/python3.10/dist-packages/torch/distributed/tensor/__init__.py)

Perhaps the guards are on the wrong version of pytorch?

it's weird, my version checks on 2.7, which means if version >= 2.7, goes the new API; else the old. But I can see in your PR you changed the pytorch from 2.5 to 2.6, both versions go the old path.

@yao-matrix - yes that is quite odd, but I was able to bisect the failure to this PR, so perhaps it is another code path that this PR is enabling that I'm hitting it from? But it does seem to be resolved by updating the torch version.

[gante]

@yao-matrix I'm going to revert part of the changes in is_torch_greater_or_equal, as it is breaking in other parts of the library. In a nutshell, we can't confirm that all dev versions for 2.X contain the features that will be release in 2.X, which is the error @loadams is seeing (2.5.0a0+b465a5843b.nv24.9 is a dev version of 2.5.0).

@yao-matrix to enable your use case I'm going to add an accept_dev flag to is_torch_greater_or_equal