AIML - Machine Learning Engineer (On-Device Machine Learning), Machine Learning Platform and Infrastructure

$135,400 - $250,600/Yr

Apple - Seattle, WA

Full-time - Mid Level

Seattle, WA

Computer and Electronic Product Manufacturing

About the position

The AIML on-device Machine Learning group at Apple is seeking a tenured software engineer to implement innovative machine learning models, including LLMs, Diffusion, and GenAI, on Apple's powerful machine learning hardware. This role focuses on enhancing user experiences through machine learning while ensuring privacy and performance. The engineer will collaborate with both first and third-party clients to optimize model execution on Apple silicon, contributing to the development of features for the on-device inference stack.

Responsibilities

• Build features for the on-device inference stack to support new and existing customers with evolving machine learning models.
• Interact with first and third-party clients to determine functional and performance requirements.
• Collaborate with backend teams to map client requirements to functional execution.
• Perform analysis to optimize memory and performance from the underlying hardware.
• Drive structural improvements in the on-device machine learning stack.

Requirements

• 5+ years of proven programming skills using standard ML tools such as C/C++, CUDA/Metal, PyTorch, TensorFlow.
• Experience building APIs and/or core components of ML frameworks with strong attention to detail.
• Ability to iterate on ideas and work with various partners across the stack, including apps, compilation, hardware architecture, and power/performance analysis.
• Excellent problem-solving, critical thinking, communication, and collaboration skills.

Nice-to-haves

• Knowledge of existing ML framework implementations (Jax, PyTorch, TensorFlow).
• Knowledge of code execution down to the hardware/accelerator level.
• Hands-on experience with LLVM, compiler technologies, and optimization techniques like quantization and sparsity-induction.
• M.S. or B.S. or equivalent experience in Computer Science.

Benefits

• Comprehensive medical and dental coverage
• Retirement benefits
• Discounted products and free services
• Reimbursement for certain educational expenses, including tuition
• Discretionary bonuses or commission payments
• Relocation assistance