Anyone AI is recruiting experiencedAWS Trainium / Neuron Kernel Interface (NKI) engineersfor a specialized project focused on evaluating and improving kernel development tasks for AI workloads.We're looking for engineers with hands-on experience building or optimizingNKI kernels on AWS Trainium or Inferentia2 hardwarewho understand how Trainium's architecture differs from traditional GPU programming.What You'll Work OnYou’ll review and evaluate technical tasks involving:NKI kernel correctness and Trainium-specific development patternsCUDA to NKI kernel migrationsTrainium performance optimization and benchmarkingMemory management across SBUF, PSUM,



and HBMTile-based computation and DMA schedulingCross-platform numerical correctness between CUDA/Triton and NKITrainium-specific performance bottlenecks and optimization opportunitiesTechnical feedback and quality assessment of kernel implementationsThe work involves determining whether implementations are not only technically correct, but alsoidiomatic and optimized for Trainium hardwarerather than simply translated from GPU-based approaches.What We're Looking For2+ years of hands-on experience developing or optimizing kernels with the Neuron Kernel Interface (NKI)Experience working with AWS Trainium and/or Inferentia2Strong understanding of:Tile-based computationSBUF / PSUM / HBM memory hierarchyPartition dimension constraintsDMA orchestration#J-18808-Ljbffr

📌 Aws Trainium / Nki Kernel Expert (Madrid)
🏢 Anyone AI
📍 Madrid