Member of Technical Staff - Research Engineer

Job summary

Responsibilities

In this role, you will be embedded in production training and help where the hardest systems and performance problems arise: attention performance, custom kernels, low-precision training, profiling, memory behavior, data movement, distributed training stability, and throughput regressions.

• Improve the performance, reliability, and numerical stability of production training runs for large multimodal generative models
• Profile full training steps across model code, attention, kernels, data loading, encoders, communication, optimizer steps, checkpointing, and memory pressure
• Implement and validate GPU-level optimizations: fused kernels, attention paths, low-precision matmuls, quantization kernels, CUDA/Triton/CuTe/CUTLASS experiments, and no-compile…
• Push lower-precision training forward, including FP8 / MXFP8 / FP4-style paths, weight and activation quantization, accumulation choices, convergence risk, and quality tradeoffs…
• Work with researchers to translate architecture changes into efficient training implementations, and help distinguish real model-quality progress from changes that only look good…

What they're looking for

• Experience working deeply on large-scale training systems, ideally as part of a training group working closely with researchers
• Strong PyTorch fluency, including comfort reading and modifying low-level training code rather than only using high-level APIs
• Experience with distributed training concepts such as FSDP, tensor/model/context/sequence parallelism, activation checkpointing, NCCL, and overlapping compute and communication
• Hands-on experience improving training throughput, memory footprint, or stability in real training runs
• Experience profiling GPU workloads with tools like Nsight Systems, Nsight Compute, torch profiler, trace viewers, or custom telemetry

Tech stack

• CUDA
• Triton
• PyTorch