ML Infrastructure and Performance Engineer, AI Expert (Remote)

Weekday AI

7h • $90 - $120 • Remote

About The Position

This role focuses on building, optimizing, debugging, and evaluating infrastructure for large language models and other frontier AI systems. The engineer will use ML systems expertise to create, evaluate, and improve technical tasks and solutions used to train and assess advanced AI models, helping improve their ability to solve complex ML infrastructure and performance engineering problems.

The position is fully remote and open to candidates worldwide. Candidates must be available for a full-time commitment of 40 hours per week on weekdays and must have no conflicting professional engagements. The engagement will be structured in accordance with applicable local employment and compliance requirements, either through employment or an appropriate local partner depending on the candidate’s country of residence.

Requirements

• At least 2 years of hands-on professional experience in ML systems, ML infrastructure, model serving, accelerator performance engineering, or a closely related field.
• Practical experience in at least one of the following areas: accelerator kernel programming, performance profiling, distributed workload debugging, or LLM inference serving.
• Production-level experience with JAX and/or PyTorch.
• Understanding of modern accelerator architectures, including TPUs (v4–v6e), NVIDIA H100/B200, AWS Trainium, or AMD MI300.
• Ability to analyze and optimize throughput, latency, memory utilization, and computational efficiency.
• Demonstrable professional growth and career progression in engineering or infrastructure-focused roles.
• Availability to commit at least 40 hours per week on weekdays, exclusively to this engagement without conflicting professional commitments.
• Excellent written communication skills, with the ability to explain complex technical concepts and engineering decisions clearly.

Nice To Haves

• Experience with JAX, XLA, and TPU-based development environments.
• Knowledge of custom operators, sharding, and distributed training frameworks such as pjit, shard_map, FSDP, Megatron, or MaxText.
• Familiarity with compiler and graph-level technologies such as XLA, StableHLO, torch.compile, and MLIR.
• Experience developing performance benchmarks, technical evaluation frameworks, or engineering assessment rubrics.
• Exposure to AI model evaluation, synthetic training data generation, or working with frontier AI systems.
• Experience developing or optimizing custom accelerator kernels using Pallas, Triton, CUDA, CUTLASS, or similar technologies.
• Experience with profiling and trace-analysis tools such as JAX/XLA Profiler, Kineto, torch.profiler, or NVIDIA Nsight.
• Experience with LLM inference serving technologies and techniques such as vLLM, SGLang, TensorRT-LLM, JetStream, KV cache optimization, paged attention, or continuous batching.

Responsibilities

• Design challenging technical tasks covering accelerator kernel development, performance profiling, distributed workload debugging, and high-throughput inference serving.
• Develop complex and adversarial engineering problems to test advanced AI models, including tasks involving accelerator programming and optimization.
• Create accurate, technically sound reference solutions that demonstrate best practices in ML systems engineering.
• Evaluate AI-generated solutions for correctness, efficiency, scalability, and practical applicability.
• Develop detailed technical guidelines and assessment rubrics for kernel optimization, profiler interpretation, distributed systems, and inference performance.
• Analyze performance bottlenecks across ML infrastructure, compilers, accelerators, and model-serving systems.
• Provide clear, structured technical feedback to help improve AI model capabilities in ML systems and infrastructure engineering.
• Collaborate with researchers and fellow subject matter experts to maintain the quality, consistency, and accuracy of AI training and evaluation data.

What This Job Offers

• Job Type: Full-time
• Career Level: Mid Level
• Education Level: No Education Listed