feat: NPU (Ascend 910B3) support for RFdiffusion inference

env/SE3Transformer/se3_transformer/model/basis.py

@@ -29,7 +29,21 @@
 import torch
 import torch.nn.functional as F
 from torch import Tensor
-from torch.cuda.nvtx import range as nvtx_range
+try:
+    from torch.cuda.nvtx import range as _nvtx_range_cuda
+    try:
+        with _nvtx_range_cuda("init"):
+            pass
+        nvtx_range = _nvtx_range_cuda
+    except Exception:
+        raise ImportError("NVTX not available at runtime")
+except (ImportError, RuntimeError):
+    from torch.autograd.profiler import record_function
+    from contextlib import contextmanager
+    @contextmanager
+    def nvtx_range(msg, *args, **kwargs):
+        with record_function(msg):
+            yield
 
 from se3_transformer.runtime.utils import degree_to_dim
 

env/SE3Transformer/se3_transformer/model/layers/attention.py

@@ -34,7 +34,21 @@
 from se3_transformer.model.layers.convolution import ConvSE3, ConvSE3FuseLevel
 from se3_transformer.model.layers.linear import LinearSE3
 from se3_transformer.runtime.utils import degree_to_dim, aggregate_residual, unfuse_features
-from torch.cuda.nvtx import range as nvtx_range
+try:
+    from torch.cuda.nvtx import range as _nvtx_range_cuda
+    try:
+        with _nvtx_range_cuda("init"):
+            pass
+        nvtx_range = _nvtx_range_cuda
+    except Exception:
+        raise ImportError("NVTX not available at runtime")
+except (ImportError, RuntimeError):
+    from torch.autograd.profiler import record_function
+    from contextlib import contextmanager
+    @contextmanager
+    def nvtx_range(msg, *args, **kwargs):
+        with record_function(msg):
+            yield
 
 
 class AttentionSE3(nn.Module):
@@ -78,7 +92,10 @@ def forward(
 
             with nvtx_range('attention dot product + softmax'):
                 # Compute attention weights (softmax of inner product between key and query)
-                edge_weights = dgl.ops.e_dot_v(graph, key, query).squeeze(-1)
+                # Use manual implementation for Ascend NPU compatibility (e_dot_v not supported)
+                dst_nodes = graph.edges()[1]
+                query_dst = query[dst_nodes]
+                edge_weights = (key * query_dst).sum(dim=-1)
                 edge_weights /= np.sqrt(self.key_fiber.num_features)
                 edge_weights = edge_softmax(graph, edge_weights)
                 edge_weights = edge_weights[..., None, None]

env/SE3Transformer/se3_transformer/model/layers/convolution.py

@@ -31,7 +31,21 @@
 import torch.nn as nn
 from dgl import DGLGraph
 from torch import Tensor
-from torch.cuda.nvtx import range as nvtx_range
+try:
+    from torch.cuda.nvtx import range as _nvtx_range_cuda
+    try:
+        with _nvtx_range_cuda("init"):
+            pass
+        nvtx_range = _nvtx_range_cuda
+    except Exception:
+        raise ImportError("NVTX not available at runtime")
+except (ImportError, RuntimeError):
+    from torch.autograd.profiler import record_function
+    from contextlib import contextmanager
+    @contextmanager
+    def nvtx_range(msg, *args, **kwargs):
+        with record_function(msg):
+            yield
 
 from se3_transformer.model.fiber import Fiber
 from se3_transformer.runtime.utils import degree_to_dim, unfuse_features

env/SE3Transformer/se3_transformer/model/layers/norm.py

@@ -27,7 +27,21 @@
 import torch
 import torch.nn as nn
 from torch import Tensor
-from torch.cuda.nvtx import range as nvtx_range
+try:
+    from torch.cuda.nvtx import range as _nvtx_range_cuda
+    try:
+        with _nvtx_range_cuda("init"):
+            pass
+        nvtx_range = _nvtx_range_cuda
+    except Exception:
+        raise ImportError("NVTX not available at runtime")
+except (ImportError, RuntimeError):
+    from torch.autograd.profiler import record_function
+    from contextlib import contextmanager
+    @contextmanager
+    def nvtx_range(msg, *args, **kwargs):
+        with record_function(msg):
+            yield
 
 from se3_transformer.model.fiber import Fiber
 

rfdiffusion/Track_module.py

@@ -233,7 +233,7 @@ def reset_parameter(self):
         nn.init.zeros_(self.embed_e1.bias)
         nn.init.zeros_(self.embed_e2.bias)
 
-    @torch.cuda.amp.autocast(enabled=False)
+    @torch.amp.autocast(device_type="npu", enabled=False)
     def forward(self, msa, pair, R_in, T_in, xyz, state, idx, motif_mask, cyclic_reses=None, top_k=64, eps=1e-5):
         B, N, L = msa.shape[:3]
 

rfdiffusion/__init__.py

@@ -0,0 +1,29 @@
+import torch
+
+try:
+    import torch_npu
+    torch.npu.config.allow_internal_format = False
+
+    # Patch torch.cdist for NPU (NPU does not support cdist natively)
+    _orig_cdist = torch.cdist
+    def _npu_cdist(x1, x2, p=2.0, compute_mode='use_mm_for_euclid_dist_if_necessary', **kwargs):
+        if x1.device.type == 'npu' or (x2 is not None and hasattr(x2, 'device') and x2.device.type == 'npu'):
+            if x1.dim() == 2:
+                x1 = x1.unsqueeze(0)
+                x2 = x2.unsqueeze(0)
+                squeeze = True
+            else:
+                squeeze = False
+            x1_sq = (x1 * x1).sum(dim=-1, keepdim=True)
+            x2_sq = (x2 * x2).sum(dim=-1, keepdim=True)
+            cross = torch.bmm(x1, x2.transpose(-1, -2))
+            dist_sq = x1_sq + x2_sq.transpose(-1, -2) - 2 * cross
+            dist_sq = dist_sq.clamp(min=0)
+            result = torch.sqrt(dist_sq + 1e-12)
+            if squeeze:
+                result = result.squeeze(0)
+            return result
+        return _orig_cdist(x1, x2, p=p, compute_mode=compute_mode, **kwargs)
+    torch.cdist = _npu_cdist
+except ImportError:
+    pass

rfdiffusion/inference/model_runners.py

@@ -48,8 +48,8 @@ def initialize(self, conf: DictConfig) -> None:
 
         """
         self._log = logging.getLogger(__name__)
-        if torch.cuda.is_available():
-            self.device = torch.device("cuda")
+        if torch.npu.is_available():
+            self.device = torch.device("npu")
         else:
             self.device = torch.device("cpu")
         needs_model_reload = (

scripts/run_inference.py

@@ -18,6 +18,8 @@
 import re
 import os, time, pickle
 import torch
+import torch_npu
+torch.npu.config.allow_internal_format = False
 from omegaconf import OmegaConf
 import hydra
 import logging
@@ -41,13 +43,16 @@ def main(conf: HydraConfig) -> None:
     if conf.inference.deterministic:
         make_deterministic()
 
-    # Check for available GPU and print result of check
-    if torch.cuda.is_available():
+    # Check for available NPU/GPU and print result of check
+    if torch.npu.is_available():
+        device_name = torch.npu.get_device_name(torch.npu.current_device())
+        log.info(f"Found NPU with device_name {device_name}. Will run RFdiffusion on {device_name}")
+    elif torch.cuda.is_available():
         device_name = torch.cuda.get_device_name(torch.cuda.current_device())
         log.info(f"Found GPU with device_name {device_name}. Will run RFdiffusion on {device_name}")
     else:
         log.info("////////////////////////////////////////////////")
-        log.info("///// NO GPU DETECTED! Falling back to CPU /////")
+        log.info("///// NO GPU/NPU DETECTED! Falling back to CPU /////")
         log.info("////////////////////////////////////////////////")
 
     # Initialize sampler and target/contig.
@@ -148,9 +153,11 @@ def main(conf: HydraConfig) -> None:
         trb = dict(
             config=OmegaConf.to_container(sampler._conf, resolve=True),
             plddt=plddt_stack.cpu().numpy(),
-            device=torch.cuda.get_device_name(torch.cuda.current_device())
-            if torch.cuda.is_available()
-            else "CPU",
+            device=torch.npu.get_device_name(torch.npu.current_device())
+            if torch.npu.is_available()
+            else (torch.cuda.get_device_name(torch.cuda.current_device())
+                  if torch.cuda.is_available()
+                  else "CPU"),
             time=time.time() - start_time,
         )
         if hasattr(sampler, "contig_map"):
@@ -188,8 +195,11 @@ def main(conf: HydraConfig) -> None:
                 chain_ids=sampler.chain_idx,
             )
 
-        if conf.inference.empty_cache_per_design and torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        if conf.inference.empty_cache_per_design:
+            if torch.npu.is_available():
+                torch.npu.empty_cache()
+            elif torch.cuda.is_available():
+                torch.cuda.empty_cache()
 
         log.info(f"Finished design in {(time.time()-start_time)/60:.2f} minutes")