Add Intel GPUs (XPU) inference support

README.md

@@ -268,6 +268,11 @@ with open("1mht_pred.cif", "w") as f:
 
 > **AMD ROCm users:** use ROCm 6.4 with PyTorch 2.9 or newer.
 
+> **Intel GPU (XPU) users:** pass `device="xpu"` instead of `"cuda"`. `pixi run -e xpu ...`
+> gives a working environment. The vendored Triton kernels run on XPU; libraries
+> with no XPU build (Transformer Engine, cuequivariance, flash-attn) fall back to
+> the reference PyTorch path automatically.
+
 ### Running ESMFold2 Through the Biohub Platform
 
 Install the `esm` Python package

esm/models/esm3.py

@@ -31,6 +31,7 @@
 from esm.utils.constants import esm3 as C
 from esm.utils.constants.models import ESM3_OPEN_SMALL, normalize_model_name
 from esm.utils.decoding import decode_protein_tensor
+from esm.utils.device import resolve_default_device, supports_amp_autocast
 from esm.utils.generation import (
     _batch_forward,
     _sample_per_prompt,
@@ -226,15 +227,17 @@ def __init__(
 
     @classmethod
     def from_pretrained(
-        cls, model_name: str = ESM3_OPEN_SMALL, device: torch.device | None = None
+        cls, model_name: str = ESM3_OPEN_SMALL, device: torch.device | str | None = None
     ) -> ESM3:
         from esm.pretrained import load_local_model
 
         model_name = normalize_model_name(model_name)
         if not model_name:
             raise ValueError(f"Model name {model_name} is not a valid ESM3 model name.")
         if device is None:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            device = resolve_default_device()
+        else:
+            device = torch.device(device)
         model = load_local_model(model_name, device=device)
         if device.type != "cpu":
             model = model.to(torch.bfloat16)
@@ -536,7 +539,7 @@ def logits(
         with (
             torch.no_grad(),  # Assume no gradients for now...
             torch.autocast(enabled=True, device_type=device.type, dtype=torch.bfloat16)
-            if device.type == "cuda"
+            if supports_amp_autocast(device.type)
             else contextlib.nullcontext(),
         ):
             output = self(

esm/models/esmc/compatibility.py

@@ -39,6 +39,7 @@
     LogitsOutput,
 )
 from esm.utils.constants.models import ESMC_6B, ESMC_300M, ESMC_600M
+from esm.utils.device import resolve_default_device, supports_amp_autocast
 from esm.utils.sampling import _BatchedESMProteinTensor
 
 _DEPRECATION_MESSAGE = (
@@ -160,7 +161,7 @@ def from_pretrained(
         use_flash_attn: bool = True,
     ) -> "ESMC":
         if device is None:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            device = resolve_default_device()
         device = torch.device(device)
         model = EsmcForMaskedLM.from_pretrained(
             _legacy_name_to_repo(model_name),
@@ -304,7 +305,9 @@ def logits(
         with (
             torch.no_grad(),
             torch.autocast(
-                device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"
+                device.type,
+                dtype=torch.bfloat16,
+                enabled=supports_amp_autocast(device.type),
             ),
         ):
             out = self(sequence_tokens=input.sequence)

esm/models/esmc/layers.py

@@ -7,6 +7,7 @@
 """
 
 import math
+from functools import lru_cache
 
 import torch
 import torch.nn as nn
@@ -26,6 +27,23 @@
 )
 
 
+@lru_cache(maxsize=1)
+def _get_xpu_rotary_kernel():
+    """Triton RoPE kernel, or ``None`` if triton isn't installed.
+
+    Only the import is guarded: a real failure inside the kernel itself
+    (compile error, bad launch, ...) must still raise, not get treated as
+    "unavailable" and silently swapped for the reference path.
+    """
+    try:
+        from esm.models.esmc.xpu_rotary import apply_rotary_emb_xpu
+    except ModuleNotFoundError as exc:
+        if exc.name != "triton":
+            raise
+        return None
+    return apply_rotary_emb_xpu
+
+
 class EsmcRotaryEmbedding(nn.Module):
     """Rotary position embeddings (RoPE) as used by ESMC.
 
@@ -35,8 +53,8 @@ class EsmcRotaryEmbedding(nn.Module):
     * The ``inv_freq`` buffer is recomputed on every device move and kept in
       fp32 even when the module is cast to bf16/fp16. CPU and CUDA ``pow``
       differ by ~1 fp32 ULP, which compounds across the deep attention stack.
-    * The ``forward`` selects the Flash-Attention Triton RoPE kernel when it is
-      available and the tensors live on CUDA, otherwise the pure-PyTorch path.
+    * CUDA uses Flash-Attention Triton RoPE when available. XPU fp32 inference
+      uses a separate Triton kernel; other calls use the pure-PyTorch path.
 
     Parameters
     ----------
@@ -187,6 +205,17 @@ def forward(
             assert apply_triton_rotary is not None
             q_rot = apply_triton_rotary(q, cos, sin, interleaved=self.interleaved)
             k_rot = apply_triton_rotary(k, cos, sin, interleaved=self.interleaved)
+        elif (
+            q.device.type == "xpu"
+            and q.dtype == torch.float32
+            and k.device.type == "xpu"
+            and k.dtype == torch.float32
+            and not self.interleaved
+            and not torch.is_grad_enabled()
+            and (apply_xpu_rotary := _get_xpu_rotary_kernel()) is not None
+        ):
+            q_rot = apply_xpu_rotary(q, cos, sin)
+            k_rot = apply_xpu_rotary(k, cos, sin)
         else:
             q_rot = apply_rotary_emb_torch(q, cos, sin, self.interleaved)
             k_rot = apply_rotary_emb_torch(k, cos, sin, self.interleaved)

esm/models/esmc/model.py

@@ -28,6 +28,7 @@
 from esm.models.esmc.layers import EsmcRotaryEmbedding, EsmcTransformerStack
 from esm.models.esmc.sae import EsmcSaeLayer
 from esm.models.hub import HubPreTrainedModel, resolve_model_dir
+from esm.utils.device import resolve_default_device
 
 _SAFETENSORS_INDEX = "model.safetensors.index.json"
 _SAFETENSORS_SINGLE = "model.safetensors"
@@ -253,7 +254,7 @@ def from_pretrained(
         if attn_implementation is not None:
             config.attn_implementation = attn_implementation
         if device is None:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
+            device = resolve_default_device()
         # Sets the default device the encoder reads to choose its kernels.
         with torch.device(device):
             return cls._load_pretrained(local_dir, config, device=device, dtype=dtype)
@@ -284,12 +285,12 @@ class EsmcModel(EsmcPreTrainedModel):
 
     def __init__(self, config: EsmcConfig):
         super().__init__(config)
-        # Every fused kernel is CUDA-only, so all of them hang off one condition:
-        # the device this model's parameters are being created on. Deriving them
-        # separately is what let the LayerNorm and attention gates disagree.
-        on_cuda = torch.get_default_device().type == "cuda"
+        # flash-attn and TE are both CUDA-only, so they hang off one condition:
+        # the device this model's parameters are being created on. Deriving
+        # them separately is what let the LayerNorm and attention gates disagree.
+        cuda_fused_kernels_available = torch.get_default_device().type == "cuda"
         self._use_flash_attn = (
-            on_cuda
+            cuda_fused_kernels_available
             and FLASH_ATTN_INSTALLED
             and config.attn_implementation == "flash_attention_2"
         )
@@ -299,7 +300,7 @@ def __init__(self, config: EsmcConfig):
             config.num_attention_heads,
             config.num_hidden_layers,
             use_flash_attn=self._use_flash_attn,
-            use_te=on_cuda and TE_INSTALLED,
+            use_te=cuda_fused_kernels_available and TE_INSTALLED,
         )
         self._sae_models: nn.ModuleDict = nn.ModuleDict()
         self._register_state_dict_hook(_drop_te_extra_state)

esm/models/esmc/xpu_rotary.py

@@ -0,0 +1,102 @@
+"""Inference-only fp32 RoPE kernel for ESMC on Intel XPU."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _rotary_kernel(
+    x_ptr,
+    cos_ptr,
+    sin_ptr,
+    out_ptr,
+    B: tl.constexpr,
+    S: tl.constexpr,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    RO: tl.constexpr,
+    X_S0: tl.constexpr,
+    X_S1: tl.constexpr,
+    X_S2: tl.constexpr,
+    X_S3: tl.constexpr,
+    C_S0: tl.constexpr,
+    C_S1: tl.constexpr,
+    S_S0: tl.constexpr,
+    S_S1: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    total = B * S * H * D
+    mask = index < total
+
+    d = index % D
+    h = (index // D) % H
+    s = (index // (D * H)) % S
+    b = index // (D * H * S)
+    x_offset = b * X_S0 + s * X_S1 + h * X_S2 + d * X_S3
+    value = tl.load(x_ptr + x_offset, mask=mask, other=0.0).to(tl.float32)
+
+    rotate_mask = mask & (d < RO)
+    frequency = tl.where(d < RO // 2, d, d - RO // 2)
+    cos_value = tl.load(
+        cos_ptr + s * C_S0 + frequency * C_S1, mask=rotate_mask, other=0.0
+    ).to(tl.float32)
+    sin_value = tl.load(
+        sin_ptr + s * S_S0 + frequency * S_S1, mask=rotate_mask, other=0.0
+    ).to(tl.float32)
+    pair_d = tl.where(d < RO // 2, d + RO // 2, d - RO // 2)
+    pair_offset = b * X_S0 + s * X_S1 + h * X_S2 + pair_d * X_S3
+    pair_value = tl.load(x_ptr + pair_offset, mask=rotate_mask, other=0.0).to(
+        tl.float32
+    )
+    pair_sign = tl.where(d < RO // 2, -1.0, 1.0)
+    rotated = value * cos_value + (pair_value * pair_sign) * sin_value
+    result = tl.where(d < RO, rotated, value)
+    tl.store(out_ptr + index, result, mask=mask)
+
+
+def apply_rotary_emb_xpu(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    """Apply ESMC's non-interleaved RoPE in fp32 inference mode."""
+    if torch.is_grad_enabled():
+        raise RuntimeError("XPU Triton RoPE is only supported with gradients disabled")
+    if x.device.type != "xpu" or cos.device != x.device or sin.device != x.device:
+        raise ValueError("XPU Triton RoPE requires all inputs on the same XPU")
+    if x.dtype != torch.float32 or cos.dtype != x.dtype or sin.dtype != x.dtype:
+        raise ValueError("XPU Triton RoPE requires fp32 inputs")
+    if x.ndim != 4 or cos.ndim != 2 or sin.shape != cos.shape:
+        raise ValueError("unexpected RoPE input shapes")
+
+    batch, sequence, heads, dim = x.shape
+    rotary_dim = cos.shape[-1] * 2
+    if rotary_dim > dim or cos.shape[0] < sequence:
+        raise ValueError("RoPE dimensions exceed the input shape")
+
+    output = torch.empty(
+        x.shape, device=x.device, dtype=x.dtype, memory_format=torch.contiguous_format
+    )
+    block = 128
+    _rotary_kernel[(triton.cdiv(x.numel(), block),)](
+        x,
+        cos,
+        sin,
+        output,
+        batch,  # ty:ignore[invalid-argument-type]
+        sequence,  # ty:ignore[invalid-argument-type]
+        heads,  # ty:ignore[invalid-argument-type]
+        dim,  # ty:ignore[invalid-argument-type]
+        rotary_dim,  # ty:ignore[invalid-argument-type]
+        x.stride(0),  # ty:ignore[invalid-argument-type]
+        x.stride(1),  # ty:ignore[invalid-argument-type]
+        x.stride(2),  # ty:ignore[invalid-argument-type]
+        x.stride(3),  # ty:ignore[invalid-argument-type]
+        cos.stride(0),  # ty:ignore[invalid-argument-type]
+        cos.stride(1),  # ty:ignore[invalid-argument-type]
+        sin.stride(0),  # ty:ignore[invalid-argument-type]
+        sin.stride(1),  # ty:ignore[invalid-argument-type]
+        BLOCK=block,  # ty:ignore[invalid-argument-type]
+        num_warps=4,  # ty:ignore[unknown-argument]
+    )
+    return output

esm/models/esmfold2/experimental.py

@@ -62,6 +62,7 @@
 # Not circular: model.py's reference to this module is inside ``from_pretrained``.
 from esm.models.esmfold2.model import _IGNORED_FEATURE_KEYS
 from esm.models.hub import HubPreTrainedModel, resolve_model_dir
+from esm.utils.device import supports_amp_autocast
 
 _EPS = 1e-5
 
@@ -988,10 +989,11 @@ def forward(
 
         atom_to_token = atom_to_token * atm_mask.long()
 
-        use_amp = ref_pos.device.type == "cuda"
+        device_type = ref_pos.device.type
+        use_amp = supports_amp_autocast(device_type)
         with (
             torch.set_grad_enabled(res_type_soft is not None),
-            torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16),
+            torch.amp.autocast(device_type, enabled=use_amp, dtype=torch.bfloat16),
         ):
             # 1. Input embeddings
             x_inputs = self.inputs_embedder(
@@ -1136,7 +1138,7 @@ def forward(
             )
 
         # 7. Diffusion sampling (always no_grad; optional seed for parity)
-        with torch.no_grad(), _seed_context(seed):
+        with torch.no_grad(), _seed_context(seed, device_type=ref_pos.device.type):
             structure_output = self.structure_head.sample(
                 z_trunk=z.float(),
                 s_inputs=x_inputs,

esm/models/esmfold2/kernels/fused_lnlin_swiglu.py

@@ -341,39 +341,46 @@ class FusedLNLinearSwiGLUFunction(torch.autograd.Function):
     """
 
     @staticmethod
-    @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.bfloat16)
     def forward(ctx, X, W12, LN_W, LN_B):
-        x_shape = X.shape
-        x_2d = X.contiguous().view(-1, x_shape[-1])
-        out, lin, mean, rstd = _lnlin_swiglu_fwd(x_2d, W12, LN_W, LN_B)
-        ctx.save_for_backward(x_2d, W12, LN_W, LN_B, mean, rstd, lin)
-        ctx.x_shape = x_shape
-        ctx.has_ln_bias = LN_B is not None
-        return out.view(*x_shape[:-1], out.shape[-1])
+        # Mirrors torch.amp.custom_fwd(cast_inputs=bfloat16), but device-derived
+        # instead of hardcoded to "cuda" so XPU autocast is respected too.
+        device_type = X.device.type
+        if torch.is_autocast_enabled(device_type):
+            X, W12, LN_W = X.bfloat16(), W12.bfloat16(), LN_W.bfloat16()
+            if LN_B is not None:
+                LN_B = LN_B.bfloat16()
+        with torch.amp.autocast(device_type, enabled=False):
+            x_shape = X.shape
+            x_2d = X.contiguous().view(-1, x_shape[-1])
+            out, lin, mean, rstd = _lnlin_swiglu_fwd(x_2d, W12, LN_W, LN_B)
+            ctx.save_for_backward(x_2d, W12, LN_W, LN_B, mean, rstd, lin)
+            ctx.x_shape = x_shape
+            ctx.has_ln_bias = LN_B is not None
+            return out.view(*x_shape[:-1], out.shape[-1])
 
     @staticmethod
-    @torch.amp.custom_bwd(device_type="cuda")
     def backward(ctx, dout):
         x_2d, W12, LN_W, LN_B, mean, rstd, lin = ctx.saved_tensors
-        dout_2d = dout.contiguous().view(-1, dout.shape[-1])
-        K = x_2d.shape[1]
-
-        # SwiGLU backward, in-place into the saved `lin` buffer (no new alloc).
-        dlin = _swiglu_bwd_inplace(dout_2d, lin)
-
-        # LN+Linear backward via cuBLAS + ATen LN bwd.
-        # x_norm needed for dW12; recomputed via F.layer_norm (cuDNN, ~100 µs at d=256).
-        x_norm = F.layer_norm(x_2d, (K,), LN_W, LN_B, eps=1e-5)
-        dW12 = x_norm.transpose(0, 1) @ dlin  # [K, 2N]
-        dx_norm = dlin @ W12.transpose(0, 1)  # [M, K]
-        del x_norm
-
-        # native_layer_norm_backward output_mask = [need_dX, need_dGamma, need_dBeta].
-        # We always need dX and dGamma; dBeta only when LN bias was used.
-        output_mask = [True, True, ctx.has_ln_bias]
-        dX, dLN_W, dLN_B = torch.ops.aten.native_layer_norm_backward(
-            dx_norm, x_2d, [K], mean.float(), rstd.float(), LN_W, LN_B, output_mask
-        )
+        with torch.amp.autocast(x_2d.device.type, enabled=False):
+            dout_2d = dout.contiguous().view(-1, dout.shape[-1])
+            K = x_2d.shape[1]
+
+            # SwiGLU backward, in-place into the saved `lin` buffer (no new alloc).
+            dlin = _swiglu_bwd_inplace(dout_2d, lin)
+
+            # LN+Linear backward via cuBLAS + ATen LN bwd.
+            # x_norm needed for dW12; recomputed via F.layer_norm (cuDNN, ~100 µs at d=256).
+            x_norm = F.layer_norm(x_2d, (K,), LN_W, LN_B, eps=1e-5)
+            dW12 = x_norm.transpose(0, 1) @ dlin  # [K, 2N]
+            dx_norm = dlin @ W12.transpose(0, 1)  # [M, K]
+            del x_norm
+
+            # native_layer_norm_backward output_mask = [need_dX, need_dGamma, need_dBeta].
+            # We always need dX and dGamma; dBeta only when LN bias was used.
+            output_mask = [True, True, ctx.has_ln_bias]
+            dX, dLN_W, dLN_B = torch.ops.aten.native_layer_norm_backward(
+                dx_norm, x_2d, [K], mean.float(), rstd.float(), LN_W, LN_B, output_mask
+            )
         # If LN_B was None (no bias), ATen returns dLN_B=None and we just pass it through.
         return dX.view(ctx.x_shape), dW12, dLN_W, dLN_B