Research Engineer Intern - AI Systems

About the role

• Location: Remote (Global)
• Type: Internship
• Company: Yotta Labs

Role Overview

We are seeking a highly motivated Research Engineer Intern to work on Trainium, GPU kernels, and LLM systems optimization.

Over a 12-16 week internship, you will own a well-scoped project at the intersection of AI Systems, Compiler and Runtime Optimization, Distributed Training & Inference, GPU/Accelerator Kernel Development, and Large Language Model Infrastructure - taking it from design to working, profiled code running on real hardware.

Your work will ship to production or open source and directly impact the performance of AI applications deployed on our platform. Strong interns receive return offers for full-time roles.

Responsibilities

• Implement and optimize compute kernels for Attention, GEMM, MoE, and quantization on NVIDIA, AMD, or AWS Trainium.
• Build custom operators using CUDA, Triton, ROCm/HIP, or the Neuron SDK with PyTorch/XLA.
• Profile and improve inference performance in vLLM, SGLang, and our custom runtimes - kernel fusion, scheduling, KV-cache and memory optimizations.
• Build benchmarks, chase down performance regressions, and turn profiler traces into concrete speedups.
• Ship code upstream to open-source AI infrastructure projects, with tests and documentation.

Qualifications

• Currently pursuing a BS, MS, or PhD in Computer Science, Computer Engineering, or a related field.
• Solid programming skills in Python and familiarity with C++.
• Understanding of GPU/accelerator architecture fundamentals (memory hierarchy, parallelism, occupancy) from coursework, research, or projects.
• Experience writing CUDA, Triton, ROCm/HIP, or Neuron kernels - class projects and personal projects count.
• 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 independently in a collaborative, remote environment.

Preferred Experience

• Contributions to open-source AI infra projects like vLLM, SGLang, PyTorch, or Triton.
• Familiarity with LLM inference internals - FlashAttention, PagedAttention, continuous batching, speculative decoding, MoE, or quantization.
• Experience with profiling tools (e.g. Nsight, ROCm Profiler, Neuron Profiler, or PyTorch Profiler) and performance debugging on real workloads.
• 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 to. Please include links to any relevant projects or contributions (GitHub, open-source PRs, course projects) - for internships, these matter more to us than a cover letter.