GPU Kernel Modifiers

NequIP GNN models support several GPU kernel modifiers that can significantly speed up both training and inference by replacing standard operations with optimized implementations.

Available GPU Kernel Modifiers

	OpenEquivariance	CuEquivariance (alpha)
Modifier Name	enable_OpenEquivariance	enable_CuEquivariance
Supported Devices	NVIDIA GPUs, AMD GPUs (HIP)	NVIDIA GPUs
Training	✅ Stable	🔨 Work in progress
ASE (TorchScript)	✅ Stable	🔨 Work in progress
ASE (AOT Inductor)	✅ Stable (requires PyTorch >= 2.9)	🔨 Work in progress
LAMMPS ML-IAP	✅ Stable	🔨 Work in progress

• OpenEquivariance Acceleration
• CuEquivariance Acceleration