Senior/Staff ML Engineer, AI-Assisted Efficiency Optimization

Nuro • Mountain View, CA

58d • $193,930 - $352,290

About The Position

We are looking for an experienced, goal-oriented Machine Learning Engineer to lead our GenAI-assisted GPU kernel optimizations project delivery. This project stands to serve the dual objectives of improving maximum flop utilization and reducing inference latency, to surpass human-engineered performance through the application of advanced AI technologies.

Requirements

• Bachelor's or Master's Degree in Computer Science, Engineering, or a related field.
• Extensive experience with AI models, including but not limited to LLMs.
• Strong understanding of retrieval-augmented generation (RAG) and Language Models (LLMs).
• Ability to self-motivate, undertake complex assignments independently, and while enjoying to balance innovative exploration with practical considerations.

Nice To Haves

• Experience with GPU programming, CUDA, Triton, and Machine Learning algorithms.
• Familiarity with optimization techniques such as neural architecture search (NAS).
• Hands-on experience with optimization programs, such as Google's AutoFDO.

Responsibilities

• Implement AI-driven methods for GPU kernel optimizations to enhance program efficiency and performance leveraging tools such as CompilerGym, KernelLLM, etc.
• Develop strategies for resource-efficient deployment of intelligent optimization processes using frameworks such as AutoFDO.
• Guide high-level optimizations that include both highly optimized kernels, as well as AI-driven neural architecture search (NAS) to optimize training efficiency and inference latency.
• Assess performance improvements via evaluation metrics and real-world feedback.
• Collaborate with internal teams to benchmark processes and strategies for future AI-assisted optimizations.
• Utilize state-of-the-art leaderboard scoring systems to test the ability of different AI models to generate efficient GPU kernels.
• Stay current with industry advances.

What This Job Offers

• Job Type: Full-time
• Career Level: Mid Level
• Industry: Publishing Industries
• Number of Employees: 501-1,000 employees