ML Systems Engineer: Training & Inference Infra

NextToken AI Inc.; Mountain View (CA); On-site; USD 180,000 - 250,000; Full time

Job summary

NextToken AI Inc. is building scalable ML infrastructure to improve model performance in production environments. You’ll work across training systems, model serving, and distributed data pipelines, shaping reliability and efficiency for researchers and engineers.

You’ll design and implement end-to-end experiments, optimize resource usage, and diagnose bottlenecks across GPUs, memory, and scheduling. You’ll collaborate with researchers and customers to turn ideas into dependable training and

Qualifications

• Experience in ML infrastructure and training/serving systems.
• Ability to explain architectural decisions and measurable improvements.
• Knowledge of PyTorch or JAX and profiling workloads.
• Strong distributed systems fundamentals: scheduling, storage, recovery.
• Experience with distributed training frameworks and container orchestration.

Responsibilities

• Develop infrastructure for post-training experiments.
• Improve resource allocation and parallel execution.
• Make model serving efficient and diagnose bottlenecks.
• Build scalable data generation and evaluation pipelines.
• Contribute to benchmarks, observability, and automation.

Skills

ML infrastructure

Training systems

Model serving

Distributed data pipelines

Tools

CUDA

Triton

PyTorch

JAX

Job description

NextToken AI Inc. is building scalable ML infrastructure to improve model performance in production environments. You’ll work across training systems, model serving, and distributed data pipelines, shaping reliability and efficiency for researchers and engineers.

You’ll design and implement end-to-end experiments, optimize resource usage, and diagnose bottlenecks across GPUs, memory, and scheduling. You’ll collaborate with researchers and customers to turn ideas into dependable training and