Senior Software Engineer — cuEquivariance

We're looking for a Senior Software Engineer to join the cuEquivariance team, which builds the computational foundation for the next generation of biological discovery. The team uses NVIDIA GPUs to accelerate geometric neural networks, enabling researchers to train and deploy equivariant models at scale.

Key Responsibilities:

• Build, implement, and optimize CUDA kernels for equivariant neural network primitives, targeting peak performance across NVIDIA GPU generations.
• Be responsible for the end-to-end delivery of GPU-accelerated geometric ML primitives, from implementation to validated, production-quality software that external frameworks depend on.
• Build and maintain the interfaces for PyTorch and JAX that expose cuEquivariance primitives to application developers and researchers.
• Drive CI/CD infrastructure for multi-GPU kernel builds, automated correctness testing, and performance regression tracking.
• Collaborate with Applied Science and research teams to evaluate new equivariant architectures and translate prototypes into production kernels.
• Engage directly with third-party framework developers and partners to align on interfaces and ensure delivered software integrates cleanly into production pipelines.

Requirements:

• 6+ years of software engineering experience with a strong background in CUDA and GPU programming.
• Deep proficiency in C++ and Python, with experience building and shipping production libraries used by external developers.
• Good foundation in GPU computing, including memory hierarchy, warp-level execution, occupancy, and performance profiling methodology.
• Experience building or contributing to production scientific software libraries, ML frameworks, or developer-facing GPU APIs.
• Familiarity with concepts in geometric machine learning, such as equivariance, group representations, irreducible representations, or tensor products.

Nice to Have:

• Chipped in to or deeply used a major neural network framework that respects equivariance, such as e3nn, MACE, NequIP, SE(3)-Transformers, or similar.
• Hands-on experience with Triton kernel development or other GPU kernel authoring tools alongside CUDA.
• Experience with mixed-precision or tensor-core-aware algorithm design for scientific or ML workloads.
• PhD or equivalent experience in computational chemistry, biophysics, physics, or computer science with a focus on geometric deep learning or HPC.
• Contributions to open-source geometric ML or GPU computing projects.

Benefits:

• Highly competitive salaries and a comprehensive benefits package
• Eligibility for equity and benefits
• Opportunity to work with a talented team of engineers and researchers
• Access to cutting-edge technology and resources
• Flexible work arrangements and professional development opportunities
• A dynamic and inclusive work environment
• Opportunities for growth and advancement within the company
• Recognition and rewards for outstanding performance and contributions
• A commitment to diversity, equity, and inclusion
• A focus on employee well-being and work-life balance
• A range of employee benefits and perks, including health insurance, retirement plans, and paid time off
• Opportunities for professional development and continuing education
• A culture of innovation and experimentation
• A commitment to social responsibility and community engagement
• A focus on customer satisfaction and feedback
• A recognition of the importance of work-life balance and flexibility
• A commitment to diversity, equity, and inclusion
• A focus on employee well-being and mental health
• A range of employee benefits and perks, including flexible work arrangements, remote work options, and employee discounts
• Opportunities for professional development and continuing education
• A culture of innovation and experimentation
• A commitment to social responsibility and community engagement
• A focus on customer satisfaction and feedback
• A recognition of the importance of work-life balance and flexibility
• A commitment to diversity, equity, and inclusion
• A focus on employee well-being and mental health
• A range of employee benefits and perks, including flexible work arrangements, remote work options, and employee discounts
• Opportunities for professional development and continuing education
• A culture of innovation and experimentation
• A commitment to social responsibility and community engagement
• A focus on customer satisfaction and feedback