cuequivariance-ops-torch

Introduction

cuequivariance_ops_torch provides CUDA kernels for the cuEquivariance project's PyTorch components. As such, it contains pytorch bindings to optimized kernels that cuEquivariance's operations map down to. In general, we advice that you access those kernels through cuEquivariance, but you may also find them useful on their own.

Installation

Please install using either pip install cuequivariance-ops-torch-cu12 or pip install cuequivariance-ops-torch-cu13 (depending on the CUDA toolkit you wish to use).

Documentation

For detailed usage information of the kernels, please refer to the doc-strings in their respective modules. For higher-level documentations, refer to cuEquivariance.

Pairformer inference primitives

The Torch package also exposes lower-level primitives for integrating Pairformer inference graphs:

  • attention_pair_bias_mask_from_normalized projects an already-normalized pair representation and applies an arbitrary Boolean [B,U,V] mask. The implementation is shape-generic; BF16 D=256, H=16, and square sequence lengths 384 and 512 are the performance-validated configurations. FP16 and BF16 inputs use native low-precision tensor-core products with FP32 accumulation; FP32 inputs retain the higher-accuracy TF32x3 path.
  • pairformer_combined_triangle_multiplication combines incoming and outgoing triangle multiplication after caller-provided non-affine LayerNorm with epsilon 1e-5. Its fused direction join applies the affine middle LayerNorm with epsilon 0.03. This operation intentionally fails closed outside BF16 inference with B=1, D=256, sequence length 384 or 512, and an SM100 GPU. The caller guarantees that the second normalized carrier is the separately materialized spatial transpose of the first; the hot path validates metadata and distinct storage but does not compare tensor contents. It returns the base update before any caller-owned peri-LayerNorm, output mask, dropout, or residual addition; model-level dropout must be disabled for inference. Use pack_pairformer_grouped_projection to convert grouped projection weights to the expected payload/gate layout.

Both operations are registered as torch.library custom operators with fake implementations for torch.compile tracing. See their Python docstrings for the complete input contracts.

Usage

You can import the library from python:

import cuequivariance_ops_torch

Kernels are primarily exposed as torch.nn.Module, but also provide a lower-level interface as torch.library operators. Generally, the module is responsible for proper input transformation and initialization, and the operator execute the kernel. This allows you to export models using this operations using torch.export, and running inference on them using TensorRT.

Support and Feedback

Please contact the cuEquivariance developers for any issues you might encounter.

Project links

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Newer release available (0.11.1)

8 maintainers

PyPI data
Data sourced directly from PyPI's database.
Avatar for fpellegrini75 from gravatar.com fpellegrini75 Avatar for hanyichou from gravatar.com hanyichou Avatar for hsadasiv from gravatar.com hsadasiv
+5 more Avatar for jomitchellnv from gravatar.com jomitchellnv Avatar for kucukben from gravatar.com kucukben Avatar for mariogeiger from gravatar.com mariogeiger Avatar for mhohnerbach_nvidia from gravatar.com mhohnerbach_nvidia Avatar for tingyu_nvidia from gravatar.com tingyu_nvidia

Credits

Author: NVIDIA Corporation

License

Other/Proprietary License (NVIDIA Proprietary Software)

Requires

Python >=3.10

Provides Extra

test plugins

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

cuequivariance_ops_torch_cu12-0.11.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (467.6 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.0-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (462.5 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

cuequivariance_ops_torch_cu12-0.11.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (467.8 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (462.2 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

cuequivariance_ops_torch_cu12-0.11.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (467.9 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (462.3 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

cuequivariance_ops_torch_cu12-0.11.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (468.2 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (463.3 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

cuequivariance_ops_torch_cu12-0.11.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (467.9 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

cuequivariance_ops_torch_cu12-0.11.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (463.1 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 8246c830d5a0eaeb72f8b6380320dd77b0a25da9ab2bfb60d9b29c239dc00c66
MD5 169179b9b70e779fcb31125c29938fb9
BLAKE2b-256 cfca85a3d6b8e8b7209fb6626909e666f4b69c299deba85b33fd91cdfabf8101

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 444a7f9d6fa76719b590dcb5cc312866c2ce1cdadc6ea48f85f94d26fb21a6f4
MD5 bf3bfbf51b96b87efe1f38ced0b117f6
BLAKE2b-256 140db7cdf3f54014e87eb192f39459232b053038fa1f0c68032ef6bc8efaef09

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 9263c873a9f782a4c06e97a43a1cc7fe02d525d0135b7e509d67c6549da9f4ca
MD5 220057ce77291ad4457ea529ed19ce58
BLAKE2b-256 d6033671bb5675e0b251b36bacdb93f06e29d2b59437ea4bee9ed35eb91fad95

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 47b64970ad098eb36f1f8d1e47239c2b5a535a188c3a2fbdba6e445be6b23439
MD5 f2c03990f2ecb3975635249a9028e189
BLAKE2b-256 ed9bd7d8b5896e113b3c4374e4e2893e7c732fa3c5d99be903f7630f7c9c0ca2

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 318a05a27266bc681263db79db1f5d6a005d45d73740876c677527b1e49b440f
MD5 8db1061d8eba7523ba2439f5b77e3549
BLAKE2b-256 bdb1762cb2938246149c540559b6daf6b0fedc18066fdfe39837a72ea1308edd

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 53c5d0d2c6c14b1a463a50505561acff8db091edc3bf08bf5163532cd7ad3c18
MD5 aff2f5f2d0f866b65b75cca454889eba
BLAKE2b-256 e5a37ad303b3a8737569e8ab0feca1efd690ecbe4b2ddec23d197123eeaad1c5

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 216d49b31179f6cd2ab5db5cdf92db17716001e27cb886ab28fe3534035e2e39
MD5 f15440f9a4d6fe6aadbb8275f167849f
BLAKE2b-256 1e1e792d4d15a8549961cde625af9e9a936e658cc3b3f756e93511e8a5742265

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 b91effa2de7680fc00542ea2755b0f9baa9a02e130468f89b4b9c9951cf2b2a1
MD5 28b0e2be87ed550f7743d15871c04778
BLAKE2b-256 285883d87567fd6db440b23ca2debed8d09ec2dbed36a922c45db09f76ed31aa

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 1fcf2c8e93859e5fa7a618f9b22c9482c1a1c318aff4e738ce1dc2a2f89499a1
MD5 b914ae66a72e5c8783ba70769a43df37
BLAKE2b-256 1b630fbbec9fdcda0d170727c3d37326fe4710424e86cb645a1e71fbb7c1cba6

See more details on using hashes here.

File details

Details for the file cuequivariance_ops_torch_cu12-0.11.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for cuequivariance_ops_torch_cu12-0.11.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 ea19ace6349ca145e7b2588bb57b0c65959e58a862df12b898ca81099b8cea6b
MD5 6e33bdab28fa3701319873e7d63dde20
BLAKE2b-256 fd4be98cbe4ae71178e5b61f8d9c75ca3a2cf9a59f8aa8a441b801f48e2ffab8

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page