Acerca del puesto
Anyone AI está reclutando ingenieros con experiencia en AWS Trainium / Neuron Kernel Interface (NKI) para un proyecto especializado centrado en evaluar y mejorar tareas de desarrollo de kernels para cargas de trabajo de IA.
Buscamos ingenieros con experiencia práctica construyendo u optimizando kernels NKI en hardware AWS Trainium o Inferentia2 que entiendan cómo la arquitectura de Trainium difiere de la programación tradicional en GPU.
En qué trabajarás
Revisarás y evaluarás tareas técnicas que involucran:
- Corrección de kernels NKI y patrones de desarrollo específicos de Trainium
- Migraciones de kernels de CUDA → NKI
- Optimización de rendimiento y benchmarking en Trainium
- Gestión de memoria en SBUF, PSUM y HBM
- Cómputo basado en tiles y planificación de DMA
- Corrección numérica entre plataformas entre CUDA/Triton y NKI
- Cuellos de botella de rendimiento específicos de Trainium y oportunidades de optimización
- Retroalimentación técnica y evaluación de calidad de implementaciones de kernels
El trabajo implica determinar si las implementaciones no solo son técnicamente correctas, sino también idiomáticas y optimizadas para el hardware Trainium en lugar de simplemente traducidas de enfoques basados en GPU.
Qué buscamos
- 2+ años de experiencia práctica desarrollando u optimizando kernels con la Neuron Kernel Interface (NKI)
- Experiencia trabajando con AWS Trainium y/o Inferentia2
- Sólido entendimiento de:
- Cómputo basado en tiles
- Jerarquía de memoria SBUF / PSUM / HBM
- Restricciones de dimensiones de partición
- Orquestación de DMA
- Técnicas de optimización específicas de Trainium
- Capacidad para evaluar migraciones de CUDA → NKI
- Experiencia perfilando y optimizando cargas de trabajo en Trainium
- Entendimiento de diferencias numéricas entre backends de GPU y Trainium
- Fuerte capacidad para analizar implementaciones técnicas complejas y proporcionar retroalimentación escrita clara
Deseable
- Experiencia con AWS Neuron SDK o Neuron Compiler
- Experiencia en desarrollo de kernels en CUDA o Triton
- Conocimiento de la arquitectura NeuronCore-v2
- Experiencia con cargas de trabajo FP32, BF16, FP8 e INT8
- Experiencia en benchmarking de cargas de trabajo en instancias Trn1 o Trn2
- Familiaridad con nki.language, @nki.jit o llamadas personalizadas de XLA
- Experiencia con evaluación técnica, proyectos de datos de IA/ML, RLHF o evaluación basada en rúbricas
Contratación
Tipo de trabajo: Remoto Contratación: Consultoría a tiempo parcial, basada en proyectos Enfoque: Ingeniería de kernels de AWS Trainium / NKI y evaluación técnica
Este es un encaje sólido para ingenieros que han trabajado profundamente con infraestructura de AWS Trainium y optimización de kernels de ML de bajo nivel y están interesados en aplicar esa experiencia a proyectos de IA técnicamente desafiantes.