Machine Learning Performance Engineer

8 Hours ago • All levels • Research Development

Job Summary

Job Description

We are seeking an engineer with expertise in low-level systems programming and optimization to join our expanding ML team. Your role will focus on enhancing the performance of our machine learning models, covering both training and inference. This includes optimizing large-scale training, achieving low-latency inference in real-time systems, and high-throughput inference in research. While basic CUDA improvements are involved, the core challenge lies in a holistic systems approach, addressing storage systems, networking, and both host and GPU-level considerations to ensure optimal efficiency and performance.
Must have:
  • Low-level GPU knowledge (PTX, SASS, warps)
  • Debugging/optimization with CUDA GDB, NSight
  • Library knowledge (Triton, CUTLASS, CUB)
  • Understanding of ML techniques and toolsets
  • Systems knowledge for end-to-end performance debugging
  • Understanding of collective algorithms (NCCL, MPI)
  • Fluency in English
Good to have:
  • Experience with Tensor Cores
  • Knowledge of Infiniband, RoCE, GPUDirect
  • Insight into CUDA graph launch latency/throughput

Job Details

We are looking for an engineer with experience in low-level systems programming and optimisation to join our growing ML team. 

Machine learning is a critical pillar of Jane Street's global business. Our ever-evolving trading environment serves as a unique, rapid-feedback platform for ML experimentation, allowing us to incorporate new ideas with relatively little friction.

Your part here is optimising the performance of our models – both training and inference. We care about efficient large-scale training, low-latency inference in real-time systems and high-throughput inference in research. Part of this is improving straightforward CUDA, but the interesting part needs a whole-systems approach, including storage systems, networking and host- and GPU-level considerations. Zooming in, we also want to ensure our platform makes sense even at the lowest level – is all that throughput actually goodput? Does loading that vector from the L2 cache really take that long?

If you’ve never thought about a career in finance, you’re in good company. Many of us were in the same position before working here. If you have a curious mind and a passion for solving interesting problems, we have a feeling you’ll fit right in. 

There’s no fixed set of skills, but here are some of the things we’re looking for:

  • An understanding of modern ML techniques and toolsets
  • The experience and systems knowledge required to debug a training run’s performance end to end
  • Low-level GPU knowledge of PTX, SASS, warps, cooperative groups, Tensor Cores and the memory hierarchy
  • Debugging and optimisation experience using tools like CUDA GDB, NSight Systems, NSight Computesight-systems and nsight-compute
  • Library knowledge of Triton, CUTLASS, CUB, Thrust, cuDNN and cuBLAS
  • Intuition about the latency and throughput characteristics of CUDA graph launch, tensor core arithmetic, warp-level synchronization and asynchronous memory loads
  • Background in Infiniband, RoCE, GPUDirect, PXN, rail optimisation and NVLink, and how to use these networking technologies to link up GPU clusters
  • An understanding of the collective algorithms supporting distributed GPU training in NCCL or MPI
  • An inventive approach and the willingness to ask hard questions about whether we're taking the right approaches and using the right tools
  • Fluency in English

 

If you're a recruiting agency and want to partner with us, please reach out to agency-partnerships@janestreet.com.

Similar Jobs

NVIDIA - Senior System Power Validation and Applications Engineer

NVIDIA

Canada (On-Site)
4 Months ago
Ion - License Key Analyst

Ion

Dublin, County Dublin, Ireland (On-Site)
3 Years ago
Relax Gaming - System Administrator

Relax Gaming

Belgrade, Serbia (Hybrid)
1 Month ago
Greenworks Sunrise Global Marketing - Field Service Technician

Greenworks Sunrise Global Marketing

Orlando, Florida, United States (On-Site)
2 Weeks ago
Capgemini - Remote Desktop Support

Capgemini

Bengaluru, Karnataka, India (Remote)
1 Month ago
Qualcomm - Machine Learning Research Engineer

Qualcomm

San Diego, California, United States (On-Site)
1 Month ago
bytedance - Research Scientist Graduate (High-Performance Computing (Inference Optimization) - Vision AI Platform)

bytedance

Seattle, Washington, United States (On-Site)
4 Months ago
playrix  - Senior Researcher

playrix

Georgia (Remote)
8 Months ago
Qingci - AI Original Artist

Qingci

Xiamen, Fujian, China (On-Site)
1 Month ago
Super.com - Machine Learning Engineer

Super.com

Bogotá, Bogota, Colombia (Hybrid)
1 Month ago

Get notifed when new similar jobs are uploaded

Similar Skill Jobs

Captions - Member of Technical Staff, Image Generation

Captions

New York, New York, United States (On-Site)
4 Days ago
Capgemini - Application Consultant

Capgemini

Bengaluru, Karnataka, India (On-Site)
1 Month ago
Ruselle Investments - Graduate Trainee, IT Infrastructure

Ruselle Investments

Mumbai, Maharashtra, India (On-Site)
2 Months ago
kaizen gaming  - Operations Data Specialist

kaizen gaming

São Paulo, Brazil (On-Site)
2 Weeks ago
Ness - Application Support Engineer

Ness

Atlanta, Georgia, United States (Hybrid)
1 Week ago
Scale AI - Senior Software Engineer

Scale AI

San Francisco, California, United States (Hybrid)
2 Months ago
Workato - AI Solutions Architect

Workato

Hyderabad, Telangana, India (On-Site)
1 Month ago
Ion - Senior Associate

Ion

London, England, United Kingdom (On-Site)
2 Months ago
bytedance - Software Engineer Intern (Privacy & Security - Global Payment)

bytedance

San Jose, California, United States (On-Site)
3 Months ago
luxsoft - Senior Python Developer

luxsoft

Poland (Remote)
1 Month ago

Get notifed when new similar jobs are uploaded

Jobs in New York, United States

Apple - SoC Silicon Debug Engineering Program Manager

Apple

Beaverton, Oregon, United States (On-Site)
1 Month ago
Uniswap Labs - Engineering Manager- Backend (Data)

Uniswap Labs

New York, United States (Hybrid)
1 Month ago
JDA - Sr. Staff Data Engineer - Snowflake Expert

JDA

Dallas, Texas, United States (Hybrid)
3 Weeks ago
Alpha Sense - Product Specialist, Financial Services

Alpha Sense

New York, United States (Hybrid)
1 Month ago
Airbyte - Support Engineering Manager

Airbyte

San Francisco, California, United States (Hybrid)
14 Hours ago
Perplexity - Developer Relations Manager

Perplexity

San Francisco, California, United States (On-Site)
1 Month ago
Whatnot - Senior IT Systems & Ops Engineer

Whatnot

San Francisco, California, United States (On-Site)
2 Months ago
Apple - Prototyping IT Administrator

Apple

Austin, Texas, United States (On-Site)
2 Weeks ago
The E.W. Scripps Company - Multimedia Journalist

The E.W. Scripps Company

Grand Rapids, Michigan, United States (On-Site)
1 Month ago
Rippling - Senior Software Engineer - Spend Management

Rippling

New York, United States (On-Site)
3 Weeks ago

Get notifed when new similar jobs are uploaded

Research Development Jobs

Vidsy - Senior AI Engineer

Vidsy

London, England, United Kingdom (Hybrid)
1 Week ago
Krafton - AI Programmer

Krafton

Montreal, Quebec, Canada (Hybrid)
5 Months ago
WongDoody - DIGITAL CREATIVE - AI WORKFLOW SPECIALIST (ALL GENDERS)

WongDoody

Stuttgart, Baden-Württemberg, Germany (Hybrid)
7 Months ago
CyberArk - Principal Program Manager - AI Excellence

CyberArk

Israel (Hybrid)
1 Month ago
bytedance - Student Researcher (Doubao (Seed) - Generative AI)

bytedance

San Jose, California, United States (Hybrid)
2 Months ago
Brillio - Agentic AI Lead

Brillio

Tampa, Florida, United States (On-Site)
1 Month ago
sony global (Games) - Artificial intelligence researcher

sony global (Games)

Shenzhen, Guangdong Province, China (On-Site)
2 Months ago
Hedra - Machine Learning Engineer (CUDA)

Hedra

San Francisco, California, United States (On-Site)
3 Months ago
Amazon games  - Senior ML Scientist, Amazon Games AI Research

Amazon games

San Diego, California, United States (On-Site)
6 Months ago
LLNL - Computational Biology - Postdoctoral Researcher

LLNL

Livermore, California, United States (On-Site)
2 Months ago

Get notifed when new similar jobs are uploaded

About The Company

Jane Street is a quantitative trading firm with offices in New York, London, Hong Kong, Singapore, and Amsterdam. We are always recruiting top candidates and we invest heavily in teaching and training. The environment at Jane Street is open, informal, intellectual, and fun. People grow into long careers here because there are always new and interesting problems to solve, systems to build, and theories to test.



Hong Kong (On-Site)

New York, United States (On-Site)

London, England, United Kingdom (On-Site)

New York, United States (On-Site)

Hong Kong (On-Site)

New York, United States (On-Site)

Hong Kong (On-Site)

London, England, United Kingdom (On-Site)

New York, United States (On-Site)

New York, United States (On-Site)

View All Jobs

Get notified when new jobs are added by Jane Street

Level Up Your Career in Game Development!

Transform Your Passion into Profession with Our Comprehensive Courses for Aspiring Game Developers.

Job Common Plug