Source-linked AI summary
E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU
Olivier Peltre, Armand Picard, Adrien Pichard, Miguel Bragança, Luca Giacomoni, Valentin Heyraud, Zachary Weller-Davies, Christoph Brunken, Jules Tilly
TL;DR
Equivariant architectures are often criticized as computationally heavy, while efficient full Clebsch-Gordan operations on GPUs and TPUs have been a practical challenge. E3J provides an open-source JAX backend with CUDA and Pallas kernels; its benchmarks report high throughput and speedups over established backends.
Problem
Equivariant architectures are often criticized as computationally heavy, motivating interest in learned-equivariance transformer alternatives.
Source: §4 · p. 9Method
Results
Across tensor-product and model-level benchmarks, e3j reports high hardware utilization and backend speedups, including 2.7 TB/s on H100 GPU for forward tensor products (81% of memory bandwidth).
Source: §1 · p. 2Takeaways & Limitations
Properly engineered full Clebsch-Gordan tensor products and message-passing convolutions can reach over 80% of HBM throughput on GPUs and TPUs.
Source: §4 · p. 9Dedicated Clebsch-Gordan kernels remain disadvantaged in L scaling relative to SO(2) convolution, although full tensor products apply to arbitrary feature-vector pairs.
Source: §4 · p. 9
Abstract
from arXiv · showhide
We present e3j, a fast Euclid-equivariance backend for geometric deep learning applications with JAX bindings for GPU and TPU. Leveraging both optimized CUDA and Pallas kernels and algorithmic improvements, the library achieves state-of-the-art throughput and runtime on both forward and backward paths. On a machine learning interatomic potential (MLIP) use case, it outperforms established backends, measuring up to 34% speed-up over cuEquivariance on water box NPT simulation using MACE, while remaining fully open source. E3j achieves over 80% efficiency over the H100 maximum memory bandwidth on tensor product operations, and in many cases more than doubles throughput of message passing convolutions forward compared to previously available backends. In addition, with the release of dedicated Pallas TPU kernel, e3j opens the possibility of large scale equivariant deep learning workloads on TPU architectures, which has so far been difficult to achieve. Our benchmarks show that e3j also achieves over 80% of a TPUv6e memory bandwidth, up to one order of magnitude more than e3nn-jax. The library is available on GitHub, PyPI and is released under an open source Apache 2.0 license.
1 INTRODUCTION
E3J provides open-source, modular JAX building blocks for E(3)-equivariant networks, with CUDA and Pallas kernels targeting GPUs and TPUs. Its reported benchmarks show high hardware utilization and speedups in tensor-product and end-to-end MLIP workloads.
1 INTRODUCTION
- E3J packages core E(3)-equivariant neural-network building blocks in an open-source, modular JAX backend.The library integrates accelerated kernels into model architectures.Source: §1 · p. 1
- CUDA kernels target deterministic GPU operations and broad JAX-version compatibility through ahead-of-time compilation.Pallas GPU kernels prioritize performance through JIT compilation, with tradeoffs in backward-pass determinism, older JAX compatibility, and support for some operations.Source: §1 · p. 1
- Pallas TPU kernels are included to deliver optimized throughput on TPUs.The library benchmarks both equivariant tensor products and full message-passing convolutions in MLIP workloads.Source: §1 · p. 2
- 34% speedup over cuEquivariance in MACE simulations, alongside 2x speedup over OpenEquivariance, is reported for GPU workloads.The benchmarks also reach 64 thousand atoms at approximately 40 connectivity degree on MACE before H100 memory overflow.Source: §1 · p. 2
- 2.7 TB/s on H100 and 1.33 TB/s on TPUv6.2e each reach 81% of memory bandwidth for forward tensor products.End-to-end benchmarks report about 5x force-inference speedup over e3nn on TPU.Source: §1 · p. 2
2 RELATED WORK
Prior work spans general equivariance libraries, specialized kernels, and algorithmic changes to tensor products and message passing. E3J’s stated distinction is a standalone JAX API for GPU and TPU execution.
2 RELATED WORK
- Existing equivariance libraries span Torch and JAX packages, CUDA kernel generation, and proprietary CUDA kernels, with different flexibility and operation coverage.OpenEquivariance focuses on tensor products and message-passing convolutions; cuEquivariance offers a toolset described as as complete as e3nn.Source: §2 · p. 2
- E3J’s stated distinction is a standalone, platform-agnostic JAX API for GPU and TPU, while near-optimal Clebsch-Gordan tensor-product solutions already exist for NVIDIA GPUs.The related-work overview includes e3nn, e3x, OpenEquivariance, and proprietary cuEquivariance.Source: §2 · p. 2
- Cubic scaling in L can involve compromises in applicability or expressivity, and practical MLIP workloads often focus on scaling in N at fixed L.Reparameterizations for improved L scaling can add overhead that may not pay off at small L.Source: §2 · p. 3
- Some cubic-in-L tensor-product methods incorporate antisymmetric terms to fully simulate Clebsch-Gordan products, but can collapse output multiplicity and reduce expressivity.The cited approaches include Gaunt tensor products, VSTP, and a closed-form formulation.Source: §2 · p. 3
- SO(2) convolution uses edge reference frames and Wigner-D rotations to collapse harmonic features to m = 0, while E2former replaces projected displacements with differences of projected positions.The E2former construction uses a binomial expansion to make the number of tensor products scale with nodes rather than edges.Source: §2 · p. 3
3 RESULTS
e3j benchmarks dedicated CUDA and Pallas kernels for tensor products and message-passing convolutions, alongside end-to-end MACE and NequIP workloads. The reported results include strong GPU and TPU throughput comparisons, while padding-edge handling addresses an overhead observed in deterministic convolution simulations.
3 RESULTS
- The evaluation combines module-level tensor-product and message-passing benchmarks with illustrative end-to-end MACE and NequIP tests using numerically equivalent backends in a fork of the mlip library.The end-to-end comparison is not an exhaustive assessment across models or application fields.Source: §3 · p. 4, §3.2 · p. 8
3.1.1 TENSOR PRODUCT
- Up to 80% HBM throughput in forward and backward, CGTP is not itself a compute bottleneck for L ≤ 3; TPU Pallas reaches an order of magnitude higher throughput than e3nn-jax.On GPU, Pallas is nearly on par with cuEquivariance, while CUDA is similar forward and slightly below Pallas and cuEquivariance backward; all three outpace OpenEquivariance.Source: §3.1.1 · p. 5, Figure 1 · p. 5
- For L ≤ 3, e3j kernels are 2–4× faster than Gaunt tensor products on symmetric paths and 4–6× faster than VSTP on skew-symmetric paths; the curves cross at higher degrees.The symmetric-path crossings occur at L = 6 forward and L = 5 backward; skew-symmetric crossings occur one degree later.Source: §3.1.1 · pp. 5–6
3.1.2 MESSAGE PASSING CONVOLUTION
- Over 2× cuEquivariance’s forward throughput in most GPU settings, Pallas GPU also slightly exceeds it backward; TPU convolution is over one order faster forward and nearly one order faster backward than e3nn-jax.The GPU comparison is shown in Table 1 and TPU/GPU convolution throughput in Figure 2.Source: Figure 2 · p. 6, §3.1.2 · p. 6, Table 1 · p. 7
- Against JAX-based SO(2) convolution, e3j CUDA and Pallas GPU are 4–6× faster through L = 3, with curves crossing at L = 4 forward and L = 5 backward.The comparison force-collapses multiplicity to maintain equivalent expressivity.Source: §3.1.2 · p. 8
3.2 END-TO-END BENCHMARKS
- Padding-edge skipping avoids work on a shared padding node, where OpenEquivariance’s deterministic convolution can slow simulations by over 30 ms/step at 25% padding.The slowdown grows with the number of padding edges.Source: §3.1.2 · p. 6, §3.2 · p. 8
4 DISCUSSION
The discussion argues that engineered Clebsch–Gordan tensor products and message-passing convolutions can run efficiently on GPUs and TPUs. The appendix supplies mathematical background on Euclidean equivariance, harmonic representations, and core operations.
4 DISCUSSION
- Full Clebsch–Gordan tensor products and associated message-passing convolutions reach over 80% of HBM throughput on GPUs and TPUs.The authors argue this could put equivariant networks on a more even footing with highly engineered transformer architectures.Source: §4 · p. 9
- Dedicated Clebsch–Gordan kernels are disadvantaged in scaling with L versus methods such as SO(2) convolution, but apply to arbitrary feature-vector pairs.The comparison is not fully even: SO(2) convolution would also need dedicated engineering.Source: §4 · p. 9
- Achieved throughput remains below the maximum bandwidth of the H100 GPU and v6e TPU, leaving room for further improvement.Source: §4 · p. 9
A MATHEMATICAL BACKGROUND
- The mathematical-background appendix situates representation theory in applications to quantum mechanics and points readers to further introductions.Source: A MATHEMATICAL BACKGROUND · p. 13
A.1 EUCLIDEAN EQUIVARIANCE
- The appendix frames Euclidean equivariance through E3, whose transformations combine translations, rotations, and reflections, and functions that transform consistently with their inputs.Source: §A.1 · p. 13
A.2 HARMONIC POLYNOMIALS
- Degree-l harmonic polynomials form SO3 sub-representations; irreducible representations have dimension 2l + 1, and E3 irreducibles also carry parity labels.Source: §A.2 · pp. 13–14
A.3 KEY OPERATIONS IN SCOPE
- The core operations include harmonic embeddings, linear mixing, Clebsch–Gordan tensor products, and message passing; the library does not yet include operations specific to particular MLIP models.The named example is MACE's Symmetric Contraction module; the appendix benchmarks its existing cuEquivariance kernel.Source: §A.3 · pp. 14–15
B.1 GAUNT AND VECTOR SIGNAL TENSOR PRODUCTS
The appendix compares e3j tensor products with integral-based Gaunt and vector signal products, then with e3x's matrix tensor product. e3j is faster in the reported practical degree ranges, although the matrix tensor product's scaling narrows its gap as L grows.
B.1 GAUNT AND VECTOR SIGNAL TENSOR PRODUCTS
- The GTP covers symmetric paths and VSTP skew-symmetric paths; both use spherical t-design quadrature, while e3j selects matching paths through parity and collapses outputs by degree.Source: §B.1 · pp. 15–16
- For L ≤3, e3j is 2 to 4 times faster than GTP and 4 to 6 times faster than VSTP; forward curves cross at L=6/7 and backward curves at L=5/6.The authors report e3j remains faster in the regime practical for MLIP applications.Source: Figure 5 · p. 15, §B.1 · p. 16
B.2 MATRIX TENSOR PRODUCTS FROM e3x
- MTP maps feature vectors to square matrices, multiplies those matrices, and maps back; its overall scaling is O(L^4), despite the matrix product's O(L^3) scaling.The comparison adds input and output linear mixings to make the e3j tensor product equivalent to the FusedTensor operation.Source: §B.2 · p. 16
- For L ≤8, e3j is faster than e3x's MTP, and remains faster over the tested domain as MTP's O(L^4) scaling narrows the gap.Source: §B.2 · p. 16
B.3 SO2 CONVOLUTION
With output multiplicities collapsed, e3j convolution is 4–6× faster than SO2 at ℓmax≤3 and matches it at ℓmax=4.
B.3 SO2 CONVOLUTION
- Collapsing output multiplicities sums isomorphic copies of each irreducible output space, reducing memory traffic at a cost to expressivity.Source: §B.3 · p. 16
- A faithful SO2 comparison would require comparable low-level engineering, since the plain JAX SO2 implementation already compiles efficiently with XLA.Source: §B.3 · p. 17
- 4–6× faster than SO2 at ℓmax≤3, e3j matches SO2 at ℓmax=4 in the comparison with collapsed output multiplicities.The benchmark compares runtime scaling on an H100 with N=2048 nodes, 45 edges per node, and C=128 channels.Source: §B.3 · p. 16, Figure 7 · p. 17
C.1 SYMMETRIC CONTRACTION BENCHMARKS
The benchmark contrasts e3j’s power-expansion and projection implementation with dedicated SymmetricContraction, which can fuse the operations and skip a large intermediate array.
C.1 SYMMETRIC CONTRACTION BENCHMARKS
- The e3j-based SymmetricContraction expands node features quadratically or cubically with tensor-product kernels, then applies species-dependent linear projection.The power expansion uses channel mixing mode MAP5.Source: §C.1 · p. 17
- Power expansion creates large feature multiplicities that collapse after projection, making the naive implementation suboptimal; fusion can skip materializing the large intermediate array.Source: §C.1 · pp. 17–18
- The original MACE device-code algorithm is not currently usable without resolving parameter transformations and Clebsch-Gordan normalization choices needed for numerical consistency.Its plain JAX implementation performs worse than naive expansion and projection when an efficient CG tensor product is available.Source: §C.1 · p. 18
- All NPT simulation benchmarks use the same e3j-based SymmetricContraction, isolating the convolution backend while allowing one model checkpoint to be reused.Source: §C.1 · p. 18
- At the time of writing, e3j lacks a dedicated SymmetricContraction kernel, and its AOT CUDA kernels do not reach cuEquivariance performance for this operation.The authors suggest coefficient inlining in a JIT-compiled Pallas or NVRTC kernel may be needed to match that performance.Source: §C.1 · p. 18
C.2 DETERMINISM AND DEVIATIONS OF END-TO-END PREDICTIONS
The section examines sources of nondeterminism in energy predictions and their implications for simulations, while noting that force prediction can mitigate some sources.
C.2 DETERMINISM AND DEVIATIONS OF END-TO-END PREDICTIONS
- For Monte Carlo barostat relaxation of periodic cells, deterministic energy predictions are a hard constraint because small energy fluctuations affect Metropolis-Hastings rejection.The passage warns that nondeterministic message aggregation would lead to exploding simulations in this setting.Source: §C.2 · p. 18
- Energy-prediction noise may come from adding large constant atomic energies E0 and from highly concurrent aggregation of node-energy contributions into graph energies.Source: §C.2 · pp. 18–19
- Force prediction may eliminate constant-energy addition and replace nondeterministic scatter operations with deterministic gather operations.Source: §C.2 · p. 19
- The mitigation is not automatic for differentiation, and a rigorous analysis of nondeterminism’s effects on model training is outside this work’s scope.Source: §C.2 · p. 19
- Figure 9 evaluates energy deviations across 100 runs on eight water boxes, varying energy aggregation schemes; missing boxes indicate rigorously deterministic runs.The plot shows 5th/95th percentiles as whiskers and 25th/75th percentiles as boxes.Source: Figure 9 · p. 19
D.1 THROUGHPUT AND HBM
The section describes e3j’s throughput framing and kernel designs, including CUDA tensor-product optimizations and differentiable rules. It also records precision conditions and limits on the MLIP comparisons.
D.2 PRECISION
- All tested backends show accuracies around 4 × 10−8 for tensor products and 1 × 10−7 for message passing, except e3nn at about 4 × 10−4 without highest matmul precision.The comparison uses a common deterministic float64 reference.Source: §D.2 · p. 20
D.3 CONSIDERATIONS REGARDING PALLAS AND CUDA KERNELS
- Pallas and CUDA target GPU and TPU execution; Pallas JIT specializes kernels, while CUDA binaries offer relative stability across JAX versions.Pallas specialization can inline Clebsch–Gordan coefficients and reduce memory traffic from unified SMEM/L1 cache.Source: §D.3 · p. 20
D.4 MLIP INTEGRATION
- The MLIP comparisons checked that models matched numerically, but did not test third-party model implementations or substitutions of additional MACE blocks.The authors cite the effort needed to establish faithful comparisons as the reason these checks were not performed.Source: §D.4 · p. 21
E TENSOR PRODUCT AND MESSAGE PASSING CONVOLUTION KERNELS
- The appendix presents CUDA, Pallas GPU, and Pallas TPU kernel sets, with CUDA emphasizing broad GPU applicability and determinism.Pallas GPU focuses on performance for recent GPU architectures and JAX versions.Source: E TENSOR PRODUCT AND MESSAGE PASSING… · p. 21
E.1 MATHEMATICAL DETAILS
- The kernels define reverse-mode rules by transposing linearized operations, mapping output cotangents back to cotangents on the primal inputs.The rules use bilinearity and index permutations for tensor-product backward operations.Source: §E.1 · p. 23
E.2 CUDA KERNELS
- CUDA’s sparse tensor-product algorithm addresses the JAX implementation’s poorly scaling final scatter-reduction step.It processes trailing channels in parallel and sequentially reduces the feature axis, avoiding inter-thread communication.Source: §E.2.1 · p. 24
E.2.1 CUDA TENSOR PRODUCT OPERATION
- Vectorizing across 2 or 4 channels nearly doubles CUDA tensor-product throughput by amortizing coefficient-load costs.Four-fold vectorization reduces the stated shared-memory traffic per batch from 24 to 12 bytes.Source: §E.2.1 · p. 25
E.2.2 CUDA MESSAGE PASSING CONVOLUTION
The CUDA convolution streams through edges and accumulates messages deterministically using CSR adjacency. The Pallas GPU design specializes sparse coefficients at trace time and orders work by receiver, trading forward aggregation properties against backward-gradient locality and atomics.
E.2.2 CUDA MESSAGE PASSING CONVOLUTION
- CUDA convolution streams through edges and uses CSR-based receiver aggregation to avoid atomics, though its operand and message buffers increase shared-memory pressure.The backward design retains relative algorithmic simplicity, which constrains occupancy at sufficiently high levels.Source: §E.2.2 · p. 26
E.3 PALLAS GPU CONVOLUTION KERNEL
- Pallas GPU specializes Clebsch–Gordan coefficients at trace time, removing runtime coefficient loads, loop overhead, and sparse-indexing overhead.Inlining makes contraction paths visible to the compiler for instruction reordering.Source: §E.3 · p. 27
- Receiver-ordered Pallas aggregation gives each receiver a fixed, atomics-free forward reduction, while backward gradients may require atomic accumulation across CTAs.Receiver ordering also improves reuse of receiver cotangents, but gradient contributions need not stay local to one receiver.Source: §E.3 · p. 28
E.4 PALLAS TPU KERNEL
The TPU kernel tiles receiver-ordered edges, computes and reduces messages while retaining reduction state across tiles, and avoids materializing the full edge-message tensor in HBM. The appendices explain how GPU and TPU execution and memory models shape kernel design, and report the TPU kernel’s VMEM use.
E.4 PALLAS TPU KERNEL
- The TPU kernel processes receiver-ordered edges in tiles, fuses message formation with segmented receiver reduction, and carries reduction state across tile boundaries.VMEM holds edge operands and messages, while SMEM holds edge indices; this avoids materializing the complete edge-leading message tensor in HBM.Source: §E.4 · pp. 29–30
- 13.1 and 10.1 MiB are the reported peak VMEM usages, against a 16 MiB per-core budget.The kernel’s own buffers occupy 4.3 MiB in the forward pass and 5.4 MiB in the backward pass at lmax = 3, Nq = 128, and fp32; compiler spills raise peak use.Source: §E.4 · p. 32, Table 4 · p. 32
F.1 GPU AND TPU ACCELERATORS COMPARISON
- GPU and TPU architectural differences lead to different kernel organizations, including distinct fusion strategies and temporary-buffer sizes.The paper highlights reducing memory traffic and avoiding intermediate materialization through operation fusion as factors that can affect end-to-end performance.Source: §F.1 · p. 32
F.2 OVERVIEW OF RELEVANT GPU CONCEPTS
- On GPUs, low-arithmetic-intensity kernels can be HBM-bandwidth-bound, while higher-intensity kernels may instead be compute-bound.The appendix describes peak HBM bandwidth as a useful speed-of-light estimate for memory-bound kernels.Source: §F.2 · p. 34
F.3 OVERVIEW OF RELEVANT TPU CONCEPTS
- The TPU v6e TensorCore uses 8 × 128 vector-register tiles, and poorly aligned or very small arrays may waste vector capacity or require rearrangement.Its VPU handles general vector operations, while its MXUs specialize in dense matrix multiplication.Source: §F.3 · p. 34
- TPU kernels stage large tensors from HBM through software-managed on-chip memories, with VMEM for arrays and separate SMEM for scalar data such as indices.TPU v6e has approximately 128 MiB of VMEM and 1 MiB of SMEM per TensorCore; asynchronous DMA can prefetch tiles during computation.Source: §F.3 · p. 34
F.3.1 RELEVANT DISTINCTION BETWEEN GPU AND TPU ARCHITECTURES
The appendix contrasts GPU and TPU execution and memory models, then outlines constructions for spherical harmonics and Clebsch–Gordan coefficients. These constructions use ladder operators and eigenvalue equations, with the latter solvable as a triangular system.
F.3.1 RELEVANT DISTINCTION BETWEEN GPU AND TPU ARCHITECTURES
- A TPU TensorCore operates on wide vector tiles and has much larger VMEM than a GPU SM’s shared memory, motivating different GPU and TPU kernel organizations.The appendix contrasts GPU CTA-local receiver reductions and streamed edge data with TPU tile-wide operations over larger VMEM-resident windows.Source: §F.3.1 · pp. 34–35
G.1 RECONSTRUCTION OF HARMONIC POLYNOMIALS
- Ladder operators raise or lower the magnetic quantum number m by one, and the appendix gives spherical harmonics forms that vary with longitude as e^imϕ.For m = ±l, the degree-l polynomial has the form c_l(x ± iy)^l; lower absolute spins are also described.Source: §G.1 · p. 35
G.2 RECONSTRUCTION OF CLEBSCH GORDAN COEFFICIENTS
- The Clebsch–Gordan construction obtains top-spin generators by solving eigenvalue equations expressed as a triangular system.Lower-spin eigenvectors are obtained by iterating the lowering operator from the top-spin eigenvector, up to normalization.Source: §G.2 · pp. 35–36