anyone-ai - AWS Trainium / NKI Kernel Expert

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