# AWS Trainium / NKI Kernel Expert at Anyone AI

> Open job at Anyone AI in Madrid, Spain. Workplace: Remote (Spain).

Source: Spain Dev Jobs, https://spain-devjobs.com/jobs/anyone-ai-aws-trainium-nki-kernel-expert-md17wxn
Posted 15 September 2026. Last checked on the employer's career site: 10 November 2026.

## Facts

- Company: [Anyone AI](https://spain-devjobs.com/companies/anyone-ai)
- Location: Madrid, Spain
- Workplace: Remote (Spain)
- Employment type: Contract
- Spanish: Not required
- Skills: AWS
- Not mentioned in the job post: salary, visa sponsorship, relocation, where you can apply from

## Job description

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 → 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
  - 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

### 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.

## How to apply

Apply on the job page, https://spain-devjobs.com/jobs/anyone-ai-aws-trainium-nki-kernel-expert-md17wxn: its Apply button opens the employer's own application page.

## More jobs from Anyone AI

- [AI Coding Evaluator](https://spain-devjobs.com/jobs/anyone-ai-ai-coding-evaluator-h1wcy7m): Anyone AI, Madrid, Spain. Remote (Spain). Posted 12 October 2026.
- [AI Agent Evaluator](https://spain-devjobs.com/jobs/anyone-ai-ai-agent-evaluator-8fy2v3j): Anyone AI, Madrid, Spain. Remote (Spain). Posted 12 October 2026.
- [Senior Software Engineer – Open Source & SWE-Bench Evaluation](https://spain-devjobs.com/jobs/anyone-ai-senior-software-engineer-open-source-swe-bench-evaluation-03cccub): Anyone AI, Madrid, Spain. Remote (Spain). Posted 10 October 2026.
- [Application Security Engineer – CVE & Vulnerability Research](https://spain-devjobs.com/jobs/anyone-ai-application-security-engineer-cve-vulnerability-research-ztjwm9q): Anyone AI, Madrid, Spain. Remote (Spain). Posted 10 October 2026.
- [GPU Kernel Engineer – CUDA, Triton & Accelerator Performance](https://spain-devjobs.com/jobs/anyone-ai-gpu-kernel-engineer-cuda-triton-accelerator-performance-e7px63v): Anyone AI, Madrid, Spain. Remote (Spain). Posted 15 September 2026.