Simplifying Equivariant GPU Kernels through Tile-based Programming

E(3)-equivariant neural networks have demonstrated success across a wide range of 3D modeling tasks. Until recently, they were bottlenecked due to their high memory and wall-time requirements. In this thesis we first provide an overview of recent GPU kernel efforts by both academia and industry that address this issue. These approaches tradeoff performance for engineering complexity, while still being algorithmically bottlenecked at 10 % GPU utilization. We instead trade off engineering complexity for performance. This not only lowers the barrier to GPU programming but also builds an abstraction layer to reason about future algorithmic innovations that can improve GPU utilization. Our kernel 𝐵3, based on the tiling- optimizations in just 100 lines of PyTorch-like code. We explore the performance-simplicity tradeoff with two case studies and demonstrate the practicality of our kernel workflow through downstream integration with a production model. We hope this work serves as inspiration to broaden and deepen existing equivariant kernel efforts.

Name thesis:degree_name

Master

Department dc:contributor.department

Massachusetts Institute of Technology. Center for Computational Science and Engineering

Grantor dc:publisher

Massachusetts Institute of Technology

Year dc:date.issued

2025

Author dc:creator

• Kotak, Mit

Advisor dc:contributor.advisor

• Smidt, Tess E.

Statement dc:rights

• Attribution-NonCommercial-NoDerivatives 4.0 International (CC BY-NC-ND 4.0)
• Copyright retained by author(s)

Licence dc:rights.uri

• https://creativecommons.org/licenses/by-nc-nd/4.0/

Handle dc:identifier.uri

https://hdl.handle.net/1721.1/164822

OAI identifier oai:identifier

oai:dspace.mit.edu:1721.1/164822