Machine Learning Research Engineer

Job summary

Responsibilities

Etched is building AI chips that are hard-coded for individual model architectures. Our first product (Sohu) only supports transformers, but has an order of magnitude more throughput and lower latency than a B200.

• Propose and conduct novel research to achieve results on Sohu that are unviable on GPUs
• Translate core mathematical operations from the most popular Transformer-based models into maximally performant instruction sequences for Sohu
• Develop deep architectural knowledge informing best-in-the-world software performance on Sohu HW, collaborating with HW architects and designers.
• Co-design and finetune emerging model architectures for highest efficiency on Sohu
• Guide and contribute to the Sohu software stack, performance characterization tools, and runtime abstractions by implementing frontier models using Python and Rust.

What they're looking for

• An ML Research background with interests in HW co-design
• Experience with Python, Pytorch, and / or JAX
• Familiarity with transformer model architectures and/or inference serving stacks (vLLM, SGLang, etc.) and/or experience working in distributed inference/training environments
• Experience working cross-functionally in diverse software and hardware organizations

Tech stack

• Transformers
• Python
• Rust
• PyTorch
• JAX
• vLLM
• CUDA