Staff Software Engineer, ML Systems Co-Design

Автоматический пересказ вакансии на русском. Полный текст на языке оригинала и форма отклика — на сайте работодателя.

Инженер в команду MARC (Model Architecture and Realizable-performance Co-design) в Google — на стыке передовых ML-нагрузок и многолетних дорожных карт TPU. Задача — определять стратегию достижения производительности TPU для open-source ML-моделей и выстраивать методологию анализа разрывов при маппинге.

Задачи

• Формировать сертифицированные performance ladders за 6 месяцев до пилота как ориентиры для ценообразования, планирования мощностей и оптимизации XLA/ядер
• Адаптировать новые open-source модели в TPU-нативные парные варианты (например, оптимизация разреженности) для доказательства асимметричного превосходства TPU
• Обрабатывать и оптимизировать многошаговые агентные workflow, циклы рассуждений и топологии с разделением prompt/decode до выхода кремния
• Создавать эталонные реализации (PyTorch/JAX/Pallas), гарантирующие достижимость целей до пилота в реальных ограничениях компилятора
• Работать с директорским и исполнительным техническим руководством TPU Hardware, XLA Compiler и Cloud AI Systems для управления многолетними дорожными картами

Требования

• Bachelor's degree или эквивалентный практический опыт
• 8 лет опыта программирования на C++ или Python
• 5 лет опыта тестирования и запуска программных продуктов
• 5 лет опыта работы с производительностью, крупномасштабным анализом системных данных, инструментами визуализации или отладкой
• 3 года опыта в проектировании и архитектуре ПО
• Глубокое техническое владение моделированием производительности ML-нагрузок на pre-silicon системах
• Знакомство с высокоуровневыми промежуточными представлениями компиляторов (MLIR dialects, XLA, HLO)
• Опыт с ML execution frameworks (JAX, PyTorch, PyTorch/XLA) или программированием ядер ускорителей (Pallas, Triton, CUDA)
• Подтверждённый опыт архитектуры, валидации или расширения высокопроизводительных инструментов симуляции, фреймворков эмуляции или пайплайнов аналитического моделирования производительности (roofline-модели, cycle-accurate или compiler-aware симуляторы)
• Подтверждённая способность уровня L6 вести сложные многоквартальные технические инициативы через границы подразделений (аппаратный дизайн, компиляторы, AI research, облачная инфраструктура)

Условия

• США: $207000 - $300000 (USD) + 20% bonus target + equity + benefits
• Индивидуальная оплата определяется навыками, опытом и образованием/подготовкой
• В большинстве случаев собеседования проводятся очно

Навыки

• C++
• Python
• MLIR
• XLA
• HLO
• JAX
• PyTorch
• PyTorch/XLA
• Pallas
• Triton
• CUDA
• TPU
• ML workload performance modeling
• roofline models
• cycle-accurate simulators
• compiler-aware simulators
• sparsity optimization
• agentic workflows