Job details

Senior Machine Learning Engineer - Hardware Abstractions & Performance Optimization

Get a free resume review

Luma’s mission is to build multimodal AI to expand human imagination and capabilities. We believe that multimodality is critical for intelligence. To go beyond language models and build more aware, capable and useful systems, the next step function change will come from vision. So, we are working on training and scaling up multimodal foundation models for systems that can see and understand, show and explain, and eventually interact with our world to effect change.

We are looking for engineers with significant experience maintaining & designing highly efficient systems and code that can be optimized to run on multiple hardware platforms, bringing our state-of-the-art models to as many people at the best performance per dollar.

Responsibilities

Ensure efficient implementation of models & systems with a focus on designing, maintaining, and writing abstractions that scale beyond NVIDIA/CUDA hardware.
Identify and remedy efficiency bottlenecks (memory, speed, utilization, communication) by profiling and implementing high-performance PyTorch code, deferring to Triton or similar kernel-level languages as necessary.
Benchmarking our products across a variety of hardware & software to help the product team understand the optimal tradeoffs between latency, throughput and cost at various degrees of parallelism.
Work together with our partners to help them identify bottlenecks and push forward new iterations of hardware and software.
Work closely together with the rest of the research team to ensure systems are planned to be as efficient as possible from start to finish and raise potential issues for hardware integration.

Must have experience

Experience optimizing for memory, latency and throughput in Pytorch.
- Bonus: experience with non-NVIDIA systems
Experience using torch.compile / torch.XLA.
Experience benchmarking and profiling GPU & CPU code in Pytorch for optimal device utilization (examples: torch profiler, memory profilers, trace viewers, custom tooling).
Experience building tools & abstractions to ensure models run optimally on different hardware and software stacks .
Experience working with transformer models and attention implementations.
Experience with parallel inference, particularly with tensor parallelism, pipeline parallelism.

Good to have experience

Experience with high-performance Triton/CUDA and writing custom PyTorch kernels and ops. Top candidates will be able to write fused kernels for common hot paths, understand when to make use of lower level features like tensor cores or warp intrinsics, and will understand where these tools can be most impactful.
Experience writing high-performance parallel C++. Bonus if done within an ML context with PyTorch, like for data loading, data processing, inference code
Experience building inference / demo prototype code (incl. Gradio, Docker etc.)

Luma AI Glassdoor Company Review

4.4

Luma AI DE&I Review

4.3

CEO of Luma AI

Unknown name

Approve of CEO

Average salary estimate

$150000 / YEARLY (est.)

min

max

$120000K

$180000K

If an employer mentions a salary or salary range on their job, we display it as an "Employer Estimate". If a job has no salary data, Rise displays an estimate if available.

What You Should Know About Senior Machine Learning Engineer - Hardware Abstractions & Performance Optimization, Luma AI

At Luma in Palo Alto, we're on a mission to revolutionize AI by integrating multimodal approaches to expand human imagination and capabilities. As a Senior Machine Learning Engineer specializing in Hardware Abstractions & Performance Optimization, you will be at the forefront of this exciting journey. Your role will focus on creating highly efficient systems and code optimized for various hardware platforms. Your expertise will ensure our cutting-edge models perform seamlessly, maximizing performance for our users. You'll work with the latest technologies to implement and maintain robust abstractions that extend beyond NVIDIA/CUDA platforms. Your keen eye for identifying bottlenecks will drive the improvement of memory utilization, speed, and overall efficiency. Collaborating with our talented research team, you'll play a pivotal role in the integration of our models into real-world applications. If you have experience optimizing PyTorch for performance and are eager to tackle the challenges of multimodal AI, we'd love to hear from you!

Frequently Asked Questions (FAQs) for Senior Machine Learning Engineer - Hardware Abstractions & Performance Optimization Role at Luma AI

What are the responsibilities of a Senior Machine Learning Engineer at Luma?

As a Senior Machine Learning Engineer at Luma, your core responsibilities include optimizing models for performance across multiple hardware platforms, profiling for efficiency bottlenecks, and collaborating closely with product and research teams. You'll need to benchmark products to assess tradeoffs and ensure that both hardware and software components work seamlessly together.