Informazioni sul ruolo
Anyone AI sta reclutando ingegneri esperti di AWS Trainium / Neuron Kernel Interface (NKI) per un progetto specializzato focalizzato sulla valutazione e sul miglioramento di attività di sviluppo di kernel per carichi di lavoro AI.
Cerchiamo ingegneri con esperienza pratica nella creazione o nell'ottimizzazione di kernel NKI su hardware AWS Trainium o Inferentia2, che comprendano come l'architettura di Trainium differisca dalla programmazione tradizionale su GPU.
Su cosa lavorerai
Esaminerai e valuterai attività tecniche che riguardano:
- Correttezza dei kernel NKI e pattern di sviluppo specifici di Trainium
- Migrazioni di kernel da CUDA a NKI
- Ottimizzazione delle prestazioni e benchmarking su Trainium
- Gestione della memoria tra SBUF, PSUM e HBM
- Calcolo basato su tile e pianificazione DMA
- Correttezza numerica cross-platform tra CUDA/Triton e NKI
- Colli di bottiglia prestazionali e opportunità di ottimizzazione specifici di Trainium
- Feedback tecnico e valutazione della qualità delle implementazioni dei kernel
Il lavoro consiste nel determinare se le implementazioni non siano solo tecnicamente corrette, ma anche idiomatiche e ottimizzate per l'hardware Trainium, anziché semplicemente tradotte da approcci basati su GPU.
Cosa cerchiamo
- 2+ anni di esperienza pratica nello sviluppo o nell'ottimizzazione di kernel con il Neuron Kernel Interface (NKI)
- Esperienza con AWS Trainium e/o Inferentia2
- Solida comprensione di:
- Calcolo basato su tile
- Gerarchia di memoria SBUF / PSUM / HBM
- Vincoli sulla dimensione delle partizioni
- Orchestrazione DMA
- Tecniche di ottimizzazione specifiche di Trainium
- Capacità di valutare migrazioni da CUDA a NKI
- Esperienza nel profiling e nell'ottimizzazione di carichi di lavoro su Trainium
- Comprensione delle differenze numeriche tra backend GPU e Trainium
- Forte capacità di analizzare implementazioni tecniche complesse e fornire feedback scritto chiaro
Plus graditi
- Esperienza con AWS Neuron SDK o Neuron Compiler
- Esperienza nello sviluppo di kernel CUDA o Triton
- Conoscenza dell'architettura NeuronCore-v2
- Esperienza con carichi di lavoro FP32, BF16, FP8 e INT8
- Esperienza nel benchmarking di carichi di lavoro su istanze Trn1 o Trn2
- Familiarità con nki.language, @nki.jit o XLA custom calls
- Esperienza con valutazione tecnica, progetti data AI/ML, RLHF o valutazione basata su rubriche
Modalità di collaborazione
Tipo di lavoro: Remoto Collaborazione: Consulenza part-time, basata su progetto Focus: Ingegneria e valutazione tecnica di kernel AWS Trainium / NKI
Questa è un'ottima opportunità per ingegneri che hanno lavorato a fondo con l'infrastruttura AWS Trainium e con l'ottimizzazione di kernel ML di basso livello e sono interessati ad applicare questa competenza a progetti AI tecnicamente impegnativi.