Experto en Kernels AWS Trainium / NKI

• Más de 2 años de experiencia práctica desarrollando u optimizando kernels con Neuron Kernel Interface (NKI)

• Experiencia trabajando con AWS Trainium y/o Inferentia2

• Sólido conocimiento de:

  • Computación basada en tiles

  • Jerarquía de memoria SBUF / PSUM / HBM

  • Restricciones de dimensión de partición

  • Orquestación de DMA

  • Técnicas de optimización específicas de Trainium

• Capacidad para evaluar migraciones de CUDA → NKI

• Experiencia en perfilado y optimización de cargas de trabajo en Trainium

• Comprensión de las diferencias numéricas entre los backends de GPU y Trainium

• Gran capacidad para analizar implementaciones técnicas complejas y proporcionar feedback escrito claro

• Experiencia con el AWS Neuron SDK o Neuron Compiler

• Experiencia en desarrollo de kernels CUDA o Triton

• Conocimiento de la arquitectura NeuronCore-v2

• Experiencia con cargas de trabajo FP32, BF16, FP8 e INT8

• Experiencia realizando benchmarking de cargas de trabajo en instancias Trn1 o Trn2

• Familiaridad con nki.language, @nki.jit o llamadas personalizadas de XLA

• Experiencia en evaluación técnica, proyectos de datos de IA/ML, RLHF o evaluación basada en rúbricas

Anyone AI está contratando ingenieros de AWS Trainium / Neuron Kernel Interface (NKI) experimentados para un proyecto especializado centrado en la evaluación y mejora de tareas de desarrollo de kernels para cargas de trabajo de IA.

Buscamos ingenieros con experiencia práctica en la creación u optimización de kernels NKI en hardware AWS Trainium o Inferentia2 que comprendan en qué se diferencia la arquitectura de Trainium de la programación tradicional de GPU.

Revisarás y evaluarás tareas técnicas que implican:

• La corrección de los kernels NKI y los patrones de desarrollo específicos de Trainium

• Migraciones de kernels de CUDA → NKI

• Optimización del rendimiento y benchmarking de Trainium

• Gestión de memoria en SBUF, PSUM y HBM

• Computación basada en tiles y planificación de DMA

• Corrección numérica multiplataforma entre CUDA/Triton y NKI

• Cuellos de botella de rendimiento específicos de Trainium y oportunidades de optimización

• Feedback técnico y evaluación de calidad de las implementaciones de kernels

El trabajo consiste en determinar si las implementaciones no solo son técnicamente correctas, sino también idiomáticas y optimizadas para el hardware Trainium, en lugar de ser una simple traducción de enfoques basados en GPU.

Tipo de trabajo: Remoto
Contratación: A tiempo parcial, consultoría basada en proyectos
Enfoque: Ingeniería de kernels AWS Trainium / NKI y evaluación técnica

Este puesto es ideal para ingenieros que hayan trabajado profundamente con la infraestructura de AWS Trainium y la optimización de kernels de ML de bajo nivel y estén interesados en aplicar esa experiencia en proyectos de IA técnicamente desafiantes.