Founding Machine Learning Infrastructure Engineer

Sourced from simplify.jobs

Model AI is building the Agent Cloud, a serving and training infrastructure platform for agentic AI workloads. This founding role focuses on optimizing high-performance LLM inference across accelerators, including batching, scheduling, KV cache management, and long-context inference. You'll work with frameworks such as PyTorch, vLLM, SGLang, and TensorRT-LLM to improve throughput, latency, and cost efficiency.

• Optimize large-scale LLM inference and serving systems.
• Improve total tokens per second, decode tokens per second, latency, throughput, and cost efficiency.
• Work on serving infrastructure for open-source models across different types of accelerators.
• Improve batching, scheduling, KV cache management, memory usage, and accelerator utilization.
• Support long-context inference, including workloads targeting up to 1M context.
• Debug performance bottlenecks across model execution, runtime, networking, and infrastructure.
• Work with frameworks such as JAX/XLA, PyTorch, vLLM, SGLang, TensorRT-LLM, or related systems.
• Collaborate closely with the application team to ensure infrastructure is optimized for agentic workloads.
• Help turn research prototypes into reliable, high-performance production systems.