Member of Technical Staff, Research Engineer (GPU Performance)
Research Engineer optimizes GPU performance for large-scale AI world model training and real-time inference, focusing on custom kernels, distributed systems, and mixed-precision techniques. Requires 4+ years in ML infrastructure and performance optimization.
270k – 370k/yr
Remote4+ YOEAI Research
About the role
Responsibilities
Optimize training throughput across large GPU clusters — improving MFU through custom kernels, mixed-precision strategies (FP8, BF16), memory-efficient attention, and activation checkpointing
Design and maintain distributed training infrastructure: tensor parallelism, context parallelism, FSDP, and fault-tolerant multi-node setups
Profile and accelerate inference pipelines for real-time multimodal generation — CUDA graph compilation, KV cache optimization, operator fusion, and latency reduction
Optimize and scale training infrastructure to improve efficiency and reliability
Contribute to the entire stack, from low-level kernel optimizations to high-level model design
Requirements
4+ years of experience in systems engineering, ML infrastructure, or performance optimization for deep learning
Familiarity with GPU kernel development (CUDA, Triton, CUTLASS) and distributed systems (NCCL, collective communication, model parallelism)
Experience with ML framework internals (PyTorch, JAX) and mixed-precision / low-precision techniques (FP8, INT8)
Experience building and operating large-scale training infrastructure, including fault tolerance and cluster orchestration
Excitement about building AI that simulates the world — and making it performant enough to run in real time
Lead a world-class research team advancing LLM scaling, post-training, RL, and inference efficiency at Databricks AI. Drive research roadmap and translate breakthroughs into production systems while collaborating closely with engineering and product teams.
270k – 340k/yr
On-siteAI Research
Member of Technical Staff, Research Engineer
RunwayUnited States
Research Engineer develops and experiments with multimodal world models, focusing on data strategies, training techniques, evaluations, and production deployment for AI simulation technologies. Requires 4+ years in ML research/engineering and proficiency in PyTorch or JAX.
270k – 370k/yr
Remote4+ YOEAI Research
Staff Machine Learning Research Scientist/ Engineer, Agents
Scale AISan Francisco, CA +2
Conducts cutting-edge research on data for state-of-the-art AI agents like browser and SWE agents, develops prototypes using LLMs and frameworks like PyTorch/JAX, and publishes in top ML venues. Requires 3+ years ML experience and strong cross-functional communication.
275k – 350k/yr
On-siteAI Research
Staff Applied Scientist - Agentic Interfaces
DatadogNew York, NY
As a Staff Applied Scientist, you will define and build measurement systems for AI agent interfaces at Datadog, focusing on evaluation strategy, metric definition, and dataset creation to improve agent performance on customer workflows.
276k – 345k/yr
On-site10+ YOEAI Research
Member of Technical Staff, Applied Research Scientist
RunwayUnited States
Applied Research Scientist develops computer vision and generative models for creative AI tools, collaborating cross-functionally to productionize research. Requires 4+ years in ML/CV/graphics, strong Python/C++/PyTorch skills, and experience mentoring researchers.