21. Lesson 12a: Foundations of cuEquivariance

What you will learn

• Why equivariant operations built from scratch in Parts II-IV are often memory-bound at scale and are limited by kernel-launch overhead at scales this course deals with. What NVIDIA’s cuEquivariance can do about it.

• The cuEquivariant’s three-layer abstraction: EquivariantPolynomial, SegmentedPolynomial, and SegmentedTensorProduct and how to read the descriptor of an operation.

• The cue.Irreps vs. o3.Irreps comparison: why the group is written down explicitly, and what the two memory layouts mul_ir and ir_mul mean for your tensors.

• How to make a first drop-in replacement of an e3nn module and verify it.

Prerequisites:

• Lesson 02a: Irreps and o3.Irreps,

• Lesson 03b: o3.TensorProduct, connection modes and instructions: the descriptors in this lesson convey the same information in a different notation.

Helpful but not required:

• Lesson 08b: The convolution,

• Lesson 10bcoming soon: Body-order, which supply the two operations we will inspect in this lesson.

Note: This lesson lays the groundwork for Lesson 12b which performs the benchmarking and analysis. While all lessons in Parts I-V of this course can run on a CPU, Lessons 12a and 12b need an NVIDIA GPU.

import sys
sys.path.insert(0, "..")

import torch
from e3nn import o3

import cuequivariance as cue
import cuequivariance_torch as cuet

torch.manual_seed(0)
assert torch.cuda.is_available(), "this lesson needs an NVIDIA GPU"
device = "cuda"

print(f"torch {torch.__version__} | cuequivariance {cue.__version__}")
print(f"GPU: {torch.cuda.get_device_name(0)}")

torch 2.7.1+cu126 | cuequivariance 0.8.1
GPU: NVIDIA A100 80GB PCIe

Installation and version pinning. cuEquivariance ships as a frontend plus a separate CUDA ops package:

• cuequivariance — the framework-independent core: irreps, layouts, descriptors.
• cuequivariance-torch — the PyTorch modules (cuet) we use here. A JAX front end, cuequivariance-jax, exists too.
• cuequivariance-ops-torch-cu12 — the compiled kernels (cu13 for CUDA 13).

The ops package tracks recent PyTorch releases closely. This course pins torch==2.7.1 / CUDA 12.6, for which the matching release is 0.8.1; installing the current 0.11.x would pull torch>=2.11 and silently upgrade the environment that produced every committed output in this repository. Install as an extra:

uv pip install "cuequivariance==0.8.1" "cuequivariance-torch==0.8.1" "cuequivariance-ops-torch-cu12==0.8.1"

21.1. Why a specialized library exists

Nothing we built in this course is arithmetically expensive. A tensor product of \(l \le 2\) features is a handful of small contractions against Clebsch-Gordan coefficients that are known at compile time and mostly zero. These operations are never compute-bound and the arithmetic intensity is low. Thus, the limit on the workflow is how fast data moves, not how fast the multiplications are happening.

The source of the imposed cost is also structural and impacted by the package design. e3nn compiles a tensor product into one weighted path per \((l_i, l_f, l_o)\) combination, and each path becomes its own einsum, reshape and tensordot operation. The resulting CUDA kernel launch can read its inputs from and write its intermediate back to the global memory. A modest-size convolution often has a handful of paths. For example, a MACE symmetric contraction at \(\nu = 3\) has dozens.

In practice, factors that make an operation memory-bound can manifest as kernels that move far more bytes than they do useful arithmetic. However, at the system sizes this course generally deals with (e.g., 8 argon atoms, 21 aspirin atoms) something else dominates the performance. Here, each kernel becomes too small to fill the device and subsequently, the GPU finishes the work in one kernel long before the next is queued, sitting idle in between. That regime is neither compute- nor bandwidth-limited; it is mainly impacted by the kernel-launch overhead. This is why a model of a few hundred atoms can take almost exactly as long per step as one of a few thousand. Lesson 12b measures both regimes.

cuEquivariance can mitigate this issue. It describes a whole operation as one segmented polynomial, which is a single algebraic object covering every path, and emits one fused CUDA kernel for it. As such, the per-path launches and the intermediates will go away together.

The library is not specific to interatomic potentials. The same machinery supports equivariant models in structural biology and docking, and the package ships fused triangle-attention and triangle-multiplication kernels of the kind AlphaFold-style architectures need.

Before proceeding with inner workings of cuEquivariance, let us familiarize ourselves with the necessary vocabulary.

21.2. Irreps, with the group written down

The o3.Irreps("32x0e + 32x1o") notation in e3nn carries the group notation, implicitly: e/o are \(O(3)\) parities. cuEquivariance makes the group an explicit argument, because the same machinery can serve multiple groups such as \(SO(3)\), \(O(3)\) and \(SU(2)\):

# Irreps in e3nn
e3nn_irreps = o3.Irreps("32x0e + 32x1o")

# Irreps in cuEquivariance with explicit group
cue_irreps = cue.Irreps("O3", "32x0e + 32x1o")

# Compare the two irreps
print(f"e3nn: {e3nn_irreps}   dim = {e3nn_irreps.dim}")
print(f"cue : {cue_irreps}   dim = {cue_irreps.dim}")
print(f"\nsame dimension: {e3nn_irreps.dim == cue_irreps.dim}")
print(f"cue also speaks: {cue.Irreps('SO3', '32x0 + 32x1')}  (SO(3): no parity label)")

e3nn: 32x0e+32x1o   dim = 128
cue : 32x0e+32x1o   dim = 128

same dimension: True
cue also speaks: 32x0+32x1  (SO(3): no parity label)

21.3. Data layouts in cuEquivariance and e3nn: An important pitfall

A block of multiplicity \(m\) and irrep dimension \(2l+1\) holds \(m \times (2l+1)\) numbers. There are two ways to arrange these numbers in memory:

layout	shape	used by
cue.mul_ir	\((m,\; 2l+1)\); multiplicity first	e3nn (hence, every model in this course)
cue.ir_mul	\((2l+1,\; m)\); irrep first	cuEquivariance’s preferred layout

The mul_ir refers to the layout the e3nn uses. Thus, every cuequivariance-torch (cuet) module takes a layout= argument, and passing cue.mul_ir lets you feed it e3nn-shaped tensors directly.

Let’s see these layouts in action.

# Create an array of six numbers = 2x1o (two vectors of three components)
x = torch.arange(6.0)

# Print the array in different layouts
print("flat:", x.tolist())
print(f"mul_ir -> shape (2, 3): {x.reshape(2, 3).tolist()}")
print(f"   vector 0 = {x.reshape(2, 3)[0].tolist()},\n   vector 1 = {x.reshape(2, 3)[1].tolist()}")
print(f"ir_mul -> shape (3, 2): {x.reshape(2, 3).T.tolist()}")
print(f"x-components = {x.reshape(2, 3).T[0].tolist()} across both vectors")

# cue.MulIrrep(2, cue.O3(1, -1)) --> 2x1o
print(f"\ncue.mul_ir.shape(...) helper: {cue.mul_ir.shape(cue.MulIrrep(2, cue.O3(1, -1)))}")
print(f"cue.ir_mul.shape(...) helper: {cue.ir_mul.shape(cue.MulIrrep(2, cue.O3(1, -1)))}")

flat: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
mul_ir -> shape (2, 3): [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
   vector 0 = [0.0, 1.0, 2.0],
   vector 1 = [3.0, 4.0, 5.0]
ir_mul -> shape (3, 2): [[0.0, 3.0], [1.0, 4.0], [2.0, 5.0]]
x-components = [0.0, 3.0] across both vectors

cue.mul_ir.shape(...) helper: (2, 3)
cue.ir_mul.shape(...) helper: (3, 2)

21.4. Descriptors: One object instead of many instructions

This is the conceptual centre of the library. Recall from Lesson 03b that an o3.TensorProduct is defined by a list of instructions, each being a weighted path between input and output blocks. cuEquivariance replaces the instruction list with a single algebraic descriptor (cue.descriptors) for all operations you have already learned about during the course.

Let us compare the two descriptions of the same channel-wise convolution in e3nn and cuEquivariance.

# Create the input and output irreps
IRS_IN, IRS_SH = "32x0e + 32x1o", "0e + 1o"

# ===============================
# e3nn: a list of weighted paths
# ===============================
# Create the irreps objects
I1, I2 = o3.Irreps(IRS_IN), o3.Irreps(IRS_SH)

# Loop over the irreps and create the tensor product
out, instr = [], []
for i, (mul, ir1) in enumerate(I1):
    for j, (_, ir2) in enumerate(I2):
        
        # Selection rule
        for ir3 in ir1 * ir2:
            
            # Only include irreps with angular momentum <= 1
            if ir3.l <= 1:
                
                # Store the instruction
                instr.append((i, j, len(out), "uvu", True))
                
                # Add the output irreps
                out.append((mul, ir3))

# Create the tensor product
tp = o3.TensorProduct(I1, I2, o3.Irreps(out), instr,
                      shared_weights=False, internal_weights=False)

# Print the instructions
print(f"e3nn describes the operation with {len(tp.instructions)} separate instructions:")
for ins in tp.instructions:
    print(f"   in1={ins.i_in1} x in2={ins.i_in2} -> out={ins.i_out}  "
          f"mode={ins.connection_mode}  path_shape={ins.path_shape}")

# ===============================
# cuEquivariance: one polynomial
# ===============================

# Create the descriptor
desc = cue.descriptors.channelwise_tensor_product(
    cue.Irreps("O3", IRS_IN),
    cue.Irreps("O3", IRS_SH),
    cue.Irreps("O3", IRS_IN)
)

# Print the descriptor
print(f"\ncuEquivariance describes the operation with one {type(desc).__name__}:")
print(desc)

e3nn describes the operation with 5 separate instructions:
   in1=0 x in2=0 -> out=0  mode=uvu  path_shape=(32, 1)
   in1=0 x in2=1 -> out=1  mode=uvu  path_shape=(32, 1)
   in1=1 x in2=0 -> out=2  mode=uvu  path_shape=(32, 1)
   in1=1 x in2=1 -> out=3  mode=uvu  path_shape=(32, 1)
   in1=1 x in2=1 -> out=4  mode=uvu  path_shape=(32, 1)

cuEquivariance describes the operation with one EquivariantPolynomial:
╭ a=128x0e b=32x0e+32x1o c=0e+1o -> D=32x0e+32x0e+32x1o+32x1o
╰─ [ijk]·a[uv]·b[iu]·c[jv]➜D[kuv] ─ num_paths=4 i={1, 3} j={1, 3} k={1, 3} u=32 v=1

In the cuEquivariance descriptor expression, [ijk]·a[uv]·b[iu]·c[jv] -> D[kuv], the Einstein summation notation is implied. Here, the bracketed [ijk] denotes the Clebsch-Gordan tensor, a stands for the weights, b and c are the two inputs, and D is the output tensor. The symbols u and v denote the channel indices and i, j, k run over irrep components. Also, num_paths=4 counts the non-zero couplings.

Note that the two counts in e3nn and cuEquivariance do not match (5 vs 4). The e3nn loop keeps every output with \(l \le 1\), which includes a 1e block from \(1o \otimes 1o\). The descriptor is given an explicit irreps_out of 32x0e + 32x1o which causes 1e to be dropped. The point here does not concern the arithmetic of the counts but its form: e3nn hands the compiler a list of paths (instructions) to walk while cuEquivariance hands it one object (descriptor) to consume as a whole.

Let’s consider the hierarchy behind behind this design in cuEquivariance:

1. EquivariantPolynomial is the operation plus the irreps of its inputs and outputs (what we just printed above),

2. SegmentedPolynomial is the algebra with the representation labels stripped off, and

3. SegmentedTensorProduct is a single contraction within the segmented polynomial.

This hierarchy allows us to access the underlying structure of the equivariant operations at different levels of abstraction. Let’s inspecting various components of the cuEquivariance descriptor object

# Accessing the polynomial component
poly = desc.polynomial

# Printing the object type
print(f"EquivariantPolynomial.polynomial -> {type(poly).__name__}")

# The number of operations in the polynomial
print(f"Operations in this polynomial: {len(poly.operations)}")

# Print the inputs and output irreps in the descriptor
print(f"inputs : {[str(o) for o in desc.inputs]}")
print(f"outputs: {[str(o) for o in desc.outputs]}")

EquivariantPolynomial.polynomial -> SegmentedPolynomial
Operations in this polynomial: 1
inputs : ['128x0e', '32x0e+32x1o', '0e+1o']
outputs: ['32x0e+32x0e+32x1o+32x1o']

21.4.1. A descriptor you have already derived for MACE

In Lesson 10b, we’ve counted MACE’s body order: a correlation order \(\nu\) message is a \((\nu+1)\)-body interaction, and the \(B\)-basis collects products of \(\nu\) copies of the atomic basis \(A\). Here, we can use a similar analogy for a descriptor: one line per body order, with the number of symmetry-allowed couplings at each body order:

# Create a symmetric contraction descriptor
sym_desc = cue.descriptors.symmetric_contraction(
    cue.Irreps("O3", "32x0e + 32x1o"), cue.Irreps("O3", "32x0e"), [1, 2, 3])

# Print the symmetric contraction descriptor
print(sym_desc)

╭ a=32x0e+64x0e+64x0e b=32x0e+32x1o -> C=32x0e
│  []·a[u]·b[u]➜C[u] ─────────── num_paths=1 u=32
│  []·a[u]·b[u]·b[u]➜C[u] ────── num_paths=4 u=32
╰─ []·a[u]·b[u]·b[u]·b[u]➜C[u] ─ num_paths=10 u=32

The three rows correspond to \(\nu = 1, 2, 3\): b[u] alone, then b[u]·b[u], then b[u]·b[u]·b[u], with num_paths 1, 4 and 10. This growth in the number of terms is precisely why Lesson 10c built its body-ordering by iterating tensor products instead. This is also the reason behind superiority of the fused kernels over the iterative approach, which evaluate all three rows in one launch which we will see in Lesson 12b.

21.5. Verifying a first drop-in swap

The cuEquivariance-torch (cuet) modules are torch.nn.Modules. Setting the argument, layout=cue.mul_ir, enables the to accept e3nn-shaped tensors. Thus, we should be able to swap an e3nn module with its cuet counterpart in one line.

let us try this with a simple example involving spherical harmonics where the verification of the quivalence of the two implementations can be numerically exact.

# Create a set of random 3D vectors
v = torch.nn.functional.normalize(
    torch.randn(500, 3, device=device, dtype=torch.float64), dim=-1
)

# Compute spherical harmonics using both e3nn and cuEquivariance
sh_e3nn = o3.spherical_harmonics(
    [0, 1, 2], v, normalize=False, normalization="component"
)
sh_cuet = cuet.SphericalHarmonics(
    [0, 1, 2], normalize=False, device=device, math_dtype=torch.float64
)(v)

# Check the difference between the two implementations
print(f"spherical harmonics, max |e3nn - cuEq| = {(sh_e3nn - sh_cuet).abs().max():.2e}")

# Create a learnable module: same irreps in, same irreps out, e3nn memory layout
cuet_lin = cuet.Linear(
    cue.Irreps("O3", "32x0e + 32x1o"),
    cue.Irreps("O3", "16x0e + 16x1o"),
    layout=cue.mul_ir,
    device=device,
    dtype=torch.float32,
)

# Create linear module in torch
torch_lin = torch.nn.Linear(
    o3.Irreps("32x0e + 32x1o").dim,
    o3.Irreps("16x0e + 16x1o").dim,
    device=device,
    dtype=torch.float32,
)

# Apply the linear transformation in cuet and torch
h = torch.randn(4, o3.Irreps("32x0e + 32x1o").dim, device=device)
print(f"\ncuet.Linear: {tuple(h.shape)} -> {tuple(cuet_lin(h).shape)}")
print(f"torch.nn.Linear: {tuple(h.shape)} -> {tuple(torch_lin(h).shape)}")

spherical harmonics, max |e3nn - cuEq| = 6.66e-16

cuet.Linear: (4, 128) -> (4, 64)
torch.nn.Linear: (4, 128) -> (4, 64)

21.6. Summary

In this lesson, we have learned:

• Equivariant tensor products are not compute-bound: e3nn emits one kernel per weighted path, which makes them memory-bound at scale and dominated by kernel-launch overhead at the small system sizes where these models are usually trained at in this course.

• cuEquivariance replaces the path list with a descriptor which is an EquivariantPolynomial over a SegmentedPolynomial. It then compiles the whole descriptor object into one fused CUDA kernel.

• cue.Irreps names its group explicitly (compared to e3nn which does not). The cue.mul_ir argument reproduces e3nn’s memory layout while cue.ir_mul is cuEquivariant’s native mode. Passing the wrong layout does not raise an error, but returns wrong numbers.

• The cue.descriptors object builds the operations we already know from this course: channelwise_tensor_product, symmetric_contraction, linear, spherical_harmonics. Printing the descriptor is the fastest way to see what a kernel is about to be computed.

• Swapping a module is a one-line change. Verifying it is not optional, and Lesson 12b shows precisely how much verification headroom the fused kernels cost.

Next: Lesson 12bcoming soon measures the impact of all issues we’ve mentioned in this lesson on the convolution, the symmetric contraction, a whole model, and the cases where kernel fusion makes things slower.

21.7. Exercises

1. Read a descriptor you have not seen (Difficulty: 🌶️🌶️): Print cue.descriptors.fully_connected_tensor_product for 32x0e+32x1o against 0e+1o, and compare its index expression and num_paths with the channel-wise descriptor above. Which index disappears, and why does that make the weight count grow as \(m^3\) instead of \(m\)?

Solution

The channel-wise form carries u through to the output (➜D[kuv]): each channel is contracted with itself. Thus, the weights scale with \(m\). The fully connected tensor product sums over the input channel indices and produces an independent output channel. So, weights carry three channel indices at once and scale as \(m^3\). That is the difference between a convolution operation NequIP can afford at every layer and a dense mixing step it can use sparingly. In Lesson 12b, we show that a fused kernel version of this operation can help improve the performance.

2. Break a layout on purpose (Difficulty: 🌶️): Build two instance of cuet.Linear, one with layout=cue.mul_ir and the other with layout=cue.ir_mul. Then, feed both instances the same e3nn-shaped tensor, and compare the outputs. Finally, check whether either result is still equivariant under a rotation.

Solution

The outputs differ but neither function call raises an exception: this is the pitfall the layout section warns about. Only the mul_ir version stays equivariant on e3nn-shaped input. The ir_mul variant reinterprets the multiplicity axis as the irrep axis. As such, a rotation operation no longer acts on the irrep components as it should. The users should treat layout= as part of the tensor’s type, not as a tuning knob.

3. Count what kernel fusion can save (Difficulty: 🌶️🌶️): For a NequIP-style convolution at \(l_\max = 2\) and 32 channels, count e3nn’s instructions, then read num_paths off the corresponding cuEquivariance descriptor. Predict from those two numbers roughly how many kernel launches the fused version can remove per interaction block. Check your prediction against the timings in Lesson 12b.

Solution

Each e3nn instruction contributes several elementary operations (a gather, one or two einsums, a reshape). Therefore, a block with a dozen instructions issues on the order of tens of kernel launches. The fused kernel descriptor issues one kernel launch. That ratio explains why the speedup in Lesson 12b is largest at small system sizes, where the cost of kernel launches rather than arithmetic operations limit the pace and why the speedups shrink once the tensors become big enough that the GPU gets sufficiently busy.

21.8. References

• cuEquivariance documentation: Overview of the main concepts in cuEquivariance, API reference and the PyTorch/CUDA compatibility matrix.

• cuequivariance.IrrepsLayout: Comparing the mul_ir / ir_mul definitions mentioned in this lesson.

• Accelerating MACE with cuEquivariance: The application of cuEquivariance to the production mace-torch package.

• cuEquivariance on GitHub: The Apache-2.0 front end.