geometric-nn-ops on ROCm documentation#
2026-10-29
1 min read time
geometric-nn-ops is a library of vendor-neutral GPU kernels, with both forward and backward passes, for geometry-aware neural networks such as protein structure prediction models. The kernels are written in Triton and have been validated and optimized for AMD GPUs.
The first kernel family, geometric_nn_ops_torch.triangle, covers the triangle
operations at the core of AlphaFold-style protein folding models:
triangle_attention, triangle_multiplicative_update, and
attention_pair_bias. The kernels are exposed through a PyTorch API, and an
optional compatibility package, hipequivariance-ops-torch, lets code written
against cuequivariance_ops_torch run on AMD GPUs without source
changes.
Note
geometric-nn-ops is under active development, and some operations do not yet have a native Triton kernel for every pass. For the current support status, see Operation support matrix.
The geometric-nn-ops project is located in ROCm/geometric-nn-ops.
To contribute to geometric-nn-ops, refer to CONTRIBUTING.md.
You can find licensing information on the Licensing page.