NVIDIA/cuEquivariance v0.8.0

v0.8.0

Repository: NVIDIA/cuEquivariance

Tag: v0.8.0

Published: 2025-12-08T21:27:57Z

Prerelease: no

Release notes:

Added

• Support for CUDA 13 on ARM
• [Torch/JAX] Blackwell-optimized BF16/FP16 forward and backward kernels for cuet.triangle_attention (runs on compute capabilities 10.0 and 10.3). These kernels provide superior performance especially for long sequences and higher head dimensions. This is only supported on cu13 builds as of this release

Bug fix

• [Torch/JAX] Fixed index overflow and out of bound issues leading to illegal memory access in cuet.triangle_attention

Notes

• [Torch/JAX] Blackwell-optimized kernels require the sequence length N to be a multiple of 8 for the forward pass; pad the sequence if necessary
• [Torch/JAX] Blackwell-optimized kernels are currently supported only for CUDA 13 builds

What's Changed

• 16bit math datatype for uniform_1d by @mariogeiger in https://github.com/NVIDIA/cuEquivariance/pull/188
• [JAX] code examples by @mariogeiger in https://github.com/NVIDIA/cuEquivariance/pull/156
• [docs] add mace perf figure in docs by @mariogeiger in https://github.com/NVIDIA/cuEquivariance/pull/192
• Adding examples to the documentation by @phiandark in https://github.com/NVIDIA/cuEquivariance/pull/191
• Small SegmentedPolynomial improvements by @phiandark in https://github.com/NVIDIA/cuEquivariance/pull/210
• fix by @mariogeiger in https://github.com/NVIDIA/cuEquivariance/pull/212
• Added wrappers for ONNX/TRT initialization by @borisfom in https://github.com/NVIDIA/cuEquivariance/pull/214
• [JAX] triatt 16bit fix by @mariogeiger in https://github.com/NVIDIA/cuEquivariance/pull/213
• [Torch] Improve error messages for uniform 1d by @mariogeiger in https://github.com/NVIDIA/cuEquivariance/pull/217
• triatt jax bwd dtype issue by @mariogeiger in https://github.com/NVIDIA/cuEquivariance/pull/219
• doc-string updates+ rel-notes by @hsadasivan in https://github.com/NVIDIA/cuEquivariance/pull/220

Full Changelog: https://github.com/NVIDIA/cuEquivariance/compare/v0.7.0...v0.8.0