GPU Benchmark Engineer for AI Cloud Infra

Nebius; United States; Remote; USD 150,000 - 190,000; Full time

Job summary

Nebius is seeking a highly skilled ML/AI Engineer to lead benchmarking of GPU platforms for ML and AI workloads. You will profile GPU performance at system and kernel levels and evaluate across platforms, architectures, and software stacks, including CUDA and ROCm.

You will debug and optimise workloads and run experiments to understand interconnect impacts. Join a team defining performance standards for next‑gen AI hardware, with dashboards and tooling to visualize trends.

Qualifications

• Strong understanding of ML theory and fundamentals.
• Deep knowledge of performance aspects of large neural networks training and inference.
• Hands-on experience with PyTorch, JAX, Megatron-LM, Tensort-LLM.
• Solid grasp of the GPU stack: CUDA, NCCL, drivers, libraries.
• Experience with containerized environments (Docker, Kubernetes).
• Excellent communication and ability to work independently.

Responsibilities

• Profile and analyze GPU performance at system and kernel levels with hardware teams.
• Evaluate and compare GPU performance across platforms, architectures and stacks (CUDA, ROCm).
• Debug and optimize ML workloads to run efficiently on GPU hardware and resolve bottlenecks.
• Perform acceptance testing for new GPU clusters ensuring performance and compatibility for AI workloads.
• Experiment with diverse GPU configurations to assess interconnect strategies and system-level optimisations.
• Develop tools and dashboards to visualize performance metrics, bottlenecks and trends.
• Contribute to internal tooling, frameworks and best practices.

Skills

ML theory

Large-scale model perf

PyTorch

JAX

Megatron-LM

Tensor-LLM

CUDA

NCCL

Docker

Kubernetes

Python

Performance profiling

Communication

Independent work

Tools

CUDA

NCCL

Nsight

nvprof

Docker

Kubernetes

perf

TensorRT

Job description

Nebius is seeking a highly skilled ML/AI Engineer to lead benchmarking of GPU platforms for ML and AI workloads. You will profile GPU performance at system and kernel levels and evaluate across platforms, architectures, and software stacks, including CUDA and ROCm.

You will debug and optimise workloads and run experiments to understand interconnect impacts. Join a team defining performance standards for next‑gen AI hardware, with dashboards and tooling to visualize trends.