Repository navigation
Add Intel GPUs (XPU) inference support - #423
Open
rajeshgangireddy wants to merge 15 commits into
Open
rajeshgangireddy wants to merge 15 commits into
rajeshgangireddy wants to merge 15 commits into
Conversation
torch/triton-xpu from the pytorch xpu wheel index, own solve-groups so they don't affect the existing cuda/cpu environments.
resolve_default_device/autocast_device_type/empty_cache/synchronize/ manual_seed_all/*_rng_state_all dispatch to cuda or xpu. Extends fp32_autocast_context with an xpu branch alongside its existing cpu/cuda/mps ones.
esm3.py, esmc/model.py, esmc/compatibility.py, widgets/utils/clients.py now resolve their default device via esm.utils.device instead of 'cuda' if torch.cuda.is_available() else 'cpu'. Autocast-enable conditions gated on device.type == 'cuda' now use supports_amp_autocast so xpu gets the same bf16 autocast cuda does.
torch.amp.autocast("cuda", ...) call sites now derive device_type
from the actual tensor/device instead of a literal 'cuda', so xpu
gets the same bf16 autocast path cuda does. The precision-guard
@autocast("cuda", enabled=False) decorators (Kabsch alignment, RMSD)
get a paired @autocast("xpu", enabled=False) so they disable autocast
on whichever device is actually active.
processor.py's _seed_context saved/restored torch.cuda's rng state directly; now goes through esm.utils.device so xpu is covered too.
- generalize the two remaining hardcoded cuda autocast sites - _seed_context saves/restores rng state via esm.utils.device - allow the vendored Triton kernels to import on xpu, not just cuda (is_cuda checks -> device.type in (cuda, xpu)) - CUE_AVAILABLE now also requires torch.cuda.is_available(): the cuequivariance import can succeed on xpu even though the library JIT-compiles via nvrtc and can't actually run there - FusedLNLinearSwiGLUFunction's forward/backward used torch.amp.custom_fwd/custom_bwd(device_type="cuda"), which silently skips the bf16 cast under an xpu autocast context since device_type is fixed at decoration time. Replaced with an explicit, device-derived equivalent. Verified forward+backward on both cuda and xpu: bf16 autocast now engages correctly and both devices land on the same max-diff vs an fp32 reference.
Rename only; this flag gates CUDA-only libs (flash-attn, TE), not whether an accelerator is present at all.
- xpu pytest marker - esmc_test: cpu-vs-xpu agreement + build matrix - esmfold2_test: kernel_backends_agree extended to xpu/fused - esmfold2_builds_test: cpu-vs-xpu agreement, use_amp, build matrix - fixed stale use_amp docstring (now supports_amp_autocast, not a cuda string check) All tolerances measured on real hardware (RTX 3090 / Arc Pro B70), not guessed.
rajeshgangireddy
marked this pull request as ready for review
October 8, 2026 15:27
rajeshgangireddy
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds Intel XPU inference for ESMC, ESMFold2, and ESM3. Pass
device="xpu"; CPUand CUDA paths remain available.
fp32 ESMC RoPE kernel. CUDA-only libraries use reference paths by default;
explicitly selecting cuequivariance on XPU raises an error.
usage note.
Validation
Compared published checkpoints on the same machine, using an RTX 3090
(CUDA) and Intel Arc Pro B70 (XPU). Each device was run separately with the
same inputs and precision; the results were then compared.
Diffusion and iterative generation are not bit-identical across devices.
ESMFold2 structures from both devices were also close to the experimental
structures: 1.03 Å to 1UBQ and 0.40–0.45 Å to 1PGA.
The default, CUDA, and XPU test runs completed without unexpected failures:
418, 32, and 23 tests passed, respectively.
On the B70, the ESMC RoPE kernel reduced warmed fp32 forward latency by
23% for a short sequence and 5% for a longer batch. Compiling from an
empty Triton cache adds about 4.9 seconds to the first call.
Scope
Multi-XPU inference is not included. The full ESMFold2 checkpoint was not
compared against CUDA because its fp32 backbone exceeds the 3090's 24 GB;
the Experimental checkpoint was used instead.
While working on this PR, I saw that support for MPS was also created (#409). Quite some changes follow similar pattern. Would rebase if that get's merged first.
Let me know if any further numbers/validation needs to be done.