Principal Machine Learning Engineer
Full Time
|
Mountain View, CA
|
Unity
We are building the next generation of AI-driven game experiences by running generative AI models directly on players' devices—including phones, tablets, laptops, and desktops.
Our games run within a modern, browser-native runtime built on technologies such as WebGPU and WebNN. As a Principal Engineer – On-Device AI Inference & Systems, you will be the technical authority responsible for bringing state-of-the-art multimodal models—including transformers and diffusion networks—from research into production. Your work will focus on making these models fast, memory-efficient, and reliable while integrating them seamlessly into a production game engine.
This is a deeply hands-on, high-impact engineering leadership role. You will own the entire inference pipeline—from model export and optimization through low-level kernel tuning—to shipping AI-powered features that run at interactive frame rates within strict memory, power, and latency budgets.
You'll establish engineering standards, define runtime architecture, mentor senior engineers, and directly influence the AI experiences delivered to millions of players worldwide.
If you're passionate about bridging the gap between cutting-edge research and production software—and enjoy profiling, frame captures, operator fusion, and eliminating milliseconds and megabytes—this role is for you.
What You'll Be Doing
Inference & On-Device Optimization
Own the complete optimization pipeline, including:
Model export
Graph transformations
Operator fusion
Memory layout optimization
Hardware-specific kernel tuning across NPUs, mobile GPUs, and desktop GPUs
Make architectural decisions around:
INT4, INT8, and FP16 quantization
Weight sharing
Structured and unstructured pruning
Knowledge distillation
Optimize models to meet aggressive latency, memory, and power targets while maintaining quality.
Perform low-level GPU optimization by writing and tuning:
WebGPU compute shaders (WGSL)
Native compute kernels using Metal, Vulkan/SPIR-V, Direct3D 12, or CUDA
Profile performance using tools such as:
Chrome/Dawn GPU traces
PIX
Apple Instruments / Metal System Trace
Snapdragon Profiler
NVIDIA Nsight
RenderDoc
Apply advanced optimization strategies including:
Dynamic resolution
Token reduction
Cross-frame caching
Reduced-step diffusion samplers
Runtime & Systems Integration
Evaluate and implement inference runtimes including:
ONNX Runtime Web
Transformers.js
WebLLM
TensorFlow.js
CoreML
TensorFlow Lite
ExecuTorch
Extend or develop runtime infrastructure where existing solutions are insufficient.
Design the integration between the ML runtime and the game engine, including:
Thread scheduling
Memory pooling
Zero-copy buffer sharing
Frame-budget management
Architect real-time inference systems capable of processing:
Images
Text
Scene primitives
Metadata
Build supporting infrastructure for:
Model packaging
Asset pipelines
Device capability tiers
Runtime fallbacks
Telemetry
Automated on-device benchmarking
Research Productionization
Partner with research scientists to deploy new model architectures efficiently.
Provide feedback to research teams regarding:
Hardware constraints
Operator support limitations
Performance cost models
Evaluate advances in efficient inference, including:
Efficient attention
Distillation
Reduced-step diffusion
Recommend engineering investments based on measurable improvements in latency, memory usage, and power consumption.
Engineering Leadership
Lead and mentor a team of engineers.
Establish engineering best practices, code review standards, benchmarking methodologies, and performance regression gates.
Define and track KPIs for:
Latency
Quality
Memory usage
Power consumption
Collaborate with platform engineering, runtime, and product teams to align AI capabilities with device constraints and product roadmaps.
What We're Looking For
8+ years of software or machine learning engineering experience.
At least 4 years focused on:
On-device AI inference
Edge AI
Performance-critical real-time systems
Proven experience shipping transformer or diffusion models on mobile, desktop, or embedded hardware.
Hands-on experience deploying models with WebGPU using technologies such as:
ONNX Runtime Web
Transformers.js
WebLLM
TensorFlow.js
Experience writing and optimizing WGSL compute shaders.
Deep knowledge of at least one major inference runtime, including:
ONNX Runtime
ORT Web
CoreML
TensorFlow Lite
ExecuTorch
Strong understanding of:
Operator fusion
Runtime scheduling
Memory layout optimization
Low-level GPU programming experience using one or more of:
WebGPU / WGSL
Metal
Vulkan
Direct3D 12
CUDA
Experience with model optimization techniques including:
INT4, INT8, and FP16 quantization
Weight sharing
Pruning
Knowledge distillation
Familiarity with modern hardware platforms, including:
Apple Neural Engine
Qualcomm Hexagon
Adreno
ARM Mali
Apple Silicon
NVIDIA GPUs
AMD GPUs
Intel GPUs
Strong proficiency in:
TypeScript
JavaScript
WGSL
Python
Ability to understand, modify, and optimize deployed model architectures.
Demonstrated technical leadership and mentoring experience.
Nice to Have
Experience with:
On-device world models, neural rendering, or real-time generative AI, including:
NeRF
Gaussian Splatting (3DGS)
Real-time diffusion
Game engine development using:
Unity
Unreal Engine
Custom rendering engines
Graphics APIs including:
Metal
Vulkan
Direct3D
OpenGL ES
Contributions to open-source inference runtimes or GPU libraries, particularly within the WebGPU ecosystem:
Dawn
wgpu
ONNX Runtime Web
Transformers.js
WebLLM
Knowledge of emerging WebGPU features, including:
Subgroups
FP16 / shader-f16
Timestamp queries
Compiler frameworks such as:
MLIR
TVM
IREE
XLA
On-device benchmarking infrastructure, performance regression testing, and large device-farm automation.
Additional Information
International relocation assistance is not available.
Employment visa sponsorship is not available.
Compensation
Base Salary: $278,100 – $417,100 USD
Final compensation will depend on geographic location, experience, technical expertise, and overall qualifications.
In addition to base salary, this position may be eligible for:
Equity awards
Annual incentive compensation
Other company bonus programs
Benefits
Unity offers a comprehensive benefits package designed to support employee well-being and work-life balance. Benefits vary by country and employment status but may include:
Medical, life, and disability insurance
Retirement and pension plans
Employee stock ownership programs
Commute subsidies
Generous vacation and personal leave
Parental leave and family care programs
Mental health and wellness resources
Employee Resource Groups (ERGs)
Employee Assistance Program (EAP)
Professional development and training
Volunteer and donation matching programs
Office snacks and workplace amenities
About Unity
Unity is the world's leading game engine, powering interactive experiences for more than 3 billion users every month. Beyond gaming, Unity enables organizations across automotive, manufacturing, healthcare, architecture, and other industries to create, simulate, and collaborate in immersive 3D environments.
Unity is committed to fostering an inclusive workplace and is proud to be an Equal Opportunity Employer.