feat: NPU (Ascend 910B3) support for RFdiffusion inference - #459
Open
xuejiakn wants to merge 1 commit into
Open
Conversation
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
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.
Description
This PR adapts RFdiffusion to run on Huawei Ascend NPU (Ascend 910B3) with
torch_npu, enabling protein structure generation on NPU hardware.Background
RFdiffusion depends on PyTorch + DGL + e3nn + SE3Transformer, all originally CUDA-only. This PR replaces CUDA-specific calls with NPU-compatible alternatives and patches operators not supported on NPU.
Changes (8 files, +114 -16 lines)
rfdiffusion/__init__.pyallow_internal_format=False+torch.cdistNPU patch (manualbmm + sqrt)rfdiffusion/inference/model_runners.pytorch.cuda→torch.npurfdiffusion/Track_module.pytorch.cuda.amp.autocast→torch.amp.autocast(device_type="npu")scripts/run_inference.pytorch.npu.empty_cache(), device metadataSE3Transformer/model/basis.pyrecord_function→ no-opSE3Transformer/model/layers/attention.pye_dot_vmanual implementationSE3Transformer/model/layers/convolution.pySE3Transformer/model/layers/norm.pyKey Adaptations
torch.cuda.*calls replaced withtorch.npu.*equivalentstorch.npu.config.allow_internal_format = Falseto fixtorch.cattensor format conflicts in SE3Transformertorch.cdistnatively; patched with||x1-x2|| = sqrt(||x1||² + ||x2||² - 2*x1·x2^T)torch.cuda.nvtx.rangeimport succeeds but throws at runtime on NPU; three-tier fallback withtorch.autograd.profiler.record_functionfor NPU profiling supportlhs_target=0; replaced with manual(key * query[dst]).sum(dim=-1)Prerequisites
copy_rhsfix (see dgl-ascend PR #24)Verification
Type