Senior ML Systems Engineer - GPU & HPC Optimizations

Sponsor Finder; Boston; Sur place; GBP 120 000 - 180 000; Plein temps

Résumé du poste

Proxima is seeking a Principal ML Performance Engineer to optimize training and inference for state-of-the-art models. You will profile PyTorch, write custom kernels in CUDA/Triton, and leverage compilers like torch.compile, TensorRT, and XLA to maximize throughput on large GPU clusters.

You\'ll scale distributed training across 32–64 nodes on GCP, manage memory scaling for large complexes, and build benchmarks and profiling tools for the research team.

Qualifications

• 6+ years in ML systems, HPC, or performance engineering.
• Deep knowledge of PyTorch internals and profiling bottlenecks.
• Experience with CUDA and Triton; reading Nsight outputs.
• Distributed multi-node training experience.
• Strong proficiency in Python and C++.
• Ability to name a model they made materially faster and quantify the improvement.

Responsabilités

• Profile and optimize training and inference for structural and generative models.
• Write and tune custom kernels (CUDA, Triton) and use compilers (torch.compile, TensorRT, XLA).
• Scale distributed training across 32–64 nodes with advanced parallelism and mixed precision.
• Reduce inference cost by memory scaling and throughput improvements for large GPU workloads.
• Manage GPU cluster efficiency on cloud, including scheduling and cost reporting.
• Develop benchmarks and profiling tools for the research team.

Connaissances

ML systems

HPC

Performance engineering

PyTorch internals

CUDA

Triton

Python

C++

Distributed training

Formation

BS/MS/PhD in CS/EE or related field

Outils

Nsight

TensorRT

XLA

torch.compile

Description du poste

Proxima is seeking a Principal ML Performance Engineer to optimize training and inference for state-of-the-art models. You will profile PyTorch, write custom kernels in CUDA/Triton, and leverage compilers like torch.compile, TensorRT, and XLA to maximize throughput on large GPU clusters.

You\'ll scale distributed training across 32–64 nodes on GCP, manage memory scaling for large complexes, and build benchmarks and profiling tools for the research team.