À propos du poste
Anyone AI recrute des ingénieurs expérimentés en AWS Trainium / Neuron Kernel Interface (NKI) pour un projet spécialisé axé sur l'évaluation et l'amélioration des tâches de développement de noyaux pour les charges de travail d'IA.
Nous recherchons des ingénieurs ayant une expérience pratique de la construction ou de l'optimisation de noyaux NKI sur du matériel AWS Trainium ou Inferentia2, et qui comprennent en quoi l'architecture de Trainium diffère de la programmation GPU traditionnelle.
Ce sur quoi vous travaillerez
Vous examinerez et évaluerez des tâches techniques portant sur :
- la correction des noyaux NKI et les modèles de développement spécifiques à Trainium
- les migrations de noyaux CUDA → NKI
- l'optimisation des performances et le benchmarking sur Trainium
- la gestion de la mémoire entre SBUF, PSUM et HBM
- le calcul par tuiles et l'ordonnancement DMA
- la correction numérique multiplateforme entre CUDA/Triton et NKI
- les goulets d'étranglement de performance spécifiques à Trainium et les opportunités d'optimisation
- les retours techniques et l'évaluation de la qualité des implémentations de noyaux
Le travail consiste à déterminer si les implémentations sont non seulement techniquement correctes, mais aussi idiomatiques et optimisées pour le matériel Trainium, plutôt que simplement transposées à partir d'approches basées sur GPU.
Ce que nous recherchons
- 2+ ans d'expérience pratique du développement ou de l'optimisation de noyaux avec le Neuron Kernel Interface (NKI)
- Expérience de travail avec AWS Trainium et/ou Inferentia2
- Solide compréhension de :
- le calcul par tuiles
- la hiérarchie mémoire SBUF / PSUM / HBM
- les contraintes de dimension de partition
- l'orchestration DMA
- les techniques d'optimisation spécifiques à Trainium
- Capacité à évaluer les migrations CUDA → NKI
- Expérience du profilage et de l'optimisation de charges de travail sur Trainium
- Compréhension des différences numériques entre les backends GPU et Trainium
- Forte capacité à analyser des implémentations techniques complexes et à fournir des retours écrits clairs
Atouts appréciés
- Expérience avec l'AWS Neuron SDK ou le Neuron Compiler
- Expérience du développement de noyaux CUDA ou Triton
- Connaissance de l'architecture NeuronCore-v2
- Expérience avec les charges de travail FP32, BF16, FP8 et INT8
- Expérience du benchmarking de charges de travail sur des instances Trn1 ou Trn2
- Familiarité avec nki.language, @nki.jit ou les appels personnalisés XLA
- Expérience de l'évaluation technique, des projets de données IA/ML, du RLHF ou de l'évaluation basée sur des rubriques
Engagement
Type de travail : À distance Engagement : Conseil à temps partiel, basé sur des projets Focus : Ingénierie de noyaux AWS Trainium / NKI et évaluation technique
Ce poste convient parfaitement aux ingénieurs qui ont travaillé en profondeur avec l'infrastructure AWS Trainium et l'optimisation de noyaux ML de bas niveau, et qui souhaitent appliquer cette expertise à des projets d'IA techniquement exigeants.