Software Engineer II - AI/ML, AWS Neuron

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

Инженер в команду Inference Enablement and Acceleration (Annapurna Labs, AWS) для разработки и оптимизации распределённого инференса в AWS Neuron SDK на ускорителях Inferentia и Trainium. Работа на стыке ML, низкоуровневой оптимизации и системной архитектуры.

Задачи

• Проектировать, разрабатывать и оптимизировать ML-модели и фреймворки для развёртывания на кастомных ML-ускорителях
• Участвовать во всех этапах жизненного цикла ML-системы: проектирование распределённой архитектуры, реализация, профилирование производительности, аппаратно-специфичные оптимизации, тестирование и продакшн-деплой
• Строить инфраструктуру для систематического анализа и онбординга множества моделей с разной архитектурой
• Проектировать и реализовывать высокопроизводительные ядра и фичи для ML-операций с использованием архитектуры Neuron
• Анализировать и оптимизировать производительность на системном уровне для нескольких поколений железа Neuron
• Проводить детальный анализ производительности профилировщиками и устранять узкие места
• Реализовывать оптимизации: fusion, sharding, tiling, scheduling
• Проводить комплексное тестирование, включая unit- и end-to-end тесты моделей, с непрерывным деплоем через пайплайны
• Работать напрямую с клиентами над включением и оптимизацией их ML-моделей на ускорителях AWS
• Сотрудничать с командами компилятора, рантайма, фреймворков и железа

Требования

• 3+ года профессионального опыта разработки ПО (не стажировки)
• Степень бакалавра по computer science или эквивалент
• 3+ года опыта проектирования или архитектуры новых и существующих систем (паттерны проектирования, надёжность, масштабирование)
• Понимание основ машинного обучения и LLM, их архитектуры, жизненного цикла обучения и инференса, а также опыт оптимизации выполнения моделей
• Опыт разработки на C++ и Python (обязателен опыт хотя бы на одном языке)
• Сильное понимание производительности систем, управления памятью и принципов параллельных вычислений
• Навыки отладки, профилирования и применения лучших практик разработки в крупномасштабных системах

Условия

• Базовая зарплата 143 700 – 194 400 USD в год
• Пакет включает sign-on выплаты и RSU
• Медицинская страховка (medical, dental, vision, рецепты, Basic Life & AD&D, опция Supplemental life), EAP, поддержка ментального здоровья, Medical Advice Line, Flexible Spending Accounts, компенсация усыновления и суррогатного материнства
• 401(k) с matching
• Оплачиваемые отпуска и parental leave
• Работа в небольшой agile-команде в среде, похожей на стартап
• Менторство, код-ревью и поддержка карьерного роста

Навыки

• Python
• C++
• PyTorch
• JAX
• LLM
• Machine Learning
• distributed inference
• CUDA
• CUTLASS
• FlashInfer
• Triton
• vLLM
• SGLang
• TensorRT
• JIT compilation
• AOT tracing
• profiling
• parallel computing
• computer architecture
• AWS Trainium
• AWS Inferentia
• Neuron SDK