Machine Learning Engineer - AI Compiler Optimization

ByteDance
San Jose
Workplace: OnsiteFull timeFunction: Data Science & Machine LearningSkills: ["Performance optimization","Collaboration","Problem-solving","Technical analysis"]

Build and implement an AI compiler optimization system for recommendation models, designing full-stack optimizations across graph, operator, and memory levels. Partner with hardware and algorithm teams on hardware-software co-design, and adapt recommendation models from PyTorch for efficient import, conversion, and code generation. Focus on GPU/NPU compilation performance improvements, including loop, memory, and operator optimization, to enable low-latency inference at scale.

Loading

Loading job details...

Preparing the role view and application actions.

FursaFursa
ByteDance
ByteDance
1 month ago

Machine Learning Engineer - AI Compiler Optimization

✓ Verified Job

Canonical indexed version, validated from employer's careers page.

Source: Company careers pageValidated by: Fursa AI
Last checked: 30 days agoStatus: Live
Reposted: similar role first listed 1 month ago

Job Summary

Build and implement an AI compiler optimization system for recommendation models, designing full-stack optimizations across graph, operator, and memory levels. Partner with hardware and algorithm teams on hardware-software co-design, and adapt recommendation models from PyTorch for efficient import, conversion, and code generation. Focus on GPU/NPU compilation performance improvements, including loop, memory, and operator optimization, to enable low-latency inference at scale.
Location: San Jose
Workplace: Onsite
Employment Type: Full time
Job Function: Data Science & Machine Learning

Key Responsibilities

  • •Build and implement a compilation optimization system for a recommendation machine learning engine, including graph, operator, and memory-level optimizations.
  • •Design compilation strategies to maximize hardware compute utilization, including graph-operator fusion and automatic operator generation.
  • •Collaborate with hardware and algorithm teams on hardware-software co-design and optimize based on hardware characteristics.
  • •Adapt recommendation models from PyTorch to the engine by optimizing model import, conversion, and code generation to simplify deployment.
  • •Improve development efficiency and support compilation adaptation for recommendation model scenarios.

Key Requirements

  • •Proficiency in an AI compiler framework (Triton, MLIR, or TVM) with practical experience in customized compilation optimization and pass development.
  • •Experience with GPU/NPU compilation optimization, including loop, memory, and operator optimization, plus independent performance bottleneck analysis.
  • •Familiarity with deep learning model structures and compilation adaptation logic in PyTorch and TensorFlow, with ability to design targeted optimization solutions.
  • •Experience with architecture design of recommendation ML engines and low-latency inference optimization for large-scale recommendation systems.
  • •Experience contributing to open-source AI compiler projects (TVM, MLIR) or expertise in compilation adaptation of large recommendation models, including sparse operator generation.
Skills:Performance optimizationCollaborationProblem-solvingTechnical analysis
Tech Stack:AI compilerTritonMLIRTVMPyTorchTensorFlowGPUNPUGraph-operator fusion

Company Brief

ByteDance
Develops consumer internet and content platforms, including TikTok and other apps for short-form video, news, and entertainment. It also builds advertising, commerce, and creator tools that connect audiences, brands, and publishers across global markets.
Industry: Digital Media
Company Size: Enterprise (1,001+ employees)
Revenue: USD 1B+
Growth: Established Company
Headquarters: Beijing, China
Founded: 2012
WebsiteLinkedIn