The English-language job board for Spain

AWS Trainium / NKI Kernel Expert

Remote, Madrid·Part-time·Added 1 month ago (last checked 29 Oct)

Anyone AI

29 open roles

Apply

Our summary

Salary: US$65 an hour
Work mode: Remote (Madrid)
Contract: Fixed-term, part-time hours
English: Required
Experience: 2+ years

Visa options

The ad doesn’t mention sponsorship
Not a refusal, but not a yes. Worth asking in your first message.

Find jobs with visa support

Skills

Nice to have

CUDA

Requirements

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
- Trainium-specific optimization techniques
Ability to evaluate CUDA → NKI migrations
Experience profiling and optimizing workloads on Trainium
Understanding of numerical differences across GPU and Trainium backends
Strong ability to analyze complex technical implementations and provide clear written feedback

Nice to have

Experience with the AWS Neuron SDK or Neuron Compiler
CUDA or Triton kernel development experience
Knowledge of NeuronCore-v2 architecture
Experience with FP32, BF16, FP8, and INT8 workloads
Experience benchmarking workloads on Trn1 or Trn2 instances
Familiarity with nki.language, @nki.jit, or XLA custom calls
Experience with technical evaluation, AI/ML data projects, RLHF, or rubric-based assessment

The role

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 do

You'll review and evaluate technical tasks involving:

NKI kernel correctness and Trainium-specific development patterns
CUDA → 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.

Engagement

Work Type: Remote
Engagement: Part-time, project-based consulting
Focus: AWS Trainium / NKI kernel engineering and technical evaluation

This is a strong fit for engineers who have worked deeply with AWS Trainium infrastructure and low-level ML kernel optimization and are interested in applying that expertise to technically challenging AI projects.

About Anyone AI

Anyone AI is an EdTech company that develops AI-powered educational tools and platforms. Based in San Francisco, the company focuses on leveraging artificial intelligence to enhance learning experiences, offering products that range from AI tutoring systems to adaptive learning technologies. While the company is relatively young, it has established a presence in the AI education space, working with experts across various domains to train and refine its AI models.

Anyone AI has a significant remote-first presence in Spain, hiring experts in fields like mathematics, physics, biology, and software engineering to contribute to AI training and development. The company offers fully remote positions within Spain, making it an attractive option for international professionals seeking flexible work arrangements. With a focus on AI model training and development, Anyone AI provides opportunities for experts to apply their knowledge in a cutting-edge field, though its brand recognition is limited outside of niche AI and EdTech circles.

Industry: EdTech
Employees: 11–50
Headquarters: San Francisco, USA

Good to know if you are moving

All Spain-based roles are fully remote, offering flexibility for international professionals.
The company hires domain experts (e.g., mathematics, physics, biology) to train AI models, making it a unique opportunity for academics and specialists.
As a US-based company, salaries and contracts may be structured under Spanish employment laws, but specifics should be confirmed during the hiring process.
The company is small, so employees may have a direct impact on AI training processes and product development.
Working with a US-based company from Spain may require managing time zone differences, though remote work mitigates this.

All open roles at Anyone AI