E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory

Equivariant Graph Neural Networks (EGNNs) have become a widely used approach for modeling 3D atomistic systems. However, mainstream architectures face critical efficiency bottlenecks due to the explicit construction of geometric features or dense tensor products on every edge. To overcome this, we introduce E2Former-V2, an efficient architecture that integrates algebraic sparsity with hardware-aware execution. We introduce Equivariant Axis-Aligned Sparsification (EAAS), which leverages an $\mathrm{SO}(3)\!\rightarrow\!\mathrm{SO}(2)$ change of basis to convert dense Wigner-$6j$ tensor contractions into sparse, parity-based re-indexing operations. Building on this representation, we propose On-the-Fly Equivariant Attention, a fully node-centric mechanism implemented via a fused Triton kernel. By eliminating materialized edge tensors and maximizing SRAM utilization, our kernel achieves up to 20$\times$ higher TFLOPS than standard implementations. Experiments on SPICE and OMol25 show that E2Former-V2 preserves predictive accuracy while substantially accelerating inference, demonstrating the practical feasibility of large equivariant transformers on commodity GPUs. Our released code can be found at https://github.com/IQuestLab/UBio-MolFM/tree/main