Research Engineer - AI Systems

About the role

• Location: Remote (Global)
• Type: Full-time
• Company: Yotta Labs

Role Overview

We are seeking a highly motivated AI Systems Research Engineer specializing in Trainium, GPU kernels, and LLM systems optimization. You will work at the intersection of AI Systems, Compiler and Runtime Optimization, Distributed Training & Inference, GPU/Accelerator Kernel Development, and Large Language Model Infrastructure. Your work will directly impact the scalability and performance of AI applications deployed on our platform.

Responsibilities

• Design and implement high-performance kernels for Attention, MoE, GEMM, collective communication, and quantization.
• Optimize kernels for NVIDIA, AMD, and AWS Trainium.
• Develop custom operators and graph optimizations using Neuron SDK, PyTorch/XLA, Torch Dynamo, and Neuron Compiler.
• Improve performance of vLLM, SGLang, TensorRT-LLM, and custom inference runtimes.
• Design scalable distributed training and inference solutions across thousands of accelerators.
• Contribute to open-source projects, publish technical findings and engage with the developer community.

Qualifications

• Proficiency in AI programming languages such as Python and C++.
• Deep understanding of GPU architecture and performance optimization.
• Experience with CUDA, Triton, ROCm/HIP, or AWS Neuron.
• Strong understanding of AI frameworks (e.g., PyTorch, Dynamo, LMCache), model architectures and profiling tools (e.g. Nsight, ROCm Profiler, or Neuron Profiler).
• Strong problem-solving skills and the ability to work in a collaborative, remote environment.
• A background in computer science, engineering, or a related field is preferred.

Preferred Experience

• Contributions to open-source AI infra projects like vLLM, SGLang, PyTorch, or Triton.
• Experience with with FlashAttention, PagedAttention, MoE, RLHF, or distributed AI systems.
• Publications in top-tier conferences like MLSys, OSDI, SOSP, NSDI, SC, HPCA, or ISCA

How to Apply

Interested candidates should apply directly or send their resume and a brief cover letter to. Please include links to any relevant projects or contributions.