Carlos ALIAGAS CASTELL, Sergio GÓMEZ JIMÉNEZ

-

Optimizing EGNNs via OpenAI Triton

Equivariant Graph Neural Networks (EGNNs) are highly effective for 3D molecular modeling, but their scalability is severely bottlenecked by hardware memory constraints. In standard frameworks like PyTorch, materializing intermediate edge tensors for the backward pass creates a crippling $\mathcal{O}(E \cdot H)$ space complexity, leading to premature Out-Of-Memory (OOM) failures. This thesis introduces a custom, highly optimized GPU kernel built with the OpenAI Triton compiler to entirely eradicate this memory wall. By fusing the network's multi-layer perceptrons and bypassing the standard automatic differentiation tape, the kernel traps intermediate states in L1 Shared Memory and utilizes on-the-fly register-level rematerialization during the backward pass. Empirical evaluations on an NVIDIA RTX A4000 GPU confirm a strict $\mathcal{O}(N \cdot H)$ memory scaling. This architectural intervention increased the absolute graph capacity limit by 13.5x (from 225,000 to 2.7 million nodes) and delivered a $\sim$7.5x computational speedup on the forward pass. Ultimately, this thesis successfully shifts the EGNN from a memory-bound limitation to a compute-bound reality.

08-DES-2025

Graph Neural Networks, PyTorch, Triton

-