Senior Manager, Software Engineering - JAX

JAX ist das Framework der Wahl für die ehrgeizigsten KI-Forschungen, und die Rechenzentren, auf denen es läuft, entwickeln sich schnell weiter. NVIDIA-Supercomputer werden heterogener und kombinieren GPUs, CPUs und LPUs, und JAX muss sich weiterentwickeln, um auf allen gut zu funktionieren. NVIDIA sucht einen Senior Engineering Manager, um die JAX-Strategie von NVIDIA zu definieren und voranzutreiben, mehrere Teams und Arbeitsströme zu koordinieren, um sicherzustellen, dass JAX eine Spitzenleistung über die Plattform hinweg liefert, von Einzelbeschleuniger-Workloads bis hin zu Hunderttausenden von Geräten auf den größten Supercomputern, die jemals gebaut wurden. Ein kritischer Teil dieser Rolle besteht darin, sicherzustellen, dass JAX ab dem ersten Tag auf Architekturen der nächsten Generation, einschließlich Rubin Ultra, LPUs und den CPUs, die sie antreiben, bereit ist. Dies erfordert, mit den Entwicklungen im Bereich KI Schritt zu halten. Das bedeutet, aufkommende Bedürfnisse in den Bereichen Training, Post-Training, Inferenz und Robotik zu unterstützen und neue Hardwarefähigkeiten mit KI-Trends zu verbinden.

Was Sie tun werden:

• Die Strategie für den Ingenieureinsatz im gesamten JAX-Ökosystem vorantreiben – einschließlich JAX-Kern, XLA, Shardy, MPMD und anwendungsbezogenen Codebasen wie MaxText – und priorisieren, wo die Beiträge von NVIDIA maximale Auswirkungen auf Leistung, Akzeptanz und Einfluss auf das Ökosystem haben.
• Teamarbeit über organisatorische Grenzen hinweg fördern, um die KI-Plattform von NVIDIA über Compiler (XLA, TileIR, CUDA), Laufzeit, Netzwerk- und Hardwareteams zu gestalten.
• Tiefe Partnerschaften mit wichtigen Open-Source-Projekten aufbauen, insbesondere mit dem JAX-Team von Google – unsere Partner unterstützen und gleichzeitig die Interessen von NVIDIA bei der Gestaltung der Projektentwicklung vertreten.
• Prozesse entwerfen und klare Erfolgskriterien definieren, um die Teams auf Ergebnisse auszurichten.
• Personalpläne erstellen, die Teamkapazität überprüfen und strategische Einstellungsentscheidungen treffen, um aktuellen und zukünftigen Bedürfnissen voraus zu sein.
• Eine leistungsstarke Ingenieurorganisation leiten, wachsen und betreuen – technische Führungskräfte entwickeln und Raum für Ihr Team schaffen, um zu experimentieren, Prototypen zu erstellen und innovative Systeme zu entwickeln, die Produkte und Prozesse verbessern.

Was wir sehen müssen:

• MS oder PhD in Informatik, Computertechnik oder einem verwandten Bereich (oder gleichwertige Erfahrung).
• Über 10 Jahre Erfahrung in der Softwareentwicklung, davon mehr als 4 Jahre in der technischen Leitung oder im Engineering-Management.
• Starker technischer Hintergrund in Systemsoftware, Compilern, Hochleistungsrechnen oder KI-Frameworks.
• Nachgewiesene Erfahrung in der Leitung von strategischen Projekten, die organisatorische Grenzen überschreiten, Konsens bilden, ohne direkte Autorität Einfluss nehmen und die Kommunikation an das Publikum anpassen können, von Ingenieuren bis hin zu Führungskräften.
• Erfahrung in der Verantwortung für Ergebnisse über mehrere Projekte und Teams hinweg, Abwägungen ausbalancieren und Fortschritte mit unvollständigen Informationen vorantreiben.
• Erfahrung in der Identifizierung von Lücken im Entwicklungszyklus und der Definition von Arbeitsabläufen über Teams und Organisationen hinweg zur Verbesserung der Effizienz.
• Erfahrung in der Leitung von Teams, die zu großangelegten Projekten mit externen Partnern beitragen.

Wie Sie sich von der Masse abheben können:

• Praktische Erfahrung mit JAX und seinen Systemen wie XLA, HLO, PJRT, Shardy oder JAXs Tracing- und Transformationssystem.
• Erfahrung in der Leitung von Teams, die an Compiler- oder Laufzeit-Infrastrukturen für Deep-Learning-Frameworks (JAX, PyTorch, TensorFlow oder gleichwertig) arbeiten.
• Vertrautheit mit den Komponenten der KI-Plattform von NVIDIA: GPU-Programmierung, CUDA, cuDNN, NCCL, TensorRT, Transformer Engine und heterogene Rechenziele (GPU, CPU, LPU).
• Erfahrung in der Durchführung von Leistungsoptimierungsmaßnahmen für großangelegte verteilte KI-Trainings-, Post-Training (RLHF, Alignment), Inferenz- oder Robotik-Workloads (Multi-Node, Multi-GPU).
• Erfahrung in der Durchführung von Framework- oder Softwareeinführungen auf neuen Hardwarearchitekturen, Koordination zwischen Silizium-, Treiber-, Compiler- und Framework-Teams.