CUDA Equivariant Operations for PyTorch

0.9.1 · active · verified Fri Apr 17

cuequivariance (version 0.9.1) provides efficient and flexible CUDA kernels for equivariant operations on 3D data, specifically designed for use with PyTorch. It aims to accelerate computations for deep learning models that require rotational, translational, or other forms of equivariance, often found in point cloud processing and molecular modeling. The library is actively developed and part of the vLLM project.

Common errors

```
RuntimeError: CUDA is not available
```
cause PyTorch is installed without CUDA support, CUDA drivers are not installed, or no CUDA-enabled GPU is detected on your system.

fix Ensure you have a CUDA-enabled GPU with up-to-date drivers. Install PyTorch with CUDA support (e.g., `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118` for CUDA 11.8). Verify `torch.cuda.is_available()` returns `True` before running any code using cuequivariance.
```
ninja: build stopped: subcommand failed. /usr/bin/ld: cannot find -lcuda
```
cause Compilation error during `pip install` because the NVIDIA CUDA Toolkit is not correctly installed, or its `nvcc` compiler and CUDA libraries (`libcua.so`) are not in the system PATH and `LD_LIBRARY_PATH` (or equivalent environment variables).

fix Install the NVIDIA CUDA Toolkit relevant to your system and ensure its `bin` directory (containing `nvcc`) and `lib64` directory (containing CUDA libraries) are added to your system's PATH and `LD_LIBRARY_PATH` environment variables respectively.
```
AttributeError: module 'cuequivariance' has no attribute 'PointCloudTransformer'
```
cause Attempting to import modules directly from the top-level `cuequivariance` package when they are nested within its submodules (e.g., `modules`).

fix Import specific modules from their correct submodule path. For example, use `from cuequivariance.modules import PointCloudTransformer` instead of `from cuequivariance import PointCloudTransformer`.

Warnings

gotcha cuequivariance is a CUDA-dependent library. It requires a CUDA-enabled GPU, NVIDIA drivers, and the CUDA Toolkit installed on your system for both compilation during `pip install` and for runtime execution. Without a proper CUDA setup, the library will not function and installation may fail.
Fix: Ensure you have a compatible CUDA-enabled GPU. Install the NVIDIA CUDA Toolkit (matching your PyTorch CUDA version if possible) and corresponding GPU drivers. Verify your PyTorch installation is compiled with CUDA support (e.g., `torch.cuda.is_available()` should be `True`).
gotcha The library explicitly requires PyTorch version 2.0.0 or newer. Using older PyTorch versions (e.g., < 2.0) will lead to API incompatibilities, compilation failures, or runtime errors.
Fix: Upgrade your PyTorch installation to version 2.0.0 or newer (e.g., `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118` for CUDA 11.8).
breaking As the library is currently in version 0.9.x (pre-1.0.0), its API is subject to change without strict backward compatibility guarantees between minor releases. Future updates may introduce breaking changes to module names, class signatures, or function parameters.
Fix: Always refer to the latest official documentation or release notes when upgrading to new minor versions. Pin your `cuequivariance` version in `requirements.txt` to mitigate unexpected breaking changes in production environments (e.g., `cuequivariance==0.9.1`).

Install

pip install cuequivariance Install latest version

Imports

PointCloudTransformer
wrong:
```
from cuequivariance import PointCloudTransformer
```
correct:
```
from cuequivariance.modules import PointCloudTransformer
```
Most core modules, including PointCloudTransformer, are nested under the 'modules' submodule.
EquivariantLinear
wrong:
```
from cuequivariance import EquivariantLinear
```
correct:
```
from cuequivariance.modules import EquivariantLinear
```
Most core modules, including EquivariantLinear, are nested under the 'modules' submodule.

Quickstart

This quickstart demonstrates how to use the PointCloudTransformer module from cuequivariance. It initializes a batch of 3D points and associated features, then processes them through the transformer on a CUDA-enabled GPU. This highlights the library's core functionality for equivariant operations on point cloud data and includes essential runtime CUDA checks.

import torch
from cuequivariance.modules import PointCloudTransformer

# --- Runtime CUDA Check ---
# cuequivariance requires a CUDA-enabled GPU and PyTorch compiled with CUDA support.
if not torch.cuda.is_available():
    raise RuntimeError(
        "CUDA is not available. cuequivariance requires a CUDA-enabled GPU and PyTorch "
        "compiled with CUDA support. Please check your PyTorch installation and CUDA setup."
    )

# Ensure PyTorch device is set to CUDA
device = torch.device("cuda")

# --- Example Data Generation ---
batch_size = 1
num_points = 100
input_feature_dim = 64
output_feature_dim = 128

# Input points (batch_size, num_points, 3) - Represents 3D coordinates
points = torch.randn(batch_size, num_points, 3, device=device)
# Input features (batch_size, num_points, input_feature_dim) - Features associated with each point
features = torch.randn(batch_size, num_points, input_feature_dim, device=device)

# --- Initialize and Use a cuequivariance Module ---
# PointCloudTransformer is an example module demonstrating 3D equivariant operations.
transformer = PointCloudTransformer(input_feature_dim, output_feature_dim).to(device)

# Perform the forward pass
output_features = transformer(points, features)

# --- Verify Output ---
print(f"Input points shape: {points.shape}")
print(f"Input features shape: {features.shape}")
print(f"Output features shape: {output_features.shape}")

assert output_features.shape == (batch_size, num_points, output_feature_dim)
print("cuequivariance quickstart example ran successfully!")

view raw JSON →