Senior AI Infrastructure Engineer, Kodiak Robotics

Model Training

$190-260k

+ Annual Bonus + Equity

Python

C++

CUDA

PyTorch

Junior and Mid level

San Francisco Bay Area

More information about location

Office located in Mountain View, CA

Self-driving technology for long-haul trucks

Be an early applicant

Self-driving technology for long-haul trucks

201-500 employees

B2BArtificial IntelligenceLogisticsFreightSupply ChainAutomation

Be an early applicant

$190-260k

+ Annual Bonus + Equity

Python

C++

CUDA

PyTorch

Junior and Mid level

San Francisco Bay Area

More information about location

Office located in Mountain View, CA

201-500 employees

B2BArtificial IntelligenceLogisticsFreightSupply ChainAutomation

Company mission

To carry freight forward - so people, partners, and the planet thrive.

Job

Company

Role

Who you are

BS, MS, or PhD in Computer Science or a related field, and at least 2-3 years of industry experience in ML systems or infrastructure
Hands-on experience with distributed training frameworks and techniques (PyTorch DDP/FSDP, DeepSpeed, Megatron, NCCL) and a strong grasp of parallelism trade-offs
Experience building high-performance data pipelines for large-scale training, including streaming dataset formats (WebDataset, MosaicML Streaming/MDS, or similar), sharding, and storage/network-aware loading
Deep understanding of GPU performance: mixed precision, memory hierarchy, kernel fusion, profiling tools (Nsight, PyTorch Profiler), and interconnects (NVLink, InfiniBand)
Strong Python skills and proficiency in PyTorch internals; systems-level experience (C++/CUDA/Triton) a plus
Passion for building the infrastructure that lets AI for the physical world train faster, scale further, and improve continuously

What the job involves

Kodiak's AI is only as good as the speed at which we can train it
Every improvement to our models – from GigaFusionNet to large-scale world models – depends on infrastructure that turns thousands of hours of multimodal driving data into training throughput
We are looking for engineers who make model training fast: streaming massive camera, LiDAR, and radar datasets without stalling a single GPU, sharding data and models efficiently across nodes, and extracting every FLOP from the latest hardware
If you measure your impact in tokens per second and GPU utilization, this role is for you
Design high-throughput data loading and streaming systems for multimodal sensor data (camera, LiDAR, radar), including dataset formats, sharding strategies, and prefetching pipelines that keep GPUs saturated
Build and optimize distributed training infrastructure across multi-node GPU clusters, applying data, tensor, pipeline, and fully sharded (FSDP/ZeRO) parallelism to models that don't fit on a single device
Maximize utilization of modern accelerators such as NVIDIA B200s through mixed-precision training (BF16/FP8), fused kernels, memory optimization, and communication/computation overlap
Profile end-to-end training pipelines to find and eliminate bottlenecks across storage, network, CPU preprocessing, and GPU compute
Develop scalable dataset construction pipelines that convert petabytes of raw driving logs into training-ready, streamable formats
Partner with ML teams to scale new architectures from prototype to full-cluster training runs efficiently and reliably

Our take

Global supply chains rely on frictionless long-haul trucking freight networks, but the repetitive and stationary nature of the job has resulted in a lack of willing drivers. Whilst this has been partially resolved through the introduction of self-driving trucks, such systems are notoriously difficult to maintain which ultimately leads to significant downtime.

Kodiak Robotics develops self-driving technology for long-haul trucks, providing wing mirrors built with sensors called Kodiak SensorPods. These sensors provide the technology necessary for the trucks to be driven safely and autonomously and have been designed to be replaced or serviced rapidly to minimize downtime.

odiak Robotics has taken a unique approach by recognizing the high-maintenance nature of existing self-driving systems, innovating upon this through its easily replaceable SensorPods. Through this approach and several successful rounds of funding, the company aims to fill empty truck driving positions with safer AI drivers.

Steph

Company Specialist at Welcome to the Jungle

Insights

Top investors

21% employee growth in 12 months

Glassdoor (3.5)

Company

Funding (last 2 of 3 rounds)

Sep 2023

$75m

CONVERTIBLE

Nov 2021

$125m

SERIES B

Total funding: $240m

Company benefits

Remote-friendly work environment
401k
Generous PTO policy
Life insurance
Assortment of medical (PPO/HMO/HDHP), dental, vision, and FSA plans
Family-friendly company events
Dog-friendly office
Employee-driven fitness classes
Beautiful, renovated facilities in Mountain View, CA and Lancaster, TX
Free EV charging
Trivia and board game nights

Company HQ

Mountain View, CA

Leadership

Don Burnette

(CEO & Founder)

Experience as Co-founder of Otto. Software Technical Lead at Uber from 2016 to 2018.

Salary benchmarks

We don't have enough data yet to provide salary benchmarks for this role.

Submit your salary to help other candidates with crowdsourced salary estimates.

Share this job

View 29 more jobs at Kodiak Robotics