Pull Request №459 RosettaCommons/RFdiffusion/main ← xuejiakn/RFdiffusion/feat/npu-ascend-support
Merge: 86507b6538f51fce57b5a72477165f03999ed7ae←bdba6e6b674b1b8ea8668fc275267b40432f4dd0
feat: NPU (Ascend 910B3) support for RFdiffusion inference
----------------
Merge commit message:
feat: NPU (Ascend 910B3) support for RFdiffusion inference
Adapt RFdiffusion to run on Huawei Ascend NPU (Ascend 910B3) with
torch_npu, enabling protein structure generation on NPU hardware.
Changes (8 files, +114 -16 lines):
1. rfdiffusion/__init__.py:
- Global torch.npu.config.allow_internal_format = False
- Patch torch.cdist for NPU (manual bmm + sqrt implementation)
2. rfdiffusion/inference/model_runners.py:
- Device selection: torch.cuda -> torch.npu
3. rfdiffusion/Track_module.py:
- autocast: torch.cuda.amp.autocast -> torch.amp.autocast(device_type='npu')
4. scripts/run_inference.py:
- NPU device detection (priority NPU > CUDA > CPU)
- torch.npu.empty_cache() support
- TRB metadata records actual device name
5-8. SE3Transformer (basis.py, attention.py, convolution.py, norm.py):
- nvtx range three-tier fallback: CUDA nvtx -> NPU record_function -> no-op
- dgl.ops.e_dot_v replaced with manual implementation (Ascend SDDMM
only supports lhs_target=0)
Verified: RFdiffusion NPU pLDDT=0.9885, CPU pLDDT=0.9885, diff=0.000005
NPU speedup: 15.35x over CPU (43.3s vs 664.7s for 20-step inference)
Environment: PyTorch 2.7.1+cpu, torch_npu 2.7.1.post4, CANN 8.2.RC1,
DGL 2.5 (Ascend), Ascend 910B3