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 On You’ll review and evaluate technical tasks involving:

NKI kernel correctness and Trainium-specific development patterns

CUDA to NKI kernel migrations

Trainium performance optimization and benchmarking

Memory management across SBUF, PSUM, and HBM

Tile-based computation and DMA scheduling

Cross-platform numerical correctness between CUDA/Triton and NKI





Trainium-specific performance bottlenecks and optimization opportunities

Technical feedback and quality assessment of kernel implementations

The 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 For

2+ years of hands-on experience developing or optimizing kernels with the Neuron Kernel Interface (NKI)

Experience working with AWS Trainium and/or Inferentia2

Strong understanding of:

Tile-based computation

SBUF / PSUM / HBM memory hierarchy

Partition dimension constraints

DMA orchestration

#J-18808-Ljbffr

📌 AWS Trainium / NKI Kernel Expert (Madrid)
🏢 Anyone AI
📍 Madrid