Skip to content

feat: Add Intel XPU (GPU) support#131

Open
Daankrol wants to merge 2 commits into
HannesStark:mainfrom
Daankrol:xpu-support
Open

feat: Add Intel XPU (GPU) support#131
Daankrol wants to merge 2 commits into
HannesStark:mainfrom
Daankrol:xpu-support

Conversation

@Daankrol
Copy link
Copy Markdown

@Daankrol Daankrol commented Dec 18, 2025

feat: Add Intel XPU (GPU) Support

Hi @HannesStark ! 👋

First of all, I want to say thank you for creating and sharing BoltzGen. It's a fantastic tool for generative modeling of biomolecular structures, and the codebase is really well-structured and a pleasure to work with!

What this PR does

This PR adds full support for Intel XPU (GPU) devices, enabling BoltzGen to run on Intel GPUs (Arc, Data Center Max Series, etc.) via PyTorch's XPU backend while maintaining full CUDA compatibility.

Features & Changes

New CLI functionality:

  • Added --accelerator argument with choices: auto, gpu, xpu, cpu, cuda
  • Implemented auto-detection of available hardware (CUDA → gpu, XPU → xpu)
  • Automatically disables CUDA-specific kernels (cuequivariance) on non-CUDA devices

PyTorch Lightning XPU Integration:

  • Custom XPUAccelerator class for PyTorch Lightning
  • SingleXPUStrategy for single-device XPU training/inference
  • XPUMixedPrecision plugin for proper mixed precision handling on XPU devices

Model & Module Updates:

  • Replaced hardcoded 'cuda' autocast device type with dynamic get_autocast_device_type()
  • Added _get_device() helper that checks for XPU availability in addition to CUDA
  • Updated tensor device handling across model layers, loss functions, and modules:
    • Attention layers
    • Confidence utilities
    • Pairformer, Miniformer, Triangular attention
    • All loss functions (diffusion, confidence, distogram, etc.)
    • Encoders and diffusion modules

Documentation:

  • Added Intel GPU (XPU) installation instructions to the README

Daan Krol and others added 2 commits December 18, 2025 16:16
- Add --accelerator CLI argument (auto, gpu, xpu, cpu, cuda)
- Add XPUAccelerator, SingleXPUStrategy, and XPUMixedPrecision for PyTorch Lightning
- Replace hardcoded 'cuda' with dynamic get_autocast_device_type()
- Update tensor device handling across layers, losses, and modules
- Add Intel GPU installation instructions to README
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant