MLOps Engineer, LLM Performance (Kernels, Profiling, Serving)

Role Overview

Help advance foundational large language models by creating and evaluating high-quality ML systems training data. This hands-on role focuses on accelerator kernels, performance profiling, distributed-workload debugging, and high-throughput inference serving, with particular value placed on JAX, XLA, and TPU expertise.

Key Responsibilities

• Design challenging and adversarial ML systems tasks covering kernels, profiling, debugging, and inference serving, including problems that reveal frontier-model failures in accelerator code such as TPU Pallas kernels.
• Write accurate, well-structured solutions for ML systems tasks used in AI model training and evaluation.
• Help research and engineering teams address knowledge gaps and improve model performance across ML systems, compilers, and training infrastructure.
• Evaluate technical tasks and solutions, providing clear written feedback suitable for expert review.
• Create detailed guidelines and rubrics for kernel optimization, profiler-output interpretation, distributed-systems reasoning, and serving throughput and latency trade-offs.
• Collaborate with subject matter experts to maintain consistent, accurate training data.

Qualifications

• At least 2 years of hands-on professional experience in ML systems, ML infrastructure, model serving, or accelerator performance engineering.
• Hands-on systems engineering background rather than an applied modeling or data science focus.
• Practical experience in at least one of the following areas, with experience across multiple areas preferred: custom kernel development or optimization using Pallas, Triton, CUDA, or CUTLASS; profiling and trace analysis using JAX/XLA Profiler, Kineto, torch.profiler, or Nsight; debugging distributed or accelerator-bound workloads; or large-scale LLM serving with vLLM, SGLang, TensorRT-LLM, JetStream, KV cache, paged attention, or continuous batching.
• Production experience with JAX and/or PyTorch. Depth in custom operators, sharding, distributed training using pjit, shard_map, FSDP, Megatron, or MaxText, or compiler and graph-level work with XLA, StableHLO, torch.compile, or MLIR is strongly valued.
• Experience with accelerators such as TPU v4 through v6e, H100, B200, Trainium, or MI300, including the ability to assess throughput, latency, and memory trade-offs.
• Demonstrable career progression, strong written communication, and the ability to explain complex technical decisions clearly.

Work Terms

• Fully remote and open to candidates worldwide.
• Full-time, 40 hours per week during weekdays. Candidates must be able to commit without conflicting engagements.
• Engagement may be administered directly or through a local employer-of-record partner, depending on the candidate''s location.

Compensation

Hourly compensation of $90 to $120.

Eligibility

Applicants are considered without discrimination based on race, religion, color, national origin, sex, pregnancy, childbirth, reproductive health decisions, related medical conditions, sexual orientation, gender identity or expression, age, protected-veteran status, disability, genetic information, political views or activity, or any other legally protected characteristic.