Anyone AI is recruiting experienced AWS Trainium / Neuron Kernel Interface (NKI) engineers for 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 optimizing NKI kernels on AWS Trainium or Inferentia2 hardware who 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 also idiomatic and optimized for Trainium hardware rather 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